|
2 | 2 |
|
3 | 3 | const { AsyncLocalStorage } = require('async_hooks')
|
4 | 4 |
|
5 |
| -class DatadogStorage { |
6 |
| - constructor () { |
7 |
| - this._storage = new AsyncLocalStorage() |
8 |
| - } |
9 |
| - |
10 |
| - disable () { |
11 |
| - this._storage.disable() |
12 |
| - } |
13 |
| - |
| 5 | +/// This is exactly the same as AsyncLocalStorage, with the exception that it |
| 6 | +/// uses a WeakMap to store the store object. This is because ALS stores the |
| 7 | +/// store object as a property of the resource object, which causes all sorts |
| 8 | +/// of problems with logging and memory. We substitute the `store` object with |
| 9 | +/// a "handle" object, which is used as a key in a WeakMap, where the values |
| 10 | +/// are the real store objects. |
| 11 | +class DatadogStorage extends AsyncLocalStorage { |
14 | 12 | enterWith (store) {
|
15 | 13 | const handle = {}
|
16 | 14 | stores.set(handle, store)
|
17 |
| - this._storage.enterWith(handle) |
18 |
| - } |
19 |
| - |
20 |
| - exit (callback, ...args) { |
21 |
| - this._storage.exit(callback, ...args) |
| 15 | + super.enterWith(handle) |
22 | 16 | }
|
23 | 17 |
|
| 18 | + // This is method is a passthrough to the real `getStore()`, so that, when we |
| 19 | + // need it, we can use the handle rather than our mapped store. |
24 | 20 | // TODO: Refactor the Scope class to use a span-only store and remove this.
|
| 21 | + // It's only here because stores are currently used for a bunch of things, |
| 22 | + // and we don't want to hold on to all of them in spans |
| 23 | + // (see opentracing/span.js). Using a namespaced storage for spans would |
| 24 | + // solve this. |
25 | 25 | getHandle () {
|
26 |
| - return this._storage.getStore() |
| 26 | + return super.getStore() |
27 | 27 | }
|
28 | 28 |
|
| 29 | + // Here, we replicate the behavior of the original `getStore()` method by |
| 30 | + // passing in the handle, which we retrieve by calling it on super. Handles |
| 31 | + // retrieved through `getHandle()` can also be passed in to be used as the |
| 32 | + // key. This is useful if you've stashed a handle somewhere and want to |
| 33 | + // retrieve the store with it. |
29 | 34 | getStore (handle) {
|
30 | 35 | if (!handle) {
|
31 |
| - handle = this._storage.getStore() |
| 36 | + handle = super.getStore() |
32 | 37 | }
|
33 | 38 |
|
34 | 39 | return stores.get(handle)
|
35 | 40 | }
|
36 | 41 |
|
| 42 | + // Here, we replicate the behavior of the original `run()` method. We ensure |
| 43 | + // that our `enterWith()` is called internally, so that the handle to the |
| 44 | + // store is set. As an optimization, we use super for getStore and enterWith |
| 45 | + // when dealing with the parent store, so that we don't have to access the |
| 46 | + // WeakMap. |
37 | 47 | run (store, fn, ...args) {
|
38 |
| - const prior = this._storage.getStore() |
| 48 | + const prior = super.getStore() |
39 | 49 | this.enterWith(store)
|
40 | 50 | try {
|
41 | 51 | return Reflect.apply(fn, null, args)
|
42 | 52 | } finally {
|
43 |
| - this._storage.enterWith(prior) |
| 53 | + super.enterWith(prior) |
44 | 54 | }
|
45 | 55 | }
|
46 | 56 | }
|
47 | 57 |
|
48 |
| -const storages = Object.create(null) |
49 |
| -const legacyStorage = new DatadogStorage() |
| 58 | +// This is the map from handles to real stores, used in the class above. |
| 59 | +const stores = new WeakMap() |
50 | 60 |
|
51 |
| -const storage = function (namespace) { |
| 61 | +// For convenience, we use the `storage` function as a registry of namespaces |
| 62 | +// corresponding to DatadogStorage instances. This lets us have separate |
| 63 | +// storages for separate purposes. |
| 64 | +const storages = Object.create(null) |
| 65 | +function storage (namespace) { |
52 | 66 | if (!storages[namespace]) {
|
53 | 67 | storages[namespace] = new DatadogStorage()
|
54 | 68 | }
|
55 | 69 | return storages[namespace]
|
56 | 70 | }
|
57 | 71 |
|
| 72 | +// Namespaces are a new concept, so for existing internal code that does not |
| 73 | +// use namespaces, we have a "legacy" storage object. |
| 74 | +const legacyStorage = new DatadogStorage() |
58 | 75 | storage.disable = legacyStorage.disable.bind(legacyStorage)
|
59 | 76 | storage.enterWith = legacyStorage.enterWith.bind(legacyStorage)
|
60 | 77 | storage.exit = legacyStorage.exit.bind(legacyStorage)
|
61 | 78 | storage.getHandle = legacyStorage.getHandle.bind(legacyStorage)
|
62 | 79 | storage.getStore = legacyStorage.getStore.bind(legacyStorage)
|
63 | 80 | storage.run = legacyStorage.run.bind(legacyStorage)
|
64 | 81 |
|
65 |
| -const stores = new WeakMap() |
66 |
| - |
67 | 82 | module.exports = storage
|
0 commit comments