Skip to content

Commit

Permalink
Merge pull request #37 from Dwolla/additional-client-attributes
Browse files Browse the repository at this point in the history
Allow additional attributes to be added to client spans
  • Loading branch information
armanbilge authored Aug 29, 2024
2 parents 9af4cab + 934c815 commit 0f0f9dd
Show file tree
Hide file tree
Showing 5 changed files with 138 additions and 71 deletions.
20 changes: 11 additions & 9 deletions build.sbt
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
ThisBuild / tlBaseVersion := "0.6"

val http4sVersion = "0.23.17"
val natchezVersion = "0.3.5"
val scala212Version = "2.12.17"
val scala213Version = "2.13.10"
val scala3Version = "3.3.3"
val slf4jVersion = "2.0.4"
val munitCEVersion = "2.0.0-M3"
val http4sVersion = "0.23.17"
val natchezVersion = "0.3.5"
val scala212Version = "2.12.17"
val scala213Version = "2.13.10"
val scala3Version = "3.3.3"
val slf4jVersion = "2.0.4"
val munitCEVersion = "2.0.0-M3"
val scalacheckEffectVersion = "2.0.0-M2"

ThisBuild / organization := "org.tpolecat"
ThisBuild / tlSonatypeUseLegacyHost := false
Expand All @@ -29,8 +30,9 @@ lazy val commonSettings = Seq(
),

libraryDependencies ++= Seq(
"org.typelevel" %%% "munit-cats-effect" % munitCEVersion % Test,
"org.http4s" %%% "http4s-dsl" % http4sVersion % Test,
"org.typelevel" %%% "munit-cats-effect" % munitCEVersion % Test,
"org.http4s" %%% "http4s-dsl" % http4sVersion % Test,
"org.typelevel" %%% "scalacheck-effect-munit" % scalacheckEffectVersion % Test,
)
)

