Add server stream events

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

@@ -120,43 +120,30 @@ struct ServerRPCExecutor {
       _ context: ServerContext
     ) async throws -> ServerResponse.Stream<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)
+          // Cancel the RPC if the timeout passes.
+          context.streamState.events.continuation.yield(.rpcCancelled)
+        } catch {
+          ()  // Sleep was cancelled, the RPC completed.
         }
-        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, if it's still running.
+      group.cancelAll()
     }
   }
 

Sources/GRPCCore/Call/Server/RPCRouter.swift

@@ -155,7 +155,7 @@ extension RPCRouter {
     context: ServerContext,
     interceptors: [any ServerInterceptor]
   ) async {
-    if let handler = self.handlers[stream.descriptor] {
+    if let handler = self.handlers[context.descriptor] {
       await handler.handle(stream: stream, context: context, interceptors: interceptors)
     } else {
       // If this throws then the stream must be closed which we can't do anything about, so ignore

Sources/GRPCCore/Call/Server/ServerContext.swift

@@ -15,12 +15,17 @@
  */
 
 /// Additional information about an RPC handled by a server.
+@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
 public struct ServerContext: Sendable {
   /// A description of the method being called.
   public var descriptor: MethodDescriptor
 
+  /// The state of the server stream.
+  public var streamState: ServerStreamState
+
   /// Create a new server context.
-  public init(descriptor: MethodDescriptor) {
+  public init(descriptor: MethodDescriptor, streamState: ServerStreamState) {
     self.descriptor = descriptor
+    self.streamState = streamState
   }
 }

Sources/GRPCCore/Call/Server/ServerStreamEvent.swift

@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+/// An out-of-band event which can happen to the underlying stream
+/// which an RPC is executing on.
+public struct ServerStreamEvent: Hashable, Sendable {
+  internal enum Value: Hashable, Sendable {
+    case rpcCancelled
+  }
+
+  internal var value: Value
+
+  private init(_ value: Value) {
+    self.value = value
+  }
+
+  /// The RPC was cancelled and the service should stop processing it.
+  ///
+  /// RPCs can be cancelled for a number of reasons including, but not limited to:
+  /// - it took too long to complete
+  /// - the client closed the underlying stream
+  /// - the stream closed unexpectedly (due to a network failure, for example)
+  /// - the server initiated a graceful shutdown
+  ///
+  /// You should stop processing the RPC and cleanup any associated state if you
+  /// receive this event.
+  public static let rpcCancelled = Self(.rpcCancelled)
+}
+
+extension ServerStreamEvent: CustomStringConvertible {
+  public var description: String {
+    String(describing: self.value)
+  }
+}

Sources/GRPCCore/Call/Server/ServerStreamState.swift

@@ -0,0 +1,213 @@
+/*
+ * 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 Synchronization
+
+@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
+public struct ServerStreamState: Sendable {
+  /// Returns whether the RPC has been cancelled.
+  ///
+  /// - SeeAlso: ``ServerStreamEvent/rpcCancelled``.
+  public var isRPCCancelled: Bool {
+    self.events.isRPCCancelled
+  }
+
+  /// Events which can happen to the underlying stream the RPC is being run on.
+  public let events: Events
+
+  private init(events: Events) {
+    self.events = events
+  }
+
+  public static func makeState() -> (streamState: Self, eventContinuation: Events.Continuation) {
+    let events = Events()
+    return (ServerStreamState(events: events), eventContinuation: events.continuation)
+  }
+}
+
+@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
+extension ServerStreamState {
+  /// An `AsyncSequence` of events which can happen to the stream.
+  ///
+  /// Each event will be delivered at most once.
+  ///
+  /// - Note: This sequence supports _multiple_ concurrent iterators.
+  public struct Events {
+    private let storage: Storage
+    @usableFromInline
+    internal let continuation: Continuation
+  }
+}
+
+@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
+extension ServerStreamState.Events: AsyncSequence, Sendable {
+  public typealias Element = ServerStreamEvent
+  public typealias Failure = Never
+
+  init() {
+    self.storage = Storage()
+    self.continuation = Continuation(storage: self.storage)
+  }
+
+  fileprivate var isRPCCancelled: Bool {
+    self.storage.eventSet(contains: .rpcCancelled)
+  }
+
+  public func makeAsyncIterator() -> AsyncIterator {
+    let streamEvents = AsyncStream.makeStream(of: ServerStreamEvent.self)
+    self.storage.registerContinuation(streamEvents.continuation)
+    return AsyncIterator(iterator: streamEvents.stream.makeAsyncIterator())
+  }
+
+  public struct AsyncIterator: AsyncIteratorProtocol {
+    private var iterator: AsyncStream<ServerStreamEvent>.AsyncIterator
+
+    fileprivate init(iterator: AsyncStream<ServerStreamEvent>.AsyncIterator) {
+      self.iterator = iterator
+    }
+
+    public mutating func next() async throws(Never) -> ServerStreamEvent? {
+      await self.next(isolation: nil)
+    }
+
+    public mutating func next(
+      isolation actor: isolated (any Actor)?
+    ) async throws(Never) -> ServerStreamEvent? {
+      return await self.iterator.next(isolation: actor)
+    }
+  }
+
+  public struct Continuation: Sendable {
+    private let storage: Storage
+
+    init(storage: Storage) {
+      self.storage = storage
+    }
+
+    /// Yield an event to the stream.
+    ///
+    /// - Important: Events are only delivered once. If the event has already been yielded
+    ///   then attempting to yield it again is a no-op.
+    /// - Parameter event: The event to yield.
+    public func yield(_ event: ServerStreamEvent) {
+      self.storage.yield(event)
+    }
+
+    /// Indicate that no more events will be delivered.
+    public func finish() {
+      self.storage.finish()
+    }
+  }
+}
+
+@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
+extension ServerStreamState.Events {
+  final class Storage: Sendable {
+    private let state: Mutex<State>
+
+    init() {
+      self.state = Mutex(State())
+    }
+
+    func eventSet(contains event: ServerStreamEvent) -> Bool {
+      self.state.withLock { $0.eventSet.contains(EventSet(event)) }
+    }
+
+    func registerContinuation(_ continuation: AsyncStream<ServerStreamEvent>.Continuation) {
+      self.state.withLock {
+        let events = $0.registerContinuation(continuation)
+
+        if events.contains(.rpcCancelled) {
+          continuation.yield(.rpcCancelled)
+        }
+
+        if events.contains(.finished) {
+          continuation.finish()
+        }
+      }
+    }
+
+    func yield(_ event: ServerStreamEvent) {
+      self.state.withLock {
+        for continuation in $0.publishStreamEvent(event) {
+          continuation.yield(event)
+        }
+      }
+    }
+
+    func finish() {
+      self.state.withLock {
+        for continuation in $0.finish() {
+          continuation.finish()
+        }
+      }
+    }
+  }
+
+  private struct EventSet: OptionSet, Hashable, Sendable {
+    var rawValue: UInt8
+
+    init(rawValue: UInt8) {
+      self.rawValue = rawValue
+    }
+
+    init(_ event: ServerStreamEvent) {
+      switch event.value {
+      case .rpcCancelled:
+        self = .rpcCancelled
+      }
+    }
+
+    static let finished = EventSet(rawValue: 1 << 0)
+    static let rpcCancelled = EventSet(rawValue: 1 << 1)
+  }
+
+  private struct State: Sendable {
+    private(set) var eventSet: EventSet
+    private var continuations: [AsyncStream<ServerStreamEvent>.Continuation]
+
+    init() {
+      self.eventSet = EventSet()
+      self.continuations = []
+    }
+
+    mutating func registerContinuation(
+      _ continuation: AsyncStream<ServerStreamEvent>.Continuation
+    ) -> EventSet {
+      if !self.eventSet.contains(.finished) {
+        self.continuations.append(continuation)
+      }
+
+      return self.eventSet
+    }
+
+    mutating func publishStreamEvent(
+      _ streamEvent: ServerStreamEvent
+    ) -> [AsyncStream<ServerStreamEvent>.Continuation] {
+      if self.eventSet.contains(.finished) {
+        return []
+      } else {
+        let (inserted, _) = self.eventSet.insert(EventSet(streamEvent))
+        return inserted ? self.continuations : []
+      }
+    }
+
+    mutating func finish() -> [AsyncStream<ServerStreamEvent>.Continuation] {
+      let (inserted, _) = self.eventSet.insert(.finished)
+      return inserted ? self.continuations : []
+    }
+  }
+}