Skip to content

Commit 064d3f7

Browse files
big-cirbig-cir
committed
Add custom connection validation to ConnectionPoolSupport
1 parent 319e315 commit 064d3f7

File tree

3 files changed

+125
-15
lines changed

3 files changed

+125
-15
lines changed

‎src/main/java/io/lettuce/core/support/ConnectionPoolSupport.java

+48-4
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import java.util.concurrent.atomic.AtomicReference;
77
import java.util.concurrent.locks.Lock;
88
import java.util.concurrent.locks.ReentrantLock;
9+
import java.util.function.Predicate;
910
import java.util.function.Supplier;
1011

1112
import org.apache.commons.pool2.BasePooledObjectFactory;
@@ -60,6 +61,7 @@
6061
* </pre>
6162
*
6263
* @author Mark Paluch
64+
* @author dae won
6365
* @since 4.3
6466
*/
6567
public abstract class ConnectionPoolSupport {
@@ -77,8 +79,8 @@ private ConnectionPoolSupport() {
7779
* @return the connection pool.
7880
*/
7981
public static <T extends StatefulConnection<?, ?>> GenericObjectPool<T> createGenericObjectPool(
80-
Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> config) {
81-
return createGenericObjectPool(connectionSupplier, config, true);
82+
Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> config, Predicate<T> connectionValidator) {
83+
return createGenericObjectPool(connectionSupplier, config, true, connectionValidator);
8284
}
8385

8486
/**
@@ -94,14 +96,17 @@ private ConnectionPoolSupport() {
9496
*/
9597
@SuppressWarnings("unchecked")
9698
public static <T extends StatefulConnection<?, ?>> GenericObjectPool<T> createGenericObjectPool(
97-
Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> config, boolean wrapConnections) {
99+
Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> config, boolean wrapConnections,
100+
Predicate<T> connectionValidator) {
98101

99102
LettuceAssert.notNull(connectionSupplier, "Connection supplier must not be null");
100103
LettuceAssert.notNull(config, "GenericObjectPoolConfig must not be null");
104+
LettuceAssert.notNull(connectionValidator, "Connection validator must not be null");
101105

102106
AtomicReference<Origin<T>> poolRef = new AtomicReference<>();
103107

104-
GenericObjectPool<T> pool = new GenericObjectPool<T>(new RedisPooledObjectFactory<T>(connectionSupplier), config) {
108+
GenericObjectPool<T> pool = new GenericObjectPool<T>(
109+
new EnhancedRedisPooledObjectFactory<T>(connectionSupplier, connectionValidator), config) {
105110

106111
@Override
107112
public T borrowObject() throws Exception {
@@ -249,4 +254,43 @@ public CompletableFuture<Void> returnObjectAsync(T o) throws Exception {
249254

250255
}
251256

257+
private static class EnhancedRedisPooledObjectFactory<T extends StatefulConnection<?, ?>>
258+
extends BasePooledObjectFactory<T> {
259+
260+
private final Supplier<T> connectionSupplier;
261+
262+
private final Predicate<T> connectionValidator;
263+
264+
EnhancedRedisPooledObjectFactory(Supplier<T> connectionSupplier, Predicate<T> connectionValidator) {
265+
this.connectionSupplier = connectionSupplier;
266+
this.connectionValidator = connectionValidator;
267+
}
268+
269+
@Override
270+
public T create() throws Exception {
271+
return connectionSupplier.get();
272+
}
273+
274+
@Override
275+
public PooledObject<T> wrap(T obj) {
276+
return new DefaultPooledObject<>(obj);
277+
}
278+
279+
@Override
280+
public boolean validateObject(PooledObject<T> p) {
281+
T connection = p.getObject();
282+
return connection.isOpen() && connectionValidator.test(connection);
283+
}
284+
285+
@Override
286+
public void destroyObject(PooledObject<T> p) throws Exception {
287+
try {
288+
p.getObject().close();
289+
} catch (Exception e) {
290+
e.printStackTrace();
291+
}
292+
}
293+
294+
}
295+
252296
}

‎src/test/java/io/lettuce/core/dynamic/RedisCommandsIntegrationTests.java

+7-1
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,13 @@ void verifierShouldCatchTooFewParametersDeclarations() {
100100
void shouldWorkWithPooledConnection() throws Exception {
101101

102102
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
103-
.createGenericObjectPool(client::connect, new GenericObjectPoolConfig<>());
103+
.createGenericObjectPool(client::connect, new GenericObjectPoolConfig<>(), connection -> {
104+
try {
105+
return "PONG".equals(connection.sync().ping());
106+
} catch (Exception e) {
107+
return false;
108+
}
109+
});
104110

105111
try (StatefulRedisConnection<String, String> connection = pool.borrowObject()) {
106112

‎src/test/java/io/lettuce/core/support/ConnectionPoolSupportIntegrationTests.java

+70-10
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,13 @@ static void afterClass() {
6666
void genericPoolShouldWorkWithWrappedConnections() throws Exception {
6767

6868
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
69-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>());
69+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), connection -> {
70+
try {
71+
return "PONG".equals(connection.sync().ping());
72+
} catch (Exception e) {
73+
return false;
74+
}
75+
});
7076

7177
borrowAndReturn(pool);
7278
borrowAndClose(pool);
@@ -91,7 +97,13 @@ void genericPoolShouldCloseConnectionsAboveMaxIdleSize() throws Exception {
9197
poolConfig.setMaxIdle(2);
9298

9399
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
94-
.createGenericObjectPool(() -> client.connect(), poolConfig);
100+
.createGenericObjectPool(() -> client.connect(), poolConfig, connection -> {
101+
try {
102+
return "PONG".equals(connection.sync().ping());
103+
} catch (Exception e) {
104+
return false;
105+
}
106+
});
95107

96108
borrowAndReturn(pool);
97109
borrowAndClose(pool);
@@ -120,7 +132,13 @@ void genericPoolShouldCloseConnectionsAboveMaxIdleSize() throws Exception {
120132
void genericPoolShouldWorkWithPlainConnections() throws Exception {
121133

122134
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
123-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), false);
135+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), false, connection -> {
136+
try {
137+
return "PONG".equals(connection.sync().ping());
138+
} catch (Exception e) {
139+
return false;
140+
}
141+
});
124142

125143
borrowAndReturn(pool);
126144

@@ -151,7 +169,13 @@ void softReferencePoolShouldWorkWithPlainConnections() throws Exception {
151169
void genericPoolUsingWrappingShouldPropagateExceptionsCorrectly() throws Exception {
152170

153171
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
154-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>());
172+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), connection -> {
173+
try {
174+
return "PONG".equals(connection.sync().ping());
175+
} catch (Exception e) {
176+
return false;
177+
}
178+
});
155179

156180
StatefulRedisConnection<String, String> connection = pool.borrowObject();
157181
RedisCommands<String, String> sync = connection.sync();
@@ -172,7 +196,13 @@ void genericPoolUsingWrappingShouldPropagateExceptionsCorrectly() throws Excepti
172196
void wrappedConnectionShouldUseWrappers() throws Exception {
173197

174198
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
175-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>());
199+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), connection -> {
200+
try {
201+
return "PONG".equals(connection.sync().ping());
202+
} catch (Exception e) {
203+
return false;
204+
}
205+
});
176206

177207
StatefulRedisConnection<String, String> connection = pool.borrowObject();
178208
RedisCommands<String, String> sync = connection.sync();
@@ -197,7 +227,13 @@ void wrappedMasterSlaveConnectionShouldUseWrappers() throws Exception {
197227

198228
GenericObjectPool<StatefulRedisMasterReplicaConnection<String, String>> pool = ConnectionPoolSupport
199229
.createGenericObjectPool(() -> MasterReplica.connect(client, new StringCodec(), RedisURI.create(host, port)),
200-
new GenericObjectPoolConfig<>());
230+
new GenericObjectPoolConfig<>(), connection -> {
231+
try {
232+
return "PONG".equals(connection.sync().ping());
233+
} catch (Exception e) {
234+
return false;
235+
}
236+
});
201237

202238
StatefulRedisMasterReplicaConnection<String, String> connection = pool.borrowObject();
203239
RedisCommands<String, String> sync = connection.sync();
@@ -223,7 +259,13 @@ void wrappedClusterConnectionShouldUseWrappers() throws Exception {
223259
RedisURI.create(TestSettings.host(), 7379));
224260

225261
GenericObjectPool<StatefulRedisClusterConnection<String, String>> pool = ConnectionPoolSupport
226-
.createGenericObjectPool(redisClusterClient::connect, new GenericObjectPoolConfig<>());
262+
.createGenericObjectPool(redisClusterClient::connect, new GenericObjectPoolConfig<>(), connection -> {
263+
try {
264+
return "PONG".equals(connection.sync().ping());
265+
} catch (Exception e) {
266+
return false;
267+
}
268+
});
227269

228270
StatefulRedisClusterConnection<String, String> connection = pool.borrowObject();
229271
RedisAdvancedClusterCommands<String, String> sync = connection.sync();
@@ -250,7 +292,13 @@ void wrappedClusterConnectionShouldUseWrappers() throws Exception {
250292
void plainConnectionShouldNotUseWrappers() throws Exception {
251293

252294
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
253-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), false);
295+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), false, connection -> {
296+
try {
297+
return "PONG".equals(connection.sync().ping());
298+
} catch (Exception e) {
299+
return false;
300+
}
301+
});
254302

255303
StatefulRedisConnection<String, String> connection = pool.borrowObject();
256304
RedisCommands<String, String> sync = connection.sync();
@@ -295,7 +343,13 @@ void softRefPoolShouldWorkWithWrappedConnections() throws Exception {
295343
void wrappedObjectClosedAfterReturn() throws Exception {
296344

297345
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
298-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), true);
346+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), true, connection -> {
347+
try {
348+
return "PONG".equals(connection.sync().ping());
349+
} catch (Exception e) {
350+
return false;
351+
}
352+
});
299353

300354
StatefulRedisConnection<String, String> connection = pool.borrowObject();
301355
RedisCommands<String, String> sync = connection.sync();
@@ -317,7 +371,13 @@ void wrappedObjectClosedAfterReturn() throws Exception {
317371
void tryWithResourcesReturnsConnectionToPool() throws Exception {
318372

319373
GenericObjectPool<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
320-
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>());
374+
.createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), connection -> {
375+
try {
376+
return "PONG".equals(connection.sync().ping());
377+
} catch (Exception e) {
378+
return false;
379+
}
380+
});
321381

322382
StatefulRedisConnection<String, String> usedConnection = null;
323383
try (StatefulRedisConnection<String, String> connection = pool.borrowObject()) {

0 commit comments

Comments
 (0)