Merge pull request #174 from trace4cats/feature/extra-attributes

Allow to inject extra attributes
trace4cats · Jun 16, 2023 · 12ab574 · 12ab574
2 parents d0959a0 + 8e4c604
commit 12ab574
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 8 deletions.
diff --git a/modules/http4s-client/src/main/scala/trace4cats/http4s/client/ClientTracer.scala b/modules/http4s-client/src/main/scala/trace4cats/http4s/client/ClientTracer.scala
@@ -2,12 +2,13 @@ package trace4cats.http4s.client
 
 import cats.effect.kernel.{MonadCancelThrow, Resource}
 import cats.syntax.flatMap._
+import cats.syntax.option._
 import org.http4s.client.{Client, UnexpectedStatus}
 import org.http4s.{Request, Response}
 import org.typelevel.ci.CIString
 import trace4cats.Span
 import trace4cats.context.Provide
-import trace4cats.http4s.common.{Http4sHeaders, Http4sSpanNamer, Http4sStatusMapping, Request_, Response_}
+import trace4cats.http4s.common._
 import trace4cats.model.{AttributeValue, SampleDecision, SpanKind, TraceHeaders}
 import trace4cats.optics.{Getter, Lens}
 
@@ -44,7 +45,8 @@ object ClientTracer {
               // only extract request attributes if the span is sampled as the host parsing is quite expensive
               val reqExtraAttrs =
                 if (isSampled)
-                  Http4sClientRequest.toAttributes(request) ++ requestAttributesGetter.get(request)
+                  Http4sClientRequest.toAttributes(request) ++ requestAttributesGetter.get(request) ++
+                    request.attributes.lookup(Http4sAttributes.Keys.ExtraRequestAttributes).map(_.value).orEmpty
                 else Map.empty
 
               for {
@@ -53,7 +55,11 @@ object ClientTracer {
                 res <- runClient(req.mapK(P.provideK(childCtx)))
                   .evalTap { resp =>
                     val respHeaderAttrs = Http4sHeaders.responseFields(resp, dropHeadersWhen)
-                    val respExtraAttrs = if (isSampled) responseAttributesGetter.get(resp.mapK(P.liftK)) else Map.empty
+                    val respExtraAttrs =
+                      if (isSampled)
+                        responseAttributesGetter.get(resp.mapK(P.liftK)) ++
+                          resp.attributes.lookup(Http4sAttributes.Keys.ExtraResponseAttributes).map(_.value).orEmpty
+                      else Map.empty
                     childSpan.setStatus(Http4sStatusMapping.toSpanStatus(resp.status)) >>
                       childSpan.putAll(respHeaderAttrs ++ respExtraAttrs: _*)
                   }

diff --git a/modules/http4s-client/src/test/scala/trace4cats/http4s/client/BaseClientTracerSpec.scala b/modules/http4s-client/src/test/scala/trace4cats/http4s/client/BaseClientTracerSpec.scala
@@ -3,7 +3,7 @@ package trace4cats.http4s.client
 import cats.data.NonEmptyList
 import cats.effect.kernel.{Async, Ref}
 import cats.implicits._
-import cats.{~>, Eq, Id}
+import cats.{~>, Eq, Eval, Id}
 import org.http4s._
 import org.http4s.client.Client
 import org.http4s.client.dsl.Http4sClientDsl
@@ -16,7 +16,7 @@ import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 import org.scalatestplus.scalacheck.ScalaCheckDrivenPropertyChecks
 import trace4cats.context.Provide
-import trace4cats.http4s.common.{Http4sHeaders, Http4sStatusMapping}
+import trace4cats.http4s.common.{Http4sAttributes, Http4sHeaders, Http4sStatusMapping}
 import trace4cats.kernel.{SpanCompleter, SpanSampler}
 import trace4cats.model.TraceHeaders
 import trace4cats._
@@ -67,7 +67,13 @@ abstract class BaseClientTracerSpec[F[_]: Async, G[_]: Async: Trace, Ctx](
 
       unsafeRunK(RefSpanCompleter[F]("test").flatMap { completer =>
         withClient(httpApp) { client =>
-          def req(body: String): G[Unit] = runReq(client, GET(body, Uri.unsafeFromString(s"/")))
+          def req(body: String): G[Unit] = {
+            val req = GET(body, Uri.unsafeFromString(s"/")).withAttribute(
+              Http4sAttributes.Keys.ExtraRequestAttributes,
+              Eval.now(Map("test" -> AttributeValue.stringToTraceValue("works")))
+            )
+            runReq(client, req)
+          }
 
           for {
             _ <- entryPoint(completer)
@@ -100,6 +106,8 @@ abstract class BaseClientTracerSpec[F[_]: Async, G[_]: Async: Trace, Ctx](
               )
             )
 
+            spans.head.attributes("test").toString() should be("works")
+
             val expectedStatus = Http4sStatusMapping.toSpanStatus(response.status)
             (spans.toList.collect {
               case span if span.name == "GET /" => span.status

diff --git a/modules/http4s-common/src/main/scala/trace4cats/http4s/common/Http4sAttributes.scala b/modules/http4s-common/src/main/scala/trace4cats/http4s/common/Http4sAttributes.scala
@@ -1,10 +1,16 @@
 package trace4cats.http4s.common
 
+import cats.Eval
 import cats.effect.SyncIO
 import org.typelevel.vault.Key
+import trace4cats.model.AttributeValue
 
 object Http4sAttributes {
   object Keys {
     val SpanName: Key[String] = Key.newKey[SyncIO, String].unsafeRunSync()
+    val ExtraRequestAttributes: Key[Eval[Map[String, AttributeValue]]] =
+      Key.newKey[SyncIO, Eval[Map[String, AttributeValue]]].unsafeRunSync()
+    val ExtraResponseAttributes: Key[Eval[Map[String, AttributeValue]]] =
+      Key.newKey[SyncIO, Eval[Map[String, AttributeValue]]].unsafeRunSync()
   }
 }
diff --git a/modules/http4s-server/src/main/scala/trace4cats/http4s/server/ServerTracer.scala b/modules/http4s-server/src/main/scala/trace4cats/http4s/server/ServerTracer.scala
@@ -5,11 +5,12 @@ import cats.effect.kernel.MonadCancelThrow
 import cats.syntax.apply._
 import cats.syntax.flatMap._
 import cats.syntax.functor._
+import cats.syntax.option._
 import cats.{FlatMap, Monad}
 import org.http4s.{HttpApp, HttpRoutes, Request, Response}
 import org.typelevel.ci.CIString
 import trace4cats.context.Provide
-import trace4cats.http4s.common.{Http4sHeaders, Http4sStatusMapping, Request_, Response_}
+import trace4cats.http4s.common.{Http4sAttributes, Http4sHeaders, Http4sStatusMapping, Request_, Response_}
 import trace4cats.{ResourceKleisli, Trace}
 
 object ServerTracer {
@@ -23,13 +24,19 @@ object ServerTracer {
         for {
           ctx <- P.ask[Ctx]
           _ <- Trace[G].putAll(Http4sHeaders.requestFields(req, dropHeadersWhen): _*)
+          _ <- Trace[G].putAll(
+            req.attributes.lookup(Http4sAttributes.Keys.ExtraRequestAttributes).map(_.value).orEmpty.toSeq: _*
+          )
           resp <-
             routes
               .run(req.mapK(P.liftK))
               .map(_.mapK(P.provideK(ctx)))
               .semiflatTap { res =>
                 Trace[G].setStatus(Http4sStatusMapping.toSpanStatus(res.status)) *>
-                  Trace[G].putAll(Http4sHeaders.responseFields(res, dropHeadersWhen): _*)
+                  Trace[G].putAll(Http4sHeaders.responseFields(res, dropHeadersWhen): _*) *>
+                  Trace[G].putAll(
+                    res.attributes.lookup(Http4sAttributes.Keys.ExtraResponseAttributes).map(_.value).orEmpty.toSeq: _*
+                  )
               }
               .value
         } yield resp
@@ -49,9 +56,15 @@ object ServerTracer {
         for {
           ctx <- P.ask[Ctx]
           _ <- Trace[G].putAll(Http4sHeaders.requestFields(req, dropHeadersWhen): _*)
+          _ <- Trace[G].putAll(
+            req.attributes.lookup(Http4sAttributes.Keys.ExtraRequestAttributes).map(_.value).orEmpty.toSeq: _*
+          )
           resp <- app.run(req.mapK(P.liftK)).map(_.mapK(P.provideK(ctx)))
           _ <- Trace[G].setStatus(Http4sStatusMapping.toSpanStatus(resp.status))
           _ <- Trace[G].putAll(Http4sHeaders.responseFields(resp, dropHeadersWhen): _*)
+          _ <- Trace[G].putAll(
+            resp.attributes.lookup(Http4sAttributes.Keys.ExtraResponseAttributes).map(_.value).orEmpty.toSeq: _*
+          )
         } yield resp
 
       k(req).use(P.provide(fa))