From 31fe59535e0c26bc2faa6bbd3040579886563a95 Mon Sep 17 00:00:00 2001
From: Ethan <meshethan@gmail.com>
Date: Wed, 27 Jan 2021 17:27:15 +0800
Subject: [PATCH 1/2] use invalid to check asyncState

---
 src/trace/context/Context.ts        | 2 ++
 src/trace/context/ContextManager.ts | 7 ++++---
 src/trace/context/DummyContext.ts   | 1 +
 src/trace/context/SpanContext.ts    | 1 +
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/trace/context/Context.ts b/src/trace/context/Context.ts
index 8a23252..ea55397 100644
--- a/src/trace/context/Context.ts
+++ b/src/trace/context/Context.ts
@@ -24,6 +24,8 @@ import { ContextCarrier } from './ContextCarrier';
 export default interface Context {
   segment: Segment;
 
+  invalid: boolean;
+
   newLocalSpan(operation: string): Span;
 
   newEntrySpan(operation: string, carrier?: ContextCarrier): Span;
diff --git a/src/trace/context/ContextManager.ts b/src/trace/context/ContextManager.ts
index a586856..92aa9b2 100644
--- a/src/trace/context/ContextManager.ts
+++ b/src/trace/context/ContextManager.ts
@@ -59,8 +59,8 @@ if (async_hooks.AsyncLocalStorage) {
 class ContextManager {
   get asyncState(): AsyncState {
     let asyncState = store.getStore();
-
-    if (asyncState === undefined) {
+    // since `AsyncLocalStorage.getStore` may get previous state, see issue https://github.com/nodejs/node/issues/35286#issuecomment-697207158, so recreate when context is invalid
+    if (asyncState === undefined || asyncState.context.invalid) {
       asyncState = { context: new SpanContext(), spans: [] };
       store.enterWith(asyncState);
     }
@@ -79,7 +79,7 @@ class ContextManager {
   spansDup(): Span[] {
     let asyncState = store.getStore();
 
-    if (asyncState === undefined) {
+    if (asyncState === undefined || asyncState.context.invalid) {
       asyncState = { context: new SpanContext(), spans: [] };
     } else {
       asyncState = { context: asyncState.context, spans: [...asyncState.spans] };
@@ -91,6 +91,7 @@ class ContextManager {
   }
 
   clear(): void {
+    this.current.invalid = true;
     store.enterWith(undefined as unknown as AsyncState);
   }
 
diff --git a/src/trace/context/DummyContext.ts b/src/trace/context/DummyContext.ts
index b64a1dd..f8587b8 100644
--- a/src/trace/context/DummyContext.ts
+++ b/src/trace/context/DummyContext.ts
@@ -32,6 +32,7 @@ export default class DummyContext implements Context {
   });
   segment: Segment = new Segment();
   depth = 0;
+  invalid = false;
 
   newEntrySpan(operation: string, carrier?: ContextCarrier): Span {
     return this.span;
diff --git a/src/trace/context/SpanContext.ts b/src/trace/context/SpanContext.ts
index 20bf96c..efb336b 100644
--- a/src/trace/context/SpanContext.ts
+++ b/src/trace/context/SpanContext.ts
@@ -39,6 +39,7 @@ export default class SpanContext implements Context {
   spanId = 0;
   nSpans = 0;
   segment: Segment = new Segment();
+  invalid = false;
 
   get parent(): Span | null {
     if (ContextManager.spans.length > 0) {

From 9f9badd1a5ebda8f674fd411e4c279f14adf1bda Mon Sep 17 00:00:00 2001
From: Ethan <meshethan@gmail.com>
Date: Wed, 27 Jan 2021 20:24:47 +0800
Subject: [PATCH 2/2] add invalid flag for AsyncState

---
 src/trace/context/Context.ts        |  2 --
 src/trace/context/ContextManager.ts | 18 +++++++++---------
 src/trace/context/DummyContext.ts   |  1 -
 src/trace/context/SpanContext.ts    |  1 -
 4 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/src/trace/context/Context.ts b/src/trace/context/Context.ts
index ea55397..8a23252 100644
--- a/src/trace/context/Context.ts
+++ b/src/trace/context/Context.ts
@@ -24,8 +24,6 @@ import { ContextCarrier } from './ContextCarrier';
 export default interface Context {
   segment: Segment;
 
-  invalid: boolean;
-
   newLocalSpan(operation: string): Span;
 
   newEntrySpan(operation: string, carrier?: ContextCarrier): Span;
diff --git a/src/trace/context/ContextManager.ts b/src/trace/context/ContextManager.ts
index 92aa9b2..ef21803 100644
--- a/src/trace/context/ContextManager.ts
+++ b/src/trace/context/ContextManager.ts
@@ -23,7 +23,7 @@ import SpanContext from '../../trace/context/SpanContext';
 
 import async_hooks from 'async_hooks';
 
-type AsyncState = { context: Context, spans: Span[] };
+type AsyncState = { context: Context, spans: Span[], invalid: boolean };
 
 let store: {
   getStore(): AsyncState | undefined;
@@ -59,9 +59,9 @@ if (async_hooks.AsyncLocalStorage) {
 class ContextManager {
   get asyncState(): AsyncState {
     let asyncState = store.getStore();
-    // since `AsyncLocalStorage.getStore` may get previous state, see issue https://github.com/nodejs/node/issues/35286#issuecomment-697207158, so recreate when context is invalid
-    if (asyncState === undefined || asyncState.context.invalid) {
-      asyncState = { context: new SpanContext(), spans: [] };
+    // since `AsyncLocalStorage.getStore` may get previous state, see issue https://github.com/nodejs/node/issues/35286#issuecomment-697207158, so recreate when asyncState is invalid
+    if (asyncState === undefined || asyncState.invalid) {
+      asyncState = { context: new SpanContext(), spans: [], invalid: false };
       store.enterWith(asyncState);
     }
 
@@ -79,10 +79,10 @@ class ContextManager {
   spansDup(): Span[] {
     let asyncState = store.getStore();
 
-    if (asyncState === undefined || asyncState.context.invalid) {
-      asyncState = { context: new SpanContext(), spans: [] };
+    if (asyncState === undefined || asyncState.invalid) {
+      asyncState = { context: new SpanContext(), spans: [], invalid: false };
     } else {
-      asyncState = { context: asyncState.context, spans: [...asyncState.spans] };
+      asyncState = { context: asyncState.context, spans: [...asyncState.spans], invalid: asyncState.invalid };
     }
 
     store.enterWith(asyncState);
@@ -91,12 +91,12 @@ class ContextManager {
   }
 
   clear(): void {
-    this.current.invalid = true;
+    this.asyncState.invalid = true;
     store.enterWith(undefined as unknown as AsyncState);
   }
 
   restore(context: Context, spans: Span[]): void {
-    store.enterWith({ context, spans: spans || [] });
+    store.enterWith({ context, spans: spans || [], invalid: this.asyncState.invalid });
   }
 
   withSpan(span: Span, callback: (...args: any[]) => any, ...args: any[]): any {
diff --git a/src/trace/context/DummyContext.ts b/src/trace/context/DummyContext.ts
index f8587b8..b64a1dd 100644
--- a/src/trace/context/DummyContext.ts
+++ b/src/trace/context/DummyContext.ts
@@ -32,7 +32,6 @@ export default class DummyContext implements Context {
   });
   segment: Segment = new Segment();
   depth = 0;
-  invalid = false;
 
   newEntrySpan(operation: string, carrier?: ContextCarrier): Span {
     return this.span;
diff --git a/src/trace/context/SpanContext.ts b/src/trace/context/SpanContext.ts
index efb336b..20bf96c 100644
--- a/src/trace/context/SpanContext.ts
+++ b/src/trace/context/SpanContext.ts
@@ -39,7 +39,6 @@ export default class SpanContext implements Context {
   spanId = 0;
   nSpans = 0;
   segment: Segment = new Segment();
-  invalid = false;
 
   get parent(): Span | null {
     if (ContextManager.spans.length > 0) {