Auto extract IDs from headers on the server

Motivation:

It's common for service authors to extract a tracing ID from a request
header and attach that value to the logger for an RPC. This only
requires a few lines of code but must be done in each RPC. gRPC should
provide configuration to do this automatically.

Modifications:

- Add server configuration which extracts a value from headers and
  attaches it to a logger at the start of each RPC.
- Wire up the config. Add tests.

Result:

Trace IDs can be automatically extracted from headers on the server and
attached to loggers.

Author: glbrntt

Sources/GRPC/AsyncAwaitSupport/GRPCAsyncServerHandler.swift

@@ -191,7 +191,10 @@ internal final class AsyncServerHandler<
 
   /// A logger.
   @usableFromInline
-  internal let logger: Logger
+  internal var logger: Logger
+
+  @usableFromInline
+  internal let traceIDExtractor: Server.Configuration.TraceIDExtractor?
 
   /// 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 +270,7 @@ internal final class AsyncServerHandler<
     self.compressionEnabledOnRPC = context.encoding.isEnabled
     self.compressResponsesIfPossible = true
     self.logger = context.logger
+    self.traceIDExtractor = context.traceIDExtractor
 
     self.userInfoRef = Ref(UserInfo())
     self.handlerStateMachine = .init()
@@ -295,6 +299,10 @@ internal final class AsyncServerHandler<
   internal func receiveMetadata(_ headers: HPACKHeaders) {
     switch self.interceptorStateMachine.interceptRequestMetadata() {
     case .intercept:
+      if let extractor = self.traceIDExtractor, let id = headers.first(name: extractor.headerName) {
+        self.logger[metadataKey: extractor.loggerKey] = "\(id)"
+        self.interceptors?.logger[metadataKey: extractor.loggerKey] = "\(id)"
+      }
       self.interceptors?.receive(.metadata(headers))
     case .cancel:
       self.cancel(error: nil)

Sources/GRPC/CallHandlers/BidirectionalStreamingServerHandler.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
 
@@ -56,6 +57,9 @@ public final class BidirectionalStreamingServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
   @usableFromInline
   internal enum State {
     // No headers have been received.
@@ -85,6 +89,7 @@ public final class BidirectionalStreamingServerHandler<
 
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
+    self.logger = context.logger
     self.interceptors = ServerInterceptorPipeline(
       logger: context.logger,
       eventLoop: context.eventLoop,
@@ -102,6 +107,11 @@ public final class BidirectionalStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
+    if let extractor = self.context.traceIDExtractor,
+       let id = headers.first(name: extractor.headerName) {
+      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
+      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    }
     self.interceptors.receive(.metadata(headers))
   }
 
@@ -164,7 +174,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,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
 
@@ -56,6 +57,9 @@ public final class ClientStreamingServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
   @usableFromInline
   internal enum State {
     // Nothing has happened yet.
@@ -86,6 +90,7 @@ public final class ClientStreamingServerHandler<
 
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
+    self.logger = context.logger
     self.interceptors = ServerInterceptorPipeline(
       logger: context.logger,
       eventLoop: context.eventLoop,
@@ -103,6 +108,11 @@ public final class ClientStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
+    if let extractor = self.context.traceIDExtractor,
+       let id = headers.first(name: extractor.headerName) {
+      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
+      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    }
     self.interceptors.receive(.metadata(headers))
   }
 
@@ -165,7 +175,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,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
 
@@ -52,6 +53,9 @@ public final class ServerStreamingServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
   @usableFromInline
   internal enum State {
     // Initial state. Nothing has happened yet.
@@ -82,6 +86,7 @@ public final class ServerStreamingServerHandler<
 
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
+    self.logger = context.logger
     self.interceptors = ServerInterceptorPipeline(
       logger: context.logger,
       eventLoop: context.eventLoop,
@@ -99,6 +104,11 @@ public final class ServerStreamingServerHandler<
 
   @inlinable
   public func receiveMetadata(_ headers: HPACKHeaders) {
+    if let extractor = self.context.traceIDExtractor,
+       let id = headers.first(name: extractor.headerName) {
+      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
+      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    }
     self.interceptors.receive(.metadata(headers))
   }
 
@@ -161,7 +171,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,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import Logging
 import NIOCore
 import NIOHPACK
 
@@ -51,6 +52,9 @@ public final class UnaryServerHandler<
   @usableFromInline
   internal var state: State = .idle
 
+  @usableFromInline
+  internal var logger: Logger
+
   @usableFromInline
   internal enum State {
     // Initial state. Nothing has happened yet.
@@ -77,6 +81,7 @@ public final class UnaryServerHandler<
     self.serializer = responseSerializer
     self.deserializer = requestDeserializer
     self.context = context
+    self.logger = context.logger
 
     let userInfoRef = Ref(UserInfo())
     self.userInfoRef = userInfoRef
@@ -97,6 +102,11 @@ public final class UnaryServerHandler<
 
   @inlinable
   public func receiveMetadata(_ metadata: HPACKHeaders) {
+    if let extractor = self.context.traceIDExtractor,
+       let id = metadata.first(name: extractor.headerName) {
+      self.logger[metadataKey: extractor.loggerKey] = "\(id)"
+      self.interceptors.logger[metadataKey: extractor.loggerKey] = "\(id)"
+    }
     self.interceptors.receive(.metadata(metadata))
   }
 
@@ -159,7 +169,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,
+      traceIDExtractor: self.configuration.traceIDExtractor
     )
   }
 

Sources/GRPC/GRPCServerRequestRoutingHandler.swift

@@ -59,6 +59,8 @@ public struct CallHandlerContext {
   internal var allocator: ByteBufferAllocator
   @usableFromInline
   internal var closeFuture: EventLoopFuture<Void>
+  @usableFromInline
+  internal var traceIDExtractor: Server.Configuration.TraceIDExtractor?
 }
 
 /// A call URI split into components.

Sources/GRPC/HTTP2ToRawGRPCServerCodec.swift

@@ -23,6 +23,7 @@ internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServe
   typealias OutboundOut = HTTP2Frame.FramePayload
 
   private var logger: Logger
+  private let traceIDExtractor: Server.Configuration.TraceIDExtractor?
   private var state: HTTP2ToRawGRPCStateMachine
   private let errorDelegate: ServerErrorDelegate?
   private var context: ChannelHandlerContext!
@@ -73,9 +74,11 @@ internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServe
     errorDelegate: ServerErrorDelegate?,
     normalizeHeaders: Bool,
     maximumReceiveMessageLength: Int,
-    logger: Logger
+    logger: Logger,
+    traceIDExtractor: Server.Configuration.TraceIDExtractor?
   ) {
     self.logger = logger
+    self.traceIDExtractor = traceIDExtractor
     self.errorDelegate = errorDelegate
     self.servicesByName = servicesByName
     self.encoding = encoding
@@ -127,7 +130,8 @@ internal final class HTTP2ToRawGRPCServerCodec: ChannelInboundHandler, GRPCServe
         closeFuture: context.channel.closeFuture,
         services: self.servicesByName,
         encoding: self.encoding,
-        normalizeHeaders: self.normalizeHeaders
+        normalizeHeaders: self.normalizeHeaders,
+        traceIDExtractor: self.traceIDExtractor
       )
 
       switch receiveHeaders {

Sources/GRPC/HTTP2ToRawGRPCStateMachine.swift

@@ -277,7 +277,8 @@ extension HTTP2ToRawGRPCStateMachine.State {
     closeFuture: EventLoopFuture<Void>,
     services: [Substring: CallHandlerProvider],
     encoding: ServerMessageEncoding,
-    normalizeHeaders: Bool
+    normalizeHeaders: Bool,
+    traceIDExtractor: Server.Configuration.TraceIDExtractor?
   ) -> HTTP2ToRawGRPCStateMachine.StateAndReceiveHeadersAction {
     // Extract and validate the content type. If it's nil we need to close.
     guard let contentType = self.extractContentType(from: headers) else {
@@ -326,7 +327,8 @@ extension HTTP2ToRawGRPCStateMachine.State {
       remoteAddress: remoteAddress,
       responseWriter: responseWriter,
       allocator: allocator,
-      closeFuture: closeFuture
+      closeFuture: closeFuture,
+      traceIDExtractor: traceIDExtractor
     )
 
     // We have a matching service, hopefully we have a provider for the method too.
@@ -865,7 +867,8 @@ extension HTTP2ToRawGRPCStateMachine {
     closeFuture: EventLoopFuture<Void>,
     services: [Substring: CallHandlerProvider],
     encoding: ServerMessageEncoding,
-    normalizeHeaders: Bool
+    normalizeHeaders: Bool,
+    traceIDExtractor: Server.Configuration.TraceIDExtractor?
   ) -> ReceiveHeadersAction {
     return self.withStateAvoidingCoWs { state in
       state.receive(
@@ -879,7 +882,8 @@ extension HTTP2ToRawGRPCStateMachine {
         closeFuture: closeFuture,
         services: services,
         encoding: encoding,
-        normalizeHeaders: normalizeHeaders
+        normalizeHeaders: normalizeHeaders,
+        traceIDExtractor: traceIDExtractor
       )
     }
   }
@@ -974,7 +978,8 @@ extension HTTP2ToRawGRPCStateMachine.State {
     closeFuture: EventLoopFuture<Void>,
     services: [Substring: CallHandlerProvider],
     encoding: ServerMessageEncoding,
-    normalizeHeaders: Bool
+    normalizeHeaders: Bool,
+    traceIDExtractor: Server.Configuration.TraceIDExtractor?
   ) -> HTTP2ToRawGRPCStateMachine.ReceiveHeadersAction {
     switch self {
     // These are the only states in which we can receive headers. Everything else is invalid.
@@ -991,7 +996,8 @@ extension HTTP2ToRawGRPCStateMachine.State {
         closeFuture: closeFuture,
         services: services,
         encoding: encoding,
-        normalizeHeaders: normalizeHeaders
+        normalizeHeaders: normalizeHeaders,
+        traceIDExtractor: traceIDExtractor
       )
       self = stateAndAction.state
       return stateAndAction.action

Sources/GRPC/Interceptor/ServerInterceptorPipeline.swift

@@ -36,7 +36,7 @@ internal final class ServerInterceptorPipeline<Request, Response> {
 
   /// A logger.
   @usableFromInline
-  internal let logger: Logger
+  internal var logger: Logger
 
   /// A reference to a 'UserInfo'.
   @usableFromInline

Sources/GRPC/Server.swift

@@ -17,6 +17,7 @@ import Foundation
 import Logging
 import NIOCore
 import NIOExtras
+import NIOHPACK
 import NIOHTTP1
 import NIOHTTP2
 import NIOPosix
@@ -272,11 +273,9 @@ extension Server {
         return Array(self.serviceProvidersByName.values)
       }
       set {
-        self
-          .serviceProvidersByName = Dictionary(
-            uniqueKeysWithValues: newValue
-              .map { ($0.serviceName, $0) }
-          )
+        self.serviceProvidersByName = Dictionary(
+          uniqueKeysWithValues: newValue.map { ($0.serviceName, $0) }
+        )
       }
     }
 
@@ -353,6 +352,10 @@ extension Server {
     /// available to service providers via `context`. Defaults to a no-op logger.
     public var logger = Logger(label: "io.grpc", factory: { _ in SwiftLogNoOpLogHandler() })
 
+    /// Configuration for extracting trace IDs from metadata and injecting the extracted ID into
+    /// a logger.
+    public var traceIDExtractor: TraceIDExtractor?
+
     /// A channel initializer which will be run after gRPC has initialized each accepted channel.
     /// This may be used to add additional handlers to the pipeline and is intended for debugging.
     /// This is analogous to `NIO.ServerBootstrap.childChannelInitializer`.
@@ -451,6 +454,22 @@ extension Server {
   }
 }
 
+extension Server.Configuration {
+  /// Configuration for extracting IDs from metadata and inserting them into a logger.
+  public struct TraceIDExtractor {
+    /// The name of the header field to extract an ID from. Header lookups are case insensitive.
+    public var headerName: String
+
+    /// Extracted IDs will be set on the logger with this key.
+    public var loggerKey: String
+
+    public init(headerName: String, loggerKey: String) {
+      self.headerName = headerName
+      self.loggerKey = loggerKey
+    }
+  }
+}
+
 extension ServerBootstrapProtocol {
   fileprivate func bind(to target: BindTarget) -> EventLoopFuture<Channel> {
     switch target.wrapped {