Skip to content

Commit 115dfda

Browse files
sergiyvamzsergiyvamz
authored
consolidate cache clean-up in a single place (#1234)
1 parent 64e3b1b commit 115dfda

30 files changed

+451
-190
lines changed

‎wrapper/src/main/java/software/amazon/jdbc/Driver.java

+43
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,21 @@
3333
import org.checkerframework.checker.nullness.qual.NonNull;
3434
import org.checkerframework.checker.nullness.qual.Nullable;
3535
import software.amazon.jdbc.authentication.AwsCredentialsManager;
36+
import software.amazon.jdbc.dialect.DialectManager;
37+
import software.amazon.jdbc.hostlistprovider.RdsHostListProvider;
38+
import software.amazon.jdbc.hostlistprovider.monitoring.MonitoringRdsHostListProvider;
39+
import software.amazon.jdbc.plugin.AwsSecretsManagerCacheHolder;
40+
import software.amazon.jdbc.plugin.DataCacheConnectionPlugin;
41+
import software.amazon.jdbc.plugin.OpenedConnectionTracker;
42+
import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl;
43+
import software.amazon.jdbc.plugin.customendpoint.CustomEndpointPlugin;
44+
import software.amazon.jdbc.plugin.efm.MonitorThreadContainer;
45+
import software.amazon.jdbc.plugin.federatedauth.FederatedAuthCacheHolder;
46+
import software.amazon.jdbc.plugin.federatedauth.OktaAuthCacheHolder;
47+
import software.amazon.jdbc.plugin.iam.IamAuthCacheHolder;
48+
import software.amazon.jdbc.plugin.limitless.LimitlessRouterServiceImpl;
49+
import software.amazon.jdbc.plugin.strategy.fastestresponse.FastestResponseStrategyPlugin;
50+
import software.amazon.jdbc.plugin.strategy.fastestresponse.HostResponseTimeServiceImpl;
3651
import software.amazon.jdbc.profile.ConfigurationProfile;
3752
import software.amazon.jdbc.profile.DriverConfigurationProfiles;
3853
import software.amazon.jdbc.states.ResetSessionStateOnCloseCallable;
@@ -277,4 +292,32 @@ public static void setPrepareHostFunc(final Function<String, String> func) {
277292
public static void resetPrepareHostFunc() {
278293
RdsUtils.resetPrepareHostFunc();
279294
}
295+
296+
public static void clearCaches() {
297+
RdsUtils.clearCache();
298+
RdsHostListProvider.clearAll();
299+
PluginServiceImpl.clearCache();
300+
DialectManager.resetEndpointCache();
301+
MonitoringRdsHostListProvider.clearCache();
302+
CustomEndpointMonitorImpl.clearCache();
303+
OpenedConnectionTracker.clearCache();
304+
AwsSecretsManagerCacheHolder.clearCache();
305+
DataCacheConnectionPlugin.clearCache();
306+
FederatedAuthCacheHolder.clearCache();
307+
OktaAuthCacheHolder.clearCache();
308+
IamAuthCacheHolder.clearCache();
309+
LimitlessRouterServiceImpl.clearCache();
310+
RoundRobinHostSelector.clearCache();
311+
FastestResponseStrategyPlugin.clearCache();
312+
}
313+
314+
public static void releaseResources() {
315+
software.amazon.jdbc.plugin.efm2.MonitorServiceImpl.closeAllMonitors();
316+
MonitorThreadContainer.releaseInstance();
317+
ConnectionProviderManager.releaseResources();
318+
CustomEndpointPlugin.closeMonitors();
319+
HikariPoolsHolder.closeAllPools();
320+
HostResponseTimeServiceImpl.closeAllMonitors();
321+
clearCaches();
322+
}
280323
}

‎wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java

+27-79
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import software.amazon.jdbc.targetdriverdialect.ConnectInfo;
4040
import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
4141
import software.amazon.jdbc.util.Messages;
42+
import software.amazon.jdbc.util.Pair;
4243
import software.amazon.jdbc.util.PropertyUtils;
4344
import software.amazon.jdbc.util.RdsUrlType;
4445
import software.amazon.jdbc.util.RdsUtils;
@@ -61,17 +62,22 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
6162
});
6263

