commit 912ad723a51ce774932ae22daa2d62c6b2421342
author Tomasz Pytel <tompytel@gmail.com> Tue Mar 30 04:27:33 2021 -0300
committer GitHub <noreply@github.com> Tue Mar 30 15:27:33 2021 +0800
tree 3c066673b30e2c4e9f9d5002adbd12017f3c12d6
parent 4670a45509e187aac77248eec5da720347830552

BugFix: wrong context during many async spans (#46)

src/trace/context/Context.ts

diff --git a/src/trace/context/Context.ts b/src/trace/context/Context.ts
index fdb6258..045a07e 100644
--- a/src/trace/context/Context.ts
+++ b/src/trace/context/Context.ts
@@ -24,6 +24,7 @@
 
 export default interface Context {
   segment: Segment;
+  nSpans: number;
 
   newLocalSpan(operation: string): Span;
 

src/trace/context/ContextManager.ts

diff --git a/src/trace/context/ContextManager.ts b/src/trace/context/ContextManager.ts
index 2d88d30..9f3205c 100644
--- a/src/trace/context/ContextManager.ts
+++ b/src/trace/context/ContextManager.ts
@@ -23,7 +23,7 @@
 
 import async_hooks from 'async_hooks';
 
-type AsyncState = { context: Context, spans: Span[], valid: boolean };
+type AsyncState = { spans: Span[], valid: boolean };
 
 let store: {
   getStore(): AsyncState | undefined;
@@ -62,7 +62,7 @@
     // Necessary because span may "finish()" in a child async task of where the asyncState was actually created and so clearing in the child would not clear in parent and invalid asyncState would be reused in new children of that parent.
     let asyncState = store.getStore();
     if (!asyncState?.valid) {
-      asyncState = { context: new SpanContext(), spans: [], valid: true };
+      asyncState = { spans: [], valid: true };
       store.enterWith(asyncState);
     }
 
@@ -80,7 +80,9 @@
   }
 
   get current(): Context {
-    return this.asyncState.context;
+    const asyncState = this.asyncState;
+
+    return !asyncState.spans.length ? new SpanContext() : asyncState.spans[asyncState.spans.length - 1].context;
   }
 
   get spans(): Span[] {
@@ -91,9 +93,9 @@
     let asyncState = store.getStore();
 
     if (!asyncState?.valid) {
-      asyncState = { context: new SpanContext(), spans: [], valid: true };
+      asyncState = { spans: [], valid: true };
     } else {
-      asyncState = { context: asyncState.context, spans: [...asyncState.spans], valid: asyncState.valid };
+      asyncState = { spans: [...asyncState.spans], valid: asyncState.valid };
     }
 
     store.enterWith(asyncState);
@@ -107,7 +109,7 @@
   }
 
   restore(context: Context, spans: Span[]): void {
-    store.enterWith({ context, spans: spans || [], valid: this.asyncState.valid });
+    store.enterWith({ spans: spans || [], valid: this.asyncState.valid });
   }
 
   withSpan(span: Span, callback: (...args: any[]) => any, ...args: any[]): any {

src/trace/context/DummyContext.ts

diff --git a/src/trace/context/DummyContext.ts b/src/trace/context/DummyContext.ts
index 99a722f..f08c93d 100644
--- a/src/trace/context/DummyContext.ts
+++ b/src/trace/context/DummyContext.ts
@@ -32,7 +32,7 @@
     type: SpanType.LOCAL,
   });
   segment: Segment = new Segment();
-  depth = 0;
+  nSpans = 0;
 
   newEntrySpan(operation: string, carrier?: ContextCarrier, inherit?: Component): Span {
     return this.span;
@@ -47,12 +47,12 @@
   }
 
   start(): Context {
-    this.depth++;
+    this.nSpans++;
     return this;
   }
 
   stop(): boolean {
-    return --this.depth === 0;
+    return --this.nSpans === 0;
   }
 
   async(span: Span) {

src/trace/context/SpanContext.ts

diff --git a/src/trace/context/SpanContext.ts b/src/trace/context/SpanContext.ts
index 5afa513..29241cd 100644
--- a/src/trace/context/SpanContext.ts
+++ b/src/trace/context/SpanContext.ts
@@ -168,7 +168,7 @@
     });
 
     this.nSpans += 1;
-    if (ContextManager.spans.every((s) => s.id !== span.id)) {
+    if (ContextManager.spans.every((s) => s.id !== span.id || s.context !== span.context)) {
       ContextManager.spans.push(span);
     }
 
@@ -226,7 +226,7 @@
 
     if (!ContextManager.hasContext || !ContextManager.spans.length) {
       ContextManager.restore(span.context, [span]);
-    } else if (ContextManager.spans.every((s) => s.id !== span.id)) {
+    } else if (ContextManager.spans.every((s) => s.id !== span.id || s.context !== span.context)) {
       ContextManager.spans.push(span);
     }
   }