Add option to include response metadata in client interceptor

Author: gjcairo

Sources/GRPCOTelTracingInterceptors/Tracing/ClientOTelTracingInterceptor.swift

@@ -36,6 +36,7 @@ public struct ClientOTelTracingInterceptor: ClientInterceptor {
 
   private let traceEachMessage: Bool
   private var includeRequestMetadata: Bool
+  private var includeResponseMetadata: Bool
 
   /// Create a new instance of a ``ClientOTelTracingInterceptor``.
   ///
@@ -45,20 +46,24 @@ public struct ClientOTelTracingInterceptor: ClientInterceptor {
   ///  `network.transport` attribute in spans.
   ///  - traceEachMessage: If `true`, each request part sent and response part received will be recorded as a separate
   ///  event in a tracing span.
-  ///  - includeRequestMetadata: if `true`, **all** metadata keys included in the request will be added to the span as attributes.
+  ///  - includeRequestMetadata: if `true`, **all** metadata keys with string values included in the request will be added to the span as attributes.
+  ///  - includeResponseMetadata: if `true`, **all** metadata keys with string values included in the response will be added to the span as attributes.
   ///
-  /// - Important: Be careful when setting `includeRequestMetadata=true`, as including all request metadata can be a security risk.
+  /// - Important: Be careful when setting `includeRequestMetadata` or `includeResponseMetadata` to `true`,
+  /// as including all request/response metadata can be a security risk.
   public init(
     serverHostname: String,
     networkTransportMethod: String,
     traceEachMessage: Bool = true,
-    includeRequestMetadata: Bool = false
+    includeRequestMetadata: Bool = false,
+    includeResponseMetadata: Bool = false
   ) {
     self.injector = ClientRequestInjector()
     self.serverHostname = serverHostname
     self.networkTransportMethod = networkTransportMethod
     self.traceEachMessage = traceEachMessage
     self.includeRequestMetadata = includeRequestMetadata
+    self.includeResponseMetadata = includeResponseMetadata
   }
 
   /// This interceptor will inject as the request's metadata whatever `ServiceContext` key-value pairs
@@ -142,6 +147,11 @@ public struct ClientOTelTracingInterceptor: ClientInterceptor {
       }
 
       var response = try await next(request, context)
+
+      if self.includeResponseMetadata {
+        span.setMetadataStringAttributesAsResponseSpanAttributes(response.metadata)
+      }
+
       switch response.accepted {
       case .success(var success):
         let hookedSequence:
@@ -150,14 +160,22 @@ public struct ClientOTelTracingInterceptor: ClientInterceptor {
           >
         if self.traceEachMessage {
           let messageReceivedCounter = Atomic(1)
-          hookedSequence = HookedRPCAsyncSequence(wrapping: success.bodyParts) { _ in
-            var event = SpanEvent(name: "rpc.message")
-            event.attributes[GRPCTracingKeys.rpcMessageType] = "RECEIVED"
-            event.attributes[GRPCTracingKeys.rpcMessageID] =
+          hookedSequence = HookedRPCAsyncSequence(wrapping: success.bodyParts) { part in
+            switch part {
+            case .message(let message):
+              var event = SpanEvent(name: "rpc.message")
+              event.attributes[GRPCTracingKeys.rpcMessageType] = "RECEIVED"
+              event.attributes[GRPCTracingKeys.rpcMessageID] =
               messageReceivedCounter
-              .wrappingAdd(1, ordering: .sequentiallyConsistent)
-              .oldValue
-            span.addEvent(event)
+                .wrappingAdd(1, ordering: .sequentiallyConsistent)
+                .oldValue
+              span.addEvent(event)
+
+            case .trailingMetadata(let trailingMetadata):
+              if self.includeResponseMetadata {
+                span.setMetadataStringAttributesAsResponseSpanAttributes(trailingMetadata)
+              }
+            }
           } onFinish: { error in
             if let error {
               if let errorCode = error.grpcErrorCode {

Tests/GRPCOTelTracingInterceptorsTests/GRPCOTelTracingInterceptorsTests.swift

@@ -180,7 +180,7 @@ struct OTelTracingClientInterceptorTests {
     }
   }
 
-  @Test("All request metadata is included if opted-in")
+  @Test("All string-valued request metadata is included if opted-in")
   func testRequestMetadataOptIn() async throws {
     var serviceContext = ServiceContext.topLevel
     let traceIDString = UUID().uuidString
@@ -199,7 +199,7 @@ struct OTelTracingClientInterceptorTests {
       )
       let methodDescriptor = MethodDescriptor(
         fullyQualifiedService: "OTelTracingClientInterceptorTests",
-        method: "testAllEventsRecorded"
+        method: "testRequestMetadataOptIn"
       )
       let response = try await interceptor.intercept(
         tracer: self.tracer,
@@ -288,6 +288,118 @@ struct OTelTracingClientInterceptorTests {
     }
   }
 
+  @Test("All string-valued response metadata is included if opted-in")
+  func testResponseMetadataOptIn() async throws {
+    var serviceContext = ServiceContext.topLevel
+    let traceIDString = UUID().uuidString
+
+    let (requestStream, requestStreamContinuation) = AsyncStream<String>.makeStream()
+    serviceContext.traceID = traceIDString
+
+    // FIXME: use 'ServiceContext.withValue(serviceContext)'
+    //
+    // This is blocked on: https://github.com/apple/swift-service-context/pull/46
+    try await ServiceContext.$current.withValue(serviceContext) {
+      let interceptor = ClientOTelTracingInterceptor(
+        serverHostname: "someserver.com",
+        networkTransportMethod: "tcp",
+        includeResponseMetadata: true
+      )
+      let methodDescriptor = MethodDescriptor(
+        fullyQualifiedService: "OTelTracingClientInterceptorTests",
+        method: "testResponseMetadataOptIn"
+      )
+      let response = try await interceptor.intercept(
+        tracer: self.tracer,
+        request: .init(
+          metadata: [
+            "some-request-metadata": "some-request-value",
+            "some-repeated-request-metadata": "some-repeated-request-value1",
+            "some-repeated-request-metadata": "some-repeated-request-value2",
+            "some-request-metadata-bin": .binary([1]),
+          ],
+          producer: { writer in
+            try await writer.write(contentsOf: ["request1"])
+            try await writer.write(contentsOf: ["request2"])
+          }
+        ),
+        context: ClientContext(
+          descriptor: methodDescriptor,
+          remotePeer: "ipv4:10.1.2.80:567",
+          localPeer: "ipv4:10.1.2.80:123"
+        )
+      ) { stream, _ in
+        // Assert the metadata contains the injected context key-value.
+        #expect(
+          stream.metadata.contains(where: {
+            ($0.key == "trace-id") && ($0.value == .string(traceIDString))
+          })
+        )
+
+        // Write into the request stream to make sure the `producer` closure's called.
+        let writer = RPCWriter(wrapping: TestWriter(streamContinuation: requestStreamContinuation))
+        try await stream.producer(writer)
+        requestStreamContinuation.finish()
+
+        return .init(
+          metadata: [
+            "some-response-metadata": "some-response-value",
+            "some-response-metadata-bin": .binary([2]),
+          ],
+          bodyParts: RPCAsyncSequence(
+            wrapping: AsyncThrowingStream<StreamingClientResponse.Contents.BodyPart, any Error> {
+              $0.yield(.message(["response"]))
+              $0.yield(.trailingMetadata([
+                "some-repeated-response-metadata": "some-repeated-response-value1",
+                "some-repeated-response-metadata": "some-repeated-response-value2"
+              ]))
+              $0.finish()
+            }
+          )
+        )
+      }
+
+      await assertStreamContentsEqual(["request1", "request2"], requestStream)
+      try await assertStreamContentsEqual([["response"]], response.messages)
+
+      assertTestSpanComponents(forMethod: methodDescriptor, tracer: self.tracer) { events in
+        #expect(
+          events == [
+            // Recorded when `request1` is sent
+            TestSpanEvent("rpc.message", ["rpc.message.type": "SENT", "rpc.message.id": 1]),
+            // Recorded when `request2` is sent
+            TestSpanEvent("rpc.message", ["rpc.message.type": "SENT", "rpc.message.id": 2]),
+            // Recorded when receiving response part
+            TestSpanEvent("rpc.message", ["rpc.message.type": "RECEIVED", "rpc.message.id": 1]),
+          ]
+        )
+      } assertAttributes: { attributes in
+        #expect(
+          attributes == [
+            "rpc.system": "grpc",
+            "rpc.method": .string(methodDescriptor.method),
+            "rpc.service": .string(methodDescriptor.service.fullyQualifiedService),
+            "rpc.grpc.status_code": 0,
+            "server.address": "someserver.com",
+            "server.port": 567,
+            "network.peer.address": "10.1.2.80",
+            "network.peer.port": 567,
+            "network.transport": "tcp",
+            "network.type": "ipv4",
+            "rpc.grpc.response.metadata.some-response-metadata": "some-response-value",
+            "rpc.grpc.response.metadata.some-repeated-response-metadata": .stringArray([
+              "some-repeated-response-value1", "some-repeated-response-value2",
+            ]),
+          ]
+        )
+      } assertStatus: { status in
+        #expect(status == nil)
+      } assertErrors: { errors in
+        #expect(errors == [])
+      }
+    }
+  }
+
   @Test("RPC that throws is correctly recorded")
   func testThrowingRPC() async throws {
     var serviceContext = ServiceContext.topLevel