fixup: Move tasks out of enum assoc data to avoid race

Signed-off-by: Si Beaumont <[email protected]>

Author: simonjbeaumont

Sources/GRPC/AsyncAwaitSupport/GRPCAsyncServerHandler.swift

@@ -195,11 +195,15 @@ internal final class AsyncServerHandler<
 
   /// The task used to run the async user function.
   ///
-  /// - TODO: Should this be part of the associated metadata in the `State` enum? Doing so would
-  /// make testing a bit cumbersome since a bunch of tests await this task finishing. Shoving it in
-  /// the enum.
+  /// - TODO: I'd like it if this was part of the assoc data for the .active state but doing so may introduce a race condition.
   @usableFromInline
-  internal var task: Task<Void, Never>? = nil
+  internal var userHandlerTask: Task<Void, Never>? = nil
+
+  /// The task used to drain the response stream writer after the user function has completed.
+  ///
+  /// - TODO: I'd like it if this was part of the assoc data for the .finishingSuccessfully state but doing so may introduce a race condition.
+  @usableFromInline
+  internal var responseStreamDrainTask: Task<Void, Never>? = nil
 
   @usableFromInline
   internal enum State {
@@ -250,9 +254,6 @@ internal final class AsyncServerHandler<
     /// handler has completed that have yet to be written. While in this state we will await the
     /// writer flushing its responses before sending `.end` to the stream.
     ///
-    /// - The `Task` is used to drain the response stream writer. It is stashed here so that we can
-    /// cancel it in the event of an error.
-    ///
     /// - The `EventLoopPromise` bridges the NIO and async-await worlds. It is the mechanism that we
     /// use to run a callback when the response stream has been flushed. It is fulfilled with the
     /// result of the async `Task` executing the user handler using `completeWithTask(_:)`.
@@ -264,7 +265,6 @@ internal final class AsyncServerHandler<
     case finishingSuccessfully(
       GRPCAsyncServerCallContext,
       GRPCAsyncResponseStreamWriter<Response>,
-      Task<Void, Never>,
       EventLoopPromise<GRPCStatus>
     )
 
@@ -341,10 +341,10 @@ internal final class AsyncServerHandler<
       self.state = .completed
 
     case .active:
-      self.task?.cancel()
+      self.userHandlerTask?.cancel()
 
-    case let .finishingSuccessfully(_, _, task, _):
-      task.cancel()
+    case .finishingSuccessfully:
+      self.responseStreamDrainTask?.cancel()
 
     case .completed:
       self.interceptors = nil
@@ -404,7 +404,7 @@ internal final class AsyncServerHandler<
       self.interceptors.send(.metadata([:]), promise: nil)
 
       // Spin up a task to call the async user handler.
-      self.task = userHandlerPromise.completeWithTask {
+      self.userHandlerTask = userHandlerPromise.completeWithTask {
         try await withTaskCancellationHandler {
           do {
             // Call the user function.
@@ -545,7 +545,7 @@ internal final class AsyncServerHandler<
         // Register callback for the response stream being drained.
         responseStreamDrainedPromise.futureResult.whenComplete(self.responseStreamDrained(_:))
 
-        let responseStreamDrainTask = responseStreamDrainedPromise.completeWithTask {
+        self.responseStreamDrainTask = responseStreamDrainedPromise.completeWithTask {
           try await withTaskCancellationHandler {
             // Await the writer finish.
             try await responseStreamWriter._asyncWriter.finish(())
@@ -561,7 +561,6 @@ internal final class AsyncServerHandler<
         self.state = .finishingSuccessfully(
           context,
           responseStreamWriter,
-          responseStreamDrainTask,
           responseStreamDrainedPromise
         )
 
@@ -587,7 +586,7 @@ internal final class AsyncServerHandler<
     case .active:
       preconditionFailure()
 
-    case let .finishingSuccessfully(context, _, _, _):
+    case let .finishingSuccessfully(context, _, _):
       switch result {
       case let .success(status):
         // Now we have drained the response stream writer from the user handler we can send end.
@@ -643,11 +642,10 @@ internal final class AsyncServerHandler<
       // which it is reading and give the user handler an opportunity to cleanup.
       //
       // NOTE: This line used to be before we explicitly fail the status promise but it was exaserbating a race condition and causing crashes. See https://bugs.swift.org/browse/SR-15108.
-      self.task?.cancel()
+      self.userHandlerTask?.cancel()
 
-    case let .finishingSuccessfully(_, _, responseStreamWriterDrainTask, _):
-      self.task?.cancel()
-      responseStreamWriterDrainTask.cancel()
+    case .finishingSuccessfully(_, _, _):
+      self.responseStreamDrainTask?.cancel()
 
     case .completed:
       ()

Tests/GRPCTests/GRPCAsyncServerHandlerTests.swift

@@ -81,8 +81,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     handler.receiveMessage(ByteBuffer(string: "3"))
     handler.receiveEnd()
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     handler.finish()
 
@@ -107,8 +108,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     handler.receiveMessage(ByteBuffer(string: "3"))
     handler.receiveEnd()
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(
       self.recorder.messages,
@@ -135,8 +137,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     handler.receiveMessage(ByteBuffer(string: "3"))
     handler.receiveEnd()
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(
       self.recorder.messages,
@@ -148,11 +151,11 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
   func testTaskOnlyCreatedAfterHeaders() { XCTAsyncTest {
     let handler = self.makeHandler(observer: self.echo(requests:responseStreamWriter:context:))
 
-    await assertThat(handler.task, .is(.nil()))
+    await assertThat(handler.userHandlerTask, .is(.nil()))
 
     handler.receiveMetadata([:])
 
-    await assertThat(handler.task, .is(.notNil()))
+    await assertThat(handler.userHandlerTask, .is(.notNil()))
   } }
 
   func testThrowingDeserializer() { XCTAsyncTest {
@@ -174,8 +177,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     let buffer = ByteBuffer(string: "hello")
     handler.receiveMessage(buffer)
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(self.recorder.messages, .isEmpty())
     await assertThat(self.recorder.status, .notNil(.hasCode(.internalError)))
@@ -197,8 +201,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     handler.receiveMessage(buffer)
     handler.receiveEnd()
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(self.recorder.messages, .isEmpty())
     await assertThat(self.recorder.status, .notNil(.hasCode(.internalError)))
@@ -210,7 +215,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
 
     handler.receiveMessage(ByteBuffer(string: "foo"))
 
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(self.recorder.metadata, .is(.nil()))
     await assertThat(self.recorder.messages, .isEmpty())
@@ -231,8 +238,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
 
     handler.receiveMetadata([:])
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(self.recorder.messages, .isEmpty())
     await assertThat(self.recorder.status, .notNil(.hasCode(.internalError)))
@@ -260,8 +268,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
 
     handler.finish()
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(self.recorder.messages, .isEmpty())
     await assertThat(self.recorder.status, .notNil(.hasCode(.unavailable)))
@@ -279,8 +288,9 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
 
     handler.finish()
 
-    // Wait for user handler to finish.
-    await handler.task?.value
+    // Wait for tasks to finish.
+    await handler.userHandlerTask?.value
+    await handler.responseStreamDrainTask?.value
 
     await assertThat(self.recorder.messages.first, .is(ByteBuffer(string: "hello")))
     await assertThat(self.recorder.status, .notNil(.hasCode(.unavailable)))
@@ -297,7 +307,7 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
     handler.receiveMetadata([:])
 
     // Wait for user handler to finish (it's gonna throw immediately).
-    await assertThat(await handler.task?.value, .notNil())
+    await assertThat(await handler.userHandlerTask?.value, .notNil())
 
     // Check the status is `.unknown`.
     await assertThat(self.recorder.status, .notNil(.hasCode(.unknown)))
@@ -316,19 +326,22 @@ class AsyncServerHandlerTests: ServerHandlerTestCaseBase {
 
     // Send two requests and end, pausing the writer in the middle.
     switch handler.state {
-    case let .active(_, _, responseStreamWriter, _):
+    case let .active(_, _, responseStreamWriter, promise):
       handler.receiveMessage(ByteBuffer(string: "diaz"))
       await responseStreamWriter._asyncWriter.toggleWritability()
       handler.receiveMessage(ByteBuffer(string: "santiago"))
       handler.receiveEnd()
       await responseStreamWriter._asyncWriter.toggleWritability()
-      await handler.task?.value
+      await handler.userHandlerTask?.value
+      _ = try await promise.futureResult.get()
     default:
       XCTFail("Unexpected handler state: \(handler.state)")
     }
 
     handler.finish()
 
+    await assertThat(handler.responseStreamDrainTask, .notNil())
+
     await assertThat(self.recorder.messages, .is([
       ByteBuffer(string: "diaz"),
       ByteBuffer(string: "santiago"),