fix: avoid memory leaks due to undisposed promise resources (#885)

kjin · web-flow · commit 8454389beaf7 · 2018-11-02T14:59:58.000-07:00
diff --git a/src/cls/async-hooks.ts b/src/cls/async-hooks.ts
@@ -58,13 +58,17 @@ export class AsyncHooksCLS<Context extends {}> implements CLS<Context> {
         // init is called when a new AsyncResource is created. We want code
         // that runs within the scope of this new AsyncResource to see the same
         // context as its "parent" AsyncResource. The criteria for the parent
-        // depends on the type of the AsyncResource.
+        // depends on the type of the AsyncResource. (If the parent doesn't have
+        // an associated context, don't do anything.)
         if (type === 'PROMISE') {
           // Opt not to use the trigger ID for Promises, as this causes context
           // confusion in applications using async/await.
           // Instead, use the ID of the AsyncResource in whose scope we are
           // currently running.
-          this.contexts[id] = this.contexts[this.ah.executionAsyncId()];
+          const currentId = this.ah.executionAsyncId();
+          if (this.contexts[currentId] !== undefined) {
+            this.contexts[id] = this.contexts[currentId];
+          }
         } else {
           // Use the trigger ID for any other type. In Node core, this is
           // usually equal the ID of the AsyncResource in whose scope we are
@@ -75,17 +79,31 @@ export class AsyncHooksCLS<Context extends {}> implements CLS<Context> {
           // AsyncResource API, because users of that API can specify their own
           // trigger ID. In this case, we choose to respect the user's
           // selection.
-          this.contexts[id] = this.contexts[triggerId];
+          if (this.contexts[triggerId] !== undefined) {
+            this.contexts[id] = this.contexts[triggerId];
+          }
         }
-        // Note that this function always assigns values in this.contexts to
-        // values under other keys, which may or may not be undefined. Consumers
-        // of the CLS API will get the sentinel (default) value if they query
-        // the current context when it is stored as undefined.
+        // Note that this function only assigns values in this.contexts to
+        // values under other keys; it never generates new context values.
+        // Consumers of the CLS API will get the sentinel (default) value if
+        // they query the current context when it is not stored in
+        // this.contexts.
       },
       destroy: (id: number) => {
         // destroy is called when the AsyncResource is no longer used, so also
         // delete its entry in the map.
         delete this.contexts[id];
+      },
+      promiseResolve: (id: number) => {
+        // Promise async resources may not get their destroy hook entered for
+        // a long time, so we listen on promiseResolve hooks as well. If this
+        // event is emitted, the async scope of the Promise will not be entered
+        // again, so it is generally safe to delete its entry in the map. (There
+        // is a possibility that a future async resource will directly reference
+        // this Promise as its trigger parent -- in this case, it will have
+        // the wrong parent, but this is still better than a potential memory
+        // leak.)
+        delete this.contexts[id];
       }
     });
   }
diff --git a/test/test-cls-ah.ts b/test/test-cls-ah.ts
@@ -90,7 +90,7 @@ maybeSkip(describe)('AsyncHooks-based CLS', () => {
     }, 'modified');
   });
 
-  describe('Using AsyncResource API', () => {
+  describe('Compatibility with AsyncResource API', () => {
     it('Supports context propagation without trigger ID', async () => {
       let res!: asyncHooksModule.AsyncResource;
       await cls.runWithContext(async () => {
@@ -115,4 +115,111 @@ maybeSkip(describe)('AsyncHooks-based CLS', () => {
       });
     });
   });
+
+  describe('Memory consumption with Promises', () => {
+    const createdPromiseIDs: number[] = [];
+    let hook: asyncHooksModule.AsyncHook;
+
+    before(() => {
+      hook = asyncHooks
+                 .createHook({
+                   init: (uid: number, type: string) => {
+                     if (type === 'PROMISE') {
+                       createdPromiseIDs.push(uid);
+                     }
+                   }
+                 })
+                 .enable();
+    });
+
+    after(() => {
+      hook.disable();
+    });
+
+    const testCases:
+        Array<{description: string; skip?: boolean; fn: () => {}}> = [
+          {description: 'a no-op async function', fn: async () => {}}, {
+            description: 'an async function that throws',
+            fn: async () => {
+              throw new Error();
+            }
+          },
+          {
+            description: 'an async function that awaits a rejected value',
+            fn: async () => {
+              await new Promise(reject => setImmediate(reject));
+            }
+          },
+          {
+            description: 'an async function with awaited values',
+            fn: async () => {
+              await 0;
+              await new Promise(resolve => resolve());
+              await new Promise(resolve => setImmediate(resolve));
+            }
+          },
+          {
+            description: 'an async function that awaits another async function',
+            fn: async () => {
+              await (async () => {
+                await Promise.resolve();
+              })();
+            }
+          },
+          {
+            description: 'a plain function that returns a Promise',
+            fn: () => Promise.resolve()
+          },
+          {
+            description:
+                'a plain function that returns a Promise that will reject',
+            fn: () => Promise.reject()
+          },
+          {
+            description: 'an async function with spread args',
+            // TODO(kjin): A possible bug in exists that causes an extra Promise
+            // async resource to be initialized when an async function with
+            // spread args is invoked. promiseResolve is not called for this
+            // async resource. Fix this bug and then remove this skip directive.
+            skip: true,
+            fn: async (...args: number[]) => args
+          }
+        ];
+
+    for (const testCase of testCases) {
+      const skipIfTestSpecifies = !!testCase.skip ? it.skip : it;
+      skipIfTestSpecifies(
+          `Doesn't retain stale references when running ${
+              testCase.description} in a context`,
+          async () => {
+            createdPromiseIDs.length = 0;
+            try {
+              // Run the test function in a new context.
+              await cls.runWithContext(testCase.fn, 'will-be-stale');
+            } catch (e) {
+              // Ignore errors; they aren't important for this test.
+            } finally {
+              // At this point, Promises created from invoking the test function
+              // should have had either their destroy or promiseResolve hook
+              // called. We observe this by iterating through the Promises that
+              // were created in this context, and checking to see that getting
+              // the current context in the scope of an async resource that
+              // references any of these Promises as its trigger parent doesn't
+              // yield the stale context value from before. The only way this is
+              // possible is if the CLS implementation internally kept a stale
+              // reference to a context-local value keyed on the ID of a PROMISE
+              // resource that should have been disposed.
+              const stalePromiseIDs = createdPromiseIDs.filter((id) => {
+                const a = new AsyncResource('test', id);
+                const result = a.runInAsyncScope(() => {
+                  return cls.getContext() === 'will-be-stale';
+                });
+                a.emitDestroy();
+                return result;
+              });
+              assert.strictEqual(stalePromiseIDs.length, 0);
+            }
+          });
+    }
+  });
 });
