Add client and server tracing interceptors

Package.swift

@@ -72,6 +72,10 @@ let packageDependencies: [Package.Dependency] = [
     url: "https://github.com/apple/swift-docc-plugin",
     from: "1.0.0"
   ),
+  .package(
+    url: "https://github.com/apple/swift-distributed-tracing.git",
+    from: "1.0.0"
+  ),
 ].appending(
   .package(
     url: "https://github.com/apple/swift-nio-ssl.git",
@@ -131,9 +135,11 @@ extension Target.Dependency {
   )
   static let dequeModule: Self = .product(name: "DequeModule", package: "swift-collections")
   static let atomics: Self = .product(name: "Atomics", package: "swift-atomics")
+  static let tracing: Self = .product(name: "Tracing", package: "swift-distributed-tracing")
 
   static let grpcCore: Self = .target(name: "GRPCCore")
   static let grpcInProcessTransport: Self = .target(name: "GRPCInProcessTransport")
+  static let grpcInterceptors: Self = .target(name: "GRPCInterceptors")
   static let grpcHTTP2Core: Self = .target(name: "GRPCHTTP2Core")
   static let grpcHTTP2TransportNIOPosix: Self = .target(name: "GRPCHTTP2TransportNIOPosix")
   static let grpcHTTP2TransportNIOTransportServices: Self = .target(name: "GRPCHTTP2TransportNIOTransportServices")
@@ -181,6 +187,14 @@ extension Target {
     ]
   )
 
+  static let grpcInterceptors: Target = .target(
+    name: "GRPCInterceptors",
+    dependencies: [
+      .grpcCore,
+      .tracing
+    ]
+  )
+
   static let grpcHTTP2Core: Target = .target(
     name: "GRPCHTTP2Core",
     dependencies: [
@@ -274,10 +288,20 @@ extension Target {
     name: "GRPCInProcessTransportTests",
     dependencies: [
       .grpcCore,
-      .grpcInProcessTransport,
+      .grpcInProcessTransport
     ]
   )
 
+  static let grpcInterceptorsTests: Target = .testTarget(
+    name: "GRPCInterceptorsTests",
+    dependencies: [
+      .grpcCore,
+      .tracing,
+      .nioCore,
+      .grpcInterceptors
+    ]
+  )
+
   static let grpcHTTP2CoreTests: Target = .testTarget(
     name: "GRPCHTTP2CoreTests",
     dependencies: [
@@ -638,6 +662,7 @@ let package = Package(
     .grpcCore,
     .grpcInProcessTransport,
     .grpcCodeGen,
+    .grpcInterceptors,
     .grpcHTTP2Core,
     .grpcHTTP2TransportNIOPosix,
     .grpcHTTP2TransportNIOTransportServices,
@@ -646,6 +671,7 @@ let package = Package(
     .grpcCoreTests,
     .grpcInProcessTransportTests,
     .grpcCodeGenTests,
+    .grpcInterceptorsTests,
     .grpcHTTP2CoreTests,
     .grpcHTTP2TransportNIOPosixTests,
     .grpcHTTP2TransportNIOTransportServicesTests

Sources/GRPCInterceptors/ClientTracingInterceptor.swift

@@ -0,0 +1,140 @@
+/*
+ * Copyright 2024, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import GRPCCore
+import Tracing
+
+/// A client interceptor that injects tracing information into the request.
+///
+/// The tracing information is taken from the current `ServiceContext`, and injected into the request's
+/// metadata. It will then be picked up by the server-side ``ServerTracingInterceptor``.
+///
+/// For more information, refer to the documentation for `swift-distributed-tracing`.
+@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+public struct ClientTracingInterceptor: ClientInterceptor {
+  private let injector: ClientRequestInjector
+  private let emitEventOnEachWrite: Bool
+
+  /// Create a new instance of a ``ClientTracingInterceptor``.
+  ///
+  /// - Parameter emitEventOnEachWrite: If `true`, each request part sent and response part
+  /// received will be recorded as a separate event in a tracing span. Otherwise, only the request/response
+  /// start and end will be recorded as events.
+  public init(emitEventOnEachWrite: Bool = false) {
+    self.injector = ClientRequestInjector()
+    self.emitEventOnEachWrite = emitEventOnEachWrite
+  }
+
+  /// This interceptor will inject as the request's metadata whatever `ServiceContext` key-value pairs
+  /// have been made available by the tracing implementation bootstrapped in your application.
+  ///
+  /// Which key-value pairs are injected will depend on the specific tracing implementation
+  /// that has been configured when bootstrapping `swift-distributed-tracing` in your application.
+  public func intercept<Input, Output>(
+    request: ClientRequest.Stream<Input>,
+    context: ClientInterceptorContext,
+    next: @Sendable (ClientRequest.Stream<Input>, ClientInterceptorContext) async throws ->
+      ClientResponse.Stream<Output>
+  ) async throws -> ClientResponse.Stream<Output> where Input: Sendable, Output: Sendable {
+    var request = request
+    let tracer = InstrumentationSystem.tracer
+    let serviceContext = ServiceContext.current ?? .topLevel
+
+    tracer.inject(
+      serviceContext,
+      into: &request.metadata,
+      using: self.injector
+    )
+
+    return try await tracer.withSpan(
+      context.descriptor.fullyQualifiedMethod,
+      context: serviceContext,
+      ofKind: .client
+    ) { span in
+      span.addEvent("Request started")
+
+      if self.emitEventOnEachWrite {
+        let wrappedProducer = request.producer
+        request.producer = { writer in
+          let eventEmittingWriter = HookedWriter(
+            wrapping: writer,
+            beforeEachWrite: {
+              span.addEvent("Sending request part")
+            },
+            afterEachWrite: {
+              span.addEvent("Sent request part")
+            }
+          )
+
+          do {
+            try await wrappedProducer(RPCWriter(wrapping: eventEmittingWriter))
+          } catch {
+            span.addEvent("Error encountered")
+            throw error
+          }
+
+          span.addEvent("Request end")
+        }
+      }
+
+      var response: ClientResponse.Stream<Output>
+      do {
+        response = try await next(request, context)
+      } catch {
+        span.addEvent("Error encountered")
+        throw error
+      }
+
+      switch response.accepted {
+      case .success(var success):
+        if self.emitEventOnEachWrite {
+          let onEachPartRecordingSequence = success.bodyParts.map { element in
+            span.addEvent("Received response part")
+            return element
+          }
+          let onFinishRecordingSequence = OnFinishAsyncSequence(
+            wrapping: onEachPartRecordingSequence
+          ) {
+            span.addEvent("Received response end")
+          }
+          success.bodyParts = RPCAsyncSequence(wrapping: onFinishRecordingSequence)
+          response.accepted = .success(success)
+        } else {
+          let onFinishRecordingSequence = OnFinishAsyncSequence(wrapping: success.bodyParts) {
+            span.addEvent("Received response end")
+          }
+          success.bodyParts = RPCAsyncSequence(wrapping: onFinishRecordingSequence)
+          response.accepted = .success(success)
+        }
+      case .failure:
+        span.addEvent("Received error response")
+      }
+
+      return response
+    }
+  }
+}
+
+/// An injector responsible for injecting the required instrumentation keys from the `ServiceContext` into
+/// the request metadata.
+@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+struct ClientRequestInjector: Instrumentation.Injector {
+  typealias Carrier = Metadata
+
+  func inject(_ value: String, forKey key: String, into carrier: inout Carrier) {
+    carrier.addString(value, forKey: key)
+  }
+}

Sources/GRPCInterceptors/HookedWriter.swift

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2024, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import GRPCCore
+import Tracing
+
+@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+struct HookedWriter<Element>: RPCWriterProtocol {
+  private let writer: any RPCWriterProtocol<Element>
+  private let beforeEachWrite: @Sendable () -> Void
+  private let afterEachWrite: @Sendable () -> Void
+
+  init(
+    wrapping other: some RPCWriterProtocol<Element>,
+    beforeEachWrite: @Sendable @escaping () -> Void,
+    afterEachWrite: @Sendable @escaping () -> Void
+  ) {
+    self.writer = other
+    self.beforeEachWrite = beforeEachWrite
+    self.afterEachWrite = afterEachWrite
+  }
+
+  func write(contentsOf elements: some Sequence<Element>) async throws {
+    self.beforeEachWrite()
+    try await self.writer.write(contentsOf: elements)
+    self.afterEachWrite()
+  }
+}

Sources/GRPCInterceptors/OnFinishAsyncSequence.swift

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2024, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+struct OnFinishAsyncSequence<Element: Sendable>: AsyncSequence, Sendable {
+  private let _makeAsyncIterator: @Sendable () -> AsyncIterator
+
+  init<S: AsyncSequence>(
+    wrapping other: S,
+    onFinish: @escaping () -> Void
+  ) where S.Element == Element {
+    self._makeAsyncIterator = {
+      AsyncIterator(wrapping: other.makeAsyncIterator(), onFinish: onFinish)
+    }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator {
+    self._makeAsyncIterator()
+  }
+
+  struct AsyncIterator: AsyncIteratorProtocol {
+    private var iterator: any AsyncIteratorProtocol
+    private var onFinish: (() -> Void)?
+
+    fileprivate init<Iterator>(
+      wrapping other: Iterator,
+      onFinish: @escaping () -> Void
+    ) where Iterator: AsyncIteratorProtocol, Iterator.Element == Element {
+      self.iterator = other
+      self.onFinish = onFinish
+    }
+
+    mutating func next() async throws -> Element? {
+      let elem = try await self.iterator.next()
+
+      if elem == nil {
+        self.onFinish?()
+        self.onFinish = nil
+      }
+
+      return elem as? Element
+    }
+  }
+}