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 = extractor.extract(from: headers) {
+        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,10 @@ 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)"
+    }
     self.interceptors.receive(.metadata(headers))
   }
 
@@ -164,7 +173,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,10 @@ 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)"
+    }
     self.interceptors.receive(.metadata(headers))
   }
 
@@ -165,7 +174,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,10 @@ 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)"
+    }
     self.interceptors.receive(.metadata(headers))
   }
 
@@ -161,7 +170,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,10 @@ 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)"
+    }
     self.interceptors.receive(.metadata(metadata))
   }
 
@@ -159,7 +168,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