async-await: Move to completed state before cancelling task during finish()

Sources/GRPC/AsyncAwaitSupport/GRPCAsyncServerHandler.swift

@@ -327,6 +327,8 @@ internal final class AsyncServerHandler<
       self.state = .completed
 
     case .active:
+      self.state = .completed
+      self.interceptors = nil
       self.userHandlerTask?.cancel()
 
     case .completed:
@@ -524,8 +526,10 @@ internal final class AsyncServerHandler<
       self.interceptors.send(.message(response, metadata), promise: nil)
 
     case .completed:
-      /// If we are in the completed state then the async writer delegate must have terminated.
-      preconditionFailure()
+      /// If we are in the completed state then the async writer delegate will have been cancelled,
+      /// however the cancellation is asynchronous so there's a chance that we receive this callback
+      /// after that has happened. We can drop the response.
+      ()
     }
   }
 

Tests/GRPCTests/AsyncAwaitSupport/AsyncIntegrationTests.swift

@@ -49,9 +49,9 @@ final class AsyncIntegrationTests: GRPCTestCase {
   }
 
   override func tearDown() {
-    XCTAssertNoThrow(try self.client.close().wait())
-    XCTAssertNoThrow(try self.server.close().wait())
-    XCTAssertNoThrow(try self.group.syncShutdownGracefully())
+    XCTAssertNoThrow(try self.client?.close().wait())
+    XCTAssertNoThrow(try self.server?.close().wait())
+    XCTAssertNoThrow(try self.group?.syncShutdownGracefully())
     super.tearDown()
   }
 
@@ -195,6 +195,17 @@ final class AsyncIntegrationTests: GRPCTestCase {
       ])
     }
   }
+
+  func testServerCloseAfterMessage() {
+    XCTAsyncTest {
+      let update = self.echo.makeUpdateCall()
+      try await update.requestStream.send(.with { $0.text = "hello" })
+      _ = try await update.responses.first(where: { _ in true })
+      XCTAssertNoThrow(try self.server.close().wait())
+      self.server = nil  // So that tearDown() does not call close() again.
+      try await update.requestStream.finish()
+    }
+  }
 }
 
 extension HPACKHeaders {

Tests/GRPCTests/GRPCAsyncServerHandlerTests.swift

@@ -278,8 +278,8 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
 
     await assertThat(self.recorder.metadata, .nil())
     await assertThat(self.recorder.messages, .isEmpty())
-    await assertThat(self.recorder.status, .notNil(.hasCode(.unavailable)))
-    await assertThat(self.recorder.trailers, .is([:]))
+    await assertThat(self.recorder.status, .nil())
+    await assertThat(self.recorder.trailers, .nil())
   } }
 
   func testFinishAfterMessage() { XCTAsyncTest {
@@ -296,6 +296,38 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     // Wait for tasks to finish.
     await handler.userHandlerTask?.value
 
+    await assertThat(self.recorder.messages.first, .is(ByteBuffer(string: "hello")))
+    await assertThat(self.recorder.status, .nil())
+    await assertThat(self.recorder.trailers, .nil())
+  } }
+
+  func testErrorAfterHeaders() { XCTAsyncTest {
+    let handler = self.makeHandler(observer: self.echo(requests:responseStreamWriter:context:))
+
+    handler.receiveMetadata([:])
+    handler.receiveError(CancellationError())
+
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+
+    await assertThat(self.recorder.status, .notNil(.hasCode(.unavailable)))
+    await assertThat(self.recorder.trailers, .is([:]))
+  } }
+
+  func testErrorAfterMessage() { XCTAsyncTest {
+    let handler = self.makeHandler(observer: self.echo(requests:responseStreamWriter:context:))
+
+    handler.receiveMetadata([:])
+    handler.receiveMessage(ByteBuffer(string: "hello"))
+
+    // Wait for the async user function to have processed the message.
+    try self.recorder.recordedMessagePromise.futureResult.wait()
+
+    handler.receiveError(CancellationError())
+
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+
     await assertThat(self.recorder.messages.first, .is(ByteBuffer(string: "hello")))
     await assertThat(self.recorder.status, .notNil(.hasCode(.unavailable)))
     await assertThat(self.recorder.trailers, .is([:]))