Skip to content

Commit cbd1b38

Browse files
kubukozclaude
andcommitted
Add Trace.fromIOLocal for building a Trace[IO] from an existing IOLocal
Extracts the anonymous IOLocal-backed instance out of `ioTrace` into a public `fromIOLocal` constructor that takes a pre-constructed `IOLocal[Span[IO]]` and returns a `Trace[IO]` directly (no `IO`). This lets callers retain ownership of the `IOLocal`, so the current `Span[IO]` remains accessible via `local.get`. `ioTrace` now delegates to `fromIOLocal`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 7f65769 commit cbd1b38

1 file changed

Lines changed: 44 additions & 39 deletions

File tree

modules/core/shared/src/main/scala/Trace.scala

Lines changed: 44 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -110,50 +110,55 @@ object Trace {
110110

111111
/** A `Trace` instance that uses `IOLocal` internally. */
112112
def ioTrace(rootSpan: Span[IO]): IO[Trace[IO]] =
113-
IOLocal(rootSpan).map { local =>
114-
new Trace[IO] {
115-
116-
override def put(fields: (String, TraceValue)*): IO[Unit] =
117-
local.get.flatMap(_.put(fields: _*))
118-
119-
override def attachError(err: Throwable, fields: (String, TraceValue)*): IO[Unit] =
120-
local.get.flatMap(_.attachError(err, fields: _*))
121-
122-
override def log(fields: (String, TraceValue)*): IO[Unit] =
123-
local.get.flatMap(_.log(fields: _*))
124-
125-
override def log(event: String): IO[Unit] =
126-
local.get.flatMap(_.log(event))
127-
128-
override def kernel: IO[Kernel] =
129-
local.get.flatMap(_.kernel)
130-
131-
override def spanR(name: String, options: Span.Options): Resource[IO, IO ~> IO] =
132-
for {
133-
parent <- Resource.eval(local.get)
134-
child <- parent.span(name, options)
135-
} yield new (IO ~> IO) {
136-
def apply[A](fa: IO[A]): IO[A] =
137-
local.get.flatMap { old =>
138-
local
139-
.set(child)
140-
.bracket(_ => fa.onError { case e => child.attachError(e) })(_ => local.set(old))
141-
}
113+
IOLocal(rootSpan).map(fromIOLocal)
142114

143-
}
115+
/** A `Trace` instance backed by an existing `IOLocal[Span[IO]]`. Unlike [[ioTrace]], this gives
116+
* the caller ownership of the `IOLocal`, so the current `Span[IO]` remains accessible via
117+
* `local.get`.
118+
*/
119+
def fromIOLocal(local: IOLocal[Span[IO]]): Trace[IO] =
120+
new Trace[IO] {
121+
122+
override def put(fields: (String, TraceValue)*): IO[Unit] =
123+
local.get.flatMap(_.put(fields: _*))
124+
125+
override def attachError(err: Throwable, fields: (String, TraceValue)*): IO[Unit] =
126+
local.get.flatMap(_.attachError(err, fields: _*))
127+
128+
override def log(fields: (String, TraceValue)*): IO[Unit] =
129+
local.get.flatMap(_.log(fields: _*))
130+
131+
override def log(event: String): IO[Unit] =
132+
local.get.flatMap(_.log(event))
133+
134+
override def kernel: IO[Kernel] =
135+
local.get.flatMap(_.kernel)
136+
137+
override def spanR(name: String, options: Span.Options): Resource[IO, IO ~> IO] =
138+
for {
139+
parent <- Resource.eval(local.get)
140+
child <- parent.span(name, options)
141+
} yield new (IO ~> IO) {
142+
def apply[A](fa: IO[A]): IO[A] =
143+
local.get.flatMap { old =>
144+
local
145+
.set(child)
146+
.bracket(_ => fa.onError { case e => child.attachError(e) })(_ => local.set(old))
147+
}
144148

145-
override def span[A](name: String, options: Span.Options)(k: IO[A]): IO[A] =
146-
spanR(name, options).use(_(k))
149+
}
147150

148-
override def traceId: IO[Option[String]] =
149-
local.get.flatMap(_.traceId)
151+
override def span[A](name: String, options: Span.Options)(k: IO[A]): IO[A] =
152+
spanR(name, options).use(_(k))
150153

151-
override def spanId(implicit F: Applicative[IO]): IO[Option[String]] =
152-
local.get.flatMap(_.spanId)
154+
override def traceId: IO[Option[String]] =
155+
local.get.flatMap(_.traceId)
153156

154-
override def traceUri: IO[Option[URI]] =
155-
local.get.flatMap(_.traceUri)
156-
}
157+
override def spanId(implicit F: Applicative[IO]): IO[Option[String]] =
158+
local.get.flatMap(_.spanId)
159+
160+
override def traceUri: IO[Option[URI]] =
161+
local.get.flatMap(_.traceUri)
157162
}
158163

159164
/** A `Trace` instance that uses `IOLocal` internally. Span creation delegates to the supplied entry point. */

0 commit comments

Comments
 (0)