8373877: QUIC connections are removed too early

Reviewed-by: dfuchs

Author: djelinski

src/java.net.http/share/classes/jdk/internal/net/http/quic/ConnectionTerminatorImpl.java

@@ -344,10 +344,6 @@ private void drain(final ConnectionCloseFrame incomingFrame) {
             }
         }
         failHandshakeCFs();
-        // remap the connection to a draining connection
-        final QuicEndpoint endpoint = this.connection.endpoint();
-        assert endpoint != null : "QUIC endpoint is null";
-        endpoint.draining(connection);
         discardConnectionState();
         connection.streams.terminate(terminationCause);
         if (Log.quic()) {
@@ -439,7 +435,7 @@ private void pushConnectionCloseFrame(final KeySpace keySpace,
         final ProtectionRecord protectionRecord = ProtectionRecord.single(packet,
                 connection::allocateDatagramForEncryption);
         // while sending the packet containing the CONNECTION_CLOSE frame, the pushDatagram will
-        // remap (or remove) the QuicConnectionImpl in QuicEndpoint.
+        // remap the QuicConnectionImpl in QuicEndpoint.
         connection.pushDatagram(protectionRecord);
     }
 

src/java.net.http/share/classes/jdk/internal/net/http/quic/QuicConnectionImpl.java

@@ -2811,7 +2811,7 @@ private void pushEncryptedDatagram(final ProtectionRecord protectionRecord) {
                 // a CONNECTION_CLOSE frame is being sent to the peer when the local
                 // connection state is in DRAINING. This implies that the local endpoint
                 // is responding to an incoming CONNECTION_CLOSE frame from the peer.
-                // we remove the connection from the endpoint for such cases.
+                // we switch this connection to one that does not respond to incoming packets.
                 endpoint.pushClosedDatagram(this, peerAddress(), datagram);
             } else if (stateHandle.isMarked(QuicConnectionState.CLOSING)) {
                 // a CONNECTION_CLOSE frame is being sent to the peer when the local

src/java.net.http/share/classes/jdk/internal/net/http/quic/QuicEndpoint.java

@@ -1509,7 +1509,8 @@ public void pushClosingDatagram(QuicConnectionImpl connection, InetSocketAddress
     /**
      * Called to schedule sending of a datagram that contains a single {@code ConnectionCloseFrame}
      * sent in response to a {@code ConnectionClose} frame.
-     * This will completely remove the connection from the connection map.
+     * This will replace the {@link QuicConnectionImpl} with a {@link DrainingConnection} that
+     * will discard all incoming packets.
      * @param connection   the connection being closed
      * @param destination  the peer address
      * @param datagram     the datagram
@@ -1518,7 +1519,7 @@ public void pushClosedDatagram(QuicConnectionImpl connection,
                                    InetSocketAddress destination,
                                    ByteBuffer datagram) {
         if (debug.on()) debug.log("Pushing closed datagram for " + connection.logTag());
-        removeConnection(connection);
+        draining(connection);
         pushDatagram(connection, destination, datagram);
     }
 

test/jdk/java/net/httpclient/quic/StatelessResetReceiptTest.java

@@ -47,6 +47,7 @@
 import jdk.internal.net.http.quic.QuicClient;
 import jdk.internal.net.http.quic.QuicConnectionId;
 import jdk.internal.net.http.quic.QuicConnectionImpl;
+import jdk.internal.net.http.quic.QuicTransportParameters;
 import jdk.internal.net.http.quic.TerminationCause;
 import jdk.internal.net.quic.QuicTLSContext;
 import jdk.internal.net.quic.QuicVersion;
@@ -58,11 +59,13 @@
 import static jdk.internal.net.quic.QuicTransportErrors.NO_VIABLE_PATH;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /*
  * @test
+ * @bug 8373877
  * @summary verify that when a QUIC (client) connection receives a stateless reset
  *          from the peer, then the connection is properly terminated
  * @library /test/lib /test/jdk/java/net/httpclient/lib
@@ -87,6 +90,10 @@ static void beforeAll() throws Exception {
                 .availableVersions(new QuicVersion[]{QuicVersion.QUIC_V1})
                 .sslContext(sslContext)
                 .build();
+        // Use a longer max ack delay to inflate the draining time (3xPTO)
+        final QuicTransportParameters transportParameters = new QuicTransportParameters();
+        transportParameters.setIntParameter(QuicTransportParameters.ParameterId.max_ack_delay,
+                (1 << 14) - 1); // 16 seconds, maximum allowed
         server.start();
         System.out.println("Server started at " + server.getAddress());
     }
@@ -179,6 +186,9 @@ public void testClosingConnection() throws Exception {
             conn.close();
             // verify connection is no longer open
             assertFalse(conn.underlyingQuicConnection().isOpen(), "QUIC connection is still open");
+            // Check that the (closing) connection is still registered with the endpoint
+            assertNotEquals(0, conn.endpoint().connectionCount(),
+                    "Expected the QUIC connection to be registered");
             // now send a stateless reset from the server connection
             sendStatelessResetFrom(serverConn);
             // wait for the stateless reset to be processed
@@ -231,12 +241,19 @@ public void testDrainingConnection() throws Exception {
                     .loggedAs("intentionally closed by server to initiate draining state" +
                             " on client connection");
             serverConn.connectionTerminator().terminate(tc);
+            // send ping in case the connection_close message gets lost
+            conn.underlyingQuicConnection().requestSendPing();
             // wait for client conn to terminate
             final TerminationCause clientTC = ((QuicConnectionImpl) conn.underlyingQuicConnection())
                     .futureTerminationCause().get();
             // verify connection closed for the right reason
             assertEquals(NO_VIABLE_PATH.code(), clientTC.getCloseCode(),
                     "unexpected termination cause");
+            // wait for a while to let the connection close completely
+            Thread.sleep(10);
+            // Check that the (draining) connection is still registered with the endpoint
+            assertNotEquals(0, conn.endpoint().connectionCount(),
+                    "Expected the QUIC connection to be registered");
             // now send a stateless reset from the server connection
             sendStatelessResetFrom(serverConn);
             // wait for the stateless reset to be processed