use Channel.acquireUseRelease for Channel.withSpan (#3134)

.changeset/stupid-islands-enjoy.md

@@ -0,0 +1,5 @@
+---
+"effect": patch
+---
+
+use Channel.acquireUseRelease for Channel.withSpan

packages/effect/src/internal/channel.ts

@@ -8,6 +8,7 @@ import * as Either from "../Either.js"
 import * as Equal from "../Equal.js"
 import * as Exit from "../Exit.js"
 import * as Fiber from "../Fiber.js"
+import * as FiberRef from "../FiberRef.js"
 import { constVoid, dual, identity, pipe } from "../Function.js"
 import type { LazyArg } from "../Function.js"
 import * as Layer from "../Layer.js"
@@ -28,6 +29,7 @@ import * as mergeDecision from "./channel/mergeDecision.js"
 import * as mergeState from "./channel/mergeState.js"
 import * as _mergeStrategy from "./channel/mergeStrategy.js"
 import * as singleProducerAsyncInput from "./channel/singleProducerAsyncInput.js"
+import * as coreEffect from "./core-effect.js"
 import * as core from "./core-stream.js"
 import * as MergeDecisionOpCodes from "./opCodes/channelMergeDecision.js"
 import * as MergeStateOpCodes from "./opCodes/channelMergeState.js"
@@ -2326,29 +2328,25 @@ export const withSpan: {
  const dataFirst = typeof arguments[0] !== "string"
  const name = dataFirst ? arguments[1] : arguments[0]
  const options = tracer.addSpanStackTrace(dataFirst ? arguments[2] : arguments[1])
+ const acquire = Effect.all([
+ Effect.makeSpan(name, options),
+ Effect.context(),
+ Effect.clock,
+ FiberRef.get(FiberRef.currentTracerTimingEnabled)
+ ])
  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 acquireUseRelease(
+ acquire,
+ ([span, context]) => core.provideContext(self, Context.add(context, tracer.spanTag, span)),
+ ([span, , clock, timingEnabled], exit) => coreEffect.endSpan(span, exit, clock, timingEnabled)
  )
  }
- 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))
- )
- )
+ return (self: Channel.Channel<any>) =>
+ acquireUseRelease(
+ acquire,
+ ([span, context]) => core.provideContext(self, Context.add(context, tracer.spanTag, span)),
+ ([span, , clock, timingEnabled], exit) => coreEffect.endSpan(span, exit, clock, timingEnabled)
  )
 } as any
 

packages/effect/src/internal/core-effect.ts

@@ -1,3 +1,4 @@
+import type { Exit } from "effect/Exit"
 import { internalCall } from "effect/Utils"
 import * as Arr from "../Array.js"
 import type * as Cause from "../Cause.js"
@@ -2122,6 +2123,18 @@ export const spanAnnotations: Effect.Effect<HashMap.HashMap<string, unknown>> =
 export const spanLinks: Effect.Effect<Chunk.Chunk<Tracer.SpanLink>> = core
  .fiberRefGet(core.currentTracerSpanLinks)
 
+/** @internal */
+export const endSpan = <A, E>(span: Tracer.Span, exit: Exit<A, E>, clock: Clock.Clock, timingEnabled: boolean) =>
+ core.sync(() => {
+ if (span.status._tag === "Ended") {
+ return
+ }
+ if (core.exitIsFailure(exit) && internalCause.spanToTrace.has(span)) {
+ span.attribute("code.stacktrace", internalCause.spanToTrace.get(span)!())
+ }
+ span.end(timingEnabled ? clock.unsafeCurrentTimeNanos() : bigint0, exit)
+ })
+
 /** @internal */
 export const useSpan: {
  <A, E, R>(name: string, evaluate: (span: Tracer.Span) => Effect.Effect<A, E, R>): Effect.Effect<A, E, R>
@@ -2144,16 +2157,7 @@ export const useSpan: {
  const span = unsafeMakeSpan(fiber, name, options)
  const timingEnabled = fiber.getFiberRef(core.currentTracerTimingEnabled)
  const clock = Context.get(fiber.getFiberRef(defaultServices.currentServices), clockTag)
- return core.onExit(evaluate(span), (exit) =>
- core.sync(() => {
- if (span.status._tag === "Ended") {
- return
- }
- if (core.exitIsFailure(exit) && internalCause.spanToTrace.has(span)) {
- span.attribute("code.stacktrace", internalCause.spanToTrace.get(span)!())
- }
- span.end(timingEnabled ? clock.unsafeCurrentTimeNanos() : bigint0, exit)
- }))
+ return core.onExit(evaluate(span), (exit) => endSpan(span, exit, clock, timingEnabled))
  })
 }
 

packages/effect/src/internal/fiberRuntime.ts

@@ -3666,13 +3666,7 @@ export const makeSpanScoped = (
  const timingEnabled = fiber.getFiberRef(core.currentTracerTimingEnabled)
  const clock_ = Context.get(fiber.getFiberRef(defaultServices.currentServices), clock.clockTag)
  return core.as(
- core.scopeAddFinalizerExit(scope, (exit) =>
- core.sync(() => {
- if (span.status._tag === "Ended") {
- return
- }
- span.end(timingEnabled ? clock_.unsafeCurrentTimeNanos() : BigInt(0), exit)
- })),
+ core.scopeAddFinalizerExit(scope, (exit) => internalEffect.endSpan(span, exit, clock_, timingEnabled)),
  span
  )
  })