IOLocal Based Invisible Trace

IOLocalInvisibleTrace.scala

object Test extends ResourceApp.Simple {

  def run = for {
    ep <- Jaeger.entryPoint("server")(_ => IO(Configuration.fromEnv().getTracerBuilder().build()))
    root <- ep.root("run")
    ioLocal <- Resource.eval(IOLocal(root))
    trace = MyTrace.fromIOLocal(ioLocal)
    spanning = MyServer.ioTraceOverride(ep, ioLocal)
    _ <- MyServer.inF(spanning)(Async[IO], trace)
  } yield ()

}

object MyServer {
  // Equivalent to you normal app
  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: natchez.Trace] = {
    val dsl = Http4sDsl[F]; import dsl._
    HttpRoutes.of[F]{
      case GET -> Root / "foo" => natchez.Trace[F].span("Woot")(Ok("bar"))
    }
  }

  def ioTraceOverride(
    ep: EntryPoint[IO],
    local: IOLocal[Span[IO]]
  ): HttpMiddleware[IO] = http => Kleisli{ req =>
    val kernel = Kernel(req.headers.headers.map(h => (h.name.toString -> h.value)).toMap)
    val spanR: Resource[IO, Span[IO]] = 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 {
  // 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)
    }
}