6364
protected static final RdsUtils rdsUtils = new RdsUtils();
64-
protected static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
65-
new SlidingExpirationCache<>(
66-
(hikariDataSource) -> hikariDataSource.getHikariPoolMXBean().getActiveConnections() == 0,
67-
HikariDataSource::close
68-
);
6965
protected static long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);
7066
protected final HikariPoolConfigurator poolConfigurator;
7167
protected final HikariPoolMapping poolMapping;
7268
protected final AcceptsUrlFunc acceptsUrlFunc;
7369
protected final LeastConnectionsHostSelector leastConnectionsHostSelector;
7470

71+
static {
72+
HikariPoolsHolder.databasePools.setShouldDisposeFunc(
73+
(hikariDataSource) -> {
74+
if (hikariDataSource instanceof HikariDataSource) {
75+
return ((HikariDataSource) hikariDataSource).getHikariPoolMXBean().getActiveConnections() == 0;
76+
}
77+
return true;
78+
});
79+
}
80+
7581
/**
7682
* {@link HikariPooledConnectionProvider} constructor. This class can be passed to
7783
* {@link ConnectionProviderManager#setConnectionProvider} to enable internal connection pools for
@@ -114,7 +120,7 @@ public HikariPooledConnectionProvider(
114120
this.poolConfigurator = hikariPoolConfigurator;
115121
this.poolMapping = mapping;
116122
this.acceptsUrlFunc = null;
117-
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
123+
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools);
118124
}
119125

120126
/**
@@ -151,8 +157,8 @@ public HikariPooledConnectionProvider(
151157
this.poolMapping = mapping;
152158
this.acceptsUrlFunc = null;
153159
poolExpirationCheckNanos = poolExpirationNanos;
154-
databasePools.setCleanupIntervalNanos(poolCleanupNanos);
155-
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
160+
HikariPoolsHolder.databasePools.setCleanupIntervalNanos(poolCleanupNanos);
161+
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools);
156162
}
157163

158164
/**
@@ -193,8 +199,8 @@ public HikariPooledConnectionProvider(
193199
this.poolMapping = mapping;
194200
this.acceptsUrlFunc = acceptsUrlFunc;
195201
poolExpirationCheckNanos = poolExpirationNanos;
196-
databasePools.setCleanupIntervalNanos(poolCleanupNanos);
197-
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
202+
HikariPoolsHolder.databasePools.setCleanupIntervalNanos(poolCleanupNanos);
203+
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools);
198204
}
199205

200206

@@ -272,8 +278,8 @@ public Connection connect(
272278
final HostSpec finalHostSpec = connectionHostSpec;
273279
dialect.prepareConnectProperties(copy, protocol, finalHostSpec);
274280

275-
final HikariDataSource ds = databasePools.computeIfAbsent(
276-
new PoolKey(hostSpec.getUrl(), getPoolKey(finalHostSpec, copy)),
281+
final HikariDataSource ds = (HikariDataSource) HikariPoolsHolder.databasePools.computeIfAbsent(
282+
Pair.create(hostSpec.getUrl(), getPoolKey(finalHostSpec, copy)),
277283
(lambdaPoolKey) -> createHikariDataSource(protocol, finalHostSpec, copy, targetDriverDialect),
278284
poolExpirationCheckNanos
279285
);
@@ -297,22 +303,7 @@ protected String getPoolKey(HostSpec hostSpec, Properties props) {
297303

298304
@Override
299305
public void releaseResources() {
300-
databasePools.getEntries().forEach((poolKey, pool) -> {
301-
if (!pool.isClosed()) {
302-
pool.close();
303-
}
304-
});
305-
databasePools.clear();
306-
}
307-
308-
// For testing purposes
309-
public static void clearCache() {
310-
databasePools.getEntries().forEach((poolKey, pool) -> {
311-
if (!pool.isClosed()) {
312-
pool.close();
313-
}
314-
});
315-
databasePools.clear();
306+
HikariPoolsHolder.closeAllPools();
316307
}
317308

318309
/**
@@ -378,7 +369,7 @@ protected void configurePool(
378369
* @return the number of active connection pools
379370
*/
380371
public int getHostCount() {
381-
return databasePools.size();
372+
return HikariPoolsHolder.databasePools.size();
382373
}
383374

