DATAES-870 - Consume response body to release connection directly.

We now ensure that response bodies from ClientResponse get released as part of our result handling. This is to prevent cancel signals issuing the connection release so that the connection release can be synchronized (awaited) before any subsequent requests get issued.

Connection release should be part of the Framework but the fallback interferes with Reactor Netty's HttpClient therefore we're ensuring proper resource disposal.

Author: mp911de

src/main/java/org/springframework/data/elasticsearch/client/reactive/DefaultReactiveElasticsearchClient.java

@@ -240,11 +240,7 @@ private static WebClientProvider getWebClientProvider(ClientConfiguration client
 		Duration connectTimeout = clientConfiguration.getConnectTimeout();
 		Duration soTimeout = clientConfiguration.getSocketTimeout();
 
-		// DATAES-870: previously: HttpClient httpClient = HttpClient.create();
-		// disable pooling of connections (see https://github.com/spring-projects/spring-framework/issues/22464)
-		// otherwise we get errors: "Connection prematurely closed BEFORE response; nested exception is
-		// java.lang.RuntimeException: Connection prematurely closed BEFORE response"
-		HttpClient httpClient = HttpClient.newConnection().compress(true);
+		HttpClient httpClient = HttpClient.create().compress(true);
 
 		if (!connectTimeout.isNegative()) {
 			httpClient = httpClient.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, Math.toIntExact(connectTimeout.toMillis()));
@@ -305,7 +301,7 @@ private static WebClientProvider getWebClientProvider(ClientConfiguration client
 	public Mono<Boolean> ping(HttpHeaders headers) {
 
 		return sendRequest(new MainRequest(), requestCreator.ping(), RawActionResponse.class, headers) //
-				.map(response -> response.statusCode().is2xxSuccessful()) //
+				.flatMap(response -> response.releaseBody().thenReturn(response.statusCode().is2xxSuccessful())) //
 				.onErrorResume(NoReachableHostException.class, error -> Mono.just(false)).next();
 	}
 
@@ -355,7 +351,7 @@ public Flux<GetResult> multiGet(HttpHeaders headers, MultiGetRequest multiGetReq
 	public Mono<Boolean> exists(HttpHeaders headers, GetRequest getRequest) {
 
 		return sendRequest(getRequest, requestCreator.exists(), RawActionResponse.class, headers) //
-				.map(response -> response.statusCode().is2xxSuccessful()) //
+				.flatMap(response -> response.releaseBody().thenReturn(response.statusCode().is2xxSuccessful())) //
 				.next();
 	}
 
@@ -555,7 +551,7 @@ public Mono<BulkResponse> bulk(HttpHeaders headers, BulkRequest bulkRequest) {
 	public Mono<Boolean> existsIndex(HttpHeaders headers, GetIndexRequest request) {
 
 		return sendRequest(request, requestCreator.indexExists(), RawActionResponse.class, headers) //
-				.map(response -> response.statusCode().is2xxSuccessful()) //
+				.flatMap(response -> response.releaseBody().thenReturn(response.statusCode().is2xxSuccessful())) //
 				.next();
 	}
 

src/main/java/org/springframework/data/elasticsearch/client/reactive/MultiNodeHostProvider.java

@@ -15,7 +15,6 @@
  */
 package org.springframework.data.elasticsearch.client.reactive;
 
-import org.springframework.http.HttpHeaders;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.util.function.Tuple2;
@@ -33,6 +32,7 @@
 import org.springframework.data.elasticsearch.client.ElasticsearchHost;
 import org.springframework.data.elasticsearch.client.ElasticsearchHost.State;
 import org.springframework.data.elasticsearch.client.NoReachableHostException;
+import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
 import org.springframework.web.reactive.function.client.ClientResponse;
 import org.springframework.web.reactive.function.client.WebClient;
@@ -121,15 +121,15 @@ private Mono<InetSocketAddress> findActiveForSate(State state) {
 				.map(ElasticsearchHost::getEndpoint).next();
 	}
 
-	private ElasticsearchHost updateNodeState(Tuple2<InetSocketAddress, ClientResponse> tuple2) {
+	private ElasticsearchHost updateNodeState(Tuple2<InetSocketAddress, State> tuple2) {
 
-		State state = tuple2.getT2().statusCode().isError() ? State.OFFLINE : State.ONLINE;
+		State state = tuple2.getT2();
 		ElasticsearchHost elasticsearchHost = new ElasticsearchHost(tuple2.getT1(), state);
 		hosts.put(tuple2.getT1(), elasticsearchHost);
 		return elasticsearchHost;
 	}
 
-	private Flux<Tuple2<InetSocketAddress, ClientResponse>> nodes(@Nullable State state) {
+	private Flux<Tuple2<InetSocketAddress, State>> nodes(@Nullable State state) {
 
 		return Flux.fromIterable(hosts()) //
 				.filter(entry -> state == null || entry.getState().equals(state)) //
@@ -144,7 +144,8 @@ private Flux<Tuple2<InetSocketAddress, ClientResponse>> nodes(@Nullable State st
 								clientProvider.getErrorListener().accept(throwable);
 							});
 
-					return Mono.just(host).zipWith(exchange);
+					return Mono.just(host).zipWith(exchange
+							.flatMap(it -> it.releaseBody().thenReturn(it.statusCode().isError() ? State.OFFLINE : State.ONLINE)));
 				}) //
 				.onErrorContinue((throwable, o) -> clientProvider.getErrorListener().accept(throwable));
 	}

src/main/java/org/springframework/data/elasticsearch/client/reactive/RawActionResponse.java

@@ -15,11 +15,12 @@
  */
 package org.springframework.data.elasticsearch.client.reactive;
 
-import org.elasticsearch.common.io.stream.StreamOutput;
+import reactor.core.publisher.Mono;
 
 import java.io.IOException;
 
 import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.common.io.stream.StreamOutput;
 
 import org.springframework.http.HttpStatus;
 import org.springframework.http.client.reactive.ClientHttpResponse;
@@ -73,4 +74,13 @@ public <T> T body(BodyExtractor<T, ? super ClientHttpResponse> extractor) {
 	@Override
 	public void writeTo(StreamOutput out) throws IOException {
 	}
+
+	/**
+	 * Ensure the response body is released to properly release the underlying connection.
+	 *
+	 * @return
+	 */
+	public Mono<Void> releaseBody() {
+		return delegate.releaseBody();
+	}
 }

src/main/java/org/springframework/data/elasticsearch/client/reactive/SingleNodeHostProvider.java

@@ -60,13 +60,13 @@ public Mono<ClusterInformation> clusterInfo() {
 				.head().uri("/")
 				.headers(httpHeaders -> httpHeaders.addAll(headersSupplier.get())) //
 				.exchange() //
-				.map(it -> {
+				.flatMap(it -> {
 					if (it.statusCode().isError()) {
 						state = ElasticsearchHost.offline(endpoint);
 					} else {
 						state = ElasticsearchHost.online(endpoint);
 					}
-					return state;
+					return it.releaseBody().thenReturn(state);
 				}).onErrorResume(throwable -> {
 
 					state = ElasticsearchHost.offline(endpoint);

src/test/java/org/springframework/data/elasticsearch/client/reactive/MultiNodeHostProviderUnitTests.java

@@ -117,9 +117,11 @@ public void triesDeadHostsIfNoActiveFound() {
 		mock.when(HOST_2).get(requestHeadersUriSpec -> {
 
 			ClientResponse response1 = mock(ClientResponse.class);
+			when(response1.releaseBody()).thenReturn(Mono.empty());
 			Receive.error(response1);
 
 			ClientResponse response2 = mock(ClientResponse.class);
+			when(response2.releaseBody()).thenReturn(Mono.empty());
 			Receive.ok(response2);
 
 			when(requestHeadersUriSpec.exchange()).thenReturn(Mono.just(response1), Mono.just(response2));

src/test/java/org/springframework/data/elasticsearch/client/reactive/ReactiveMockClientTestsUtils.java

@@ -243,6 +243,7 @@ public WebClient get(InetSocketAddress endpoint) {
 					Mockito.when(headersUriSpec.exchange()).thenReturn(Mono.just(response));
 					Mockito.when(bodySpy.exchange()).thenReturn(Mono.just(response));
 					Mockito.when(response.statusCode()).thenReturn(HttpStatus.ACCEPTED);
+					Mockito.when(response.releaseBody()).thenReturn(Mono.empty());
 
 					headersUriSpecMap.putIfAbsent(key, headersUriSpec);
 					bodyUriSpecMap.putIfAbsent(key, bodySpy);