+tracing implement TraceID handling using swift-distributed-tracing

Author: ktoso

Package.swift

@@ -51,19 +51,25 @@ let packageDependencies: [Package.Dependency] = [
     from: "1.20.2"
   ),
   .package(
-    url: "https://github.com/apple/swift-log.git",
-    from: "1.4.4"
+    // url: "https://github.com/apple/swift-log.git",
+    // from: "1.4.4"
+    url: "https://github.com/slashmo/swift-log.git",
+    branch: "feature/baggage"
   ),
   .package(
     url: "https://github.com/apple/swift-argument-parser.git",
     // Version is higher than in other Package@swift manifests: 1.1.0 raised the minimum Swift
-    // version and indluded async support.
+    // version and included async support.
     from: "1.1.1"
   ),
   .package(
     url: "https://github.com/apple/swift-docc-plugin",
     from: "1.0.0"
   ),
+  .package(
+    url: "https://github.com/apple/swift-distributed-tracing",
+    branch: "main"
+  ),
 ].appending(
   .package(
     url: "https://github.com/apple/swift-nio-ssl.git",
@@ -114,6 +120,7 @@ extension Target.Dependency {
     package: "swift-nio-transport-services"
   )
   static let logging: Self = .product(name: "Logging", package: "swift-log")
+  static let tracing: Self = .product(name: "Tracing", package: "swift-distributed-tracing")
   static let protobuf: Self = .product(name: "SwiftProtobuf", package: "swift-protobuf")
   static let protobufPluginLibrary: Self = .product(
     name: "SwiftProtobufPluginLibrary",
@@ -139,6 +146,7 @@ extension Target {
       .nioHTTP2,
       .nioExtras,
       .logging,
+      .tracing,
       .protobuf,
     ].appending(
       .nioSSL, if: includeNIOSSL

Sources/GRPC/AsyncAwaitSupport/GRPCAsyncServerHandler.swift

@@ -18,6 +18,7 @@ import DequeModule
 import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
 public struct GRPCAsyncServerHandler<
@@ -193,8 +194,13 @@ internal final class AsyncServerHandler<
   @usableFromInline
   internal var logger: Logger
 
+  /// Contextual baggage which can be used to start tracing spans,
+  // or carry additional contextual information through the handler.
   @usableFromInline
-  internal let traceIDExtractor: Server.Configuration.TraceIDExtractor?
+  internal var baggage: Baggage
+
+  @usableFromInline
+  internal let tracer: Tracing.Tracer? // FIXME: tri state "don't trace" / "global" / "configured"
 
   /// A reference to the user info. This is shared with the interceptor pipeline and may be accessed
   /// from the async call context. `UserInfo` is _not_ `Sendable` and must always be accessed from
@@ -270,7 +276,8 @@ internal final class AsyncServerHandler<
     self.compressionEnabledOnRPC = context.encoding.isEnabled
     self.compressResponsesIfPossible = true
     self.logger = context.logger
-    self.traceIDExtractor = context.traceIDExtractor
+    self.baggage = .topLevel // TODO: or carry from context?
+    self.tracer = context.tracer
 
     self.userInfoRef = Ref(UserInfo())
     self.handlerStateMachine = .init()
@@ -299,9 +306,17 @@ internal final class AsyncServerHandler<
   internal func receiveMetadata(_ headers: HPACKHeaders) {
     switch self.interceptorStateMachine.interceptRequestMetadata() {
     case .intercept:
-      if let extractor = self.traceIDExtractor, let id = extractor.extract(from: headers) {
-        self.logger[metadataKey: extractor.loggerKey] = "\(id)"
-        self.interceptors?.logger[metadataKey: extractor.loggerKey] = "\(id)"
+      if let tracer = self.tracer {
+        tracer.extract(headers, into: &baggage, using: HPACKHeadersExtractor())
+
+        if let metadata: Logger.Metadata = self.logger.metadataProvider?.metadata(baggage) { // FIXME: function naming a bit ugly here
+          for (k, v) in metadata {
+            self.logger[metadataKey: k] = v
+            self.interceptors?.logger[metadataKey: k] = v
+            // self.logger[metadataKey: extractor.loggerKey] = "\(id)"
+            // self.interceptors?.logger[metadataKey: extractor.loggerKey] = "\(id)"
+          }
+        }
       }
       self.interceptors?.receive(.metadata(headers))
     case .cancel:
@@ -323,7 +338,7 @@ internal final class AsyncServerHandler<
 
     switch self.interceptorStateMachine.interceptRequestMessage() {
     case .intercept:
-      self.interceptors?.receive(.message(request))
+      self.interceptors?.receive(.message(request)) // TODO: if we wrapped this in a withSpan, would that be correct?
     case .cancel:
       self.cancel(error: nil)
     case .drop:

Sources/GRPC/CallHandlers/BidirectionalStreamingServerHandler.swift

@@ -16,6 +16,7 @@
 import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class BidirectionalStreamingServerHandler<
   Serializer: MessageSerializer,
@@ -60,6 +61,9 @@ public final class BidirectionalStreamingServerHandler<
   @usableFromInline
   internal var logger: Logger
 
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // No headers have been received.
@@ -90,6 +94,7 @@ public final class BidirectionalStreamingServerHandler<
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
     self.logger = context.logger
+    self.baggage = .topLevel
     self.interceptors = ServerInterceptorPipeline(
       logger: context.logger,
       eventLoop: context.eventLoop,
@@ -107,9 +112,15 @@ public final class BidirectionalStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
-    if let extractor = self.context.traceIDExtractor, let id = extractor.extract(from: headers) {
-      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
-      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    if let tracer = self.context.tracer {
+      tracer.extract(headers, into: &self.baggage, using: HPACKHeadersExtractor())
+
+      if let metadata: Logger.Metadata = self.logger.metadataProvider?.metadata(baggage) { // FIXME: function naming a bit ugly here
+        for (k, v) in metadata {
+          self.logger[metadataKey: k] = v
+          self.interceptors?.logger[metadataKey: k] = v
+        }
+      }
     }
     self.interceptors.receive(.metadata(headers))
   }

Sources/GRPC/CallHandlers/ClientStreamingServerHandler.swift

@@ -16,6 +16,7 @@
 import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class ClientStreamingServerHandler<
   Serializer: MessageSerializer,
@@ -60,6 +61,9 @@ public final class ClientStreamingServerHandler<
   @usableFromInline
   internal var logger: Logger
 
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // Nothing has happened yet.
@@ -91,6 +95,7 @@ public final class ClientStreamingServerHandler<
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
     self.logger = context.logger
+    self.baggage = .topLevel
     self.interceptors = ServerInterceptorPipeline(
       logger: context.logger,
       eventLoop: context.eventLoop,
@@ -108,9 +113,17 @@ public final class ClientStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
-    if let extractor = self.context.traceIDExtractor, let id = extractor.extract(from: headers) {
-      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
-      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    if let tracer = self.context.tracer {
+      var baggage = Baggage.topLevel
+      tracer.extract(headers, into: &baggage, using: HPACKHeadersExtractor())
+      self.baggage = baggage
+
+      if let metadata: Logger.Metadata = self.logger.metadataProvider?.metadata(baggage) { // FIXME: function naming a bit ugly here
+        for (k, v) in metadata {
+          self.logger[metadataKey: k] = v
+          self.interceptors?.logger[metadataKey: k] = v
+        }
+      }
     }
     self.interceptors.receive(.metadata(headers))
   }

Sources/GRPC/CallHandlers/ServerStreamingServerHandler.swift

@@ -16,6 +16,7 @@
 import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class ServerStreamingServerHandler<
   Serializer: MessageSerializer,
@@ -56,6 +57,9 @@ public final class ServerStreamingServerHandler<
   @usableFromInline
   internal var logger: Logger
 
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // Initial state. Nothing has happened yet.
@@ -87,6 +91,7 @@ public final class ServerStreamingServerHandler<
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
     self.logger = context.logger
+    self.baggage = .topLevel
     self.interceptors = ServerInterceptorPipeline(
       logger: context.logger,
       eventLoop: context.eventLoop,
@@ -104,9 +109,15 @@ public final class ServerStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
-    if let extractor = self.context.traceIDExtractor, let id = extractor.extract(from: headers) {
-      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
-      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    if let tracer = self.context.tracer {
+      tracer.extract(headers, into: &self.baggage, using: HPACKHeadersExtractor())
+
+      if let metadata: Logger.Metadata = self.logger.metadataProvider?.metadata(baggage) { // FIXME: function naming a bit ugly here
+        for (k, v) in metadata {
+          self.logger[metadataKey: k] = v
+          self.interceptors?.logger[metadataKey: k] = v
+        }
+      }
     }
     self.interceptors.receive(.metadata(headers))
   }

Sources/GRPC/CallHandlers/UnaryServerHandler.swift

@@ -16,6 +16,7 @@
 import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class UnaryServerHandler<
   Serializer: MessageSerializer,
@@ -55,6 +56,9 @@ public final class UnaryServerHandler<
   @usableFromInline
   internal var logger: Logger
 
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // Initial state. Nothing has happened yet.
@@ -81,6 +85,7 @@ public final class UnaryServerHandler<
     self.serializer = responseSerializer
     self.deserializer = requestDeserializer
     self.context = context
+    self.baggage = .topLevel // TODO: decide if rather we set it from somewhere here?
     self.logger = context.logger
 
     let userInfoRef = Ref(UserInfo())
@@ -102,9 +107,15 @@ public final class UnaryServerHandler<
 
   @inlinable
   public func receiveMetadata(_ metadata: HPACKHeaders) {
-    if let extractor = self.context.traceIDExtractor, let id = extractor.extract(from: metadata) {
-      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
-      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    if let tracer = self.context.tracer {
+      tracer.extract(metadata, into: &baggage, using: HPACKHeadersExtractor())
+
+      if let metadata: Logger.Metadata = self.logger.metadataProvider?.metadata(baggage) { // FIXME: function naming a bit ugly here
+        for (k, v) in metadata {
+          self.logger[metadataKey: k] = v
+          self.interceptors?.logger[metadataKey: k] = v
+        }
+      }
     }
     self.interceptors.receive(.metadata(metadata))
   }

Sources/GRPC/GRPCServerPipelineConfigurator.swift

@@ -143,7 +143,7 @@ final class GRPCServerPipelineConfigurator: ChannelInboundHandler, RemovableChan
       normalizeHeaders: normalizeHeaders,
       maximumReceiveMessageLength: self.configuration.maximumReceiveMessageLength,
       logger: logger,
-      traceIDExtractor: self.configuration.traceIDExtractor
+      tracer: self.configuration.tracer
     )
   }
 

Sources/GRPC/GRPCServerRequestRoutingHandler.swift

@@ -19,6 +19,7 @@ import NIOHPACK
 import NIOHTTP1
 import NIOHTTP2
 import SwiftProtobuf
+import Tracing
 
 /// Provides ``GRPCServerHandlerProtocol`` objects for the methods on a particular service name.
 ///
@@ -45,6 +46,7 @@ public struct CallHandlerContext {
   internal var errorDelegate: ServerErrorDelegate?
   @usableFromInline
   internal var logger: Logger
+  // TODO: does this need baggage?
   @usableFromInline
   internal var encoding: ServerMessageEncoding
   @usableFromInline
@@ -60,7 +62,7 @@ public struct CallHandlerContext {
   @usableFromInline
   internal var closeFuture: EventLoopFuture<Void>
   @usableFromInline
-  internal var traceIDExtractor: Server.Configuration.TraceIDExtractor?
+  internal var tracer: Tracing.Tracer?
 }
 
 /// A call URI split into components.

Sources/GRPC/HTTP2ToRawGRPCServerCodec.swift

@@ -17,13 +17,14 @@ import Logging
 import NIOCore
 import NIOHPACK
 import NIOHTTP2
+import Tracing
 
 internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServerResponseWriter {
   typealias InboundIn = HTTP2Frame.FramePayload
   typealias OutboundOut = HTTP2Frame.FramePayload
 
   private var logger: Logger
-  private let traceIDExtractor: Server.Configuration.TraceIDExtractor?
+  private var tracer: Tracing.Tracer?
   private var state: HTTP2ToRawGRPCStateMachine
   private let errorDelegate: ServerErrorDelegate?
   private var context: ChannelHandlerContext!
@@ -75,10 +76,10 @@ internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServe
     normalizeHeaders: Bool,
     maximumReceiveMessageLength: Int,
     logger: Logger,
-    traceIDExtractor: Server.Configuration.TraceIDExtractor?
+    tracer: Tracing.Tracer?
   ) {
     self.logger = logger
-    self.traceIDExtractor = traceIDExtractor
+    self.tracer = tracer
     self.errorDelegate = errorDelegate
     self.servicesByName = servicesByName
     self.encoding = encoding
@@ -131,7 +132,7 @@ internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServe
         services: self.servicesByName,
         encoding: self.encoding,
         normalizeHeaders: self.normalizeHeaders,
-        traceIDExtractor: self.traceIDExtractor
+        tracer: self.tracer
       )
 
       switch receiveHeaders {