feat(otel): add properties of DecoratorTraceData for TraceLog()

- before/after callback of TraceLog() support return traceContext to update traceContext and span
- endParentSpan to end current and parent span

Author: waitingsong

packages/otel/src/lib/decorator.helper.async.ts

@@ -9,10 +9,13 @@ import type {
   DecoratorTraceDataResp,
   TraceDecoratorOptions,
 } from './trace.service/index.trace.service.js'
-import { isSpanEnded } from './util.js'
+import { getSpan, isSpanEnded } from './util.js'
 
 // #region processDecoratorBeforeAfterAsync
 
+/**
+ * @description options.{traceContext|span} may be modified with `traceContext` returned by `before()` or `after()`
+ */
 export async function processDecoratorBeforeAfterAsync(
   type: 'before' | 'after' | 'afterThrow',
   options: DecoratorExecutorParam<TraceDecoratorOptions>,
@@ -59,17 +62,24 @@ export async function processDecoratorBeforeAfterAsync(
       data = await func2(options.methodArgs, options.error, decoratorContext)
     }
 
-    if (data) {
+    if (data && Object.keys(data).length) {
       const eventName = type
       if (data.events && ! data.events['event']) {
         data.events['event'] = eventName
       }
 
+      if (data.traceContext) {
+        options.traceContext = data.traceContext
+        options.span = getSpan(options.traceContext)
+      }
+      assert(options.span, `processDecoratorBeforeAfterAsync(): span is required with new traceContext returned by "${type}()"`)
+
       assert(options.traceScope, 'processDecoratorBeforeAfterAsync(): traceScope is required')
-      processDecoratorSpanData(options.traceScope, traceService, span, data)
+      processDecoratorSpanData(options.traceScope, traceService, options.span, data)
+      return data
     }
 
-    return data
+    return null
   }
 }
 

packages/otel/src/lib/decorator.helper.sync.ts

@@ -11,7 +11,7 @@ import type {
   TraceDecoratorOptions,
 } from './trace.service/index.trace.service.js'
 import { AttrNames } from './types.js'
-import { isSpanEnded } from './util.js'
+import { getSpan, isSpanEnded } from './util.js'
 
 
 // #region processDecoratorBeforeAfterSync
@@ -62,7 +62,7 @@ export function processDecoratorBeforeAfterSync(
       data = func2(options.methodArgs, options.error, decoratorContext)
     }
 
-    if (data) {
+    if (data && Object.keys(data).length) {
       if (isPromise(data)) {
         const err = new Error(`processDecoratorBeforeAfterSync() decorator ${type}() return value is a promise,
       class: ${options.callerAttr[AttrNames.CallerClass]}, method: ${options.callerAttr[AttrNames.CallerMethod]}`)
@@ -75,11 +75,18 @@ export function processDecoratorBeforeAfterSync(
         data.events['event'] = eventName
       }
 
+      if (data.traceContext) {
+        options.traceContext = data.traceContext
+        options.span = getSpan(options.traceContext)
+      }
+      assert(options.span, `processDecoratorBeforeAfterSync(): span is required with new traceContext returned by "${type}()"`)
+
       assert(options.traceScope, 'processDecoratorBeforeAfterSync(): traceScope is required')
-      processDecoratorSpanData(options.traceScope, traceService, span, data)
+      processDecoratorSpanData(options.traceScope, traceService, options.span, data)
+      return data
     }
 
-    return data
+    return null
   }
 
 }

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

@@ -2,6 +2,7 @@ import assert from 'node:assert'
 
 import { processDecoratorBeforeAfterAsync } from '../decorator.helper.async.js'
 import { genTraceScopeFrom } from '../decorator.helper.base.js'
+import { endTraceSpan } from '../trace.helper.js'
 import type { DecoratorExecutorParam, DecoratorTraceDataResp } from '../trace.service/index.trace.service.js'
 import { ConfigKey } from '../types.js'
 
@@ -16,16 +17,29 @@ export async function beforeAsync(options: DecoratorExecutorParam): Promise<void
   }
 
   if (! options.span) {
-    options.span = traceService.getActiveSpan(options.traceScope)
+    const info = traceService.getActiveTraceInfo(options.traceScope)
+    if (info) {
+      options.span = info.span
+      options.traceContext = info.traceContext
+    }
   }
