Move span creation to middleware

msosnicki · msosnicki · commit a80c6a964677 · 2023-06-06T14:42:02.000+02:00
diff --git a/modules/http4s/src/main/scala/natchez/http4s/NatchezMiddleware.scala b/modules/http4s/src/main/scala/natchez/http4s/NatchezMiddleware.scala
@@ -10,13 +10,27 @@ import cats.effect.{MonadCancel, Outcome}
 import cats.effect.syntax.all._
 import Outcome._
 import org.http4s.HttpRoutes
-import natchez.{Trace, TraceValue, Tags}
+import org.typelevel.ci.CIString
+import natchez.{Kernel, Span, Trace, TraceValue, Tags}
+import org.http4s.Request
 import org.http4s.Response
 import org.http4s.client.Client
 import java.io.ByteArrayOutputStream
 import java.io.PrintStream
 import cats.effect.Resource
 
+/**
+ * @define excludedHeaders
+ *         All headers except security (Authorization, Cookie, Set-Cookie)
+ *         and payload (Content-Length, ContentType, Content-Range, Trailer, Transfer-Encoding)
+ *         are passed to Kernel by default.
+ *
+ * @define isKernelHeader should an HTTP header be passed to Kernel or not
+ *
+ * @define spanName compute the span name from the request
+ *
+ * @define modifySpanOptions modify default span creation options
+ */
 object NatchezMiddleware {
   import syntax.kernel._
 
@@ -27,7 +41,8 @@ object NatchezMiddleware {
     server(routes)
 
   /**
-   * A middleware that adds the following standard fields to the current span:
+   * A middleware that creates a per request span. 
+   * It also adds following standard fields to newly created span:
    *
    * - "http.method"      -> "GET", "PUT", etc.
    * - "http.url"         -> request URI (not URL)
@@ -39,8 +54,20 @@ object NatchezMiddleware {
    * - "error.message"    -> Exception message
    * - "error.stacktrace" -> Exception stack trace as a multi-line string
    * - "cancelled"        -> true // only present in case of cancellation
+   *
+   * @note $excludedHeaders
+   *
+   * @param isKernelHeader $isKernelHeader
+   * @param spanName $spanName
+   * @param modifySpanOptions $modifySpanOptions
    */
-  def server[F[_]: Trace](routes: HttpRoutes[F])(
+
+  def server[F[_]: Trace](
+    routes: HttpRoutes[F],
+    isKernelHeader: CIString => Boolean = name => !ExcludedHeaders.contains(name),
+    spanName: Request[F] => String = (req: Request[F]) => req.uri.path.toString,
+    modifySpanOptions: Span.Options => Span.Options = identity
+  )(
     implicit ev: MonadCancel[F, Throwable]
   ): HttpRoutes[F] =
     Kleisli { req =>
@@ -71,17 +98,25 @@ object NatchezMiddleware {
             new String(baos.toByteArray, "UTF-8")
           }
         )
-
-      routes(req).guaranteeCase {
-        case Canceled() => OptionT.liftF(addRequestFields *> Trace[F].put(("cancelled", TraceValue.BooleanValue(true)), Tags.error(true)))
-        case Errored(e) => OptionT.liftF(addRequestFields *> addErrorFields(e))
-        case Succeeded(fa) => OptionT.liftF {
-          fa.value.flatMap {
-            case Some(resp) => addRequestFields *> addResponseFields(resp)
-            case None => MonadCancel[F].unit
+      val kernelHeaders = req.headers.headers
+        .collect {
+          case header if isKernelHeader(header.name) => header.name -> header.value
+        }
+        .toMap
+      val kernel = Kernel(kernelHeaders)
+      OptionT(
+        Trace[F].span(spanName(req), modifySpanOptions(Span.Options.Defaults.withParentKernel(kernel).withSpanKind(Span.SpanKind.Server))) {
+          addRequestFields >> routes(req).value.guaranteeCase {
+            case Canceled() => Trace[F].put(("cancelled", TraceValue.BooleanValue(true)), Tags.error(true))
+            case Errored(e) => addErrorFields(e)
+            case Succeeded(fa) =>
+              fa.flatMap {
+                case Some(resp) => addResponseFields(resp)
+                case None => ev.unit
+              }
           }
         }
-      }
+      )
    }
 
   /**
@@ -113,4 +148,25 @@ object NatchezMiddleware {
       }
     }
 
+  val ExcludedHeaders: Set[CIString] = {
+    import org.http4s.headers._
+    import org.typelevel.ci._
+
+    val payload = Set(
+      `Content-Length`.name,
+      ci"Content-Type",
+      `Content-Range`.name,
+      ci"Trailer",
+      `Transfer-Encoding`.name,
+    )
+
+    val security = Set(
+      Authorization.name,
+      Cookie.name,
+      `Set-Cookie`.name,
+    )
+
+    payload ++ security
+  }
+
 }
diff --git a/modules/http4s/src/main/scala/natchez/http4s/syntax/EntryPointOps.scala b/modules/http4s/src/main/scala/natchez/http4s/syntax/EntryPointOps.scala
@@ -6,140 +6,69 @@ package natchez.http4s.syntax
 
 import cats.~>
 import cats.data.{ Kleisli, OptionT }
-import cats.data.Kleisli.applyK
 import cats.effect.MonadCancel
-import natchez.{ EntryPoint, Kernel, Span }
 import org.http4s.HttpRoutes
 import cats.effect.Resource
 import org.http4s.server.websocket.WebSocketBuilder2
 import org.typelevel.ci.CIString
+import natchez.EntryPoint
+import natchez.Span
 
-/**
- * @define excludedHeaders
- *         All headers except security (Authorization, Cookie, Set-Cookie)
- *         and payload (Content-Length, ContentType, Content-Range, Trailer, Transfer-Encoding)
- *         are passed to Kernel by default.
- *
- * @define isKernelHeader should an HTTP header be passed to Kernel or not
- *
- * @define spanName compute the span name from the request
- *
- * @define spanOptions options used in span creation
- */
 trait EntryPointOps[F[_]] { outer =>
 
   def self: EntryPoint[F]
 
   /**
-   * Given an entry point and HTTP Routes in Kleisli[F, Span[F], *] return routes in F. A new span
-   * is created with by default the URI path as the name, either as a continuation of the incoming trace, if
-   * any, or as a new root.
+   * Given an entry point and HTTP Routes in Kleisli[F, Span[F], *] return routes in F. 
+   * A span that is injected is a RootsSpan (smilarly to Trace.ioTraceForEntryPoint)
    *
-   * @note $excludedHeaders
-   *
-   * @param isKernelHeader $isKernelHeader
-   * @param spanName $spanName
-   * @param spanOptions $spanOptions
    */
   def liftT(
     routes: HttpRoutes[Kleisli[F, Span[F], *]],
-    isKernelHeader: CIString => Boolean = name => !EntryPointOps.ExcludedHeaders.contains(name),
-    spanName: org.http4s.Request[F] => String = _.uri.path.toString,
-    spanOptions: Span.Options = Span.Options.Defaults,
   )(implicit ev: MonadCancel[F, Throwable]): HttpRoutes[F] =
     Kleisli { req =>
-      val kernelHeaders = req.headers.headers
-        .collect {
-          case header if isKernelHeader(header.name) => header.name -> header.value
-        }
-        .toMap
-
-      val kernel = Kernel(kernelHeaders)
-      val spanR  = self.continueOrElseRoot(spanName(req), kernel, spanOptions)
-      OptionT {
-        spanR.use { span =>
-          routes.run(req.mapK(lift)).mapK(applyK(span)).map(_.mapK(applyK(span))).value
-        }
-      }
+      val root = Span.makeRoots(self)
+      OptionT(routes.run(req.mapK(Kleisli.liftK)).mapK(Kleisli.applyK(root)).map(_.mapK(Kleisli.applyK(root))).value)
     }
 
+
   /**
    * 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.
    *
-   * @note $excludedHeaders
-   *
-   * @param isKernelHeader $isKernelHeader
    */
   def liftR(
     routes: Resource[Kleisli[F, Span[F], *], HttpRoutes[Kleisli[F, Span[F], *]]],
-    isKernelHeader: CIString => Boolean = name => !EntryPointOps.ExcludedHeaders.contains(name)
   )(implicit ev: MonadCancel[F, Throwable]): Resource[F, HttpRoutes[F]] =
-    routes.map(liftT(_, isKernelHeader)).mapK(Span.dropTracing)
+    routes.map(liftT).mapK(Span.dropTracing)
 
   /**
    * Given an entry point and a function from `WebSocketBuilder2` to HTTP Routes in
-   * Kleisli[F, Span[F], *] return a function from `WebSocketBuilder2` to routes in F. A new span
-   * is created with the URI path as the name, either as a continuation of the incoming trace, if
-   * any, or as a new root.
-   *
-   * @note $excludedHeaders
-   *
-   * @param isKernelHeader $isKernelHeader
+   * Kleisli[F, Span[F], *] return a function from `WebSocketBuilder2` to routes in F.    
    */
   def wsLiftT(
-    routes: WebSocketBuilder2[Kleisli[F, Span[F], *]] => HttpRoutes[Kleisli[F, Span[F], *]],
-    isKernelHeader: CIString => Boolean = name => !EntryPointOps.ExcludedHeaders.contains(name),
-    spanName: org.http4s.Request[F] => String = _.uri.path.toString,
-    spanOptions: Span.Options = Span.Options.Defaults
+    routes: WebSocketBuilder2[Kleisli[F, Span[F], *]] => HttpRoutes[Kleisli[F, Span[F], *]]
    )(implicit ev: MonadCancel[F, Throwable]): WebSocketBuilder2[F] => HttpRoutes[F] = wsb =>
-    liftT(routes(wsb.imapK(lift)(Span.dropTracing)), isKernelHeader, spanName, spanOptions)
+    liftT(routes(wsb.imapK(lift)(Span.dropTracing)))
 
   /**
    * Lift a `WebSocketBuilder2 => 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.
-   *
-   * @note $excludedHeaders
-   *
-   * @param isKernelHeader $isKernelHeader
    */
   def wsLiftR(
-    routes: Resource[Kleisli[F, Span[F], *], WebSocketBuilder2[Kleisli[F, Span[F], *]] => HttpRoutes[Kleisli[F, Span[F], *]]],
-    isKernelHeader: CIString => Boolean = name => !EntryPointOps.ExcludedHeaders.contains(name),
-    spanName: org.http4s.Request[F] => String = _.uri.path.toString,
-    spanOptions: Span.Options = Span.Options.Defaults
+    routes: Resource[Kleisli[F, Span[F], *], WebSocketBuilder2[Kleisli[F, Span[F], *]] => HttpRoutes[Kleisli[F, Span[F], *]]]
   )(implicit ev: MonadCancel[F, Throwable]): Resource[F, WebSocketBuilder2[F] => HttpRoutes[F]] =
-    routes.map(wsLiftT(_, isKernelHeader, spanName, spanOptions)).mapK(Span.dropTracing)
+    routes.map(wsLiftT).mapK(Span.dropTracing)
 
   private val lift: F ~> Kleisli[F, Span[F], *] =
     Kleisli.liftK
 }
 
 object EntryPointOps {
-  val ExcludedHeaders: Set[CIString] = {
-    import org.http4s.headers._
-    import org.typelevel.ci._
-
-    val payload = Set(
-      `Content-Length`.name,
-      ci"Content-Type",
-      `Content-Range`.name,
-      ci"Trailer",
-      `Transfer-Encoding`.name,
-    )
-
-    val security = Set(
-      Authorization.name,
-      Cookie.name,
-      `Set-Cookie`.name,
-    )
-
-    payload ++ security
-  }
 }
 
 trait ToEntryPointOps {
diff --git a/modules/http4s/src/test/scala/natchez/http4s/NatchezMiddlewareSuite.scala b/modules/http4s/src/test/scala/natchez/http4s/NatchezMiddlewareSuite.scala
@@ -97,14 +97,14 @@ class NatchezMiddlewareSuite extends InMemorySuite {
 
       List(
         (Lineage.Root,                                          NatchezCommand.CreateRootSpan("/hello/some-name", requestKernel, Span.Options.Defaults)),
+        (Lineage.Root,                                          NatchezCommand.Put(requestTags)),
         (Lineage.Root,                                          NatchezCommand.CreateSpan("call-proxy", None, Span.Options.Defaults)),
         (Lineage.Root / "call-proxy",                           NatchezCommand.CreateSpan("http4s-client-request", None, Span.Options.Defaults)),
         (Lineage.Root / "call-proxy" / "http4s-client-request", NatchezCommand.AskKernel(requestKernel)),
         (Lineage.Root / "call-proxy" / "http4s-client-request", NatchezCommand.Put(clientRequestTags)),
         (Lineage.Root / "call-proxy" / "http4s-client-request", NatchezCommand.Put(clientResponseTags)),
         (Lineage.Root / "call-proxy",                           NatchezCommand.ReleaseSpan("http4s-client-request")),
         (Lineage.Root,                                          NatchezCommand.ReleaseSpan("call-proxy")),
-        (Lineage.Root,                                          NatchezCommand.Put(requestTags)),
         (Lineage.Root,                                          NatchezCommand.Put(responseTags)),
         (Lineage.Root,                                          NatchezCommand.ReleaseRootSpan("/hello/some-name"))
       )
