Full FiberLocal Tracing

FiberLocalTracing.scala

object Test extends ResourceApp.Simple {

  def run = for {
    ep <- Resource.eval(
      Jaeger.globalTracerEntryPoint[IO](None).flatMap(_.toRight(new Throwable("No Global")).liftTo[IO])
    )
    root <- ep.root("run")
    fiberLocal <- Resource.eval(GenFiberLocal[IO].local(root))
    trace = MyTrace.fromFiberLocal(fiberLocal)
    spanning = MyServer.myMiddleware(ep, fiberLocal)
    _ <- MyServer.inF(spanning)(Async[IO], trace)
  } yield ()
  

}


object MyServer {
  def inF[F[_]: Async: natchez.Trace](spanning: org.http4s.server.HttpMiddleware[F]): Resource[F,Unit] = {
    val app = spanning(myApp[F]).orNotFound
    Resource.eval(
      EmberServerBuilder.default[F]
      .withHttpApp(app)
      .build
      .useForever
    )
  }

  def myApp[F[_]: Monad] = {
    val dsl = Http4sDsl[F]; import dsl._
    HttpRoutes.of[F]{
      case GET -> Root / "foo" => Ok("bar")
    }
  }

  def myMiddleware[F[_]: MonadCancelThrow](
    ep: EntryPoint[F],
    local: FiberLocal[F, Span[F]],
  ): HttpMiddleware[F] = http => Kleisli{req =>
    val kernel = Kernel(req.headers.headers.map(h => (h.name.toString -> h.value)).toMap)
    val spanR: Resource[F, Span[F]]  = ep.continueOrElseRoot(req.uri.path.toString, kernel)
    OptionT(spanR.use(span => 
      Resource.make(local.set(span))(_ => local.reset).use(_ => 
        http.run(req).value
      )
    ))
  }

}

object MyTrace {

  def fromFiberLocal[F[_]: MonadCancelThrow](local: FiberLocal[F, Span[F]]): natchez.Trace[F] = 
    new natchez.Trace[F] {
      def put(fields: (String, TraceValue)*): F[Unit] =
        local.get.flatMap(_.put(fields: _*))

      def kernel: F[Kernel] =
        local.get.flatMap(_.kernel)

      def span[A](name: String)(k: F[A]): F[A] =
        local.get.flatMap { parent =>
          parent.span(name).flatMap { child =>
            Resource.make(local.set(child))(_ => local.set(parent))
          } .use { _ => k }
        }

      def traceId: F[Option[String]] =
        local.get.flatMap(_.traceId)

      def traceUri: F[Option[URI]] =
        local.get.flatMap(_.traceUri)
    }

  // Proves this can all be done without FiberLocal if we are ok
  // being precisely IO at the end of the world.
  def fromIOLocal(local: IOLocal[Span[IO]]): natchez.Trace[IO] = 
    new natchez.Trace[IO] {
      def put(fields: (String, TraceValue)*): IO[Unit] =
        local.get.flatMap(_.put(fields: _*))

      def kernel: IO[Kernel] =
        local.get.flatMap(_.kernel)

      def span[A](name: String)(k: IO[A]): IO[A] =
        local.get.flatMap { parent =>
          parent.span(name).flatMap { child =>
            Resource.make(local.set(child))(_ => local.set(parent))
          } .use { _ => k }
        }

      def traceId: IO[Option[String]] =
        local.get.flatMap(_.traceId)

      def traceUri: IO[Option[URI]] =
        local.get.flatMap(_.traceUri)
    }
}

// FiberLocal is my proposal to exist

trait FiberLocal[F[_], A]{
  def get: F[A]

  def set(value: A): F[Unit]

  def reset: F[Unit]

  def update(f: A => A): F[Unit]

  def modify[B](f: A => (A, B)): F[B]

  def getAndSet(value: A): F[A]

  def getAndReset: F[A]
}
object FiberLocal {
  def fromIOLocal[F[_]: LiftIO, A](local: IOLocal[A]): FiberLocal[F, A] = new FiberLocal[F, A] {
    def get: F[A] = LiftIO[F].liftIO(local.get)
    
    def set(value: A): F[Unit] = LiftIO[F].liftIO(local.set(value))
    
    def reset: F[Unit] = LiftIO[F].liftIO(local.reset)
    
    def update(f: A => A): F[Unit] = LiftIO[F].liftIO(local.update(f))
    
    def modify[B](f: A => (A, B)): F[B] = LiftIO[F].liftIO(local.modify(f))
    
    def getAndSet(value: A): F[A] = LiftIO[F].liftIO(local.getAndSet(value))
    
    def getAndReset: F[A] = LiftIO[F].liftIO(local.getAndReset)
  }
}
trait GenFiberLocal[F[_]]{
  def local[A](default: A): F[FiberLocal[F, A]]
}
object GenFiberLocal {
  def apply[F[_]](implicit ev: GenFiberLocal[F]): GenFiberLocal[F] = ev

  implicit def fromLiftIO[F[_]: LiftIO: Functor]: GenFiberLocal[F] = new GenFiberLocal[F] {
    def local[A](default: A): F[FiberLocal[F,A]] = LiftIO[F].liftIO(IOLocal(default)).map( local =>
      FiberLocal.fromIOLocal(local)
    )
  }
}