From adb2b0a775f3f4ebac98876066c9ce0ee7a68fe6 Mon Sep 17 00:00:00 2001
From: Asmir Mustafic <asmir.mustafic@getyourguide.com>
Date: Thu, 12 Dec 2024 22:32:09 +0100
Subject: [PATCH] Allow to prepare the pool

---
 .../DefaultLettucePoolingClientConfiguration.java  | 14 ++++++++++++++
 .../lettuce/LettucePoolingClientConfiguration.java | 10 +++++++++-
 .../lettuce/LettucePoolingConnectionProvider.java  | 10 ++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/springframework/data/redis/connection/lettuce/DefaultLettucePoolingClientConfiguration.java b/src/main/java/org/springframework/data/redis/connection/lettuce/DefaultLettucePoolingClientConfiguration.java
index 6a89a2bd3d..16477aff93 100644
--- a/src/main/java/org/springframework/data/redis/connection/lettuce/DefaultLettucePoolingClientConfiguration.java
+++ b/src/main/java/org/springframework/data/redis/connection/lettuce/DefaultLettucePoolingClientConfiguration.java
@@ -38,12 +38,26 @@ class DefaultLettucePoolingClientConfiguration implements LettucePoolingClientCo
 
 	private final LettuceClientConfiguration clientConfiguration;
 	private final GenericObjectPoolConfig poolConfig;
+  private final boolean preparePool;
 
 	DefaultLettucePoolingClientConfiguration(LettuceClientConfiguration clientConfiguration,
 			GenericObjectPoolConfig poolConfig) {
 
 		this.clientConfiguration = clientConfiguration;
 		this.poolConfig = poolConfig;
+    this.preparePool = false;
+	}
+
+  DefaultLettucePoolingClientConfiguration(LettuceClientConfiguration clientConfiguration,
+                                           GenericObjectPoolConfig poolConfig, boolean preparePool) {
+    this.clientConfiguration = clientConfiguration;
+		this.poolConfig = poolConfig;
+    this.preparePool = preparePool;
+  }
+
+  @Override
+	public boolean preparePool() {
+		return preparePool;
 	}
 
 	@Override
diff --git a/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingClientConfiguration.java b/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingClientConfiguration.java
index c9fa6c242e..a4155dcef6 100644
--- a/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingClientConfiguration.java
+++ b/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingClientConfiguration.java
@@ -36,6 +36,7 @@
  */
 public interface LettucePoolingClientConfiguration extends LettuceClientConfiguration {
 
+	boolean preparePool();
 	/**
 	 * @return the {@link GenericObjectPoolConfig}. Never {@literal null}.
 	 */
@@ -92,6 +93,7 @@ static LettucePoolingClientConfiguration defaultConfiguration() {
 	class LettucePoolingClientConfigurationBuilder extends LettuceClientConfigurationBuilder {
 
 		GenericObjectPoolConfig poolConfig = new GenericObjectPoolConfig();
+		boolean preparePool = false;
 
 		LettucePoolingClientConfigurationBuilder() {
 			super();
@@ -125,6 +127,12 @@ public LettucePoolingClientConfigurationBuilder commandTimeout(Duration timeout)
 			return this;
 		}
 
+		public LettucePoolingClientConfigurationBuilder preparePool(boolean preparePool) {
+
+			this.preparePool = preparePool;
+			return this;
+    }
+
 		@Override
 		public LettucePoolingClientConfigurationBuilder shutdownTimeout(Duration shutdownTimeout) {
 
@@ -173,7 +181,7 @@ public LettucePoolingClientConfigurationBuilder poolConfig(GenericObjectPoolConf
 
 		@Override
 		public LettucePoolingClientConfiguration build() {
-			return new DefaultLettucePoolingClientConfiguration(super.build(), poolConfig);
+			return new DefaultLettucePoolingClientConfiguration(super.build(), poolConfig, preparePool);
 		}
 	}
 
diff --git a/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingConnectionProvider.java b/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingConnectionProvider.java
index 79e6b5e836..5b9a042335 100644
--- a/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingConnectionProvider.java
+++ b/src/main/java/org/springframework/data/redis/connection/lettuce/LettucePoolingConnectionProvider.java
@@ -63,6 +63,7 @@ class LettucePoolingConnectionProvider implements LettuceConnectionProvider, Red
 	private static final Log log = LogFactory.getLog(LettucePoolingConnectionProvider.class);
 
 	private final LettuceConnectionProvider connectionProvider;
+  private final LettucePoolingClientConfiguration clientConfiguration;
 	private final GenericObjectPoolConfig poolConfig;
 	private final Map<StatefulConnection<?, ?>, GenericObjectPool<StatefulConnection<?, ?>>> poolRef = new ConcurrentHashMap<>(
 			32);
@@ -81,6 +82,7 @@ class LettucePoolingConnectionProvider implements LettuceConnectionProvider, Red
 		Assert.notNull(connectionProvider, "ConnectionProvider must not be null");
 		Assert.notNull(clientConfiguration, "ClientConfiguration must not be null");
 
+		this.clientConfiguration = clientConfiguration;
 		this.connectionProvider = connectionProvider;
 		this.poolConfig = clientConfiguration.getPoolConfig();
 		this.asyncPoolConfig = CommonsPool2ConfigConverter.bounded(this.poolConfig);
@@ -93,6 +95,14 @@ class LettucePoolingConnectionProvider implements LettuceConnectionProvider, Red
 			return ConnectionPoolSupport.createGenericObjectPool(() -> connectionProvider.getConnection(connectionType),
 					poolConfig, false);
 		});
+		if (clientConfiguration.preparePool()) {
+			try {
+				pool.preparePool();
+
+			} catch (Exception ex) {
+				throw new PoolException("Could not prepare the pool", ex);
+			}
+    }
 
 		try {