diff --git a/README.md b/README.md index 2b42bea00b..83a955ec86 100644 --- a/README.md +++ b/README.md @@ -253,16 +253,82 @@ does not re-enter routing or write shared timeout fields. The streaming ExternalServiceWorkerDispatchTest forces read-listener continuation and verifies that an unrelated request progresses while the upstream is blocked; it also checks body preservation, one rewrite, direct GETs, and unchanged unmatched routing. -The worker continuation currently constructs a new default ExceptionHandler per -dispatched request. It returns the default structured error response (ERR10011 / -HTTP 400 for an uncaught IOException) after the original exception-handler frame -has returned. This does not reuse the chain's configured handler or its enabled -flag: the dispatched path applies default exception handling even when the chain -intentionally disables it. Other worker-entry paths still use their surrounding -chain policy. This is an accepted limitation of this fix, tracked in -[issue #2820](https://github.com/networknt/light-4j/issues/2820) for separate -exception-module continuation support and review. The current implementation -retains ingress-proxy's compile dependency on the exception module. +HandlerContinuation in the handler module captures the active chain exception +boundary before asynchronous body reads and matched-request worker dispatch. +ExceptionHandler supplies its existing policy without rebuilding middleware or +changing its next handler. Named handlers are retained; when exception handling +is disabled and omitted from the chain, failures remain under Undertow's default +policy. Nested boundaries restore their enclosing scope, and state is isolated +per exchange. Capturing a boundary does not snapshot thread-local MDC. At resume, +one scope reads the directly captured request holder (or an empty context when +absent), then restores the executing thread's previous MDC on every exit. Nested +frames reuse that scope. Policy identity and exchange identity share one frame +tracker; proxy callbacks do not maintain a second thread-local guard. Framework +MDC updates use one stable RequestMdcContext holder with an atomic immutable-map +reference. The first boundary execution or capture initializes that holder +before its asynchronous hand-off; addHandlerMDCContext can also initialize it on +first serial use. Later updates replace only the holder's snapshot, never its +attachment. RequestState stores only MDC in a concurrent weak-identity index; it never stores +an exception policy. Entries are removed through bounded ReferenceQueue cleanup +after their exchange is collected; initialization never registers a completion listener. +Captured scopes read the holder directly, avoiding index +lookups on callback resumes. An explicitly retained HandlerContinuation.Scope carries +the lexical policy across hand-offs; a new capture outside its frame has no policy. +This prevents an earlier sub-pipeline boundary from shadowing a later one. Native +Undertow HandlerWrapper policies are applied at execution, not capture time. +Custom asynchronous code must initialize/capture before hand-off. +The Map-typed attachment key remains compatible with read-only consumers; direct +attachment replacement or mutation is unsupported after initialization. +Proxy callbacks use this scope rather than the unscoped +LightThreadExecutor. ExceptionHandler retains its original response-channel +availability policy after commitment; it does not forcibly close general-service +responses. Its initial worker frame still clears request MDC at chain exit. +BodyHandler, DumpHandler and McpHandler redispatch, HMAC replay/admin completion, +and ProxyHandler acquisition, retries, timeout cancellation and client send/response +callbacks carry the selected boundary. The body-read listener itself uses its +pre-suspension scope and scopes interceptor +execution as well as downstream handling, and preserves pending-read lifecycle. +Proxy actions reuse one captured policy; internal synchronous calls do not add +boundary frames or reset MDC again. Unexpected proxy failures cancel only the +upstream request's channels on multiplexed connections. Before response commitment +an unhandled failure returns 500; after commitment it closes the failed response +stream without appending an error. Undertow's HTTP/2 ServerConnection.close() +sends RST_STREAM for that stream, preserving healthy siblings. Late failures on +completed exchanges do not rewrite a response or close a reused connection. +The neutral Runnable adapter attempts a 500 before commitment or a stream reset +after commitment before propagating the original Error. Failed writes/resets are +attached as suppressed failures and logged without replacing the original Error; +selected exception policies keep +their existing handling behavior. Ordinary exceptions use the same guarded cleanup. +Proxy callbacks execute on the downstream exchange's I/O thread, without a +callback-wide monitor or shared mutable callback slot. A per-request FIFO queues +both foreign-thread and owner-thread arrivals, preserving submission order. The queue +is private to proxy-handler. Fatal cleanup runs on its I/O owner before +the queue stops. Executor rejection closes an immutable resource snapshot +published through a volatile field +and reports the failure through the selected policy once; it never reads or mutates +the owner's request/connection/attempt/stopping fields. Later stopped submissions +only close their own incoming resources. Accepted work that fails replays only its +failure through that policy. The public `CONNECTION` +attachment is populated when a connection is accepted. Retry generations discard +late acquisition/send/response callbacks and release the previous attempt's resources +before accepting a new connection. A proxy-generated terminal response is marked +before sending, so benign late callbacks leave it draining. Pushed exchanges inherit +the selected policy and an independent MDC snapshot, with their own cancellation +state. Unexpected committed-response callbacks reset only the affected HTTP/2 stream. +Local allocation probes are separate qualification artifacts, not repository source +or CI qualification; gateway performance remains an explicit review gate. +The review-11 ThreadMXBean probe includes request setup and four callbacks: both +branches measured 1.25-1.31 KB with empty MDC and 2.97-3.40 KB with three fields +(20k warm-up, three 20k samples). A separate real HTTP/1.1 loopback probe includes +ExceptionHandler, ProxyHandler, upstream transport and the Java HTTP client; its +129-133 KB/request totals are not attributable solely to continuation overhead. +Neither probe is a production throughput qualification. +Custom asynchronous handlers must explicitly capture before registering their +callbacks; this contract cannot intercept arbitrary third-party hand-offs. +Body-reading requests capture lazily only when a read suspends. ingress-proxy uses +exception only for tests. Interrupted upstream sends restore the interrupt flag +and stop retrying immediately. Regressions compare worker-entry and streamed-continuation failures, cover 405 rejection and metrics injection, and check that concurrent initialization shares one HTTP client. diff --git a/body/src/main/java/com/networknt/body/BodyHandler.java b/body/src/main/java/com/networknt/body/BodyHandler.java index 2246c91e21..77dc9ec2d8 100644 --- a/body/src/main/java/com/networknt/body/BodyHandler.java +++ b/body/src/main/java/com/networknt/body/BodyHandler.java @@ -19,6 +19,7 @@ import tools.jackson.core.type.TypeReference; import com.networknt.config.Config; import com.networknt.handler.Handler; +import com.networknt.handler.HandlerContinuation; import com.networknt.handler.MiddlewareHandler; import com.networknt.httpstring.AttachmentConstants; import com.networknt.utility.StringUtils; @@ -105,7 +106,7 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception { String contentType = exchange.getRequestHeaders().getFirst(Headers.CONTENT_TYPE); if (contentType != null) { if (exchange.isInIoThread()) { - exchange.dispatch(this); + exchange.dispatch(HandlerContinuation.capture(exchange, this)); return; } exchange.startBlocking(); diff --git a/body/src/test/java/com/networknt/body/BodyContinuationTest.java b/body/src/test/java/com/networknt/body/BodyContinuationTest.java new file mode 100644 index 0000000000..ba2c1f4d0c --- /dev/null +++ b/body/src/test/java/com/networknt/body/BodyContinuationTest.java @@ -0,0 +1,38 @@ +package com.networknt.body; + +import com.networknt.handler.HandlerContinuation; +import io.undertow.Undertow; +import io.undertow.server.HttpHandler; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.*; +import java.time.Duration; +import static org.junit.jupiter.api.Assertions.*; + +class BodyContinuationTest { + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void redispatchRetainsBoundary(boolean enabled) throws Exception { + BodyHandler body = new BodyHandler(); + body.setNext(exchange -> { + assertFalse(exchange.isInIoThread()); + throw new IllegalStateException("downstream failure"); + }); + HttpHandler root = enabled ? exchange -> HandlerContinuation.execute(exchange, task -> ex -> { + try { task.handleRequest(ex); } + catch (IllegalStateException _) { ex.setStatusCode(409).getResponseSender().send("selected"); } + }, body) : body; + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(root).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/")) + .timeout(Duration.ofSeconds(5)).header("Content-Type", "text/plain") + .POST(HttpRequest.BodyPublishers.ofString("body")).build(), HttpResponse.BodyHandlers.ofString()); + assertEquals(enabled ? 409 : 500, response.statusCode(), response.body()); + assertEquals(enabled, response.body().contains("selected")); + } finally { server.stop(); } + } +} diff --git a/dump/src/main/java/com/networknt/dump/DumpHandler.java b/dump/src/main/java/com/networknt/dump/DumpHandler.java index 8a1d696ce5..f63cd932c4 100644 --- a/dump/src/main/java/com/networknt/dump/DumpHandler.java +++ b/dump/src/main/java/com/networknt/dump/DumpHandler.java @@ -17,6 +17,7 @@ package com.networknt.dump; import com.networknt.handler.Handler; +import com.networknt.handler.HandlerContinuation; import com.networknt.handler.MiddlewareHandler; import io.undertow.Handlers; import io.undertow.server.HttpHandler; @@ -77,7 +78,7 @@ public boolean isEnabled() { @Override public void handleRequest(final HttpServerExchange exchange) throws Exception { if (exchange.isInIoThread()) { - exchange.dispatch(this); + exchange.dispatch(HandlerContinuation.capture(exchange, this)); return; } if(isEnabled()) { diff --git a/dump/src/test/java/com/networknt/dump/DumpContinuationTest.java b/dump/src/test/java/com/networknt/dump/DumpContinuationTest.java new file mode 100644 index 0000000000..74e6c0e53e --- /dev/null +++ b/dump/src/test/java/com/networknt/dump/DumpContinuationTest.java @@ -0,0 +1,39 @@ +package com.networknt.dump; + +import com.networknt.handler.HandlerContinuation; +import io.undertow.Undertow; +import io.undertow.server.*; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.*; +import java.time.Duration; +import static org.junit.jupiter.api.Assertions.*; + +class DumpContinuationTest { + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void redispatchRetainsSelectedBoundary(boolean enabled) throws Exception { + DumpHandler handler = new DumpHandler() { + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { + if (!exchange.isInIoThread()) throw new IllegalStateException("resumed failure"); + super.handleRequest(exchange); + } + }; + HttpHandler root = enabled ? exchange -> HandlerContinuation.execute(exchange, task -> ex -> { + try { task.handleRequest(ex); } + catch (IllegalStateException _) { ex.setStatusCode(409).getResponseSender().send("selected"); } + }, handler) : handler; + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(root).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/mcp")) + .timeout(Duration.ofSeconds(5)).POST(HttpRequest.BodyPublishers.ofString("{}")) + .build(), HttpResponse.BodyHandlers.ofString()); + assertEquals(enabled ? 409 : 500, response.statusCode(), response.body()); + assertEquals(enabled, response.body().contains("selected")); + } finally { server.stop(); } + } +} diff --git a/exception/src/main/java/com/networknt/exception/ExceptionHandler.java b/exception/src/main/java/com/networknt/exception/ExceptionHandler.java index 955e3791d1..2351104b3b 100644 --- a/exception/src/main/java/com/networknt/exception/ExceptionHandler.java +++ b/exception/src/main/java/com/networknt/exception/ExceptionHandler.java @@ -17,7 +17,9 @@ package com.networknt.exception; import com.networknt.handler.Handler; +import com.networknt.handler.HandlerContinuation; import com.networknt.handler.MiddlewareHandler; +import com.networknt.status.Status; import io.undertow.Handlers; import io.undertow.server.HttpHandler; import io.undertow.server.HttpServerExchange; @@ -79,58 +81,52 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception { return; } - // These two lines of code should be placed here to avoid the log being printed twice - // when switching from the IO thread to the worker thread. + try { + HandlerContinuation.execute(exchange, continuation -> ex -> handleContinuation(ex, continuation), + ex -> Handler.next(ex, next)); + } finally { + // Initial chain execution owns request-thread cleanup. Resumed frames + // scope and restore the executing thread's MDC themselves. + MDC.clear(); + } + } + + private void handleContinuation(HttpServerExchange exchange, HttpHandler continuation) throws Exception { + // Log each executing frame, including asynchronous continuations. String path = exchange.getRequestPath(); - if(logger.isDebugEnabled()) logger.debug("ExceptionHandler starting - path: {}", path); + if(logger.isDebugEnabled()) logger.debug("ExceptionHandler frame starting - path: {}", path); try { - Handler.next(exchange, next); + continuation.handleRequest(exchange); } catch (Throwable e) { logger.error("ExceptionHandler error processing - path: {}", path, e); - if(exchange.isResponseChannelAvailable()) { - //handle exceptions - if(e instanceof RuntimeException) { - // check if it is FrameworkException which is subclass of RuntimeException. - if(e instanceof FrameworkException) { - FrameworkException fe = (FrameworkException)e; - exchange.setStatusCode(fe.getStatus().getStatusCode()); - exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json"); - exchange.getResponseSender().send(fe.getStatus().toString()); - logger.error(fe.getStatus().toString(), e); - } else { - setExchangeStatus(exchange, STATUS_RUNTIME_EXCEPTION); - } - } else { - if(e instanceof ApiException) { - ApiException ae = (ApiException)e; - exchange.setStatusCode(ae.getStatus().getStatusCode()); - exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json"); - exchange.getResponseSender().send(ae.getStatus().toString()); - logger.error(ae.getStatus().toString(), e); - } else if(e instanceof ClientException){ - ClientException ce = (ClientException)e; - if(ce.getStatus().getStatusCode() == 0){ - setExchangeStatus(exchange, STATUS_UNCAUGHT_EXCEPTION); - } else { - exchange.setStatusCode(ce.getStatus().getStatusCode()); - exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json"); - exchange.getResponseSender().send(ce.getStatus().toString()); - } - - } else { - setExchangeStatus(exchange, STATUS_UNCAUGHT_EXCEPTION); - } - } - } + if (exchange.isResponseChannelAvailable()) writeFailure(exchange, e); } finally { - if(logger.isDebugEnabled()) logger.debug("ExceptionHandler completed - path: {}", path); - // at last, clean the MDC. Most likely, correlationId in side. - //logger.debug("Clear MDC"); - MDC.clear(); + if(logger.isDebugEnabled()) logger.debug("ExceptionHandler frame completed - path: {}", path); } } + private void writeFailure(HttpServerExchange exchange, Throwable failure) { + if (failure instanceof FrameworkException framework) { + sendTypedStatus(exchange, framework.getStatus(), failure, true); + } else if (failure instanceof RuntimeException) { + setExchangeStatus(exchange, STATUS_RUNTIME_EXCEPTION); + } else if (failure instanceof ApiException api) { + sendTypedStatus(exchange, api.getStatus(), failure, true); + } else if (failure instanceof ClientException) { + Status status = ClientException.getStatus(); + if (status.getStatusCode() == 0) setExchangeStatus(exchange, STATUS_UNCAUGHT_EXCEPTION); + else sendTypedStatus(exchange, status, failure, false); + } else setExchangeStatus(exchange, STATUS_UNCAUGHT_EXCEPTION); + } + + private static void sendTypedStatus(HttpServerExchange exchange, Status status, Throwable failure, boolean log) { + exchange.setStatusCode(status.getStatusCode()); + exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json"); + exchange.getResponseSender().send(status.toString()); + if (log && logger.isErrorEnabled()) logger.error(status.toString(), failure); + } + @Override public HttpHandler getNext() { return next; diff --git a/exception/src/test/java/com/networknt/exception/ExceptionContinuationTest.java b/exception/src/test/java/com/networknt/exception/ExceptionContinuationTest.java new file mode 100644 index 0000000000..57d7c72df3 --- /dev/null +++ b/exception/src/test/java/com/networknt/exception/ExceptionContinuationTest.java @@ -0,0 +1,302 @@ +package com.networknt.exception; + +import com.networknt.config.Config; +import ch.qos.logback.classic.Logger; +import ch.qos.logback.classic.spi.ILoggingEvent; +import ch.qos.logback.core.AppenderBase; +import com.networknt.handler.HandlerContinuation; +import com.networknt.status.Status; +import io.undertow.Undertow; +import io.undertow.server.HttpHandler; +import io.undertow.server.HttpServerExchange; +import io.undertow.util.HttpString; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.*; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.*; + +@Timeout(20) +class ExceptionContinuationTest { + static Stream cases() { + return Stream.of("exception", "continuation-named", "continuation-disabled") + .flatMap(name -> Stream.of(false, true).flatMap(async -> + Stream.of(false, true).flatMap(runtime -> Stream.of(false, true) + .map(mdc -> Arguments.of(name, async, runtime, mdc))))); + } + + @ParameterizedTest + @MethodSource("cases") + void selectedPolicyHandlesInitialAndResumedFailures(String name, boolean async, boolean runtime, boolean hasMdc) throws Exception { + boolean enabled = !name.endsWith("disabled"); + Map previous = Config.getInstance().getJsonMapConfig("exception"); + Config.getInstance().putInConfigCache(name, Map.of("enabled", enabled)); + ExceptionHandler selected = new ExceptionHandler(name) { + @Override public void setExchangeStatus(HttpServerExchange exchange, String code, Object... args) { + exchange.getResponseHeaders().put(new HttpString("X-Selected"), name); + super.setExchangeStatus(exchange, code, args); + } + }; + HttpHandler failing = exchange -> { + if (async) { + assertEquals(hasMdc ? "canonical" : null, MDC.get("request")); + assertNull(MDC.get("executor")); + } + if (runtime) throw new IllegalStateException("runtime failure"); + throw new IOException("checked failure"); + }; + var resumedMdc = new AtomicReference(); + var resumedDone = new CountDownLatch(1); + HttpHandler work = exchange -> { + if (hasMdc) selected.addHandlerMDCContext(exchange, "request", "canonical"); + if (async) { + HttpHandler resumed = HandlerContinuation.capture(exchange, failing); + exchange.dispatch(() -> { + MDC.put("executor", "owns-context"); + try { + HandlerContinuation.task(exchange, resumed).run(); + resumedMdc.set(MDC.get("executor")); + } finally { + MDC.clear(); + resumedDone.countDown(); + } + }); + } else failing.handleRequest(exchange); + }; + selected.setNext(work); + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1") + .setHandler(selected.isEnabled() ? selected : work).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + HttpResponse response = client.send(HttpRequest.newBuilder( + URI.create("http://127.0.0.1:" + port + "/failure")).timeout(Duration.ofSeconds(5)).build(), + HttpResponse.BodyHandlers.ofString()); + assertEquals(enabled ? (runtime ? 500 : 400) : 500, response.statusCode(), response.body()); + assertEquals(enabled ? name : "", response.headers().firstValue("X-Selected").orElse("")); + assertEquals(enabled, response.body().contains(runtime ? "ERR10010" : "ERR10011"), response.body()); + if (async) { + assertTrue(resumedDone.await(3, TimeUnit.SECONDS)); + assertEquals("owns-context", resumedMdc.get(), "Resumed exception handling must leave executor MDC alone"); + } + } finally { + server.stop(); + if (name.equals("exception") && previous != null) Config.getInstance().putInConfigCache(name, previous); + else Config.getInstance().clearConfigCache(name); + } + } + @ParameterizedTest + @ValueSource(strings = {"rejected", "accepted-io", "accepted-error"}) + void dispatcherFailuresRetainFrameworkPolicy(String mode) throws Exception { + String name = "continuation-dispatch-failure"; + Config.getInstance().putInConfigCache(name, Map.of("enabled", true)); + var calls = new AtomicInteger(); + ExceptionHandler selected = new ExceptionHandler(name) { + @Override public void setExchangeStatus(HttpServerExchange exchange, String code, Object... args) { + exchange.getResponseHeaders().put(new HttpString("X-Selected"), name); + super.setExchangeStatus(exchange, code, args); + } + }; + selected.setNext(exchange -> { + var scope = HandlerContinuation.capture(exchange); + exchange.dispatch(() -> { + scope.resume(ex -> { + if (mode.equals("rejected")) throw new RejectedExecutionException("executor rejected transport"); + calls.incrementAndGet(); + if (mode.equals("accepted-error")) throw new StackOverflowError("injected fatal sentinel"); + throw new IOException("accepted I/O failure"); + }); + }); + }); + var server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(selected).build(); + try (var client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/dispatch")) + .timeout(Duration.ofSeconds(5)).build(), HttpResponse.BodyHandlers.ofString()); + boolean rejected = mode.equals("rejected"); + assertEquals(rejected ? 500 : 400, response.statusCode(), response.body()); + assertTrue(response.body().contains(rejected ? "ERR10010" : "ERR10011"), response.body()); + assertEquals(name, response.headers().firstValue("X-Selected").orElseThrow()); + assertEquals(rejected ? 0 : 1, calls.get()); + } finally { + server.stop(); + Config.getInstance().clearConfigCache(name); + } + } + + static Stream typedFailureCases() { + return Stream.of("framework", "api", "client", "client-without-status") + .flatMap(kind -> Stream.of(false, true).map(resumed -> Arguments.of(kind, resumed))); + } + + @ParameterizedTest + @MethodSource("typedFailureCases") + void typedFailuresRetainResponseStatusAndBody(String kind, boolean resumed) throws Exception { + Status previousClientStatus = ClientException.getStatus(); + Status status = new Status(429, "TEST_TYPED_FAILURE", "Typed failure", "Preserve the supplied status"); + Exception failure = switch (kind) { + case "framework" -> new FrameworkException(status); + case "api" -> new ApiException(status); + case "client" -> new ClientException(status); + case "client-without-status" -> new ClientException(new Status()); + default -> throw new IllegalArgumentException(kind); + }; + var logged = new CompletableFuture(); + String expectedMessage = status.toString(); + var events = new AppenderBase() { + @Override protected void append(ILoggingEvent event) { + if (expectedMessage.equals(event.getMessage()) && event.getThrowableProxy() != null + && failure.getClass().getName().equals(event.getThrowableProxy().getClassName())) { + event.prepareForDeferredProcessing(); + logged.complete(event); + } + } + }; + var log = (Logger) ExceptionHandler.logger; + events.start(); + log.addAppender(events); + ExceptionHandler selected = new ExceptionHandler(); + HttpHandler failing = exchange -> { throw failure; }; + selected.setNext(exchange -> { + if (resumed) exchange.dispatch(HandlerContinuation.capture(exchange, failing)); + else failing.handleRequest(exchange); + }); + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(selected).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port)) + .timeout(Duration.ofSeconds(5)).build(), HttpResponse.BodyHandlers.ofString()); + boolean missingStatus = kind.equals("client-without-status"); + assertEquals(missingStatus ? 400 : 429, response.statusCode(), response.body()); + assertTrue(response.body().contains(missingStatus ? "ERR10011" : "TEST_TYPED_FAILURE"), response.body()); + assertEquals("application/json", response.headers().firstValue("Content-Type").orElse("")); + if (kind.equals("framework") || kind.equals("api")) + assertEquals(expectedMessage, logged.get(5, TimeUnit.SECONDS).getMessage(), + "Await the server's typed log event, independently of response delivery"); + } finally { + server.stop(); + ClientException.setStatus(previousClientStatus); + log.detachAppender(events); + events.stop(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void committedResponseIsNotTruncatedByInitialOrResumedBoundary(boolean resumed) throws Exception { + var release = new CountDownLatch(1); + ExceptionHandler selected = new ExceptionHandler(); + HttpHandler streaming = exchange -> { + exchange.startBlocking(); + var output = exchange.getOutputStream(); + output.write("first".getBytes(StandardCharsets.UTF_8)); + output.flush(); + exchange.dispatch(() -> { + try { + assertTrue(release.await(5, TimeUnit.SECONDS)); + output.write("second".getBytes(StandardCharsets.UTF_8)); + output.close(); + } catch (Exception e) { throw new RuntimeException(e); } + }); + throw new IOException("failure after sending response headers"); + }; + selected.setNext(exchange -> { + if (resumed) exchange.dispatch(HandlerContinuation.capture(exchange, streaming)); + else streaming.handleRequest(exchange); + }); + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(selected).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port)) + .timeout(Duration.ofSeconds(8)).build(), HttpResponse.BodyHandlers.ofInputStream()); + assertEquals(200, response.statusCode()); + try (var body = response.body()) { + assertEquals("first", new String(body.readNBytes(5), StandardCharsets.UTF_8)); + release.countDown(); + assertEquals("second", new String(body.readAllBytes(), StandardCharsets.UTF_8)); + } + } finally { + release.countDown(); + server.stop(); + } + } + + @Test + void lateResumedFailureDoesNotRewriteCompletedResponseOrCloseKeepAliveConnection() throws Exception { + var saved = new AtomicReference(); + var completed = new AtomicReference(); + var done = new CountDownLatch(1); + ExceptionHandler selected = new ExceptionHandler(); + selected.setNext(exchange -> { + saved.set(HandlerContinuation.capture(exchange, ex -> { throw new IOException("late failure"); })); + exchange.addExchangeCompleteListener((ex, next) -> { + completed.set(ex); + done.countDown(); + next.proceed(); + }); + exchange.getResponseSender().send("accepted"); + }); + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(selected).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var request = HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port)).timeout(Duration.ofSeconds(5)).build(); + assertEquals("accepted", client.send(request, HttpResponse.BodyHandlers.ofString()).body()); + assertTrue(done.await(3, TimeUnit.SECONDS)); + assertTrue(completed.get().isComplete()); + saved.get().handleRequest(completed.get()); + assertTrue(completed.get().getConnection().isOpen()); + assertEquals("accepted", client.send(request, HttpResponse.BodyHandlers.ofString()).body()); + } finally { server.stop(); } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void installedBoundaryRetainsExistingFatalErrorMappingOnHttp11(boolean resumed) throws Exception { + var calls = new AtomicInteger(); + Error sentinel = new StackOverflowError("installed boundary sentinel"); + ExceptionHandler selected = new ExceptionHandler(); + HttpHandler failing = exchange -> { calls.incrementAndGet(); throw sentinel; }; + selected.setNext(exchange -> { + if (exchange.getRequestPath().equals("/after")) exchange.getResponseSender().send("accepted"); + else if (resumed) exchange.dispatch(HandlerContinuation.capture(exchange, failing)); + else failing.handleRequest(exchange); + }); + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(selected).build(); + try (HttpClient client = HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1).build()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + URI endpoint = URI.create("http://127.0.0.1:" + port); + var response = client.send(HttpRequest.newBuilder(endpoint).timeout(Duration.ofSeconds(5)).build(), + HttpResponse.BodyHandlers.ofString()); + assertEquals(400, response.statusCode(), response.body()); + assertTrue(response.body().contains("ERR10011"), response.body()); + assertEquals(1, calls.get()); + assertEquals("accepted", client.send(HttpRequest.newBuilder(endpoint.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString()).body()); + } finally { server.stop(); } + } + +} diff --git a/handler/pom.xml b/handler/pom.xml index 6ef2321cb9..995f96422b 100644 --- a/handler/pom.xml +++ b/handler/pom.xml @@ -90,6 +90,11 @@ com.networknt common + + org.mockito + mockito-core + test + diff --git a/handler/src/main/java/com/networknt/handler/HandlerContinuation.java b/handler/src/main/java/com/networknt/handler/HandlerContinuation.java new file mode 100644 index 0000000000..5d8f29d7ef --- /dev/null +++ b/handler/src/main/java/com/networknt/handler/HandlerContinuation.java @@ -0,0 +1,203 @@ +package com.networknt.handler; + +import io.undertow.server.HandlerWrapper; +import io.undertow.server.HttpHandler; +import io.undertow.server.HttpServerExchange; +import io.undertow.util.StatusCodes; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.slf4j.MDC; + +import java.util.Map; +import java.util.Objects; + +/** Lexical exception policy, carried explicitly across asynchronous hand-offs. */ +public final class HandlerContinuation { + private static final Logger LOG = LoggerFactory.getLogger(HandlerContinuation.class); + private record Frame(HttpServerExchange exchange, RequestState state, + Scope scope, boolean mdcScoped) { } + private static final ThreadLocal CURRENT = new ThreadLocal<>(); + + private HandlerContinuation() { } + + /** Execute within a native Undertow wrapper; restore the enclosing lexical scope on exit. */ + public static void execute(HttpServerExchange exchange, HandlerWrapper boundary, HttpHandler continuation) throws Exception { + Objects.requireNonNull(boundary); + Frame previous = CURRENT.get(); + boolean sameExchange = previous != null && previous.exchange() == exchange; + RequestState state = sameExchange ? previous.state() : RequestState.initialize(exchange); + Scope scope = new Scope(exchange, boundary, state); + Frame entry = new Frame(exchange, state, scope, sameExchange && previous.mdcScoped()); + CURRENT.set(entry); + try { boundary.wrap(new EntryContinuation(scope, continuation, entry)).handleRequest(exchange); } + finally { restore(previous); } + } + + /** + * Capture the current lexical policy. Outside a frame the policy is absent: + * an exchange cannot identify a caller's position in the chain. Off-thread + * callers must retain the Scope captured before their hand-off. + */ + public static Scope capture(HttpServerExchange exchange) { + Frame frame = CURRENT.get(); + boolean sameExchange = frame != null && frame.exchange() == exchange; + return sameExchange ? frame.scope() : new Scope(exchange, null, RequestState.initialize(exchange)); + } + + public static HttpHandler capture(HttpServerExchange exchange, HttpHandler continuation) { + return capture(exchange).wrap(continuation); + } + + public static Runnable task(HttpServerExchange exchange, HttpHandler continuation) { + if (continuation instanceof BoundHandler bound) { + if (exchange != bound.scope.exchange) throw new IllegalArgumentException("Fork the captured scope for a different exchange"); + return bound.scope.task(bound.continuation); + } + return capture(exchange).task(continuation); + } + + /** A reusable, typed request scope. It never guesses a policy from exchange history. */ + public static final class Scope implements HandlerWrapper { + private final HttpServerExchange exchange; + private final HandlerWrapper boundary; + private final RequestState state; + private final Frame resumedFrame; + + private Scope(HttpServerExchange exchange, HandlerWrapper boundary, RequestState state) { + this.exchange = exchange; + this.boundary = boundary; + this.state = state; + this.resumedFrame = new Frame(exchange, state, this, true); + } + + public Scope fork(HttpServerExchange child) { + return new Scope(child, boundary, RequestState.initialize(child, state.mdc.snapshot())); + } + + @Override public HttpHandler wrap(HttpHandler continuation) { return new BoundHandler(this, continuation); } + + public Runnable task(HttpHandler continuation) { return () -> execute(continuation); } + + /** Selected policy handles ordinary failures; unhandled Errors are cleaned up then propagated. */ + public void execute(HttpHandler continuation) { + try { + inContext(continuation, true); + } catch (Exception failure) { + finishFailedRequest(failure); + LOG.error("Unhandled continuation failure for {}", exchange.getRequestPath(), failure); + } catch (Error failure) { + finishFailedRequest(failure); + throw failure; + } + } + + /** Resume accepted work; replay only a failure through the selected policy, never the work itself. */ + public void resume(HttpHandler continuation) { + try { inContext(continuation, false); } + catch (Exception | Error failure) { execute(new FailureHandler(failure)); } + } + + private void inContext(HttpHandler continuation, boolean enterBoundary) throws Exception { + Frame previous = CURRENT.get(); + boolean reuse = previous != null && previous.exchange() == exchange && previous.mdcScoped(); + boolean same = reuse && previous.scope() == this; + Map before = reuse ? null : MDC.getCopyOfContextMap(); + if (!same) CURRENT.set(resumedFrame); + try { + if (!reuse) { + Map request = state.mdc.snapshot(); + if (!request.equals(before == null ? Map.of() : before)) setMdc(request); + } + invoke(continuation, enterBoundary); + } finally { + if (!reuse) setMdc(before); + if (!same) restore(previous); + } + } + + private void invoke(HttpHandler continuation, boolean enterBoundary) throws Exception { + if (!enterBoundary || boundary == null) { + continuation.handleRequest(exchange); + return; + } + HttpHandler contextual = continuation instanceof ScopedCallback callback && callback.scope() == this + ? continuation : new ContextualHandler(this, continuation); + boundary.wrap(contextual).handleRequest(exchange); + } + + private static void setMdc(Map context) { + if (context == null || context.isEmpty()) MDC.clear(); + else MDC.setContextMap(context); + } + + /** Best-effort cleanup must never mask an already propagating failure. */ + private void finishFailedRequest(Throwable primary) { + if (exchange.isComplete()) return; + if (!exchange.isResponseStarted()) { + try { + exchange.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR); + exchange.endExchange(); + return; + } catch (Throwable cleanupFailure) { + rememberCleanupFailure(primary, cleanupFailure); + } + } + try { + if (exchange.getConnection() != null) exchange.getConnection().close(); + } catch (Throwable cleanupFailure) { + rememberCleanupFailure(primary, cleanupFailure); + } + } + private static void rememberCleanupFailure(Throwable primary, Throwable cleanupFailure) { + if (primary == cleanupFailure) return; + try { + primary.addSuppressed(cleanupFailure); + LOG.error("Continuation cleanup failed; retaining the original failure", cleanupFailure); + } catch (Throwable _) { + // Diagnostics can fail under memory pressure; never replace the primary failure. + } + } + + } + + private interface ScopedCallback extends HttpHandler { Scope scope(); } + + private record EntryContinuation(Scope scope, HttpHandler continuation, Frame entry) implements HttpHandler { + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { + if (CURRENT.get() == entry) continuation.handleRequest(exchange); + else scope.inContext(continuation, false); + } + } + + private record BoundHandler(Scope scope, HttpHandler continuation, ContextualHandler contextual) implements HttpHandler { + private BoundHandler(Scope scope, HttpHandler continuation) { + this(scope, continuation, new ContextualHandler(scope, continuation)); + } + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { + if (exchange != scope.exchange) throw new IllegalArgumentException("Fork the captured scope for a different exchange"); + scope.inContext(contextual, true); + } + } + + private record ContextualHandler(Scope scope, HttpHandler continuation) implements ScopedCallback { + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { + if (exchange != scope.exchange) throw new IllegalArgumentException("Fork the captured scope for a different exchange"); + scope.inContext(continuation, false); + } + } + + private record FailureHandler(Throwable failure) implements HttpHandler { + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { + if (failure instanceof Exception exception) throw exception; + if (failure instanceof Error error) throw error; + throw new IllegalArgumentException("Unsupported failure", failure); + } + } + + private static void restore(Frame frame) { + if (frame == null) CURRENT.remove(); + else CURRENT.set(frame); + } + + +} diff --git a/handler/src/main/java/com/networknt/handler/MiddlewareHandler.java b/handler/src/main/java/com/networknt/handler/MiddlewareHandler.java index 7c5c370edd..894a34826b 100644 --- a/handler/src/main/java/com/networknt/handler/MiddlewareHandler.java +++ b/handler/src/main/java/com/networknt/handler/MiddlewareHandler.java @@ -16,11 +16,9 @@ package com.networknt.handler; -import com.networknt.httpstring.AttachmentConstants; import io.undertow.server.HttpHandler; import io.undertow.server.HttpServerExchange; -import java.util.HashMap; /** * A interface for middleware handlers. All middleware handlers must implement this interface @@ -60,20 +58,15 @@ public interface MiddlewareHandler extends LightHttpHandler { boolean isEnabled(); /** - * Adds a field and value to the MDC context for the handler. + * Adds a field and value to the MDC context for the handler. First use must + * occur in the serial frame before asynchronous hand-off. Subsequent updates + * publish immutable snapshots through the stable holder. * * @param exchange HttpServerExchange to add the context to. * @param field The field name. * @param value The value for the field. */ default void addHandlerMDCContext(final HttpServerExchange exchange, String field, String value) { - var context = exchange.getAttachment(AttachmentConstants.MDC_CONTEXT); - if (context != null) { - context.put(field, value); - } else { - var newContext = new HashMap(); - newContext.put(field, value); - exchange.putAttachment(AttachmentConstants.MDC_CONTEXT, newContext); - } + RequestMdcContext.initialize(exchange).putField(field, value); } } diff --git a/handler/src/main/java/com/networknt/handler/RequestInterceptorInjectionHandler.java b/handler/src/main/java/com/networknt/handler/RequestInterceptorInjectionHandler.java index d772b72cc5..9d65086992 100644 --- a/handler/src/main/java/com/networknt/handler/RequestInterceptorInjectionHandler.java +++ b/handler/src/main/java/com/networknt/handler/RequestInterceptorInjectionHandler.java @@ -167,7 +167,9 @@ private enum ReadOutcome { private final class RequestBodyReadState implements ChannelListener { private final HttpServerExchange exchange; private final RequestInjectionConfig config; - private final HttpHandler continuation; + private final HttpHandler nextHandler; + private HttpHandler resumedHandler; + private Runnable readContinuation; private final StreamSourceChannel channel; private final PooledByteBuffer[] bufferedData; private final boolean exactLimit; @@ -175,10 +177,10 @@ private final class RequestBodyReadState implements ChannelListener 0; @@ -241,6 +243,11 @@ private ReadOutcome probeForEndOfStream() throws Exception { } private ReadOutcome awaitMoreData() { + if (resumedHandler == null) { + HandlerContinuation.Scope scope = HandlerContinuation.capture(exchange); + resumedHandler = scope.wrap(ex -> Handler.next(ex, nextHandler)); + readContinuation = scope.task(ex -> resumeRead()); + } channel.getReadSetter().set(this); channel.resumeReads(); return ReadOutcome.PENDING; @@ -265,23 +272,28 @@ private void closeBuffers() { @Override public void handleEvent(StreamSourceChannel ignored) { + // A pending read must not be auto-ended by a root-handler wrapper. + readContinuation.run(); + } + + private void resumeRead() throws Exception { + boolean complete; try { ReadOutcome outcome = readAvailable(); - if (outcome == ReadOutcome.PENDING) - return; - + if (outcome == ReadOutcome.PENDING) return; stopReading(); - if (outcome == ReadOutcome.COMPLETE) { - saveBufferAndResetUndertowConnector(exchange, bufferedData); - if (!exchange.isResponseStarted()) - Connectors.executeRootHandler(ex -> Handler.next(ex, continuation), exchange); - } + complete = outcome == ReadOutcome.COMPLETE; + if (complete) saveBufferAndResetUndertowConnector(exchange, bufferedData); } catch (Throwable e) { LOG.error(e.getMessage(), e); stopReading(); closeBuffers(); setExchangeStatus(exchange, GENERIC_EXCEPTION, e.getMessage()); + return; } + // Downstream failures belong to the selected chain boundary, not body-read errors. + if (complete && !exchange.isResponseStarted()) + Connectors.executeRootHandler(resumedHandler, exchange); } } diff --git a/handler/src/main/java/com/networknt/handler/RequestMdcContext.java b/handler/src/main/java/com/networknt/handler/RequestMdcContext.java new file mode 100644 index 0000000000..7f1321b8cd --- /dev/null +++ b/handler/src/main/java/com/networknt/handler/RequestMdcContext.java @@ -0,0 +1,55 @@ +package com.networknt.handler; + +import io.undertow.server.HttpServerExchange; + +import java.util.AbstractMap; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiConsumer; + +/** + * Stable, read-only Map view of request MDC. Initialize in the first serial + * request frame before handing the exchange to another thread. Subsequent + * publication changes only the atomic snapshot, never the attachment map. + */ +public final class RequestMdcContext extends AbstractMap { + private final AtomicReference> current; + + RequestMdcContext(Map initial) { + current = new AtomicReference<>(initial == null || initial.isEmpty() + ? Map.of() : Collections.unmodifiableMap(new HashMap<>(initial))); + } + + /** + * First use must be serial. Boundary execution and capture initialize before + * registering asynchronous work; repeated calls do not write any attachment. + * Legacy Map-valued context is migrated once in that serial frame. + */ + public static RequestMdcContext initialize(HttpServerExchange exchange) { + return RequestState.initialize(exchange).mdc; + } + + /** One immutable generation suitable for coherent iteration by external readers. */ + public Map snapshot() { return current.get(); } + + void putField(String field, String value) { + current.updateAndGet(previous -> { + Map next = new HashMap<>(previous); + next.put(field, value); + return Collections.unmodifiableMap(next); + }); + } + + @Override public String get(Object key) { return current.get().get(key); } + @Override public boolean containsKey(Object key) { return current.get().containsKey(key); } + @Override public boolean equals(Object other) { return other == this || current.get().equals(other); } + @Override public int hashCode() { return current.get().hashCode(); } + @Override public int size() { return current.get().size(); } + @Override public Set> entrySet() { return current.get().entrySet(); } + @Override public void forEach(BiConsumer consumer) { + current.get().forEach(consumer); + } +} diff --git a/handler/src/main/java/com/networknt/handler/RequestState.java b/handler/src/main/java/com/networknt/handler/RequestState.java new file mode 100644 index 0000000000..64a21d1c27 --- /dev/null +++ b/handler/src/main/java/com/networknt/handler/RequestState.java @@ -0,0 +1,82 @@ +package com.networknt.handler; + +import com.networknt.httpstring.AttachmentConstants; +import io.undertow.server.HttpServerExchange; + +import java.lang.ref.Reference; +import java.lang.ref.ReferenceQueue; +import java.lang.ref.WeakReference; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** MDC only: exception policy is carried explicitly by captured scopes. */ +final class RequestState { + private interface Key { HttpServerExchange get(); } + + private static final class WeakKey extends WeakReference implements Key { + private final int hash; + private WeakKey(HttpServerExchange exchange) { + super(exchange, CLEARED); + hash = System.identityHashCode(exchange); + } + @Override public int hashCode() { return hash; } + @Override public boolean equals(Object other) { + HttpServerExchange exchange = get(); + return this == other || exchange != null && other instanceof Key key && exchange == key.get(); + } + } + + /** Reusable lookup key; never inserted and never retains an exchange after get(). */ + private static final class LookupKey implements Key { + private HttpServerExchange exchange; + @Override public HttpServerExchange get() { return exchange; } + @Override public int hashCode() { return System.identityHashCode(exchange); } + @Override public boolean equals(Object other) { + return other instanceof Key key && exchange == key.get(); + } + } + + private static final ReferenceQueue CLEARED = new ReferenceQueue<>(); + private static final ConcurrentHashMap INDEX = new ConcurrentHashMap<>(); + private static final ThreadLocal LOOKUP = ThreadLocal.withInitial(LookupKey::new); + // Values contain strings/MDC only, never a boundary closure or an exchange. + final RequestMdcContext mdc; + + private RequestState(RequestMdcContext mdc) { this.mdc = mdc; } + + static RequestState find(HttpServerExchange exchange) { + LookupKey key = LOOKUP.get(); + key.exchange = exchange; + try { return INDEX.get(key); } + finally { key.exchange = null; } + } + + static RequestState initialize(HttpServerExchange exchange) { return initialize(exchange, Map.of()); } + + static RequestState initialize(HttpServerExchange exchange, Map inheritedMdc) { + RequestState existing = find(exchange); + if (existing != null) return existing; + if (exchange.isComplete()) return detached(exchange, inheritedMdc); + // Bounded weak-key cleanup; no exchange-completion listener is registered. + for (int count = 0; count < 16; count++) { + var key = CLEARED.poll(); + if (key == null) break; + INDEX.remove(key); + } + WeakKey key = new WeakKey(exchange); + try { + return INDEX.computeIfAbsent(key, ignored -> { + RequestState created = detached(exchange, inheritedMdc); + exchange.putAttachment(AttachmentConstants.MDC_CONTEXT, created.mdc); + return created; + }); + } finally { Reference.reachabilityFence(exchange); } + } + + private static RequestState detached(HttpServerExchange exchange, Map inheritedMdc) { + Map attached = exchange.getAttachment(AttachmentConstants.MDC_CONTEXT); + Map initial = attached == null ? inheritedMdc : attached; + RequestMdcContext context = attached instanceof RequestMdcContext holder ? holder : new RequestMdcContext(initial); + return new RequestState(context); + } +} diff --git a/handler/src/test/java/com/networknt/handler/HandlerContinuationTest.java b/handler/src/test/java/com/networknt/handler/HandlerContinuationTest.java new file mode 100644 index 0000000000..8c088a3f3b --- /dev/null +++ b/handler/src/test/java/com/networknt/handler/HandlerContinuationTest.java @@ -0,0 +1,483 @@ +package com.networknt.handler; + +import com.networknt.httpstring.AttachmentConstants; +import io.undertow.server.HttpHandler; +import io.undertow.server.HandlerWrapper; +import io.undertow.server.HttpServerExchange; +import io.undertow.server.ServerConnection; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; + +import java.io.IOException; +import java.lang.ref.WeakReference; +import java.lang.ref.ReferenceQueue; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +class HandlerContinuationTest { + @Test + void offFrameCaptureRetainsEntryPolicyAndHolderOnAnotherThread() throws Exception { + var exchange = spy(new HttpServerExchange(null)); + var selected = new AtomicInteger(); + var saved = new AtomicReference(); + HandlerContinuation.execute(exchange, next -> ex -> { + try { next.handleRequest(ex); } + catch (IOException _) { selected.incrementAndGet(); } + }, ex -> { + publicationMiddleware().addHandlerMDCContext(ex, "request", "entry"); + saved.set(HandlerContinuation.capture(ex)); + }); + var holder = RequestMdcContext.initialize(exchange); + clearInvocations(exchange); + try (var worker = Executors.newSingleThreadExecutor()) { + worker.submit(() -> { + assertSame(holder, RequestMdcContext.initialize(exchange)); + HttpHandler failing = ex -> { + assertEquals("entry", MDC.get("request")); + throw new IOException("off-frame failure"); + }; + assertThrows(IOException.class, () -> HandlerContinuation.capture(exchange, failing).handleRequest(exchange), + "An exchange alone cannot infer lexical policy outside its frame"); + saved.get().wrap(failing).handleRequest(exchange); + return null; + }).get(5, TimeUnit.SECONDS); + } + assertEquals(1, selected.get()); + verify(exchange, never()).getAttachment(any()); + verify(exchange, never()).putAttachment(any(), any()); + } + + @Test + void deferredBoundaryRetainsItsContinuationAndRestoresWorkerMdc() throws Exception { + var exchange = new HttpServerExchange(null); + var pending = new AtomicReference(); + var saved = new AtomicReference(); + HandlerContinuation.execute(exchange, next -> ex -> pending.set(next), ex -> { + publicationMiddleware().addHandlerMDCContext(ex, "request", "deferred"); + saved.set(HandlerContinuation.capture(ex)); + }); + pending.getAndSet(null).handleRequest(exchange); + var calls = new AtomicInteger(); + saved.get().task(ex -> { + assertEquals("deferred", MDC.get("request")); + calls.incrementAndGet(); + }).run(); + assertEquals(0, calls.get()); + try (var worker = Executors.newSingleThreadExecutor()) { + worker.submit(() -> { + MDC.put("worker", "prior"); + try { + pending.getAndSet(null).handleRequest(exchange); + assertEquals("prior", MDC.get("worker")); + } finally { MDC.clear(); } + return null; + }).get(5, TimeUnit.SECONDS); + } + assertEquals(1, calls.get()); + } + + @Test + void inheritedExchangeOwnsItsMdcAndUsesTheSelectedPolicy() throws Exception { + var parent = new HttpServerExchange(null); + var child = new HttpServerExchange(null); + var policy = new AtomicReference(); + var seen = new ArrayList(); + HandlerContinuation.execute(parent, next -> ex -> { seen.add(ex); next.handleRequest(ex); }, ex -> { + publicationMiddleware().addHandlerMDCContext(ex, "request", "parent"); + policy.set(HandlerContinuation.capture(ex)); + }); + seen.clear(); + policy.get().fork(child).wrap(ex -> { + assertSame(child, ex); + assertEquals("parent", MDC.get("request")); + publicationMiddleware().addHandlerMDCContext(ex, "request", "child"); + }).handleRequest(child); + assertEquals(List.of(child), seen); + assertEquals("parent", RequestMdcContext.initialize(parent).snapshot().get("request")); + assertEquals("child", RequestMdcContext.initialize(child).snapshot().get("request")); + } + + @Test + void activePolicyTracksIdentityAndExchangeAndRestoresAfterNestedFailure() throws Exception { + HttpServerExchange first = new HttpServerExchange(null); + HttpServerExchange second = new HttpServerExchange(null); + var policy = HandlerContinuation.capture(first); + var otherPolicy = HandlerContinuation.capture(first); + assertNotSame(policy, HandlerContinuation.capture(first)); + policy.wrap(ex -> { + assertSame(policy, HandlerContinuation.capture(first)); + assertNotSame(otherPolicy, HandlerContinuation.capture(first)); + assertNotSame(policy, HandlerContinuation.capture(second)); + var child = policy.fork(second); + assertThrows(IOException.class, () -> child.wrap(nested -> { + assertSame(child, HandlerContinuation.capture(second)); + assertNotSame(policy, HandlerContinuation.capture(first)); + throw new IOException("nested failure"); + }).handleRequest(second)); + assertSame(policy, HandlerContinuation.capture(first)); + }).handleRequest(first); + assertNotSame(policy, HandlerContinuation.capture(first)); + assertNotSame(policy, HandlerContinuation.capture(second)); + } + + @Test + void mdcPublicationRetainsImmutableSnapshotsAndConcurrentUpdates() throws Exception { + HttpServerExchange exchange = spy(new HttpServerExchange(null)); + MiddlewareHandler middleware = new MiddlewareHandler() { + @Override public boolean isEnabled() { return true; } + @Override public HttpHandler getNext() { return null; } + @Override public MiddlewareHandler setNext(HttpHandler next) { return this; } + @Override public void handleRequest(HttpServerExchange ex) { /* Publication-only fixture. */ } + }; + middleware.addHandlerMDCContext(exchange, "first", "one"); + var holder = assertInstanceOf(RequestMdcContext.class, + exchange.getAttachment(AttachmentConstants.MDC_CONTEXT)); + var first = holder.snapshot(); + assertThrows(UnsupportedOperationException.class, () -> first.put("bad", "mutation")); + try (var workers = Executors.newFixedThreadPool(4)) { + var updates = java.util.stream.IntStream.range(0, 20).mapToObj(i -> workers.submit(() -> + middleware.addHandlerMDCContext(exchange, "field" + i, "value" + i))).toList(); + for (var update : updates) update.get(); + } + assertEquals(Map.of("first", "one"), first); + assertEquals(21, holder.size()); + assertSame(holder, exchange.getAttachment(AttachmentConstants.MDC_CONTEXT)); + verify(exchange, times(1)).putAttachment( + eq(AttachmentConstants.MDC_CONTEXT), any()); + } + + @Test + void productionResumesReadStableSnapshotsWhileWritersPublish() throws Exception { + HttpServerExchange exchange = spy(new HttpServerExchange(null)); + exchange.putAttachment(AttachmentConstants.MDC_CONTEXT, Map.of("epoch", "0", "legacy", "retained")); + MiddlewareHandler middleware = publicationMiddleware(); + AtomicReference context = new AtomicReference<>(); + AtomicReference writer = new AtomicReference<>(); + AtomicReference reader = new AtomicReference<>(); + CountDownLatch readerHoldingSnapshot = new CountDownLatch(1); + CountDownLatch updated = new CountDownLatch(1); + HandlerContinuation.execute(exchange, next -> ex -> next.handleRequest(ex), ex -> { + context.set(RequestMdcContext.initialize(ex)); + reader.set(HandlerContinuation.capture(ex, resumed -> { + assertEquals("retained", MDC.get("legacy")); + var generation = context.get().snapshot(); + var iterator = generation.entrySet().iterator(); + assertTrue(iterator.hasNext()); + var first = iterator.next(); + readerHoldingSnapshot.countDown(); + assertTrue(updated.await(5, TimeUnit.SECONDS)); + var copy = new HashMap(); + copy.put(first.getKey(), first.getValue()); + iterator.forEachRemaining(entry -> copy.put(entry.getKey(), entry.getValue())); + assertEquals("0", copy.get("epoch"), "A retained generation must not change during publication"); + assertThrows(UnsupportedOperationException.class, () -> generation.put("bad", "mutation")); + })); + writer.set(HandlerContinuation.capture(ex, resumed -> { + assertTrue(readerHoldingSnapshot.await(5, TimeUnit.SECONDS)); + for (int i = 1; i <= 2000; i++) middleware.addHandlerMDCContext(resumed, "epoch", Integer.toString(i)); + updated.countDown(); + })); + }); + clearInvocations(exchange); + try (var workers = Executors.newFixedThreadPool(2)) { + var reading = workers.submit(() -> { reader.get().handleRequest(exchange); return null; }); + var writing = workers.submit(() -> { writer.get().handleRequest(exchange); return null; }); + reading.get(10, TimeUnit.SECONDS); + writing.get(10, TimeUnit.SECONDS); + } + // Both production resumes and helper updates reuse the captured holder without attachment access. + verify(exchange, never()).getAttachment(AttachmentConstants.MDC_CONTEXT); + verify(exchange, never()).putAttachment(eq(AttachmentConstants.MDC_CONTEXT), any()); + reader.set(HandlerContinuation.capture(exchange, resumed -> assertEquals("2000", MDC.get("epoch")))); + reader.get().handleRequest(exchange); + assertEquals("retained", context.get().snapshot().get("legacy")); + } + + @Test + void earlierNarrowPolicyCannotShadowLaterPolicyOrOffFrameCapture() throws Exception { + var exchange = new HttpServerExchange(null); + var narrow = new AtomicReference(); + var later = new AtomicReference(); + var seen = new ArrayList(); + HandlerWrapper first = next -> ex -> { + try { next.handleRequest(ex); } + catch (IOException _) { seen.add("narrow"); } + }; + HandlerWrapper second = next -> ex -> { + try { next.handleRequest(ex); } + catch (IOException _) { seen.add("later"); } + }; + HandlerContinuation.execute(exchange, first, ex -> narrow.set(HandlerContinuation.capture(ex))); + var outside = HandlerContinuation.capture(exchange); + HandlerContinuation.execute(exchange, second, ex -> later.set(HandlerContinuation.capture(ex))); + HttpHandler fail = ex -> { throw new IOException("selected failure"); }; + assertThrows(IOException.class, () -> outside.wrap(fail).handleRequest(exchange)); + narrow.get().wrap(fail).handleRequest(exchange); + later.get().wrap(fail).handleRequest(exchange); + assertEquals(List.of("narrow", "later"), seen); + } + + @Test + void acceptedWorkRunsOnceAndItsFailureUsesSelectedPolicy() throws Exception { + var exchange = new HttpServerExchange(null); + var scope = new AtomicReference(); + var failure = new IOException("accepted failure"); + var observed = new AtomicReference(); + var calls = new AtomicInteger(); + HandlerContinuation.execute(exchange, next -> ex -> { + try { next.handleRequest(ex); } + catch (IOException caught) { observed.set(caught); } + }, ex -> scope.set(HandlerContinuation.capture(ex))); + scope.get().resume(ex -> { calls.incrementAndGet(); throw failure; }); + assertSame(failure, observed.get()); + assertEquals(1, calls.get()); + } + + @Test + void captureDoesNotApplyWrapperUntilExecution() throws Exception { + var exchange = new HttpServerExchange(null); + var applications = new AtomicInteger(); + var captured = new AtomicReference(); + HandlerWrapper boundary = next -> { + applications.incrementAndGet(); + return next; + }; + HandlerContinuation.execute(exchange, boundary, ex -> { + captured.set(HandlerContinuation.capture(ex, resumed -> assertSame(exchange, resumed))); + assertEquals(1, applications.get()); + }); + captured.get().handleRequest(exchange); + assertEquals(2, applications.get()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void initializationNeverRegistersACompletionListener(boolean alreadyComplete) { + var exchange = spy(new HttpServerExchange(null)); + doReturn(alreadyComplete, true).when(exchange).isComplete(); + doThrow(new IllegalStateException("exchange completed during registration")) + .when(exchange).addExchangeCompleteListener(any()); + var state = assertDoesNotThrow(() -> RequestState.initialize(exchange)); + assertNotNull(state.mdc); + assertTrue(exchange.isComplete()); + verify(exchange, never()).addExchangeCompleteListener(any()); + if (alreadyComplete) verify(exchange, never()).putAttachment(any(), any()); + else assertSame(state, RequestState.initialize(exchange)); + } + + @Test + void weakIndexReleasesAbandonedExchangeAndMdcAfterReferenceQueueCleanup() throws Exception { + var collected = new ReferenceQueue(); + var abandoned = abandonedRequest(collected); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (abandoned.get(1).get() != null && System.nanoTime() < deadline) { + System.gc(); + // Each new registration performs bounded ReferenceQueue expunging. + RequestState.initialize(new HttpServerExchange(null)); + collected.remove(10); + } + assertNull(abandoned.get(0).get(), "The index must not retain the exchange"); + assertNull(abandoned.get(1).get(), "A queued weak key must be removed, releasing its MDC state"); + } + + private static List> abandonedRequest(ReferenceQueue collected) { + var exchange = new HttpServerExchange(null); + var state = RequestState.initialize(exchange); + return List.of(new WeakReference<>(exchange, collected), new WeakReference<>(state, collected)); + } + + private static MiddlewareHandler publicationMiddleware() { + return new MiddlewareHandler() { + @Override public boolean isEnabled() { return true; } + @Override public HttpHandler getNext() { return null; } + @Override public MiddlewareHandler setNext(HttpHandler next) { return this; } + @Override public void handleRequest(HttpServerExchange ex) { /* Publication-only fixture. */ } + }; + } + + @Test + void legacyContextInRawRunnableIsMigratedBeforeItRuns() throws Exception { + HttpServerExchange exchange = spy(new HttpServerExchange(null)); + exchange.putAttachment(AttachmentConstants.MDC_CONTEXT, Map.of("legacy", "value")); + var callback = HandlerContinuation.task(exchange, ex -> assertEquals("value", MDC.get("legacy"))); + clearInvocations(exchange); + try (var workers = Executors.newSingleThreadExecutor()) { workers.submit(callback).get(5, TimeUnit.SECONDS); } + verify(exchange, never()).getAttachment(AttachmentConstants.MDC_CONTEXT); + verify(exchange, never()).putAttachment(eq(AttachmentConstants.MDC_CONTEXT), any()); + } + + @Test + void exceptionCleanupPreservesFailureWhenStatusWriteAndResetThrow() throws Exception { + IOException original = new IOException("callback failed"); + RuntimeException writeFailure = new IllegalStateException("write failed"); + RuntimeException closeFailure = new IllegalStateException("close failed"); + var connection = mock(ServerConnection.class); + doThrow(closeFailure).when(connection).close(); + var exchange = spy(new HttpServerExchange(connection)); + doThrow(writeFailure).when(exchange).setStatusCode(500); + assertDoesNotThrow(() -> HandlerContinuation.task(exchange, ex -> { throw original; }).run()); + assertArrayEquals(new Throwable[]{writeFailure, closeFailure}, original.getSuppressed()); + verify(exchange).setStatusCode(500); + verify(connection).close(); + } + + @Test + void overlappingFramesKeepTheirOwnPolicyWithoutAttachmentAccess() throws Exception { + var exchange = spy(new HttpServerExchange(null)); + var aEntered = new CountDownLatch(1); + var bEntered = new CountDownLatch(1); + var aCaptured = new CountDownLatch(1); + var callsA = new AtomicInteger(); + var callsB = new AtomicInteger(); + var a = new AtomicReference(); + var b = new AtomicReference(); + HandlerContinuation.execute(exchange, next -> ex -> { callsA.incrementAndGet(); next.handleRequest(ex); }, ex -> + a.set(HandlerContinuation.capture(ex, resumed -> { + aEntered.countDown(); + assertTrue(bEntered.await(5, TimeUnit.SECONDS)); + HandlerContinuation.capture(resumed, child -> { /* Exercise the captured policy. */ }).handleRequest(resumed); + aCaptured.countDown(); + }))); + HandlerContinuation.execute(exchange, next -> ex -> { callsB.incrementAndGet(); next.handleRequest(ex); }, ex -> + b.set(HandlerContinuation.capture(ex, resumed -> { + assertTrue(aEntered.await(5, TimeUnit.SECONDS)); + bEntered.countDown(); + assertTrue(aCaptured.await(5, TimeUnit.SECONDS)); + }))); + callsA.set(0); + callsB.set(0); + clearInvocations(exchange); + try (var workers = Executors.newFixedThreadPool(2)) { + var first = workers.submit(() -> { a.get().handleRequest(exchange); return null; }); + var second = workers.submit(() -> { b.get().handleRequest(exchange); return null; }); + first.get(10, TimeUnit.SECONDS); + second.get(10, TimeUnit.SECONDS); + } + assertEquals(2, callsA.get()); + assertEquals(1, callsB.get()); + verify(exchange, never()).getAttachment(any()); + verify(exchange, never()).putAttachment(any(), any()); + verify(exchange, never()).removeAttachment(any()); + } + + @Test + void failedFatalStatusWriteCannotReplaceOriginalError() throws Exception { + Error fatal = new StackOverflowError("original"); + RuntimeException writeFailure = new IllegalStateException("status write failed"); + IOException closeFailure = new IOException("reset failed too"); + var connection = mock(ServerConnection.class); + doThrow(closeFailure).when(connection).close(); + HttpServerExchange exchange = spy(new HttpServerExchange(connection)); + doThrow(writeFailure).when(exchange).endExchange(); + Runnable callback = HandlerContinuation.task(exchange, ex -> { throw fatal; }); + Error observed = assertThrows(StackOverflowError.class, callback::run); + assertSame(fatal, observed); + assertArrayEquals(new Throwable[]{writeFailure, closeFailure}, fatal.getSuppressed()); + assertEquals(500, exchange.getStatusCode()); + verify(exchange).endExchange(); + verify(connection).close(); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void resumeScopesCanonicalRequestMdcAndRestoresWorker(boolean present) throws Exception { + var original = MDC.getCopyOfContextMap(); + try { + HttpServerExchange exchange = new HttpServerExchange(null); + if (present) exchange.putAttachment(AttachmentConstants.MDC_CONTEXT, + Map.of("request", "canonical")); + MDC.put("request", "stale-at-capture"); + HttpHandler resumed = HandlerContinuation.capture(exchange, ex -> { + assertEquals(present ? "canonical" : null, MDC.get("request")); + assertNull(MDC.get("worker"), "Do not borrow the executing thread's MDC"); + HandlerContinuation.capture(ex, nested -> { + assertEquals(present ? "canonical" : null, MDC.get("request")); + }).handleRequest(ex); + throw new IOException("failure"); + }); + MDC.clear(); + MDC.put("worker", "unrelated"); + assertThrows(IOException.class, () -> resumed.handleRequest(exchange)); + assertEquals(Map.of("worker", "unrelated"), MDC.getCopyOfContextMap()); + } finally { + if (original == null) MDC.clear(); + else MDC.setContextMap(original); + } + } + + @Test + void runnableAdapterPropagatesFatalErrorsAndRestoresMdc() { + var original = MDC.getCopyOfContextMap(); + try { + HttpServerExchange exchange = spy(new HttpServerExchange(null)); + doReturn(exchange).when(exchange).endExchange(); + MDC.put("worker", "prior"); + Error error = new OutOfMemoryError("injected sentinel, no memory pressure"); + Runnable task = HandlerContinuation.task(exchange, + HandlerContinuation.capture(exchange, ex -> { throw error; })); + assertSame(error, assertThrows(OutOfMemoryError.class, task::run)); + assertEquals("prior", MDC.get("worker")); + verify(exchange).endExchange(); + assertEquals(500, exchange.getStatusCode()); + } finally { + if (original == null) MDC.clear(); + else MDC.setContextMap(original); + } + } + + @Test + void nestedCaptureRestoresEnclosingBoundaryAndRetainsSelectedInstances() throws Exception { + HttpServerExchange exchange = new HttpServerExchange(null); + List seen = new ArrayList<>(); + HandlerWrapper outer = task -> ex -> { + seen.add("outer"); + task.handleRequest(ex); + }; + HandlerWrapper inner = task -> ex -> { + seen.add("inner"); + task.handleRequest(ex); + }; + AtomicReference nested = new AtomicReference<>(); + AtomicReference enclosing = new AtomicReference<>(); + HandlerContinuation.execute(exchange, outer, ex -> { + HandlerContinuation.execute(ex, inner, innerExchange -> + nested.set(HandlerContinuation.capture(innerExchange, ignored -> seen.add("nested")))); + enclosing.set(HandlerContinuation.capture(ex, ignored -> seen.add("enclosing"))); + }); + seen.clear(); + nested.get().handleRequest(exchange); + enclosing.get().handleRequest(exchange); + assertEquals(List.of("inner", "nested", "outer", "enclosing"), seen); + seen.clear(); + HandlerContinuation.capture(exchange, ignored -> seen.add("unwrapped")).handleRequest(exchange); + assertEquals(List.of("unwrapped"), seen, "Lexical policy ends with its frame; retain a Scope across hand-off"); + } + + @Test + void absentBoundaryStaysAbsentAndFailuresRestoreScope() throws Exception { + HttpServerExchange exchange = new HttpServerExchange(null); + IOException failure = new IOException("failure"); + HttpHandler unwrapped = HandlerContinuation.capture(exchange, ignored -> { throw failure; }); + HandlerWrapper selected = task -> ex -> { + assertDoesNotThrow(() -> task.handleRequest(ex), + "A captured absent boundary must hide the active boundary"); + }; + HandlerContinuation.execute(exchange, selected, ex -> + assertSame(failure, assertThrows(IOException.class, () -> unwrapped.handleRequest(ex)))); + assertSame(failure, assertThrows(IOException.class, () -> + HandlerContinuation.execute(exchange, task -> ex -> { throw failure; }, ignored -> { /* The boundary throws before this continuation runs. */ }))); + assertSame(failure, assertThrows(IOException.class, () -> unwrapped.handleRequest(exchange))); + } +} diff --git a/handler/src/test/java/com/networknt/handler/RequestInjectionContinuationTest.java b/handler/src/test/java/com/networknt/handler/RequestInjectionContinuationTest.java new file mode 100644 index 0000000000..65ca5f57f6 --- /dev/null +++ b/handler/src/test/java/com/networknt/handler/RequestInjectionContinuationTest.java @@ -0,0 +1,170 @@ +package com.networknt.handler; + +import com.networknt.config.Config; +import com.networknt.service.SingletonServiceFactory; +import io.undertow.Undertow; +import io.undertow.server.DefaultByteBufferPool; +import io.undertow.server.HttpServerExchange; +import io.undertow.server.ServerConnection; +import io.undertow.util.HeaderMap; +import io.undertow.util.Methods; +import org.junit.jupiter.api.Test; +import org.xnio.ChannelListener; +import org.xnio.channels.StreamSourceChannel; +import io.undertow.server.HttpHandler; +import io.undertow.server.HandlerWrapper; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.junit.jupiter.api.Timeout; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.net.Socket; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +@Timeout(20) +class RequestInjectionContinuationTest { + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void partialBodyRetainsBoundaryOrItsAbsence(boolean enabled) throws Exception { + String config = "request-continuation-module-test"; + RequestInterceptor[] previous = SingletonServiceFactory.getBeans(RequestInterceptor.class); + var interceptorMdc = new java.util.concurrent.atomic.AtomicReference(); + var downstreamMdc = new java.util.concurrent.atomic.AtomicReference(); + RequestInterceptor interceptor = new RequestInterceptor() { + @Override public boolean isRequiredContent() { return true; } + @Override public boolean isEnabled() { return true; } + @Override public HttpHandler getNext() { return null; } + @Override public MiddlewareHandler setNext(HttpHandler next) { return this; } + @Override public void handleRequest(io.undertow.server.HttpServerExchange exchange) { + interceptorMdc.set(org.slf4j.MDC.get("request")); + } + }; + SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), new RequestInterceptor[]{interceptor}); + Config.getInstance().putInConfigCache(config, Map.of("enabled", true, "maxBuffers", 4, + "maxBodyBytes", 1024, "appliedBodyInjectionPathPrefixes", List.of("/"))); + RequestInterceptorInjectionHandler reader = new RequestInterceptorInjectionHandler(config); + AtomicBoolean onIo = new AtomicBoolean(); + reader.setNext(exchange -> { + onIo.set(exchange.isInIoThread()); + downstreamMdc.set(org.slf4j.MDC.get("request")); + throw new IOException("downstream"); + }); + CountDownLatch suspended = new CountDownLatch(1); + HandlerWrapper boundary = task -> exchange -> { + try { task.handleRequest(exchange); } + catch (IOException _) { exchange.setStatusCode(409).getResponseSender().send("selected boundary"); } + }; + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setIoThreads(1) + .setHandler(exchange -> { + if (exchange.getRequestPath().equals("/probe")) { + exchange.getResponseSender().send(org.slf4j.MDC.get("worker") + ":" + org.slf4j.MDC.get("request")); + org.slf4j.MDC.clear(); + return; + } + exchange.putAttachment(com.networknt.httpstring.AttachmentConstants.MDC_CONTEXT, + Map.of("request", "canonical")); + org.slf4j.MDC.put("worker", "prior-io"); + if (enabled) HandlerContinuation.execute(exchange, boundary, reader); + else reader.handleRequest(exchange); + suspended.countDown(); + }).build(); + try { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + try (Socket socket = new Socket("127.0.0.1", port)) { + socket.setSoTimeout(5000); + socket.getOutputStream().write((""" + POST / HTTP/1.1\r + Host: localhost\r + Content-Type: application/json\r + Transfer-Encoding: chunked\r + Connection: close\r + \r + 1\r + {\r + """).getBytes(StandardCharsets.UTF_8)); + socket.getOutputStream().flush(); + assertTrue(suspended.await(5, TimeUnit.SECONDS)); + socket.getOutputStream().write("1\r\n}\r\n0\r\n\r\n".getBytes(StandardCharsets.UTF_8)); + socket.getOutputStream().flush(); + String response = new String(socket.getInputStream().readAllBytes(), StandardCharsets.UTF_8); + assertTrue(response.startsWith(enabled ? "HTTP/1.1 409" : "HTTP/1.1 500"), response); + assertEquals(enabled, response.contains("selected boundary"), response); + assertTrue(onIo.get(), "Exercise the actual channel-listener continuation"); + assertEquals("canonical", interceptorMdc.get()); + assertEquals("canonical", downstreamMdc.get()); + try (var client = java.net.http.HttpClient.newHttpClient()) { + var probe = client.send(java.net.http.HttpRequest.newBuilder( + java.net.URI.create("http://127.0.0.1:" + port + "/probe")).build(), + java.net.http.HttpResponse.BodyHandlers.ofString()); + assertEquals("prior-io:null", probe.body(), "Restore the I/O thread's context after the callback"); + } + } + } finally { + server.stop(); + SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), previous); + Config.getInstance().clearConfigCache(config); + } + } + @Test + void listenerFailureOutsideBodyReadTryUsesCapturedPolicy() throws Exception { + String config = "request-listener-failure"; + RequestInterceptor[] previous = SingletonServiceFactory.getBeans(RequestInterceptor.class); + var interceptor = mock(RequestInterceptor.class); + when(interceptor.isRequiredContent()).thenReturn(true); + SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), new RequestInterceptor[]{interceptor}); + Config.getInstance().putInConfigCache(config, Map.of("enabled", true, "maxBuffers", 4, + "maxBodyBytes", 1024, "appliedBodyInjectionPathPrefixes", List.of("/"))); + var exchange = mock(HttpServerExchange.class); + var connection = mock(ServerConnection.class); + var channel = mock(StreamSourceChannel.class); + ChannelListener.Setter setter = new ChannelListener.SimpleSetter<>(); + doReturn(setter).when(channel).getReadSetter(); + when(channel.read(any(java.nio.ByteBuffer.class))).thenReturn(0, -1); + var headers = new HeaderMap(); + when(exchange.getRequestHeaders()).thenReturn(headers); + when(exchange.getRequestMethod()).thenReturn(Methods.POST); + when(exchange.getRequestPath()).thenReturn("/"); + when(exchange.getRequestContentLength()).thenReturn(-1L); + when(exchange.getRequestChannel()).thenReturn(channel); + when(exchange.getConnection()).thenReturn(connection); + var failure = new IllegalStateException("outside resumeRead's body-read catch"); + var responseChecks = new AtomicInteger(); + when(exchange.isResponseStarted()).thenAnswer(invocation -> { + if (responseChecks.getAndIncrement() == 1) throw failure; // outside the body-read catch + return false; + }); + var selected = new AtomicReference(); + try (var pool = new DefaultByteBufferPool(false, 16)) { + when(connection.getByteBufferPool()).thenReturn(pool); + var reader = new RequestInterceptorInjectionHandler(config); + reader.setNext(ex -> fail("Failure must occur before downstream execution")); + HandlerContinuation.execute(exchange, next -> ex -> { + try { next.handleRequest(ex); } + catch (IllegalStateException caught) { selected.set(caught); } + }, reader); + ChannelListener listener = + ((ChannelListener.SimpleSetter) setter).get(); + assertNotNull(listener); + listener.handleEvent(channel); + assertSame(failure, selected.get(), "The listener itself must retain the lexical policy"); + verify(channel).suspendReads(); + verify(exchange, never()).setStatusCode(500); + } finally { + SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), previous); + Config.getInstance().clearConfigCache(config); + } + } + +} diff --git a/hmac/src/main/java/com/networknt/hmac/HmacHandler.java b/hmac/src/main/java/com/networknt/hmac/HmacHandler.java index 4d250974ab..c919a4c477 100644 --- a/hmac/src/main/java/com/networknt/hmac/HmacHandler.java +++ b/hmac/src/main/java/com/networknt/hmac/HmacHandler.java @@ -18,6 +18,7 @@ import com.networknt.config.ConfigException; import com.networknt.handler.Handler; +import com.networknt.handler.HandlerContinuation; import com.networknt.handler.HandlerChainValidator; import com.networknt.handler.MiddlewareHandler; import com.networknt.handler.RequestInterceptor; @@ -129,8 +130,9 @@ public void handleRequest(HttpServerExchange exchange) throws Exception { return; } + var continuation = HandlerContinuation.capture(exchange); exchange.dispatch(); - stage.whenComplete((outcome, failure) -> exchange.dispatch(completed -> { + stage.whenComplete((outcome, failure) -> exchange.dispatch(continuation.wrap(completed -> { if (failure != null || outcome == null) { unavailable(completed, rule, store, "reserve_failed"); } else if (outcome instanceof ReserveOutcome.Duplicate) { @@ -156,7 +158,7 @@ public void handleRequest(HttpServerExchange exchange) throws Exception { } else { unavailable(completed, rule, store, "unknown_reserve_outcome"); } - })); + }))); } private void addReservationCompletionListener(HttpServerExchange exchange, diff --git a/hmac/src/main/java/com/networknt/hmac/HmacReplayAdminHandler.java b/hmac/src/main/java/com/networknt/hmac/HmacReplayAdminHandler.java index 0612d02343..105f480e35 100644 --- a/hmac/src/main/java/com/networknt/hmac/HmacReplayAdminHandler.java +++ b/hmac/src/main/java/com/networknt/hmac/HmacReplayAdminHandler.java @@ -20,6 +20,7 @@ import com.networknt.config.Config; import com.networknt.handler.AuditAttachmentUtil; import com.networknt.handler.LightHttpHandler; +import com.networknt.handler.HandlerContinuation; import com.networknt.hmac.config.HmacConfig; import com.networknt.httpstring.AttachmentConstants; import com.networknt.security.UnifiedSecurityConfig; @@ -102,8 +103,9 @@ public void handleRequest(HttpServerExchange exchange) { setExchangeStatus(exchange, HmacVerifier.HMAC_UNAVAILABLE); return; } + var continuation = HandlerContinuation.capture(exchange); exchange.dispatch(); - removalStage.whenComplete((removed, failure) -> exchange.dispatch(completed -> { + removalStage.whenComplete((removed, failure) -> exchange.dispatch(continuation.wrap(completed -> { if (failure != null) { audit(completed, "unavailable", store); setExchangeStatus(completed, HmacVerifier.HMAC_UNAVAILABLE); @@ -120,7 +122,7 @@ public void handleRequest(HttpServerExchange exchange) { } catch (Exception e) { setExchangeStatus(completed, HmacVerifier.HMAC_UNAVAILABLE); } - })); + }))); } private Removal parse(Object requestBody) { diff --git a/hmac/src/test/java/com/networknt/hmac/HmacHandlerIntegrationTest.java b/hmac/src/test/java/com/networknt/hmac/HmacHandlerIntegrationTest.java index da53949dbb..c943b07218 100644 --- a/hmac/src/test/java/com/networknt/hmac/HmacHandlerIntegrationTest.java +++ b/hmac/src/test/java/com/networknt/hmac/HmacHandlerIntegrationTest.java @@ -59,6 +59,29 @@ import java.util.concurrent.atomic.AtomicReference; class HmacHandlerIntegrationTest { + private java.util.function.UnaryOperator chainWrapper = handler -> handler; + + @org.junit.jupiter.params.ParameterizedTest + @org.junit.jupiter.params.provider.ValueSource(booleans = {true, false}) + void replayCompletionRetainsExceptionBoundary(boolean enabled) throws Exception { + Context context = context("continuation-" + enabled); + if (enabled) chainWrapper = handler -> exchange -> + com.networknt.handler.HandlerContinuation.execute(exchange, task -> ex -> { + try { task.handleRequest(ex); } + catch (java.io.IOException _) { ex.setStatusCode(409).getResponseSender().send("selected"); } + }, handler); + try (RunningServer server = start(context, new RequestInterceptor[]{context.interceptor}, + exchange -> { throw new java.io.IOException("callee failed"); })) { + var response = server.post("{}".getBytes(StandardCharsets.UTF_8), "continuation"); + Assertions.assertTrue(context.store.reserveStarted.await(5, TimeUnit.SECONDS)); + Assertions.assertFalse(response.isDone()); + context.store.allowReservation.complete(null); + var completed = response.get(5, TimeUnit.SECONDS); + Assertions.assertEquals(enabled ? 409 : 500, completed.statusCode()); + Assertions.assertEquals(enabled, new String(completed.body(), StandardCharsets.UTF_8).contains("selected")); + } + } + private static final String SECRET = "phase-3-exact-forwarding-secret"; private static final String INJECTION_CONFIG = "hmac-phase3-request-injection"; private final List configNames = new java.util.ArrayList<>(); @@ -298,7 +321,7 @@ private RunningServer start(Context context, hmacHandler.setNext(callee); RequestInterceptorInjectionHandler injection = new RequestInterceptorInjectionHandler(INJECTION_CONFIG); injection.setNext(hmacHandler); - return new RunningServer(injection, context, version); + return new RunningServer(chainWrapper.apply(injection), context, version); } private Context context(String suffix) { diff --git a/ingress-proxy/pom.xml b/ingress-proxy/pom.xml index 5de587361a..1677ea9c23 100644 --- a/ingress-proxy/pom.xml +++ b/ingress-proxy/pom.xml @@ -19,6 +19,7 @@ com.networknt exception + test com.networknt diff --git a/ingress-proxy/src/main/java/com/networknt/proxy/ExternalServiceHandler.java b/ingress-proxy/src/main/java/com/networknt/proxy/ExternalServiceHandler.java index 61729db2f6..08e7865acf 100644 --- a/ingress-proxy/src/main/java/com/networknt/proxy/ExternalServiceHandler.java +++ b/ingress-proxy/src/main/java/com/networknt/proxy/ExternalServiceHandler.java @@ -4,7 +4,6 @@ import com.networknt.client.Http2Client; import com.networknt.client.ssl.TLSConfig; import com.networknt.config.Config; -import com.networknt.exception.ExceptionHandler; import com.networknt.handler.*; import com.networknt.handler.config.UrlRewriteRule; import com.networknt.httpstring.AttachmentConstants; @@ -133,9 +132,8 @@ public void handleRequest(HttpServerExchange exchange) throws Exception { requestHost, currentTimeout, connectTimeout, startTime); // Capture the match once: asynchronous body reads can resume on an I/O thread. if (exchange.isInIoThread()) { - exchange.dispatch(workerExchange -> new ExceptionHandler() - .setNext(forwardExchange -> forwardMatchedRequest(forwardExchange, context)) - .handleRequest(workerExchange)); + exchange.dispatch(HandlerContinuation.capture(exchange, + forwardExchange -> forwardMatchedRequest(forwardExchange, context))); } else { forwardMatchedRequest(exchange, context); } @@ -229,7 +227,10 @@ private HttpResponse sendWithRetries(HttpClient forwardingClient, HttpRe HttpRequest finalRequest = attempt == 0 ? request : buildRetryRequest(request, attempt); try { return forwardingClient.send(finalRequest, HttpResponse.BodyHandlers.ofByteArray()); - } catch (IOException | InterruptedException e) { + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } catch (IOException e) { if (attempt >= maxRetries - 1) throw e; logger.warn("Attempt {} failed ({}). Retrying...", attempt + 1, e.getMessage()); } diff --git a/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceInterruptTest.java b/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceInterruptTest.java new file mode 100644 index 0000000000..906cf30cd8 --- /dev/null +++ b/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceInterruptTest.java @@ -0,0 +1,66 @@ +package com.networknt.proxy; + +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import java.io.IOException; +import java.lang.reflect.InvocationTargetException; +import java.net.*; +import java.net.http.*; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.*; +import javax.net.ssl.*; +import static org.junit.jupiter.api.Assertions.*; + +class ExternalServiceInterruptTest { + @ParameterizedTest + @CsvSource({"1,1", "3,1", "3,3"}) + void interruptionStopsRetriesAndPreservesFlag(int attempts, int interruptAt) throws Exception { + var supported = ExternalServiceHandler.class.getDeclaredField("connectionCloseHeaderSupported"); + var logged = ExternalServiceHandler.class.getDeclaredField("connectionCloseHeaderFallbackLogged"); + supported.setAccessible(true); + logged.setAccessible(true); + Object previousSupported = supported.get(null); + var loggedFlag = (java.util.concurrent.atomic.AtomicBoolean) logged.get(null); + boolean previousLogged = loggedFlag.get(); + InterruptingClient client = new InterruptingClient(interruptAt); + var method = ExternalServiceHandler.class.getDeclaredMethod("sendWithRetries", HttpClient.class, HttpRequest.class, int.class); + method.setAccessible(true); + try { + var thrown = assertThrows(InvocationTargetException.class, () -> method.invoke( + new ExternalServiceHandler(), client, HttpRequest.newBuilder(URI.create("http://localhost/")).build(), attempts)); + assertInstanceOf(InterruptedException.class, thrown.getCause()); + assertEquals(interruptAt, client.calls); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + supported.set(null, previousSupported); + loggedFlag.set(previousLogged); + } + } + + private static final class InterruptingClient extends HttpClient { + private final int interruptAt; + private int calls; + InterruptingClient(int interruptAt) { this.interruptAt = interruptAt; } + @Override public HttpResponse send(HttpRequest request, HttpResponse.BodyHandler handler) + throws IOException, InterruptedException { + if (++calls == interruptAt) throw new InterruptedException("cancelled"); + throw new IOException("retryable"); + } + @Override public Optional cookieHandler() { return Optional.empty(); } + @Override public Optional connectTimeout() { return Optional.empty(); } + @Override public Redirect followRedirects() { return Redirect.NEVER; } + @Override public Optional proxy() { return Optional.empty(); } + @Override public SSLContext sslContext() { throw new UnsupportedOperationException(); } + @Override public SSLParameters sslParameters() { return new SSLParameters(); } + @Override public Optional authenticator() { return Optional.empty(); } + @Override public Version version() { return Version.HTTP_1_1; } + @Override public Optional executor() { return Optional.empty(); } + @Override public CompletableFuture> sendAsync(HttpRequest request, HttpResponse.BodyHandler handler) { + throw new UnsupportedOperationException(); + } + @Override public CompletableFuture> sendAsync(HttpRequest request, HttpResponse.BodyHandler handler, + HttpResponse.PushPromiseHandler push) { throw new UnsupportedOperationException(); } + } +} diff --git a/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceWorkerDispatchTest.java b/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceWorkerDispatchTest.java index aef4e2addf..fac39a771b 100644 --- a/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceWorkerDispatchTest.java +++ b/ingress-proxy/src/test/java/com/networknt/proxy/ExternalServiceWorkerDispatchTest.java @@ -22,6 +22,11 @@ import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.Arguments; +import io.undertow.util.HttpString; +import java.util.stream.Stream; import java.io.ByteArrayOutputStream; import java.net.InetSocketAddress; @@ -55,9 +60,44 @@ class ExternalServiceWorkerDispatchTest { private RequestInterceptor[] previousInterceptors; private String previousTlsProperty; private HttpHandler previousMetrics; + private Map previousExceptionConfig; + enum Policy { + ENABLED, NAMED, DISABLED, NAMED_DISABLED; + + String configName() { return this == ENABLED ? "exception" : "external-worker-exception-" + name(); } + boolean enabled() { return this == ENABLED || this == NAMED; } + } + + static Stream streamedPolicies() { + return Stream.of(Policy.values()).flatMap(policy -> + Stream.of(false, true).map(failure -> Arguments.of(failure, policy))); + } + + private HttpHandler withPolicy(Policy policy, HttpHandler next) { + ExceptionHandler selected = new ExceptionHandler(policy.configName()) { + @Override + public void setExchangeStatus(HttpServerExchange exchange, String code, Object... args) { + exchange.getResponseHeaders().put(new HttpString("X-Exception-Policy"), policy.name()); + super.setExchangeStatus(exchange, code, args); + } + }; + selected.setNext(next); + // Match handler-chain materialization: disabled middleware is omitted. + return selected.isEnabled() ? selected : next; + } + + private static void assertFailure(Policy policy, int status, String body, String selected) { + boolean enabled = policy == Policy.ENABLED || policy == Policy.NAMED; + assertEquals(enabled ? 400 : 500, status, body); + assertEquals(enabled ? policy.name() : "", selected, "Use the selected instance only"); + assertEquals(enabled, body.contains("ERR10011"), body); + } @BeforeEach void rememberSharedState() { + previousExceptionConfig = Config.getInstance().getJsonMapConfig("exception"); + for (Policy policy : Policy.values()) + Config.getInstance().putInConfigCache(policy.configName(), Map.of("enabled", policy.enabled())); previousInterceptors = SingletonServiceFactory.getBeans(RequestInterceptor.class); previousTlsProperty = System.getProperty(TLS_PROPERTY); previousMetrics = Handler.getHandlers().get(MetricsConfig.CONFIG_NAME); @@ -68,6 +108,10 @@ void restoreSharedState() { SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), previousInterceptors); if (previousMetrics == null) Handler.getHandlers().remove(MetricsConfig.CONFIG_NAME); else Handler.getHandlers().put(MetricsConfig.CONFIG_NAME, previousMetrics); + if (previousExceptionConfig == null) Config.getInstance().clearConfigCache("exception"); + else Config.getInstance().putInConfigCache("exception", previousExceptionConfig); + for (Policy policy : Policy.values()) + if (policy != Policy.ENABLED) Config.getInstance().clearConfigCache(policy.configName()); Config.getInstance().clearConfigCache(EXTERNAL_CONFIG); Config.getInstance().clearConfigCache(INJECTION_CONFIG); if (previousTlsProperty == null) System.clearProperty(TLS_PROPERTY); @@ -75,8 +119,8 @@ void restoreSharedState() { } @ParameterizedTest - @ValueSource(booleans = {false, true}) - void streamedBodyContinuationCannotBlockTheFrontEndIoThread(boolean upstreamFailure) throws Exception { + @MethodSource("streamedPolicies") + void streamedBodyContinuationCannotBlockTheFrontEndIoThread(boolean upstreamFailure, Policy policy) throws Exception { BodyCaptureInterceptor interceptor = new BodyCaptureInterceptor(); SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), new RequestInterceptor[]{interceptor}); @@ -89,7 +133,7 @@ void streamedBodyContinuationCannotBlockTheFrontEndIoThread(boolean upstreamFail RequestInterceptorInjectionHandler reader = new RequestInterceptorInjectionHandler(INJECTION_CONFIG); reader.setNext(injection); - return new ExceptionHandler().setNext(exchange -> { + return withPolicy(policy, exchange -> { reader.handleRequest(exchange); initialReadReturned.countDown(); }); @@ -116,9 +160,7 @@ void streamedBodyContinuationCannotBlockTheFrontEndIoThread(boolean upstreamFail fixture.releaseUpstream.countDown(); String response = new String(socket.getInputStream().readAllBytes(), StandardCharsets.UTF_8); if (upstreamFailure) { - assertTrue(response.startsWith("HTTP/1.1 400"), response); - assertTrue(response.contains("application/json"), response); - assertTrue(response.contains("ERR10011"), "Use framework structured errors: " + response); + assertRawFailure(policy, response); } else { assertTrue(response.startsWith("HTTP/1.1 200"), response); assertTrue(response.endsWith("accepted"), response); @@ -171,18 +213,71 @@ void unmatchedRequestRetainsItsOriginalIoThread() throws Exception { } } - @Test - void workerThreadFailureUsesTheSameFrameworkError() throws Exception { - try (ProxyFixture fixture = new ProxyFixture(handler -> new ExceptionHandler().setNext(handler)); + @ParameterizedTest + @EnumSource(Policy.class) + void workerThreadFailureUsesTheSelectedPolicy(Policy policy) throws Exception { + try (ProxyFixture fixture = new ProxyFixture(handler -> new BlockingHandler(withPolicy(policy, handler))); HttpClient client = HttpClient.newHttpClient()) { fixture.failUpstream.set(true); fixture.releaseUpstream.countDown(); HttpResponse response = client.send( HttpRequest.newBuilder(fixture.uri("/proxy/failure")) .timeout(Duration.ofSeconds(5)).build(), HttpResponse.BodyHandlers.ofString()); - assertEquals(400, response.statusCode(), "Preserve the framework ERR10011 status"); - assertTrue(response.headers().firstValue("Content-Type").orElse("").contains("application/json")); - assertTrue(response.body().contains("ERR10011"), response.body()); + assertFailure(policy, response.statusCode(), response.body(), + response.headers().firstValue("X-Exception-Policy").orElse("")); + } + } + + private static void assertRawFailure(Policy policy, String response) { + assertTrue(response.matches("(?s)HTTP/1\\.[01] [0-9]{3} .*"), + "Expected a complete HTTP status line, received: " + response); + int status = Integer.parseInt(response.substring(9, 12)); + String selected = response.lines().filter(line -> line.toLowerCase().startsWith("x-exception-policy:")) + .map(line -> line.substring(line.indexOf(':') + 1).trim()).findFirst().orElse(""); + assertFailure(policy, status, response, selected); + if (status == 400) assertTrue(response.contains("application/json"), response); + } + + @Test + void concurrentStreamedRequestsKeepTheirOwnBoundary() throws Exception { + BodyCaptureInterceptor interceptor = new BodyCaptureInterceptor(); + SingletonServiceFactory.setBean(RequestInterceptor.class.getName(), new RequestInterceptor[]{interceptor}); + Config.getInstance().putInConfigCache(INJECTION_CONFIG, Map.of( + "enabled", true, "maxBuffers", 4, "maxBodyBytes", 1024, + "appliedBodyInjectionPathPrefixes", List.of("/proxy"))); + int count = 8; + CountDownLatch initialReads = new CountDownLatch(count); + try (ProxyFixture fixture = new ProxyFixture(handler -> { + var chains = new java.util.EnumMap(Policy.class); + for (Policy policy : Policy.values()) { + RequestInterceptorInjectionHandler reader = new RequestInterceptorInjectionHandler(INJECTION_CONFIG); + reader.setNext(handler); + chains.put(policy, withPolicy(policy, exchange -> { + reader.handleRequest(exchange); + initialReads.countDown(); + })); + } + return exchange -> chains.get(Policy.valueOf(exchange.getRequestHeaders().getFirst("X-Policy"))) + .handleRequest(exchange); + }); var executor = Executors.newFixedThreadPool(count)) { + fixture.failUpstream.set(true); + fixture.releaseUpstream.countDown(); + var futures = IntStream.range(0, count).mapToObj(i -> executor.submit(() -> { + Policy policy = Policy.values()[i % Policy.values().length]; + try (Socket socket = new Socket("localhost", fixture.port())) { + socket.setSoTimeout(10000); + write(socket, "POST /proxy/concurrent HTTP/1.1\r\nHost: localhost\r\n" + + "X-Policy: " + policy + "\r\nContent-Type: application/json\r\n" + + "Transfer-Encoding: chunked\r\nConnection: close\r\n\r\n1\r\n{\r\n"); + assertTrue(initialReads.await(5, TimeUnit.SECONDS), "All original frames must return before resuming"); + write(socket, "1\r\n}\r\n0\r\n\r\n"); + assertRawFailure(policy, new String(socket.getInputStream().readAllBytes(), StandardCharsets.UTF_8)); + } + return null; + })).toList(); + for (var future : futures) future.get(15, TimeUnit.SECONDS); + assertEquals(count, fixture.handlerEntries.get()); + assertTrue(interceptor.onIoThread.get()); } } diff --git a/mcp-router/src/main/java/com/networknt/mcp/McpHandler.java b/mcp-router/src/main/java/com/networknt/mcp/McpHandler.java index 258423c41d..0c51c8472c 100644 --- a/mcp-router/src/main/java/com/networknt/mcp/McpHandler.java +++ b/mcp-router/src/main/java/com/networknt/mcp/McpHandler.java @@ -6,6 +6,7 @@ import com.networknt.config.Config; import com.networknt.config.JsonMapper; import com.networknt.handler.Handler; +import com.networknt.handler.HandlerContinuation; import com.networknt.handler.MiddlewareHandler; import com.networknt.httpstring.AttachmentConstants; import com.networknt.rule.RuleConstants; @@ -173,7 +174,7 @@ private void handleSse(HttpServerExchange exchange, McpConfig config) throws Exc private void handleMessage(HttpServerExchange exchange, McpConfig config) { if (exchange.isInIoThread()) { - exchange.dispatch(this); + exchange.dispatch(HandlerContinuation.capture(exchange, this)); return; } diff --git a/mcp-router/src/test/java/com/networknt/mcp/McpContinuationTest.java b/mcp-router/src/test/java/com/networknt/mcp/McpContinuationTest.java new file mode 100644 index 0000000000..568aa53d60 --- /dev/null +++ b/mcp-router/src/test/java/com/networknt/mcp/McpContinuationTest.java @@ -0,0 +1,39 @@ +package com.networknt.mcp; + +import com.networknt.handler.HandlerContinuation; +import io.undertow.Undertow; +import io.undertow.server.*; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.*; +import java.time.Duration; +import static org.junit.jupiter.api.Assertions.*; + +class McpContinuationTest { + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void redispatchRetainsSelectedBoundary(boolean enabled) throws Exception { + McpHandler handler = new McpHandler() { + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { + if (!exchange.isInIoThread()) throw new IllegalStateException("resumed failure"); + super.handleRequest(exchange); + } + }; + HttpHandler root = enabled ? exchange -> HandlerContinuation.execute(exchange, task -> ex -> { + try { task.handleRequest(ex); } + catch (IllegalStateException _) { ex.setStatusCode(409).getResponseSender().send("selected"); } + }, handler) : handler; + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setHandler(root).build(); + try (HttpClient client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/mcp")) + .timeout(Duration.ofSeconds(5)).POST(HttpRequest.BodyPublishers.ofString("{}")) + .build(), HttpResponse.BodyHandlers.ofString()); + assertEquals(enabled ? 409 : 500, response.statusCode(), response.body()); + assertEquals(enabled, response.body().contains("selected")); + } finally { server.stop(); } + } +} diff --git a/proxy-handler/src/main/java/com/networknt/handler/ProxyCallbackDispatcher.java b/proxy-handler/src/main/java/com/networknt/handler/ProxyCallbackDispatcher.java new file mode 100644 index 0000000000..c812a1ac04 --- /dev/null +++ b/proxy-handler/src/main/java/com/networknt/handler/ProxyCallbackDispatcher.java @@ -0,0 +1,103 @@ +package com.networknt.handler; + +import io.undertow.server.HttpHandler; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Executor; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; +import java.util.function.Consumer; + +/** Proxy-local FIFO. Transport callbacks and ownerFailure run only on the owning I/O thread. */ +final class ProxyCallbackDispatcher { + private static final Logger LOG = LoggerFactory.getLogger(ProxyCallbackDispatcher.class); + private static final Runnable NO_RESOURCE = () -> { }; + private record Event(HttpHandler callback, Runnable discardedResource, boolean accepted) { } + + private final HandlerContinuation.Scope scope; + private final Executor executor; + private final BooleanSupplier inOwner; + private final Runnable ownerFailure; + private final Consumer rejectedExecutor; + private final ConcurrentLinkedQueue queue = new ConcurrentLinkedQueue<>(); + private final AtomicInteger work = new AtomicInteger(); + private final AtomicReference terminal = new AtomicReference<>(); + private final Runnable drainTask = this::drain; + + ProxyCallbackDispatcher(HandlerContinuation.Scope scope, Executor executor, BooleanSupplier inOwner, + Runnable ownerFailure, Consumer rejectedExecutor) { + this.scope = scope; + this.executor = executor; + this.inOwner = inOwner; + this.ownerFailure = ownerFailure; + this.rejectedExecutor = rejectedExecutor; + } + + void execute(HttpHandler callback) { execute(callback, NO_RESOURCE); } + + /** discardedResource may run on any submitting thread: close only this incoming event's resource. */ + void execute(HttpHandler callback, Runnable discardedResource) { submit(callback, discardedResource, false); } + + void resume(HttpHandler callback, Runnable discardedResource) { submit(callback, discardedResource, true); } + + private void submit(HttpHandler callback, Runnable discardedResource, boolean accepted) { + queue.add(new Event(callback, discardedResource, accepted)); + Throwable failure = terminal.get(); + if (failure != null) { discardQueued(failure); return; } + if (work.getAndIncrement() != 0) return; + if (inOwner.getAsBoolean()) drain(); + else { + try { executor.execute(drainTask); } + catch (RejectedExecutionException rejected) { + if (terminal.compareAndSet(null, rejected)) { + discardQueued(rejected); + // Exactly the rejecting submission owns this terminal notification. + // It must use a published resource snapshot, never mutable proxy fields. + rejectedExecutor.accept(rejected); + } + } + } + } + + private void drain() { + int missed = 1; + try { + do { + Event event; + while ((event = queue.poll()) != null) { + if (event.accepted()) scope.resume(event.callback()); + else scope.execute(event.callback()); + } + missed = work.addAndGet(-missed); + } while (missed != 0); + } catch (Error failure) { + // This path is still on the owner. Stop before retiring queued resources. + cleanup(ownerFailure, failure); + terminal.compareAndSet(null, failure); + discardQueued(failure); + throw failure; + } + } + + private void discardQueued(Throwable failure) { + Event event; + while ((event = queue.poll()) != null) cleanup(event.discardedResource(), failure); + } + + static void cleanup(Runnable cleanup, Throwable primary) { + try { cleanup.run(); } + catch (Throwable failure) { + if (primary == failure) return; + try { + primary.addSuppressed(failure); + LOG.error("Proxy callback resource cleanup failed", failure); + } catch (Throwable _) { + // Best-effort diagnostics must not mask the original terminal failure. + } + } + } +} diff --git a/proxy-handler/src/main/java/com/networknt/handler/ProxyHandler.java b/proxy-handler/src/main/java/com/networknt/handler/ProxyHandler.java index 99aa28d8c0..2c5e2db65c 100644 --- a/proxy-handler/src/main/java/com/networknt/handler/ProxyHandler.java +++ b/proxy-handler/src/main/java/com/networknt/handler/ProxyHandler.java @@ -25,6 +25,7 @@ import java.nio.channels.FileChannel; import java.security.cert.Certificate; import java.util.*; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import javax.net.ssl.SSLPeerUnverifiedException; @@ -33,7 +34,6 @@ import com.networknt.handler.config.MethodRewriteRule; import com.networknt.handler.config.QueryHeaderRewriteRule; import com.networknt.handler.config.UrlRewriteRule; -import com.networknt.handler.thread.LightThreadExecutor; import com.networknt.httpstring.AttachmentConstants; import com.networknt.httpstring.HttpStringConstants; import com.networknt.metrics.AbstractMetricsHandler; @@ -74,6 +74,7 @@ import io.undertow.server.protocol.http.HttpAttachments; import io.undertow.server.protocol.http.HttpContinue; +import java.util.function.Consumer; import java.util.stream.Collectors; /** @@ -103,11 +104,10 @@ public class ProxyHandler implements HttpHandler { private static final Logger LOG = LoggerFactory.getLogger(ProxyHandler.class); - private LightThreadExecutor lightThreadExecutor; + private static final Executor INLINE_EXECUTOR = Runnable::run; public static final AttachmentKey CONNECTION = AttachmentKey.create(ProxyConnection.class); private static final AttachmentKey EXCHANGE = AttachmentKey.create(HttpServerExchange.class); - private static final AttachmentKey TIMEOUT_KEY = AttachmentKey.create(XnioExecutor.Key.class); private final ProxyClient proxyClient; private int maxRequestTime; @@ -164,7 +164,6 @@ private ProxyHandler(Builder builder) { } public void handleRequest(final HttpServerExchange exchange) throws Exception { - this.lightThreadExecutor = new LightThreadExecutor(exchange); final ProxyClient.ProxyTarget target = proxyClient.findTarget(exchange); if (target == null) { @@ -198,8 +197,8 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception { if (timeout > 0) { LOG.trace("Adding timeout to exchange of {} ms to exchange: {}", effectiveMaxRequestTime, exchange); - final XnioExecutor.Key key = WorkerUtils.executeAfter(exchange.getIoThread(), () -> clientHandler.cancel(exchange), effectiveMaxRequestTime, TimeUnit.MILLISECONDS); - exchange.putAttachment(TIMEOUT_KEY, key); + final XnioExecutor.Key key = WorkerUtils.executeAfter(exchange.getIoThread(), () -> clientHandler.cancel(), effectiveMaxRequestTime, TimeUnit.MILLISECONDS); + clientHandler.continuation.timeoutKey = key; exchange.addExchangeCompleteListener((exchange1, nextListener) -> { LOG.trace("Exchange completed for exchange removing timeout key: {}", exchange1); key.remove(); @@ -208,8 +207,7 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception { } - if (exchange.isInIoThread()) exchange.dispatch(lightThreadExecutor, clientHandler); - + if (exchange.isInIoThread()) exchange.dispatch(INLINE_EXECUTOR, clientHandler); else exchange.dispatch(exchange.getIoThread(), clientHandler); } @@ -349,14 +347,6 @@ private static void removeHeaders(HeaderMap headers, List headerName } } - private static void cancelExchangeTimeout(HttpServerExchange exchange) { - XnioExecutor.Key key = exchange.getAttachment(TIMEOUT_KEY); - if (key != null) { - key.remove(); - exchange.removeAttachment(TIMEOUT_KEY); - } - } - @Override public String toString() { var proxyTargets = proxyClient.getAllTargets(); @@ -377,7 +367,101 @@ public String toString() { } } - private final class ProxyClientHandler implements ProxyCallback, Runnable { + /** Mutable transport state is confined to the downstream exchange's I/O thread. */ + private static final class ProxyContinuation { + private static final Runnable NO_REJECTED_RESOURCE = () -> { }; + private final HandlerContinuation.Scope scope; + private final ProxyCallbackDispatcher dispatcher; + private ProxyConnection connection; + private ClientExchange request; + private boolean stopping; + private volatile XnioExecutor.Key timeoutKey; + private long attempt; + private volatile UpstreamResources publishedResources = UpstreamResources.EMPTY; + + private ProxyContinuation(HttpServerExchange exchange, HandlerContinuation.Scope scope) { + this.scope = scope; + this.dispatcher = new ProxyCallbackDispatcher(scope, exchange.getIoThread(), + () -> exchange.isInIoThread(), () -> abortUpstream(this), failure -> { + ProxyCallbackDispatcher.cleanup(publishedResources::close, failure); + scope.execute(ex -> { throw failure; }); + }); + } + + private void run(Consumer callback) { run(callback, NO_REJECTED_RESOURCE); } + + private void run(Consumer callback, Runnable rejectedCleanup) { + dispatcher.execute(ex -> acceptedCallback(ex, callback, rejectedCleanup), rejectedCleanup); + } + + private void acceptedCallback(HttpServerExchange exchange, Consumer callback, + Runnable rejectedCleanup) { + if (!exchange.isInIoThread()) { + dispatcher.resume(ex -> acceptedCallback(ex, callback, rejectedCleanup), rejectedCleanup); + return; + } + boolean returned = false; + try { callback.accept(exchange); returned = true; } + finally { if (!returned) abortUpstream(this); } + } + + private void publishResources() { + publishedResources = new UpstreamResources(connection, request); + } + + private void cancelTimeout() { + XnioExecutor.Key key = timeoutKey; + if (key != null) key.remove(); + } + } + + /** Immutable shutdown lease; closing XNIO channels is safe from a foreign thread. */ + private record UpstreamResources(ProxyConnection connection, ClientExchange request) { + private static final UpstreamResources EMPTY = new UpstreamResources(null, null); + private void close() { + if (request != null) cancelClientExchange(request); + if (connection != null && (request == null || request.getConnection() != connection.getConnection()) + && !connection.getConnection().isMultiplexingSupported()) IoUtils.safeClose(connection.getConnection()); + } + } + + private static void cancelClientExchange(ClientExchange request) { + if (request.getConnection().isMultiplexingSupported()) { + // Close only this request's channels, never its shared HTTP/2 connection. + IoUtils.safeClose(request.getRequestChannel(), request.getResponseChannel()); + } else IoUtils.safeClose(request.getConnection()); + } + + private static void releaseAttempt(ProxyContinuation state) { + ClientExchange request = state.request; + ProxyConnection connection = state.connection; + state.request = null; + state.connection = null; + state.publishedResources = UpstreamResources.EMPTY; + if (request != null) cancelClientExchange(request); + // A stream can belong to the preceding pooled connection. Always release + // a distinct current HTTP/1 connection as well; never close an HTTP/2 parent. + if (connection != null && (request == null || request.getConnection() != connection.getConnection()) + && !connection.getConnection().isMultiplexingSupported()) IoUtils.safeClose(connection.getConnection()); + } + + private static void abortUpstream(ProxyContinuation state) { + state.stopping = true; + releaseAttempt(state); + } + + /** Ignore callbacks owned by a terminal response; terminate only unexpected commitment. */ + private static boolean stopCallbackIfTerminal(HttpServerExchange exchange, ProxyContinuation state) { + if (exchange.isComplete() || state.stopping) return true; + if (!exchange.isResponseStarted()) return false; + LOG.warn("Unexpected proxy callback after response commitment for {}; terminating the stream", exchange.getRequestPath()); + abortUpstream(state); + IoUtils.safeClose(exchange.getConnection()); + return true; + } + + private final class ProxyClientHandler implements Runnable { + private final ProxyContinuation continuation; private int tries; private final long timeout; private final int maxRetryAttempts; @@ -388,6 +472,7 @@ private final class ProxyClientHandler implements ProxyCallback ProxyClientHandler(HttpServerExchange exchange, ProxyClient.ProxyTarget target, long timeout, int maxRetryAttempts, Predicate idempotentPredicate, boolean streamingRequest) { this.exchange = exchange; + this.continuation = new ProxyContinuation(exchange, HandlerContinuation.capture(exchange)); this.timeout = timeout; this.maxRetryAttempts = maxRetryAttempts; this.target = target; @@ -397,57 +482,81 @@ private final class ProxyClientHandler implements ProxyCallback @Override public void run() { - proxyClient.getConnection(this.target, this.exchange, this, -1, TimeUnit.MILLISECONDS); + continuation.run(ex -> { + if (!stopCallbackIfTerminal(ex, continuation)) + acquire(this.target, -1); + }); } - @Override - public void completed(final HttpServerExchange exchange, final ProxyConnection connection) { - LOG.trace("ProxyCallback completed for exchange: {}", exchange); - exchange.putAttachment(CONNECTION, connection); - exchange.dispatch(lightThreadExecutor, new ProxyAction(connection, exchange, requestHeaders, rewriteHostHeader, reuseXForwarded, exchange.isRequestComplete() ? this : null, idempotentPredicate, urlRewriteRules, methodRewriteRules, queryParamRewriteRules, headerRewriteRules, streamingRequest, streamResponseContentTypes, streamIdleTimeout, streamResponseHeaderOverwrite)); + private void acquire(ProxyClient.ProxyTarget target, long remaining) { + // A retry owns new resources. Drop and cancel the previous stream before + // accepting any resources from this attempt. + long attempt = ++continuation.attempt; + releaseAttempt(continuation); + proxyClient.getConnection(target, exchange, new ProxyCallback<>() { + @Override public void completed(HttpServerExchange ex, ProxyConnection connection) { + continuation.run(resumed -> { + if (attempt != continuation.attempt) { + if (!connection.getConnection().isMultiplexingSupported()) IoUtils.safeClose(connection.getConnection()); + } else ProxyClientHandler.this.acceptConnection(resumed, connection); + }, () -> { if (!connection.getConnection().isMultiplexingSupported()) IoUtils.safeClose(connection.getConnection()); }); + } + @Override public void failed(HttpServerExchange ex) { failedAttempt(attempt); } + @Override public void queuedRequestFailed(HttpServerExchange ex) { failedAttempt(attempt); } + @Override public void couldNotResolveBackend(HttpServerExchange ex) { + continuation.run(resumed -> { + if (attempt == continuation.attempt) unresolvedBackend(resumed); + }); + } + }, remaining, TimeUnit.MILLISECONDS); } - @Override - public void failed(final HttpServerExchange exchange) { - LOG.debug("Failed calling backend with tries = {} maxRetryAttempts = {} for exchange: {}", this.tries, this.maxRetryAttempts, exchange); - - final long time = System.currentTimeMillis(); - - if (this.tries++ < this.maxRetryAttempts) { - - if (this.timeout > 0 && time > this.timeout) { - LOG.trace("Current time = {} passes timeout {} for exchange: {}", time, this.timeout, exchange); - cancel(exchange); - - } else { - this.target = proxyClient.findTarget(exchange); - LOG.trace("Retry target = {}", target); - - if (this.target != null) { - final long remaining = this.timeout > 0 ? this.timeout - time : -1; - LOG.trace("Retry with remaining = {} for exchange: {}", remaining, exchange); - proxyClient.getConnection(target, exchange, this, remaining, TimeUnit.MILLISECONDS); - - } else { - LOG.trace("Target is null, cannot resolve the backend"); - couldNotResolveBackend(exchange); // The context was registered when we started, so return 503 - } - } + private void failedAttempt(long attempt) { + continuation.run(resumed -> { + if (attempt == continuation.attempt) handleFailed(resumed); + }); + } - } else { - LOG.trace("Max number fo retry attempts reached for exchange: {}", exchange); - couldNotResolveBackend(exchange); + private void acceptConnection(final HttpServerExchange ex, final ProxyConnection connection) { + if (stopCallbackIfTerminal(ex, continuation)) { + if (!connection.getConnection().isMultiplexingSupported()) IoUtils.safeClose(connection.getConnection()); + return; } + LOG.trace("ProxyCallback completed for exchange: {}", ex); + continuation.connection = connection; + continuation.publishResources(); + ex.putAttachment(CONNECTION, connection); + new ProxyAction(connection, ex, requestHeaders, rewriteHostHeader, reuseXForwarded, ex.isRequestComplete() ? this : null, idempotentPredicate, urlRewriteRules, methodRewriteRules, queryParamRewriteRules, headerRewriteRules, streamingRequest, streamResponseContentTypes, streamIdleTimeout, streamResponseHeaderOverwrite, continuation).run(); } - @Override - public void queuedRequestFailed(HttpServerExchange exchange) { - failed(exchange); + private void failed() { + continuation.run(this::handleFailed); } - @Override - public void couldNotResolveBackend(HttpServerExchange ex) { + private void handleFailed(final HttpServerExchange exchange) { + if (stopCallbackIfTerminal(exchange, continuation)) return; + LOG.debug("Failed calling backend with tries = {} maxRetryAttempts = {} for exchange: {}", tries, maxRetryAttempts, exchange); + if (tries++ >= maxRetryAttempts) { + unresolvedBackend(exchange); + return; + } + long time = System.currentTimeMillis(); + if (timeout > 0 && time > timeout) { + cancelRequest(exchange); + return; + } + target = proxyClient.findTarget(exchange); + if (target == null) { + unresolvedBackend(exchange); + return; + } + long remaining = timeout > 0 ? timeout - time : -1; + acquire(target, remaining); + } + private void unresolvedBackend(HttpServerExchange ex) { + if (ex.isComplete() || continuation.stopping) return; + abortUpstream(continuation); if (ex.isResponseStarted()) IoUtils.safeClose(ex.getConnection()); @@ -466,15 +575,19 @@ public void couldNotResolveBackend(HttpServerExchange ex) { } } - void cancel(final HttpServerExchange exchange) { + void cancel() { + continuation.run(this::cancelRequest); + } + private void cancelRequest(final HttpServerExchange exchange) { + if (exchange.isComplete() || continuation.stopping) return; + final ProxyConnection connectionAttachment = continuation.connection; + abortUpstream(continuation); //NOTE: this method is called only in context of timeouts. - final ProxyConnection connectionAttachment = exchange.getAttachment(CONNECTION); if (connectionAttachment != null) { ClientConnection clientConnection = connectionAttachment.getConnection(); UndertowLogger.PROXY_REQUEST_LOGGER.timingOutRequest(clientConnection.getPeerAddress() + "" + exchange.getRequestURI()); - IoUtils.safeClose(clientConnection); } else UndertowLogger.PROXY_REQUEST_LOGGER.timingOutRequest(exchange.getRequestURI()); @@ -499,9 +612,10 @@ void cancel(final HttpServerExchange exchange) { } private static class ProxyAction implements Runnable { + private final ProxyContinuation continuation; private final ProxyConnection clientConnection; + private final long attempt; private final HttpServerExchange exchange; - private final LightThreadExecutor lightThreadExecutor; private final Map requestHeaders; private final boolean rewriteHostHeader; private final boolean reuseXForwarded; @@ -520,7 +634,10 @@ private static class ProxyAction implements Runnable { boolean rewriteHostHeader, boolean reuseXForwarded, ProxyClientHandler proxyClientHandler, Predicate idempotentPredicate, List urlRewriteRules, List methodRewriteRules, Map> queryParamRewriteRules, Map> headerRewriteRules, - boolean streamingRequest, List streamResponseContentTypes, int streamIdleTimeout, List streamResponseHeaderOverwrite) { + boolean streamingRequest, List streamResponseContentTypes, int streamIdleTimeout, List streamResponseHeaderOverwrite, + ProxyContinuation continuation) { + this.continuation = continuation; + this.attempt = continuation.attempt; this.clientConnection = clientConnection; this.exchange = exchange; this.requestHeaders = requestHeaders; @@ -536,7 +653,6 @@ private static class ProxyAction implements Runnable { this.streamResponseContentTypes = streamResponseContentTypes; this.streamIdleTimeout = streamIdleTimeout; this.streamResponseHeaderOverwrite = streamResponseHeaderOverwrite; - this.lightThreadExecutor = new LightThreadExecutor(exchange); } @Override @@ -890,8 +1006,29 @@ private void attachRemoteHost(ClientRequest r) { } } + private ClientCallback boundaryCallback(ClientCallback callback) { + return new ClientCallback<>() { + @Override public void completed(ClientExchange result) { + continuation.run(ex -> { + if (ex.isComplete() || continuation.stopping || attempt != continuation.attempt) { + cancelClientExchange(result); + return; + } + continuation.request = result; + continuation.publishResources(); + callback.completed(result); + }, () -> cancelClientExchange(result)); + } + @Override public void failed(IOException failure) { + continuation.run(ex -> { + if (!ex.isComplete() && !continuation.stopping && attempt == continuation.attempt) callback.failed(failure); + }); + } + }; + } + private void sendWithCallback(ClientRequest r, String host) { - this.clientConnection.getConnection().sendRequest(r, new ClientCallback<>() { + this.clientConnection.getConnection().sendRequest(r, boundaryCallback(new ClientCallback<>() { @Override public void completed(final ClientExchange result) { @@ -908,16 +1045,16 @@ public void completed(final ClientExchange result) { if (exchange.getConnection().isPushSupported() && result.getConnection().isPushSupported()) this.handleServerPush(result); - result.setResponseListener(new ResponseCallback(exchange, proxyClientHandler, idempotentPredicate, headerRewriteRules, streamingRequest, streamResponseContentTypes, streamIdleTimeout, streamResponseHeaderOverwrite)); + result.setResponseListener(boundaryCallback(new ResponseCallback(ProxyAction.this))); - final IoExceptionHandler handler = new IoExceptionHandler(exchange, clientConnection.getConnection()); + final IoExceptionHandler handler = new IoExceptionHandler(exchange, continuation, attempt); if (requiresContinueResponse) { this.prepRequestChannelForContinue(result, handler); return; } - HTTPTrailerChannelListener trailerListener = new HTTPTrailerChannelListener(exchange, result, exchange, proxyClientHandler, idempotentPredicate); + HTTPTrailerChannelListener trailerListener = new HTTPTrailerChannelListener(exchange, result, exchange, continuation, proxyClientHandler, idempotentPredicate); if (!exchange.isRequestComplete()) { @@ -930,7 +1067,7 @@ public void completed(final ClientExchange result) { @Override public void failed(IOException e) { - handleFailure(exchange, proxyClientHandler, idempotentPredicate, e); + continuation.run(ex -> handleFailure(ex, continuation, proxyClientHandler, idempotentPredicate, e)); } private void prepContinueHandler(ClientExchange result) { @@ -962,7 +1099,7 @@ private void prepRequestChannelForContinue(ClientExchange ce, IoExceptionHandler exchange.getRequestChannel(), ce.getRequestChannel(), ChannelListeners.closingChannelListener(), - new HTTPTrailerChannelListener(exchange, ce, exchange, proxyClientHandler, idempotentPredicate), + new HTTPTrailerChannelListener(exchange, ce, exchange, continuation, proxyClientHandler, idempotentPredicate), io, io, exchange.getConnection().getByteBufferPool()), @@ -976,25 +1113,32 @@ private void prepRequestChannelForContinue(ClientExchange ce, IoExceptionHandler private void handleServerPush(ClientExchange ce) { ce.setPushHandler((originalRequest, pushedRequest) -> { - - LOG.debug("Sending push request {} received from {} to target {} for exchange {}", pushedRequest.getRequest(), r, host, exchange); - - final ClientRequest req = pushedRequest.getRequest(); - - exchange.getConnection().pushResource(req.getPath(), req.getMethod(), req.getRequestHeaders(), exchange -> { - String path = req.getPath(); - int i = path.indexOf("?"); - - if (i > 0) - path = path.substring(0, i); - - exchange.dispatch(lightThreadExecutor, new ProxyAction(new ProxyConnection(pushedRequest.getConnection(), path), exchange, requestHeaders, rewriteHostHeader, reuseXForwarded, null, idempotentPredicate, urlRewriteRules, methodRewriteRules, queryParamRewriteRules, headerRewriteRules, false, streamResponseContentTypes, streamIdleTimeout, streamResponseHeaderOverwrite)); - }); + continuation.run(parent -> { + if (parent.isComplete() || continuation.stopping) { + cancelClientExchange(pushedRequest); + return; + } + LOG.debug("Sending push request {} received from {} to target {} for exchange {}", pushedRequest.getRequest(), r, host, parent); + final ClientRequest req = pushedRequest.getRequest(); + boolean accepted = parent.getConnection().pushResource(req.getPath(), req.getMethod(), req.getRequestHeaders(), child -> { + String path = req.getPath(); + int i = path.indexOf("?"); + if (i > 0) path = path.substring(0, i); + String targetPath = path; + ProxyContinuation pushed = new ProxyContinuation(child, continuation.scope.fork(child)); + ProxyConnection pushedConnection = new ProxyConnection(pushedRequest.getConnection(), targetPath); + child.putAttachment(CONNECTION, pushedConnection); + pushed.connection = pushedConnection; + pushed.publishResources(); + child.dispatch(INLINE_EXECUTOR, () -> pushed.run(ex -> new ProxyAction(pushedConnection, ex, requestHeaders, rewriteHostHeader, reuseXForwarded, null, idempotentPredicate, urlRewriteRules, methodRewriteRules, queryParamRewriteRules, headerRewriteRules, false, streamResponseContentTypes, streamIdleTimeout, streamResponseHeaderOverwrite, pushed).run())); + }); + if (!accepted) cancelClientExchange(pushedRequest); + }, () -> cancelClientExchange(pushedRequest)); return true; }); - } - }); + + })); } private void attachSslInfo(ClientRequest r) { @@ -1018,7 +1162,7 @@ private void attachSslInfo(ClientRequest r) { } } - static void handleFailure(HttpServerExchange exchange, ProxyClientHandler proxyClientHandler, Predicate idempotentRequestPredicate, IOException e) { + private static void handleFailure(HttpServerExchange exchange, ProxyContinuation continuation, ProxyClientHandler proxyClientHandler, Predicate idempotentRequestPredicate, IOException e) { UndertowLogger.PROXY_REQUEST_LOGGER.proxyRequestFailed(exchange.getRequestURI(), e); if (exchange.isResponseStarted()) { @@ -1026,9 +1170,10 @@ static void handleFailure(HttpServerExchange exchange, ProxyClientHandler proxyC IoUtils.safeClose(exchange.getConnection()); } else if (idempotentRequestPredicate.resolve(exchange) && proxyClientHandler != null) - proxyClientHandler.failed(exchange); //this will attempt a retry if configured to do so + proxyClientHandler.failed(); //this will attempt a retry if configured to do so else { + abortUpstream(continuation); exchange.setStatusCode(StatusCodes.SERVICE_UNAVAILABLE); exchange.endExchange(); } @@ -1037,6 +1182,8 @@ static void handleFailure(HttpServerExchange exchange, ProxyClientHandler proxyC private static final class ResponseCallback implements ClientCallback { private final HttpServerExchange exchange; + private final ProxyContinuation continuation; + private final long attempt; private final ProxyClientHandler proxyClientHandler; private final Predicate idempotentPredicate; private final Map> headerRewriteRules; @@ -1045,16 +1192,17 @@ private static final class ResponseCallback implements ClientCallback streamResponseHeaderOverwrite; - private ResponseCallback(HttpServerExchange exchange, ProxyClientHandler proxyClientHandler, Predicate idempotentPredicate, Map> headerRewriteRules, - boolean streamingRequest, List streamResponseContentTypes, int streamIdleTimeout, List streamResponseHeaderOverwrite) { - this.exchange = exchange; - this.proxyClientHandler = proxyClientHandler; - this.idempotentPredicate = idempotentPredicate; - this.headerRewriteRules = headerRewriteRules; - this.streamingRequest = streamingRequest; - this.streamResponseContentTypes = streamResponseContentTypes; - this.streamIdleTimeout = streamIdleTimeout; - this.streamResponseHeaderOverwrite = streamResponseHeaderOverwrite; + private ResponseCallback(ProxyAction action) { + this.exchange = action.exchange; + this.continuation = action.continuation; + this.attempt = action.attempt; + this.proxyClientHandler = action.proxyClientHandler; + this.idempotentPredicate = action.idempotentPredicate; + this.headerRewriteRules = action.headerRewriteRules; + this.streamingRequest = action.streamingRequest; + this.streamResponseContentTypes = action.streamResponseContentTypes; + this.streamIdleTimeout = action.streamIdleTimeout; + this.streamResponseHeaderOverwrite = action.streamResponseHeaderOverwrite; } @Override @@ -1070,7 +1218,7 @@ public void completed(final ClientExchange result) { boolean streamingResponse = isStreamingResponse(inbound, streamResponseContentTypes); if (streamingResponse) { - if (!streamingRequest) cancelExchangeTimeout(exchange); + if (!streamingRequest) continuation.cancelTimeout(); removeHeaders(outbound, streamResponseHeaderOverwrite); } @@ -1080,14 +1228,14 @@ public void completed(final ClientExchange result) { if (exchange.isUpgrade()) this.handleUpgradeChannelOnComplete(result); - final IoExceptionHandler handler = new IoExceptionHandler(exchange, result.getConnection()); + final IoExceptionHandler handler = new IoExceptionHandler(exchange, continuation, attempt); StreamSourceChannel responseChannel = result.getResponseChannel(); if (streamingResponse && streamIdleTimeout > 0) { responseChannel = new IdleTimeoutStreamSourceChannel(responseChannel, exchange, result.getConnection(), streamIdleTimeout); } - Transfer.initiateTransfer(responseChannel, exchange.getResponseChannel(), ChannelListeners.closingChannelListener(), new HTTPTrailerChannelListener(result, exchange, exchange, proxyClientHandler, idempotentPredicate), handler, handler, exchange.getConnection().getByteBufferPool()); + Transfer.initiateTransfer(responseChannel, exchange.getResponseChannel(), ChannelListeners.closingChannelListener(), new HTTPTrailerChannelListener(result, exchange, exchange, continuation, proxyClientHandler, idempotentPredicate), handler, handler, exchange.getConnection().getByteBufferPool()); AbstractMetricsHandler metricsHandler = (AbstractMetricsHandler) exchange.getAttachment(AttachmentConstants.METRICS_HANDLER); if(metricsHandler != null) { @@ -1122,7 +1270,7 @@ private void handleUpgradeChannelOnComplete(ClientExchange result) { @Override public void failed(IOException e) { - handleFailure(exchange, proxyClientHandler, idempotentPredicate, e); + handleFailure(exchange, continuation, proxyClientHandler, idempotentPredicate, e); } } @@ -1302,13 +1450,17 @@ private static final class HTTPTrailerChannelListener implements ChannelListener private final Attachable source; private final Attachable target; private final HttpServerExchange exchange; + private final ProxyContinuation continuation; + private final long attempt; private final ProxyClientHandler proxyClientHandler; private final Predicate idempotentPredicate; - private HTTPTrailerChannelListener(final Attachable source, final Attachable target, HttpServerExchange exchange, ProxyClientHandler proxyClientHandler, Predicate idempotentPredicate) { + private HTTPTrailerChannelListener(final Attachable source, final Attachable target, HttpServerExchange exchange, ProxyContinuation continuation, ProxyClientHandler proxyClientHandler, Predicate idempotentPredicate) { this.source = source; this.target = target; this.exchange = exchange; + this.continuation = continuation; + this.attempt = continuation.attempt; this.proxyClientHandler = proxyClientHandler; this.idempotentPredicate = idempotentPredicate; } @@ -1340,47 +1492,48 @@ public void handleEvent(final StreamSinkChannel channel) { LOG.error("IOException: ", e); - handleFailure(exchange, proxyClientHandler, idempotentPredicate, e); + continuation.run(ex -> { + if (attempt == continuation.attempt && !continuation.stopping) + handleFailure(ex, continuation, proxyClientHandler, idempotentPredicate, e); + }); } catch (Exception e) { LOG.error("Exception: ", e); - handleFailure(exchange, proxyClientHandler, idempotentPredicate, new IOException(e)); + continuation.run(ex -> { + if (attempt == continuation.attempt && !continuation.stopping) + handleFailure(ex, continuation, proxyClientHandler, idempotentPredicate, new IOException(e)); + }); } } } private static final class IoExceptionHandler implements ChannelExceptionHandler { - private final HttpServerExchange exchange; - private final ClientConnection clientConnection; + private final ProxyContinuation continuation; + private final long attempt; - private IoExceptionHandler(HttpServerExchange exchange, ClientConnection clientConnection) { + private IoExceptionHandler(HttpServerExchange exchange, ProxyContinuation continuation, long attempt) { this.exchange = exchange; - this.clientConnection = clientConnection; + this.continuation = continuation; + this.attempt = attempt; } @Override - public void handleException(Channel channel, IOException e) { - IoUtils.safeClose(channel); - IoUtils.safeClose(clientConnection); - - if (exchange.isResponseStarted()) { - UndertowLogger.REQUEST_IO_LOGGER.debug("Exception reading from target server", e); - - if (!exchange.isResponseStarted()) { - exchange.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR); - exchange.endExchange(); - - } else IoUtils.safeClose(exchange.getConnection()); - - } else { - UndertowLogger.REQUEST_IO_LOGGER.ioException(e); - exchange.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR); - exchange.endExchange(); - } + public void handleException(Channel channel, IOException failure) { + continuation.run(ex -> { + IoUtils.safeClose(channel); + if (attempt != continuation.attempt || continuation.stopping || ex.isComplete()) return; + abortUpstream(continuation); + UndertowLogger.REQUEST_IO_LOGGER.ioException(failure); + if (ex.isResponseStarted()) IoUtils.safeClose(ex.getConnection()); + else { + ex.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR); + ex.endExchange(); + } + }); } } diff --git a/proxy-handler/src/test/java/com/networknt/handler/ProxyCallbackDispatcherTest.java b/proxy-handler/src/test/java/com/networknt/handler/ProxyCallbackDispatcherTest.java new file mode 100644 index 0000000000..195915206e --- /dev/null +++ b/proxy-handler/src/test/java/com/networknt/handler/ProxyCallbackDispatcherTest.java @@ -0,0 +1,175 @@ +package com.networknt.handler; + +import io.undertow.Undertow; +import io.undertow.server.HttpServerExchange; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.*; + +@Timeout(30) +class ProxyCallbackDispatcherTest { + @Test + void foreignArrivalCannotBeOvertakenByOwnerArrival() throws Exception { + var exchange = new HttpServerExchange(null); + var scheduled = new LinkedBlockingQueue(); + Thread owner = Thread.currentThread(); + var queue = new ProxyCallbackDispatcher(HandlerContinuation.capture(exchange), scheduled::add, + () -> Thread.currentThread() == owner, () -> fail("Unexpected owner failure"), + failure -> fail("Unexpected rejection", failure)); + var seen = new ArrayList(); + try (var worker = Executors.newSingleThreadExecutor()) { + worker.submit(() -> queue.execute(ex -> seen.add("foreign"))).get(5, TimeUnit.SECONDS); + } + queue.execute(ex -> seen.add("owner")); + assertTrue(seen.isEmpty()); + scheduled.take().run(); + assertEquals(List.of("foreign", "owner"), seen); + assertTrue(scheduled.isEmpty()); + } + + @Test + void onlyRejectingSubmissionNotifiesPolicyAndNeverCallsOwnerCleanup() throws Exception { + var exchange = new HttpServerExchange(null); + var selected = new AtomicReference(); + var policyCalls = new AtomicInteger(); + HandlerContinuation.execute(exchange, next -> ex -> { + try { next.handleRequest(ex); } + catch (RejectedExecutionException _) { policyCalls.incrementAndGet(); } + }, ex -> selected.set(HandlerContinuation.capture(ex))); + var rejectThread = new AtomicReference(); + var retired = new AtomicInteger(); + var cleanupFailure = new IllegalStateException("isolated resource cleanup"); + var rejection = new RejectedExecutionException("owner executor stopped"); + var queue = new ProxyCallbackDispatcher(selected.get(), task -> { throw rejection; }, + () -> false, () -> fail("Mutable owner cleanup ran on a rejecting submitter"), failure -> { + assertTrue(rejectThread.compareAndSet(null, Thread.currentThread())); + selected.get().execute(ex -> { throw failure; }); + }); + try (var workers = Executors.newFixedThreadPool(4)) { + Thread first = workers.submit(() -> { + queue.execute(ex -> fail("Rejected work ran"), () -> { + retired.incrementAndGet(); + throw cleanupFailure; + }); + return Thread.currentThread(); + }).get(5, TimeUnit.SECONDS); + var futureSubmissions = new ArrayList>(); + for (int n = 0; n < 100; n++) + futureSubmissions.add(workers.submit(() -> queue.execute(ex -> fail("Stopped work ran"), + retired::incrementAndGet))); + for (var future : futureSubmissions) future.get(5, TimeUnit.SECONDS); + assertSame(first, rejectThread.get()); + } + assertEquals(1, policyCalls.get()); + assertEquals(101, retired.get()); + assertArrayEquals(new Throwable[]{cleanupFailure}, rejection.getSuppressed()); + } + + @Test + void concurrentProducersPreserveTheirOrderAndDoNotLoseIdleWakeups() throws Exception { + var exchange = new HttpServerExchange(null); + var owner = new AtomicReference(); + var problem = new AtomicReference(); + try (var executor = Executors.newSingleThreadExecutor(task -> new Thread(() -> { + owner.set(Thread.currentThread()); + task.run(); + })); var producers = Executors.newFixedThreadPool(4)) { + var queue = new ProxyCallbackDispatcher(HandlerContinuation.capture(exchange), executor, + () -> Thread.currentThread() == owner.get(), + () -> problem.compareAndSet(null, "owner failure"), + failure -> problem.compareAndSet(null, "unexpected rejection")); + int[] expected = new int[4]; + for (int round = 0; round < 40; round++) { + var submitted = new ArrayList>(); + int base = round * 200; + for (int producer = 0; producer < 4; producer++) { + int id = producer; + submitted.add(producers.submit(() -> { + for (int n = base; n < base + 200; n++) { + int sequence = n; + queue.execute(ex -> { + if (Thread.currentThread() != owner.get()) + problem.compareAndSet(null, "callback escaped owner"); + if (expected[id]++ != sequence) + problem.compareAndSet(null, "producer order changed"); + }); + if ((n & 7) == 0) Thread.yield(); + } + })); + } + for (var future : submitted) future.get(5, TimeUnit.SECONDS); + var drained = new CompletableFuture(); + queue.execute(ex -> drained.complete(null)); + drained.get(5, TimeUnit.SECONDS); + } + assertArrayEquals(new int[]{8000, 8000, 8000, 8000}, expected); + assertNull(problem.get(), problem.get()); + } + } + + @Test + void fatalOwnerFailureRetiresQueuedAndForeignEventsAndRetainsDiagnostics() throws Exception { + var fatal = new StackOverflowError("injected fatal callback"); + var cleanupFailure = new IllegalStateException("queued resource cleanup"); + var observed = new CompletableFuture(); + var dispatcher = new CompletableFuture(); + var ownerThread = new AtomicReference(); + var cleanupThread = new AtomicReference(); + var ownerAborts = new AtomicInteger(); + var retired = new AtomicInteger(); + var server = Undertow.builder().addHttpListener(0, "127.0.0.1").setIoThreads(1).setHandler(exchange -> { + ownerThread.set(Thread.currentThread()); + var queue = new ProxyCallbackDispatcher(HandlerContinuation.capture(exchange), exchange.getIoThread(), + exchange::isInIoThread, () -> { + cleanupThread.set(Thread.currentThread()); + ownerAborts.incrementAndGet(); + }, failure -> observed.completeExceptionally(failure)); + dispatcher.complete(queue); + try { + queue.execute(ex -> { + queue.execute(queued -> fail("Queued transport work survived fatal error"), () -> { + retired.incrementAndGet(); + throw cleanupFailure; + }); + throw fatal; + }); + } catch (Error failure) { observed.complete(failure); } + }).build(); + try (var client = HttpClient.newHttpClient()) { + server.start(); + int port = ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort(); + var response = client.send(HttpRequest.newBuilder(URI.create("http://127.0.0.1:" + port + "/fatal")) + .timeout(Duration.ofSeconds(5)).build(), HttpResponse.BodyHandlers.ofString()); + assertEquals(500, response.statusCode()); + assertSame(fatal, observed.get(5, TimeUnit.SECONDS)); + try (var producers = Executors.newFixedThreadPool(4)) { + var futures = new ArrayList>(); + for (int n = 0; n < 100; n++) + futures.add(producers.submit(() -> dispatcher.join().execute(ex -> fail("Stopped transport ran"), + retired::incrementAndGet))); + for (var future : futures) future.get(5, TimeUnit.SECONDS); + } + assertSame(ownerThread.get(), cleanupThread.get()); + assertEquals(1, ownerAborts.get()); + assertEquals(101, retired.get()); + assertArrayEquals(new Throwable[]{cleanupFailure}, fatal.getSuppressed()); + } finally { server.stop(); } + } +} diff --git a/proxy-handler/src/test/java/com/networknt/handler/ProxyContinuationTest.java b/proxy-handler/src/test/java/com/networknt/handler/ProxyContinuationTest.java new file mode 100644 index 0000000000..5fe49d2ec4 --- /dev/null +++ b/proxy-handler/src/test/java/com/networknt/handler/ProxyContinuationTest.java @@ -0,0 +1,911 @@ +package com.networknt.handler; + + +import com.networknt.httpstring.AttachmentConstants; +import com.networknt.metrics.AbstractMetricsHandler; +import com.networknt.metrics.MetricsConfig; +import com.networknt.metrics.TimeSeriesDbSender; +import io.undertow.Undertow; +import io.undertow.attribute.ExchangeAttribute; +import io.undertow.util.HttpString; +import io.undertow.util.HeaderMap; +import io.undertow.util.Methods; +import io.undertow.UndertowOptions; +import io.undertow.client.*; +import io.undertow.server.*; +import io.undertow.server.handlers.proxy.*; +import io.undertow.util.Headers; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.MDC; +import org.xnio.channels.Channels; + +import java.io.IOException; +import java.io.InputStream; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.net.*; +import java.net.http.*; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Supplier; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.*; + +@Timeout(20) +class ProxyContinuationTest { + enum Phase { ACQUIRE, RETRY, SEND_CALLBACK, RESPONSE_CALLBACK, TIMEOUT, BEFORE_SEND } + private static final class ControlledFailure extends RuntimeException { + final Phase phase; + ControlledFailure(Phase phase) { super("Injected failure at " + phase); this.phase = phase; } + } + + static Stream failureCases() { + return Stream.of(Phase.values()).flatMap(phase -> Stream.of(false, true) + .flatMap(enabled -> Stream.of(false, true).map(mdc -> Arguments.of(phase, enabled, mdc)))); + } + + @ParameterizedTest + @MethodSource("failureCases") + void asynchronousFailuresUseSelectedBoundaryAndCloseUpstream(Phase phase, boolean enabled, boolean hasMdc) throws Exception { + var previousMdc = MDC.getCopyOfContextMap(); + AtomicInteger injected = new AtomicInteger(); + AtomicInteger selected = new AtomicInteger(); + ThreadLocal depth = ThreadLocal.withInitial(() -> 0); + AtomicInteger maxDepth = new AtomicInteger(); + AtomicInteger acquisitions = new AtomicInteger(); + var upstream = new AtomicReference(); + CountDownLatch release = new CountDownLatch(1); + CompletableFuture retry = new CompletableFuture<>(); + Supplier failure = () -> { + assertEquals(hasMdc ? "canonical-request" : null, MDC.get("request"), + "Callbacks must install canonical request MDC for this frame"); + injected.incrementAndGet(); + return new ControlledFailure(phase); + }; + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newHttpClient()) { + URI backend = fixture.start(exchange -> exchange.dispatch(() -> { + try { + if (phase == Phase.TIMEOUT && !release.await(5, TimeUnit.SECONDS)) return; + exchange.getResponseSender().send("accepted"); + } catch (InterruptedException _) { Thread.currentThread().interrupt(); } + })); + LoadBalancingProxyClient real = new LoadBalancingProxyClient().addHost(backend); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange exchange) { return real.findTarget(exchange); } + @Override public void getConnection(ProxyTarget target, HttpServerExchange exchange, + ProxyCallback callback, long timeout, TimeUnit unit) { + int attempt = acquisitions.incrementAndGet(); + if (phase == Phase.ACQUIRE || (phase == Phase.RETRY && attempt > 1)) throw failure.get(); + if (phase == Phase.RETRY) { + retry.complete(() -> callback.failed(exchange)); + return; + } + real.getConnection(target, exchange, new ProxyCallback<>() { + @Override public void completed(HttpServerExchange ex, ProxyConnection connection) { + upstream.set(connection.getConnection()); + callback.completed(ex, new ProxyConnection(decorateConnection(connection.getConnection(), phase, failure), + connection.getTargetPath())); + } + @Override public void failed(HttpServerExchange ex) { callback.failed(ex); } + @Override public void couldNotResolveBackend(HttpServerExchange ex) { callback.couldNotResolveBackend(ex); } + @Override public void queuedRequestFailed(HttpServerExchange ex) { callback.queuedRequestFailed(ex); } + }, timeout, unit); + } + }; + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(client) + .setMaxConnectionRetries(phase == Phase.RETRY ? 1 : 0) + .setMaxRequestTime(phase == Phase.TIMEOUT ? 300 : 0).build(); + HttpHandler work = exchange -> { + if (hasMdc) RequestMdcContext.initialize(exchange).putField("request", "canonical-request"); + if (phase == Phase.TIMEOUT) { + exchange.putAttachment(AttachmentConstants.METRICS_HANDLER, new ThrowingMetrics(failure)); + exchange.putAttachment(AttachmentConstants.DOWNSTREAM_METRICS_NAME, "test"); + exchange.putAttachment(AttachmentConstants.DOWNSTREAM_METRICS_START, System.nanoTime()); + } + proxy.handleRequest(exchange); + }; + URI front = fixture.start(enabled ? exchange -> HandlerContinuation.execute(exchange, task -> ex -> { + depth.set(depth.get() + 1); + maxDepth.accumulateAndGet(depth.get(), Math::max); + try { task.handleRequest(ex); } + catch (ControlledFailure e) { + assertEquals(phase, e.phase); + selected.incrementAndGet(); + ex.setStatusCode(409).getResponseSender().send("selected"); + } finally { if (depth.get() == 1) depth.remove(); else depth.set(depth.get() - 1); } + }, work) : work); + var pending = http.sendAsync(HttpRequest.newBuilder(front).timeout(Duration.ofSeconds(8)).build(), + HttpResponse.BodyHandlers.ofString()); + if (phase == Phase.RETRY) retry.get(5, TimeUnit.SECONDS).run(); + var response = pending.get(10, TimeUnit.SECONDS); + assertEquals(enabled ? 409 : 500, response.statusCode(), response.body()); + assertEquals(1, injected.get(), "The intentional failure must execute exactly once"); + assertEquals(enabled ? 1 : 0, selected.get(), "Do not accept an incidental exception"); + if (enabled) assertEquals(1, maxDepth.get(), "Internal calls must not nest boundary frames"); + assertEquals(phase == Phase.RETRY ? 2 : 1, acquisitions.get()); + if (phase != Phase.ACQUIRE && phase != Phase.RETRY) + assertFalse(upstream.get().isOpen(), "Close upstream resources on callback failure"); + release.countDown(); + } finally { + release.countDown(); + if (previousMdc == null) MDC.clear(); + else MDC.setContextMap(previousMdc); + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void successfulProxyActionRemainsPendingUntilUpstreamCompletes(boolean enabled) throws Exception { + CountDownLatch release = new CountDownLatch(1); + CountDownLatch complete = new CountDownLatch(1); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newHttpClient()) { + URI backend = fixture.start(exchange -> exchange.dispatch(() -> { + try { + exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "text/event-stream"); + var channel = exchange.getResponseChannel(); + Channels.writeBlocking(channel, ByteBuffer.wrap("first".getBytes(StandardCharsets.UTF_8))); + Channels.flushBlocking(channel); + assertTrue(release.await(5, TimeUnit.SECONDS)); + Channels.writeBlocking(channel, ByteBuffer.wrap("second".getBytes(StandardCharsets.UTF_8))); + channel.shutdownWrites(); + Channels.flushBlocking(channel); + } catch (Exception e) { throw new RuntimeException(e); } + })); + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(new LoadBalancingProxyClient().addHost(backend)) + .setMaxRequestTime(0).build(); + HttpHandler work = exchange -> { + exchange.addExchangeCompleteListener((ex, next) -> { complete.countDown(); next.proceed(); }); + proxy.handleRequest(exchange); + }; + URI front = fixture.start(enabled ? exchange -> HandlerContinuation.execute(exchange, + task -> ex -> task.handleRequest(ex), work) : work); + var response = http.send(HttpRequest.newBuilder(front).timeout(Duration.ofSeconds(8)).build(), + HttpResponse.BodyHandlers.ofInputStream()); + assertEquals(200, response.statusCode()); + try (var body = response.body()) { + assertEquals("first", new String(body.readNBytes(5), StandardCharsets.UTF_8)); + assertEquals(1, complete.getCount(), "A successful action must leave streaming I/O pending"); + release.countDown(); + assertEquals("second", new String(body.readAllBytes(), StandardCharsets.UTF_8)); + } + assertTrue(complete.await(3, TimeUnit.SECONDS)); + } finally { release.countDown(); } + } + + @Test + void unhandledTaskFailureAfterHeadersClosesStreamWithoutAppendingAnError() throws Exception { + try (Fixture fixture = new Fixture()) { + URI front = fixture.start(exchange -> exchange.dispatch(() -> { + try { + var channel = exchange.getResponseChannel(); + Channels.writeBlocking(channel, ByteBuffer.wrap("first".getBytes(StandardCharsets.UTF_8))); + Channels.flushBlocking(channel); + assertTrue(exchange.isResponseStarted()); + HandlerContinuation.task(exchange, ex -> { throw new ControlledFailure(Phase.RESPONSE_CALLBACK); }).run(); + } catch (IOException e) { throw new RuntimeException(e); } + })); + try (Socket socket = new Socket(front.getHost(), front.getPort())) { + socket.setSoTimeout(5000); + socket.getOutputStream().write("GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n" + .getBytes(StandardCharsets.UTF_8)); + String raw = new String(socket.getInputStream().readAllBytes(), StandardCharsets.UTF_8); + assertTrue(raw.startsWith("HTTP/1.1 200"), raw); + assertTrue(raw.contains("first"), raw); + assertFalse(raw.contains("500"), raw); + assertFalse(raw.endsWith("0\r\n\r\n"), "Do not mark a failed stream as successfully complete"); + } + } + } + + @Test + void lateFailureDoesNotCloseACompletedKeepAliveConnection() throws Exception { + var completed = new AtomicReference(); + CountDownLatch done = new CountDownLatch(1); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newHttpClient()) { + URI front = fixture.start(exchange -> { + exchange.addExchangeCompleteListener((ex, next) -> { + completed.set(ex); + done.countDown(); + next.proceed(); + }); + exchange.getResponseSender().send("accepted"); + }); + var request = HttpRequest.newBuilder(front).timeout(Duration.ofSeconds(5)).build(); + assertEquals("accepted", http.send(request, HttpResponse.BodyHandlers.ofString()).body()); + assertTrue(done.await(3, TimeUnit.SECONDS)); + assertTrue(completed.get().isComplete()); + assertTrue(completed.get().getConnection().isOpen()); + HandlerContinuation.task(completed.get(), ex -> { + throw new ControlledFailure(Phase.RESPONSE_CALLBACK); + }).run(); + assertTrue(completed.get().getConnection().isOpen(), "A late callback must not close a reusable connection"); + assertEquals("accepted", http.send(request, HttpResponse.BodyHandlers.ofString()).body()); + } + } + + @ParameterizedTest + @ValueSource(strings = {"response", "timeout", "before-send", "late"}) + void upstreamHttp2FailureResetsOnlyTheFailedStream(String mode) throws Exception { + CompletableFuture late = new CompletableFuture<>(); + AtomicInteger failedAcquisitions = new AtomicInteger(); + CountDownLatch healthyEntered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Set backendPeers = ConcurrentHashMap.newKeySet(); + Set frontendPeers = ConcurrentHashMap.newKeySet(); + var shared = new AtomicReference(); + AtomicInteger injected = new AtomicInteger(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build()) { + URI backend = fixture.start(exchange -> { + assertEquals("HTTP/2.0", exchange.getProtocol().toString()); + backendPeers.add(exchange.getSourceAddress().getPort()); + if (exchange.getRequestPath().equals("/healthy") + || (mode.equals("timeout") && exchange.getRequestPath().equals("/fail"))) exchange.dispatch(() -> { + if (exchange.getRequestPath().equals("/healthy")) healthyEntered.countDown(); + try { + assertTrue(release.await(8, TimeUnit.SECONDS)); + exchange.getResponseSender().send("healthy"); + } catch (InterruptedException _) { Thread.currentThread().interrupt(); } + }); + else exchange.getResponseSender().send("accepted"); + }, true); + URI upstream = URI.create("h2c-prior://" + backend.getHost() + ":" + backend.getPort()); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange exchange) { return new ProxyTarget() {}; } + @Override public void getConnection(ProxyTarget target, HttpServerExchange exchange, + ProxyCallback callback, long timeout, TimeUnit unit) { + if (exchange.getRequestPath().equals("/fail")) failedAcquisitions.incrementAndGet(); + Consumer connected = connection -> { + ClientConnection selected = connection; + if (exchange.getRequestPath().equals("/fail") && !mode.equals("timeout")) selected = decorateConnection(connection, + mode.equals("before-send") || mode.equals("late") ? Phase.BEFORE_SEND : Phase.RESPONSE_CALLBACK, () -> { + injected.incrementAndGet(); + return new ControlledFailure(Phase.RESPONSE_CALLBACK); + }); + callback.completed(exchange, new ProxyConnection(selected, "/")); + }; + if (mode.equals("late") && exchange.getRequestPath().equals("/fail")) { + late.complete(() -> { + connected.accept(shared.get()); + callback.failed(exchange); + callback.queuedRequestFailed(exchange); + }); + return; + } + if (shared.get() != null) connected.accept(shared.get()); + else UndertowClient.getInstance().connect(new ClientCallback<>() { + @Override public void completed(ClientConnection connection) { + shared.set(connection); + connected.accept(connection); + } + @Override public void failed(IOException e) { callback.failed(exchange); } + }, upstream, exchange.getIoThread(), exchange.getConnection().getByteBufferPool(), org.xnio.OptionMap.EMPTY); + } + }; + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(client).setMaxRequestTime(0) + .setPathPrefixMaxRequestTime(Map.of("/fail", mode.equals("timeout") || mode.equals("late") ? 300 : 0)).build(); + URI front = fixture.start(exchange -> { + frontendPeers.add(exchange.getSourceAddress().getPort()); + HandlerContinuation.execute(exchange, task -> ex -> { + try { task.handleRequest(ex); } + catch (ControlledFailure _) { ex.setStatusCode(409).getResponseSender().send("selected"); } + }, proxy); + }, true); + var warm = http.send(HttpRequest.newBuilder(front.resolve("/warm")).build(), HttpResponse.BodyHandlers.ofString()); + assertEquals(HttpClient.Version.HTTP_2, warm.version()); + assertEquals("accepted", warm.body()); + assertTrue(shared.get().isMultiplexingSupported()); + var healthy = http.sendAsync(HttpRequest.newBuilder(front.resolve("/healthy")).timeout(Duration.ofSeconds(10)).build(), + HttpResponse.BodyHandlers.ofString()); + assertTrue(healthyEntered.await(5, TimeUnit.SECONDS)); + var failed = http.send(HttpRequest.newBuilder(front.resolve("/fail")).timeout(Duration.ofSeconds(5)).build(), + HttpResponse.BodyHandlers.ofString()); + boolean timeout = mode.equals("timeout") || mode.equals("late"); + assertEquals(timeout ? 504 : 409, failed.statusCode(), failed.body()); + if (mode.equals("late")) late.get(3, TimeUnit.SECONDS).run(); + assertEquals(timeout ? 0 : 1, injected.get()); + assertEquals(1, failedAcquisitions.get(), "Late failure callbacks must not retry"); + assertTrue(shared.get().isOpen(), "Keep the shared upstream HTTP/2 connection open"); + assertFalse(healthy.isDone(), "Do not cancel the healthy sibling stream"); + release.countDown(); + assertEquals("healthy", healthy.get(5, TimeUnit.SECONDS).body()); + assertEquals("accepted", http.send(HttpRequest.newBuilder(front.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString()).body()); + assertEquals(1, backendPeers.size(), "All upstream requests must share one physical connection"); + assertEquals(1, frontendPeers.size(), "All downstream requests must share one physical connection"); + } finally { + release.countDown(); + if (shared.get() != null) shared.get().close(); + } + } + + @ParameterizedTest + @CsvSource({"false,true", "true,true", "true,false"}) + void downstreamHttp2FailureDoesNotCancelSiblingStream(boolean fatal, boolean committed) throws Exception { + CompletableFuture propagated = new CompletableFuture<>(); + Error sentinel = new StackOverflowError("injected fatal error"); + CountDownLatch healthyEntered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Set peers = ConcurrentHashMap.newKeySet(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build()) { + URI front = fixture.start(exchange -> { + peers.add(exchange.getSourceAddress().getPort()); + if (exchange.getRequestPath().equals("/healthy")) exchange.dispatch(() -> { + healthyEntered.countDown(); + try { + assertTrue(release.await(8, TimeUnit.SECONDS)); + exchange.getResponseSender().send("healthy"); + } catch (InterruptedException _) { Thread.currentThread().interrupt(); } + }); + else if (exchange.getRequestPath().equals("/fail")) exchange.dispatch(() -> { + try { + if (committed) { + var channel = exchange.getResponseChannel(); + Channels.writeBlocking(channel, ByteBuffer.wrap("first".getBytes(StandardCharsets.UTF_8))); + Channels.flushBlocking(channel); + } + assertEquals(committed, exchange.isResponseStarted()); + try { + HandlerContinuation.task(exchange, ex -> { + if (fatal) throw sentinel; + throw new ControlledFailure(Phase.RESPONSE_CALLBACK); + }).run(); + } catch (Error error) { propagated.complete(error); } + } catch (IOException e) { throw new RuntimeException(e); } + }); + else exchange.getResponseSender().send("accepted"); + }, true); + assertEquals(HttpClient.Version.HTTP_2, http.send(HttpRequest.newBuilder(front.resolve("/warm")).build(), + HttpResponse.BodyHandlers.ofString()).version()); + var healthy = http.sendAsync(HttpRequest.newBuilder(front.resolve("/healthy")).timeout(Duration.ofSeconds(10)).build(), + HttpResponse.BodyHandlers.ofString()); + assertTrue(healthyEntered.await(5, TimeUnit.SECONDS)); + if (fatal && !committed) { + var failed = http.send(HttpRequest.newBuilder(front.resolve("/fail")).timeout(Duration.ofSeconds(5)).build(), + HttpResponse.BodyHandlers.ofString()); + assertEquals(500, failed.statusCode()); + } else assertThrows(IOException.class, () -> http.send(HttpRequest.newBuilder(front.resolve("/fail")) + .timeout(Duration.ofSeconds(5)).build(), HttpResponse.BodyHandlers.ofString())); + if (fatal) assertSame(sentinel, propagated.get(3, TimeUnit.SECONDS)); + assertFalse(healthy.isDone()); + release.countDown(); + assertEquals("healthy", healthy.get(5, TimeUnit.SECONDS).body()); + assertEquals("accepted", http.send(HttpRequest.newBuilder(front.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString()).body()); + assertEquals(1, peers.size(), "RST_STREAM must not kill the parent HTTP/2 connection"); + } finally { release.countDown(); } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void fatalHttp11FailureTerminatesRequestAndPropagatesOriginal(boolean committed) throws Exception { + Error fatal = new StackOverflowError("HTTP/1.1 fatal sentinel"); + CompletableFuture propagated = new CompletableFuture<>(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1).build()) { + URI front = fixture.start(exchange -> { + if (exchange.getRequestPath().equals("/after")) { + exchange.getResponseSender().send("accepted"); + return; + } + exchange.dispatch(() -> { + try { + if (committed) { + var channel = exchange.getResponseChannel(); + Channels.writeBlocking(channel, ByteBuffer.wrap("first".getBytes(StandardCharsets.UTF_8))); + Channels.flushBlocking(channel); + } + HandlerContinuation.task(exchange, ex -> { throw fatal; }).run(); + } catch (Error error) { propagated.complete(error); } + catch (IOException e) { propagated.completeExceptionally(e); } + }); + }); + var request = HttpRequest.newBuilder(front).timeout(Duration.ofSeconds(5)).build(); + if (committed) assertThrows(IOException.class, () -> http.send(request, HttpResponse.BodyHandlers.ofString())); + else assertEquals(500, http.send(request, HttpResponse.BodyHandlers.ofString()).statusCode()); + assertSame(fatal, propagated.get(3, TimeUnit.SECONDS)); + assertEquals("accepted", http.send(HttpRequest.newBuilder(front.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString()).body()); + } + } + + @ParameterizedTest + @ValueSource(strings = {"completed", "failed", "queued"}) + void callbackAfterCommitTerminatesIncompleteStreamWithoutForwarding(String kind) throws Exception { + CompletableFuture late = new CompletableFuture<>(); + CountDownLatch committed = new CountDownLatch(1); + CountDownLatch delivered = new CountDownLatch(1); + AtomicInteger acquisitions = new AtomicInteger(); + var closed = new AtomicBoolean(); + ClientConnection unusedConnection = (ClientConnection) Proxy.newProxyInstance(ClientConnection.class.getClassLoader(), + new Class[]{ClientConnection.class}, (proxy, method, args) -> switch (method.getName()) { + case "isMultiplexingSupported" -> false; + case "close" -> { closed.set(true); yield null; } + case "isOpen" -> !closed.get(); + case "toString" -> "unused late connection"; + default -> throw new AssertionError("A late callback must not start a ProxyAction: " + method.getName()); + }); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange exchange) { return new ProxyTarget() {}; } + @Override public void getConnection(ProxyTarget target, HttpServerExchange exchange, + ProxyCallback callback, long timeout, TimeUnit unit) { + acquisitions.incrementAndGet(); + late.complete(() -> { + try { + assertTrue(committed.await(5, TimeUnit.SECONDS)); + assertTrue(exchange.isResponseStarted()); + assertFalse(exchange.isComplete()); + switch (kind) { + case "completed" -> callback.completed(exchange, new ProxyConnection(unusedConnection, "/")); + case "failed" -> callback.failed(exchange); + case "queued" -> callback.queuedRequestFailed(exchange); + default -> throw new AssertionError(kind); + } + awaitIoCallbacks(exchange); + } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); } + finally { delivered.countDown(); } + }); + exchange.dispatch(worker -> { + worker.startBlocking(); + worker.getOutputStream().write("first".getBytes(StandardCharsets.UTF_8)); + worker.getOutputStream().flush(); + committed.countDown(); + assertTrue(delivered.await(5, TimeUnit.SECONDS)); + }); + } + }; + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(client).setMaxRequestTime(0).build(); + try (Fixture fixture = new Fixture()) { + URI front = fixture.start(proxy); + try (Socket socket = new Socket(front.getHost(), front.getPort())) { + socket.setSoTimeout(5000); + socket.getOutputStream().write("GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n" + .getBytes(StandardCharsets.UTF_8)); + late.get(3, TimeUnit.SECONDS).run(); + String raw = new String(socket.getInputStream().readAllBytes(), StandardCharsets.UTF_8); + assertTrue(raw.startsWith("HTTP/1.1 200"), raw); + assertTrue(raw.contains("first"), raw); + assertEquals(1, acquisitions.get()); + if (kind.equals("completed")) assertTrue(closed.get()); + } + } + } + + @ParameterizedTest + @ValueSource(strings = {"completed", "failed", "queued"}) + void unexpectedCommittedHttp2CallbackResetsOnlyItsStream(String kind) throws Exception { + CompletableFuture late = new CompletableFuture<>(); + CountDownLatch committed = new CountDownLatch(1); + CountDownLatch delivered = new CountDownLatch(1); + AtomicInteger acquisitions = new AtomicInteger(); + CountDownLatch healthyEntered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Set peers = ConcurrentHashMap.newKeySet(); + ClientConnection reusable = (ClientConnection) Proxy.newProxyInstance(ClientConnection.class.getClassLoader(), + new Class[]{ClientConnection.class}, (proxy, method, args) -> switch (method.getName()) { + case "isMultiplexingSupported", "isOpen" -> true; + case "toString" -> "shared late HTTP/2 connection"; + default -> throw new AssertionError("Do not use or close the late shared connection: " + method.getName()); + }); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange exchange) { return new ProxyTarget() {}; } + @Override public void getConnection(ProxyTarget target, HttpServerExchange exchange, + ProxyCallback callback, long timeout, TimeUnit unit) { + acquisitions.incrementAndGet(); + late.complete(() -> { + try { + assertTrue(committed.await(5, TimeUnit.SECONDS)); + assertTrue(exchange.isResponseStarted()); + assertFalse(exchange.isComplete()); + switch (kind) { + case "completed" -> callback.completed(exchange, new ProxyConnection(reusable, "/")); + case "failed" -> callback.failed(exchange); + case "queued" -> callback.queuedRequestFailed(exchange); + default -> throw new AssertionError(kind); + } + awaitIoCallbacks(exchange); + } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); } + finally { delivered.countDown(); } + }); + exchange.dispatch(worker -> { + worker.startBlocking(); + worker.getOutputStream().write("first".getBytes(StandardCharsets.UTF_8)); + worker.getOutputStream().flush(); + committed.countDown(); + assertTrue(delivered.await(5, TimeUnit.SECONDS)); + }); + } + }; + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(client).setMaxRequestTime(0).build(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build()) { + try { + URI front = fixture.start(exchange -> { + peers.add(exchange.getSourceAddress().getPort()); + if (exchange.getRequestPath().equals("/healthy")) exchange.dispatch(() -> { + healthyEntered.countDown(); + try { + assertTrue(release.await(8, TimeUnit.SECONDS)); + exchange.getResponseSender().send("healthy"); + } catch (InterruptedException _) { Thread.currentThread().interrupt(); } + }); + else if (exchange.getRequestPath().equals("/fail")) proxy.handleRequest(exchange); + else exchange.getResponseSender().send("accepted"); + }, true); + assertEquals(HttpClient.Version.HTTP_2, http.send(HttpRequest.newBuilder(front.resolve("/warm")).build(), + HttpResponse.BodyHandlers.ofString()).version()); + var healthy = http.sendAsync(HttpRequest.newBuilder(front.resolve("/healthy")).timeout(Duration.ofSeconds(10)).build(), + HttpResponse.BodyHandlers.ofString()); + assertTrue(healthyEntered.await(5, TimeUnit.SECONDS)); + var failed = http.send(HttpRequest.newBuilder(front.resolve("/fail")).timeout(Duration.ofSeconds(5)).build(), + HttpResponse.BodyHandlers.ofInputStream()); + assertEquals(200, failed.statusCode()); + try (var body = failed.body()) { + assertEquals("first", new String(body.readNBytes(5), StandardCharsets.UTF_8)); + late.get(3, TimeUnit.SECONDS).run(); + assertThrows(IOException.class, body::readAllBytes); + } + assertEquals(1, acquisitions.get()); + assertFalse(healthy.isDone()); + release.countDown(); + assertEquals("healthy", healthy.get(5, TimeUnit.SECONDS).body()); + assertEquals("accepted", http.send(HttpRequest.newBuilder(front.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString()).body()); + assertEquals(1, peers.size(), "Only the failed HTTP/2 stream may be reset"); + } finally { release.countDown(); http.shutdownNow(); } + } finally { release.countDown(); } + } + + @Test + void lateCallbacksAfterProxy503LeaveDrainingHttp2ResponseAlone() throws Exception { + CompletableFuture callbacksDone = new CompletableFuture<>(); + CountDownLatch releaseBody = new CountDownLatch(1); + CountDownLatch healthyEntered = new CountDownLatch(1); + CountDownLatch releaseHealthy = new CountDownLatch(1); + Set peers = ConcurrentHashMap.newKeySet(); + ClientConnection reusable = (ClientConnection) Proxy.newProxyInstance(ClientConnection.class.getClassLoader(), + new Class[]{ClientConnection.class}, (proxy, method, args) -> switch (method.getName()) { + case "isMultiplexingSupported", "isOpen" -> true; + default -> throw new AssertionError("Late callbacks must not use or close the shared upstream: " + method.getName()); + }); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange exchange) { return new ProxyTarget() {}; } + @Override public void getConnection(ProxyTarget target, HttpServerExchange exchange, + ProxyCallback callback, long timeout, TimeUnit unit) { + assertFalse(exchange.isRequestComplete()); + callback.couldNotResolveBackend(exchange); + exchange.getIoThread().execute(() -> { + try { + // The dispatcher processes the queued terminal callback + // before this I/O task commits and probes its response. + var responseChannel = exchange.getResponseChannel(); + if (responseChannel != null) { responseChannel.shutdownWrites(); responseChannel.flush(); } + assertTrue(exchange.isResponseStarted()); + assertFalse(exchange.isComplete(), "The 503 is still draining the request body"); + callback.completed(exchange, new ProxyConnection(reusable, "/")); + callback.failed(exchange); + callback.queuedRequestFailed(exchange); + assertTrue(exchange.getConnection().isOpen(), "Keep the draining 503 stream open"); + callbacksDone.complete(null); + } catch (Exception | Error failure) { callbacksDone.completeExceptionally(failure); } + }); + } + }; + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(client).setMaxRequestTime(0).build(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build()) { + try { + URI front = fixture.start(exchange -> { + peers.add(exchange.getSourceAddress().getPort()); + if (exchange.getRequestPath().equals("/fail")) proxy.handleRequest(exchange); + else if (exchange.getRequestPath().equals("/healthy")) exchange.dispatch(() -> { + healthyEntered.countDown(); + try { + assertTrue(releaseHealthy.await(8, TimeUnit.SECONDS)); + exchange.getResponseSender().send("healthy"); + } catch (InterruptedException _) { Thread.currentThread().interrupt(); } + }); + else exchange.getResponseSender().send("accepted"); + }, true); + assertEquals(HttpClient.Version.HTTP_2, http.send(HttpRequest.newBuilder(front.resolve("/warm")).build(), + HttpResponse.BodyHandlers.ofString()).version()); + var healthy = http.sendAsync(HttpRequest.newBuilder(front.resolve("/healthy")).build(), HttpResponse.BodyHandlers.ofString()); + assertTrue(healthyEntered.await(5, TimeUnit.SECONDS)); + var body = HttpRequest.BodyPublishers.ofInputStream(() -> new InputStream() { + @Override public int read() throws IOException { + try { if (!releaseBody.await(8, TimeUnit.SECONDS)) throw new IOException("Body was not released"); } + catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IOException(e); } + return -1; + } + }); + var failed = http.sendAsync(HttpRequest.newBuilder(front.resolve("/fail")).timeout(Duration.ofSeconds(10)).POST(body).build(), + HttpResponse.BodyHandlers.ofString()); + try { callbacksDone.get(5, TimeUnit.SECONDS); } + catch (Exception failure) { releaseBody.countDown(); releaseHealthy.countDown(); throw failure; } + assertFalse(healthy.isDone()); + releaseBody.countDown(); + assertEquals(503, failed.get(5, TimeUnit.SECONDS).statusCode()); + releaseHealthy.countDown(); + assertEquals("healthy", healthy.get(5, TimeUnit.SECONDS).body()); + assertEquals("accepted", http.send(HttpRequest.newBuilder(front.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString()).body()); + assertEquals(1, peers.size()); + } finally { releaseBody.countDown(); releaseHealthy.countDown(); http.shutdownNow(); } + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void http2PushUsesItsOwnExchangeAndSelectedBoundary(boolean rejectPush) throws Exception { + var pushedResponses = new ConcurrentHashMap>>(); + var pushObserved = new CountDownLatch(1); + HttpResponse.PushPromiseHandler acceptPush = (request, pushed, acceptor) -> { + pushedResponses.put(pushed, acceptor.apply(HttpResponse.BodyHandlers.ofString())); + pushObserved.countDown(); + }; + var parentExchange = new AtomicReference(); + var childExchange = new AtomicReference(); + var failures = new CompletableFuture(); + Set peers = ConcurrentHashMap.newKeySet(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build()) { + try { + URI backend = fixture.start(exchange -> { + if (exchange.getRequestPath().equals("/parent")) { + assertTrue(exchange.getConnection().pushResource("/pushed", Methods.GET, + new HeaderMap(), child -> child.getResponseSender().send("pushed"))); + exchange.getResponseSender().send("parent"); + } else exchange.getResponseSender().send("pushed"); + }, true); + URI h2 = URI.create("h2c-prior://" + backend.getHost() + ":" + backend.getPort()); + ProxyHandler proxy = ProxyHandler.builder().setProxyClient(new LoadBalancingProxyClient().addHost(h2)) + .setMaxRequestTime(0).addRequestHeader(new HttpString("X-Push-Test"), new ExchangeAttribute() { + @Override public String readAttribute(HttpServerExchange ex) { + if (rejectPush && ex.getRequestPath().equals("/pushed")) + throw new ControlledFailure(Phase.BEFORE_SEND); + return "value"; + } + @Override public void writeAttribute(HttpServerExchange ex, String value) { + throw new UnsupportedOperationException("Read-only test attribute"); + } + }).build(); + URI front = fixture.start(exchange -> { + peers.add(exchange.getSourceAddress().getPort()); + if (exchange.getRequestPath().equals("/warm")) exchange.getResponseSender().send("warm"); + else HandlerContinuation.execute(exchange, next -> ex -> { + if (ex.getRequestPath().equals("/pushed")) childExchange.set(ex); + else parentExchange.set(ex); + try { next.handleRequest(ex); } + catch (ControlledFailure failure) { + assertEquals("/pushed", ex.getRequestPath()); + assertEquals(Phase.BEFORE_SEND, failure.phase); + ex.setStatusCode(409).getResponseSender().send("selected child"); + } catch (Exception | Error failure) { failures.complete(failure); throw failure; } + }, proxy); + }, true); + assertEquals(HttpClient.Version.HTTP_2, http.sendAsync(HttpRequest.newBuilder(front.resolve("/warm")).build(), + HttpResponse.BodyHandlers.ofString(), acceptPush).get(5, TimeUnit.SECONDS).version()); + var parent = http.sendAsync(HttpRequest.newBuilder(front.resolve("/parent")).timeout(Duration.ofSeconds(6)).build(), + HttpResponse.BodyHandlers.ofString(), acceptPush); + assertEquals("parent", parent.get(8, TimeUnit.SECONDS).body()); + assertTrue(pushObserved.await(5, TimeUnit.SECONDS), "Must receive a real HTTP/2 push promise"); + assertEquals(1, pushedResponses.size(), "Must exercise actual HTTP/2 push"); + var pushed = pushedResponses.values().iterator().next().get(5, TimeUnit.SECONDS); + assertEquals(rejectPush ? 409 : 200, pushed.statusCode(), pushed.body()); + assertEquals(rejectPush ? "selected child" : "pushed", pushed.body()); + assertNotNull(parentExchange.get()); + assertNotNull(childExchange.get()); + assertNotSame(parentExchange.get(), childExchange.get()); + assertFalse(failures.isDone(), "Push must not corrupt the parent exchange"); + assertEquals("pushed", http.sendAsync(HttpRequest.newBuilder(front.resolve("/after")).build(), + HttpResponse.BodyHandlers.ofString(), acceptPush).get(5, TimeUnit.SECONDS).body()); + assertEquals(1, peers.size(), "The parent connection must survive its pushed response"); + } finally { http.shutdownNow(); } + } + } + + @Test + void deferredBoundaryHasNoCallbackSlotAndPreservesPublicConnectionAttachment() throws Exception { + var frontendExchange = new AtomicReference(); + var acquired = new AtomicReference(); + var boundaryCalls = new AtomicInteger(); + var callbackFailure = new CompletableFuture(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newHttpClient()) { + URI backend = fixture.start(exchange -> exchange.getResponseSender().send("accepted")); + var real = new LoadBalancingProxyClient().addHost(backend); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange ex) { return real.findTarget(ex); } + @Override public void getConnection(ProxyTarget target, HttpServerExchange ex, + ProxyCallback callback, long timeout, TimeUnit unit) { + assertTrue(ex.isInIoThread(), "No callback-wide monitor: transport work belongs to the I/O thread"); + real.getConnection(target, ex, new ProxyCallback<>() { + @Override public void completed(HttpServerExchange exchange, ProxyConnection connection) { + acquired.set(connection); + callback.completed(exchange, connection); + } + @Override public void failed(HttpServerExchange exchange) { callback.failed(exchange); } + @Override public void couldNotResolveBackend(HttpServerExchange exchange) { callback.couldNotResolveBackend(exchange); } + @Override public void queuedRequestFailed(HttpServerExchange exchange) { callback.queuedRequestFailed(exchange); } + }, timeout, unit); + } + }; + var proxy = ProxyHandler.builder().setProxyClient(client).setMaxRequestTime(0).build(); + URI front = fixture.start(exchange -> { + frontendExchange.set(exchange); + HandlerContinuation.execute(exchange, next -> ex -> { + if (boundaryCalls.getAndIncrement() == 0) next.handleRequest(ex); + else ex.dispatch(ex.getConnection().getWorker(), () -> { + try { next.handleRequest(ex); } + catch (Throwable failure) { callbackFailure.complete(failure); } + }); + }, proxy); + }); + var response = http.send(HttpRequest.newBuilder(front).timeout(Duration.ofSeconds(8)).build(), HttpResponse.BodyHandlers.ofString()); + assertEquals("accepted", response.body()); + assertTrue(boundaryCalls.get() > 1); + assertFalse(callbackFailure.isDone()); + awaitIoCallbacks(frontendExchange.get()); + assertSame(acquired.get(), frontendExchange.get().getAttachment(ProxyHandler.CONNECTION)); + } + } + + @Test + void retryTimeoutReleasesCurrentConnectionAndIgnoresOldAcquisitionFailure() throws Exception { + var firstResponseFailure = new CompletableFuture(); + var staleSendCompletion = new CompletableFuture(); + var staleResponseCompletion = new CompletableFuture(); + var secondSend = new CountDownLatch(1); + var connections = new CopyOnWriteArrayList(); + var acquired = new CopyOnWriteArrayList(); + var callbacks = new CopyOnWriteArrayList>(); + var frontend = new AtomicReference(); + var attempts = new AtomicInteger(); + try (Fixture fixture = new Fixture(); HttpClient http = HttpClient.newHttpClient()) { + URI backend = fixture.start(HttpServerExchange::dispatch); + ProxyClient client = new ProxyClient() { + @Override public ProxyTarget findTarget(HttpServerExchange ex) { return new ProxyTarget() {}; } + @Override public void getConnection(ProxyTarget target, HttpServerExchange exchange, + ProxyCallback callback, long timeout, TimeUnit unit) { + assertTrue(exchange.isInIoThread()); + callbacks.add(callback); + int attempt = attempts.incrementAndGet(); + UndertowClient.getInstance().connect(new ClientCallback<>() { + @Override public void completed(ClientConnection connection) { + connections.add(connection); + ClientConnection decorated = (ClientConnection) Proxy.newProxyInstance(ClientConnection.class.getClassLoader(), + new Class[]{ClientConnection.class}, (proxy, method, args) -> { + if (method.getName().equals("sendRequest")) { + if (attempt == 2) { secondSend.countDown(); return null; } + @SuppressWarnings("unchecked") var send = (ClientCallback) args[1]; + connection.sendRequest((ClientRequest) args[0], new ClientCallback<>() { + @Override public void completed(ClientExchange request) { + ClientExchange decoratedRequest = (ClientExchange) Proxy.newProxyInstance(ClientExchange.class.getClassLoader(), + new Class[]{ClientExchange.class}, (proxy, method, args) -> { + if (method.getName().equals("setResponseListener")) { + @SuppressWarnings("unchecked") var listener = (ClientCallback) args[0]; + firstResponseFailure.complete(() -> listener.failed(new IOException("first response failed"))); + staleResponseCompletion.complete(() -> listener.completed(request)); + return null; + } + try { return method.invoke(request, args); } + catch (InvocationTargetException e) { throw e.getCause(); } + }); + staleSendCompletion.complete(() -> send.completed(decoratedRequest)); + send.completed(decoratedRequest); + } + @Override public void failed(IOException e) { send.failed(e); } + }); + return null; + } + try { return method.invoke(connection, args); } + catch (InvocationTargetException e) { throw e.getCause(); } + }); + var selected = new ProxyConnection(decorated, "/"); + acquired.add(selected); + callback.completed(exchange, selected); + } + @Override public void failed(IOException e) { callback.failed(exchange); } + }, backend, exchange.getIoThread(), exchange.getConnection().getByteBufferPool(), org.xnio.OptionMap.EMPTY); + } + }; + var proxy = ProxyHandler.builder().setProxyClient(client).setMaxConnectionRetries(1).setMaxRequestTime(1500).build(); + URI front = fixture.start(exchange -> { frontend.set(exchange); proxy.handleRequest(exchange); }); + var response = http.sendAsync(HttpRequest.newBuilder(front).timeout(Duration.ofSeconds(6)).build(), HttpResponse.BodyHandlers.ofString()); + firstResponseFailure.get(5, TimeUnit.SECONDS).run(); + assertTrue(secondSend.await(5, TimeUnit.SECONDS)); + staleSendCompletion.get(5, TimeUnit.SECONDS).run(); + staleResponseCompletion.get(5, TimeUnit.SECONDS).run(); + callbacks.getFirst().failed(frontend.get()); + callbacks.getFirst().queuedRequestFailed(frontend.get()); + callbacks.getFirst().couldNotResolveBackend(frontend.get()); + callbacks.getFirst().completed(frontend.get(), acquired.getFirst()); + awaitIoCallbacks(frontend.get()); + assertSame(acquired.getLast(), frontend.get().getAttachment(ProxyHandler.CONNECTION)); + assertEquals(2, attempts.get(), "An old acquisition callback cannot retry the current attempt"); + assertEquals(504, response.get(5, TimeUnit.SECONDS).statusCode()); + assertEquals(2, connections.size()); + assertFalse(connections.getFirst().isOpen()); + assertFalse(connections.getLast().isOpen(), "Timeout must release the current attempt's connection"); + } + } + + private static void awaitIoCallbacks(HttpServerExchange exchange) { + assertFalse(exchange.isInIoThread(), "Only the test thread waits for I/O ownership"); + var drained = new CompletableFuture(); + exchange.getIoThread().execute(() -> drained.complete(null)); + try { drained.get(5, TimeUnit.SECONDS); } + catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); } + catch (ExecutionException | TimeoutException e) { throw new AssertionError(e); } + } + + private static ClientConnection decorateConnection(ClientConnection real, Phase phase, + Supplier failure) { + return (ClientConnection) Proxy.newProxyInstance(ClientConnection.class.getClassLoader(), + new Class[]{ClientConnection.class}, (proxy, method, args) -> { + if (method.getName().equals("sendRequest")) { + if (phase == Phase.BEFORE_SEND) throw failure.get(); + @SuppressWarnings("unchecked") var callback = (ClientCallback) args[1]; + real.sendRequest((ClientRequest) args[0], new ClientCallback<>() { + @Override public void completed(ClientExchange result) { callback.completed(decorateExchange(result, phase, failure)); } + @Override public void failed(IOException e) { callback.failed(e); } + }); + return null; + } + try { return method.invoke(real, args); } + catch (InvocationTargetException e) { throw e.getCause(); } + }); + } + + private static ClientExchange decorateExchange(ClientExchange real, Phase phase, + Supplier failure) { + return (ClientExchange) Proxy.newProxyInstance(ClientExchange.class.getClassLoader(), + new Class[]{ClientExchange.class}, (proxy, method, args) -> { + if (method.getName().equals("setResponseListener")) { + if (phase == Phase.SEND_CALLBACK) throw failure.get(); + @SuppressWarnings("unchecked") var callback = (ClientCallback) args[0]; + real.setResponseListener(new ClientCallback<>() { + @Override public void completed(ClientExchange result) { callback.completed(decorateExchange(result, phase, failure)); } + @Override public void failed(IOException e) { callback.failed(e); } + }); + return null; + } + if (method.getName().equals("getResponse") && phase == Phase.RESPONSE_CALLBACK) throw failure.get(); + try { return method.invoke(real, args); } + catch (InvocationTargetException e) { throw e.getCause(); } + }); + } + + private static final class ThrowingMetrics extends AbstractMetricsHandler { + private final Supplier failure; + ThrowingMetrics(Supplier failure) { this.failure = failure; } + @Override public void injectMetrics(HttpServerExchange exchange, long startTime, String name, String endpoint) { throw failure.get(); } + @Override protected void createMetricsReporter(TimeSeriesDbSender sender, MetricsConfig config) { /* No reporter in this fixture. */ } + @Override public void handleRequest(HttpServerExchange exchange) { /* Lookup-only metrics fixture. */ } + @Override public HttpHandler getNext() { return null; } + @Override public MiddlewareHandler setNext(HttpHandler next) { return this; } + } + + private static final class Fixture implements AutoCloseable { + private final List servers = new ArrayList<>(); + URI start(HttpHandler handler) { return start(handler, false); } + URI start(HttpHandler handler, boolean http2) { + Undertow server = Undertow.builder().addHttpListener(0, "127.0.0.1").setIoThreads(1) + .setServerOption(UndertowOptions.ENABLE_HTTP2, http2).setHandler(handler).build(); + server.start(); + servers.add(server); + return URI.create("http://127.0.0.1:" + ((InetSocketAddress) server.getListenerInfo().getFirst().getAddress()).getPort()); + } + @Override public void close() { + for (int i = servers.size() - 1; i >= 0; i--) servers.get(i).stop(); + } + } +}