@@ -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