trace4cats / trace4cats-http4s Goto Github PK

package example

import cats.data.Kleisli
import cats.effect.kernel.Ref
import cats.effect.{Async, ExitCode, IO, IOApp, Resource}
import cats.syntax.all._
import com.comcast.ip4s._
import fs2.Stream
import org.http4s.Http
import org.http4s.ember.server.EmberServerBuilder
import org.http4s.implicits._
import org.http4s.server.middleware.{GZip, Logger}
import org.http4s.server.{Router, Server}
import sttp.tapir.server.http4s.{Http4sServerInterpreter, Http4sServerOptions}
import trace4cats._
import trace4cats.http4s.common.Http4sRequestFilter
import trace4cats.http4s.server.syntax._
import trace4cats.jaeger.JaegerSpanCompleter


object Main extends IOApp {
  type G[x] = Kleisli[IO, Span[IO], x]

  def entryPoint[F[_]: Async](process: TraceProcess): Resource[F, EntryPoint[F]] =
    JaegerSpanCompleter[F](process = process).map { completer =>
      EntryPoint[F](SpanSampler.always, completer)
    }

  override def run(args: List[String]): IO[ExitCode] = {
    for {
      counterRef <- Ref[IO].of(0)

      serverOptions: Http4sServerOptions[G] = Http4sServerOptions.default[G]
      serverLogic: ServerLogic[G] = new ServerLogic[G](counterRef)
      resource: Resource[IO, Server] = for {
        ep <- entryPoint[IO](TraceProcess("http4s-demo"))
        r: HttpRoutes[G] = Http4sServerInterpreter[G](serverOptions).toRoutes(serverLogic.all)
        routes = r.inject(ep, requestFilter = Http4sRequestFilter.kubernetesPrometheus)
        app: Http[IO, IO] = Router("/" -> routes).orNotFound
        finalApp = Logger.httpApp(logHeaders = true, logBody = false)(GZip(app))
        s <- EmberServerBuilder.default[IO].withHost(ipv4"0.0.0.0")
          .withPort(port"8080").withHttpApp(finalApp).build >> Resource.eval(Async[IO].never)
      } yield s
      s <- Stream.resource(resource).drain.compile.drain.as(ExitCode.Success)
    } yield s
  }
}

  def liftTrace[F[_]: MonadCancelThrow, G[_]: MonadCancelThrow](
      client: Client[F],
      toHeaders: ToHeaders = ToHeaders.standard
  )(implicit P: Provide[F, G, Span[F]]): Client[G] =
    Client { (request: Request[G]) =>
      Resource.eval(P.ask).flatMap { span =>
        val traceHeaders = toHeaders
          .fromContext(span.context)
          .values
          .map(Header.Raw.apply)
          .toList

        client
          .run(
            request
              .transformHeaders(_ ++ Headers(traceHeaders))
              .mapK(P.provideK(span))
          )
          .mapK(P.liftK)
          .map(_.mapK(P.liftK))
      }
    }