Allow adding ClientInterceptors to specific services and methods

Sources/GRPCCore/Call/Client/ClientInterceptor.swift

@@ -21,10 +21,11 @@
 /// received from the transport. They are typically used for cross-cutting concerns like injecting
 /// metadata, validating messages, logging additional data, and tracing.
 ///
-/// Interceptors are registered with a client and apply to all RPCs. If you need to modify the
-/// behavior of an interceptor on a per-RPC basis then you can use the
-/// ``ClientContext/descriptor`` to determine which RPC is being called and
-/// conditionalise behavior accordingly.
+/// Interceptors are registered with the server via ``ClientInterceptorPipelineOperation``s.
+/// You may register them for all services registered with a server, for RPCs directed to specific services, or
+/// for RPCs directed to specific methods. If you need to modify the behavior of an interceptor on a
+/// per-RPC basis in more detail, then you can use the ``ClientContext/descriptor`` to determine
+/// which RPC is being called and conditionalise behavior accordingly.
 ///
 /// - TODO: Update example and documentation to show how to register an interceptor.
 ///

Sources/GRPCCore/Call/Client/ClientInterceptorPipelineOperation.swift

@@ -0,0 +1,99 @@
+/*
+ * 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.
+ */
+
+/// A `ClientInterceptorPipelineOperation` describes to which RPCs a client interceptor should be applied.
+///
+/// You can configure a client interceptor to be applied to:
+/// - all RPCs and services;
+/// - requests directed only to specific services; or
+/// - requests directed only to specific methods (of a specific service).
+///
+/// - SeeAlso: ``ClientInterceptor`` for more information on client interceptors, and
+///  ``ServerInterceptorPipelineOperation`` for the server-side version of this type.
+public struct ClientInterceptorPipelineOperation: Sendable {
+  /// The subject of a ``ClientInterceptorPipelineOperation``.
+  /// The subject of an interceptor can either be all services and methods, only specific services, or only specific methods.
+  public struct Subject: Sendable {
+    internal enum Wrapped: Sendable {
+      case all
+      case services(Set<ServiceDescriptor>)
+      case methods(Set<MethodDescriptor>)
+    }
+
+    private let wrapped: Wrapped
+
+    /// An operation subject specifying an interceptor that applies to all RPCs across all services will be registered with this client.
+    public static var all: Self { .init(wrapped: .all) }
+
+    /// An operation subject specifying an interceptor that will be applied only to RPCs directed to the specified services.
+    /// - Parameters:
+    ///   - services: The list of service names for which this interceptor should intercept RPCs.
+    /// - Returns: A ``ClientInterceptorPipelineOperation``.
+    public static func services(_ services: Set<ServiceDescriptor>) -> Self {
+      Self(wrapped: .services(services))
+    }
+
+    /// An operation subject specifying an interceptor that will be applied only to RPCs directed to the specified service methods.
+    /// - Parameters:
+    ///   - methods: The list of method descriptors for which this interceptor should intercept RPCs.
+    /// - Returns: A ``ClientInterceptorPipelineOperation``.
+    public static func methods(_ methods: Set<MethodDescriptor>) -> Self {
+      Self(wrapped: .methods(methods))
+    }
+
+    @usableFromInline
+    internal func applies(to descriptor: MethodDescriptor) -> Bool {
+      switch self.wrapped {
+      case .all:
+        return true
+
+      case .services(let services):
+        return services.map({ $0.fullyQualifiedService }).contains(descriptor.service)
+
+      case .methods(let methods):
+        return methods.contains(descriptor)
+      }
+    }
+  }
+
+  /// The interceptor specified for this operation.
+  public let interceptor: any ClientInterceptor
+
+  @usableFromInline
+  internal let subject: Subject
+
+  private init(interceptor: any ClientInterceptor, appliesTo: Subject) {
+    self.interceptor = interceptor
+    self.subject = appliesTo
+  }
+
+  /// Create an operation, specifying which ``ClientInterceptor`` to apply and to which ``Subject``.
+  /// - Parameters:
+  ///   - interceptor: The ``ClientInterceptor`` to register with the client.
+  ///   - subject: The ``Subject`` to which the `interceptor` applies.
+  /// - Returns: A ``ClientInterceptorPipelineOperation``.
+  public static func apply(_ interceptor: any ClientInterceptor, to subject: Subject) -> Self {
+    Self(interceptor: interceptor, appliesTo: subject)
+  }
+
+  /// Returns whether this ``ClientInterceptorPipelineOperation`` applies to the given `descriptor`.
+  /// - Parameter descriptor: A ``MethodDescriptor`` for which to test whether this interceptor applies.
+  /// - Returns: `true` if this interceptor applies to the given `descriptor`, or `false` otherwise.
+  @inlinable
+  internal func applies(to descriptor: MethodDescriptor) -> Bool {
+    self.subject.applies(to: descriptor)
+  }
+}

