Implement two-phase graceful shutdown for HTTP/2 server connections, as recommended by RFC 9113 section 6.8.

httpcore5-h2/src/main/java/org/apache/hc/core5/http2/impl/nio/AbstractH2StreamMultiplexer.java

@@ -129,6 +129,12 @@ enum SettingsHandshake { READY, TRANSMITTED, ACKED }
     private final AtomicInteger connOutputWindow;
     private final AtomicInteger outputRequests;
     private final H2StreamListener streamListener;
+    private final boolean serverSide;
+
+    // A server sends an initial GOAWAY before determining its final last-stream-id.
+    private boolean gracefulShutdownAwaitingPing;
+    private boolean gracefulShutdownPingAcknowledged;
+    private long gracefulShutdownPingStartedNanos;
 
     private ConnectionHandshake connState = ConnectionHandshake.READY;
     private SettingsHandshake localSettingState = SettingsHandshake.READY;
@@ -209,6 +215,8 @@ enum SettingsHandshake { READY, TRANSMITTED, ACKED }
         this.frameFactory = Args.notNull(frameFactory, "Frame factory");
         this.httpProcessor = Args.notNull(httpProcessor, "HTTP processor");
         this.streams = new H2Streams(idGenerator);
+        // Server-generated stream IDs are even (RFC 9113, section 5.1.1).
+        this.serverSide = idGenerator.isSameSide(2);
         this.localConfig = h2Config != null ? h2Config : H2Config.DEFAULT;
         this.inputMetrics = new BasicH2TransportMetrics();
         this.outputMetrics = new BasicH2TransportMetrics();
