From 7f93d36b79cf2be26778e22631356c9367e3735f Mon Sep 17 00:00:00 2001 From: Bryan English Date: Mon, 7 Oct 2024 14:31:53 -0400 Subject: [PATCH] use AsyncLocalStorage instead of our home-grown solutions (#4201) * use AsyncLocalStorage instead of our home-grown solutions The comment in the file that selected a storage implementation suggested just using AsyncLocalStorage once it supports triggerAsyncResource(). That said, literally zero of our code uses triggerAsyncResource(), so this is assumed to be historical and no longer relevant. Switching to stock AsyncLocalStorage will enable the usage of TracingChannel in the future. * self-contain profiling's AsyncLocalStorage channel usage * remove flag detection --- packages/datadog-core/index.js | 4 +- .../src/storage/async_resource.js | 108 ------------ packages/datadog-core/src/storage/index.js | 5 - packages/datadog-core/test/setup.js | 8 - .../test/storage/async_resource.spec.js | 20 --- packages/datadog-core/test/storage/test.js | 160 ------------------ .../dd-trace/src/profiling/profilers/wall.js | 40 +++++ 7 files changed, 42 insertions(+), 303 deletions(-) delete mode 100644 packages/datadog-core/src/storage/async_resource.js delete mode 100644 packages/datadog-core/src/storage/index.js delete mode 100644 packages/datadog-core/test/setup.js delete mode 100644 packages/datadog-core/test/storage/async_resource.spec.js delete mode 100644 packages/datadog-core/test/storage/test.js diff --git a/packages/datadog-core/index.js b/packages/datadog-core/index.js index 72b0403aa75..9819b32f3ba 100644 --- a/packages/datadog-core/index.js +++ b/packages/datadog-core/index.js @@ -1,7 +1,7 @@ 'use strict' -const LocalStorage = require('./src/storage') +const { AsyncLocalStorage } = require('async_hooks') -const storage = new LocalStorage() +const storage = new AsyncLocalStorage() module.exports = { storage } diff --git a/packages/datadog-core/src/storage/async_resource.js b/packages/datadog-core/src/storage/async_resource.js deleted file mode 100644 index 4738845e415..00000000000 --- a/packages/datadog-core/src/storage/async_resource.js +++ /dev/null @@ -1,108 +0,0 @@ -'use strict' - -const { createHook, executionAsyncResource } = require('async_hooks') -const { channel } = require('dc-polyfill') - -const beforeCh = channel('dd-trace:storage:before') -const afterCh = channel('dd-trace:storage:after') -const enterCh = channel('dd-trace:storage:enter') - -let PrivateSymbol = Symbol -function makePrivateSymbol () { - // eslint-disable-next-line no-new-func - PrivateSymbol = new Function('name', 'return %CreatePrivateSymbol(name)') -} - -try { - makePrivateSymbol() -} catch (e) { - try { - const v8 = require('v8') - v8.setFlagsFromString('--allow-natives-syntax') - makePrivateSymbol() - v8.setFlagsFromString('--no-allow-natives-syntax') - // eslint-disable-next-line no-empty - } catch (e) {} -} - -class AsyncResourceStorage { - constructor () { - this._ddResourceStore = PrivateSymbol('ddResourceStore') - this._enabled = false - this._hook = createHook(this._createHook()) - } - - disable () { - if (!this._enabled) return - - this._hook.disable() - this._enabled = false - } - - getStore () { - if (!this._enabled) return - - const resource = this._executionAsyncResource() - - return resource[this._ddResourceStore] - } - - enterWith (store) { - this._enable() - - const resource = this._executionAsyncResource() - - resource[this._ddResourceStore] = store - enterCh.publish() - } - - run (store, callback, ...args) { - this._enable() - - const resource = this._executionAsyncResource() - const oldStore = resource[this._ddResourceStore] - - resource[this._ddResourceStore] = store - enterCh.publish() - - try { - return callback(...args) - } finally { - resource[this._ddResourceStore] = oldStore - enterCh.publish() - } - } - - _createHook () { - return { - init: this._init.bind(this), - before () { - beforeCh.publish() - }, - after () { - afterCh.publish() - } - } - } - - _enable () { - if (this._enabled) return - - this._enabled = true - this._hook.enable() - } - - _init (asyncId, type, triggerAsyncId, resource) { - const currentResource = this._executionAsyncResource() - - if (Object.prototype.hasOwnProperty.call(currentResource, this._ddResourceStore)) { - resource[this._ddResourceStore] = currentResource[this._ddResourceStore] - } - } - - _executionAsyncResource () { - return executionAsyncResource() || {} - } -} - -module.exports = AsyncResourceStorage diff --git a/packages/datadog-core/src/storage/index.js b/packages/datadog-core/src/storage/index.js deleted file mode 100644 index e522e61ced2..00000000000 --- a/packages/datadog-core/src/storage/index.js +++ /dev/null @@ -1,5 +0,0 @@ -'use strict' - -// TODO: default to AsyncLocalStorage when it supports triggerAsyncResource - -module.exports = require('./async_resource') diff --git a/packages/datadog-core/test/setup.js b/packages/datadog-core/test/setup.js deleted file mode 100644 index 2f8af45cdd2..00000000000 --- a/packages/datadog-core/test/setup.js +++ /dev/null @@ -1,8 +0,0 @@ -'use strict' - -require('tap').mochaGlobals() - -const chai = require('chai') -const sinonChai = require('sinon-chai') - -chai.use(sinonChai) diff --git a/packages/datadog-core/test/storage/async_resource.spec.js b/packages/datadog-core/test/storage/async_resource.spec.js deleted file mode 100644 index ce19b216260..00000000000 --- a/packages/datadog-core/test/storage/async_resource.spec.js +++ /dev/null @@ -1,20 +0,0 @@ -'use strict' - -require('../setup') - -const StorageBackend = require('../../src/storage/async_resource') -const testStorage = require('./test') - -describe('storage/async_resource', () => { - let storage - - beforeEach(() => { - storage = new StorageBackend() - }) - - afterEach(() => { - storage.disable() - }) - - testStorage(() => storage) -}) diff --git a/packages/datadog-core/test/storage/test.js b/packages/datadog-core/test/storage/test.js deleted file mode 100644 index 0f69a43d9f0..00000000000 --- a/packages/datadog-core/test/storage/test.js +++ /dev/null @@ -1,160 +0,0 @@ -'use strict' - -const { expect } = require('chai') -const { inspect } = require('util') -const { - AsyncResource, - executionAsyncId, - executionAsyncResource -} = require('async_hooks') - -module.exports = factory => { - let storage - let store - - beforeEach(() => { - storage = factory() - store = {} - }) - - describe('getStore()', () => { - it('should return undefined by default', () => { - expect(storage.getStore()).to.be.undefined - }) - }) - - describe('run()', () => { - it('should return the value returned by the callback', () => { - expect(storage.run(store, () => 'test')).to.equal('test') - }) - - it('should preserve the surrounding scope', () => { - expect(storage.getStore()).to.be.undefined - - storage.run(store, () => {}) - - expect(storage.getStore()).to.be.undefined - }) - - it('should run the span on the current scope', () => { - expect(storage.getStore()).to.be.undefined - - storage.run(store, () => { - expect(storage.getStore()).to.equal(store) - }) - - expect(storage.getStore()).to.be.undefined - }) - - it('should persist through setTimeout', done => { - storage.run(store, () => { - setTimeout(() => { - expect(storage.getStore()).to.equal(store) - done() - }, 0) - }) - }) - - it('should persist through setImmediate', done => { - storage.run(store, () => { - setImmediate(() => { - expect(storage.getStore()).to.equal(store) - done() - }, 0) - }) - }) - - it('should persist through setInterval', done => { - storage.run(store, () => { - let shouldReturn = false - - const timer = setInterval(() => { - expect(storage.getStore()).to.equal(store) - - if (shouldReturn) { - clearInterval(timer) - return done() - } - - shouldReturn = true - }, 0) - }) - }) - - it('should persist through process.nextTick', done => { - storage.run(store, () => { - process.nextTick(() => { - expect(storage.getStore()).to.equal(store) - done() - }, 0) - }) - }) - - it('should persist through promises', () => { - const promise = Promise.resolve() - - return storage.run(store, () => { - return promise.then(() => { - expect(storage.getStore()).to.equal(store) - }) - }) - }) - - it('should handle concurrency', done => { - storage.run(store, () => { - setImmediate(() => { - expect(storage.getStore()).to.equal(store) - done() - }) - }) - - storage.run(store, () => {}) - }) - - it('should not break propagation for nested resources', done => { - storage.run(store, () => { - const asyncResource = new AsyncResource( - 'TEST', { triggerAsyncId: executionAsyncId(), requireManualDestroy: false } - ) - - asyncResource.runInAsyncScope(() => {}) - - expect(storage.getStore()).to.equal(store) - - done() - }) - }) - - it('should not log ddResourceStore contents', done => { - function getKeys (output) { - return output.split('\n').slice(1, -1).map(line => { - return line.split(':').map(v => v.trim())[0] - }) - } - - setImmediate(() => { - const withoutStore = getKeys(inspect(executionAsyncResource(), { depth: 0 })) - storage.run(store, () => { - setImmediate(() => { - const withStore = getKeys(inspect(executionAsyncResource(), { depth: 0 })) - expect(withStore).to.deep.equal(withoutStore) - done() - }) - }) - }) - }) - }) - - describe('enterWith()', () => { - it('should transition into the context for the remainder of the current execution', () => { - const newStore = {} - - storage.run(store, () => { - storage.enterWith(newStore) - expect(storage.getStore()).to.equal(newStore) - }) - - expect(storage.getStore()).to.be.undefined - }) - }) -} diff --git a/packages/dd-trace/src/profiling/profilers/wall.js b/packages/dd-trace/src/profiling/profilers/wall.js index ee23b1145b0..39af4ca2bfc 100644 --- a/packages/dd-trace/src/profiling/profilers/wall.js +++ b/packages/dd-trace/src/profiling/profilers/wall.js @@ -76,6 +76,44 @@ function getWebTags (startedSpans, i, span) { return memoize(null) } +let channelsActivated = false +function ensureChannelsActivated () { + if (channelsActivated) return + + const { AsyncLocalStorage, createHook } = require('async_hooks') + const shimmer = require('../../../../datadog-shimmer') + + createHook({ before: () => beforeCh.publish() }).enable() + + let inRun = false + shimmer.wrap(AsyncLocalStorage.prototype, 'enterWith', function (original) { + return function (...args) { + const retVal = original.apply(this, args) + if (!inRun) enterCh.publish() + return retVal + } + }) + + shimmer.wrap(AsyncLocalStorage.prototype, 'run', function (original) { + return function (store, callback, ...args) { + const wrappedCb = shimmer.wrapFunction(callback, cb => function (...args) { + inRun = false + enterCh.publish() + const retVal = cb.apply(this, args) + inRun = true + return retVal + }) + inRun = true + const retVal = original.call(this, store, wrappedCb, ...args) + enterCh.publish() + inRun = false + return retVal + } + }) + + channelsActivated = true +} + class NativeWallProfiler { constructor (options = {}) { this.type = 'wall' @@ -121,6 +159,8 @@ class NativeWallProfiler { start ({ mapper } = {}) { if (this._started) return + ensureChannelsActivated() + this._mapper = mapper this._pprof = require('@datadog/pprof') kSampleCount = this._pprof.time.constants.kSampleCount