Merge pull request #20994 from vpavic

* improve-es-config:
  Polish
  Improve Elasticsearch RestClient customization capabilities

Closes gh-20994

Author: bclozel

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/elasticsearch/ElasticsearchRestClientConfigurations.java

@@ -23,7 +23,9 @@
 import org.apache.http.auth.Credentials;
 import org.apache.http.auth.UsernamePasswordCredentials;
 import org.apache.http.client.CredentialsProvider;
+import org.apache.http.client.config.RequestConfig;
 import org.apache.http.impl.client.BasicCredentialsProvider;
+import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.client.RestClientBuilder;
 import org.elasticsearch.client.RestHighLevelClient;
@@ -40,32 +42,30 @@
  *
  * @author Brian Clozel
  * @author Stephane Nicoll
+ * @author Vedran Pavic
  */
 class ElasticsearchRestClientConfigurations {
 
 	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnMissingBean(RestClientBuilder.class)
 	static class RestClientBuilderConfiguration {
 
 		@Bean
-		@ConditionalOnMissingBean
+		RestClientBuilderCustomizer defaultRestClientBuilderCustomizer(ElasticsearchRestClientProperties properties) {
+			return new DefaultRestClientBuilderCustomizer(properties);
+		}
+
+		@Bean
 		RestClientBuilder elasticsearchRestClientBuilder(ElasticsearchRestClientProperties properties,
 				ObjectProvider<RestClientBuilderCustomizer> builderCustomizers) {
 			HttpHost[] hosts = properties.getUris().stream().map(HttpHost::create).toArray(HttpHost[]::new);
 			RestClientBuilder builder = RestClient.builder(hosts);
-			PropertyMapper map = PropertyMapper.get();
-			map.from(properties::getUsername).whenHasText().to((username) -> {
-				CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
-				Credentials credentials = new UsernamePasswordCredentials(properties.getUsername(),
-						properties.getPassword());
-				credentialsProvider.setCredentials(AuthScope.ANY, credentials);
-				builder.setHttpClientConfigCallback(
-						(httpClientBuilder) -> httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
+			builder.setHttpClientConfigCallback((httpClientBuilder) -> {
+				builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(httpClientBuilder));
+				return httpClientBuilder;
 			});
 			builder.setRequestConfigCallback((requestConfigBuilder) -> {
-				map.from(properties::getConnectionTimeout).whenNonNull().asInt(Duration::toMillis)
-						.to(requestConfigBuilder::setConnectTimeout);
-				map.from(properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
-						.to(requestConfigBuilder::setSocketTimeout);
+				builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(requestConfigBuilder));
 				return requestConfigBuilder;
 			});
 			builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(builder));
@@ -108,4 +108,39 @@ RestClient elasticsearchRestClient(RestClientBuilder builder) {
 
 	}
 
+	static class DefaultRestClientBuilderCustomizer implements RestClientBuilderCustomizer {
+
+		private static final PropertyMapper map = PropertyMapper.get();
+
+		private final ElasticsearchRestClientProperties properties;
+
+		DefaultRestClientBuilderCustomizer(ElasticsearchRestClientProperties properties) {
+			this.properties = properties;
+		}
+
+		@Override
+		public void customize(RestClientBuilder builder) {
+		}
+
+		@Override
+		public void customize(HttpAsyncClientBuilder builder) {
+			map.from(this.properties::getUsername).whenHasText().to((username) -> {
+				CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
+				Credentials credentials = new UsernamePasswordCredentials(this.properties.getUsername(),
+						this.properties.getPassword());
+				credentialsProvider.setCredentials(AuthScope.ANY, credentials);
+				builder.setDefaultCredentialsProvider(credentialsProvider);
+			});
+		}
+
+		@Override
+		public void customize(RequestConfig.Builder builder) {
+			map.from(this.properties::getConnectionTimeout).whenNonNull().asInt(Duration::toMillis)
+					.to(builder::setConnectTimeout);
+			map.from(this.properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
+					.to(builder::setSocketTimeout);
+		}
+
+	}
+
 }

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/elasticsearch/RestClientBuilderCustomizer.java

@@ -16,6 +16,8 @@
 
 package org.springframework.boot.autoconfigure.elasticsearch;
 
+import org.apache.http.client.config.RequestConfig;
+import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
 import org.elasticsearch.client.RestClientBuilder;
 
 /**
@@ -24,15 +26,37 @@
  * retaining default auto-configuration.
  *
  * @author Brian Clozel
+ * @author Vedran Pavic
  * @since 2.1.0
  */
 @FunctionalInterface
 public interface RestClientBuilderCustomizer {
 
 	/**
 	 * Customize the {@link RestClientBuilder}.
+	 * <p>
+	 * Possibly overrides customizations made with the {@code "spring.elasticsearch.rest"}
+	 * configuration properties namespace. For more targeted changes, see
+	 * {@link #customize(HttpAsyncClientBuilder)} and
+	 * {@link #customize(RequestConfig.Builder)}.
 	 * @param builder the builder to customize
 	 */
 	void customize(RestClientBuilder builder);
 
+	/**
+	 * Customize the {@link HttpAsyncClientBuilder}.
+	 * @param builder the builder
+	 * @since 2.3.0
+	 */
+	default void customize(HttpAsyncClientBuilder builder) {
+	}
+
+	/**
+	 * Customize the {@link RequestConfig.Builder}.
+	 * @param builder the builder
+	 * @since 2.3.0
+	 */
+	default void customize(RequestConfig.Builder builder) {
+	}
+
 }

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/elasticsearch/ElasticsearchRestClientAutoConfigurationTests.java

@@ -20,6 +20,8 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import org.apache.http.client.config.RequestConfig;
+import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
 import org.elasticsearch.action.get.GetRequest;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.client.RequestOptions;
@@ -36,7 +38,6 @@
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.test.util.ReflectionTestUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
@@ -45,6 +46,7 @@
  * Tests for {@link ElasticsearchRestClientAutoConfiguration}.
  *
  * @author Brian Clozel
+ * @author Vedran Pavic
  */
 @Testcontainers(disabledWithoutDocker = true)
 class ElasticsearchRestClientAutoConfigurationTests {
@@ -53,7 +55,7 @@ class ElasticsearchRestClientAutoConfigurationTests {
 	static final ElasticsearchContainer elasticsearch = new ElasticsearchContainer().withStartupAttempts(5)
 			.withStartupTimeout(Duration.ofMinutes(10));
 
-	private ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
 			.withConfiguration(AutoConfigurations.of(ElasticsearchRestClientAutoConfiguration.class));
 
 	@Test
@@ -106,6 +108,8 @@ void configureWhenBuilderCustomizerShouldApply() {
 			assertThat(context).hasSingleBean(RestClient.class);
 			RestClient restClient = context.getBean(RestClient.class);
 			assertThat(restClient).hasFieldOrPropertyWithValue("pathPrefix", "/test");
+			assertThat(restClient).extracting("client.connmgr.pool.maxTotal").isEqualTo(100);
+			assertThat(restClient).extracting("client.defaultConfig.cookieSpec").isEqualTo("rfc6265-lax");
 		});
 	}
 
@@ -130,10 +134,10 @@ void configureWithCustomTimeouts() {
 	}
 
 	private static void assertTimeouts(RestClient restClient, Duration connectTimeout, Duration readTimeout) {
-		Object client = ReflectionTestUtils.getField(restClient, "client");
-		Object config = ReflectionTestUtils.getField(client, "defaultConfig");
-		assertThat(config).hasFieldOrPropertyWithValue("socketTimeout", Math.toIntExact(readTimeout.toMillis()));
-		assertThat(config).hasFieldOrPropertyWithValue("connectTimeout", Math.toIntExact(connectTimeout.toMillis()));
+		assertThat(restClient).extracting("client.defaultConfig.socketTimeout")
+				.isEqualTo(Math.toIntExact(readTimeout.toMillis()));
+		assertThat(restClient).extracting("client.defaultConfig.connectTimeout")
+				.isEqualTo(Math.toIntExact(connectTimeout.toMillis()));
 	}
 
 	@Test
@@ -167,7 +171,24 @@ static class BuilderCustomizerConfiguration {
 
 		@Bean
 		RestClientBuilderCustomizer myCustomizer() {
-			return (builder) -> builder.setPathPrefix("/test");
+			return new RestClientBuilderCustomizer() {
+
+				@Override
+				public void customize(RestClientBuilder builder) {
+					builder.setPathPrefix("/test");
+				}
+
+				@Override
+				public void customize(HttpAsyncClientBuilder builder) {
+					builder.setMaxConnTotal(100);
+				}
+
+				@Override
+				public void customize(RequestConfig.Builder builder) {
+					builder.setCookieSpec("rfc6265-lax");
+				}
+
+			};
 		}
 
 	}