Skip to content

Commit 0d4cc2a

Browse files
cquirozcquiroz
committed
Add W3C trace propagation for WebSocket passing the parent on extensions.
1 parent 30632e3 commit 0d4cc2a

7 files changed

Lines changed: 114 additions & 40 deletions

File tree

‎core/src/main/scala/clue/TraceHeaderInjector.scala‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// Copyright (c) 2016-2025 Association of Universities for Research in Astronomy, Inc. (AURA)
2+
// For license information see LICENSE or https://opensource.org/licenses/BSD-3-Clause
3+
4+
package clue
5+
6+
trait TraceHeaderInjector[P]:
7+
def addHeaders(params: P, headers: Map[String, String]): P
8+
9+
object TraceHeaderInjector:
10+
def apply[P](using ev: TraceHeaderInjector[P]): TraceHeaderInjector[P] = ev
11+
12+
given TraceHeaderInjector[Unit] with
13+
def addHeaders(params: Unit, headers: Map[String, String]): Unit = ()

‎core/src/main/scala/clue/clients.scala‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,8 @@ trait StreamingClient[F[_], S] extends FetchClientWithPars[F, Unit, S] {
9191
protected[clue] def subscribeInternal[D: Decoder](
9292
document: GraphQLQuery,
9393
operationName: Option[String] = none,
94-
variables: Option[JsonObject] = none
94+
variables: Option[JsonObject] = none,
95+
extensions: Option[JsonObject] = none
9596
): Resource[F, fs2.Stream[F, GraphQLResponse[D]]]
9697
}
9798

@@ -109,11 +110,12 @@ case class SubscriptionApplied[
109110
client.subscribeInternal(
110111
GraphQLQuery(subscription.document),
111112
operationName,
112-
variables.asJsonObject.some
113+
variables.asJsonObject.some,
114+
none
113115
)
114116

115117
def apply: Resource[F, fs2.Stream[F, GraphQLResponse[D]]] =
116-
client.subscribeInternal(GraphQLQuery(subscription.document), operationName, none)
118+
client.subscribeInternal(GraphQLQuery(subscription.document), operationName, none, none)
117119
}
118120