@@ -490,6 +498,9 @@ public final void onConnect() throws HttpException, IOException {
     }
 
     public final void onInput(final ByteBuffer src) throws HttpException, IOException {
+        if (gracefulShutdownDeadlineReached()) {
+            completeGracefulShutdown();
+        }
         if (connState == ConnectionHandshake.SHUTDOWN) {
             ioSession.clearEvent(SelectionKey.OP_READ);
         } else {
@@ -519,6 +530,10 @@ public final void onInput(final ByteBuffer src) throws HttpException, IOExceptio
     }
 
     public final void onOutput() throws HttpException, IOException {
+        if (gracefulShutdownAwaitingPing
+                && (gracefulShutdownPingAcknowledged || gracefulShutdownDeadlineReached())) {
+            completeGracefulShutdown();
+        }
         ioSession.getLock().lock();
         try {
             if (!outputBuffer.isEmpty()) {
@@ -666,6 +681,10 @@ public final void onOutput() throws HttpException, IOException {
     }
 
     public final void onTimeout(final Timeout timeout) throws HttpException, IOException {
+        // This callback represents the application's socket inactivity timeout,
+        // NOT expiry of the optional GOAWAY RTT probe. Preserve its semantics.
+        gracefulShutdownAwaitingPing = false;
+        gracefulShutdownPingAcknowledged = false;
         connState = ConnectionHandshake.SHUTDOWN;
 
         final RawFrame goAway;
@@ -685,6 +704,8 @@ public final void onTimeout(final Timeout timeout) throws HttpException, IOExcep
     }
 
     public final void onDisconnect() {
+        gracefulShutdownAwaitingPing = false;
+        gracefulShutdownPingAcknowledged = false;
         for (;;) {
             final AsyncPingHandler pingHandler = pingHandlers.poll();
             if (pingHandler != null) {
@@ -699,17 +720,48 @@ public final void onDisconnect() {
 
     private void executeShutdown(final ShutdownCommand shutdownCommand) throws IOException {
         if (shutdownCommand.getType() == CloseMode.IMMEDIATE) {
+            gracefulShutdownAwaitingPing = false;
+            gracefulShutdownPingAcknowledged = false;
             streams.shutdownAndReleaseAll();
             connState = ConnectionHandshake.SHUTDOWN;
-        } else {
-            if (connState.compareTo(ConnectionHandshake.ACTIVE) <= 0) {
+        } else if (connState.compareTo(ConnectionHandshake.ACTIVE) <= 0 && !gracefulShutdownAwaitingPing) {
+            if (serverSide) {
+                // RFC 9113, section 6.8: announce shutdown first, then allow one RTT
+                // for requests that were already in flight before fixing last-stream-id.
+                final RawFrame goAway = frameFactory.createGoAway(Integer.MAX_VALUE, H2Error.NO_ERROR, "Graceful shutdown");
+                commitFrame(goAway);
+                gracefulShutdownAwaitingPing = true;
+                // Do not replace the session's configured inactivity timeout.
+                // Probe expiry is checked on subsequent I/O events.
+                gracefulShutdownPingStartedNanos = System.nanoTime();
+                executePing(new PingCommand(new BasicPingHandler(result -> {
+                    if (result && gracefulShutdownAwaitingPing) {
+                        gracefulShutdownPingAcknowledged = true;
+                        requestSessionOutput();
+                    }
+                })));
+            } else {
+                // Client-side shutdown keeps its existing single-GOAWAY behavior.
                 final RawFrame goAway = frameFactory.createGoAway(streams.getLastRemoteId(), H2Error.NO_ERROR, "Graceful shutdown");
                 commitFrame(goAway);
                 connState = streams.isEmpty() ? ConnectionHandshake.SHUTDOWN : ConnectionHandshake.GRACEFUL_SHUTDOWN;
             }
         }
     }
 
+    private boolean gracefulShutdownDeadlineReached() {
+        return gracefulShutdownAwaitingPing && !pingAckTimeout.isDisabled()
+                && System.nanoTime() - gracefulShutdownPingStartedNanos >= pingAckTimeout.toNanoseconds();
+    }
+
+    private void completeGracefulShutdown() throws IOException {
+        gracefulShutdownAwaitingPing = false;
+        gracefulShutdownPingAcknowledged = false;
+        final RawFrame goAway = frameFactory.createGoAway(streams.getLastRemoteId(), H2Error.NO_ERROR, "Graceful shutdown");
+        commitFrame(goAway);
+        connState = streams.isEmpty() ? ConnectionHandshake.SHUTDOWN : ConnectionHandshake.GRACEFUL_SHUTDOWN;
+    }
+
     private void executePing(final PingCommand pingCommand) throws IOException {
         final AsyncPingHandler handler = pingCommand.getHandler();
         pingHandlers.add(handler);

httpcore5-h2/src/test/java/org/apache/hc/core5/http2/impl/nio/TestAbstractH2StreamMultiplexer.java

@@ -1451,6 +1451,191 @@
 
         Assertions.assertDoesNotThrow(() -> streamMultiplexer.onInput(ByteBuffer.wrap(writableChannel.toByteArray())));
     }
+
+    // RFC 9113 section 6.8: the server first announces the maximum stream ID,
+    // then waits for an RTT probe ACK before sending the actual last-stream-id.
+    @Test
+    void testServerGracefulShutdownSendsTwoGoAwayFrames() throws Exception {
+        final List<byte[]> writes = new ArrayList<>();
+        Mockito.when(protocolIOSession.write(ArgumentMatchers.any(ByteBuffer.class)))
+                .thenAnswer(inv -> {
+                    final ByteBuffer buffer = inv.getArgument(0, ByteBuffer.class);
+                    final byte[] copy = new byte[buffer.remaining()];
+                    buffer.get(copy);
+                    writes.add(copy);
+                    return copy.length;
+                });
+        Mockito.when(protocolIOSession.poll()).thenReturn(
+                org.apache.hc.core5.http.nio.command.ShutdownCommand.GRACEFUL, null);
+
+        final H2StreamMultiplexerImpl mux = new H2StreamMultiplexerImpl(
+                protocolIOSession, FRAME_FACTORY, StreamIdGenerator.EVEN,
+                httpProcessor, CharCodingConfig.DEFAULT, H2Config.DEFAULT,
+                h2StreamListener, () -> streamHandler);
+        mux.onConnect();
+        completeSettingsHandshake(mux);
+        writes.clear();
+
+        mux.onOutput();
+        final List<FrameStub> firstFrames = parseFrames(concat(writes));
+        final List<FrameStub> firstGoAway = firstFrames.stream().filter(FrameStub::isGoAway)
+                .collect(java.util.stream.Collectors.toList());
+        Assertions.assertEquals(1, firstGoAway.size());
+        Assertions.assertEquals(Integer.MAX_VALUE,
+                ByteBuffer.wrap(firstGoAway.get(0).payload).getInt() & Integer.MAX_VALUE);
+        final FrameStub ping = firstFrames.stream().filter(f -> f.isPing() && !f.isAck())
+                .findFirst().orElseThrow(AssertionError::new);
+        Assertions.assertTrue(mux.isOpen(), "Server must accept in-flight streams until RTT barrier");
+        Mockito.verify(protocolIOSession, Mockito.never())
+                .setSocketTimeout(ArgumentMatchers.any());
+
+        // Simulate an in-flight peer stream reaching the server after the first GOAWAY.
+        final H2StreamChannel channel = mux.createChannel(1);
+        mux.createStream(channel, streamHandler);
+
+        feedFrame(mux, new RawFrame(FrameType.PING.getValue(), FrameFlag.ACK.getValue(),
+                0, ByteBuffer.wrap(ping.payload)));
+        mux.onOutput();
+        final List<FrameStub> finalGoAway = parseFrames(concat(writes)).stream()
+                .filter(FrameStub::isGoAway).collect(java.util.stream.Collectors.toList());
+        Assertions.assertEquals(2, finalGoAway.size());
+        Assertions.assertEquals(1,
+                ByteBuffer.wrap(finalGoAway.get(1).payload).getInt() & Integer.MAX_VALUE);
+        Assertions.assertFalse(mux.isOpen());
+    }
+
+    @Test
+    void testClientGracefulShutdownStillUsesSingleGoAway() throws Exception {
+        final List<byte[]> writes = new ArrayList<>();
+        Mockito.when(protocolIOSession.write(ArgumentMatchers.any(ByteBuffer.class)))
+                .thenAnswer(inv -> {
+                    final ByteBuffer buffer = inv.getArgument(0, ByteBuffer.class);
+                    final byte[] copy = new byte[buffer.remaining()];
+                    buffer.get(copy);
+                    writes.add(copy);
+                    return copy.length;
+                });
+        Mockito.when(protocolIOSession.poll()).thenReturn(
+                org.apache.hc.core5.http.nio.command.ShutdownCommand.GRACEFUL, null);
+
+        final H2StreamMultiplexerImpl mux = new H2StreamMultiplexerImpl(
+                protocolIOSession, FRAME_FACTORY, StreamIdGenerator.ODD,
+                httpProcessor, CharCodingConfig.DEFAULT, H2Config.DEFAULT,
+                h2StreamListener, () -> streamHandler);
+        mux.onConnect();
+        completeSettingsHandshake(mux);
+        writes.clear();
+
+        mux.onOutput();
+        final List<FrameStub> frames = parseFrames(concat(writes));
+        Assertions.assertEquals(1, frames.stream().filter(FrameStub::isGoAway).count());
+        Assertions.assertFalse(frames.stream().anyMatch(f -> f.isPing() && !f.isAck()));
+    }
+
+    @Test
+    void testServerGracefulShutdownDeadlineOnOutputSendsFinalGoAway() throws Exception {
+        final List<byte[]> writes = new ArrayList<>();
+        Mockito.when(protocolIOSession.write(ArgumentMatchers.any(ByteBuffer.class)))
+                .thenAnswer(inv -> {
+                    final ByteBuffer buffer = inv.getArgument(0, ByteBuffer.class);
+                    final byte[] copy = new byte[buffer.remaining()];
+                    buffer.get(copy);
+                    writes.add(copy);
+                    return copy.length;
+                });
+        Mockito.when(protocolIOSession.poll()).thenReturn(
+                org.apache.hc.core5.http.nio.command.ShutdownCommand.GRACEFUL, null);
+        final H2StreamMultiplexerImpl mux = new H2StreamMultiplexerImpl(
+                protocolIOSession, FRAME_FACTORY, StreamIdGenerator.EVEN,
+                httpProcessor, CharCodingConfig.DEFAULT, H2Config.DEFAULT,
+                h2StreamListener, () -> streamHandler, null, Timeout.ofNanoseconds(1));
+        mux.onConnect();
+        completeSettingsHandshake(mux);
+        writes.clear();
+
+        mux.onOutput();
+        final H2Stream active = mux.createStream(mux.createChannel(1), streamHandler);
+        Assertions.assertTrue(active.isActive());
+        // A subsequent I/O event sees the expired monotonic deadline, without
+        // changing the reactor socket timeout or needing a background timer.
+        mux.onOutput();
+
+        final List<FrameStub> goAways = parseFrames(concat(writes)).stream()
+                .filter(FrameStub::isGoAway)
+                .collect(java.util.stream.Collectors.toList());
+        Assertions.assertEquals(2, goAways.size());
+        Assertions.assertEquals(Integer.MAX_VALUE,
+                ByteBuffer.wrap(goAways.get(0).payload).getInt() & Integer.MAX_VALUE);
+        Assertions.assertEquals(1,
+                ByteBuffer.wrap(goAways.get(1).payload).getInt() & Integer.MAX_VALUE);
+        Assertions.assertTrue(active.isActive(), "Accepted requests must survive probe expiry");
+        Assertions.assertFalse(mux.isOpen());
+        Mockito.verify(protocolIOSession, Mockito.never()).setSocketTimeout(ArgumentMatchers.any());
+    }
+
+    @Test
+    void testServerGracefulShutdownInvalidPingAckDoesNotComplete() throws Exception {
+        final List<byte[]> writes = new ArrayList<>();
+        Mockito.when(protocolIOSession.write(ArgumentMatchers.any(ByteBuffer.class)))
+                .thenAnswer(inv -> {
+                    final ByteBuffer buffer = inv.getArgument(0, ByteBuffer.class);
+                    final byte[] copy = new byte[buffer.remaining()];
+                    buffer.get(copy);
+                    writes.add(copy);
+                    return copy.length;
+                });
+        Mockito.when(protocolIOSession.poll()).thenReturn(
+                org.apache.hc.core5.http.nio.command.ShutdownCommand.GRACEFUL, null);
+
+        final H2StreamMultiplexerImpl mux = new H2StreamMultiplexerImpl(
+                protocolIOSession, FRAME_FACTORY, StreamIdGenerator.EVEN,
+                httpProcessor, CharCodingConfig.DEFAULT, H2Config.DEFAULT,
+                h2StreamListener, () -> streamHandler);
+        mux.onConnect();
+        completeSettingsHandshake(mux);
+        writes.clear();
+
+        mux.onOutput();
+        final H2Stream active = mux.createStream(mux.createChannel(1), streamHandler);
+        feedFrame(mux, new RawFrame(FrameType.PING.getValue(), FrameFlag.ACK.getValue(),
+                0, ByteBuffer.wrap(new byte[8])));
+        mux.onOutput();
+        Assertions.assertEquals(1, parseFrames(concat(writes)).stream()
+                .filter(FrameStub::isGoAway).count(),
+                "Wrong ACK must not finish the GOAWAY exchange before the deadline");
+        Assertions.assertTrue(mux.isOpen());
+        Assertions.assertTrue(active.isActive());
+        Mockito.verify(protocolIOSession, Mockito.never()).setSocketTimeout(ArgumentMatchers.any());
+    }
+
+    @Test
+    void testServerGracefulShutdownPreservesConfiguredSocketTimeout() throws Exception {
+        final List<byte[]> writes = new ArrayList<>();
+        Mockito.when(protocolIOSession.write(ArgumentMatchers.any(ByteBuffer.class)))
+                .thenAnswer(inv -> {
+                    final ByteBuffer buffer = inv.getArgument(0, ByteBuffer.class);
+                    final byte[] copy = new byte[buffer.remaining()];
+                    buffer.get(copy);
+                    writes.add(copy);
+                    return copy.length;
+                });
+        Mockito.when(protocolIOSession.poll()).thenReturn(
+                org.apache.hc.core5.http.nio.command.ShutdownCommand.GRACEFUL, null);
+        final H2StreamMultiplexerImpl mux = new H2StreamMultiplexerImpl(
+                protocolIOSession, FRAME_FACTORY, StreamIdGenerator.EVEN,
+                httpProcessor, CharCodingConfig.DEFAULT, H2Config.DEFAULT,
+                h2StreamListener, () -> streamHandler);
+        mux.onConnect();
+        completeSettingsHandshake(mux);
+        mux.onOutput();
+
+        // An actual reactor inactivity timeout retains its original semantics,
+        // even if a GOAWAY RTT probe was pending.
+        mux.onTimeout(Timeout.ofSeconds(1));
+        Assertions.assertFalse(mux.isOpen());
+        Mockito.verify(protocolIOSession, Mockito.never()).setSocketTimeout(ArgumentMatchers.any());
+    }
+
     private static byte[] encodeFrame(final RawFrame frame) throws IOException {
         final WritableByteChannelMock writableChannel = new WritableByteChannelMock(256);
         final FrameOutputBuffer outBuffer = new FrameOutputBuffer(16 * 1024);