Add an RPC cancellation handler

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

@@ -0,0 +1,254 @@
+/*
+ * 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.
+ */
+
+private import Synchronization
+
+/// Stores cancellation state for an RPC on the server .
+package final class ServerCancellationManager: Sendable {
+  private let state: Mutex<State>
+
+  package init() {
+    self.state = Mutex(State())
+  }
+
+  /// Returns whether the RPC has been marked as cancelled.
+  package var isRPCCancelled: Bool {
+    self.state.withLock {
+      return $0.isRPCCancelled
+    }
+  }
+
+  /// Marks the RPC as cancelled, potentially running any cancellation handlers.
+  package func cancelRPC() {
+    switch self.state.withLock({ $0.cancelRPC() }) {
+    case .executeAndResume(let onCancelHandlers, let onCancelWaiters):
+      for handler in onCancelHandlers {
+        handler.handler()
+      }
+
+      for onCancelWaiter in onCancelWaiters {
+        switch onCancelWaiter {
+        case .taskCancelled:
+          ()
+        case .waiting(_, let continuation):
+          continuation.resume(returning: .rpc)
+        }
+      }
+
+    case .doNothing:
+      ()
+    }
+  }
+
+  /// Adds a handler which is invoked when the RPC is cancelled.
+  ///
+  /// - Returns: The ID of the handler, if it was added, or `nil` if the RPC is already cancelled.
+  package func addRPCCancelledHandler(_ handler: @Sendable @escaping () -> Void) -> UInt64? {
+    return self.state.withLock { state -> UInt64? in
+      state.addRPCCancelledHandler(handler)
+    }
+  }
+
+  /// Removes a handler by its ID.
+  package func removeRPCCancelledHandler(withID id: UInt64) {
+    self.state.withLock { state in
+      state.removeRPCCancelledHandler(withID: id)
+    }
+  }
+
+  /// Suspends until the RPC is cancelled or the `Task` is cancelled.
+  package func suspendUntilRPCIsCancelled() async throws(CancellationError) {
+    let id = self.state.withLock { $0.nextID() }
+
+    let source = await withTaskCancellationHandler {
+      await withCheckedContinuation { continuation in
+        let onAddWaiter = self.state.withLock {
+          $0.addRPCIsCancelledWaiter(continuation: continuation, withID: id)
+        }
+
+        switch onAddWaiter {
+        case .doNothing:
+          ()
+        case .complete(let continuation, let result):
+          continuation.resume(returning: result)
+        }
+      }
+    } onCancel: {
+      switch self.state.withLock({ $0.cancelRPCCancellationWaiter(withID: id) }) {
+      case .resume(let continuation, let result):
+        continuation.resume(returning: result)
+      case .doNothing:
+        ()
+      }
+    }
+
+    switch source {
+    case .rpc:
+      ()
+    case .task:
+      throw CancellationError()
+    }
+  }
+}
+
+extension ServerCancellationManager {
+  enum CancellationSource {
+    case rpc
+    case task
+  }
+
+  struct Handler: Sendable {
+    var id: UInt64
+    var handler: @Sendable () -> Void
+  }
+
+  enum Waiter: Sendable {
+    case waiting(UInt64, CheckedContinuation<CancellationSource, Never>)
+    case taskCancelled(UInt64)
+
+    var id: UInt64 {
+      switch self {
+      case .waiting(let id, _):
+        return id
+      case .taskCancelled(let id):
+        return id
+      }
+    }
+  }
+
+  struct State {
+    private var handlers: [Handler]
+    private var waiters: [Waiter]
+    private var _nextID: UInt64
+    var isRPCCancelled: Bool
+
+    mutating func nextID() -> UInt64 {
+      let id = self._nextID
+      self._nextID &+= 1
+      return id
+    }
+
+    init() {
+      self.handlers = []
+      self.waiters = []
+      self._nextID = 0
+      self.isRPCCancelled = false
+    }
+
+    mutating func cancelRPC() -> OnCancelRPC {
+      let onCancel: OnCancelRPC
+
+      if self.isRPCCancelled {
+        onCancel = .doNothing
+      } else {
+        self.isRPCCancelled = true
+        onCancel = .executeAndResume(self.handlers, self.waiters)
+        self.handlers = []
+        self.waiters = []
+      }
+
+      return onCancel
+    }
+
+    mutating func addRPCCancelledHandler(_ handler: @Sendable @escaping () -> Void) -> UInt64? {
+      if self.isRPCCancelled {
+        handler()
+        return nil
+      } else {
+        let id = self.nextID()
+        self.handlers.append(.init(id: id, handler: handler))
+        return id
+      }
+    }
+
+    mutating func removeRPCCancelledHandler(withID id: UInt64) {
+      if let index = self.handlers.firstIndex(where: { $0.id == id }) {
+        self.handlers.remove(at: index)
+      }
+    }
+
+    enum OnCancelRPC {
+      case executeAndResume([Handler], [Waiter])
+      case doNothing
+    }
+
+    enum OnAddWaiter {
+      case complete(CheckedContinuation<CancellationSource, Never>, CancellationSource)
+      case doNothing
+    }
+
+    mutating func addRPCIsCancelledWaiter(
+      continuation: CheckedContinuation<CancellationSource, Never>,
+      withID id: UInt64
+    ) -> OnAddWaiter {
+      let onAddWaiter: OnAddWaiter
+
+      if self.isRPCCancelled {
+        onAddWaiter = .complete(continuation, .rpc)
+      } else if let index = self.waiters.firstIndex(where: { $0.id == id }) {
+        switch self.waiters[index] {
+        case .taskCancelled:
+          onAddWaiter = .complete(continuation, .task)
+        case .waiting:
+          // There's already a continuation enqueued.
+          fatalError("Inconsistent state")
+        }
+      } else {
+        self.waiters.append(.waiting(id, continuation))
+        onAddWaiter = .doNothing
+      }
+
+      return onAddWaiter
+    }
+
+    enum OnCancelRPCCancellationWaiter {
+      case resume(CheckedContinuation<CancellationSource, Never>, CancellationSource)
+      case doNothing
+    }
+
+    mutating func cancelRPCCancellationWaiter(withID id: UInt64) -> OnCancelRPCCancellationWaiter {
+      let onCancelWaiter: OnCancelRPCCancellationWaiter
+
+      if let index = self.waiters.firstIndex(where: { $0.id == id }) {
+        let waiter = self.waiters.removeWithoutMaintainingOrder(at: index)
+        switch waiter {
+        case .taskCancelled:
+          onCancelWaiter = .doNothing
+        case .waiting(_, let continuation):
+          onCancelWaiter = .resume(continuation, .task)
+        }
+      } else {
+        self.waiters.append(.taskCancelled(id))
+        onCancelWaiter = .doNothing
+      }
+
+      return onCancelWaiter
+    }
+  }
+}
+
+extension Array {
+  fileprivate mutating func removeWithoutMaintainingOrder(at index: Int) -> Element {
+    let lastElementIndex = self.index(before: self.endIndex)
+
+    if index == lastElementIndex {
+      return self.remove(at: index)
+    } else {
+      self.swapAt(index, lastElementIndex)
+      return self.removeLast()
+    }
+  }
+}

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