Sources/GRPCCore/Call/Server/Internal/ServerRPCExecutor.swift

@@ -23,7 +23,8 @@ struct ServerRPCExecutor {
   ///   - stream: The accepted stream to execute the RPC on.
   ///   - deserializer: A deserializer for messages received from the client.
   ///   - serializer: A serializer for messages to send to the client.
-  ///   - interceptors: Server interceptors to apply to this RPC.
+  ///   - interceptors: Server interceptors to apply to this RPC.  The
+  ///       interceptors will be called in the order of the array.
   ///   - handler: A handler which turns the request into a response.
   @inlinable
   static func execute<Input, Output>(

Sources/GRPCCore/Call/Server/ServerInterceptorPipelineOperation.swift

@@ -21,7 +21,8 @@
 /// - requests directed only to specific services registered with your server; or
 /// - requests directed only to specific methods (of a specific service).
 ///
-/// - SeeAlso: ``ServerInterceptor`` for more information on server interceptors.
+/// - SeeAlso: ``ServerInterceptor`` for more information on server interceptors, and
+///  ``ClientInterceptorPipelineOperation`` for the client-side version of this type.
 public struct ServerInterceptorPipelineOperation: Sendable {
   /// The subject of a ``ServerInterceptorPipelineOperation``.
   /// The subject of an interceptor can either be all services and methods, only specific services, or only specific methods.

Sources/GRPCCore/GRPCClient.swift

@@ -112,34 +112,53 @@ public final class GRPCClient: Sendable {
   /// The transport which provides a bidirectional communication channel with the server.
   private let transport: any ClientTransport
 
-  /// A collection of interceptors providing cross-cutting functionality to each accepted RPC.
-  ///
-  /// The order in which interceptors are added reflects the order in which they are called. The
-  /// first interceptor added will be the first interceptor to intercept each request. The last
-  /// interceptor added will be the final interceptor to intercept each request before calling
-  /// the appropriate handler.
-  private let interceptors: [any ClientInterceptor]
+  private let interceptorPipeline: [ClientInterceptorPipelineOperation]
 
   /// The current state of the client.
   private let state: Mutex<State>
 
   /// The state of the client.
   private enum State: Sendable {
+
     /// The client hasn't been started yet. Can transition to `running` or `stopped`.
-    case notStarted
+    case notStarted(
+      /// A collection of interceptors providing cross-cutting functionality to each accepted RPC, keyed by the method to which they apply.
+      ///
+      /// The list of interceptors for each method is computed from `interceptorsPipeline` when calling a method for the first time.
+      /// This caching is done to avoid having to compute the applicable interceptors for each request made.
+      ///
+      /// The order in which interceptors are added reflects the order in which they are called. The
+      /// first interceptor added will be the first interceptor to intercept each request. The last
+      /// interceptor added will be the final interceptor to intercept each request before calling
+      /// the appropriate handler.
+      interceptorsPerMethod: [MethodDescriptor: [any ClientInterceptor]]
+    )
     /// The client is running and can send RPCs. Can transition to `stopping`.
-    case running
+    case running(
+      /// A collection of interceptors providing cross-cutting functionality to each accepted RPC, keyed by the method to which they apply.
+      ///
+      /// The list of interceptors for each method is computed from `interceptorsPipeline` when calling a method for the first time.
+      /// This caching is done to avoid having to compute the applicable interceptors for each request made.
+      ///
+      /// The order in which interceptors are added reflects the order in which they are called. The
+      /// first interceptor added will be the first interceptor to intercept each request. The last
+      /// interceptor added will be the final interceptor to intercept each request before calling
+      /// the appropriate handler.
+      interceptorsPerMethod: [MethodDescriptor: [any ClientInterceptor]]
+    )
     /// The client is stopping and no new RPCs will be sent. Existing RPCs may run to
     /// completion. May transition to `stopped`.
     case stopping
     /// The client has stopped, no RPCs are in flight and no more will be accepted. This state
     /// is terminal.
     case stopped
+    /// Temporary state to avoid CoWs.
+    case _modifying
 
     mutating func run() throws {
       switch self {
-      case .notStarted:
-        self = .running
+      case .notStarted(let interceptorsPerMethod):
+        self = .running(interceptorsPerMethod: interceptorsPerMethod)
 
       case .running:
         throw RuntimeError(
@@ -152,6 +171,9 @@ public final class GRPCClient: Sendable {
           code: .clientIsStopped,
           message: "The client has stopped and can only be started once."
         )
+
+      case ._modifying:
+        fatalError("Internal inconsistency")
       }
     }
 
@@ -169,6 +191,8 @@ public final class GRPCClient: Sendable {
         return true
       case .stopping, .stopped:
         return false
+      case ._modifying:
+        fatalError("Internal inconsistency")
       }
     }
 
@@ -183,6 +207,8 @@ public final class GRPCClient: Sendable {
           code: .clientIsStopped,
           message: "Client has been stopped. Can't make any more RPCs."
         )
+      case ._modifying:
+        fatalError("Internal inconsistency")
       }
     }
   }
@@ -191,18 +217,37 @@ public final class GRPCClient: Sendable {
   ///
   /// - Parameters:
   ///   - transport: The transport used to establish a communication channel with a server.
-  ///   - interceptors: A collection of interceptors providing cross-cutting functionality to each
+  ///   - interceptors: A collection of ``ClientInterceptor``s providing cross-cutting functionality to each
   ///       accepted RPC. The order in which interceptors are added reflects the order in which they
   ///       are called. The first interceptor added will be the first interceptor to intercept each
   ///       request. The last interceptor added will be the final interceptor to intercept each
   ///       request before calling the appropriate handler.
-  public init(
+  convenience public init(
     transport: some ClientTransport,
     interceptors: [any ClientInterceptor] = []
+  ) {
+    self.init(
+      transport: transport,
+      interceptorPipeline: interceptors.map { .apply($0, to: .all) }
+    )
+  }
+
+  /// Creates a new client with the given transport, interceptors and configuration.
+  ///
+  /// - Parameters:
+  ///   - transport: The transport used to establish a communication channel with a server.
+  ///   - interceptorPipeline: A collection of ``ClientInterceptorPipelineOperation`` providing cross-cutting
+  ///       functionality to each accepted RPC. Only applicable interceptors from the pipeline will be applied to each RPC.
+  ///       The order in which interceptors are added reflects the order in which they are called.
+  ///       The first interceptor added will be the first interceptor to intercept each request.
+  ///       The last interceptor added will be the final interceptor to intercept each request before calling the appropriate handler.
+  public init(
+    transport: some ClientTransport,
+    interceptorPipeline: [ClientInterceptorPipelineOperation]
   ) {
     self.transport = transport
-    self.interceptors = interceptors
-    self.state = Mutex(.notStarted)
+    self.interceptorPipeline = interceptorPipeline
+    self.state = Mutex(.notStarted(interceptorsPerMethod: [:]))
   }
 
   /// Start the client.
@@ -361,14 +406,50 @@ public final class GRPCClient: Sendable {
     var options = options
     options.formUnion(with: methodConfig)
 
+    let applicableInterceptors = self.state.withLock {
+      switch $0 {
+      case .notStarted(var interceptorsPerMethod):
+        if let interceptors = interceptorsPerMethod[descriptor] {
+          return interceptors
+        } else {
+          $0 = ._modifying
+          let interceptors = self.interceptorPipeline
+            .filter { $0.applies(to: descriptor) }
+            .map { $0.interceptor }
+          interceptorsPerMethod[descriptor] = interceptors
+          $0 = .notStarted(interceptorsPerMethod: interceptorsPerMethod)
+          return interceptors
+        }
+
+      case .running(var interceptorsPerMethod):
+        if let interceptors = interceptorsPerMethod[descriptor] {
+          return interceptors
+        } else {
+          $0 = ._modifying
+          let interceptors = self.interceptorPipeline
+            .filter { $0.applies(to: descriptor) }
+            .map { $0.interceptor }
+          interceptorsPerMethod[descriptor] = interceptors
+          $0 = .running(interceptorsPerMethod: interceptorsPerMethod)
+          return interceptors
+        }
+
+      case .stopping, .stopped:
+        fatalError("The checkExecutable call should have failed.")
+
+      case ._modifying:
+        fatalError("Internal inconsistency")
+      }
+    }
+
     return try await ClientRPCExecutor.execute(
       request: request,
       method: descriptor,
       options: options,
       serializer: serializer,
       deserializer: deserializer,
       transport: self.transport,
-      interceptors: self.interceptors,
+      interceptors: applicableInterceptors,
       handler: handler
     )
   }

Tests/GRPCCoreTests/Call/Client/ClientInterceptorPipelineOperationTests.swift

@@ -0,0 +1,66 @@
+/*
+ * 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 Testing
+
+@testable import GRPCCore
+
+@Suite("ClientInterceptorPipelineOperation")
+struct ClientInterceptorPipelineOperationTests {
+  @Suite("Applies to")
+  struct AppliesToTests {
+    @Test
+    func all() async throws {
+      let operation = ClientInterceptorPipelineOperation.apply(
+        .requestCounter(.init()),
+        to: .all
+      )
+
+      #expect(operation.applies(to: MethodDescriptor(service: "foo", method: "bar")))
+      #expect(operation.applies(to: MethodDescriptor(service: "foo", method: "baz")))
+      #expect(operation.applies(to: MethodDescriptor(service: "bar", method: "foo")))
+      #expect(operation.applies(to: MethodDescriptor(service: "bar", method: "baz")))
+    }
+
+    @Test
+    func serviceSpecific() async throws {
+      let operation = ClientInterceptorPipelineOperation.apply(
+        .requestCounter(.init()),
+        to: .services(Set([ServiceDescriptor(package: "pkg", service: "foo")]))
+      )
+
+      #expect(operation.applies(to: MethodDescriptor(service: "pkg.foo", method: "bar")))
+      #expect(operation.applies(to: MethodDescriptor(service: "pkg.foo", method: "baz")))
+
+      #expect(!operation.applies(to: MethodDescriptor(service: "pkg.bar", method: "foo")))
+      #expect(!operation.applies(to: MethodDescriptor(service: "pkg.bar", method: "baz")))
+    }
+
+    @Test
+    func methodSpecific() async throws {
+      let operation = ClientInterceptorPipelineOperation.apply(
+        .requestCounter(.init()),
+        to: .methods(Set([MethodDescriptor(service: "bar", method: "foo")]))
+      )
+
+      #expect(operation.applies(to: MethodDescriptor(service: "bar", method: "foo")))
+
+      #expect(!operation.applies(to: MethodDescriptor(service: "foo", method: "bar")))
+      #expect(!operation.applies(to: MethodDescriptor(service: "foo", method: "baz")))
+      #expect(!operation.applies(to: MethodDescriptor(service: "bar", method: "baz")))
+    }
+  }
+}