capture stack trace for tracing spans

Effect-TS · May 1, 2024 · 142048e · 142048e
1 parent 81c8bff
commit 142048e
Show file tree

Hide file tree

Showing 18 changed files with 264 additions and 139 deletions.
diff --git a/.changeset/tiny-zebras-work.md b/.changeset/tiny-zebras-work.md
@@ -0,0 +1,9 @@
+---
+"effect": minor
+"@effect/experimental": patch
+"@effect/platform": patch
+"@effect/rpc": patch
+"@effect/sql": patch
+---
+
+capture stack trace for tracing spans
diff --git a/packages/effect/src/Tracer.ts b/packages/effect/src/Tracer.ts
@@ -92,6 +92,7 @@ export interface SpanOptions {
   readonly root?: boolean | undefined
   readonly context?: Context.Context<never> | undefined
   readonly kind?: SpanKind | undefined
+  readonly captureStackTrace?: boolean | Error | undefined
 }
 
 /**

diff --git a/packages/effect/src/internal/channel.ts b/packages/effect/src/internal/channel.ts
@@ -2310,29 +2310,47 @@ export const updateService = dual<
     )))
 
 /** @internal */
-export const withSpan = dual<
+export const withSpan: {
   (
     name: string,
     options?: Tracer.SpanOptions
-  ) => <OutElem, InElem, OutErr, InErr, OutDone, InDone, Env>(
+  ): <OutElem, InElem, OutErr, InErr, OutDone, InDone, Env>(
     self: Channel.Channel<OutElem, InElem, OutErr, InErr, OutDone, InDone, Env>
-  ) => Channel.Channel<OutElem, InElem, OutErr, InErr, OutDone, InDone, Exclude<Env, Tracer.ParentSpan>>,
+  ) => Channel.Channel<OutElem, InElem, OutErr, InErr, OutDone, InDone, Exclude<Env, Tracer.ParentSpan>>
   <OutElem, InElem, OutErr, InErr, OutDone, InDone, Env>(
     self: Channel.Channel<OutElem, InElem, OutErr, InErr, OutDone, InDone, Env>,
     name: string,
     options?: Tracer.SpanOptions
-  ) => Channel.Channel<OutElem, InElem, OutErr, InErr, OutDone, InDone, Exclude<Env, Tracer.ParentSpan>>
->(3, (self, name, options) =>
-  unwrapScoped(
-    Effect.flatMap(
-      Effect.context(),
-      (context) =>
-        Effect.map(
-          Effect.makeSpanScoped(name, options),
-          (span) => core.provideContext(self, Context.add(context, tracer.spanTag, span))
-        )
+  ): Channel.Channel<OutElem, InElem, OutErr, InErr, OutDone, InDone, Exclude<Env, Tracer.ParentSpan>>
+} = function() {
+  const dataFirst = typeof arguments[0] !== "string"
+  const name = dataFirst ? arguments[1] : arguments[0]
+  const options = tracer.addSpanStackTrace(dataFirst ? arguments[2] : arguments[1])
+  if (dataFirst) {
+    const self = arguments[0]
+    return unwrapScoped(
+      Effect.flatMap(
+        Effect.context(),
+        (context) =>
+          Effect.map(
+            Effect.makeSpanScoped(name, options),
+            (span) => core.provideContext(self, Context.add(context, tracer.spanTag, span))
+          )
+      )
+    )
+  }
+  return (self: Effect.Effect<any, any, any>) =>
+    unwrapScoped(
+      Effect.flatMap(
+        Effect.context(),
+        (context) =>
+          Effect.map(
+            Effect.makeSpanScoped(name, options),
+            (span) => core.provideContext(self, Context.add(context, tracer.spanTag, span))
+          )
+      )
     )
-  ) as any)
+} as any
 
 /** @internal */
 export const writeAll = <OutElem>(

diff --git a/packages/effect/src/internal/core-effect.ts b/packages/effect/src/internal/core-effect.ts
@@ -2027,7 +2027,7 @@ const bigint0 = BigInt(0)
 export const unsafeMakeSpan = <XA, XE>(
   fiber: FiberRuntime<XA, XE>,
   name: string,
-  options?: Tracer.SpanOptions
+  options: Tracer.SpanOptions
 ) => {
   const enabled = fiber.getFiberRef(core.currentTracerEnabled)
   if (enabled === false) {
@@ -2045,34 +2045,38 @@ export const unsafeMakeSpan = <XA, XE>(
   const annotationsFromEnv = FiberRefs.get(fiberRefs, core.currentTracerSpanAnnotations)
   const linksFromEnv = FiberRefs.get(fiberRefs, core.currentTracerSpanLinks)
 
-  const parent = options?.parent
+  const parent = options.parent
     ? Option.some(options.parent)
-    : options?.root
+    : options.root
     ? Option.none()
     : Context.getOption(context, internalTracer.spanTag)
 
   const links = linksFromEnv._tag === "Some" ?
-    options?.links !== undefined ?
+    options.links !== undefined ?
       [
         ...Chunk.toReadonlyArray(linksFromEnv.value),
-        ...(options?.links ?? [])
+        ...(options.links ?? [])
       ] :
       Chunk.toReadonlyArray(linksFromEnv.value) :
-    options?.links ?? Arr.empty()
+    options.links ?? Arr.empty()
 
   const span = tracer.span(
     name,
     parent,
-    options?.context ?? Context.empty(),
+    options.context ?? Context.empty(),
     links,
     timingEnabled ? clock.unsafeCurrentTimeNanos() : bigint0,
-    options?.kind ?? "internal"
+    options.kind ?? "internal"
   )
 
+  if (typeof options.captureStackTrace === "object" && typeof options.captureStackTrace.stack === "string") {
+    span.attribute("code.stacktrace", options.captureStackTrace.stack)
+  }
+
   if (annotationsFromEnv._tag === "Some") {
     HashMap.forEach(annotationsFromEnv.value, (value, key) => span.attribute(key, value))
   }
-  if (options?.attributes !== undefined) {
+  if (options.attributes !== undefined) {
     Object.entries(options.attributes).forEach(([k, v]) => span.attribute(k, v))
   }
 
@@ -2083,7 +2087,10 @@ export const unsafeMakeSpan = <XA, XE>(
 export const makeSpan = (
   name: string,
   options?: Tracer.SpanOptions
-): Effect.Effect<Tracer.Span> => core.withFiberRuntime((fiber) => core.succeed(unsafeMakeSpan(fiber, name, options)))
+): Effect.Effect<Tracer.Span> => {
+  options = internalTracer.addSpanStackTrace(options)
+  return core.withFiberRuntime((fiber) => core.succeed(unsafeMakeSpan(fiber, name, options)))
+}
 
 /* @internal */
 export const spanAnnotations: Effect.Effect<HashMap.HashMap<string, unknown>> = core
@@ -2108,7 +2115,7 @@ export const useSpan: {
     evaluate: (span: Tracer.Span) => Effect.Effect<A, E, R>
   ]
 ) => {
-  const options: Tracer.SpanOptions | undefined = args.length === 1 ? undefined : args[0]
+  const options = internalTracer.addSpanStackTrace(args.length === 1 ? undefined : args[0])
   const evaluate: (span: Tracer.Span) => Effect.Effect<A, E, R> = args[args.length - 1]
 
   return core.withFiberRuntime<A, E, R>((fiber) => {
@@ -2134,25 +2141,26 @@ export const withParentSpan = dual<
 >(2, (self, span) => provideService(self, internalTracer.spanTag, span))
 
 /** @internal */
-export const withSpan = dual<
+export const withSpan: {
   (
     name: string,
-    options?: Tracer.SpanOptions
-  ) => <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, Exclude<R, Tracer.ParentSpan>>,
+    options?: Tracer.SpanOptions | undefined
+  ): <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, Exclude<R, Tracer.ParentSpan>>
   <A, E, R>(
     self: Effect.Effect<A, E, R>,
     name: string,
-    options?: Tracer.SpanOptions
-  ) => Effect.Effect<A, E, Exclude<R, Tracer.ParentSpan>>
->(
-  (args) => typeof args[0] !== "string",
-  (self, name, options) =>
-    useSpan(
-      name,
-      options ?? {},
-      (span) => withParentSpan(self, span)
-    )
-)
+    options?: Tracer.SpanOptions | undefined
+  ): Effect.Effect<A, E, Exclude<R, Tracer.ParentSpan>>
+} = function() {
+  const dataFirst = typeof arguments[0] !== "string"
+  const name = dataFirst ? arguments[1] : arguments[0]
+  const options = internalTracer.addSpanStackTrace(dataFirst ? arguments[2] : arguments[1])
+  if (dataFirst) {
+    const self = arguments[0]
+    return useSpan(name, options, (span) => withParentSpan(self, span))
+  }
+  return (self: Effect.Effect<any, any, any>) => useSpan(name, options, (span) => withParentSpan(self, span))
+} as any
 
 // -------------------------------------------------------------------------------------
 // optionality

diff --git a/packages/effect/src/internal/fiberRuntime.ts b/packages/effect/src/internal/fiberRuntime.ts
@@ -3622,8 +3622,9 @@ export const interruptWhenPossible = dual<
 export const makeSpanScoped = (
   name: string,
   options?: Tracer.SpanOptions | undefined
-): Effect.Effect<Tracer.Span, never, Scope.Scope> =>
-  core.uninterruptible(
+): Effect.Effect<Tracer.Span, never, Scope.Scope> => {
+  options = tracer.addSpanStackTrace(options)
+  return core.uninterruptible(
     core.withFiberRuntime((fiber) => {
       const scope = Context.unsafeGet(fiber.getFiberRef(core.currentContext), scopeTag)
       const span = internalEffect.unsafeMakeSpan(fiber, name, options)
@@ -3641,27 +3642,37 @@ export const makeSpanScoped = (
       )
     })
   )
+}
 
 /* @internal */
 export const withTracerScoped = (value: Tracer.Tracer): Effect.Effect<void, never, Scope.Scope> =>
   fiberRefLocallyScopedWith(defaultServices.currentServices, Context.add(tracer.tracerTag, value))
 
 /** @internal */
-export const withSpanScoped = dual<
+export const withSpanScoped: {
   (
     name: string,
     options?: Tracer.SpanOptions
-  ) => <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, Exclude<R, Tracer.ParentSpan> | Scope.Scope>,
+  ): <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, Scope.Scope | Exclude<R, Tracer.ParentSpan>>
   <A, E, R>(
     self: Effect.Effect<A, E, R>,
     name: string,
     options?: Tracer.SpanOptions
-  ) => Effect.Effect<A, E, Exclude<R, Tracer.ParentSpan> | Scope.Scope>
->(
-  (args) => typeof args[0] !== "string",
-  (self, name, options) =>
+  ): Effect.Effect<A, E, Scope.Scope | Exclude<R, Tracer.ParentSpan>>
+} = function() {
+  const dataFirst = typeof arguments[0] !== "string"
+  const name = dataFirst ? arguments[1] : arguments[0]
+  const options = tracer.addSpanStackTrace(dataFirst ? arguments[2] : arguments[1])
+  if (dataFirst) {
+    const self = arguments[0]
+    return core.flatMap(
+      makeSpanScoped(name, tracer.addSpanStackTrace(options)),
+      (span) => internalEffect.provideService(self, tracer.spanTag, span)
+    )
+  }
+  return (self: Effect.Effect<any, any, any>) =>
     core.flatMap(
-      makeSpanScoped(name, options),
+      makeSpanScoped(name, tracer.addSpanStackTrace(options)),
       (span) => internalEffect.provideService(self, tracer.spanTag, span)
     )
-)
+} as any
diff --git a/packages/effect/src/internal/layer.ts b/packages/effect/src/internal/layer.ts
@@ -1113,15 +1113,15 @@ export const unwrapScoped = <A, E1, R1, E, R>(
 // -----------------------------------------------------------------------------
 
 /** @internal */
-export const withSpan = dual<
+export const withSpan: {
   (
     name: string,
     options?: Tracer.SpanOptions & {
       readonly onEnd?:
         | ((span: Tracer.Span, exit: Exit.Exit<unknown, unknown>) => Effect.Effect<void>)
         | undefined
     }
-  ) => <A, E, R>(self: Layer.Layer<A, E, R>) => Layer.Layer<A, E, Exclude<R, Tracer.ParentSpan>>,
+  ): <A, E, R>(self: Layer.Layer<A, E, R>) => Layer.Layer<A, E, Exclude<R, Tracer.ParentSpan>>
   <A, E, R>(
     self: Layer.Layer<A, E, R>,
     name: string,
@@ -1130,19 +1130,42 @@ export const withSpan = dual<
         | ((span: Tracer.Span, exit: Exit.Exit<unknown, unknown>) => Effect.Effect<void>)
         | undefined
     }
-  ) => Layer.Layer<A, E, Exclude<R, Tracer.ParentSpan>>
->((args) => isLayer(args[0]), (self, name, options) =>
-  unwrapScoped(
-    core.map(
-      options?.onEnd
-        ? core.tap(
-          fiberRuntime.makeSpanScoped(name, options),
-          (span) => fiberRuntime.addFinalizer((exit) => options.onEnd!(span, exit))
-        )
-        : fiberRuntime.makeSpanScoped(name, options),
-      (span) => withParentSpan(self, span)
+  ): Layer.Layer<A, E, Exclude<R, Tracer.ParentSpan>>
+} = function() {
+  const dataFirst = typeof arguments[0] !== "string"
+  const name = dataFirst ? arguments[1] : arguments[0]
+  const options = tracer.addSpanStackTrace(dataFirst ? arguments[2] : arguments[1]) as Tracer.SpanOptions & {
+    readonly onEnd?:
+      | ((span: Tracer.Span, exit: Exit.Exit<unknown, unknown>) => Effect.Effect<void>)
+      | undefined
+  }
+  if (dataFirst) {
+    const self = arguments[0]
+    return unwrapScoped(
+      core.map(
+        options?.onEnd
+          ? core.tap(
+            fiberRuntime.makeSpanScoped(name, options),
+            (span) => fiberRuntime.addFinalizer((exit) => options.onEnd!(span, exit))
+          )
+          : fiberRuntime.makeSpanScoped(name, options),
+        (span) => withParentSpan(self, span)
+      )
     )
-  ))
+  }
+  return (self: Layer.Layer<any, any, any>) =>
+    unwrapScoped(
+      core.map(
+        options?.onEnd
+          ? core.tap(
+            fiberRuntime.makeSpanScoped(name, options),
+            (span) => fiberRuntime.addFinalizer((exit) => options.onEnd!(span, exit))
+          )
+          : fiberRuntime.makeSpanScoped(name, options),
+        (span) => withParentSpan(self, span)
+      )
+    )
+} as any
 
 /** @internal */
 export const withParentSpan = dual<

diff --git a/packages/effect/src/internal/layer/circular.ts b/packages/effect/src/internal/layer/circular.ts
@@ -196,8 +196,9 @@ export const span = (
       | ((span: Tracer.Span, exit: Exit.Exit<unknown, unknown>) => Effect.Effect<void>)
       | undefined
   }
-): Layer.Layer<Tracer.ParentSpan> =>
-  layer.scoped(
+): Layer.Layer<Tracer.ParentSpan> => {
+  options = tracer.addSpanStackTrace(options) as any
+  return layer.scoped(
     tracer.spanTag,
     options?.onEnd
       ? core.tap(
@@ -206,6 +207,7 @@ export const span = (
       )
       : fiberRuntime.makeSpanScoped(name, options)
   )
+}
 
 /** @internal */
 export const setTracer = (tracer: Tracer.Tracer): Layer.Layer<never> =>

diff --git a/packages/effect/src/internal/stream.ts b/packages/effect/src/internal/stream.ts
@@ -49,6 +49,7 @@ import * as SinkEndReason from "./stream/sinkEndReason.js"
 import * as ZipAllState from "./stream/zipAllState.js"
 import * as ZipChunksState from "./stream/zipChunksState.js"
 import * as InternalTake from "./take.js"
+import * as InternalTracer from "./tracer.js"
 
 /** @internal */
 const StreamSymbolKey = "effect/Stream"
@@ -6800,17 +6801,26 @@ export const whenEffect = dual<
 )
 
 /** @internal */
-export const withSpan = dual<
+export const withSpan: {
   (
     name: string,
     options?: Tracer.SpanOptions
-  ) => <A, E, R>(self: Stream.Stream<A, E, R>) => Stream.Stream<A, E, Exclude<R, Tracer.ParentSpan>>,
+  ): <A, E, R>(self: Stream.Stream<A, E, R>) => Stream.Stream<A, E, Exclude<R, Tracer.ParentSpan>>
   <A, E, R>(
     self: Stream.Stream<A, E, R>,
     name: string,
     options?: Tracer.SpanOptions
-  ) => Stream.Stream<A, E, Exclude<R, Tracer.ParentSpan>>
->(3, (self, name, options) => new StreamImpl(channel.withSpan(toChannel(self), name, options)))
+  ): Stream.Stream<A, E, Exclude<R, Tracer.ParentSpan>>
+} = function() {
+  const dataFirst = typeof arguments[0] !== "string"
+  const name = dataFirst ? arguments[1] : arguments[0]
+  const options = InternalTracer.addSpanStackTrace(dataFirst ? arguments[2] : arguments[1])
+  if (dataFirst) {
+    const self = arguments[0]
+    return new StreamImpl(channel.withSpan(toChannel(self), name, options))
+  }
+  return (self: Stream.Stream<any, any, any>) => new StreamImpl(channel.withSpan(toChannel(self), name, options))
+} as any
 
 /** @internal */
 export const zip = dual<