add namespace support for async storage (#4775)

rochdev · web-flow · commit 83e11a3e137e · 2024-11-13T15:43:52.000-05:00
diff --git a/packages/datadog-core/index.js b/packages/datadog-core/index.js
@@ -1,7 +1,5 @@
 'use strict'
 
-const { AsyncLocalStorage } = require('async_hooks')
-
-const storage = new AsyncLocalStorage()
+const storage = require('./src/storage')
 
 module.exports = { storage }
diff --git a/packages/datadog-core/src/storage.js b/packages/datadog-core/src/storage.js
@@ -0,0 +1,21 @@
+'use strict'
+
+const { AsyncLocalStorage } = require('async_hooks')
+
+const storages = Object.create(null)
+const legacyStorage = new AsyncLocalStorage()
+
+const storage = function (namespace) {
+  if (!storages[namespace]) {
+    storages[namespace] = new AsyncLocalStorage()
+  }
+  return storages[namespace]
+}
+
+storage.disable = legacyStorage.disable.bind(legacyStorage)
+storage.enterWith = legacyStorage.enterWith.bind(legacyStorage)
+storage.exit = legacyStorage.exit.bind(legacyStorage)
+storage.getStore = legacyStorage.getStore.bind(legacyStorage)
+storage.run = legacyStorage.run.bind(legacyStorage)
+
+module.exports = storage
diff --git a/packages/datadog-core/test/storage.spec.js b/packages/datadog-core/test/storage.spec.js
@@ -0,0 +1,50 @@
+'use strict'
+
+require('../../dd-trace/test/setup/tap')
+
+const { expect } = require('chai')
+const storage = require('../src/storage')
+
+describe('storage', () => {
+  let testStorage
+  let testStorage2
+
+  beforeEach(() => {
+    testStorage = storage('test')
+    testStorage2 = storage('test2')
+  })
+
+  afterEach(() => {
+    testStorage.enterWith(undefined)
+    testStorage2.enterWith(undefined)
+  })
+
+  it('should enter a store', done => {
+    const store = 'foo'
+
+    testStorage.enterWith(store)
+
+    setImmediate(() => {
+      expect(testStorage.getStore()).to.equal(store)
+      done()
+    })
+  })
+
+  it('should enter stores by namespace', done => {
+    const store = 'foo'
+    const store2 = 'bar'
+
+    testStorage.enterWith(store)
+    testStorage2.enterWith(store2)
+
+    setImmediate(() => {
+      expect(testStorage.getStore()).to.equal(store)
+      expect(testStorage2.getStore()).to.equal(store2)
+      done()
+    })
+  })
+
+  it('should return the same storage for a namespace', () => {
+    expect(storage('test')).to.equal(testStorage)
+  })
+})