119121
object SubscriptionApplied {

‎core/src/main/scala/clue/websocket/ApolloClient.scala‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -115,9 +115,10 @@ class ApolloClient[F[_], P, S](
115115
override protected[clue] def subscribeInternal[D: Decoder](
116116
subscription: GraphQLQuery,
117117
operationName: Option[String],
118-
variables: Option[JsonObject]
118+
variables: Option[JsonObject],
119+
extensions: Option[JsonObject]
119120
): Resource[F, fs2.Stream[F, GraphQLResponse[D]]] =
120-
subscriptionResource(subscription, operationName, variables)
121+
subscriptionResource(subscription, operationName, variables, extensions)
121122

122123
// <FetchClient>
123124
override protected[clue] def requestInternal[D: Decoder](
@@ -127,7 +128,7 @@ class ApolloClient[F[_], P, S](
127128
modParams: Unit => Unit // This is ignored here.
128129
): F[GraphQLResponse[D]] =
129130
F.async(cb =>
130-
startSubscription[D](document, operationName, variables)
131+
startSubscription[D](document, operationName, variables, none)
131132
.flatMap(subscription =>
132133
subscription.stream.attempt.head.compile.onlyOrError.attempt
133134
.map(onlyOrError => cb(onlyOrError.flatten))
@@ -450,10 +451,11 @@ class ApolloClient[F[_], P, S](
450451
private def subscriptionResource[D: Decoder, R](
451452
subscription: GraphQLQuery,
452453
operationName: Option[String],
453-
variables: Option[JsonObject]
454+
variables: Option[JsonObject],
455+
extensions: Option[JsonObject]
454456
): Resource[F, fs2.Stream[F, GraphQLResponse[D]]] =
455457
Resource
456-
.make(startSubscription[D](subscription, operationName, variables))(
458+
.make(startSubscription[D](subscription, operationName, variables, extensions))(
457459
_.stop()
458460
.handleErrorWith(_.logF("Error stopping subscription"))
459461
)
@@ -462,11 +464,12 @@ class ApolloClient[F[_], P, S](
462464
private def startSubscription[D: Decoder](
463465
subscription: GraphQLQuery,
464466
operationName: Option[String],
465-
variables: Option[JsonObject]
467+
variables: Option[JsonObject],
468+
extensions: Option[JsonObject]
466469
): F[GraphQLSubscription[F, GraphQLResponse[D]]] =
467470
state.get.flatMap {
468471
case Connected(_, _, _, _) =>
469-
val request = GraphQLRequest(subscription, operationName, variables)
472+
val request = GraphQLRequest(subscription, operationName, variables, extensions)
470473

471474
buildQueue[D](request).flatMap { case (id, emitter) =>
472475
def acquire: F[Unit] =
@@ -522,7 +525,7 @@ class ApolloClient[F[_], P, S](
522525
}
523526
case Connecting(_, _, _, _, latch) =>
524527
latch.resolve >>
525-
startSubscription(subscription, operationName, variables)
528+
startSubscription(subscription, operationName, variables, extensions)
526529
case _ =>
527530
ConnectionNotInitializedException.logAndRaiseF_
528531
}

‎http4s/src/main/scala/clue/http4s/package.scala‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,16 @@
33

44
package clue
55

6+
import org.http4s.Header
67
import org.http4s.Request
8+
import org.typelevel.ci.CIString
79

810
package object http4s {
911
type Http4sHttpClient[F[_], S] = FetchClientWithPars[F, Request[F], S]
12+
13+
given [F[_]]: TraceHeaderInjector[Request[F]] with
14+
def addHeaders(params: Request[F], headers: Map[String, String]): Request[F] =
15+
headers.foldLeft(params) { case (req, (k, v)) =>
16+
req.putHeaders(Header.Raw(CIString(k), v))
17+
}
1018
}

‎natchez/src/main/scala/clue/natchez/NatchezMiddleware.scala‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,8 @@ class NatchezStreamingClient[F[_]: Trace: MonadCancelThrow, S](
136136
protected[clue] def subscribeInternal[D: Decoder](
137137
document: GraphQLQuery,
138138
operationName: Option[String] = none,
139-
variables: Option[JsonObject] = none
139+
variables: Option[JsonObject] = none,
140+
extensions: Option[JsonObject] = none
140141
): Resource[F, fs2.Stream[F, GraphQLResponse[D]]] =
141142
val resource: Resource[F, fs2.Stream[F, GraphQLResponse[D]]] =
142143
Resource.applyFull: poll =>
@@ -146,7 +147,7 @@ class NatchezStreamingClient[F[_]: Trace: MonadCancelThrow, S](
146147
_ <- Trace[F].put(additionalAttributes*)
147148
result <- poll:
148149
wrapped
149-
.subscribeInternal[D](document, operationName, variables)
150+
.subscribeInternal[D](document, operationName, variables, extensions)
150151
.allocatedCase
151152
yield result
152153
resource.onFinalizeCase: exitCase =>

‎otel4s/src/main/scala/clue/otel4s/Otel4sMiddleware.scala‎

Lines changed: 67 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package clue.otel4s
55

66
import cats.Applicative
7+
import cats.effect.Concurrent
78
import cats.effect.MonadCancelThrow
89
import cats.effect.Resource
910
import cats.syntax.all.*
@@ -12,12 +13,14 @@ import clue.FetchClientWithPars
1213
import clue.PersistentClientStatus
1314
import clue.PersistentStreamingClient
1415
import clue.StreamingClient
16+
import clue.TraceHeaderInjector
1517
import clue.model.GraphQLQuery
1618
import clue.model.GraphQLResponse
1719
import io.circe.Decoder
1820
import io.circe.Encoder
1921
import io.circe.JsonObject
2022
import org.typelevel.otel4s.Attribute
23+
import org.typelevel.otel4s.context.propagation.TextMapUpdater
2124
import org.typelevel.otel4s.semconv.attributes.HttpAttributes
2225
import org.typelevel.otel4s.semconv.experimental.attributes.GraphqlExperimentalAttributes
2326
import org.typelevel.otel4s.trace.Span
@@ -45,40 +48,40 @@ object Otel4sMiddleware:
4548
(_, _) => List.empty[Attribute[?]].pure[F]
4649

4750
// Fetch client
48-
def apply[F[_]: Tracer: MonadCancelThrow, P, S](
51+
def apply[F[_]: Tracer: MonadCancelThrow, P: TraceHeaderInjector, S](
4952
client: FetchClientWithPars[F, P, S],
5053
spanBuilderMod: SpanBuilderMod[F],
5154
additionalAttributesF: (GraphQLQuery, Option[JsonObject]) => F[List[Attribute[?]]]
5255
): FetchClientWithPars[F, P, S] =
5356
Otel4sFetchClient[F, P, S](client, spanBuilderMod, additionalAttributesF)
5457

55-
def apply[F[_]: Tracer: MonadCancelThrow, P, S](
58+
def apply[F[_]: Tracer: MonadCancelThrow, P: TraceHeaderInjector, S](
5659
client: FetchClientWithPars[F, P, S]
5760
): FetchClientWithPars[F, P, S] =
5861
apply(client, identityMod[F], emptyAttrs[F])
5962

6063
// Streaming client
61-
def apply[F[_]: Tracer: MonadCancelThrow, S](
64+
def apply[F[_]: Tracer: Concurrent, S](
6265
client: StreamingClient[F, S],
6366
spanBuilderMod: SpanBuilderMod[F],
6467
additionalAttributesF: (GraphQLQuery, Option[JsonObject]) => F[List[Attribute[?]]]
6568
): StreamingClient[F, S] =
6669
Otel4sStreamingClient(client, spanBuilderMod, additionalAttributesF)
6770

68-
def apply[F[_]: Tracer: MonadCancelThrow, S](
71+
def apply[F[_]: Tracer: Concurrent, S](
6972
client: StreamingClient[F, S]
7073
): StreamingClient[F, S] =
7174
apply(client, identityMod[F], emptyAttrs[F])
7275

7376
// Persistent streaming client
74-
def apply[F[_]: Tracer: MonadCancelThrow, S, CP, CE](
77+
def apply[F[_]: Tracer: Concurrent, S, CP, CE](
7578
client: PersistentStreamingClient[F, S, CP, CE],
7679
spanBuilderMod: SpanBuilderMod[F],
7780
additionalAttributesF: (GraphQLQuery, Option[JsonObject]) => F[List[Attribute[?]]]
7881
): PersistentStreamingClient[F, S, CP, CE] =
7982
Otel4sPersistentStreamingClient(client, spanBuilderMod, additionalAttributesF)
8083

81-
def apply[F[_]: Tracer: MonadCancelThrow, S, CP, CE](
84+
def apply[F[_]: Tracer: Concurrent, S, CP, CE](
8285
client: PersistentStreamingClient[F, S, CP, CE]
8386
): PersistentStreamingClient[F, S, CP, CE] =
8487
apply(client, identityMod[F], emptyAttrs[F])
@@ -113,7 +116,7 @@ object Otel4sMiddleware:
113116
Attribute("clue.response.errors", errorsStr)
114117
) *> span.setStatus(StatusCode.Error, "GraphQL request returned errors")
115118

116-
class Otel4sFetchClient[F[_]: Tracer: MonadCancelThrow, P, S](
119+
class Otel4sFetchClient[F[_]: Tracer: MonadCancelThrow, P: TraceHeaderInjector, S](
117120
wrapped: FetchClientWithPars[F, P, S],
118121
spanBuilderMod: Otel4sMiddleware.SpanBuilderMod[F],
119122
additionalAttributesF: (GraphQLQuery, Option[JsonObject]) => F[List[Attribute[?]]]
@@ -132,13 +135,16 @@ class Otel4sFetchClient[F[_]: Tracer: MonadCancelThrow, P, S](
132135
.addAttributes(Otel4sMiddleware.commonAttributes(document, operationName)*)
133136
).build.use: span =>
134137
for
135-
additional <- additionalAttributesF(document, variables)
136-
_ <- span.addAttributes(additional*)
137-
result <- poll(wrapped.requestInternal[D](document, operationName, variables, modParams))
138-
_ <- Otel4sMiddleware.recordResponseAttributes(span, result)
138+
additional <- additionalAttributesF(document, variables)
139+
_ <- span.addAttributes(additional*)
140+
traceHeaders <- Tracer[F].propagate(Map.empty[String, String])
141+
modWithTrace = modParams.andThen(p => TraceHeaderInjector[P].addHeaders(p, traceHeaders))
142+
result <-
143+
poll(wrapped.requestInternal[D](document, operationName, variables, modWithTrace))
144+
_ <- Otel4sMiddleware.recordResponseAttributes(span, result)
139145
yield result
140146

141-
class Otel4sStreamingClient[F[_]: Tracer: MonadCancelThrow, S](
147+
class Otel4sStreamingClient[F[_]: Tracer: Concurrent, S](
142148
wrapped: StreamingClient[F, S],
143149
spanBuilderMod: Otel4sMiddleware.SpanBuilderMod[F],
144150
additionalAttributesF: (GraphQLQuery, Option[JsonObject]) => F[List[Attribute[?]]]
@@ -147,29 +153,63 @@ class Otel4sStreamingClient[F[_]: Tracer: MonadCancelThrow, S](
147153
protected[clue] def subscribeInternal[D: Decoder](
148154
document: GraphQLQuery,
149155
operationName: Option[String] = none,
150-
variables: Option[JsonObject] = none
156+
variables: Option[JsonObject] = none,
157+
extensions: Option[JsonObject] = none
151158
): Resource[F, fs2.Stream[F, GraphQLResponse[D]]] =
152159
for
153-
res <- spanBuilderMod(
154-
Tracer[F]
155-
.spanBuilder(s"clue-client-subscribe-${document.querySummary}")
156-
.withSpanKind(SpanKind.Client)
157-
.addAttributes(Otel4sMiddleware.commonAttributes(document, operationName)*)
158-
).build.resource
159-
span = res.span
160-
_ <- Resource.eval:
161-
additionalAttributesF(document, variables).flatMap: attrs =>
162-
span.addAttributes(attrs*)
163-
stream <- wrapped.subscribeInternal[D](document, operationName, variables)
160+
res <- spanBuilderMod(
161+
Tracer[F]
162+
.spanBuilder(s"clue-client-subscribe-${document.querySummary}")
163+
.withSpanKind(SpanKind.Client)
164+
.addAttributes(Otel4sMiddleware.commonAttributes(document, operationName)*)
165+
).build.resource
166+
span = res.span
167+
_ <- Resource.eval:
168+
additionalAttributesF(document, variables).flatMap: attrs =>
169+
span.addAttributes(attrs*)
170+
traceHeaders <- Resource.eval(Tracer[F].propagate(Map.empty[String, String]))
171+
traceExt = JsonObject.fromMap(traceHeaders.map((k, v) => k -> io.circe.Json.fromString(v)))
172+
mergedExt = extensions
173+
.map(ext => JsonObject.fromIterable(ext.toIterable ++ traceExt.toIterable))
174+
.orElse(Some(traceExt))
175+
.filterNot(_.isEmpty)
176+
stream <- wrapped.subscribeInternal[D](document, operationName, variables, mergedExt)
164177
yield stream.onFinalizeCase: exitCase =>
165178
span.addAttribute(Attribute("clue.exitCase", exitCase.toOutcome.toString)) *>
166179
(exitCase match
167180
case Resource.ExitCase.Errored(e) =>
168181
span.setStatus(StatusCode.Error, Option(e.getMessage).getOrElse(e.toString))
169-
case _ => Applicative[F].unit
170-
)
182+
case _ => Applicative[F].unit)
171183

172-
class Otel4sPersistentStreamingClient[F[_]: Tracer: MonadCancelThrow, S, CP, CE](
184+
override protected[clue] def requestInternal[D: Decoder](
185+
document: GraphQLQuery,
186+
operationName: Option[String],
187+
variables: Option[JsonObject],
188+
modParams: Unit => Unit
189+
): F[GraphQLResponse[D]] =
190+
MonadCancelThrow[F].uncancelable: poll =>
191+
spanBuilderMod(
192+
Tracer[F]
193+
.spanBuilder(s"clue-client-request-${document.querySummary}")
194+
.withSpanKind(SpanKind.Client)
195+
.addAttributes(Otel4sMiddleware.commonAttributes(document, operationName)*)
196+
).build.use: span =>
197+
for
198+
additional <- additionalAttributesF(document, variables)
199+
_ <- span.addAttributes(additional*)
200+
traceHeaders <- Tracer[F].propagate(Map.empty[String, String])
201+
traceExt =
202+
JsonObject.fromMap(traceHeaders.map((k, v) => k -> io.circe.Json.fromString(v)))
203+
mergedExt = Some(traceExt).filterNot(_.isEmpty)
204+
result <- poll(
205+
wrapped
206+
.subscribeInternal[D](document, operationName, variables, mergedExt)
207+
.use(_.head.compile.onlyOrError)
208+
)
209+
_ <- Otel4sMiddleware.recordResponseAttributes(span, result)
210+
yield result
211+
212+
class Otel4sPersistentStreamingClient[F[_]: Tracer: Concurrent, S, CP, CE](
173213
wrapped: PersistentStreamingClient[F, S, CP, CE],
174214
spanBuilderMod: Otel4sMiddleware.SpanBuilderMod[F],
175215
additionalAttributesF: (GraphQLQuery, Option[JsonObject]) => F[List[Attribute[?]]]

‎scalajs/src/main/scala/clue/js/package.scala‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,16 @@
44
package clue
55

66
import clue.websocket.ApolloClient
7+
import org.scalajs.dom.Headers
78

89
package object js {
910
type FetchJsClient[F[_], S] = FetchClientWithPars[F, FetchJsRequest, S]
1011

1112
type WebSocketJsClient[F[_], S] = ApolloClient[F, String, S]
13+
14+
given TraceHeaderInjector[FetchJsRequest] with
15+
def addHeaders(params: FetchJsRequest, headers: Map[String, String]): FetchJsRequest =
16+
val newHeaders = new Headers(params.headers)
17+
headers.foreach((k, v) => newHeaders.set(k, v))
18+
params.copy(headers = newHeaders)
1219
}

0 commit comments

Comments
 (0)