Auto extract IDs from headers on the server

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<
@@ -191,7 +192,15 @@ internal final class AsyncServerHandler<
 
   /// A logger.
   @usableFromInline
-  internal let logger: Logger
+  internal var logger: Logger
+
+  /// Contextual baggage which can be used to start tracing spans,
+  /// or carry additional contextual information through the handler.
+  @usableFromInline
+  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
@@ -267,6 +276,8 @@ internal final class AsyncServerHandler<
     self.compressionEnabledOnRPC = context.encoding.isEnabled
     self.compressResponsesIfPossible = true
     self.logger = context.logger
+    self.baggage = .topLevel // TODO: or carry from context?
+    self.tracer = context.tracer
 
     self.userInfoRef = Ref(UserInfo())
     self.handlerStateMachine = .init()
@@ -295,6 +306,18 @@ internal final class AsyncServerHandler<
   internal func receiveMetadata(_ headers: HPACKHeaders) {
     switch self.interceptorStateMachine.interceptRequestMetadata() {
     case .intercept:
+      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:
       self.cancel(error: nil)
@@ -315,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

@@ -13,8 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class BidirectionalStreamingServerHandler<
   Serializer: MessageSerializer,
@@ -56,6 +58,12 @@ public final class BidirectionalStreamingServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // No headers have been received.
@@ -85,6 +93,8 @@ 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,
@@ -102,6 +112,16 @@ public final class BidirectionalStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
+    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))
   }
 
@@ -164,7 +184,7 @@ public final class BidirectionalStreamingServerHandler<
       let context = _StreamingResponseCallContext<Request, Response>(
         eventLoop: self.context.eventLoop,
         headers: headers,
-        logger: self.context.logger,
+        logger: self.logger,
         userInfoRef: self.userInfoRef,
         compressionIsEnabled: self.context.encoding.isEnabled,
         closeFuture: self.context.closeFuture,

Sources/GRPC/CallHandlers/ClientStreamingServerHandler.swift

@@ -13,8 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class ClientStreamingServerHandler<
   Serializer: MessageSerializer,
@@ -56,6 +58,12 @@ public final class ClientStreamingServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // Nothing has happened yet.
@@ -86,6 +94,8 @@ 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,
@@ -103,6 +113,18 @@ public final class ClientStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
+    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))
   }
 
@@ -165,7 +187,7 @@ public final class ClientStreamingServerHandler<
       let context = UnaryResponseCallContext<Response>(
         eventLoop: self.context.eventLoop,
         headers: headers,
-        logger: self.context.logger,
+        logger: self.logger,
         userInfoRef: self.userInfoRef,
         closeFuture: self.context.closeFuture
       )

Sources/GRPC/CallHandlers/ServerStreamingServerHandler.swift

@@ -13,8 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class ServerStreamingServerHandler<
   Serializer: MessageSerializer,
@@ -52,6 +54,12 @@ public final class ServerStreamingServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // Initial state. Nothing has happened yet.
@@ -82,6 +90,8 @@ 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,
@@ -99,6 +109,16 @@ public final class ServerStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
+    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))
   }
 
@@ -161,7 +181,7 @@ public final class ServerStreamingServerHandler<
       let context = _StreamingResponseCallContext<Request, Response>(
         eventLoop: self.context.eventLoop,
         headers: headers,
-        logger: self.context.logger,
+        logger: self.logger,
         userInfoRef: self.userInfoRef,
         compressionIsEnabled: self.context.encoding.isEnabled,
         closeFuture: self.context.closeFuture,

Sources/GRPC/CallHandlers/UnaryServerHandler.swift

@@ -13,8 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
+import Tracing
 
 public final class UnaryServerHandler<
   Serializer: MessageSerializer,
@@ -51,6 +53,12 @@ public final class UnaryServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
+  @usableFromInline
+  internal var baggage: Baggage
+
   @usableFromInline
   internal enum State {
     // Initial state. Nothing has happened yet.
@@ -77,6 +85,8 @@ 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())
     self.userInfoRef = userInfoRef
@@ -97,6 +107,16 @@ public final class UnaryServerHandler<
 
   @inlinable
   public func receiveMetadata(_ metadata: HPACKHeaders) {
+    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))
   }
 
@@ -159,7 +179,7 @@ public final class UnaryServerHandler<
       let context = UnaryResponseCallContext<Response>(
         eventLoop: self.context.eventLoop,
         headers: headers,
-        logger: self.context.logger,
+        logger: self.logger,
         userInfoRef: self.userInfoRef,
         closeFuture: self.context.closeFuture
       )

Sources/GRPC/GRPCServerPipelineConfigurator.swift

@@ -142,7 +142,8 @@ final class GRPCServerPipelineConfigurator: ChannelInboundHandler, RemovableChan
       errorDelegate: self.configuration.errorDelegate,
       normalizeHeaders: normalizeHeaders,
       maximumReceiveMessageLength: self.configuration.maximumReceiveMessageLength,
-      logger: logger
+      logger: logger,
+      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
@@ -59,6 +61,8 @@ public struct CallHandlerContext {
   internal var allocator: ByteBufferAllocator
   @usableFromInline
   internal var closeFuture: EventLoopFuture<Void>
+  @usableFromInline
+  internal var tracer: Tracing.Tracer?
 }
 
 /// A call URI split into components.