zio-telemetry middleware

Middleware.scala

import io.opentelemetry.api.baggage.Baggage
import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator
import io.opentelemetry.context.propagation.TextMapGetter
import io.wolverinebeach.zio.telemetry.Tracing
import zio.Trace
import zio.ZIO
import zio.ZIOAspect
import zio.http.*

import java.lang
import scala.jdk.CollectionConverters.IterableHasAsJava

object Middleware {
  val RequestIdHeader: String = "x-request-id"

  class CustomAspectZIO[LowerR <: UpperR, UpperR](
      aspect: Request => ZIOAspect[LowerR, UpperR, Response, Response, Response, Response],
  ) extends RequestHandlerMiddleware.Contextual[LowerR, UpperR, Response, Response] {

    override type OutEnv[Env] = Env & LowerR
    override type OutErr[Err] = Err

    override def apply[Env >: LowerR <: UpperR, Err >: Response <: Response](
        handler: Handler[Env, Err, Request, Response],
    )(implicit trace: Trace): Handler[Env, Err, Request, Response] = {
      Handler.fromFunctionZIO[Request] { request =>
        handler.runZIO(request) @@ aspect(request)
      }
    }

    override def apply[Env >: LowerR <: UpperR, Err >: Response <: Response](http: Http[Env, Err, Request, Response])(
        implicit trace: Trace,
    ): Http[Env, Err, Request, Response] = Http.fromHandler { apply(http.toHandler(Handler.status(Status.Ok))) }

  }

  private class WithBaggage[R0](request: Request)
      extends ZIOAspect[R0, Tracing, Response, Response, Response, Response] {

    override def apply[R >: R0 <: Tracing, E >: Response <: Response, A >: Response <: Response](
        zio: ZIO[R, E, A],
    )(implicit trace: Trace): ZIO[R, E, A] =
      for {
        tracing <- ZIO.service[Tracing]
        context <- tracing.getCurrentContext.map {
          W3CBaggagePropagator.getInstance().extract(
            _,
            null,
            new TextMapGetter[Null] {
              override def keys(carrier: Null): lang.Iterable[String] = request.headers.map(_.headerName).asJava

              override def get(carrier: Null, key: String): String = request.headers.get(key).orNull
            },
          )
        }
        _      <- tracing.setBaggage(Baggage.fromContext(context))
        result <- zio
      } yield result

  }

  def withTracing(rootSpanName: String) = new CustomAspectZIO[Tracing, Tracing]({ request =>
    val base = Tracing.root[Response](rootSpanName) @@ new WithBaggage[Tracing](request)
    request.headers.get(RequestIdHeader) match {
      case Some(id) => Tracing.setAttribute(Tracing.Attributes.RequestId, id) @@ base
      case None     => base
    }
  })

}