Natchez Tracing

TracingMiddleware.scala

package com.me

import cats.data.{Kleisli, OptionT}
import cats.effect.{Bracket, Resource}
import cats.{Defer, Monad, ~>}
import natchez._
import org.http4s.{HttpRoutes, Request, Response}
import cats.implicits._

import scala.util.control.NonFatal

object TracingMiddleware {
  private def liftTE[F[_] : Bracket[?[_], Throwable]](
    entryPoint: EntryPoint[F]
  )(routes: HttpRoutes[Kleisli[F, Span[F], ?]]): HttpRoutes[F] =
    Kleisli { req =>
      type G[A] = Kleisli[F, Span[F], A]
      val lift = λ[F ~> G](fa => Kleisli(_ => fa))
      val kernel = Kernel(req.headers.toList.map(h => (h.name.value -> h.value)).toMap)
      val spanR = entryPoint.continueOrElseRoot(req.uri.path, kernel)
      OptionT {
        spanR.use { span =>
          val lower = λ[G ~> F](_(span))
          routes.run(req.mapK(lift)).mapK(lower).map(_.mapK(lower)).value
        }
      }
    }

  /**
    * A middleware that adds the following standard fields to the current span:
    *
    * - "http.method"      -> "GET", "PUT", etc.
    * - "http.url"         -> request URI (not URL)
    * - "http.status_code" -> "200", "403", etc. // why is this a string?
    * - "error"            -> true // not present if no error
    *
    * In addition the following non-standard fields are added in case of error:
    *
    * - "error.message"    -> Exception message
    * - "error.stacktrace" -> Exception stack trace as a multi-line string
    */
  def natchezMiddleware[F[_] : Bracket[?[_], Throwable] : Trace](routes: HttpRoutes[F]): HttpRoutes[F] =
    Kleisli { req =>
      val addRequestFields: F[Unit] =
        Trace[F].put(Tags.http.method(req.method.name), Tags.http.url(req.uri.renderString))

      def addResponseFields(res: Response[F]): F[Unit] =
        Trace[F].put(Tags.http.status_code(res.status.code.toString))

      def addErrorFields(e: Throwable): F[Unit] =
        Trace[F].put(
          Tags.error(true),
          "error.message"    -> e.getMessage,
          "error.stacktrace" -> e.getStackTrace.mkString("\n")
        )

      OptionT {
        routes(req)
          .onError {
            case NonFatal(e) => OptionT.liftF(addRequestFields *> addErrorFields(e))
          }
          .value
          .flatMap {
            case Some(handler) => addRequestFields *> addResponseFields(handler).as(Some(handler))
            case None          => Option.empty[Response[F]].pure[F]
          }
      }
    }

  implicit class EntryPointOps[F[_]](self: EntryPoint[F]) {

    private def dummySpan(implicit ev: Monad[F]): Span[F] =
      new Span[F] {
        val kernel: F[Kernel] = Kernel(Map.empty).pure[F]
        def put(fields: (String, TraceValue)*): F[Unit] = Monad[F].unit
        def span(name: String): Resource[F, Span[F]] = Monad[Resource[F, ?]].pure(this)
      }

    def liftT(routes: HttpRoutes[Kleisli[F, Span[F], ?]])(implicit ev: Bracket[F, Throwable]): HttpRoutes[F] =
      liftTE(self)(routes)

    /**
      * Lift an `HttpRoutes`-yielding resource that consumes `Span`s into the bare effect. We do this
      * by ignoring any tracing that happens during allocation and freeing of the `HttpRoutes`
      * resource. The reasoning is that such a resource typically lives for the lifetime of the
      * application and it's of little use to keep a span open that long.
      */
    def liftR(
      routes: Resource[Kleisli[F, Span[F], ?], HttpRoutes[Kleisli[F, Span[F], ?]]]
    )(implicit ev: Bracket[F, Throwable], d: Defer[F]): Resource[F, HttpRoutes[F]] =
      routes
        .map(liftT)
        .mapK(λ[Kleisli[F, Span[F], ?] ~> F] { fa =>
          fa.run(dummySpan)
        })

    private def unliftTE(
      entryPoint: EntryPoint[F]
    )(routes: HttpRoutes[F])(implicit ev: Bracket[F, Throwable], d: Defer[F]): HttpRoutes[Kleisli[F, Span[F], ?]] =
      Kleisli { req =>
        type G[A] = Kleisli[F, Span[F], A]
        val lift = λ[G ~> F](fa => fa.run(dummySpan))
        val kernel = Kernel(req.headers.toList.map(h => (h.name.value -> h.value)).toMap)
        OptionT {
          val lower = λ[F ~> G] { (fa: F[Any]) =>
            Kleisli(_ => fa)
          }
          routes.run(req.mapK(lift)).mapK(lower).map(_.mapK(lower)).value
        }
      }

    def unliftT(
      routes: HttpRoutes[F]
    )(implicit d: Defer[F], ev: Bracket[F, Throwable]): HttpRoutes[Kleisli[F, Span[F], ?]] =
      unliftTE(self)(routes)

    def unliftR(routes: Resource[F, HttpRoutes[F]])(
      implicit ev: Bracket[F, Throwable],
      d: Defer[F]
    ): Resource[Kleisli[F, Span[F], ?], HttpRoutes[Kleisli[F, Span[F], ?]]] =
      routes
        .mapK(λ[F ~> Kleisli[F, Span[F], ?]] { (fa: F[Any]) =>
          Kleisli(_ => fa)
        })
        .map(unliftT)
  }
}