use weakmap to avoid references from node to datadog stores (#4953)

packages/datadog-core/src/storage.js

@@ -2,12 +2,47 @@
 
 const { AsyncLocalStorage } = require('async_hooks')
 
+class DatadogStorage {
+  constructor () {
+    this._storage = new AsyncLocalStorage()
+  }
+
+  disable () {
+    this._storage.disable()
+  }
+
+  enterWith (store) {
+    const handle = {}
+    stores.set(handle, store)
+    this._storage.enterWith(handle)
+  }
+
+  exit (callback, ...args) {
+    this._storage.exit(callback, ...args)
+  }
+
+  getStore () {
+    const handle = this._storage.getStore()
+    return stores.get(handle)
+  }
+
+  run (store, fn, ...args) {
+    const prior = this._storage.getStore()
+    this.enterWith(store)
+    try {
+      return Reflect.apply(fn, null, args)
+    } finally {
+      this._storage.enterWith(prior)
+    }
+  }
+}
+
 const storages = Object.create(null)
-const legacyStorage = new AsyncLocalStorage()
+const legacyStorage = new DatadogStorage()
 
 const storage = function (namespace) {
   if (!storages[namespace]) {
-    storages[namespace] = new AsyncLocalStorage()
+    storages[namespace] = new DatadogStorage()
   }
   return storages[namespace]
 }
@@ -18,4 +53,6 @@ storage.exit = legacyStorage.exit.bind(legacyStorage)
 storage.getStore = legacyStorage.getStore.bind(legacyStorage)
 storage.run = legacyStorage.run.bind(legacyStorage)
 
+const stores = new WeakMap()
+
 module.exports = storage

packages/datadog-core/test/storage.spec.js

@@ -3,6 +3,7 @@
 require('../../dd-trace/test/setup/tap')
 
 const { expect } = require('chai')
+const { executionAsyncResource } = require('async_hooks')
 const storage = require('../src/storage')
 
 describe('storage', () => {
@@ -47,4 +48,16 @@ describe('storage', () => {
   it('should return the same storage for a namespace', () => {
     expect(storage('test')).to.equal(testStorage)
   })
+
+  it('should not have its store referenced by the underlying async resource', () => {
+    const resource = executionAsyncResource()
+
+    testStorage.enterWith({ internal: 'internal' })
+
+    for (const sym of Object.getOwnPropertySymbols(resource)) {
+      if (sym.toString() === 'Symbol(kResourceStore)' && resource[sym]) {
+        expect(resource[sym]).to.not.have.property('internal')
+      }
+    }
+  })
 })

packages/dd-trace/src/llmobs/storage.js

@@ -1,7 +1,6 @@
 'use strict'
 
-// TODO: remove this and use namespaced storage once available
-const { AsyncLocalStorage } = require('async_hooks')
-const storage = new AsyncLocalStorage()
+const { storage: createStorage } = require('../../../datadog-core')
+const storage = createStorage('llmobs')
 
 module.exports = { storage }

packages/dd-trace/src/opentelemetry/context_manager.js

@@ -1,6 +1,6 @@
 'use strict'
 
-const { AsyncLocalStorage } = require('async_hooks')
+const { storage } = require('../../../datadog-core')
 const { trace, ROOT_CONTEXT } = require('@opentelemetry/api')
 const DataDogSpanContext = require('../opentracing/span_context')
 
@@ -9,7 +9,7 @@ const tracer = require('../../')
 
 class ContextManager {
   constructor () {
-    this._store = new AsyncLocalStorage()
+    this._store = storage('opentelemetry')
   }
 
   active () {

packages/dd-trace/src/profiling/profilers/event_plugins/event.js

@@ -1,4 +1,4 @@
-const { AsyncLocalStorage } = require('async_hooks')
+const { storage } = require('../../../../../datadog-core')
 const TracingPlugin = require('../../../plugins/tracing')
 const { performance } = require('perf_hooks')
 
@@ -8,7 +8,7 @@ class EventPlugin extends TracingPlugin {
   constructor (eventHandler) {
     super()
     this.eventHandler = eventHandler
-    this.store = new AsyncLocalStorage()
+    this.store = storage('profiling')
     this.entryType = this.constructor.entryType
   }