fix(cloudfront): release origin connections when viewers disconnect (#5564)

(cherry picked from commit ee229af97569d5876f756f65ce7956726e2ed696)

src/main/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontServingController.java

@@ -32,6 +32,7 @@ import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.StreamingOutput;
 import jakarta.ws.rs.core.UriInfo;
 import org.jboss.logging.Logger;
+import org.jboss.resteasy.reactive.server.core.CurrentRequestManager;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -954,6 +955,9 @@ public class CloudFrontServingController {
                 }
             }
             if (includeBody && origin.stream() != null) {
+                // A disconnected viewer can skip StreamingOutput entirely, so its finally block
+                // alone cannot release the origin connection.
+                CurrentRequestManager.get().registerCompletionCallback(ignored -> origin.closeStream());
                 InputStream body = origin.stream();
                 StreamingOutput entity = output -> {
                     try (InputStream in = body) {

src/test/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontCustomOriginServingTest.java

@@ -39,7 +39,11 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -49,6 +53,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 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.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -390,6 +395,45 @@ class CloudFrontCustomOriginServingTest {
         assertArrayEquals(upload, received.get().body());
     }
 
+    @Test
+    void releasesOriginConnectionsWhenViewersDisconnectBeforeOriginHeaders() throws Exception {
+        BlockingQueue<CountDownLatch> pendingResponses = new LinkedBlockingQueue<>();
+        originServer = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+        originServer.createContext("/", exchange -> {
+            CountDownLatch sendResponse = new CountDownLatch(1);
+            pendingResponses.add(sendResponse);
+            try {
+                if (!sendResponse.await(10, TimeUnit.SECONDS)) {
+                    throw new IOException("Timed out waiting for viewer disconnect");
+                }
+                byte[] body = "origin-ok".getBytes(StandardCharsets.UTF_8);
+                exchange.sendResponseHeaders(200, body.length);
+                exchange.getResponseBody().write(body);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new IOException("Origin response interrupted", e);
+            } finally {
+                exchange.close();
+            }
+        });
+        originServer.createContext("/after-disconnects", exchange -> {
+            byte[] body = "origin-ok".getBytes(StandardCharsets.UTF_8);
+            exchange.sendResponseHeaders(200, body.length);
+            exchange.getResponseBody().write(body);
+            exchange.close();
+        });
+        originServer.start();
+        Distribution created = cloudFrontService.createDistribution(
+                customOriginDistribution(allMethodsBehavior()), Map.of());
+
+        for (int i = 0; i < CloudFrontOriginHttpClient.MAX_CONNECTIONS + 10; i++) {
+            disconnectViewerBeforeOriginHeaders(created.getDomainName(), i, pendingResponses);
+        }
+
+        given().header("Host", created.getDomainName()).when().get("/after-disconnects")
+                .then().statusCode(200).body(equalTo("origin-ok"));
+    }
+
     @Test
     void releasesTheOriginConnectionOfEveryResponseItDoesNotSend() throws Exception {
         // The origin client keeps a fixed number of connections. An error response replaced by a custom
@@ -789,6 +833,23 @@ class CloudFrontCustomOriginServingTest {
         assertEquals("127.0.0.1:" + originServer.getAddress().getPort(), request.header("Host"));
     }
 
+    private static void disconnectViewerBeforeOriginHeaders(String host, int requestIndex,
+                                                            BlockingQueue<CountDownLatch> pendingResponses)
+            throws Exception {
+        CountDownLatch sendResponse = null;
+        try (Socket viewer = new Socket("127.0.0.1", RestAssured.port)) {
+            viewer.setSoLinger(true, 0);
+            viewer.getOutputStream().write(("GET /disconnect/" + requestIndex + " HTTP/1.1\r\nHost: "
+                    + host + "\r\n\r\n").getBytes(StandardCharsets.US_ASCII));
+            sendResponse = pendingResponses.poll(5, TimeUnit.SECONDS);
+            assertNotNull(sendResponse, "No origin connection for viewer " + requestIndex);
+        } finally {
+            if (sendResponse != null) {
+                sendResponse.countDown();
+            }
+        }
+    }
+
     private void startRecordingOrigin(AtomicReference<ReceivedRequest> received, AtomicInteger hits)
             throws IOException {
         originServer = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);