diff --git a/packages/datadog-core/src/storage.js b/packages/datadog-core/src/storage.js index d28420ed259..15c9fff239c 100644 --- a/packages/datadog-core/src/storage.js +++ b/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 diff --git a/packages/datadog-core/test/storage.spec.js b/packages/datadog-core/test/storage.spec.js index 89839f1fca3..e5bca4e7d5d 100644 --- a/packages/datadog-core/test/storage.spec.js +++ b/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') + } + } + }) }) diff --git a/packages/dd-trace/src/llmobs/storage.js b/packages/dd-trace/src/llmobs/storage.js index 1362aaf966e..82202c18174 100644 --- a/packages/dd-trace/src/llmobs/storage.js +++ b/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 } diff --git a/packages/dd-trace/src/opentelemetry/context_manager.js b/packages/dd-trace/src/opentelemetry/context_manager.js index fba84eef9f4..430626bbd7e 100644 --- a/packages/dd-trace/src/opentelemetry/context_manager.js +++ b/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 () { diff --git a/packages/dd-trace/src/profiling/profilers/event_plugins/event.js b/packages/dd-trace/src/profiling/profilers/event_plugins/event.js index f47a3468f78..73d3214e231 100644 --- a/packages/dd-trace/src/profiling/profilers/event_plugins/event.js +++ b/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 }