|
| 1 | +import { cacheable, cached } from '@seedcompany/common'; |
| 2 | +import { patchDecoratedMethod as patchMethod } from '@seedcompany/nest'; |
| 3 | +import { AsyncLocalStorage } from 'async_hooks'; |
| 4 | +import { AbstractClass } from 'type-fest'; |
| 5 | +import { getParentTypes as getHierarchyList } from './parent-types'; |
| 6 | + |
| 7 | +export type TraceNames = Readonly<{ |
| 8 | + cls: string; |
| 9 | + /** The parent class that declares the method when it is not the current class */ |
| 10 | + ownCls?: string; |
| 11 | + method: string; |
| 12 | +}>; |
| 13 | + |
| 14 | +/** |
| 15 | + * A performant way to identify certain methods in the call stack. |
| 16 | + * |
| 17 | + * This class helps to wrap class methods |
| 18 | + * and store their names in an AsyncLocalStorage. |
| 19 | + * This allows helper methods to identify who called them, |
| 20 | + * assuming they've been preconfigured. |
| 21 | + * |
| 22 | + * Traces are wrapped with a "layer"/"group" name. |
| 23 | + * This allows different groups of methods to be intertwined in the call |
| 24 | + * stack without conflicting with each other. |
| 25 | + * Multiple layers can be applied to the same method, independently. |
| 26 | + * |
| 27 | + * Best starting point is {@link applyToInstance} |
| 28 | + * ```ts |
| 29 | + * class Foo { |
| 30 | + * constructor() { |
| 31 | + * TraceLayer.as('logic').applyToInstance(this); |
| 32 | + * } |
| 33 | + * } |
| 34 | + * ``` |
| 35 | + * This capture approach is best because it automatically |
| 36 | + * applies to subclasses and applies to inherited methods. |
| 37 | + * |
| 38 | + * Next best is {@link applyToClass} |
| 39 | + * ```ts |
| 40 | + * @TraceLayer.as('logic').applyToClass() |
| 41 | + * class Foo {} |
| 42 | + * ``` |
| 43 | + * This applies to all owned & inherited methods. |
| 44 | + * |
| 45 | + * The current stack can be pulled anywhere with |
| 46 | + * ```ts |
| 47 | + * TraceLayer.as('logic').currentStack; |
| 48 | + * ``` |
| 49 | + * This gives a list of call names ordered by most recent. |
| 50 | + * ```ts |
| 51 | + * const { cls, method } = currentStack?.[0] ?? {}; |
| 52 | + * ``` |
| 53 | + */ |
| 54 | +export class TraceLayer { |
| 55 | + private static readonly layers = new AsyncLocalStorage<{ |
| 56 | + readonly [layer in string]?: readonly TraceNames[]; |
| 57 | + }>(); |
| 58 | + private static readonly instances = new Map<string, TraceLayer>(); |
| 59 | + private static readonly getNameCacheForInstance = cacheable< |
| 60 | + object, |
| 61 | + Map<string, TraceNames> |
| 62 | + >(new WeakMap(), () => new Map()); |
| 63 | + |
| 64 | + private readonly seenClasses = new WeakSet<AbstractClass<unknown>>(); |
| 65 | + |
| 66 | + private constructor(readonly layer: string) {} |
| 67 | + |
| 68 | + static as(layer: string) { |
| 69 | + return cached(TraceLayer.instances, layer, () => new TraceLayer(layer)); |
| 70 | + } |
| 71 | + |
| 72 | + /** |
| 73 | + * This gives a list of call names, if any, ordered by most recent. |
| 74 | + * The identity of these entries is maintained, |
| 75 | + * allowing them to be used as cache keys. |
| 76 | + */ |
| 77 | + get currentStack() { |
| 78 | + const layers = TraceLayer.layers.getStore(); |
| 79 | + return layers?.[this.layer]; |
| 80 | + } |
| 81 | + |
| 82 | + /** |
| 83 | + * A shortcut to create a memoized function that takes the current trace |
| 84 | + * and converts it to another shape. |
| 85 | + */ |
| 86 | + makeGetter<T>(mapFnCached: (names: TraceNames) => T) { |
| 87 | + const cache = new WeakMap(); |
| 88 | + return () => { |
| 89 | + const names = this.currentStack?.[0]; |
| 90 | + return names ? cached(cache, names, mapFnCached) : undefined; |
| 91 | + }; |
| 92 | + } |
| 93 | + |
| 94 | + applyToClass(): ClassDecorator { |
| 95 | + return (target) => { |
| 96 | + this.applyToStaticClass(target as any); |
| 97 | + }; |
| 98 | + } |
| 99 | + |
| 100 | + applyToInstance(obj: InstanceType<AbstractClass<any>>) { |
| 101 | + this.applyToStaticClass(obj.constructor); |
| 102 | + } |
| 103 | + |
| 104 | + applyToStaticClass(cls: AbstractClass<unknown>) { |
| 105 | + for (const aClass of getHierarchyList(cls)) { |
| 106 | + this.applyToOwnStaticClass(aClass); |
| 107 | + } |
| 108 | + } |
| 109 | + |
| 110 | + applyToOwnStaticClass(cls: AbstractClass<unknown>) { |
| 111 | + if (this.seenClasses.has(cls)) { |
| 112 | + return; |
| 113 | + } |
| 114 | + this.seenClasses.add(cls); |
| 115 | + |
| 116 | + const proto = cls.prototype; |
| 117 | + const descriptors = Object.getOwnPropertyDescriptors(proto); |
| 118 | + const methods = Object.entries(descriptors).flatMap(([key, descriptor]) => { |
| 119 | + return key !== 'constructor' && typeof descriptor.value === 'function' |
| 120 | + ? [key] |
| 121 | + : []; |
| 122 | + }); |
| 123 | + |
| 124 | + const layer = this.layer; |
| 125 | + for (const name of methods) { |
| 126 | + patchMethod(proto as any, name, (original) => { |
| 127 | + return function (...args) { |
| 128 | + const nameCache = TraceLayer.getNameCacheForInstance(this); |
| 129 | + const names = cached(nameCache, name, (): TraceNames => { |
| 130 | + const cls = this.constructor.name; |
| 131 | + const ownCls = proto.constructor.name; |
| 132 | + return { |
| 133 | + cls: cls, |
| 134 | + ...(cls !== ownCls && { ownCls }), |
| 135 | + method: name, |
| 136 | + }; |
| 137 | + }); |
| 138 | + |
| 139 | + const prev = TraceLayer.layers.getStore(); |
| 140 | + const next = { |
| 141 | + ...prev, |
| 142 | + [layer]: [names, ...(prev?.[layer] ?? [])], |
| 143 | + }; |
| 144 | + return TraceLayer.layers.run(next, original, ...args); |
| 145 | + }; |
| 146 | + }); |
| 147 | + } |
| 148 | + } |
| 149 | +} |
0 commit comments