remove destroy hook by using current store for the sets

monti-apm · Nov 30, 2023 · ad2daf4 · ad2daf4
1 parent bd46806
commit ad2daf4
Showing 1 changed file with 17 additions and 27 deletions.
diff --git a/src/await-detector.ts b/src/await-detector.ts
@@ -17,8 +17,6 @@ export class AwaitDetector {
 
   start = Date.now();
 
-  asyncFunctions = new Set();
-  awaits = new Set();
   afterAwaits = new Map();
   ignoreNextPromise = 0;
 
@@ -48,7 +46,6 @@ export class AwaitDetector {
     this.hook = createHook({
       init: this.init.bind(this),
       before: this.before.bind(this),
-      destroy: this.destroy.bind(this),
       promiseResolve: this.promiseResolve.bind(this),
     });
 
@@ -130,22 +127,24 @@ export class AwaitDetector {
 
     if (AwaitDetector.IgnoreStorage.getStore()) return;
 
+    const store = AwaitDetector.Storage.getStore() as any;
+
     const isAsyncFunction = triggerAsyncId === executionAsyncId();
 
     if (isAsyncFunction) {
-      this.asyncFunctions.add(asyncId);
+      store.asyncFunctions.add(asyncId);
       this.log(`${type}(${asyncId}): async function start`);
       return;
     }
 
-    if (this.asyncFunctions.has(triggerAsyncId)) {
+    if (store.asyncFunctions.has(triggerAsyncId)) {
       this.onAwaitStart(asyncId, triggerAsyncId);
-      this.awaits.add(asyncId);
+      store.awaits.add(asyncId);
       this.awaitData.set(asyncId, [triggerAsyncId]);
       this.log(
         `${type}(${asyncId}): await start - async function: ${triggerAsyncId}`,
       );
-    } else if (this.awaits.has(triggerAsyncId)) {
+    } else if (store.awaits.has(triggerAsyncId)) {
       this.afterAwaits.set(asyncId, triggerAsyncId);
     }
   }
@@ -155,6 +154,8 @@ export class AwaitDetector {
       return;
     }
 
+    const store = AwaitDetector.Storage.getStore() as any;
+
     if (this.afterAwaits.has(asyncId)) {
       const awaitAsyncId = this.afterAwaits.get(asyncId);
 
@@ -164,7 +165,7 @@ export class AwaitDetector {
       this.awaitData.delete(awaitAsyncId);
       // Awaited a thenable or non-promise value
       this.log(`await end:  ${this.afterAwaits.get(asyncId)} (A)`);
-    } else if (this.awaits.has(asyncId)) {
+    } else if (store.awaits.has(asyncId)) {
       if (!this.awaitData.has(asyncId)) return;
       const [triggerAsyncId] = this.awaitData.get(asyncId) as [number];
       this.onAwaitEnd(asyncId, triggerAsyncId);
@@ -174,30 +175,17 @@ export class AwaitDetector {
     }
   }
 
-  destroy(asyncId: number) {
-    if (!this.isWithinContext()) {
-      return;
-    }
-
-    if (this.asyncFunctions.has(asyncId)) {
-      this.asyncFunctions.delete(asyncId);
-      return;
-    }
-
-    if (this.awaits.has(asyncId)) {
-      this.awaits.delete(asyncId);
-    }
-  }
-
   promiseResolve(asyncId: number) {
     if (!this.isWithinContext()) {
       return;
     }
 
-    if (this.asyncFunctions.has(asyncId)) {
-      this.asyncFunctions.delete(asyncId);
-    } else if (this.awaits.has(asyncId)) {
-      this.awaits.delete(asyncId); // Added later
+    const store = AwaitDetector.Storage.getStore() as any;
+
+    if (store.asyncFunctions.has(asyncId)) {
+      store.asyncFunctions.delete(asyncId);
+    } else if (store.awaits.has(asyncId)) {
+      store.awaits.delete(asyncId); // Added later
     } else if (this.afterAwaits.has(asyncId)) {
       this.log(`promise resolve: ${asyncId} (C)`);
       this.afterAwaits.delete(asyncId); // Added later
@@ -208,6 +196,8 @@ export class AwaitDetector {
     return AwaitDetector.Storage.run(
       {
         [AwaitDetector.Symbol]: this,
+        asyncFunctions: new Set(),
+        awaits: new Set(),
       },
       callback,
     );