From edb8135a7960815e8b112aae3a2c3c34e9b3d812 Mon Sep 17 00:00:00 2001
From: Kelvin Jin <kelvinjin@google.com>
Date: Thu, 19 Apr 2018 08:45:22 -0700
Subject: [PATCH] fix: restore context when a function run with a given context
 throws (#727)

PR-URL: #727
---
 src/cls/async-hooks.ts | 18 ++++++++++--------
 test/test-cls.ts       | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/cls/async-hooks.ts b/src/cls/async-hooks.ts
index 3d23d357d..59551ef69 100644
--- 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
index d749f198c..a3071f983 100644
--- 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');