+
   const res: DecoratorTraceDataResp = await processDecoratorBeforeAfterAsync(type, options)
   if (res?.endSpanAfterTraceLog) {
     assert(options.span, 'span is required')
-    if (res.spanStatusOptions) {
-      traceService.endSpan({ span: options.span, spanStatusOptions: res.spanStatusOptions })
+    endTraceSpan(traceService, options.span, res.spanStatusOptions)
+  }
+
+  if (res?.endParentSpan) {
+    assert(options.span, 'span is required')
+
+    if (! res.endSpanAfterTraceLog) {
+      endTraceSpan(traceService, options.span, res.spanStatusOptions)
     }
-    else {
-      traceService.endSpan({ span: options.span })
+
+    const parentSpan = traceService.retrieveParentTraceInfoBySpan(options.span, options.traceScope)?.span
+    if (parentSpan) {
+      endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
     }
   }
 }
@@ -42,13 +56,20 @@ export async function afterReturnAsync(options: DecoratorExecutorParam): Promise
 
   const res: DecoratorTraceDataResp = await processDecoratorBeforeAfterAsync('after', options)
   if (res?.endSpanAfterTraceLog) {
-    if (res.spanStatusOptions) {
-      traceService.endSpan({ span, spanStatusOptions: res.spanStatusOptions })
+    endTraceSpan(traceService, span, res.spanStatusOptions)
+  }
+
+  if (res?.endParentSpan) {
+    if (! res.endSpanAfterTraceLog) {
+      endTraceSpan(traceService, span, res.spanStatusOptions)
     }
-    else {
-      traceService.endSpan({ span })
+
+    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

@@ -5,6 +5,7 @@ import { ConfigKey } from '@mwcp/share'
 
 import { genTraceScopeFrom } from '../decorator.helper.base.js'
 import { processDecoratorBeforeAfterSync } from '../decorator.helper.sync.js'
+import { endTraceSpan } from '../trace.helper.js'
 import type { DecoratorExecutorParam, DecoratorTraceDataResp } from '../trace.service/index.trace.service.js'
 import { AttrNames } from '../types.js'
 
@@ -27,16 +28,29 @@ export function beforeSync(options: DecoratorExecutorParam): void {
   }
 
   if (! options.span) {
-    options.span = traceService.getActiveSpan(options.traceScope)
+    const info = traceService.getActiveTraceInfo(options.traceScope)
+    if (info) {
+      options.span = info.span
+      options.traceContext = info.traceContext
+    }
   }
+
   const res: DecoratorTraceDataResp = processDecoratorBeforeAfterSync(type, options)
   if (res?.endSpanAfterTraceLog) {
     assert(options.span, 'span is required')
-    if (res.spanStatusOptions) {
-      traceService.endSpan({ span: options.span, spanStatusOptions: res.spanStatusOptions })
+    endTraceSpan(traceService, options.span, res.spanStatusOptions)
+  }
+
+  if (res?.endParentSpan) {
+    assert(options.span, 'span is required')
+
+    if (! res.endSpanAfterTraceLog) {
+      endTraceSpan(traceService, options.span, res.spanStatusOptions)
     }
-    else {
-      traceService.endSpan({ span: options.span })
+
+    const parentSpan = traceService.retrieveParentTraceInfoBySpan(options.span, options.traceScope)?.span
+    if (parentSpan) {
+      endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
     }
   }
 }
@@ -54,14 +68,19 @@ export function afterReturnSync(options: DecoratorExecutorParam): unknown {
 
   const res: DecoratorTraceDataResp = processDecoratorBeforeAfterSync('after', options)
   if (res?.endSpanAfterTraceLog) {
-    if (res.spanStatusOptions) {
-      traceService.endSpan({ span, spanStatusOptions: res.spanStatusOptions })
+    endTraceSpan(traceService, span, res.spanStatusOptions)
+  }
+
+  if (res?.endParentSpan) {
+    if (! res.endSpanAfterTraceLog) {
+      endTraceSpan(traceService, span, res.spanStatusOptions)
     }
-    else {
-      traceService.endSpan({ span })
+
+    const parentSpan = traceService.retrieveParentTraceInfoBySpan(span, options.traceScope)?.span
+    if (parentSpan) {
+      endTraceSpan(traceService, parentSpan, res.spanStatusOptions)
     }
   }
-
   return options.methodResult
 }
 

packages/otel/src/lib/trace.helper.ts

@@ -4,13 +4,16 @@ import assert from 'node:assert'
 import type { DecoratorExecutorParamBase } from '@mwcp/share'
 import { isArrowFunction } from '@waiting/shared-core'
 
+import type { Span, SpanStatusOptions } from '##/index.js'
+
 import type {
   DecoratorContext,
   DecoratorExecutorParam,
   GenDecoratorExecutorOptions,
   KeyGenerator,
   TraceDecoratorOptions,
 } from './trace.service/index.trace.service.js'
+import type { TraceServiceSpan } from './trace.service/trace.service.span.js'
 import { AttrNames, ConfigKey } from './types.js'
 
 
@@ -213,3 +216,13 @@ export function genDecoratorExecutorOptions(
 
   return ret
 }
+
+
+export function endTraceSpan(traceService: TraceServiceSpan, span: Span, spanStatusOptions: SpanStatusOptions | undefined): void {
+  if (spanStatusOptions) {
+    traceService.endSpan({ span, spanStatusOptions })
+  }
+  else {
+    traceService.endSpan({ span })
+  }
+}

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

@@ -37,15 +37,27 @@ export interface DecoratorTraceData {
   rootEvents?: Attributes
 
   /**
-   * End then span after method `before()` or `after()` called
+   * End the span after method `before()` or `after()` called
    * used by TraceLog decorator, ignored by TraceInit/Trace decorator
+   * @description Current span also will be ended if endParentSpan:true, regardless of this value
    * @default false
    */
   endSpanAfterTraceLog?: boolean
+  /**
+   * End the parent span after method `before()` or `after()` called
+   * used by TraceLog decorator, ignored by TraceInit/Trace decorator
+   * @description Current span also will be ended if endParentSpan:true, regardless of endSpanAfterTraceLog
+   * @default false
+   */
+  endParentSpan?: boolean
   /**
    * Used by TraceLog decorator and endSpanAfterTraceLog:true, ignored by TraceInit/Trace decorator
    */
   spanStatusOptions?: SpanStatusOptions
+  /**
+   * options.traceContext will be overwritten by this value, and options.span also will be updated
+   */
+  traceContext?: TraceContext
 }
 export type DecoratorTraceDataResp = DecoratorTraceData | undefined | null
 export type DecoratorTraceDataRespAsync = Promise<DecoratorTraceData | undefined | null>