com.datastax.oss.driver.internal.core.session.PoolManager Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.datastax.oss.driver.internal.core.session;
import com.datastax.oss.driver.api.core.AsyncAutoCloseable;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.InvalidKeyspaceException;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.NodeState;
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
import com.datastax.oss.driver.internal.core.metadata.DistanceEvent;
import com.datastax.oss.driver.internal.core.metadata.NodeStateEvent;
import com.datastax.oss.driver.internal.core.metadata.TopologyEvent;
import com.datastax.oss.driver.internal.core.pool.ChannelPool;
import com.datastax.oss.driver.internal.core.pool.ChannelPoolFactory;
import com.datastax.oss.driver.internal.core.util.Loggers;
import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
import com.datastax.oss.driver.internal.core.util.concurrent.ReplayingEventFilter;
import com.datastax.oss.driver.internal.core.util.concurrent.RunOrSchedule;
import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
import com.datastax.oss.driver.shaded.guava.common.collect.Lists;
import com.datastax.oss.driver.shaded.guava.common.collect.MapMaker;
import edu.umd.cs.findbugs.annotations.NonNull;
import com.datastax.oss.driver.shaded.netty.util.concurrent.EventExecutor;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.WeakHashMap;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import net.jcip.annotations.ThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Maintains the connection pools of a session.
*
* Logically this belongs to {@link DefaultSession}, but it's extracted here in order to be
* accessible from the context (notably for metrics).
*/
@ThreadSafe
public class PoolManager implements AsyncAutoCloseable {
private static final Logger LOG = LoggerFactory.getLogger(PoolManager.class);
// This is read concurrently, but only updated from adminExecutor
private volatile CqlIdentifier keyspace;
private final ConcurrentMap pools =
new ConcurrentHashMap<>(
16,
0.75f,
// the map will only be updated from adminExecutor
1);
// The raw data to reprepare requests on the fly, if we hit a node that doesn't have them in
// its cache.
// This is raw protocol-level data, as opposed to the actual instances returned to the client
// (e.g. DefaultPreparedStatement) which are handled at the protocol level (e.g.
// CqlPrepareAsyncProcessor). We keep the two separate to avoid introducing a dependency from the
// session to a particular processor implementation.
private final ConcurrentMap repreparePayloads;
private final String logPrefix;
private final EventExecutor adminExecutor;
private final DriverExecutionProfile config;
private final SingleThreaded singleThreaded;
public PoolManager(InternalDriverContext context) {
this.logPrefix = context.getSessionName();
this.adminExecutor = context.getNettyOptions().adminEventExecutorGroup().next();
this.config = context.getConfig().getDefaultProfile();
this.singleThreaded = new SingleThreaded(context);
if (config.getBoolean(DefaultDriverOption.PREPARED_CACHE_WEAK_VALUES, true)) {
LOG.debug("[{}] Prepared statements cache configured to use weak values", logPrefix);
this.repreparePayloads = new MapMaker().weakValues().makeMap();
} else {
LOG.debug("[{}] Prepared statements cache configured to use strong values", logPrefix);
this.repreparePayloads = new MapMaker().makeMap();
}
}
public CompletionStage init(CqlIdentifier keyspace) {
RunOrSchedule.on(adminExecutor, () -> singleThreaded.init(keyspace));
return singleThreaded.initFuture;
}
public CqlIdentifier getKeyspace() {
return keyspace;
}
public CompletionStage setKeyspace(CqlIdentifier newKeyspace) {
CqlIdentifier oldKeyspace = this.keyspace;
if (Objects.equals(oldKeyspace, newKeyspace)) {
return CompletableFuture.completedFuture(null);
}
if (config.getBoolean(DefaultDriverOption.REQUEST_WARN_IF_SET_KEYSPACE)) {
LOG.warn(
"[{}] Detected a keyspace change at runtime ({} => {}). "
+ "This is an anti-pattern that should be avoided in production "
+ "(see '{}' in the configuration).",
logPrefix,
(oldKeyspace == null) ? "" : oldKeyspace.asInternal(),
newKeyspace.asInternal(),
DefaultDriverOption.REQUEST_WARN_IF_SET_KEYSPACE.getPath());
}
this.keyspace = newKeyspace;
CompletableFuture result = new CompletableFuture<>();
RunOrSchedule.on(adminExecutor, () -> singleThreaded.setKeyspace(newKeyspace, result));
return result;
}
public Map getPools() {
return pools;
}
public ConcurrentMap getRepreparePayloads() {
return repreparePayloads;
}
@NonNull
@Override
public CompletionStage closeFuture() {
return singleThreaded.closeFuture;
}
@NonNull
@Override
public CompletionStage closeAsync() {
RunOrSchedule.on(adminExecutor, singleThreaded::close);
return singleThreaded.closeFuture;
}
@NonNull
@Override
public CompletionStage forceCloseAsync() {
RunOrSchedule.on(adminExecutor, singleThreaded::forceClose);
return singleThreaded.closeFuture;
}
private class SingleThreaded {
private final InternalDriverContext context;
private final ChannelPoolFactory channelPoolFactory;
private final CompletableFuture initFuture = new CompletableFuture<>();
private boolean initWasCalled;
private final CompletableFuture closeFuture = new CompletableFuture<>();
private boolean closeWasCalled;
private boolean forceCloseWasCalled;
private final Object distanceListenerKey;
private final ReplayingEventFilter distanceEventFilter =
new ReplayingEventFilter<>(this::processDistanceEvent);
private final Object stateListenerKey;
private final ReplayingEventFilter stateEventFilter =
new ReplayingEventFilter<>(this::processStateEvent);
private final Object topologyListenerKey;
// The pools that we have opened but have not finished initializing yet
private final Map> pending = new HashMap<>();
// If we receive events while a pool is initializing, the last one is stored here
private final Map pendingDistanceEvents = new WeakHashMap<>();
private final Map pendingStateEvents = new WeakHashMap<>();
private SingleThreaded(InternalDriverContext context) {
this.context = context;
this.channelPoolFactory = context.getChannelPoolFactory();
this.distanceListenerKey =
context
.getEventBus()
.register(
DistanceEvent.class, RunOrSchedule.on(adminExecutor, this::onDistanceEvent));
this.stateListenerKey =
context
.getEventBus()
.register(NodeStateEvent.class, RunOrSchedule.on(adminExecutor, this::onStateEvent));
this.topologyListenerKey =
context
.getEventBus()
.register(
TopologyEvent.class, RunOrSchedule.on(adminExecutor, this::onTopologyEvent));
}
private void init(CqlIdentifier keyspace) {
assert adminExecutor.inEventLoop();
if (initWasCalled) {
return;
}
initWasCalled = true;
LOG.debug("[{}] Starting initialization", logPrefix);
PoolManager.this.keyspace = keyspace;
// Make sure we don't miss any event while the pools are initializing
distanceEventFilter.start();
stateEventFilter.start();
Collection nodes = context.getMetadataManager().getMetadata().getNodes().values();
List> poolStages = new ArrayList<>(nodes.size());
for (Node node : nodes) {
NodeDistance distance = node.getDistance();
if (distance == NodeDistance.IGNORED) {
LOG.debug("[{}] Skipping {} because it is IGNORED", logPrefix, node);
} else if (node.getState() == NodeState.FORCED_DOWN) {
LOG.debug("[{}] Skipping {} because it is FORCED_DOWN", logPrefix, node);
} else {
LOG.debug("[{}] Creating a pool for {}", logPrefix, node);
poolStages.add(channelPoolFactory.init(node, keyspace, distance, context, logPrefix));
}
}
CompletableFutures.whenAllDone(poolStages, () -> this.onPoolsInit(poolStages), adminExecutor);
}
private void onPoolsInit(List> poolStages) {
assert adminExecutor.inEventLoop();
LOG.debug("[{}] All pools have finished initializing", logPrefix);
// We will only propagate an invalid keyspace error if all pools get it
boolean allInvalidKeyspaces = poolStages.size() > 0;
for (CompletionStage poolStage : poolStages) {
// Note: pool init always succeeds
ChannelPool pool = CompletableFutures.getCompleted(poolStage.toCompletableFuture());
boolean invalidKeyspace = pool.isInvalidKeyspace();
if (invalidKeyspace) {
LOG.debug("[{}] Pool to {} reports an invalid keyspace", logPrefix, pool.getNode());
}
allInvalidKeyspaces &= invalidKeyspace;
pools.put(pool.getNode(), pool);
}
if (allInvalidKeyspaces) {
initFuture.completeExceptionally(
new InvalidKeyspaceException("Invalid keyspace " + keyspace.asCql(true)));
forceClose();
} else {
LOG.debug("[{}] Initialization complete, ready", logPrefix);
initFuture.complete(null);
distanceEventFilter.markReady();
stateEventFilter.markReady();
}
}
private void onDistanceEvent(DistanceEvent event) {
assert adminExecutor.inEventLoop();
distanceEventFilter.accept(event);
}
private void onStateEvent(NodeStateEvent event) {
assert adminExecutor.inEventLoop();
stateEventFilter.accept(event);
}
private void processDistanceEvent(DistanceEvent event) {
assert adminExecutor.inEventLoop();
// no need to check closeWasCalled, because we stop listening for events one closed
DefaultNode node = event.node;
NodeDistance newDistance = event.distance;
if (pending.containsKey(node)) {
pendingDistanceEvents.put(node, event);
} else if (newDistance == NodeDistance.IGNORED) {
ChannelPool pool = pools.remove(node);
if (pool != null) {
LOG.debug("[{}] {} became IGNORED, destroying pool", logPrefix, node);
pool.closeAsync()
.exceptionally(
error -> {
Loggers.warnWithException(LOG, "[{}] Error closing pool", logPrefix, error);
return null;
});
}
} else {
NodeState state = node.getState();
if (state == NodeState.FORCED_DOWN) {
LOG.warn(
"[{}] {} became {} but it is FORCED_DOWN, ignoring", logPrefix, node, newDistance);
return;
}
ChannelPool pool = pools.get(node);
if (pool == null) {
LOG.debug(
"[{}] {} became {} and no pool found, initializing it", logPrefix, node, newDistance);
CompletionStage poolFuture =
channelPoolFactory.init(node, keyspace, newDistance, context, logPrefix);
pending.put(node, poolFuture);
poolFuture
.thenAcceptAsync(this::onPoolInitialized, adminExecutor)
.exceptionally(UncaughtExceptions::log);
} else {
LOG.debug("[{}] {} became {}, resizing it", logPrefix, node, newDistance);
pool.resize(newDistance);
}
}
}
private void processStateEvent(NodeStateEvent event) {
assert adminExecutor.inEventLoop();
// no need to check closeWasCalled, because we stop listening for events once closed
DefaultNode node = event.node;
NodeState oldState = event.oldState;
NodeState newState = event.newState;
if (pending.containsKey(node)) {
pendingStateEvents.put(node, event);
} else if (newState == null || newState == NodeState.FORCED_DOWN) {
ChannelPool pool = pools.remove(node);
if (pool != null) {
LOG.debug(
"[{}] {} was {}, destroying pool",
logPrefix,
node,
newState == null ? "removed" : newState.name());
pool.closeAsync()
.exceptionally(
error -> {
Loggers.warnWithException(LOG, "[{}] Error closing pool", logPrefix, error);
return null;
});
}
} else if (oldState == NodeState.FORCED_DOWN
&& newState == NodeState.UP
&& node.getDistance() != NodeDistance.IGNORED) {
LOG.debug("[{}] {} was forced back UP, initializing pool", logPrefix, node);
createOrReconnectPool(node);
}
}
private void onTopologyEvent(TopologyEvent event) {
assert adminExecutor.inEventLoop();
if (event.type == TopologyEvent.Type.SUGGEST_UP) {
context
.getMetadataManager()
.getMetadata()
.findNode(event.broadcastRpcAddress)
.ifPresent(
node -> {
if (node.getDistance() != NodeDistance.IGNORED) {
LOG.debug(
"[{}] Received a SUGGEST_UP event for {}, reconnecting pool now",
logPrefix,
node);
ChannelPool pool = pools.get(node);
if (pool != null) {
pool.reconnectNow();
}
}
});
}
}
private void createOrReconnectPool(Node node) {
ChannelPool pool = pools.get(node);
if (pool == null) {
CompletionStage poolFuture =
channelPoolFactory.init(node, keyspace, node.getDistance(), context, logPrefix);
pending.put(node, poolFuture);
poolFuture
.thenAcceptAsync(this::onPoolInitialized, adminExecutor)
.exceptionally(UncaughtExceptions::log);
} else {
pool.reconnectNow();
}
}
private void onPoolInitialized(ChannelPool pool) {
assert adminExecutor.inEventLoop();
Node node = pool.getNode();
if (closeWasCalled) {
LOG.debug(
"[{}] Session closed while a pool to {} was initializing, closing it", logPrefix, node);
pool.forceCloseAsync();
} else {
LOG.debug("[{}] New pool to {} initialized", logPrefix, node);
if (Objects.equals(keyspace, pool.getInitialKeyspaceName())) {
reprepareStatements(pool);
} else {
// The keyspace changed while the pool was being initialized, switch it now.
pool.setKeyspace(keyspace)
.handleAsync(
(result, error) -> {
if (error != null) {
Loggers.warnWithException(
LOG, "Error while switching keyspace to " + keyspace, error);
}
reprepareStatements(pool);
return null;
},
adminExecutor);
}
}
}
private void reprepareStatements(ChannelPool pool) {
assert adminExecutor.inEventLoop();
if (config.getBoolean(DefaultDriverOption.REPREPARE_ENABLED)) {
new ReprepareOnUp(
logPrefix + "|" + pool.getNode().getEndPoint(),
pool,
adminExecutor,
repreparePayloads,
context,
() -> RunOrSchedule.on(adminExecutor, () -> onPoolReady(pool)))
.start();
} else {
LOG.debug("[{}] Reprepare on up is disabled, skipping", logPrefix);
onPoolReady(pool);
}
}
private void onPoolReady(ChannelPool pool) {
assert adminExecutor.inEventLoop();
Node node = pool.getNode();
pending.remove(node);
pools.put(node, pool);
DistanceEvent distanceEvent = pendingDistanceEvents.remove(node);
NodeStateEvent stateEvent = pendingStateEvents.remove(node);
if (stateEvent != null && stateEvent.newState == NodeState.FORCED_DOWN) {
LOG.debug(
"[{}] Received {} while the pool was initializing, processing it now",
logPrefix,
stateEvent);
processStateEvent(stateEvent);
} else if (distanceEvent != null) {
LOG.debug(
"[{}] Received {} while the pool was initializing, processing it now",
logPrefix,
distanceEvent);
processDistanceEvent(distanceEvent);
}
}
private void setKeyspace(CqlIdentifier newKeyspace, CompletableFuture doneFuture) {
assert adminExecutor.inEventLoop();
if (closeWasCalled) {
doneFuture.complete(null);
return;
}
LOG.debug("[{}] Switching to keyspace {}", logPrefix, newKeyspace);
List> poolReadyFutures = Lists.newArrayListWithCapacity(pools.size());
for (ChannelPool pool : pools.values()) {
poolReadyFutures.add(pool.setKeyspace(newKeyspace));
}
CompletableFutures.completeFrom(CompletableFutures.allDone(poolReadyFutures), doneFuture);
}
private void close() {
assert adminExecutor.inEventLoop();
if (closeWasCalled) {
return;
}
closeWasCalled = true;
LOG.debug("[{}] Starting shutdown", logPrefix);
// Stop listening for events
context.getEventBus().unregister(distanceListenerKey, DistanceEvent.class);
context.getEventBus().unregister(stateListenerKey, NodeStateEvent.class);
context.getEventBus().unregister(topologyListenerKey, TopologyEvent.class);
List> closePoolStages = new ArrayList<>(pools.size());
for (ChannelPool pool : pools.values()) {
closePoolStages.add(pool.closeAsync());
}
CompletableFutures.whenAllDone(
closePoolStages, () -> onAllPoolsClosed(closePoolStages), adminExecutor);
}
private void forceClose() {
assert adminExecutor.inEventLoop();
if (forceCloseWasCalled) {
return;
}
forceCloseWasCalled = true;
LOG.debug(
"[{}] Starting forced shutdown (was {}closed before)",
logPrefix,
(closeWasCalled ? "" : "not "));
if (closeWasCalled) {
for (ChannelPool pool : pools.values()) {
pool.forceCloseAsync();
}
} else {
List> closePoolStages = new ArrayList<>(pools.size());
for (ChannelPool pool : pools.values()) {
closePoolStages.add(pool.forceCloseAsync());
}
CompletableFutures.whenAllDone(
closePoolStages, () -> onAllPoolsClosed(closePoolStages), adminExecutor);
}
}
private void onAllPoolsClosed(List> closePoolStages) {
assert adminExecutor.inEventLoop();
Throwable firstError = null;
for (CompletionStage closePoolStage : closePoolStages) {
CompletableFuture closePoolFuture = closePoolStage.toCompletableFuture();
assert closePoolFuture.isDone();
if (closePoolFuture.isCompletedExceptionally()) {
Throwable error = CompletableFutures.getFailed(closePoolFuture);
if (firstError == null) {
firstError = error;
} else {
firstError.addSuppressed(error);
}
}
}
if (firstError != null) {
closeFuture.completeExceptionally(firstError);
} else {
LOG.debug("[{}] Shutdown complete", logPrefix);
closeFuture.complete(null);
}
}
}
}