@@ -119,43 +119,29 @@ struct ServerRPCExecutor {
       _ context: ServerContext
     ) async throws -> StreamingServerResponse<Output>
   ) async {
-    await withTaskGroup(of: ServerExecutorTask.self) { group in
+    await withTaskGroup(of: Void.self) { group in
       group.addTask {
-        let result = await Result {
+        do {
           try await Task.sleep(for: timeout, clock: .continuous)
+          context.cancellation.cancel()
+        } catch {
+          ()  // Only cancel the RPC if the timeout completes.
         }
-        return .timedOut(result)
       }
 
-      group.addTask {
-        await Self._processRPC(
-          context: context,
-          metadata: metadata,
-          inbound: inbound,
-          outbound: outbound,
-          deserializer: deserializer,
-          serializer: serializer,
-          interceptors: interceptors,
-          handler: handler
-        )
-        return .executed
-      }
-
-      while let next = await group.next() {
-        switch next {
-        case .timedOut(.success):
-          // Timeout expired; cancel the work.
-          group.cancelAll()
-
-        case .timedOut(.failure):
-          // Timeout failed (because it was cancelled). Wait for more tasks to finish.
-          ()
+      await Self._processRPC(
+        context: context,
+        metadata: metadata,
+        inbound: inbound,
+        outbound: outbound,
+        deserializer: deserializer,
+        serializer: serializer,
+        interceptors: interceptors,
+        handler: handler
+      )
 
-        case .executed:
-          // The work finished. Cancel any remaining tasks.
-          group.cancelAll()
-        }
-      }
+      // Cancel the timeout
+      group.cancelAll()
     }
   }