Wait for writes to drain before closing (#183)

Motivation:

The server can drop pending writes during graceful shutdown. The
connection management handler waits for signals from SwiftNIO to
determine which streams are currently open, and importantly for graceful
shutdown, when all streams are closed. However these notifications
happen when the stream is closed in NIOs HTTP/2 state machine, not when
all data for that stream has been written. This is problematic if there
peer is waiting for a WINDOW_UPDATE or the bytes haven't yet been
flushed.

Modifications:

- When the stream count drops to zero, half-close the connection.
- When half-close complete (and thus all previously comitted data has
been written) full close the connection.

Result:

Fewer failed RPCs

Author: glbrntt

Sources/GRPCNIOTransportCore/Server/Connection/ServerConnectionManagementHandler.swift

@@ -537,22 +537,14 @@ extension ServerConnectionManagementHandler {
     case .startIdleTimer:
       self.maxIdleTimerHandler?.start()
     case .close:
-      // Defer closing until the next tick of the event loop.
+      // Defer until the next tick so the HTTP/2 handler's pending flush (which carries
+      // the END_STREAM frame's bytes) propagates first; then half-close to wait for those
+      // bytes to drain before fully closing.
       //
-      // This point is reached because the server is shutting down gracefully and the stream count
-      // has dropped to zero, meaning the connection is no longer required and can be closed.
-      // However, the stream would've been closed by writing and flushing a frame with end stream
-      // set. These are two distinct events in the channel pipeline. The HTTP/2 handler updates the
-      // state machine when a frame is written, which in this case results in the stream closed
-      // event which we're reacting to here.
-      //
-      // Importantly the HTTP/2 handler hasn't yet seen the flush event, so the bytes of the frame
-      // with end-stream set - and potentially some other frames - are sitting in a buffer in the
-      // HTTP/2 handler. If we close on this event loop tick then those frames will be dropped.
-      // Delaying the close by a loop tick will allow the flush to happen before the close.
-      let loopBound = NIOLoopBound(context, eventLoop: context.eventLoop)
-      context.eventLoop.execute {
-        loopBound.value.close(mode: .all, promise: nil)
+      // Without the half-close barrier any bytes still pending would be dropped on
+      // full close.
+      context.eventLoop.assumeIsolated().execute {
+        Self.gracefullyClose(context: context)
       }
 
     case .none:
@@ -654,7 +646,7 @@ extension ServerConnectionManagementHandler {
       self.maybeFlush(context: context)
 
       if close {
-        context.close(promise: nil)
+        Self.gracefullyClose(context: context)
       } else {
         // RPCs may have a grace period for finishing once the second GOAWAY frame has finished.
         // If this is set close the connection abruptly once the grace period passes.
@@ -666,6 +658,20 @@ extension ServerConnectionManagementHandler {
     }
   }
 
+  /// Closes the connection gracefully.
+  ///
+  /// `close(mode: .all)` would drop any writes still queued in the channel's pending-writes
+  /// buffer or the kernel send buffer. Half-closing the output side first acts as a drain
+  /// barrier: the close-output promise only succeeds once all preceding writes have been
+  /// written, after which the full close is safe.
+  private static func gracefullyClose(context: ChannelHandlerContext) {
+    let promise = context.eventLoop.makePromise(of: Void.self)
+    context.close(mode: .output, promise: promise)
+    promise.futureResult.assumeIsolated().whenComplete { _ in
+      context.close(mode: .all, promise: nil)
+    }
+  }
+
   private func keepaliveTimerFired() {
     guard let context = self.context else { return }
     let ping = HTTP2Frame(streamID: .rootStream, payload: .ping(self.keepalivePingData, ack: false))

Tests/GRPCNIOTransportCoreTests/Server/Connection/ServerConnectionManagementHandlerTests.swift

@@ -140,6 +140,61 @@ struct ServerConnectionManagementHandlerTests {
     try connection.waitUntilClosed()
   }
 
+  @Test("Graceful close half-closes output before full close (no streams)")
+  @available(gRPCSwiftNIOTransport 2.0, *)
+  func gracefulCloseHalfClosesBeforeFullCloseWithNoStreams() throws {
+    let connection = try Connection(maxIdleTime: .minutes(1))
+    let recorder = CloseRecorder()
+    try connection.channel.pipeline.syncOperations.addHandler(recorder, position: .first)
+    try connection.activate()
+
+    connection.advanceTime(by: .minutes(1))
+    try self.testGracefulShutdown(connection: connection, lastStreamID: 0)
+
+    // The ping-ack is the trigger for closing output. The recording handler stores the promise
+    // which will be succeeded in a moment.
+    #expect(recorder.modes == [.output])
+
+    // Complete the half-close promise. Doing that triggers the full close.
+    recorder.completeOutput()
+    #expect(recorder.modes == [.output, .all])
+
+    try connection.waitUntilClosed()
+  }
+
+  @Test("Graceful close half-closes output before full close (last stream closes)")
+  @available(gRPCSwiftNIOTransport 2.0, *)
+  func gracefulCloseHalfClosesBeforeFullCloseOnLastStreamClose() throws {
+    // Path 2: _streamClosed → .close (PING ACK first; stream closes after second GOAWAY).
+    let connection = try Connection(maxIdleTime: .minutes(1))
+    let recorder = CloseRecorder()
+    try connection.channel.pipeline.syncOperations.addHandler(recorder, position: .first)
+    try connection.activate()
+
+    connection.advanceTime(by: .minutes(1))
+    try self.testGracefulShutdown(
+      connection: connection,
+      lastStreamID: 1,
+      streamToOpenBeforePingAck: 1
+    )
+
+    // Stream still open after second GOAWAY: no close yet.
+    #expect(recorder.modes.isEmpty)
+
+    // Closing the last stream schedules a graceful close on the event loop.
+    connection.streamClosed(1)
+    connection.loop.run()
+
+    // The recording handler stores the promise which will be succeeded in a moment.
+    #expect(recorder.modes == [.output])
+
+    // Complete the half-close promise. Doing that triggers the full close.
+    recorder.completeOutput()
+    #expect(recorder.modes == [.output, .all])
+
+    try connection.waitUntilClosed()
+  }
+
   @Test("Keepalive works on new connection")
   @available(gRPCSwiftNIOTransport 2.0, *)
   func keepaliveOnNewConnection() throws {
@@ -480,3 +535,25 @@ extension ServerConnectionManagementHandlerTests {
     }
   }
 }
+
+private final class CloseRecorder: ChannelOutboundHandler {
+  typealias OutboundIn = NIOAny
+
+  private(set) var modes: [CloseMode] = []
+  private var halfClosePromise: EventLoopPromise<Void>?
+
+  func close(context: ChannelHandlerContext, mode: CloseMode, promise: EventLoopPromise<Void>?) {
+    self.modes.append(mode)
+    switch mode {
+    case .output:
+      self.halfClosePromise = promise
+    default:
+      context.close(mode: mode, promise: promise)
+    }
+  }
+
+  func completeOutput() {
+    self.halfClosePromise?.succeed(())
+    self.halfClosePromise = nil
+  }
+}