Allow server handlers to send response headers directly (#1599)

Motivation:

The async server call context allows users to set headers which are sent
when the first message is sent. In many cases this is fine, however,
some use cases require the headers to be sent immediately.

Modifications:

- Add `sendHeaders(_:)` to the `GRPCAsyncServerCallContext` which sends
  headers to the client and throws if headers have already been written
  or it's too late to send them.

Result:

Headers can be sent directly from a server call handler.

Author: glbrntt

Sources/GRPC/AsyncAwaitSupport/AsyncServerHandler/ServerHandlerStateMachine/ServerHandlerStateMachine+Actions.swift

@@ -70,12 +70,13 @@ extension ServerHandlerStateMachine {
 
     /// Update the metadata. It must not have been written yet.
     @inlinable
-    mutating func update(_ metadata: HPACKHeaders) {
+    mutating func update(_ metadata: HPACKHeaders) -> Bool {
       switch self {
       case .notWritten:
         self = .notWritten(metadata)
+        return true
       case .written:
-        assertionFailure("Metadata must not be set after it has been sent")
+        return false
       }
     }
 

Sources/GRPC/AsyncAwaitSupport/AsyncServerHandler/ServerHandlerStateMachine/ServerHandlerStateMachine+Draining.swift

@@ -50,16 +50,16 @@ extension ServerHandlerStateMachine {
     @inlinable
     mutating func setResponseHeaders(
       _ metadata: HPACKHeaders
-    ) -> Self.NextStateAndOutput<Void> {
-      self.responseHeaders.update(metadata)
-      return .init(nextState: .draining(self))
+    ) -> Self.NextStateAndOutput<Bool> {
+      let output = self.responseHeaders.update(metadata)
+      return .init(nextState: .draining(self), output: output)
     }
 
     @inlinable
     mutating func setResponseTrailers(
       _ metadata: HPACKHeaders
     ) -> Self.NextStateAndOutput<Void> {
-      self.responseTrailers.update(metadata)
+      _ = self.responseTrailers.update(metadata)
       return .init(nextState: .draining(self))
     }
 

Sources/GRPC/AsyncAwaitSupport/AsyncServerHandler/ServerHandlerStateMachine/ServerHandlerStateMachine+Finished.swift

@@ -35,8 +35,8 @@ extension ServerHandlerStateMachine {
     @inlinable
     mutating func setResponseHeaders(
       _ headers: HPACKHeaders
-    ) -> Self.NextStateAndOutput<Void> {
-      return .init(nextState: .finished(self))
+    ) -> Self.NextStateAndOutput<Bool> {
+      return .init(nextState: .finished(self), output: false)
     }
 
     @inlinable

Sources/GRPC/AsyncAwaitSupport/AsyncServerHandler/ServerHandlerStateMachine/ServerHandlerStateMachine+Handling.swift

@@ -50,16 +50,16 @@ extension ServerHandlerStateMachine {
     @inlinable
     mutating func setResponseHeaders(
       _ metadata: HPACKHeaders
-    ) -> Self.NextStateAndOutput<Void> {
-      self.responseHeaders.update(metadata)
-      return .init(nextState: .handling(self))
+    ) -> Self.NextStateAndOutput<Bool> {
+      let output = self.responseHeaders.update(metadata)
+      return .init(nextState: .handling(self), output: output)
     }
 
     @inlinable
     mutating func setResponseTrailers(
       _ metadata: HPACKHeaders
     ) -> Self.NextStateAndOutput<Void> {
-      self.responseTrailers.update(metadata)
+      _ = self.responseTrailers.update(metadata)
       return .init(nextState: .handling(self))
     }
 

Sources/GRPC/AsyncAwaitSupport/AsyncServerHandler/ServerHandlerStateMachine/ServerHandlerStateMachine.swift

@@ -27,7 +27,7 @@ internal struct ServerHandlerStateMachine {
   }
 
   @inlinable
-  mutating func setResponseHeaders(_ headers: HPACKHeaders) {
+  mutating func setResponseHeaders(_ headers: HPACKHeaders) -> Bool {
     switch self.state {
     case var .handling(handling):
       let nextStateAndOutput = handling.setResponseHeaders(headers)

Sources/GRPC/AsyncAwaitSupport/GRPCAsyncServerCallContext.swift

@@ -30,6 +30,21 @@ public struct GRPCAsyncServerCallContext: Sendable {
     Response(contextProvider: self.contextProvider)
   }
 
+  /// Notifies the client that the RPC has been accepted for processing by the server.
+  ///
+  /// On accepting the RPC the server will send the given headers (which may be empty) along with
+  /// any transport specific headers (such the ":status" pseudo header) to the client.
+  ///
+  /// It is not necessary to call this function: the RPC is implicitly accepted when the first
+  /// response message is sent, however this may be useful when clients require an early indication
+  /// that the RPC has been accepted.
+  ///
+  /// If the RPC has already been accepted (either implicitly or explicitly) then this function is
+  /// a no-op.
+  public func acceptRPC(headers: HPACKHeaders) async {
+    await self.contextProvider.acceptRPC(headers)
+  }
+
   /// Access the ``UserInfo`` dictionary which is shared with the interceptor contexts for this RPC.
   ///
   /// - Important: While ``UserInfo`` has value-semantics, this function accesses a reference
@@ -87,8 +102,8 @@ extension GRPCAsyncServerCallContext {
     /// Set the metadata to return at the start of the RPC.
     ///
     /// - Important: If this is required it should be updated _before_ the first response is sent
-    ///   via the response stream writer. Updates must not be made after the first response has
-    ///   been sent.
+    ///   via the response stream writer. Updates must not be made after the RPC has been accepted
+    ///   or the first response has been sent otherwise this method will throw an error.
     public func setHeaders(_ headers: HPACKHeaders) async throws {
       try await self.contextProvider.setResponseHeaders(headers)
     }

Sources/GRPC/AsyncAwaitSupport/GRPCAsyncServerHandler.swift

@@ -207,6 +207,16 @@ internal final class AsyncServerHandler<
   @usableFromInline
   internal var compressResponsesIfPossible: Bool
 
+  /// The interceptor pipeline does not track flushing as a separate event. The flush decision is
+  /// included with metadata alongside each message. For the status and trailers the flush is
+  /// implicit. For headers we track whether to flush here.
+  ///
+  /// In most cases the flush will be delayed until the first message is flushed and this will
+  /// remain unset. However, this may be set when the server handler
+  /// uses ``GRPCAsyncServerCallContext/sendHeaders(_:)``.
+  @usableFromInline
+  internal var flushNextHeaders: Bool
+
   /// A state machine for the interceptor pipeline.
   @usableFromInline
   internal private(set) var interceptorStateMachine: ServerInterceptorStateMachine
@@ -265,6 +275,7 @@ internal final class AsyncServerHandler<
     self.errorDelegate = context.errorDelegate
     self.compressionEnabledOnRPC = context.encoding.isEnabled
     self.compressResponsesIfPossible = true
+    self.flushNextHeaders = false
     self.logger = context.logger
 
     self.userInfoRef = Ref(UserInfo())
@@ -685,7 +696,9 @@ internal final class AsyncServerHandler<
     switch self.interceptorStateMachine.interceptedResponseMetadata() {
     case .forward:
       if let responseWriter = self.responseWriter {
-        responseWriter.sendMetadata(metadata, flush: false, promise: promise)
+        let flush = self.flushNextHeaders
+        self.flushNextHeaders = false
+        responseWriter.sendMetadata(metadata, flush: flush, promise: promise)
       } else if let promise = promise {
         promise.fail(GRPCStatus.processingError)
       }
@@ -747,11 +760,44 @@ extension AsyncServerHandler: AsyncServerCallContextProvider {
   @usableFromInline
   internal func setResponseHeaders(_ headers: HPACKHeaders) async throws {
     let completed = self.eventLoop.submit {
-      self.handlerStateMachine.setResponseHeaders(headers)
+      if !self.handlerStateMachine.setResponseHeaders(headers) {
+        throw GRPCStatus(
+          code: .failedPrecondition,
+          message: "Tried to send response headers in an invalid state"
+        )
+      }
     }
     try await completed.get()
   }
 
+  @usableFromInline
+  internal func acceptRPC(_ headers: HPACKHeaders) async {
+    let completed = self.eventLoop.submit {
+      guard self.handlerStateMachine.setResponseHeaders(headers) else { return }
+
+      // Shh,it's a lie! We don't really have a message to send but the state machine doesn't know
+      // (or care) about that. It will, however, tell us if we can send the headers or not.
+      switch self.handlerStateMachine.sendMessage() {
+      case let .intercept(.some(headers)):
+        switch self.interceptorStateMachine.interceptResponseMetadata() {
+        case .intercept:
+          self.flushNextHeaders = true
+          self.interceptors?.send(.metadata(headers), promise: nil)
+        case .cancel:
+          return self.cancel(error: nil)
+        case .drop:
+          ()
+        }
+
+      case .intercept(.none), .drop:
+        // intercept(.none) means headers have already been sent; we should never hit this because
+        // we guard on setting the response headers above.
+        ()
+      }
+    }
+    try? await completed.get()
+  }
+
   @usableFromInline
   internal func setResponseTrailers(_ headers: HPACKHeaders) async throws {
     let completed = self.eventLoop.submit {
@@ -798,6 +844,7 @@ extension AsyncServerHandler: AsyncServerCallContextProvider {
 @usableFromInline
 protocol AsyncServerCallContextProvider: Sendable {
   func setResponseHeaders(_ headers: HPACKHeaders) async throws
+  func acceptRPC(_ headers: HPACKHeaders) async
   func setResponseTrailers(_ trailers: HPACKHeaders) async throws
   func setResponseCompression(_ enabled: Bool) async throws
 

Tests/GRPCTests/AsyncAwaitSupport/AsyncServerHandler/ServerHandlerStateMachine/ServerHandlerStateMachineTests.swift

@@ -210,15 +210,15 @@ internal final class ServerHandlerStateMachineTests: GRPCTestCase {
 
   func testSetResponseHeadersWhenHandling() {
     var stateMachine = self.makeStateMachine(inState: .handling)
-    stateMachine.setResponseHeaders(["foo": "bar"])
+    XCTAssertTrue(stateMachine.setResponseHeaders(["foo": "bar"]))
     stateMachine.sendMessage().assertInterceptHeadersThenMessage { headers in
       XCTAssertEqual(headers, ["foo": "bar"])
     }
   }
 
   func testSetResponseHeadersWhenHandlingAreMovedToDraining() {
     var stateMachine = self.makeStateMachine(inState: .handling)
-    stateMachine.setResponseHeaders(["foo": "bar"])
+    XCTAssertTrue(stateMachine.setResponseHeaders(["foo": "bar"]))
     stateMachine.handleEnd().assertForward()
     stateMachine.sendMessage().assertInterceptHeadersThenMessage { headers in
       XCTAssertEqual(headers, ["foo": "bar"])
@@ -227,16 +227,15 @@ internal final class ServerHandlerStateMachineTests: GRPCTestCase {
 
   func testSetResponseHeadersWhenDraining() {
     var stateMachine = self.makeStateMachine(inState: .draining)
-    stateMachine.setResponseHeaders(["foo": "bar"])
+    XCTAssertTrue(stateMachine.setResponseHeaders(["foo": "bar"]))
     stateMachine.sendMessage().assertInterceptHeadersThenMessage { headers in
       XCTAssertEqual(headers, ["foo": "bar"])
     }
   }
 
   func testSetResponseHeadersWhenFinished() {
     var stateMachine = self.makeStateMachine(inState: .finished)
-    stateMachine.setResponseHeaders(["foo": "bar"])
-    // Nothing we can assert on, only that we don't crash.
+    XCTAssertFalse(stateMachine.setResponseHeaders(["foo": "bar"]))
   }
 
   func testSetResponseTrailersWhenHandling() {