Extend lifetime of client interceptor pipeline (grpc#1265)

Motivation:

A client call (i.e. the object the user holds) may live longer than the
transport associated with it (roughly speaking, the http/2 stream channel). An
example of this is when interceptors are use to retry and RPC and
redirect responses back to the original call.

However, the interceptor pipeline is held by the transport and is
currently set to nil when the transport is removed from the channel.
This means events invoked from the call object (such as cancellation)
which go via the transport (holding the interceptor pipeline) are
incorrectly failed.

Modifications:

- Have the client interceptor pipeline break the ref cycle between the
  transport and itself when the interceptor pipeline closes rather than
  when the transport is closed
- Emit a cancellation status rater than error on cancellation
- Update the ordering of when close is called in the interceptor
  pipeline.
- Add and update tests

Result:

"sub"-RPCs may be cancelled.

Author: glbrntt

Sources/GRPC/Interceptor/ClientInterceptorPipeline.swift

@@ -79,16 +79,17 @@ internal final class ClientInterceptorPipeline<Request, Response> {
   internal let _errorDelegate: ClientErrorDelegate?
 
   @usableFromInline
-  internal let _onError: (Error) -> Void
+  internal private(set) var _onError: ((Error) -> Void)?
 
   @usableFromInline
-  internal let _onCancel: (EventLoopPromise<Void>?) -> Void
+  internal private(set) var _onCancel: ((EventLoopPromise<Void>?) -> Void)?
 
   @usableFromInline
-  internal let _onRequestPart: (GRPCClientRequestPart<Request>, EventLoopPromise<Void>?) -> Void
+  internal private(set) var _onRequestPart:
+    ((GRPCClientRequestPart<Request>, EventLoopPromise<Void>?) -> Void)?
 
   @usableFromInline
-  internal let _onResponsePart: (GRPCClientResponsePart<Response>) -> Void
+  internal private(set) var _onResponsePart: ((GRPCClientResponsePart<Response>) -> Void)?
 
   /// The index after the last user interceptor context index. (i.e. `_userContexts.endIndex`).
   @usableFromInline
@@ -217,9 +218,13 @@ internal final class ClientInterceptorPipeline<Request, Response> {
 
     case self._tailIndex:
       if part.isEnd {
+        // Update our state before handling the response part.
+        self._isOpen = false
+        self._onResponsePart?(part)
         self.close()
+      } else {
+        self._onResponsePart?(part)
       }
-      self._onResponsePart(part)
 
     default:
       self._userContexts[index].invokeReceive(part)
@@ -275,9 +280,8 @@ internal final class ClientInterceptorPipeline<Request, Response> {
   /// Handles a caught error which has traversed the interceptor pipeline.
   @usableFromInline
   internal func _errorCaught(_ error: Error) {
-    // We're about to complete, close the pipeline.
-    self.close()
-
+    // We're about to call out to an error handler: update our state first.
+    self._isOpen = false
     var unwrappedError: Error
 
     // Unwrap the error, if possible.
@@ -295,7 +299,10 @@ internal final class ClientInterceptorPipeline<Request, Response> {
     }
 
     // Emit the unwrapped error.
-    self._onError(unwrappedError)
+    self._onError?(unwrappedError)
+
+    // Close the pipeline.
+    self.close()
   }
 
   /// Writes a request message into the interceptor pipeline.
@@ -351,7 +358,7 @@ internal final class ClientInterceptorPipeline<Request, Response> {
   ) {
     switch index {
     case self._headIndex:
-      self._onRequestPart(part, promise)
+      self._onRequestPart?(part, promise)
 
     case self._tailIndex:
       self._invokeSend(
@@ -407,7 +414,7 @@ internal final class ClientInterceptorPipeline<Request, Response> {
   ) {
     switch index {
     case self._headIndex:
-      self._onCancel(promise)
+      self._onCancel?(promise)
 
     case self._tailIndex:
       self._invokeCancel(
@@ -425,7 +432,7 @@ internal final class ClientInterceptorPipeline<Request, Response> {
 
 extension ClientInterceptorPipeline {
   /// Closes the pipeline. This should be called once, by the tail interceptor, to indicate that
-  /// the RPC has completed.
+  /// the RPC has completed. If this is not called, we will leak.
   /// - Important: This *must* to be called from the `eventLoop`.
   @inlinable
   internal func close() {
@@ -437,7 +444,14 @@ extension ClientInterceptorPipeline {
     self._scheduledClose = nil
 
     // Cancel the transport.
-    self._onCancel(nil)
+    self._onCancel?(nil)
+
+    // `ClientTransport` holds a reference to us and references to itself via these callbacks. Break
+    // these references now by replacing the callbacks.
+    self._onError = nil
+    self._onCancel = nil
+    self._onRequestPart = nil
+    self._onResponsePart = nil
   }
 
   /// Sets up a deadline for the pipeline.

Sources/GRPC/Interceptor/ClientTransport.swift

@@ -85,8 +85,8 @@ internal final class ClientTransport<Request, Response> {
   // trailers here and only forward them when we receive the status.
   private var trailers: HPACKHeaders?
 
-  /// The interceptor pipeline connected to this transport. This must be set to `nil` when removed
-  /// from the `ChannelPipeline` in order to break reference cycles.
+  /// The interceptor pipeline connected to this transport. The pipeline also holds references
+  /// to `self` which are dropped when the interceptor pipeline is closed.
   @usableFromInline
   internal var _pipeline: ClientInterceptorPipeline<Request, Response>?
 
@@ -118,6 +118,7 @@ internal final class ClientTransport<Request, Response> {
     self.logger = logger
     self.serializer = serializer
     self.deserializer = deserializer
+    // The references to self held by the pipeline are dropped when it is closed.
     self._pipeline = ClientInterceptorPipeline(
       eventLoop: eventLoop,
       details: details,
@@ -241,7 +242,8 @@ extension ClientTransport {
 
     if self.state.cancel() {
       let error = GRPCError.RPCCancelledByClient()
-      self.forwardErrorToInterceptors(error)
+      let status = error.makeGRPCStatus()
+      self.forwardToInterceptors(.end(status, [:]))
       self.failBufferedWrites(with: error)
       self.channel?.close(mode: .all, promise: nil)
       self.channelPromise?.fail(error)
@@ -363,11 +365,9 @@ extension ClientTransport {
   private func dropReferences() {
     if self.callEventLoop.inEventLoop {
       self.channel = nil
-      self._pipeline = nil
     } else {
       self.callEventLoop.execute {
         self.channel = nil
-        self._pipeline = nil
       }
     }
   }

Tests/GRPCTests/ClientCallTests.swift

@@ -197,11 +197,7 @@ class ClientCallTests: GRPCTestCase {
     // Cancellation should succeed.
     assertThat(try get.cancel().wait(), .doesNotThrow())
 
-    // The status promise will fail.
-    assertThat(
-      try promise.futureResult.wait(),
-      .throws(.instanceOf(GRPCError.RPCCancelledByClient.self))
-    )
+    assertThat(try promise.futureResult.wait(), .hasCode(.cancelled))
 
     // Cancellation should now fail, we've already cancelled.
     assertThat(try get.cancel().wait(), .throws(.instanceOf(GRPCError.AlreadyComplete.self)))

Tests/GRPCTests/ClientCancellingTests.swift

@@ -27,7 +27,7 @@ class ClientCancellingTests: EchoTestCaseBase {
     call.cancel(promise: nil)
 
     call.response.whenFailure { error in
-      XCTAssertTrue(error is GRPCError.RPCCancelledByClient)
+      XCTAssertEqual((error as? GRPCStatus)?.code, .cancelled)
       responseReceived.fulfill()
     }
 
@@ -47,7 +47,7 @@ class ClientCancellingTests: EchoTestCaseBase {
     call.cancel(promise: nil)
 
     call.response.whenFailure { error in
-      XCTAssertTrue(error is GRPCError.RPCCancelledByClient)
+      XCTAssertEqual((error as? GRPCStatus)?.code, .cancelled)
       responseReceived.fulfill()
     }
 

Tests/GRPCTests/EchoHelpers/Interceptors/EchoInterceptorFactories.swift

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2021, 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 EchoModel
+import GRPC
+
+// MARK: - Client
+
+internal final class EchoClientInterceptors: Echo_EchoClientInterceptorFactoryProtocol {
+  internal typealias Factory = () -> ClientInterceptor<Echo_EchoRequest, Echo_EchoResponse>
+  private var factories: [Factory] = []
+
+  internal init(_ factories: Factory...) {
+    self.factories = factories
+  }
+
+  internal func register(_ factory: @escaping Factory) {
+    self.factories.append(factory)
+  }
+
+  private func makeInterceptors() -> [ClientInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.factories.map { $0() }
+  }
+
+  func makeGetInterceptors() -> [ClientInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+
+  func makeExpandInterceptors() -> [ClientInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+
+  func makeCollectInterceptors() -> [ClientInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+
+  func makeUpdateInterceptors() -> [ClientInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+}
+
+// MARK: - Server
+
+internal final class EchoServerInterceptors: Echo_EchoServerInterceptorFactoryProtocol {
+  internal typealias Factory = () -> ServerInterceptor<Echo_EchoRequest, Echo_EchoResponse>
+  private var factories: [Factory] = []
+
+  internal init(_ factories: Factory...) {
+    self.factories = factories
+  }
+
+  internal func register(_ factory: @escaping Factory) {
+    self.factories.append(factory)
+  }
+
+  private func makeInterceptors() -> [ServerInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.factories.map { $0() }
+  }
+
+  func makeGetInterceptors() -> [ServerInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+
+  func makeExpandInterceptors() -> [ServerInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+
+  func makeCollectInterceptors() -> [ServerInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+
+  func makeUpdateInterceptors() -> [ServerInterceptor<Echo_EchoRequest, Echo_EchoResponse>] {
+    return self.makeInterceptors()
+  }
+}