diff --git a/src/main/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontServingController.java b/src/main/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontServingController.java index 2a1b4e0..76a781f 100644 --- a/src/main/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontServingController.java +++ b/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) { diff --git a/src/test/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontCustomOriginServingTest.java b/src/test/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontCustomOriginServingTest.java index 93cc5bc..94a3ce1 100644 --- a/src/test/java/io/github/hectorvent/floci/services/cloudfront/CloudFrontCustomOriginServingTest.java +++ b/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 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 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 received, AtomicInteger hits) throws IOException { originServer = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);