Remove legacyStorage in favor of namespaced storages

packages/datadog-core/index.js

@@ -1,5 +1,5 @@
 'use strict'
 
-const storage = require('./src/storage')
+const { storage } = require('./src/storage')
 
 module.exports = { storage }

packages/datadog-core/src/storage.js

@@ -2,35 +2,54 @@
 
 const { AsyncLocalStorage } = require('async_hooks')
 
-/// This is exactly the same as AsyncLocalStorage, with the exception that it
-/// uses a WeakMap to store the store object. This is because ALS stores the
-/// store object as a property of the resource object, which causes all sorts
-/// of problems with logging and memory. We substitute the `store` object with
-/// a "handle" object, which is used as a key in a WeakMap, where the values
-/// are the real store objects.
+/**
+ * This is exactly the same as AsyncLocalStorage, with the exception that it
+ * uses a WeakMap to store the store object. This is because ALS stores the
+ * store object as a property of the resource object, which causes all sorts
+ * of problems with logging and memory. We substitute the `store` object with
+ * a "handle" object, which is used as a key in a WeakMap, where the values
+ * are the real store objects.
+ *
+ * @template T
+ */
 class DatadogStorage extends AsyncLocalStorage {
+  /**
+   *
+   * @param store {DatadogStorage}
+   */
   enterWith (store) {
     const handle = {}
     stores.set(handle, store)
     super.enterWith(handle)
   }
 
-  // This is method is a passthrough to the real `getStore()`, so that, when we
-  // need it, we can use the handle rather than our mapped store.
-  // TODO: Refactor the Scope class to use a span-only store and remove this.
-  // It's only here because stores are currently used for a bunch of things,
-  // and we don't want to hold on to all of them in spans
-  // (see opentracing/span.js). Using a namespaced storage for spans would
-  // solve this.
+  /**
+   * This is method is a passthrough to the real `getStore()`, so that, when we
+   * need it, we can use the handle rather than our mapped store.
+   *
+   * It's only here because stores are currently used for a bunch of things,
+   * and we don't want to hold on to all of them in spans
+   * (see opentracing/span.js). Using a namespaced storage for spans would
+   * solve this.
+   *
+   * TODO: Refactor the Scope class to use a span-only store and remove this.
+   *
+   * @returns {{}}
+   */
   getHandle () {
     return super.getStore()
   }
 
-  // Here, we replicate the behavior of the original `getStore()` method by
-  // passing in the handle, which we retrieve by calling it on super. Handles
-  // retrieved through `getHandle()` can also be passed in to be used as the
-  // key. This is useful if you've stashed a handle somewhere and want to
-  // retrieve the store with it.
+  /**
+   * Here, we replicate the behavior of the original `getStore()` method by
+   * passing in the handle, which we retrieve by calling it on super. Handles
+   * retrieved through `getHandle()` can also be passed in to be used as the
+   * key. This is useful if you've stashed a handle somewhere and want to
+   * retrieve the store with it.
+   *
+   * @param handle {{}}
+   * @returns {T | undefined}
+   */
   getStore (handle) {
     if (!handle) {
       handle = super.getStore()
@@ -39,11 +58,19 @@ class DatadogStorage extends AsyncLocalStorage {
     return stores.get(handle)
   }
 
-  // Here, we replicate the behavior of the original `run()` method. We ensure
-  // that our `enterWith()` is called internally, so that the handle to the
-  // store is set. As an optimization, we use super for getStore and enterWith
-  // when dealing with the parent store, so that we don't have to access the
-  // WeakMap.
+  /**
+   * Here, we replicate the behavior of the original `run()` method. We ensure
+   * that our `enterWith()` is called internally, so that the handle to the
+   * store is set. As an optimization, we use super for getStore and enterWith
+   * when dealing with the parent store, so that we don't have to access the
+   * WeakMap.
+   * @template R
+   * @template TArgs extends any[]
+   * @param store {DatadogStorage}
+   * @param fn {() => R}
+   * @param args {TArgs}
+   * @returns {void}
+   */
   run (store, fn, ...args) {
     const prior = super.getStore()
     this.enterWith(store)
@@ -55,28 +82,31 @@ class DatadogStorage extends AsyncLocalStorage {
   }
 }
 
-// This is the map from handles to real stores, used in the class above.
+/**
+ * This is the map from handles to real stores, used in the class above.
+ * @template T
+ * @type {WeakMap<WeakKey, T>}
+ */
 const stores = new WeakMap()
 
-// For convenience, we use the `storage` function as a registry of namespaces
-// corresponding to DatadogStorage instances. This lets us have separate
-// storages for separate purposes.
+/**
+ * For convenience, we use the `storage` function as a registry of namespaces
+ * corresponding to DatadogStorage instances. This lets us have separate
+ * storages for separate purposes.
+ * @type {Map<string, DatadogStorage>}
+ */
 const storages = Object.create(null)
+
+/**
+ *
+ * @param namespace {string} the namespace to use
+ * @returns {DatadogStorage}
+ */
 function storage (namespace) {
   if (!storages[namespace]) {
     storages[namespace] = new DatadogStorage()
   }
   return storages[namespace]
 }
 
-// Namespaces are a new concept, so for existing internal code that does not
-// use namespaces, we have a "legacy" storage object.
-const legacyStorage = new DatadogStorage()
-storage.disable = legacyStorage.disable.bind(legacyStorage)
-storage.enterWith = legacyStorage.enterWith.bind(legacyStorage)
-storage.exit = legacyStorage.exit.bind(legacyStorage)
-storage.getHandle = legacyStorage.getHandle.bind(legacyStorage)
-storage.getStore = legacyStorage.getStore.bind(legacyStorage)
-storage.run = legacyStorage.run.bind(legacyStorage)
-
-module.exports = storage
+module.exports = { storage }

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

@@ -16,17 +16,17 @@ describe('storage', () => {
   })
 
   afterEach(() => {
-    testStorage.enterWith(undefined)
+    testStorage('legacy').enterWith(undefined)
     testStorage2.enterWith(undefined)
   })
 
   it('should enter a store', done => {
     const store = 'foo'
 
-    testStorage.enterWith(store)
+    testStorage('legacy').enterWith(store)
 
     setImmediate(() => {
-      expect(testStorage.getStore()).to.equal(store)
+      expect(testStorage('legacy').getStore()).to.equal(store)
       done()
     })
   })
@@ -35,11 +35,11 @@ describe('storage', () => {
     const store = 'foo'
     const store2 = 'bar'
 
-    testStorage.enterWith(store)
+    testStorage('legacy').enterWith(store)
     testStorage2.enterWith(store2)
 
     setImmediate(() => {
-      expect(testStorage.getStore()).to.equal(store)
+      expect(testStorage('legacy').getStore()).to.equal(store)
       expect(testStorage2.getStore()).to.equal(store2)
       done()
     })
@@ -52,7 +52,7 @@ describe('storage', () => {
   it('should not have its store referenced by the underlying async resource', () => {
     const resource = executionAsyncResource()
 
-    testStorage.enterWith({ internal: 'internal' })
+    testStorage('legacy').enterWith({ internal: 'internal' })
 
     for (const sym of Object.getOwnPropertySymbols(resource)) {
       if (sym.toString() === 'Symbol(kResourceStore)' && resource[sym]) {

packages/datadog-instrumentations/test/body-parser.spec.js

@@ -77,7 +77,7 @@ withVersions('body-parser', 'body-parser', version => {
       let payload
 
       function handler (data) {
-        store = storage.getStore()
+        store = storage('legacy').getStore()
         payload = data
       }
       bodyParserReadCh.subscribe(handler)

packages/datadog-instrumentations/test/generic-pool.spec.js

@@ -27,11 +27,11 @@ describe('Instrumentation', () => {
       it('should run the acquire() callback in context where acquire() was called', done => {
         const store = 'store'
 
-        storage.run(store, () => {
+        storage('legacy').run(store, () => {
           // eslint-disable-next-line n/handle-callback-err
           pool.acquire((err, resource) => {
             pool.release(resource)
-            expect(storage.getStore()).to.equal(store)
+            expect(storage('legacy').getStore()).to.equal(store)
             done()
           })
         })
@@ -56,20 +56,20 @@ describe('Instrumentation', () => {
         const store = 'store'
         const store2 = 'store2'
 
-        storage.run(store, () => {
+        storage('legacy').run(store, () => {
           pool.acquire()
             .then(resource => {
               pool.release(resource)
-              expect(storage.getStore()).to.equal(store)
+              expect(storage('legacy').getStore()).to.equal(store)
             })
             .catch(done)
         })
 
-        storage.run(store2, () => {
+        storage('legacy').run(store2, () => {
           pool.acquire()
             .then(resource => {
               pool.release(resource)
-              expect(storage.getStore()).to.equal(store2)
+              expect(storage('legacy').getStore()).to.equal(store2)
               done()
             })
             .catch(done)

packages/datadog-instrumentations/test/helpers/instrument.spec.js

@@ -10,12 +10,12 @@ const { AsyncResource } = require('../../src/helpers/instrument')
 describe('helpers/instrument', () => {
   describe('AsyncResource', () => {
     it('should bind statically', () => {
-      storage.run('test1', () => {
+      storage('legacy').run('test1', () => {
         const tested = AsyncResource.bind(() => {
-          expect(storage.getStore()).to.equal('test1')
+          expect(storage('legacy').getStore()).to.equal('test1')
         })
 
-        storage.run('test2', () => {
+        storage('legacy').run('test2', () => {
           tested()
         })
       })
@@ -34,12 +34,12 @@ describe('helpers/instrument', () => {
     })
 
     it('should bind a specific instance', () => {
-      storage.run('test1', () => {
+      storage('legacy').run('test1', () => {
         const asyncResource = new AsyncResource('test')
 
-        storage.run('test2', () => {
+        storage('legacy').run('test2', () => {
           const tested = asyncResource.bind((a, b, c) => {
-            expect(storage.getStore()).to.equal('test1')
+            expect(storage('legacy').getStore()).to.equal('test1')
             expect(test.asyncResource).to.equal(asyncResource)
             expect(test).to.have.length(3)
           })

packages/datadog-instrumentations/test/helpers/promise.js

@@ -32,18 +32,18 @@ module.exports = (name, factory, versionRange) => {
 
           let promise = new Promise((resolve, reject) => {
             setImmediate(() => {
-              storage.run('promise', () => {
+              storage('legacy').run('promise', () => {
                 resolve()
               })
             })
           })
 
-          storage.run(store, () => {
+          storage('legacy').run(store, () => {
             for (let i = 0; i < promise.then.length; i++) {
               const args = new Array(i + 1)
 
               args[i] = () => {
-                expect(storage.getStore()).to.equal(store)
+                expect(storage('legacy').getStore()).to.equal(store)
               }
 
               promise = promise.then.apply(promise, args)
@@ -54,31 +54,31 @@ module.exports = (name, factory, versionRange) => {
         })
 
         it('should run the catch() callback in the context where catch() was called', () => {
-          const store = storage.getStore()
+          const store = storage('legacy').getStore()
 
           let promise = new Promise((resolve, reject) => {
             setImmediate(() => {
-              storage.run('promise', () => {
+              storage('legacy').run('promise', () => {
                 reject(new Error())
               })
             })
           })
 
-          storage.run(store, () => {
+          storage('legacy').run(store, () => {
             promise = promise
               .catch(err => {
                 throw err
               })
               .catch(() => {
-                expect(storage.getStore()).to.equal(store)
+                expect(storage('legacy').getStore()).to.equal(store)
               })
           })
 
           return promise
         })
 
         it('should allow to run without a scope if not available when calling then()', () => {
-          storage.run(null, () => {
+          storage('legacy').run(null, () => {
             const promise = new Promise((resolve, reject) => {
               setImmediate(() => {
                 resolve()
@@ -87,7 +87,7 @@ module.exports = (name, factory, versionRange) => {
 
             return promise
               .then(() => {
-                expect(storage.getStore()).to.be.null
+                expect(storage('legacy').getStore()).to.be.null
               })
           })
         })

packages/datadog-instrumentations/test/knex.spec.js

@@ -24,10 +24,10 @@ describe('Instrumentation', () => {
         afterEach(() => client.destroy())
 
         it('should propagate context', () =>
-          storage.run(store, () =>
+          storage('legacy').run(store, () =>
             client.raw('PRAGMA user_version')
               .finally(() => {
-                expect(storage.getStore()).to.equal(store)
+                expect(storage('legacy').getStore()).to.equal(store)
               })
               .catch(() => {})
           )

packages/datadog-instrumentations/test/multer.spec.js

@@ -86,7 +86,7 @@ withVersions('multer', 'multer', version => {
       let payload
 
       function handler (data) {
-        store = storage.getStore()
+        store = storage('legacy').getStore()
         payload = data
       }
       multerReadCh.subscribe(handler)

packages/datadog-instrumentations/test/passport-http.spec.js

@@ -174,7 +174,7 @@ withVersions('passport-http', 'passport-http', version => {
 
     it('should block when subscriber aborts', async () => {
       subscriberStub = sinon.spy(({ abortController }) => {
-        storage.getStore().req.res.writeHead(403).end('Blocked')
+        storage('legacy').getStore().req.res.writeHead(403).end('Blocked')
         abortController.abort()
       })
 

packages/datadog-instrumentations/test/passport-local.spec.js

@@ -154,7 +154,7 @@ withVersions('passport-local', 'passport-local', version => {
 
     it('should block when subscriber aborts', async () => {
       subscriberStub = sinon.spy(({ abortController }) => {
-        storage.getStore().req.res.writeHead(403).end('Blocked')
+        storage('legacy').getStore().req.res.writeHead(403).end('Blocked')
         abortController.abort()
       })
 

packages/datadog-instrumentations/test/passport.spec.js

@@ -145,7 +145,7 @@ withVersions('passport', 'passport', version => {
       const cookie = login.headers['set-cookie'][0]
 
       subscriberStub.callsFake(({ abortController }) => {
-        const res = storage.getStore().req.res
+        const res = storage('legacy').getStore().req.res
         res.writeHead(403)
         res.constructor.prototype.end.call(res, 'Blocked')
         abortController.abort()