Improved zio-telemetry OpenTracing

OT2.scala

package zio.telemetry.opentracing

import java.util.concurrent.TimeUnit

import io.opentracing.{Span, Tracer}
import zio.clock.Clock
import zio._

object OpenTracing2 {
  type OT = Has[Service]

  trait Service {
    def tag(key: String, value: String): UIO[Unit]
    def locally[R, A](name: String, zio: RIO[R, A]): RIO[R, A]
  }
  
  def layer(tracer: Tracer, defaultRoot: String = "ROOT"): ZManaged[Clock, Nothing, Service] = {
    for {
      clock <- Managed.fromEffect(ZIO.access[Clock](_.get))
      span <- Managed.make(UIO(tracer.buildSpan(defaultRoot).start()))(s => UIO(s.finish()))
      ref <- Managed.fromEffect(FiberRef.make(span))
    } yield new Service {
      private def now: UIO[Long] = clock.currentTime(TimeUnit.MICROSECONDS)
      def tag(key: String, value: String): UIO[Unit]                      = ref.update(_.setTag(key, value))
      def locally[R, A](name: String, zio: RIO[R, A]): RIO[R, A] = {

        def child: Task[Span] = ref.get.flatMap(current => ZIO(tracer.buildSpan(name).asChildOf(current).start()))

        ZIO.bracket(child)(c => now.flatMap(at => UIO(c.finish(at))))(span => ref.locally(span)(zio))
      }
    }
  }

  def tag(key: String, value: String): RIO[OT, Unit] = ZIO.accessM(_.get.tag(key, value))
}