fix: restore context when a function run with a given context throws (#727)

kjin · web-flow · commit edb8135a7960 · 2018-04-19T08:45:22.000-07:00
PR-URL: #727
diff --git a/src/cls/async-hooks.ts b/src/cls/async-hooks.ts
@@ -86,10 +86,11 @@ export class AsyncHooksCLS<Context extends {}> implements CLS<Context> {
   runWithNewContext<T>(fn: Func<T>): T {
     const oldContext = this.currentContext.value;
     this.currentContext.value = this.defaultContext;
-    // TODO(kjin) Handle the case where fn throws. Context: PR #708
-    const res = fn();
-    this.currentContext.value = oldContext;
-    return res;
+    try {
+      return fn();
+    } finally {
+      this.currentContext.value = oldContext;
+    }
   }
 
   bindWithCurrentContext<T>(fn: Func<T>): Func<T> {
@@ -101,10 +102,11 @@ export class AsyncHooksCLS<Context extends {}> implements CLS<Context> {
     const contextWrapper: ContextWrapped<Func<T>> = function(this: {}) {
       const oldContext = current.value;
       current.value = boundContext;
-      // TODO(kjin) Handle the case where fn throws. Context: PR #708
-      const res = fn.apply(this, arguments) as T;
-      current.value = oldContext;
-      return res;
+      try {
+        return fn.apply(this, arguments) as T;
+      } finally {
+        current.value = oldContext;
+      }
     };
     contextWrapper[WRAPPED] = true;
     Object.defineProperty(contextWrapper, 'length', {
diff --git a/test/test-cls.ts b/test/test-cls.ts
@@ -130,6 +130,32 @@ describe('Continuation-Local Storage', () => {
           });
         });
 
+        it('Corrects context when function run with new context throws', () => {
+          try {
+            c.runWithNewContext(() => {
+              c.setContext('modified');
+              throw new Error();
+            });
+          } catch (e) {
+            assert.strictEqual(c.getContext(), 'default');
+          }
+        });
+
+        it('Corrects context when function bound to a context throws', () => {
+          let runLater = () => {
+            c.setContext('modified');
+            throw new Error();
+          };
+          c.runWithNewContext(() => {
+            runLater = c.bindWithCurrentContext(runLater);
+          });
+          try {
+            runLater();
+          } catch (e) {
+            assert.strictEqual(c.getContext(), 'default');
+          }
+        });
+
         it('Can be used to patch event emitters to propagate context', () => {
           const ee = new EventEmitter();
           assert.strictEqual(c.getContext(), 'default');