384375
/**
@@ -388,8 +379,8 @@ public int getHostCount() {
388379
*/
389380
public Set<String> getHosts() {
390381
return Collections.unmodifiableSet(
391-
databasePools.getEntries().keySet().stream()
392-
.map(poolKey -> poolKey.url)
382+
HikariPoolsHolder.databasePools.getEntries().keySet().stream()
383+
.map(poolKey -> (String) poolKey.getValue1())
393384
.collect(Collectors.toSet()));
394385
}
395386

@@ -398,8 +389,8 @@ public Set<String> getHosts() {
398389
*
399390
* @return a set containing every key associated with an active connection pool
400391
*/
401-
public Set<PoolKey> getKeys() {
402-
return databasePools.getEntries().keySet();
392+
public Set<Pair> getKeys() {
393+
return HikariPoolsHolder.databasePools.getEntries().keySet();
403394
}
404395

405396
@Override
@@ -413,7 +404,7 @@ public String getTargetName() {
413404
public void logConnections() {
414405
LOGGER.finest(() -> {
415406
final StringBuilder builder = new StringBuilder();
416-
databasePools.getEntries().forEach((key, dataSource) -> {
407+
HikariPoolsHolder.databasePools.getEntries().forEach((key, dataSource) -> {
417408
builder.append("\t[ ");
418409
builder.append(key).append(":");
419410
builder.append("\n\t {");
@@ -436,51 +427,8 @@ HikariDataSource createHikariDataSource(
436427
return new HikariDataSource(config);
437428
}
438429

439-
public static class PoolKey {
440-
private final @NonNull String url;
441-
private final @NonNull String extraKey;
442-
443-
public PoolKey(final @NonNull String url, final @NonNull String extraKey) {
444-
this.url = url;
445-
this.extraKey = extraKey;
446-
}
447-
448-
public String getUrl() {
449-
return this.url;
450-
}
451-
452-
@Override
453-
public int hashCode() {
454-
final int prime = 31;
455-
int result = 1;
456-
result = prime * result + ((url == null) ? 0 : url.hashCode()) + ((extraKey == null) ? 0 : extraKey.hashCode());
457-
return result;
458-
}
459-
460-
@Override
461-
public boolean equals(final Object obj) {
462-
if (this == obj) {
463-
return true;
464-
}
465-
if (obj == null) {
466-
return false;
467-
}
468-
if (getClass() != obj.getClass()) {
469-
return false;
470-
}
471-
final PoolKey other = (PoolKey) obj;
472-
return this.url.equals(other.url) && this.extraKey.equals(other.extraKey);
473-
}
474-
475-
@Override
476-
public String toString() {
477-
return "PoolKey [url=" + url + ", extraKey=" + extraKey + "]";
478-
}
479-
480-
}
481-
482430
// For testing purposes only
483-
void setDatabasePools(SlidingExpirationCache<PoolKey, HikariDataSource> connectionPools) {
484-
databasePools = connectionPools;
431+
void setDatabasePools(SlidingExpirationCache<Pair, AutoCloseable> connectionPools) {
432+
HikariPoolsHolder.databasePools = connectionPools;
485433
}
486434
}

‎wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java

+46
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package software.amazon.jdbc;
18+
19+
import software.amazon.jdbc.util.Pair;
20+
import software.amazon.jdbc.util.SlidingExpirationCache;
21+
22+
public class HikariPoolsHolder {
23+
static SlidingExpirationCache<Pair, AutoCloseable> databasePools =
24+
new SlidingExpirationCache<>(
25+
null,
26+
(hikariDataSource) -> {
27+
try {
28+
hikariDataSource.close();
29+
} catch (Exception ex) {
30+
// ignore
31+
}
32+
}
33+
);
34+
35+
public static void closeAllPools() {
36+
databasePools.getEntries().forEach((poolKey, pool) -> {
37+
try {
38+
pool.close();
39+
} catch (Exception ex) {
40+
// ignore
41+
}
42+
});
43+
databasePools.clear();
44+
45+
}
46+
}

‎wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java

+10-6
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,15 @@
2626
import org.checkerframework.checker.nullness.qual.Nullable;
2727
import software.amazon.jdbc.hostavailability.HostAvailability;
2828
import software.amazon.jdbc.util.Messages;
29+
import software.amazon.jdbc.util.Pair;
2930
import software.amazon.jdbc.util.SlidingExpirationCache;
3031

3132
public class LeastConnectionsHostSelector implements HostSelector {
3233
public static final String STRATEGY_LEAST_CONNECTIONS = "leastConnections";
33-
private final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools;
34+
private final SlidingExpirationCache<Pair, AutoCloseable> databasePools;
3435

3536
public LeastConnectionsHostSelector(
36-
SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
37+
SlidingExpirationCache<Pair, AutoCloseable> databasePools) {
3738
this.databasePools = databasePools;
3839
}
3940

@@ -58,15 +59,18 @@ public HostSpec getHost(
5859

5960
private int getNumConnections(
6061
final HostSpec hostSpec,
61-
final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
62+
final SlidingExpirationCache<Pair, AutoCloseable> databasePools) {
6263
int numConnections = 0;
6364
final String url = hostSpec.getUrl();
64-
for (final Map.Entry<HikariPooledConnectionProvider.PoolKey, HikariDataSource> entry :
65+
for (final Map.Entry<Pair, AutoCloseable> entry :
6566
databasePools.getEntries().entrySet()) {
66-
if (!url.equals(entry.getKey().getUrl())) {
67+
if (!url.equals(entry.getKey().getValue1())) {
6768
continue;
6869
}
69-
numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections();
70+
if (!(entry.getValue() instanceof HikariDataSource)) {
71+
continue;
72+
}
73+
numConnections += ((HikariDataSource) entry.getValue()).getHikariPoolMXBean().getActiveConnections();
7074
}
7175
return numConnections;
7276
}

‎wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java

+4
Original file line numberDiff line numberDiff line change
@@ -754,4 +754,8 @@ public <T> T getPlugin(final Class<T> pluginClazz) {
754754
}
755755
return null;
756756
}
757+
758+
public static void clearCache() {
759+
hostAvailabilityExpiringCache.clear();
760+
}
757761
}

‎wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java

+1-2
Original file line numberDiff line numberDiff line change
@@ -251,8 +251,7 @@ private void updateCachedHostWeightPairsPropertiesForRoundRobinClusterInfo(
251251
}
252252
}
253253

254-
// For testing purposes only
255-
public void clearCache() {
254+
public static void clearCache() {
256255
roundRobinCache.clear();
257256
}
258257

‎wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java

+7-3
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,12 @@ public void clear() {
9393
}
9494

9595
public static void clearCache() {
96+
topologyCache.clear();
97+
primaryClusterIdCache.clear();
98+
suggestedPrimaryClusterIdCache.clear();
99+
}
100+
101+
public static void closeAllMonitors() {
96102
monitors.getEntries().values().forEach(monitor -> {
97103
try {
98104
monitor.close();
@@ -101,9 +107,7 @@ public static void clearCache() {
101107
}
102108
});
103109
monitors.clear();
104-
topologyCache.clear();
105-
primaryClusterIdCache.clear();
106-
suggestedPrimaryClusterIdCache.clear();
110+
clearCache();
107111
}
108112

109113
@Override

0 commit comments

Comments
 (0)