Expand Down
3 changes: 2 additions & 1 deletion modules/examples/src/main/scala/Example1.scala
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ package example
import cats.effect._
import cats.syntax.all._
import com.comcast.ip4s.Port
import fs2.io.net.Network
import natchez.http4s.implicits._
import org.http4s.ember.server.EmberServerBuilder
import org.http4s.server.Server
Expand All @@ -34,7 +35,7 @@ import natchez.http4s.NatchezMiddleware
object Http4sExample extends IOApp with Common {

// Our main app resource
def server[F[_]: Async]: Resource[F, Server] =
def server[F[_]: Async : Network]: Resource[F, Server] =
for {
ep <- entryPoint[F]
ap = ep.liftT(NatchezMiddleware.server(routes)).orNotFound // liftT discharges the Trace constraint
Expand Down
4 changes: 3 additions & 1 deletion modules/examples/src/main/scala/Example2.scala
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,15 @@ import fs2.Stream
import org.http4s.ember.server.EmberServerBuilder
import org.http4s.implicits._
import org.http4s.server.middleware.Logger

import scala.io.StdIn
import com.comcast.ip4s.Host
import com.comcast.ip4s.Port
import fs2.io.net.Network

object Http4sExampleStreamed extends IOApp with Common {

def stream[F[_]: Async]: Stream[F, Nothing] = {
def stream[F[_]: Async : Network]: Stream[F, Nothing] = {
for {
ep <- Stream.resource(entryPoint[F])
finalRoutes = ep.liftT(NatchezMiddleware.server(routes))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,20 @@

package natchez.http4s

import cats.data.{ Kleisli, OptionT }
import cats.data.{Kleisli, OptionT}
import cats.syntax.all._
import cats.effect.{MonadCancel, Outcome}
import cats.effect.{MonadCancel, MonadCancelThrow, Outcome, Resource}
import cats.effect.syntax.all._
import Outcome._
import org.http4s.HttpRoutes
import natchez.{Trace, TraceValue, Tags}
import org.http4s.Response
import natchez.{Tags, Trace, TraceValue}
import natchez.Span.Options.Defaults
import natchez.Span.SpanKind
import org.http4s.client.Client
import org.http4s.HttpRoutes
import org.http4s.{Request, Response}

import java.io.ByteArrayOutputStream
import java.io.PrintStream
import cats.effect.Resource

object NatchezMiddleware {
import syntax.kernel._
Expand Down Expand Up @@ -92,18 +94,53 @@ object NatchezMiddleware {
* - "client.http.status_code" -> "200", "403", etc. // why is this a string?
*
*/
def client[F[_]: Trace](client: Client[F])(
implicit ev: MonadCancel[F, Throwable]
): Client[F] =
def client[F[_] : Trace : MonadCancelThrow](client: Client[F]): Client[F] =
NatchezMiddleware.client(client, _ => Seq.empty[(String, TraceValue)].pure[F])

/**
* A middleware that adds the current span's kernel to outgoing requests, performs requests in
* a span called `http4s-client-request`, and adds the following fields to that span.
*
* - "client.http.method" -> "GET", "PUT", etc.
* - "client.http.uri" -> request URI
* - "client.http.status_code" -> "200", "403", etc. // why is this a string?
*
* @param client the `Client[F]` to be enhanced
* @param additionalAttributes additional attributes to be added to the span
* @tparam F An effect with instances of `Trace[F]` and `MonadCancelThrow[F]`
* @return the enhanced `Client[F]`
*/
def clientWithAttributes[F[_] : Trace : MonadCancelThrow](client: Client[F])
(additionalAttributes: (String, TraceValue)*): Client[F] =
NatchezMiddleware.client(client, (_: Request[F]) => additionalAttributes.pure[F])

/**
* A middleware that adds the current span's kernel to outgoing requests, performs requests in
* a span called `http4s-client-request`, and adds the following fields to that span.
*
* - "client.http.method" -> "GET", "PUT", etc.
* - "client.http.uri" -> request URI
* - "client.http.status_code" -> "200", "403", etc. // why is this a string?
*
* @param client the `Client[F]` to be enhanced
* @param additionalAttributesF a function that takes the `Request[F]` and returns any additional attributes to be added to the span
* @tparam F An effect with instances of `Trace[F]` and `MonadCancelThrow[F]`
* @return the enhanced `Client[F]`
*/
def client[F[_] : Trace : MonadCancelThrow](client: Client[F],
additionalAttributesF: Request[F] => F[Seq[(String, TraceValue)]],
): Client[F] =
Client { req =>
Resource.applyFull {poll =>
Trace[F].span("http4s-client-request") {
Trace[F].span("http4s-client-request", Defaults.withSpanKind(SpanKind.Client)) {
for {
knl <- Trace[F].kernel
_ <- Trace[F].put(
"client.http.uri" -> req.uri.toString(),
"client.http.method" -> req.method.toString
)
additionalAttributes <- additionalAttributesF(req)
_ <- Trace[F].put(additionalAttributes: _*)
reqʹ = req.withHeaders(knl.toHttp4sHeaders ++ req.headers) // prioritize request headers over kernel ones
rsrc <- poll(client.run(reqʹ).allocatedCase)
_ <- Trace[F].put("client.http.status_code" -> rsrc._1.status.code.toString())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,43 @@ package natchez.http4s
import cats.Monad
import cats.data.{Chain, Kleisli}
import cats.effect.{IO, MonadCancelThrow, Resource}
import natchez.{InMemory, Kernel, Span, Trace}
import munit.ScalaCheckEffectSuite
import natchez.Span.SpanKind
import natchez.{InMemory, Kernel, Span, Trace, TraceValue}
import natchez.TraceValue.StringValue
import natchez.http4s.syntax.entrypoint._
import org.http4s._
import org.http4s.headers._
import org.http4s.client.Client
import org.http4s.dsl.request._
import org.http4s.syntax.literals._
import org.scalacheck.{Arbitrary, Gen}
import org.scalacheck.Arbitrary.arbitrary
import org.scalacheck.effect.PropF
import org.typelevel.ci._

class NatchezMiddlewareSuite extends InMemorySuite {
class NatchezMiddlewareSuite
extends InMemorySuite
with ScalaCheckEffectSuite {

private val CustomHeaderName = ci"X-Custom-Header"
private val CorrelationIdName = ci"X-Correlation-Id"

private implicit val arbTraceValue: Arbitrary[TraceValue] = Arbitrary {
Gen.oneOf(
arbitrary[String].map(TraceValue.StringValue(_)),
arbitrary[Boolean].map(TraceValue.BooleanValue(_)),
arbitrary[Number].map(TraceValue.NumberValue(_)),
)
}

private implicit val arbAttribute: Arbitrary[(String, TraceValue)] = Arbitrary {
for {
key <- arbitrary[String]
value <- arbitrary[TraceValue]
} yield key -> value
}

test("do not leak security and payload headers to the client request") {
val headers = Headers(
// security
Expand Down Expand Up @@ -54,7 +76,7 @@ class NatchezMiddlewareSuite extends InMemorySuite {
for {
ref <- IO.ref(Chain.empty[(Lineage, NatchezCommand)])
ep <- IO.pure(new InMemory.EntryPoint(ref))
routes <- IO.pure(ep.liftT(httpRoutes[Kleisli[IO, natchez.Span[IO], *]]))
routes <- IO.pure(ep.liftT(httpRoutes[Kleisli[IO, natchez.Span[IO], *]]()))
response <- routes.orNotFound.run(request)
} yield {
assertEquals(response.status.code, 200)
Expand All @@ -63,63 +85,66 @@ class NatchezMiddlewareSuite extends InMemorySuite {
}

test("generate proper tracing history") {
val request = Request[IO](
method = Method.GET,
uri = uri"/hello/some-name",
headers = Headers(
Header.Raw(CustomHeaderName, "external"),
Header.Raw(CorrelationIdName, "id-123")
PropF.forAllF { (userSpecifiedTags: List[(String, TraceValue)]) =>
val request = Request[IO](
method = Method.GET,
uri = uri"/hello/some-name",
headers = Headers(
Header.Raw(CustomHeaderName, "external"),
Header.Raw(CorrelationIdName, "id-123")
)
)
)

val expectedHistory = {
val requestKernel = Kernel(
Map(CustomHeaderName -> "external", CorrelationIdName -> "id-123")
)
val expectedHistory = {
val requestKernel = Kernel(
Map(CustomHeaderName -> "external", CorrelationIdName -> "id-123")
)

val clientRequestTags = List(
"client.http.uri" -> StringValue("/some-name"),
"client.http.method" -> StringValue("GET")
)
val clientRequestTags = List(
"client.http.uri" -> StringValue("/some-name"),
"client.http.method" -> StringValue("GET")
)

val clientResponseTags = List(
"client.http.status_code" -> StringValue("200")
)
val clientResponseTags = List(
"client.http.status_code" -> StringValue("200")
)

val requestTags = List(
"http.method" -> StringValue("GET"),
"http.url" -> StringValue("/hello/some-name")
)
val requestTags = List(
"http.method" -> StringValue("GET"),
"http.url" -> StringValue("/hello/some-name")
)

val responseTags = List(
"http.status_code" -> StringValue("200")
)
val responseTags = List(
"http.status_code" -> StringValue("200")
)

List(
(Lineage.Root, NatchezCommand.CreateRootSpan("/hello/some-name", requestKernel, Span.Options.Defaults)),
(Lineage.Root("/hello/some-name"), NatchezCommand.CreateSpan("call-proxy", None, Span.Options.Defaults)),
(Lineage.Root("/hello/some-name") / "call-proxy", NatchezCommand.CreateSpan("http4s-client-request", None, Span.Options.Defaults)),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.AskKernel(requestKernel)),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.Put(clientRequestTags)),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.Put(clientResponseTags)),
(Lineage.Root("/hello/some-name") / "call-proxy", NatchezCommand.ReleaseSpan("http4s-client-request")),
(Lineage.Root("/hello/some-name"), NatchezCommand.ReleaseSpan("call-proxy")),
(Lineage.Root("/hello/some-name"), NatchezCommand.Put(requestTags)),
(Lineage.Root("/hello/some-name"), NatchezCommand.Put(responseTags)),
(Lineage.Root, NatchezCommand.ReleaseRootSpan("/hello/some-name"))
)
List(
(Lineage.Root, NatchezCommand.CreateRootSpan("/hello/some-name", requestKernel, Span.Options.Defaults)),
(Lineage.Root("/hello/some-name"), NatchezCommand.CreateSpan("call-proxy", None, Span.Options.Defaults)),
(Lineage.Root("/hello/some-name") / "call-proxy", NatchezCommand.CreateSpan("http4s-client-request", None, Span.Options.Defaults.withSpanKind(SpanKind.Client))),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.AskKernel(requestKernel)),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.Put(clientRequestTags)),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.Put(userSpecifiedTags)),
(Lineage.Root("/hello/some-name") / "call-proxy" / "http4s-client-request", NatchezCommand.Put(clientResponseTags)),
(Lineage.Root("/hello/some-name") / "call-proxy", NatchezCommand.ReleaseSpan("http4s-client-request")),
(Lineage.Root("/hello/some-name"), NatchezCommand.ReleaseSpan("call-proxy")),
(Lineage.Root("/hello/some-name"), NatchezCommand.Put(requestTags)),
(Lineage.Root("/hello/some-name"), NatchezCommand.Put(responseTags)),
(Lineage.Root, NatchezCommand.ReleaseRootSpan("/hello/some-name"))
)
}

for {
ep <- InMemory.EntryPoint.create[IO]
routes <- IO.pure(ep.liftT(httpRoutes[Kleisli[IO, natchez.Span[IO], *]](userSpecifiedTags: _*)))
_ <- routes.orNotFound.run(request)
history <- ep.ref.get
} yield assertEquals(history.toList, expectedHistory)
}

for {
ep <- InMemory.EntryPoint.create[IO]
routes <- IO.pure(ep.liftT(httpRoutes[Kleisli[IO, natchez.Span[IO], *]]))
_ <- routes.orNotFound.run(request)
history <- ep.ref.get
} yield assertEquals(history.toList, expectedHistory)
}

private def httpRoutes[F[_]: MonadCancelThrow: Trace]: HttpRoutes[F] = {
val client = NatchezMiddleware.client(echoHeadersClient[F])
private def httpRoutes[F[_]: MonadCancelThrow: Trace](additionalAttributes: (String, TraceValue)*): HttpRoutes[F] = {
val client = NatchezMiddleware.clientWithAttributes(echoHeadersClient[F])(additionalAttributes: _*)
val server = NatchezMiddleware.server(proxyRoutes(client))
server
}
Expand Down

0 comments on commit 0f0f9dd

Please sign in to comment.