fix handling GatewayTimeout + fix flushing servlet request in async thread (#2018)

* fix handling GatewayTimeout + fix flushing servlet request in async thread

* Update spring-boot-admin-server/src/test/java/de/codecentric/boot/admin/server/web/AbstractInstancesProxyControllerIntegrationTest.java

Co-authored-by: Erik Petzold <[email protected]>

Author: Fuud

spring-boot-admin-server/src/main/java/de/codecentric/boot/admin/server/web/InstanceWebProxy.java

@@ -104,8 +104,11 @@ private <V> Mono<V> forward(Instance instance, ForwardRequest forwardRequest,
 			log.trace("No Endpoint found for Proxy-Request for instance {} with URL '{}'", instance.getId(),
 					forwardRequest.getUri());
 			return responseHandler.apply(ClientResponse.create(HttpStatus.NOT_FOUND, this.strategies).build());
-		}).onErrorResume(WebClientRequestException.class, (ex) -> {
-			Throwable cause = ex.getCause();
+		}).onErrorResume((ex) -> {
+			Throwable cause = ex;
+			if (ex instanceof WebClientRequestException) {
+				cause = ex.getCause();
+			}
 			if (cause instanceof ReadTimeoutException || cause instanceof TimeoutException) {
 				log.trace("Timeout for Proxy-Request for instance {} with URL '{}'", instance.getId(),
 						forwardRequest.getUri());

spring-boot-admin-server/src/main/java/de/codecentric/boot/admin/server/web/servlet/InstancesProxyController.java

@@ -21,6 +21,7 @@
 import java.net.URI;
 import java.util.Set;
 
+import javax.servlet.AsyncContext;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
@@ -84,24 +85,30 @@ public InstancesProxyController(String adminContextPath, Set<String> ignoredHead
 	@ResponseBody
 	@RequestMapping(path = INSTANCE_MAPPED_PATH, method = { RequestMethod.GET, RequestMethod.HEAD, RequestMethod.POST,
 			RequestMethod.PUT, RequestMethod.PATCH, RequestMethod.DELETE, RequestMethod.OPTIONS })
-	public void endpointProxy(@PathVariable("instanceId") String instanceId, HttpServletRequest servletRequest,
-			HttpServletResponse servletResponse) {
-		ServletServerHttpRequest request = new ServletServerHttpRequest(servletRequest);
-		Flux<DataBuffer> requestBody = DataBufferUtils.readInputStream(request::getBody, this.bufferFactory, 4096);
-		InstanceWebProxy.ForwardRequest fwdRequest = createForwardRequest(request, requestBody,
+	public void instanceProxy(@PathVariable("instanceId") String instanceId, HttpServletRequest servletRequest) {
+		// start async because we will commit from different thread.
+		// otherwise incorrect thread local objects (session and security context) will be stored.
+		// check for example org.springframework.security.web.context.HttpSessionSecurityContextRepository.SaveToSessionRequestWrapper.startAsync()
+		AsyncContext asyncContext = servletRequest.startAsync();
+		asyncContext.setTimeout(-1); // no timeout because instanceWebProxy will handle it for us
+		try {
+			ServletServerHttpRequest request = new ServletServerHttpRequest((HttpServletRequest) asyncContext.getRequest());
+			Flux<DataBuffer> requestBody = DataBufferUtils.readInputStream(request::getBody, this.bufferFactory, 4096);
+			InstanceWebProxy.ForwardRequest fwdRequest = createForwardRequest(request, requestBody,
 				this.adminContextPath + INSTANCE_MAPPED_PATH);
 
-		this.instanceWebProxy
+
+			this.instanceWebProxy
 				.forward(this.registry.getInstance(InstanceId.of(instanceId)), fwdRequest, (clientResponse) -> {
-					ServerHttpResponse response = new ServletServerHttpResponse(servletResponse);
+					ServerHttpResponse response = new ServletServerHttpResponse((HttpServletResponse) asyncContext.getResponse());
 					response.setStatusCode(clientResponse.statusCode());
 					response.getHeaders()
-							.addAll(this.httpHeadersFilter.filterHeaders(clientResponse.headers().asHttpHeaders()));
+						.addAll(this.httpHeadersFilter.filterHeaders(clientResponse.headers().asHttpHeaders()));
 					try {
 						OutputStream responseBody = response.getBody();
 						response.flush();
 						return clientResponse.body(BodyExtractors.toDataBuffers()).window(1)
-								.concatMap((body) -> writeAndFlush(body, responseBody)).then();
+							.concatMap((body) -> writeAndFlush(body, responseBody)).then();
 					}
 					catch (IOException ex) {
 						return Mono.error(ex);
@@ -111,6 +118,10 @@ public void endpointProxy(@PathVariable("instanceId") String instanceId, HttpSer
 				// before any async dispatch otherwise the FrameworkServlet will add wrong
 				// Allow header for OPTIONS request
 				.block();
+		}
+		finally {
+			asyncContext.complete();
+		}
 	}
 
 	@ResponseBody

spring-boot-admin-server/src/test/java/de/codecentric/boot/admin/server/web/AbstractInstancesProxyControllerIntegrationTest.java

@@ -72,7 +72,7 @@ public abstract class AbstractInstancesProxyControllerIntegrationTest {
 
 	@BeforeAll
 	public static void setUp() {
-		StepVerifier.setDefaultTimeout(Duration.ofSeconds(60));
+		StepVerifier.setDefaultTimeout(Duration.ofSeconds(600));
 	}
 
 	@AfterAll
@@ -124,10 +124,10 @@ public void should_return_status_502() {
 
 	@Test
 	public void should_return_status_504() {
-		// 502 on invalid response
-		this.client.get().uri("/instances/{instanceId}/actuator/invalid", this.instanceId)
-				.accept(new MediaType(ApiVersion.LATEST.getProducedMimeType())).exchange().expectStatus()
-				.isEqualTo(HttpStatus.BAD_GATEWAY);
+		// 504 on read timeout
+		this.client.get().uri("/instances/{instanceId}/actuator/timeout", this.instanceId)
+    .accept(new MediaType(ApiVersion.LATEST.getProducedMimeType()))
+		.exchange().expectStatus().isEqualTo(HttpStatus.GATEWAY_TIMEOUT);
 	}
 
 	@Test