Fix reactive connection handling.

Original Pull Request #1766
Closes #1759

(cherry picked from commit 58bca88)

Author: sothawo

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

@@ -134,7 +134,7 @@
  */
 public class DefaultReactiveElasticsearchClient implements ReactiveElasticsearchClient, Indices {
 
-	private final HostProvider hostProvider;
+	private final HostProvider<?> hostProvider;
 	private final RequestCreator requestCreator;
 	private Supplier<HttpHeaders> headersSupplier = () -> HttpHeaders.EMPTY;
 
@@ -144,7 +144,7 @@ public class DefaultReactiveElasticsearchClient implements ReactiveElasticsearch
 	 *
 	 * @param hostProvider must not be {@literal null}.
 	 */
-	public DefaultReactiveElasticsearchClient(HostProvider hostProvider) {
+	public DefaultReactiveElasticsearchClient(HostProvider<?> hostProvider) {
 		this(hostProvider, new DefaultRequestCreator());
 	}
 
@@ -155,7 +155,7 @@ public DefaultReactiveElasticsearchClient(HostProvider hostProvider) {
 	 * @param hostProvider must not be {@literal null}.
 	 * @param requestCreator must not be {@literal null}.
 	 */
-	public DefaultReactiveElasticsearchClient(HostProvider hostProvider, RequestCreator requestCreator) {
+	public DefaultReactiveElasticsearchClient(HostProvider<?> hostProvider, RequestCreator requestCreator) {
 
 		Assert.notNull(hostProvider, "HostProvider must not be null");
 		Assert.notNull(requestCreator, "RequestCreator must not be null");
@@ -639,8 +639,7 @@ public Mono<ClientResponse> execute(ReactiveElasticsearchClientCallback callback
 				.flatMap(callback::doWithClient) //
 				.onErrorResume(throwable -> {
 
-					if (throwable instanceof ConnectException) {
-
+					if (isCausedByConnectionException(throwable)) {
 						return hostProvider.getActive(Verification.ACTIVE) //
 								.flatMap(callback::doWithClient);
 					}
@@ -649,6 +648,27 @@ public Mono<ClientResponse> execute(ReactiveElasticsearchClientCallback callback
 				});
 	}
 
+	/**
+	 * checks if the given throwable is a {@link ConnectException} or has one in it's cause chain
+	 *
+	 * @param throwable the throwable to check
+	 * @return true if throwable is caused by a {@link ConnectException}
+	 */
+	private boolean isCausedByConnectionException(Throwable throwable) {
+
+		Throwable t = throwable;
+		do {
+
+			if (t instanceof ConnectException) {
+				return true;
+			}
+
+			t = t.getCause();
+		} while (t != null);
+
+		return false;
+	}
+
 	@Override
 	public Mono<Status> status() {
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2020 the original author or authors.
+ * Copyright 2018-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -34,9 +34,10 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
-public interface HostProvider {
+public interface HostProvider<T extends HostProvider<T>> {
 
 	/**
 	 * Create a new {@link HostProvider} best suited for the given {@link WebClientProvider} and number of hosts.
@@ -46,7 +47,7 @@ public interface HostProvider {
 	 * @param endpoints must not be {@literal null} nor empty.
 	 * @return new instance of {@link HostProvider}.
 	 */
-	static HostProvider provider(WebClientProvider clientProvider, Supplier<HttpHeaders> headersSupplier,
+	static HostProvider<?> provider(WebClientProvider clientProvider, Supplier<HttpHeaders> headersSupplier,
 			InetSocketAddress... endpoints) {
 
 		Assert.notNull(clientProvider, "WebClientProvider must not be null");

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2020 the original author or authors.
+ * Copyright 2018-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
 import reactor.util.function.Tuple2;
 
 import java.net.InetSocketAddress;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -30,6 +31,8 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.data.elasticsearch.client.ElasticsearchHost;
 import org.springframework.data.elasticsearch.client.ElasticsearchHost.State;
 import org.springframework.data.elasticsearch.client.NoReachableHostException;
@@ -42,22 +45,28 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
-class MultiNodeHostProvider implements HostProvider {
+class MultiNodeHostProvider implements HostProvider<MultiNodeHostProvider> {
+
+	private final static Logger LOG = LoggerFactory.getLogger(MultiNodeHostProvider.class);
 
 	private final WebClientProvider clientProvider;
 	private final Supplier<HttpHeaders> headersSupplier;
 	private final Map<InetSocketAddress, ElasticsearchHost> hosts;
 
-	MultiNodeHostProvider(WebClientProvider clientProvider, Supplier<HttpHeaders> headersSupplier, InetSocketAddress... endpoints) {
+	MultiNodeHostProvider(WebClientProvider clientProvider, Supplier<HttpHeaders> headersSupplier,
+			InetSocketAddress... endpoints) {
 
 		this.clientProvider = clientProvider;
 		this.headersSupplier = headersSupplier;
 		this.hosts = new ConcurrentHashMap<>();
 		for (InetSocketAddress endpoint : endpoints) {
 			this.hosts.put(endpoint, new ElasticsearchHost(endpoint, State.UNKNOWN));
 		}
+
+		LOG.debug("initialized with " + hosts);
 	}
 
 	/*
@@ -66,7 +75,7 @@ class MultiNodeHostProvider implements HostProvider {
 	 */
 	@Override
 	public Mono<ClusterInformation> clusterInfo() {
-		return nodes(null).map(this::updateNodeState).buffer(hosts.size())
+		return checkNodes(null).map(this::updateNodeState).buffer(hosts.size())
 				.then(Mono.just(new ClusterInformation(new LinkedHashSet<>(this.hosts.values()))));
 	}
 
@@ -86,14 +95,19 @@ public WebClient createWebClient(InetSocketAddress endpoint) {
 	@Override
 	public Mono<InetSocketAddress> lookupActiveHost(Verification verification) {
 
+		LOG.trace("lookupActiveHost " + verification + " from " + hosts());
+
 		if (Verification.LAZY.equals(verification)) {
 			for (ElasticsearchHost entry : hosts()) {
 				if (entry.isOnline()) {
+					LOG.trace("lookupActiveHost returning " + entry);
 					return Mono.just(entry.getEndpoint());
 				}
 			}
+			LOG.trace("no online host found with LAZY");
 		}
 
+		LOG.trace("searching for active host");
 		return findActiveHostInKnownActives() //
 				.switchIfEmpty(findActiveHostInUnresolved()) //
 				.switchIfEmpty(findActiveHostInDead()) //
@@ -105,20 +119,30 @@ Collection<ElasticsearchHost> getCachedHostState() {
 	}
 
 	private Mono<InetSocketAddress> findActiveHostInKnownActives() {
-		return findActiveForSate(State.ONLINE);
+		return findActiveForState(State.ONLINE);
 	}
 
 	private Mono<InetSocketAddress> findActiveHostInUnresolved() {
-		return findActiveForSate(State.UNKNOWN);
+		return findActiveForState(State.UNKNOWN);
 	}
 
 	private Mono<InetSocketAddress> findActiveHostInDead() {
-		return findActiveForSate(State.OFFLINE);
+		return findActiveForState(State.OFFLINE);
 	}
 
-	private Mono<InetSocketAddress> findActiveForSate(State state) {
-		return nodes(state).map(this::updateNodeState).filter(ElasticsearchHost::isOnline)
-				.map(ElasticsearchHost::getEndpoint).next();
+	private Mono<InetSocketAddress> findActiveForState(State state) {
+
+		LOG.trace("findActiveForState state " + state + ", current hosts: " + hosts);
+
+		return checkNodes(state) //
+				.map(this::updateNodeState) //
+				.filter(ElasticsearchHost::isOnline) //
+				.map(elasticsearchHost -> {
+					LOG.trace("findActiveForState returning host " + elasticsearchHost);
+					return elasticsearchHost;
+				}).map(ElasticsearchHost::getEndpoint) //
+				.takeLast(1) //
+				.next();
 	}
 
 	private ElasticsearchHost updateNodeState(Tuple2<InetSocketAddress, ClientResponse> tuple2) {
@@ -129,17 +153,19 @@ private ElasticsearchHost updateNodeState(Tuple2<InetSocketAddress, ClientRespon
 		return elasticsearchHost;
 	}
 
-	private Flux<Tuple2<InetSocketAddress, ClientResponse>> nodes(@Nullable State state) {
+	private Flux<Tuple2<InetSocketAddress, ClientResponse>> checkNodes(@Nullable State state) {
 
 		return Flux.fromIterable(hosts()) //
 				.filter(entry -> state == null || entry.getState().equals(state)) //
 				.map(ElasticsearchHost::getEndpoint) //
-				.flatMap(host -> {
+				.concatMap(host -> {
 
 					Mono<ClientResponse> exchange = createWebClient(host) //
 							.head().uri("/") //
 							.headers(httpHeaders -> httpHeaders.addAll(headersSupplier.get())) //
-							.exchange().doOnError(throwable -> {
+							.exchange() //
+							.timeout(Duration.ofSeconds(1)) //
+							.doOnError(throwable -> {
 								hosts.put(host, new ElasticsearchHost(host, State.OFFLINE));
 								clientProvider.getErrorListener().accept(throwable);
 							});

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2020 the original author or authors.
+ * Copyright 2018-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,9 +32,10 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
-class SingleNodeHostProvider implements HostProvider {
+class SingleNodeHostProvider implements HostProvider<SingleNodeHostProvider> {
 
 	private final WebClientProvider clientProvider;
 	private final Supplier<HttpHeaders> headersSupplier;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2020 the original author or authors.
+ * Copyright 2018-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -62,7 +62,7 @@ public class ReactiveElasticsearchClientUnitTests {
 
 	static final String HOST = ":9200";
 
-	MockDelegatingElasticsearchHostProvider<HostProvider> hostProvider;
+	MockDelegatingElasticsearchHostProvider<? extends HostProvider<?>> hostProvider;
 	ReactiveElasticsearchClient client;
 
 	@BeforeEach

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

@@ -186,7 +186,7 @@ public T getDelegate() {
 			return delegate;
 		}
 
-		public MockDelegatingElasticsearchHostProvider<T> withActiveDefaultHost(String host) {
+		public MockDelegatingElasticsearchHostProvider<? extends HostProvider<?>> withActiveDefaultHost(String host) {
 			return new MockDelegatingElasticsearchHostProvider(HttpHeaders.EMPTY, clientProvider, errorCollector, delegate,
 					host);
 		}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2020 the original author or authors.
+ * Copyright 2018-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -30,6 +30,7 @@
 
 /**
  * @author Christoph Strobl
+ * @author Peter-Josef Meisch
  */
 public class SingleNodeHostProviderUnitTests {