StorageManager supports centralized exception handling

compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java

@@ -4,6 +4,8 @@ import jet.Function0;
 import jet.Function1;
 import jet.Unit;
 import junit.framework.TestCase;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.util.ReenteringLazyValueComputationException;
 
 import java.util.ArrayList;
@@ -426,6 +428,60 @@ public class StorageManagerTest extends TestCase {
         assertEquals(2, c.getCount());
     }
 
+    // ExceptionHandlingStrategy
+
+    public void testExceptionHandlingStrategyForLazyValues() throws Exception {
+        class RethrownException extends RuntimeException {}
+
+        LockBasedStorageManager m = LockBasedStorageManager.createWithExceptionHandling(new LockBasedStorageManager.ExceptionHandlingStrategy() {
+            @NotNull
+            @Override
+            public RuntimeException handleException(@NotNull Throwable throwable) {
+                throw new RethrownException();
+            }
+        });
+        try {
+            m.createLazyValue(
+                    new Function0<Object>() {
+                        @Nullable
+                        @Override
+                        public Object invoke() {
+                            throw new RuntimeException();
+                        }
+                    }
+            ).invoke();
+            fail("Exception should have occurred");
+        }
+        catch (RethrownException ignored) {
+        }
+    }
+
+    public void testExceptionHandlingStrategyForMemoizedFunctions() throws Exception {
+        class RethrownException extends RuntimeException {}
+
+        LockBasedStorageManager m = LockBasedStorageManager.createWithExceptionHandling(new LockBasedStorageManager.ExceptionHandlingStrategy() {
+            @NotNull
+            @Override
+            public RuntimeException handleException(@NotNull Throwable throwable) {
+                throw new RethrownException();
+            }
+        });
+        try {
+            m.createMemoizedFunction(
+                    new Function1<Object, Object>() {
+                        @Nullable
+                        @Override
+                        public Object invoke(@Nullable Object o) {
+                            throw new RuntimeException();
+                        }
+                    }
+            ).invoke("");
+            fail("Exception should have occurred");
+        }
+        catch (RethrownException ignored) {
+        }
+    }
+
     // toString()
 
     public void testToString() throws Exception {

core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java

@@ -31,7 +31,27 @@ import java.util.concurrent.locks.ReentrantLock;
 
 public class LockBasedStorageManager implements StorageManager {
 
-    public static final StorageManager NO_LOCKS = new LockBasedStorageManager("NO_LOCKS", NoLock.INSTANCE) {
+    public interface ExceptionHandlingStrategy {
+        ExceptionHandlingStrategy THROW = new ExceptionHandlingStrategy() {
+            @NotNull
+            @Override
+            public RuntimeException handleException(@NotNull Throwable throwable) {
+                throw ExceptionUtils.rethrow(throwable);
+            }
+        };
+
+        /*
+         * The signature of this method is a trick: it is used as
+         *
+         *     throw strategy.handleException(...)
+         *
+         * most implementations of this method throw exceptions themselves, so it does not matter what they return
+         */
+        @NotNull
+        RuntimeException handleException(@NotNull Throwable throwable);
+    }
+
+    public static final StorageManager NO_LOCKS = new LockBasedStorageManager("NO_LOCKS", ExceptionHandlingStrategy.THROW, NoLock.INSTANCE) {
         @NotNull
         @Override
         protected <T> RecursionDetectedResult<T> recursionDetectedDefault() {
@@ -39,16 +59,30 @@ public class LockBasedStorageManager implements StorageManager {
         }
     };
 
+    public static LockBasedStorageManager createWithExceptionHandling(@NotNull ExceptionHandlingStrategy exceptionHandlingStrategy) {
+        return new LockBasedStorageManager(exceptionHandlingStrategy);
+    }
+
     protected final Lock lock;
+    private final ExceptionHandlingStrategy exceptionHandlingStrategy;
     private final String debugText;
 
-    private LockBasedStorageManager(@NotNull String debugText, @NotNull Lock lock) {
+    private LockBasedStorageManager(
+            @NotNull String debugText,
+            @NotNull ExceptionHandlingStrategy exceptionHandlingStrategy,
+            @NotNull Lock lock
+    ) {
         this.lock = lock;
+        this.exceptionHandlingStrategy = exceptionHandlingStrategy;
         this.debugText = debugText;
     }
 
     public LockBasedStorageManager() {
-        this(getPointOfConstruction(), new ReentrantLock());
+        this(getPointOfConstruction(), ExceptionHandlingStrategy.THROW, new ReentrantLock());
+    }
+
+    private LockBasedStorageManager(@NotNull ExceptionHandlingStrategy exceptionHandlingStrategy) {
+        this(getPointOfConstruction(), exceptionHandlingStrategy, new ReentrantLock());
     }
 
     private static String getPointOfConstruction() {
@@ -276,7 +310,7 @@ public class LockBasedStorageManager implements StorageManager {
                         // Store only if it's a genuine result, not something thrown through recursionDetected()
                         value = WrappedValues.escapeThrowable(throwable);
                     }
-                    throw ExceptionUtils.rethrow(throwable);
+                    throw exceptionHandlingStrategy.handleException(throwable);
                 }
             }
             finally {
@@ -352,12 +386,12 @@ public class LockBasedStorageManager implements StorageManager {
                     return typedValue;
                 }
                 catch (Throwable throwable) {
-                    if (throwable == error) throw error;
+                    if (throwable == error) throw exceptionHandlingStrategy.handleException(throwable);
 
                     Object oldValue = cache.put(input, WrappedValues.escapeThrowable(throwable));
                     assert oldValue == NotValue.COMPUTING : "Race condition detected on input " + input + ". Old value is " + oldValue;
 
-                    throw ExceptionUtils.rethrow(throwable);
+                    throw exceptionHandlingStrategy.handleException(throwable);
                 }
             }
             finally {