Add support for RPC cancellation (#13)

Motivation:

grpc-swift has support for RPC cancellation via a cancellation handler.
It should be supported here as well.

Modifications:

- Modify the server stream handler so that it holds the cancellation
handle and triggers cancellation appropriately.
- Modify the server stream handling to set a cancellation handle on the
appropriate stream handler.
- Update tests

Result:

HTTP/2 server transport respects stream cancellation

Author: glbrntt

Package.swift

@@ -35,15 +35,15 @@ let products: [Product] = [
 let dependencies: [Package.Dependency] = [
   .package(
     url: "https://github.com/grpc/grpc-swift.git",
-    exact: "2.0.0-alpha.1"
+    branch: "main"
   ),
   .package(
     url: "https://github.com/apple/swift-nio.git",
     from: "2.65.0"
   ),
   .package(
     url: "https://github.com/apple/swift-nio-http2.git",
-    from: "1.32.0"
+    from: "1.34.1"
   ),
   .package(
     url: "https://github.com/apple/swift-nio-transport-services.git",

Sources/GRPCNIOTransportCore/Internal/NIOChannelPipeline+GRPC.swift

@@ -85,7 +85,8 @@ extension ChannelPipeline.SynchronousOperations {
           scheme: scheme,
           acceptedEncodings: compressionConfig.enabledAlgorithms,
           maxPayloadSize: rpcConfig.maxRequestPayloadSize,
-          methodDescriptorPromise: methodDescriptorPromise
+          methodDescriptorPromise: methodDescriptorPromise,
+          eventLoop: streamChannel.eventLoop
         )
         try streamChannel.pipeline.syncOperations.addHandler(streamHandler)
 

Sources/GRPCNIOTransportCore/Server/CommonHTTP2ServerTransport.swift

@@ -241,19 +241,35 @@ package final class CommonHTTP2ServerTransport<
         return
       }
 
-      let rpcStream = RPCStream(
-        descriptor: descriptor,
-        inbound: RPCAsyncSequence(wrapping: inbound),
-        outbound: RPCWriter.Closable(
-          wrapping: ServerConnection.Stream.Outbound(
-            responseWriter: outbound,
-            http2Stream: stream
+      await withServerContextRPCCancellationHandle { handle in
+        stream.channel.eventLoop.execute {
+          // Sync is safe: this is on the right event loop.
+          let sync = stream.channel.pipeline.syncOperations
+
+          do {
+            let handler = try sync.handler(type: GRPCServerStreamHandler.self)
+            handler.setCancellationHandle(handle)
+          } catch {
+            // Looking up the handler can fail if the channel is already closed, in which case
+            // don't execute the RPC, just return early.
+            return
+          }
+        }
+
+        let rpcStream = RPCStream(
+          descriptor: descriptor,
+          inbound: RPCAsyncSequence(wrapping: inbound),
+          outbound: RPCWriter.Closable(
+            wrapping: ServerConnection.Stream.Outbound(
+              responseWriter: outbound,
+              http2Stream: stream
+            )
           )
         )
-      )
 
-      let context = ServerContext(descriptor: descriptor)
-      await streamHandler(rpcStream, context)
+        let context = ServerContext(descriptor: descriptor, cancellation: handle)
+        await streamHandler(rpcStream, context)
+      }
     }
   }
 

Sources/GRPCNIOTransportCore/Server/GRPCServerStreamHandler.swift

@@ -26,9 +26,11 @@ package final class GRPCServerStreamHandler: ChannelDuplexHandler, RemovableChan
   package typealias OutboundOut = HTTP2Frame.FramePayload
 
   private var stateMachine: GRPCStreamStateMachine
+  private let eventLoop: any EventLoop
 
   private var isReading = false
   private var flushPending = false
+  private var isCancelled = false
 
   // We buffer the final status + trailers to avoid reordering issues (i.e.,
   // if there are messages still not written into the channel because flush has
@@ -38,6 +40,8 @@ package final class GRPCServerStreamHandler: ChannelDuplexHandler, RemovableChan
 
   private let methodDescriptorPromise: EventLoopPromise<MethodDescriptor>
 
+  private var cancellationHandle: Optional<ServerContext.RPCCancellationHandle>
+
   // Existential errors unconditionally allocate, avoid this per-use allocation by doing it
   // statically.
   private static let handlerRemovedBeforeDescriptorResolved: any Error = RPCError(
@@ -50,6 +54,8 @@ package final class GRPCServerStreamHandler: ChannelDuplexHandler, RemovableChan
     acceptedEncodings: CompressionAlgorithmSet,
     maxPayloadSize: Int,
     methodDescriptorPromise: EventLoopPromise<MethodDescriptor>,
+    eventLoop: any EventLoop,
+    cancellationHandler: ServerContext.RPCCancellationHandle? = nil,
     skipStateMachineAssertions: Bool = false
   ) {
     self.stateMachine = .init(
@@ -58,12 +64,54 @@ package final class GRPCServerStreamHandler: ChannelDuplexHandler, RemovableChan
       skipAssertions: skipStateMachineAssertions
     )
     self.methodDescriptorPromise = methodDescriptorPromise
+    self.cancellationHandle = cancellationHandler
+    self.eventLoop = eventLoop
+  }
+
+  package func setCancellationHandle(_ handle: ServerContext.RPCCancellationHandle) {
+    if self.eventLoop.inEventLoop {
+      self.syncSetCancellationHandle(handle)
+    } else {
+      let loopBoundSelf = NIOLoopBound(self, eventLoop: self.eventLoop)
+      self.eventLoop.execute {
+        loopBoundSelf.value.syncSetCancellationHandle(handle)
+      }
+    }
+  }
+
+  private func syncSetCancellationHandle(_ handle: ServerContext.RPCCancellationHandle) {
+    assert(self.cancellationHandle == nil, "\(#function) must only be called once")
+
+    if self.isCancelled {
+      handle.cancel()
+    } else {
+      self.cancellationHandle = handle
+    }
+  }
+
+  private func cancelRPC() {
+    if let handle = self.cancellationHandle.take() {
+      handle.cancel()
+    } else {
+      self.isCancelled = true
+    }
   }
 }
 
 // - MARK: ChannelInboundHandler
 
 extension GRPCServerStreamHandler {
+  package func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+    switch event {
+    case is ChannelShouldQuiesceEvent:
+      self.cancelRPC()
+    default:
+      ()
+    }
+
+    context.fireUserInboundEventTriggered(event)
+  }
+
   package func channelRead(context: ChannelHandlerContext, data: NIOAny) {
     self.isReading = true
     let frame = self.unwrapInboundIn(data)
@@ -186,6 +234,7 @@ extension GRPCServerStreamHandler {
   ) {
     switch self.stateMachine.unexpectedInboundClose(reason: reason) {
     case .fireError_serverOnly(let wrappedError):
+      self.cancelRPC()
       context.fireErrorCaught(wrappedError)
     case .doNothing:
       ()

Tests/GRPCNIOTransportCoreTests/Client/Connection/Utilities/ConnectionTest.swift

@@ -117,7 +117,8 @@ extension ConnectionTest {
                 scheme: .http,
                 acceptedEncodings: .none,
                 maxPayloadSize: .max,
-                methodDescriptorPromise: channel.eventLoop.makePromise(of: MethodDescriptor.self)
+                methodDescriptorPromise: channel.eventLoop.makePromise(of: MethodDescriptor.self),
+                eventLoop: stream.eventLoop
               )
 
               return stream.eventLoop.makeCompletedFuture {

Tests/GRPCNIOTransportCoreTests/Client/Connection/Utilities/TestServer.swift

@@ -74,7 +74,8 @@ final class TestServer: Sendable {
               scheme: .http,
               acceptedEncodings: .all,
               maxPayloadSize: .max,
-              methodDescriptorPromise: channel.eventLoop.makePromise(of: MethodDescriptor.self)
+              methodDescriptorPromise: channel.eventLoop.makePromise(of: MethodDescriptor.self),
+              eventLoop: stream.eventLoop
             )
 
             try stream.pipeline.syncOperations.addHandlers(handler)