fix(otel): afterReturnSync(), afterReturnAsync() of TraceLog() missing trace context

Author: waitingsong

packages/otel/src/lib/decorator.trace-log/trace-log.helper.async.ts

@@ -12,7 +12,7 @@ export async function beforeAsync(options: DecoratorExecutorParam): Promise<void
   const { traceService } = options
 
   const type = 'before'
-  if (! options.traceContext) {
+  if (! options.traceContext || ! options.span) {
     const info = traceService.getActiveTraceInfo()
     options.span = info.span
     options.traceContext = info.traceContext
@@ -50,21 +50,28 @@ export async function afterReturnAsync(options: DecoratorExecutorParam): Promise
   assert(! options.error, `[@mwcp/${ConfigKey.namespace}] options.error is not undefined in afterReturnAsync().
   Error: ${options.error?.message}`)
 
-  const res: DecoratorTraceDataResp = await processDecoratorBeforeAfterAsync('after', options)
-  if (res?.endSpanAfterTraceLog) {
-    endTraceSpan(traceService, span, res.spanStatusOptions)
+  if (! options.traceContext || ! options.span) {
+    const info = traceService.getActiveTraceInfo()
+    options.span = info.span
+    options.traceContext = info.traceContext
   }
-
-  if (res?.endParentSpan) {
-    if (! res.endSpanAfterTraceLog) {
+  await context.with(options.traceContext, async () => {
+    const res: DecoratorTraceDataResp = await processDecoratorBeforeAfterAsync('after', options)
+    if (res?.endSpanAfterTraceLog) {
       endTraceSpan(traceService, span, res.spanStatusOptions)
     }
 
-    const parentSpan = traceService.retrieveParentTraceInfoBySpan(span, options.traceScope)?.span
-    if (parentSpan) {
-      endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
+    if (res?.endParentSpan) {
+      if (! res.endSpanAfterTraceLog) {
+        endTraceSpan(traceService, span, res.spanStatusOptions)
+      }
+
+      const parentSpan = traceService.retrieveParentTraceInfoBySpan(span, options.traceScope)?.span
+      if (parentSpan) {
+        endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
+      }
     }
-  }
+  })
 
   return options.methodResult
 }

packages/otel/src/lib/decorator.trace-log/trace-log.helper.sync.ts

@@ -21,7 +21,7 @@ export function beforeSync(options: DecoratorExecutorParam): void {
   )
   assert(spanName, 'spanName is empty')
 
-  if (! options.traceContext) {
+  if (! options.traceContext || ! options.span) {
     const info = traceService.getActiveTraceInfo()
     options.span = info.span
     options.traceContext = info.traceContext
@@ -61,21 +61,29 @@ export function afterReturnSync(options: DecoratorExecutorParam): unknown {
   assert(! options.error, `[@mwcp/${ConfigKey.namespace}] options.error is not undefined in afterReturnSync().
   Error: ${options.error?.message}`)
 
-  const res: DecoratorTraceDataResp = processDecoratorBeforeAfterSync('after', options)
-  if (res?.endSpanAfterTraceLog) {
-    endTraceSpan(traceService, span, res.spanStatusOptions)
+  if (! options.traceContext || ! options.span) {
+    const info = traceService.getActiveTraceInfo()
+    options.span = info.span
+    options.traceContext = info.traceContext
   }
-
-  if (res?.endParentSpan) {
-    if (! res.endSpanAfterTraceLog) {
+  context.with(options.traceContext, () => {
+    const res: DecoratorTraceDataResp = processDecoratorBeforeAfterSync('after', options)
+    if (res?.endSpanAfterTraceLog) {
       endTraceSpan(traceService, span, res.spanStatusOptions)
     }
 
-    const parentSpan = traceService.retrieveParentTraceInfoBySpan(span, options.traceScope)?.span
-    if (parentSpan) {
-      endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
+    if (res?.endParentSpan) {
+      if (! res.endSpanAfterTraceLog) {
+        endTraceSpan(traceService, span, res.spanStatusOptions)
+      }
+
+      const parentSpan = traceService.retrieveParentTraceInfoBySpan(span, options.traceScope)?.span
+      if (parentSpan) {
+        endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
+      }
     }
-  }
+  })
+
   return options.methodResult
 }
 

packages/otel/src/lib/trace.service/trace.service.span.ts

@@ -79,7 +79,6 @@ export class TraceServiceSpan extends TraceServiceBase {
     assert(ret, 'startScopeActiveSpan() ret should not be null')
 
     const scope = options.scope ?? this.getWebContext()
-    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
     if (scope) {
       this.setActiveContext(ret.traceContext, scope)
     }