tom-pytel commented on a change in pull request #2: URL: https://github.com/apache/skywalking-nodejs/pull/2#discussion_r532060352
########## File path: src/trace/context/ContextManager.ts ########## @@ -18,19 +18,74 @@ */ import Context from '../../trace/context/Context'; +import Span from '../../trace/span/Span'; import SpanContext from '../../trace/context/SpanContext'; import { AsyncLocalStorage } from 'async_hooks'; -const store = new AsyncLocalStorage<Context>(); +type AsyncState = {context: Context, spans: Span[]}; + +const store = new AsyncLocalStorage<AsyncState>(); class ContextManager { - get current(): Context { - return store.getStore() || new SpanContext(); + get asyncState(): AsyncState { + let asyncState = store.getStore(); + + if (asyncState === undefined) { + asyncState = {context: new SpanContext(), spans: []}; + store.enterWith(asyncState); + } + + return asyncState; + } + + get current(): Context { return this.asyncState.context; } + get spans(): Span[] { return this.asyncState.spans; } + + spans_dup(): Span[] { + let asyncState = store.getStore(); + + if (asyncState === undefined) { + asyncState = {context: new SpanContext(), spans: []}; + } else { + asyncState = {context: asyncState.context, spans: [...asyncState.spans]}; + } + + store.enterWith(asyncState); + + return asyncState.spans; } - withContext(callback: (...args: any[]) => void, ...args: any[]) { - return store.run(new SpanContext(), callback); + clear(): void { + store.enterWith(undefined as unknown as AsyncState); } + + restore(context: Context, spans: Span[]): void { + store.enterWith({context, spans: spans || []}); + } + + withSpan(stopOnSuccess: boolean, span: Span, callback: (...args: any[]) => any, ...args: any[]): any { + try { + return callback(span, ...args); + + } catch (e) { + span.error(e); + + if (!stopOnSuccess) { + span.stop(); + } + + throw e; + + } finally { + if (stopOnSuccess) { + span.stop(); + } + } + } + + // withContext(callback: (...args: any[]) => void, ...args: any[]) { + // return store.run({context: new SpanContext(), spans: []}, callback); + // } Review comment: Actually that is what I did at first but realized this wrapper could be very handy especially for users to avoid errors (if people don't have to deal with .start() and .stop() then less possibility for screwups). Yes this is a callback but it is a cheap synchronous callback (just a function call), not the more expensive needs-its-own-async-context type of callback. Have a look at how I am using it in the examples I send up and if you do not think it is a net positive then I will remove. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org