From ad6b01d6d0b25e6de62363730f1f6b127f84b481 Mon Sep 17 00:00:00 2001
From: Stephane Nicoll <snicoll@vmware.com>
Date: Tue, 22 Dec 2020 15:20:44 +0100
Subject: [PATCH] Add control connection timeout property for Cassandra

Closes gh-24189
---
 .../cassandra/CassandraAutoConfiguration.java |  9 +++++++
 .../cassandra/CassandraProperties.java        | 26 +++++++++++++++++++
 .../CassandraAutoConfigurationTests.java      | 10 +++++++
 3 files changed, 45 insertions(+)

diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfiguration.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfiguration.java
index 7e267f8f74..bf10929eea 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfiguration.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfiguration.java
@@ -40,6 +40,7 @@ import com.typesafe.config.ConfigFactory;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.cassandra.CassandraProperties.Connection;
+import org.springframework.boot.autoconfigure.cassandra.CassandraProperties.Controlconnection;
 import org.springframework.boot.autoconfigure.cassandra.CassandraProperties.Request;
 import org.springframework.boot.autoconfigure.cassandra.CassandraProperties.Throttler;
 import org.springframework.boot.autoconfigure.cassandra.CassandraProperties.ThrottlerType;
@@ -127,6 +128,7 @@ public class CassandraAutoConfiguration {
 		mapConnectionOptions(properties, options);
 		mapPoolingOptions(properties, options);
 		mapRequestOptions(properties, options);
+		mapControlConnectionOptions(properties, options);
 		map.from(mapContactPoints(properties))
 				.to((contactPoints) -> options.add(DefaultDriverOption.CONTACT_POINTS, contactPoints));
 		map.from(properties.getLocalDatacenter()).to(
@@ -178,6 +180,13 @@ public class CassandraAutoConfiguration {
 				(drainInterval) -> options.add(DefaultDriverOption.REQUEST_THROTTLER_DRAIN_INTERVAL, drainInterval));
 	}
 
+	private void mapControlConnectionOptions(CassandraProperties properties, CassandraDriverOptions options) {
+		PropertyMapper map = PropertyMapper.get().alwaysApplyingWhenNonNull();
+		Controlconnection controlProperties = properties.getControlconnection();
+		map.from(controlProperties::getTimeout).asInt(Duration::toMillis)
+				.to((timeout) -> options.add(DefaultDriverOption.CONTROL_CONNECTION_TIMEOUT, timeout));
+	}
+
 	private List<String> mapContactPoints(CassandraProperties properties) {
 		return properties.getContactPoints().stream()
 				.map((candidate) -> formatContactPoint(candidate, properties.getPort())).collect(Collectors.toList());
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraProperties.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraProperties.java
index df6a983746..6facba1a43 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraProperties.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/cassandra/CassandraProperties.java
@@ -105,6 +105,11 @@ public class CassandraProperties {
 	 */
 	private final Request request = new Request();
 
+	/**
+	 * Control connection configuration.
+	 */
+	private final Controlconnection controlconnection = new Controlconnection();
+
 	public String getKeyspaceName() {
 		return this.keyspaceName;
 	}
@@ -259,6 +264,10 @@ public class CassandraProperties {
 		return this.request;
 	}
 
+	public Controlconnection getControlconnection() {
+		return this.controlconnection;
+	}
+
 	public static class Connection {
 
 		/**
@@ -386,6 +395,23 @@ public class CassandraProperties {
 
 	}
 
+	public static class Controlconnection {
+
+		/**
+		 * Timeout to use for control queries.
+		 */
+		private Duration timeout = Duration.ofSeconds(5);
+
+		public Duration getTimeout() {
+			return this.timeout;
+		}
+
+		public void setTimeout(Duration timeout) {
+			this.timeout = timeout;
+		}
+
+	}
+
 	public static class Throttler {
 
 		/**
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfigurationTests.java
index e2e041b272..9851f237c2 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfigurationTests.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/cassandra/CassandraAutoConfigurationTests.java
@@ -166,6 +166,16 @@ class CassandraAutoConfigurationTests {
 				});
 	}
 
+	@Test
+	void driverConfigLoaderCustomizeControlConnectionOptions() {
+		this.contextRunner.withPropertyValues("spring.data.cassandra.controlconnection.timeout=200ms")
+				.run((context) -> {
+					DriverExecutionProfile config = context.getBean(DriverConfigLoader.class).getInitialConfig()
+							.getDefaultProfile();
+					assertThat(config.getInt(DefaultDriverOption.CONTROL_CONNECTION_TIMEOUT)).isEqualTo(200);
+				});
+	}
+
 	@Test
 	void driverConfigLoaderUsePassThroughLimitingRequestThrottlerByDefault() {
 		this.contextRunner.withPropertyValues().run((context) -> {