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


Reply via email to