diff --git a/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java b/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java index c17ccd50ff8..cd21c3c8c8c 100644 --- a/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java +++ b/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java @@ -149,6 +149,8 @@ public abstract class ContinuousRequestHandlerBase chosenCallback = new CompletableFuture<>(); + private final AtomicBoolean terminal = new AtomicBoolean(); + /** * How many speculative executions are currently running (including the initial execution). We * track this in order to know when to fail the request if all executions have reached the end of @@ -254,22 +256,32 @@ protected abstract ResultSetT createResultSet( @Override public void onThrottleReady(boolean wasDelayed) { - DriverExecutionProfile executionProfile = - Conversions.resolveExecutionProfile(initialStatement, context); - if (wasDelayed - // avoid call to nanoTime() if metric is disabled: - && sessionMetricUpdater.isEnabled( - DefaultSessionMetric.THROTTLING_DELAY, executionProfile.getName())) { - session - .getMetricUpdater() - .updateTimer( - DefaultSessionMetric.THROTTLING_DELAY, - executionProfile.getName(), - System.nanoTime() - startTimeNanos, - TimeUnit.NANOSECONDS); + try { + DriverExecutionProfile executionProfile = + Conversions.resolveExecutionProfile(initialStatement, context); + if (wasDelayed + // avoid call to nanoTime() if metric is disabled: + && sessionMetricUpdater.isEnabled( + DefaultSessionMetric.THROTTLING_DELAY, executionProfile.getName())) { + session + .getMetricUpdater() + .updateTimer( + DefaultSessionMetric.THROTTLING_DELAY, + executionProfile.getName(), + System.nanoTime() - startTimeNanos, + TimeUnit.NANOSECONDS); + } + activeExecutionsCount.incrementAndGet(); + sendRequest(initialStatement, null, 0, 0, specExecEnabled); + } catch (Throwable t) { + // The concurrency throttler contains ready-callback exceptions. Complete and release here so + // failures before sendRequest's guarded setup path cannot strand an admitted request. + if (abortGlobalRequestOrChosenCallback(t) && !(t instanceof CancellationException)) { + // Cancellation is propagated through the fetched result, whose cancellation listener owns + // the single signalCancel call. + throttler.signalError(this, t); + } } - activeExecutionsCount.incrementAndGet(); - sendRequest(initialStatement, null, 0, 0, specExecEnabled); } @Override @@ -283,6 +295,8 @@ public void onThrottleFailure(@NonNull RequestThrottlingException error) { } private boolean abortGlobalRequestOrChosenCallback(@NonNull Throwable error) { + terminal.set(true); + cancelGlobalTimeout(); boolean completedChosenCallback = chosenCallback.completeExceptionally(error); if (!completedChosenCallback) { chosenCallback.thenAccept(callback -> callback.abort(error, false)); @@ -291,7 +305,16 @@ private boolean abortGlobalRequestOrChosenCallback(@NonNull Throwable error) { } public CompletionStage handle() { - globalTimeout = scheduleGlobalTimeout(); + // Immediate admission happens in the continuous graph handler's constructor. If setup failed + // there, chosenCallback is already terminal and there is no live request to time out. + if (!terminal.get()) { + globalTimeout = scheduleGlobalTimeout(); + // Admission can race with handle() after the check above but before globalTimeout is + // assigned. Ensure a synchronous terminal setup failure cannot leave that timeout behind. + if (terminal.get()) { + cancelGlobalTimeout(); + } + } return fetchNextPage(); } @@ -370,7 +393,7 @@ private void sendRequest( } } else if (!chosenCallback.isDone()) { boolean writeSubmitted = false; - Throwable terminalPreWriteFailure = null; + Throwable terminalSetupFailure = null; NodeResponseCallback nodeResponseCallback = null; try { nodeResponseCallback = @@ -392,12 +415,15 @@ private void sendRequest( writeSubmitted = true; writeFuture.addListener(nodeResponseCallback); } catch (Throwable t) { - if (!writeSubmitted && activeExecutionsCount.decrementAndGet() == 0) { - if (abortGlobalRequestOrChosenCallback(t)) { - terminalPreWriteFailure = t; + recordError(node, t); + if (activeExecutionsCount.decrementAndGet() == 0) { + if (abortGlobalRequestOrChosenCallback(t) && !(t instanceof CancellationException)) { + terminalSetupFailure = t; } + } else { + Loggers.warnWithException( + LOG, "[{}] Request setup failed, another execution is still active", logPrefix, t); } - throw t; } finally { if (!writeSubmitted) { if (nodeResponseCallback != null) { @@ -405,13 +431,19 @@ private void sendRequest( } try { channel.cancelPreAcquireId(); - } finally { - if (terminalPreWriteFailure != null) { - throttler.signalError(this, terminalPreWriteFailure); + } catch (Throwable cleanupFailure) { + if (terminalSetupFailure != null && terminalSetupFailure != cleanupFailure) { + terminalSetupFailure.addSuppressed(cleanupFailure); + } else { + Loggers.warnWithException( + LOG, "[{}] Failed to cancel stream ID reservation", logPrefix, cleanupFailure); } } } } + if (terminalSetupFailure != null) { + throttler.signalError(this, terminalSetupFailure); + } } else { channel.cancelPreAcquireId(); } diff --git a/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java b/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java index cf6875dd362..7fe69680794 100644 --- a/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java +++ b/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java @@ -19,6 +19,7 @@ import static com.datastax.oss.driver.api.core.DriverTimeoutException.UNAVAILABLE; +import com.datastax.dse.driver.api.core.config.DseDriverOption; import com.datastax.dse.driver.api.core.graph.AsyncGraphResultSet; import com.datastax.dse.driver.api.core.graph.GraphNode; import com.datastax.dse.driver.api.core.graph.GraphStatement; @@ -131,8 +132,10 @@ public class GraphRequestHandler implements Throttled { private final RequestThrottler throttler; private final RequestTracker requestTracker; private final SessionMetricUpdater sessionMetricUpdater; + private final DriverExecutionProfile initialExecutionProfile; private final GraphBinaryModule graphBinaryModule; private final GraphSupportChecker graphSupportChecker; + private volatile boolean admitted; // The errors on the nodes that were already tried (lazily initialized on the first error). // We don't use a map because nodes can appear multiple times. @@ -176,7 +179,12 @@ public class GraphRequestHandler implements Throttled { this.requestTracker = context.getRequestTracker(); this.sessionMetricUpdater = session.getMetricUpdater(); - Duration timeout = GraphConversions.resolveGraphRequestTimeout(statement, context); + this.initialExecutionProfile = Conversions.resolveExecutionProfile(statement, context); + Duration statementTimeout = statement.getTimeout(); + Duration timeout = + statementTimeout != null + ? statementTimeout + : initialExecutionProfile.getDuration(DseDriverOption.GRAPH_TIMEOUT); this.scheduledTimeout = scheduleTimeout(timeout); this.throttler = context.getRequestThrottler(); @@ -185,25 +193,31 @@ public class GraphRequestHandler implements Throttled { @Override public void onThrottleReady(boolean wasDelayed) { - DriverExecutionProfile executionProfile = - Conversions.resolveExecutionProfile(initialStatement, context); - if (wasDelayed - // avoid call to nanoTime() if metric is disabled: - && sessionMetricUpdater.isEnabled( - DefaultSessionMetric.THROTTLING_DELAY, executionProfile.getName())) { - sessionMetricUpdater.updateTimer( - DefaultSessionMetric.THROTTLING_DELAY, - executionProfile.getName(), - System.nanoTime() - startTimeNanos, - TimeUnit.NANOSECONDS); + admitted = true; + try { + if (wasDelayed + // avoid call to nanoTime() if metric is disabled: + && sessionMetricUpdater.isEnabled( + DefaultSessionMetric.THROTTLING_DELAY, initialExecutionProfile.getName())) { + sessionMetricUpdater.updateTimer( + DefaultSessionMetric.THROTTLING_DELAY, + initialExecutionProfile.getName(), + System.nanoTime() - startTimeNanos, + TimeUnit.NANOSECONDS); + } + Queue queryPlan = + initialStatement.getNode() != null + ? new SimpleQueryPlan(initialStatement.getNode()) + : context + .getLoadBalancingPolicyWrapper() + .newQueryPlan(initialStatement, initialExecutionProfile.getName(), session); + sendRequest(initialStatement, null, queryPlan, 0, 0, true); + } catch (Throwable t) { + // Concurrency-limiting throttlers contain exceptions from ready callbacks so that a failed + // request can't interrupt queue draining. Complete through the normal terminal path to + // release this request's permit before returning to the throttler. + setFinalError(initialStatement, t, null, NO_SUCCESSFUL_EXECUTION); } - Queue queryPlan = - initialStatement.getNode() != null - ? new SimpleQueryPlan(initialStatement.getNode()) - : context - .getLoadBalancingPolicyWrapper() - .newQueryPlan(initialStatement, executionProfile.getName(), session); - sendRequest(initialStatement, null, queryPlan, 0, 0, true); } public CompletionStage handle() { @@ -483,7 +497,9 @@ public void onThrottleFailure(@NonNull RequestThrottlingException error) { private void setFinalError( GraphStatement statement, Throwable error, Node node, int execution) { DriverExecutionProfile executionProfile = - Conversions.resolveExecutionProfile(statement, context); + statement == initialStatement + ? initialExecutionProfile + : Conversions.resolveExecutionProfile(statement, context); if (error instanceof DriverException) { ((DriverException) error) .setExecutionInfo( @@ -510,7 +526,8 @@ private void setFinalError( throttler.signalTimeout(this); sessionMetricUpdater.incrementCounter( DseSessionMetric.GRAPH_CLIENT_TIMEOUTS, executionProfile.getName()); - } else if (!(error instanceof RequestThrottlingException)) { + } else if (!(error instanceof CancellationException) + && (admitted || !(error instanceof RequestThrottlingException))) { throttler.signalError(this, error); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java index 16ad89f6126..b4ff8d314a2 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java @@ -24,6 +24,7 @@ import com.datastax.oss.driver.api.core.type.codec.TypeCodecs; import com.datastax.oss.driver.internal.core.channel.DriverChannel; import com.datastax.oss.driver.internal.core.channel.ResponseCallback; +import com.datastax.oss.driver.internal.core.util.Loggers; import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions; import com.datastax.oss.driver.shaded.guava.common.collect.Maps; import com.datastax.oss.protocol.internal.Frame; @@ -101,8 +102,8 @@ public static AdminRequestHandler query( private final Class expectedResponseType; protected final CompletableFuture result = new CompletableFuture<>(); - // This is only ever accessed on the channel's event loop, so it doesn't need to be volatile - private ScheduledFuture timeoutFuture; + private volatile ScheduledFuture timeoutFuture; + private volatile boolean writeSubmitted; protected AdminRequestHandler( DriverChannel channel, @@ -134,7 +135,7 @@ public CompletionStage start() { boolean writeSubmitted = false; try { Future writeFuture = channel.write(message, false, customPayload, this); - writeSubmitted = true; + this.writeSubmitted = writeSubmitted = true; writeFuture.addListener(this::onWriteComplete); } finally { if (!writeSubmitted) { @@ -161,13 +162,22 @@ protected final void cancelCallerOwnedPreAcquireId() { private void onWriteComplete(Future future) { if (future.isSuccess()) { + if (result.isDone()) { + cancelSubmittedRequest(); + return; + } LOG.debug("[{}] Successfully wrote {}, waiting for response", logPrefix, this); if (timeout.toNanos() > 0) { - timeoutFuture = + ScheduledFuture timeoutFuture = channel .eventLoop() .schedule(this::fireTimeout, timeout.toNanos(), TimeUnit.NANOSECONDS); + this.timeoutFuture = timeoutFuture; timeoutFuture.addListener(UncaughtExceptions::log); + // Terminal completion can race with timeout installation after the check above. + if (result.isDone()) { + cancelSubmittedRequest(); + } } } else { setFinalError(future.cause()); @@ -177,8 +187,27 @@ private void onWriteComplete(Future future) { private void fireTimeout() { setFinalError( new DriverTimeoutException(String.format("%s timed out after %s", debugString, timeout))); - if (!channel.closeFuture().isDone()) { - channel.cancel(this); + cancelSubmittedRequest(); + } + + /** Cancels this callback if its write was submitted, along with any installed timeout. */ + protected final void cancelSubmittedRequest() { + if (!writeSubmitted) { + return; + } + try { + if (timeoutFuture != null) { + timeoutFuture.cancel(true); + } + } catch (Throwable t) { + Loggers.warnWithException(LOG, "[{}] Error cancelling timeout for {}", logPrefix, this, t); + } + try { + if (!channel.closeFuture().isDone()) { + channel.cancel(this); + } + } catch (Throwable t) { + Loggers.warnWithException(LOG, "[{}] Error cancelling {}", logPrefix, this, t); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java index d14f4ee08c4..45ee5c63df1 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java @@ -105,6 +105,7 @@ public static ThrottledAdminRequestHandler prepare( private final long startTimeNanos; private final RequestThrottler throttler; private final SessionMetricUpdater metricUpdater; + private volatile boolean admitted; private final AtomicBoolean holdsExternalReservation; protected ThrottledAdminRequestHandler( @@ -140,9 +141,13 @@ public CompletionStage start() { throttler.register(this); } catch (Throwable t) { cancelExternalReservation(); - // Registration can fail before the throttler admits this request, so complete the result - // without calling this class's override, which would signal a permit that was never acquired. - super.setFinalError(t); + if (admitted) { + setFinalError(t); + cancelSubmittedRequest(); + } else { + // Registration failed before admission, so there is no throttler permit to release. + super.setFinalError(t); + } throw t; } return result; @@ -150,6 +155,7 @@ public CompletionStage start() { @Override public void onThrottleReady(boolean wasDelayed) { + admitted = true; try { if (wasDelayed) { metricUpdater.updateTimer( @@ -163,7 +169,7 @@ public void onThrottleReady(boolean wasDelayed) { } catch (Throwable t) { cancelExternalReservation(); setFinalError(t); - throw t; + cancelSubmittedRequest(); } } @@ -175,8 +181,6 @@ public void onThrottleFailure(@NonNull RequestThrottlingException error) { } private void cancelExternalReservation() { - // register() can invoke onThrottleReady() synchronously. If that callback throws, both - // onThrottleReady() and start() catch the same failure, so cancellation must be idempotent. if (holdsExternalReservation.compareAndSet(true, false)) { cancelCallerOwnedPreAcquireId(); } @@ -197,7 +201,7 @@ protected boolean setFinalError(Throwable error) { if (wasSet) { if (error instanceof DriverTimeoutException) { throttler.signalTimeout(this); - } else if (!(error instanceof RequestThrottlingException)) { + } else if (admitted) { throttler.signalError(this, error); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java index b05fec2cb7e..ab0b60d5b31 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java @@ -98,6 +98,7 @@ public class CqlPrepareHandler implements Throttled { private final RequestThrottler throttler; private final Boolean prepareOnAllNodes; private final DriverExecutionProfile executionProfile; + private volatile boolean admitted; private volatile InitialPrepareCallback initialCallback; // The errors on the nodes that were already tried (lazily initialized on the first error). @@ -148,18 +149,24 @@ protected CqlPrepareHandler( @Override public void onThrottleReady(boolean wasDelayed) { - DriverExecutionProfile executionProfile = - Conversions.resolveExecutionProfile(initialRequest, context); - if (wasDelayed) { - session - .getMetricUpdater() - .updateTimer( - DefaultSessionMetric.THROTTLING_DELAY, - executionProfile.getName(), - System.nanoTime() - startTimeNanos, - TimeUnit.NANOSECONDS); + admitted = true; + try { + if (wasDelayed) { + session + .getMetricUpdater() + .updateTimer( + DefaultSessionMetric.THROTTLING_DELAY, + executionProfile.getName(), + System.nanoTime() - startTimeNanos, + TimeUnit.NANOSECONDS); + } + sendRequest(initialRequest, null, 0); + } catch (Throwable t) { + // Concurrency-limiting throttlers contain exceptions from ready callbacks so that a failed + // request can't interrupt queue draining. Complete through the normal terminal path to + // release this request's permit before returning to the throttler. + setFinalError(t); } - sendRequest(initialRequest, null, 0); } public CompletableFuture handle() { @@ -370,7 +377,8 @@ private void setFinalError(Throwable error) { cancelTimeout(); if (error instanceof DriverTimeoutException) { throttler.signalTimeout(this); - } else if (!(error instanceof RequestThrottlingException)) { + } else if (!(error instanceof CancellationException) + && (admitted || !(error instanceof RequestThrottlingException))) { throttler.signalError(this, error); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index ce4b40d6d29..bfcf3d1fc4d 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -46,6 +46,7 @@ import com.datastax.oss.driver.api.core.metadata.token.Token; import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; +import com.datastax.oss.driver.api.core.retry.RetryDecision; import com.datastax.oss.driver.api.core.retry.RetryPolicy; import com.datastax.oss.driver.api.core.retry.RetryVerdict; import com.datastax.oss.driver.api.core.servererrors.BootstrappingException; @@ -110,7 +111,9 @@ import java.util.concurrent.CompletionStage; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import net.jcip.annotations.ThreadSafe; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -142,7 +145,7 @@ public class CqlRequestHandler implements Throttled { */ private final AtomicInteger startedSpeculativeExecutionsCount; - final Timeout scheduledTimeout; + volatile Timeout scheduledTimeout; final List scheduledExecutions; private final List inFlightCallbacks; private final RequestThrottler throttler; @@ -150,6 +153,10 @@ public class CqlRequestHandler implements Throttled { private final Optional requestIdGenerator; private final SessionMetricUpdater sessionMetricUpdater; private final DriverExecutionProfile executionProfile; + private volatile boolean admitted; + private final AtomicBoolean registrationComplete = new AtomicBoolean(); + private final AtomicBoolean throttlerReleased = new AtomicBoolean(); + private final AtomicReference terminalError = new AtomicReference<>(); // The errors on the nodes that were already tried (lazily initialized on the first error). // We don't use a map because nodes can appear multiple times. @@ -179,13 +186,15 @@ protected CqlRequestHandler( this.session = session; this.keyspace = session.getKeyspace().orElse(null); this.context = context; + this.throttler = context.getRequestThrottler(); this.result = new CompletableFuture<>(); this.result.exceptionally( t -> { try { if (t instanceof CancellationException) { + terminalError.compareAndSet(null, t); cancelScheduledTasks(); - context.getRequestThrottler().signalCancel(this); + releaseThrottler(t); } } catch (Throwable t2) { Loggers.warnWithException(LOG, "[{}] Uncaught exception", handlerLogPrefix, t2); @@ -205,34 +214,63 @@ protected CqlRequestHandler( this.executionProfile = Conversions.resolveExecutionProfile(initialStatement, context); Duration timeout = Conversions.resolveRequestTimeout(statement, executionProfile); this.scheduledTimeout = scheduleTimeout(timeout); - - this.throttler = context.getRequestThrottler(); - this.throttler.register(this); + if (!result.isDone()) { + try { + this.throttler.register(this); + registrationComplete.set(true); + } catch (Throwable t) { + if (admitted) { + setFinalError(initialStatement, t, null, -1); + } else { + // Registration failed before admission, so there is no throttler permit to release. + setFinalErrorWithoutThrottlerRelease(initialStatement, t, null, -1); + } + throw t; + } + // The timeout can fire while register() is still deciding whether to admit or enqueue the + // request. Once registration returns, finish any release that had to be deferred until the + // throttler was known to own the request. + if (result.isDone()) { + releaseThrottler(terminalError.get()); + } + } } @Override public void onThrottleReady(boolean wasDelayed) { - if (wasDelayed - // avoid call to nanoTime() if metric is disabled: - && sessionMetricUpdater.isEnabled( - DefaultSessionMetric.THROTTLING_DELAY, executionProfile.getName())) { - sessionMetricUpdater.updateTimer( - DefaultSessionMetric.THROTTLING_DELAY, - executionProfile.getName(), - System.nanoTime() - startTimeNanos, - TimeUnit.NANOSECONDS); - } - Queue queryPlan; - if (this.initialStatement.getNode() != null) { - queryPlan = new SimpleQueryPlan(this.initialStatement.getNode()); - } else { - queryPlan = - context - .getLoadBalancingPolicyWrapper() - .newQueryPlan(initialStatement, executionProfile.getName(), session); + admitted = true; + if (result.isDone()) { + releaseThrottler(terminalError.get()); + return; } + try { + if (wasDelayed + // avoid call to nanoTime() if metric is disabled: + && sessionMetricUpdater.isEnabled( + DefaultSessionMetric.THROTTLING_DELAY, executionProfile.getName())) { + sessionMetricUpdater.updateTimer( + DefaultSessionMetric.THROTTLING_DELAY, + executionProfile.getName(), + System.nanoTime() - startTimeNanos, + TimeUnit.NANOSECONDS); + } + Queue queryPlan; + if (this.initialStatement.getNode() != null) { + queryPlan = new SimpleQueryPlan(this.initialStatement.getNode()); + } else { + queryPlan = + context + .getLoadBalancingPolicyWrapper() + .newQueryPlan(initialStatement, executionProfile.getName(), session); + } - sendRequest(initialStatement, null, queryPlan, 0, 0, true); + sendRequest(initialStatement, null, queryPlan, 0, 0, true); + } catch (Throwable t) { + // Throttlers invoke this callback synchronously after admitting the request. Contain setup + // failures so the normal terminal path releases the permit and scheduled work, and so a + // delayed request can't throw through the completion path of the request that admitted it. + setFinalError(initialStatement, t, null, -1); + } } public CompletionStage handle() { @@ -255,12 +293,15 @@ private Timeout scheduleTimeout(Duration timeoutDuration) { timeoutDuration.toNanos(), TimeUnit.NANOSECONDS); } catch (IllegalStateException e) { - // If we raced with session shutdown the timer might be closed already, rethrow with a more - // explicit message - result.completeExceptionally( + // Timeout installation precedes throttler registration. Complete the request without + // signaling a throttler that does not own it yet. + setFinalErrorWithoutThrottlerRelease( + initialStatement, "cannot be started once stopped".equals(e.getMessage()) ? new IllegalStateException("Session is closed") - : e); + : e, + null, + -1); } } return null; @@ -380,27 +421,43 @@ private void sendRequest( } Node node = retriedNode; DriverChannel channel = null; - if (node == null - || (channel = + try { + Token routingToken = getRoutingToken(statement); + if (node != null) { + try { + channel = + session.getChannel( + node, + handlerLogPrefix, + routingToken, + getShardFromTabletMap(statement, node, routingToken)); + } catch (Throwable t) { + recordError(node, t); + } + } + if (channel == null) { + while (!result.isDone() && (node = queryPlan.poll()) != null) { + try { + channel = session.getChannel( node, handlerLogPrefix, - getRoutingToken(statement), - getShardFromTabletMap(statement, node, getRoutingToken(statement)))) - == null) { - while (!result.isDone() && (node = queryPlan.poll()) != null) { - channel = - session.getChannel( - node, - handlerLogPrefix, - getRoutingToken(statement), - getShardFromTabletMap(statement, node, getRoutingToken(statement))); - if (channel != null) { - break; - } else { - recordError(node, new NodeUnavailableException(node)); + routingToken, + getShardFromTabletMap(statement, node, routingToken)); + } catch (Throwable t) { + recordError(node, t); + continue; + } + if (channel != null) { + break; + } else { + recordError(node, new NodeUnavailableException(node)); + } } } + } catch (Throwable t) { + handleRequestSetupFailure(statement, t, currentExecutionIndex); + return; } if (channel == null) { // We've reached the end of the query plan without finding any node to write to @@ -410,6 +467,7 @@ private void sendRequest( } } else { boolean writeSubmitted = false; + Throwable setupFailure = null; try { Statement finalStatement = statement; String nodeRequestId = @@ -437,11 +495,46 @@ private void sendRequest( message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback); writeSubmitted = true; writeFuture.addListener(nodeResponseCallback); + } catch (Throwable t) { + setupFailure = t; } finally { if (!writeSubmitted) { - channel.cancelPreAcquireId(); + try { + channel.cancelPreAcquireId(); + } catch (Throwable t) { + if (setupFailure == null) { + setupFailure = t; + } else if (setupFailure != t) { + setupFailure.addSuppressed(t); + } + } } } + if (setupFailure != null) { + handleRequestSetupFailure(statement, setupFailure, currentExecutionIndex); + } + } + } + + private void handleRequestSetupFailure(Statement statement, Throwable error, int execution) { + if (result.isDone()) { + Loggers.warnWithException( + LOG, + "[{}] Request setup failed after the request had completed", + handlerLogPrefix, + error); + return; + } + if (activeExecutionsCount.decrementAndGet() == 0) { + // Setup failed before channel.write() accepted a node request, so this is a session-level + // failure rather than a failed node attempt. + setFinalError(statement, error, null, execution); + } else { + Loggers.warnWithException( + LOG, + "[{}] Request setup failed, another execution is still active", + handlerLogPrefix, + error); } } @@ -485,7 +578,9 @@ private void setFinalResult( Conversions.toResultSet(resultMessage, executionInfo, session, context); if (result.complete(resultSet)) { cancelScheduledTasks(); - throttler.signalSuccess(this); + if (throttlerReleased.compareAndSet(false, true)) { + throttler.signalSuccess(this); + } // Only call nanoTime() if we're actually going to use it long completionTimeNanos = NANOTIME_NOT_MEASURED_YET, @@ -606,10 +701,23 @@ private ExecutionInfo buildExecutionInfo( public void onThrottleFailure(@NonNull RequestThrottlingException error) { sessionMetricUpdater.incrementCounter( DefaultSessionMetric.THROTTLING_ERRORS, executionProfile.getName()); + // The throttler rejected this request, so there is no admission to release. + throttlerReleased.set(true); setFinalError(initialStatement, error, null, -1); } private void setFinalError(Statement statement, Throwable error, Node node, int execution) { + setFinalError(statement, error, node, execution, true); + } + + private void setFinalErrorWithoutThrottlerRelease( + Statement statement, Throwable error, Node node, int execution) { + setFinalError(statement, error, node, execution, false); + } + + private void setFinalError( + Statement statement, Throwable error, Node node, int execution, boolean releaseThrottler) { + terminalError.compareAndSet(null, error); if (error instanceof DriverException) { ((DriverException) error) .setExecutionInfo( @@ -633,16 +741,39 @@ private void setFinalError(Statement statement, Throwable error, Node node, i requestTracker.onError( statement, error, latencyNanos, executionProfile, node, handlerLogPrefix); } + if (releaseThrottler) { + releaseThrottler(error); + } if (error instanceof DriverTimeoutException) { - throttler.signalTimeout(this); sessionMetricUpdater.incrementCounter( DefaultSessionMetric.CQL_CLIENT_TIMEOUTS, executionProfile.getName()); - } else if (!(error instanceof RequestThrottlingException)) { - throttler.signalError(this, error); } } } + private void releaseThrottler(@Nullable Throwable error) { + if (!admitted && !registrationComplete.get()) { + // A timeout can race with register(), but signaling before registration has completed can + // corrupt throttler accounting if the request has not been enqueued yet. onThrottleReady() + // releases immediately once admission proves ownership; otherwise the constructor retries + // after register() returns. + return; + } + if (!throttlerReleased.compareAndSet(false, true)) { + return; + } + if (error instanceof DriverTimeoutException) { + throttler.signalTimeout(this); + } else if (error instanceof CancellationException || !admitted) { + // Before admission this removes a queued request. If admission raced with this call, the + // throttler treats it as completion of the transferred permit. + throttler.signalCancel(this); + } else { + throttler.signalError( + this, error == null ? new IllegalStateException("Request failed") : error); + } + } + /** * Handles the interaction with a single node in the query plan. * @@ -946,6 +1077,9 @@ private void processErrorResponse(Error errorMessage) { } else { RetryPolicy retryPolicy = Conversions.resolveRetryPolicy(context, executionProfile); RetryVerdict verdict; + DefaultNodeMetric errorMetric; + DefaultNodeMetric retriesOnErrorMetric; + DefaultNodeMetric ignoresOnErrorMetric; if (error instanceof ReadTimeoutException) { ReadTimeoutException readTimeout = (ReadTimeoutException) error; verdict = @@ -956,12 +1090,9 @@ private void processErrorResponse(Error errorMessage) { readTimeout.getReceived(), readTimeout.wasDataPresent(), retryCount); - updateErrorMetrics( - metricUpdater, - verdict, - DefaultNodeMetric.READ_TIMEOUTS, - DefaultNodeMetric.RETRIES_ON_READ_TIMEOUT, - DefaultNodeMetric.IGNORES_ON_READ_TIMEOUT); + errorMetric = DefaultNodeMetric.READ_TIMEOUTS; + retriesOnErrorMetric = DefaultNodeMetric.RETRIES_ON_READ_TIMEOUT; + ignoresOnErrorMetric = DefaultNodeMetric.IGNORES_ON_READ_TIMEOUT; } else if (error instanceof WriteTimeoutException) { WriteTimeoutException writeTimeout = (WriteTimeoutException) error; verdict = @@ -974,12 +1105,9 @@ private void processErrorResponse(Error errorMessage) { writeTimeout.getReceived(), retryCount) : RetryVerdict.RETHROW; - updateErrorMetrics( - metricUpdater, - verdict, - DefaultNodeMetric.WRITE_TIMEOUTS, - DefaultNodeMetric.RETRIES_ON_WRITE_TIMEOUT, - DefaultNodeMetric.IGNORES_ON_WRITE_TIMEOUT); + errorMetric = DefaultNodeMetric.WRITE_TIMEOUTS; + retriesOnErrorMetric = DefaultNodeMetric.RETRIES_ON_WRITE_TIMEOUT; + ignoresOnErrorMetric = DefaultNodeMetric.IGNORES_ON_WRITE_TIMEOUT; } else if (error instanceof UnavailableException) { UnavailableException unavailable = (UnavailableException) error; verdict = @@ -989,52 +1117,70 @@ private void processErrorResponse(Error errorMessage) { unavailable.getRequired(), unavailable.getAlive(), retryCount); - updateErrorMetrics( - metricUpdater, - verdict, - DefaultNodeMetric.UNAVAILABLES, - DefaultNodeMetric.RETRIES_ON_UNAVAILABLE, - DefaultNodeMetric.IGNORES_ON_UNAVAILABLE); + errorMetric = DefaultNodeMetric.UNAVAILABLES; + retriesOnErrorMetric = DefaultNodeMetric.RETRIES_ON_UNAVAILABLE; + ignoresOnErrorMetric = DefaultNodeMetric.IGNORES_ON_UNAVAILABLE; } else { verdict = Conversions.resolveIdempotence(statement, executionProfile) ? retryPolicy.onErrorResponseVerdict(statement, error, retryCount) : RetryVerdict.RETHROW; + errorMetric = DefaultNodeMetric.OTHER_ERRORS; + retriesOnErrorMetric = DefaultNodeMetric.RETRIES_ON_OTHER_ERROR; + ignoresOnErrorMetric = DefaultNodeMetric.IGNORES_ON_OTHER_ERROR; + } + RetryDecision decision = getRetryDecisionOrHandleFailure(verdict, error); + if (decision != null) { updateErrorMetrics( - metricUpdater, - verdict, - DefaultNodeMetric.OTHER_ERRORS, - DefaultNodeMetric.RETRIES_ON_OTHER_ERROR, - DefaultNodeMetric.IGNORES_ON_OTHER_ERROR); + metricUpdater, decision, errorMetric, retriesOnErrorMetric, ignoresOnErrorMetric); + processRetryVerdict(verdict, decision, error); + } + } + } + + @Nullable + private RetryDecision getRetryDecisionOrHandleFailure( + RetryVerdict verdict, Throwable requestError) { + try { + RetryDecision decision = verdict.getRetryDecision(); + if (decision == null) { + throw new NullPointerException("Retry verdict returned a null decision"); } - processRetryVerdict(verdict, error); + return decision; + } catch (Throwable setupFailure) { + handleRetrySetupFailure(requestError, setupFailure); + return null; } } - private void processRetryVerdict(RetryVerdict verdict, Throwable error) { - LOG.trace("[{}] Processing retry decision {}", logPrefix, verdict); - switch (verdict.getRetryDecision()) { + private void handleRetrySetupFailure(Throwable requestError, Throwable setupFailure) { + recordError(node, requestError); + trackNodeError(node, requestError, NANOTIME_NOT_MEASURED_YET); + handleRequestSetupFailure(statement, setupFailure, execution); + } + + private void processRetryVerdict( + RetryVerdict verdict, RetryDecision decision, Throwable error) { + LOG.trace("[{}] Processing retry decision {}", logPrefix, decision); + Statement retryStatement = null; + if (decision == RetryDecision.RETRY_SAME || decision == RetryDecision.RETRY_NEXT) { + try { + retryStatement = verdict.getRetryRequest(statement); + } catch (Throwable t) { + handleRetrySetupFailure(error, t); + return; + } + } + switch (decision) { case RETRY_SAME: recordError(node, error); trackNodeError(node, error, NANOTIME_NOT_MEASURED_YET); - sendRequest( - verdict.getRetryRequest(statement), - node, - queryPlan, - execution, - retryCount + 1, - false); + sendRequest(retryStatement, node, queryPlan, execution, retryCount + 1, false); break; case RETRY_NEXT: recordError(node, error); trackNodeError(node, error, NANOTIME_NOT_MEASURED_YET); - sendRequest( - verdict.getRetryRequest(statement), - null, - queryPlan, - execution, - retryCount + 1, - false); + sendRequest(retryStatement, null, queryPlan, execution, retryCount + 1, false); break; case RETHROW: trackNodeError(node, error, NANOTIME_NOT_MEASURED_YET); @@ -1048,12 +1194,12 @@ private void processRetryVerdict(RetryVerdict verdict, Throwable error) { private void updateErrorMetrics( NodeMetricUpdater metricUpdater, - RetryVerdict verdict, + RetryDecision decision, DefaultNodeMetric error, DefaultNodeMetric retriesOnError, DefaultNodeMetric ignoresOnError) { metricUpdater.incrementCounter(error, executionProfile.getName()); - switch (verdict.getRetryDecision()) { + switch (decision) { case RETRY_SAME: case RETRY_NEXT: metricUpdater.incrementCounter(DefaultNodeMetric.RETRIES, executionProfile.getName()); @@ -1076,26 +1222,30 @@ public void onFailure(Throwable error) { } LOG.trace("[{}] Request failure, processing: {}", logPrefix, error); RetryVerdict verdict; - if (!Conversions.resolveIdempotence(statement, executionProfile) - || error instanceof FrameTooLongException) { - verdict = RetryVerdict.RETHROW; - } else { - try { + try { + if (!Conversions.resolveIdempotence(statement, executionProfile) + || error instanceof FrameTooLongException) { + verdict = RetryVerdict.RETHROW; + } else { RetryPolicy retryPolicy = Conversions.resolveRetryPolicy(context, executionProfile); verdict = retryPolicy.onRequestAbortedVerdict(statement, error, retryCount); - } catch (Throwable cause) { - setFinalError( - statement, - new IllegalStateException("Unexpected error while invoking the retry policy", cause), - null, - execution); - return; } + } catch (Throwable cause) { + setFinalError( + statement, + new IllegalStateException("Unexpected error while invoking the retry policy", cause), + null, + execution); + return; + } + RetryDecision decision = getRetryDecisionOrHandleFailure(verdict, error); + if (decision == null) { + return; } - processRetryVerdict(verdict, error); + processRetryVerdict(verdict, decision, error); updateErrorMetrics( ((DefaultNode) node).getMetricUpdater(), - verdict, + decision, DefaultNodeMetric.ABORTED_REQUESTS, DefaultNodeMetric.RETRIES_ON_ABORTED, DefaultNodeMetric.IGNORES_ON_ABORTED); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottler.java index 8146c5b113a..6e183079c1b 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottler.java @@ -23,9 +23,11 @@ import com.datastax.oss.driver.api.core.context.DriverContext; import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; import com.datastax.oss.driver.api.core.session.throttling.Throttled; +import com.datastax.oss.driver.internal.core.util.Loggers; import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting; import edu.umd.cs.findbugs.annotations.NonNull; import edu.umd.cs.findbugs.annotations.Nullable; +import java.util.ArrayDeque; import java.util.Deque; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.atomic.AtomicInteger; @@ -56,6 +58,9 @@ public class ConcurrencyLimitingRequestThrottler implements RequestThrottler { private static final Logger LOG = LoggerFactory.getLogger(ConcurrencyLimitingRequestThrottler.class); + // Completion can synchronously admit another request. Trampoline those callbacks to keep queue + // draining iterative instead of growing the call stack once per failed request. + private static final ThreadLocal READY_CALLBACKS = new ThreadLocal<>(); private final String logPrefix; private final int maxConcurrentRequests; @@ -99,7 +104,7 @@ public void register(@NonNull Throttled request) { int newConcurrent = concurrentRequests.incrementAndGet(); if (newConcurrent <= maxConcurrentRequests) { LOG.trace("[{}] Starting newly registered request", logPrefix); - request.onThrottleReady(false); + notifyReady(request, false); return; } else { // We exceeded the limit, decrement the count and fall through to the queuing logic @@ -139,7 +144,7 @@ public void register(@NonNull Throttled request) { public void signalSuccess(@NonNull Throttled request) { Throttled nextRequest = onRequestDoneAndDequeNext(); if (nextRequest != null) { - nextRequest.onThrottleReady(true); + notifyReady(nextRequest, true); } } @@ -161,7 +166,7 @@ public void signalTimeout(@NonNull Throttled request) { } if (nextRequest != null) { - nextRequest.onThrottleReady(true); + notifyReady(nextRequest, true); } } @@ -178,7 +183,45 @@ public void signalCancel(@NonNull Throttled request) { } if (nextRequest != null) { - nextRequest.onThrottleReady(true); + notifyReady(nextRequest, true); + } + } + + private void notifyReady(Throttled request, boolean wasDelayed) { + ReadyCallbackState previous = READY_CALLBACKS.get(); + for (ReadyCallbackState state = previous; state != null; state = state.previous) { + if (state.throttler == this) { + state.add(request, wasDelayed); + return; + } + } + + ReadyCallbackState state = new ReadyCallbackState(this, previous); + READY_CALLBACKS.set(state); + try { + invokeReady(request, wasDelayed); + ReadyCallback callback; + while ((callback = state.poll()) != null) { + invokeReady(callback.request, callback.wasDelayed); + } + } finally { + if (previous == null) { + READY_CALLBACKS.remove(); + } else { + READY_CALLBACKS.set(previous); + } + } + } + + private void invokeReady(Throttled request, boolean wasDelayed) { + try { + request.onThrottleReady(wasDelayed); + } catch (Throwable t) { + // A callback can synchronously complete its request and enqueue more ready callbacks before + // throwing. Keep draining those already-admitted requests, and don't propagate a request's + // failure through the unrelated request whose completion triggered the drain. + Loggers.warnWithException( + LOG, "[{}] Uncaught exception in throttled request callback", logPrefix, t); } } @@ -228,4 +271,38 @@ Deque getQueue() { private static void fail(Throttled request, String message) { request.onThrottleFailure(new RequestThrottlingException(message)); } + + private static final class ReadyCallback { + private final Throttled request; + private final boolean wasDelayed; + + private ReadyCallback(Throttled request, boolean wasDelayed) { + this.request = request; + this.wasDelayed = wasDelayed; + } + } + + private static final class ReadyCallbackState { + private final ConcurrencyLimitingRequestThrottler throttler; + @Nullable private final ReadyCallbackState previous; + @Nullable private Deque callbacks; + + private ReadyCallbackState( + ConcurrencyLimitingRequestThrottler throttler, @Nullable ReadyCallbackState previous) { + this.throttler = throttler; + this.previous = previous; + } + + private void add(Throttled request, boolean wasDelayed) { + if (callbacks == null) { + callbacks = new ArrayDeque<>(); + } + callbacks.addLast(new ReadyCallback(request, wasDelayed)); + } + + @Nullable + private ReadyCallback poll() { + return callbacks == null ? null : callbacks.pollFirst(); + } + } } diff --git a/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java b/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java index e24579e7fe6..c4bb97d9d8d 100644 --- a/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java @@ -51,18 +51,21 @@ import com.datastax.oss.driver.api.core.cql.ExecutionInfo; import com.datastax.oss.driver.api.core.cql.Row; import com.datastax.oss.driver.api.core.cql.SimpleStatement; +import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; import com.datastax.oss.driver.api.core.servererrors.BootstrappingException; import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; import com.datastax.oss.driver.api.core.tracker.RequestTracker; import com.datastax.oss.driver.internal.core.ProtocolFeature; import com.datastax.oss.driver.internal.core.cql.PoolBehavior; import com.datastax.oss.driver.internal.core.cql.RequestHandlerTestHarness; +import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures; import com.datastax.oss.protocol.internal.Message; import com.datastax.oss.protocol.internal.ProtocolConstants; import com.tngtech.java.junit.dataprovider.UseDataProvider; import java.util.Iterator; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletionStage; import java.util.concurrent.TimeUnit; import java.util.regex.Pattern; @@ -231,7 +234,7 @@ public void should_unwind_execution_if_request_setup_fails_before_write() { statement, harness.getSession(), harness.getContext(), "test"); CompletionStage resultSetFuture = handler.handle(); - assertThatThrownBy(() -> handler.onThrottleReady(false)).isSameAs(failure); + handler.onThrottleReady(false); assertThatStage(resultSetFuture).isFailed(error -> assertThat(error).isSameAs(failure)); assertThat(handler.getActiveExecutionsCount()).isZero(); @@ -242,6 +245,66 @@ public void should_unwind_execution_if_request_setup_fails_before_write() { } } + @Test + public void should_release_permit_if_delayed_throttling_metric_update_fails() { + RuntimeException failure = new RuntimeException("mock failure"); + RequestThrottler throttler = mock(RequestThrottler.class); + RequestHandlerTestHarness.Builder builder = + continuousHarnessBuilder().withProtocolVersion(DSE_V2); + PoolBehavior node1Behavior = builder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = builder.build()) { + SessionMetricUpdater metricUpdater = harness.getSession().getMetricUpdater(); + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + when(metricUpdater.isEnabled( + DefaultSessionMetric.THROTTLING_DELAY, DriverExecutionProfile.DEFAULT_NAME)) + .thenReturn(true); + doThrow(failure) + .when(metricUpdater) + .updateTimer( + eq(DefaultSessionMetric.THROTTLING_DELAY), + eq(DriverExecutionProfile.DEFAULT_NAME), + anyLong(), + eq(TimeUnit.NANOSECONDS)); + ContinuousCqlRequestHandler handler = + new ContinuousCqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + CompletionStage resultSetFuture = handler.handle(); + + handler.onThrottleReady(true); + + assertThatStage(resultSetFuture).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(handler.getActiveExecutionsCount()).isZero(); + node1Behavior.verifyNoWrite(); + verify(throttler).signalError(handler, failure); + } + } + + @Test + public void should_release_cancelled_request_only_once() { + CancellationException failure = new CancellationException("mock cancellation"); + SimpleStatement statement = Mockito.spy(SimpleStatement.newInstance("mock query")); + doThrow(failure).when(statement).getCustomPayload(); + RequestThrottler throttler = mock(RequestThrottler.class); + RequestHandlerTestHarness.Builder builder = + continuousHarnessBuilder().withProtocolVersion(DSE_V2); + builder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = builder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + ContinuousCqlRequestHandler handler = + new ContinuousCqlRequestHandler( + statement, harness.getSession(), harness.getContext(), "test"); + CompletionStage resultSetFuture = handler.handle(); + + handler.onThrottleReady(false); + + assertThat(resultSetFuture.toCompletableFuture()).isCancelled(); + verify(throttler).signalCancel(handler); + verify(throttler, never()).signalError(eq(handler), any()); + } + } + @Test @UseDataProvider(value = "allDseProtocolVersions", location = DseTestDataProviders.class) public void should_time_out_if_first_page_takes_too_long(DseProtocolVersion version) diff --git a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerTest.java b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerTest.java index b374539f12e..71795a45d9c 100644 --- a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerTest.java @@ -23,8 +23,11 @@ import static com.datastax.oss.driver.Assertions.assertThat; import static com.datastax.oss.driver.Assertions.assertThatStage; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; @@ -40,8 +43,11 @@ import com.datastax.dse.driver.api.core.metrics.DseSessionMetric; import com.datastax.dse.driver.internal.core.graph.binary.GraphBinaryModule; import com.datastax.oss.driver.api.core.DriverTimeoutException; +import com.datastax.oss.driver.api.core.RequestThrottlingException; import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; import com.datastax.oss.driver.api.core.cql.ExecutionInfo; +import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; import com.datastax.oss.driver.internal.core.context.DefaultDriverContext; import com.datastax.oss.driver.internal.core.cql.PoolBehavior; import com.datastax.oss.driver.internal.core.cql.RequestHandlerTestHarness; @@ -55,6 +61,7 @@ import java.time.Duration; import java.util.concurrent.CompletionStage; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -193,6 +200,115 @@ public void should_honor_default_timeout() throws Exception { } } + @Test + public void should_not_schedule_timeout_after_immediate_setup_failure() { + RuntimeException failure = new RuntimeException("mock failure"); + Duration defaultTimeout = Duration.ofSeconds(1); + GraphSupportChecker supportChecker = mock(GraphSupportChecker.class); + when(supportChecker.inferGraphProtocol(any(), any(), any())).thenThrow(failure); + + RequestHandlerTestHarness.Builder builder = + GraphRequestHandlerTestHarness.builder().withGraphTimeout(defaultTimeout); + PoolBehavior node1Behavior = builder.customBehavior(node); + + try (RequestHandlerTestHarness harness = builder.build()) { + ContinuousGraphRequestHandler handler = + new ContinuousGraphRequestHandler( + ScriptGraphStatement.newInstance("mockQuery"), + harness.getSession(), + harness.getContext(), + "test", + createGraphBinaryModule(mockContext), + supportChecker); + + CompletionStage result = handler.handle(); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(harness.nextScheduledTimeout()).isNull(); + node1Behavior.verifyNoWrite(); + node1Behavior.verifyPreAcquireCancelled(); + } + } + + @Test + public void should_cancel_timeout_after_delayed_setup_failure() { + RuntimeException failure = new RuntimeException("mock failure"); + Duration defaultTimeout = Duration.ofSeconds(1); + GraphSupportChecker supportChecker = mock(GraphSupportChecker.class); + when(supportChecker.inferGraphProtocol(any(), any(), any())).thenThrow(failure); + RequestThrottler throttler = mock(RequestThrottler.class); + AtomicReference registeredRequest = new AtomicReference<>(); + doAnswer( + invocation -> { + registeredRequest.set(invocation.getArgument(0)); + return null; + }) + .when(throttler) + .register(any()); + + RequestHandlerTestHarness.Builder builder = + GraphRequestHandlerTestHarness.builder().withGraphTimeout(defaultTimeout); + PoolBehavior node1Behavior = builder.customBehavior(node); + + try (RequestHandlerTestHarness harness = builder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + ContinuousGraphRequestHandler handler = + new ContinuousGraphRequestHandler( + ScriptGraphStatement.newInstance("mockQuery"), + harness.getSession(), + harness.getContext(), + "test", + createGraphBinaryModule(mockContext), + supportChecker); + CompletionStage result = handler.handle(); + CapturedTimeout globalTimeout = harness.nextScheduledTimeout(); + + registeredRequest.get().onThrottleReady(true); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(globalTimeout.isCancelled()).isTrue(); + node1Behavior.verifyNoWrite(); + node1Behavior.verifyPreAcquireCancelled(); + verify(throttler).signalError(handler, failure); + } + } + + @Test + public void should_cancel_timeout_when_queued_request_is_rejected() { + Duration defaultTimeout = Duration.ofSeconds(1); + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + RequestThrottler throttler = mock(RequestThrottler.class); + AtomicReference registeredRequest = new AtomicReference<>(); + doAnswer( + invocation -> { + registeredRequest.set(invocation.getArgument(0)); + return null; + }) + .when(throttler) + .register(any()); + + RequestHandlerTestHarness.Builder builder = + GraphRequestHandlerTestHarness.builder().withGraphTimeout(defaultTimeout); + try (RequestHandlerTestHarness harness = builder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + ContinuousGraphRequestHandler handler = + new ContinuousGraphRequestHandler( + ScriptGraphStatement.newInstance("mockQuery"), + harness.getSession(), + harness.getContext(), + "test", + createGraphBinaryModule(mockContext), + new GraphSupportChecker()); + + CompletionStage result = handler.handle(); + CapturedTimeout globalTimeout = harness.nextScheduledTimeout(); + registeredRequest.get().onThrottleFailure(failure); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(globalTimeout.isCancelled()).isTrue(); + } + } + @Test public void should_honor_statement_timeout() throws Exception { // given diff --git a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java index 852991b9607..61f5351b038 100644 --- a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java @@ -24,14 +24,15 @@ import static com.datastax.dse.driver.internal.core.graph.GraphTestUtils.defaultDseFrameOf; import static com.datastax.dse.driver.internal.core.graph.GraphTestUtils.serialize; import static com.datastax.dse.driver.internal.core.graph.GraphTestUtils.singleGraphRow; +import static com.datastax.oss.driver.Assertions.assertThatStage; import static com.datastax.oss.driver.api.core.type.codec.TypeCodecs.BIGINT; import static com.datastax.oss.driver.api.core.type.codec.TypeCodecs.TEXT; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.matches; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -58,6 +59,8 @@ import com.datastax.oss.driver.api.core.Version; 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.session.throttling.RequestThrottler; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; import com.datastax.oss.driver.api.core.tracker.RequestTracker; import com.datastax.oss.driver.api.core.uuid.Uuids; import com.datastax.oss.driver.internal.core.cql.Conversions; @@ -585,28 +588,37 @@ public void should_cancel_pre_acquired_id_if_graph_payload_conversion_fails() { ScriptGraphStatement graphStatement = Mockito.spy(ScriptGraphStatement.newInstance("mock query")); doThrow(failure).when(graphStatement).getCustomPayload(); + RequestThrottler throttler = mock(RequestThrottler.class); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(any()); GraphRequestHandlerTestHarness.Builder builder = GraphRequestHandlerTestHarness.builder(); PoolBehavior nodeBehavior = builder.customBehavior(node); try (GraphRequestHandlerTestHarness harness = builder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); GraphSupportChecker graphSupportChecker = mock(GraphSupportChecker.class); when(graphSupportChecker.inferGraphProtocol(any(), any(), any())) .thenReturn(GRAPH_BINARY_1_0); GraphBinaryModule module = createGraphBinaryModule(harness.getContext()); - assertThatThrownBy( - () -> - new GraphRequestHandler( - graphStatement, - harness.getSession(), - harness.getContext(), - "test", - module, - graphSupportChecker)) - .isSameAs(failure); + GraphRequestHandler handler = + new GraphRequestHandler( + graphStatement, + harness.getSession(), + harness.getContext(), + "test", + module, + graphSupportChecker); + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); nodeBehavior.verifyNoWrite(); nodeBehavior.verifyPreAcquireCancelled(); + verify(throttler).signalError(handler, failure); } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java index 52f5b4a80ef..b05d5b83c6e 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java @@ -26,12 +26,15 @@ import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.datastax.oss.driver.api.core.RequestThrottlingException; import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; import com.datastax.oss.driver.api.core.session.throttling.Throttled; @@ -39,6 +42,10 @@ import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater; import com.datastax.oss.protocol.internal.Frame; import com.datastax.oss.protocol.internal.request.Query; +import io.netty.channel.ChannelFuture; +import io.netty.channel.EventLoop; +import io.netty.util.concurrent.ImmediateEventExecutor; +import io.netty.util.concurrent.Promise; import java.time.Duration; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -58,6 +65,8 @@ public class ThrottledAdminRequestHandlerTest { public void setup() { MockitoAnnotations.initMocks(this); when(channel.preAcquireId()).thenAnswer(invocation -> availableIds.compareAndSet(1, 0)); + when(channel.closeFuture()).thenReturn(mock(ChannelFuture.class)); + when(channel.eventLoop()).thenReturn(mock(EventLoop.class)); doAnswer( invocation -> { if (!availableIds.compareAndSet(0, 1)) { @@ -82,7 +91,7 @@ public void should_release_permit_and_reservation_when_metric_update_throws() { anyLong(), eq(TimeUnit.NANOSECONDS)); - assertThatThrownBy(() -> handler.onThrottleReady(true)).isSameAs(failure); + handler.onThrottleReady(true); assertThat(availableIds.get()).isEqualTo(1); verify(throttler).signalError(handler, failure); @@ -103,7 +112,7 @@ public void should_release_permit_when_synchronous_write_throws() { .register(handler); doThrow(failure).when(channel).write(any(), anyBoolean(), anyMap(), eq(handler)); - assertThatThrownBy(handler::start).isSameAs(failure); + handler.start(); assertThat(availableIds.get()).isEqualTo(1); verify(throttler).signalError(handler, failure); @@ -123,7 +132,58 @@ public void should_complete_result_without_releasing_permit_when_registration_th assertThatStage(handler.result).isFailed(error -> assertThat(error).isSameAs(failure)); } + @Test + public void should_cancel_submitted_write_when_registration_throws_after_admission() { + RuntimeException failure = new RuntimeException("mock failure"); + ThrottledAdminRequestHandler handler = newHandler(Duration.ofSeconds(1)); + Promise writeFuture = ImmediateEventExecutor.INSTANCE.newPromise(); + when(channel.write(any(), anyBoolean(), anyMap(), eq(handler))).thenReturn(writeFuture); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + throw failure; + }) + .when(throttler) + .register(handler); + + assertThatThrownBy(handler::start).isSameAs(failure); + + assertThat(availableIds.get()).isZero(); + verify(channel, never()).cancelPreAcquireId(); + verify(throttler).signalError(handler, failure); + assertThatStage(handler.result).isFailed(error -> assertThat(error).isSameAs(failure)); + + writeFuture.setSuccess(null); + + verify(channel, atLeastOnce()).cancel(handler); + verify(channel.eventLoop(), never()).schedule(any(Runnable.class), anyLong(), any()); + } + + @Test + public void should_release_permit_for_throttling_exception_after_admission() { + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + ThrottledAdminRequestHandler handler = newHandler(); + Promise writeFuture = ImmediateEventExecutor.INSTANCE.newPromise(); + when(channel.write(any(), anyBoolean(), anyMap(), eq(handler))).thenReturn(writeFuture); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + throw failure; + }) + .when(throttler) + .register(handler); + + assertThatThrownBy(handler::start).isSameAs(failure); + + verify(throttler).signalError(handler, failure); + assertThatStage(handler.result).isFailed(error -> assertThat(error).isSameAs(failure)); + } + private ThrottledAdminRequestHandler newHandler() { + return newHandler(Duration.ZERO); + } + + private ThrottledAdminRequestHandler newHandler(Duration timeout) { assertThat(channel.preAcquireId()).isTrue(); assertThat(availableIds.get()).isZero(); return ThrottledAdminRequestHandler.query( @@ -131,7 +191,7 @@ private ThrottledAdminRequestHandler newHandler() { false, new Query("mock query"), Frame.NO_PAYLOAD, - Duration.ZERO, + timeout, throttler, metricUpdater, "test", diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java index 24280113697..a286f61cf02 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java @@ -20,12 +20,13 @@ import static com.datastax.oss.driver.Assertions.assertThat; import static com.datastax.oss.driver.Assertions.assertThatStage; import static com.datastax.oss.driver.internal.core.cql.CqlRequestHandlerTestBase.defaultFrameOf; -import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -41,6 +42,8 @@ import com.datastax.oss.driver.api.core.retry.RetryPolicy; import com.datastax.oss.driver.api.core.retry.RetryVerdict; import com.datastax.oss.driver.api.core.servererrors.OverloadedException; +import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; import com.datastax.oss.driver.internal.core.channel.ResponseCallback; import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap; @@ -364,18 +367,26 @@ public void should_cancel_pre_acquired_id_if_initial_prepare_payload_access_fail RuntimeException failure = new RuntimeException("mock failure"); DefaultPrepareRequest prepareRequest = spy(new DefaultPrepareRequest("mock query")); doThrow(failure).when(prepareRequest).getCustomPayload(); + RequestThrottler throttler = mock(RequestThrottler.class); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(any()); RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); try (RequestHandlerTestHarness harness = harnessBuilder.build()) { - assertThatThrownBy( - () -> - new CqlPrepareHandler( - prepareRequest, harness.getSession(), harness.getContext(), "test")) - .isSameAs(failure); + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + CqlPrepareHandler handler = + new CqlPrepareHandler(prepareRequest, harness.getSession(), harness.getContext(), "test"); + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); node1Behavior.verifyNoWrite(); node1Behavior.verifyPreAcquireCancelled(); + verify(throttler).signalError(handler, failure); } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerRetryTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerRetryTest.java index ccac873c616..31225a18aba 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerRetryTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerRetryTest.java @@ -23,8 +23,11 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.atMost; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -39,6 +42,7 @@ import com.datastax.oss.driver.api.core.cql.Statement; import com.datastax.oss.driver.api.core.metadata.Node; import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; +import com.datastax.oss.driver.api.core.retry.RetryDecision; import com.datastax.oss.driver.api.core.retry.RetryPolicy; import com.datastax.oss.driver.api.core.retry.RetryVerdict; import com.datastax.oss.driver.api.core.servererrors.BootstrappingException; @@ -50,6 +54,7 @@ import com.datastax.oss.driver.api.core.servererrors.WriteTimeoutException; import com.datastax.oss.driver.api.core.session.Request; import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; +import com.datastax.oss.driver.api.core.tracker.RequestTracker; import com.datastax.oss.protocol.internal.ProtocolConstants; import com.datastax.oss.protocol.internal.response.Error; import com.datastax.oss.protocol.internal.response.error.ReadTimeout; @@ -68,6 +73,115 @@ public class CqlRequestHandlerRetryTest extends CqlRequestHandlerTestBase { + @Test + public void should_contain_failure_while_building_retry_request() { + HeartbeatException requestFailure = mock(HeartbeatException.class); + RuntimeException setupFailure = new RuntimeException("mock failure"); + RetryVerdict verdict = mock(RetryVerdict.class); + when(verdict.getRetryDecision()).thenReturn(RetryDecision.RETRY_NEXT); + when(verdict.getRetryRequest(any(Statement.class))).thenThrow(setupFailure); + + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + harnessBuilder.withResponseFailure(node1, requestFailure); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + RequestTracker requestTracker = mock(RequestTracker.class); + when(harness.getContext().getRequestTracker()).thenReturn(requestTracker); + when(harness + .getContext() + .getRetryPolicy(anyString()) + .onRequestAbortedVerdict(any(), eq(requestFailure), eq(0))) + .thenReturn(verdict); + + CompletionStage result = + new CqlRequestHandler( + IDEMPOTENT_STATEMENT, harness.getSession(), harness.getContext(), "test") + .handle(); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(setupFailure)); + verify(verdict).getRetryDecision(); + verify(requestTracker, times(1)) + .onNodeError( + eq(IDEMPOTENT_STATEMENT), + eq(requestFailure), + anyLong(), + any(DriverExecutionProfile.class), + eq(node1), + anyString()); + verify(requestTracker, never()) + .onNodeError( + any(), + eq(setupFailure), + anyLong(), + any(DriverExecutionProfile.class), + any(), + anyString()); + verify(requestTracker) + .onError( + eq(IDEMPOTENT_STATEMENT), + eq(setupFailure), + anyLong(), + any(DriverExecutionProfile.class), + isNull(), + anyString()); + verifyNoMoreInteractions(requestTracker); + } + } + + @Test + public void should_contain_failure_while_reading_retry_decision() { + HeartbeatException requestFailure = mock(HeartbeatException.class); + RuntimeException setupFailure = new RuntimeException("mock failure"); + RetryVerdict verdict = mock(RetryVerdict.class); + when(verdict.getRetryDecision()).thenThrow(setupFailure); + + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + harnessBuilder.withResponseFailure(node1, requestFailure); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + RequestTracker requestTracker = mock(RequestTracker.class); + when(harness.getContext().getRequestTracker()).thenReturn(requestTracker); + when(harness + .getContext() + .getRetryPolicy(anyString()) + .onRequestAbortedVerdict(any(), eq(requestFailure), eq(0))) + .thenReturn(verdict); + + CompletionStage result = + new CqlRequestHandler( + IDEMPOTENT_STATEMENT, harness.getSession(), harness.getContext(), "test") + .handle(); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(setupFailure)); + verify(verdict).getRetryDecision(); + verify(requestTracker, times(1)) + .onNodeError( + eq(IDEMPOTENT_STATEMENT), + eq(requestFailure), + anyLong(), + any(DriverExecutionProfile.class), + eq(node1), + anyString()); + verify(requestTracker, never()) + .onNodeError( + any(), + eq(setupFailure), + anyLong(), + any(DriverExecutionProfile.class), + any(), + anyString()); + verify(requestTracker) + .onError( + eq(IDEMPOTENT_STATEMENT), + eq(setupFailure), + anyLong(), + any(DriverExecutionProfile.class), + isNull(), + anyString()); + verifyNoMoreInteractions(requestTracker); + } + } + @Test @UseDataProvider("allIdempotenceConfigs") public void should_always_try_next_node_if_bootstrapping( diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerSpeculativeExecutionTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerSpeculativeExecutionTest.java index a09a9eb3d5a..fb31c357316 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerSpeculativeExecutionTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerSpeculativeExecutionTest.java @@ -22,6 +22,8 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -35,6 +37,7 @@ import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; import com.datastax.oss.driver.api.core.servererrors.BootstrappingException; import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy; +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; import com.datastax.oss.protocol.internal.ProtocolConstants; import com.datastax.oss.protocol.internal.response.Error; @@ -47,6 +50,48 @@ public class CqlRequestHandlerSpeculativeExecutionTest extends CqlRequestHandlerTestBase { + @Test + public void should_keep_initial_execution_running_if_speculative_setup_fails() throws Exception { + RequestIdGenerator requestIdGenerator = mock(RequestIdGenerator.class); + RuntimeException setupFailure = new RuntimeException("mock failure"); + when(requestIdGenerator.getSessionRequestId()).thenReturn("session"); + when(requestIdGenerator.getNodeRequestId(any(), eq("session"))).thenReturn("node1", "node2"); + doReturn(IDEMPOTENT_STATEMENT) + .doThrow(setupFailure) + .when(requestIdGenerator) + .getDecoratedStatement(any(), any()); + + RequestHandlerTestHarness.Builder harnessBuilder = + RequestHandlerTestHarness.builder().withRequestIdGenerator(requestIdGenerator); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + PoolBehavior node2Behavior = harnessBuilder.customBehavior(node2); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + SpeculativeExecutionPolicy speculativeExecutionPolicy = + harness.getContext().getSpeculativeExecutionPolicy(DriverExecutionProfile.DEFAULT_NAME); + when(speculativeExecutionPolicy.nextExecution( + any(Node.class), eq(null), eq(IDEMPOTENT_STATEMENT), eq(1))) + .thenReturn(100L); + + CompletionStage result = + new CqlRequestHandler( + IDEMPOTENT_STATEMENT, harness.getSession(), harness.getContext(), "test") + .handle(); + node1Behavior.setWriteSuccess(); + + harness.nextScheduledTimeout(); // Discard the request timeout. + CapturedTimeout speculativeExecution = harness.nextScheduledTimeout(); + speculativeExecution.task().run(speculativeExecution); + + node2Behavior.verifyNoWrite(); + node2Behavior.verifyPreAcquireCancelled(); + assertThatStage(result).isNotDone(); + + node1Behavior.setResponseSuccess(defaultFrameOf(singleRow())); + assertThatStage(result).isSuccess(); + } + } + @Test @UseDataProvider("nonIdempotentConfig") public void should_not_schedule_speculative_executions_if_not_idempotent( diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java index f9068d137f2..5754a13895b 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java @@ -21,8 +21,13 @@ import static com.datastax.oss.driver.Assertions.assertThatStage; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -31,6 +36,7 @@ import com.datastax.oss.driver.api.core.DriverTimeoutException; import com.datastax.oss.driver.api.core.NoNodeAvailableException; import com.datastax.oss.driver.api.core.NodeUnavailableException; +import com.datastax.oss.driver.api.core.RequestThrottlingException; import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import com.datastax.oss.driver.api.core.cql.AsyncResultSet; import com.datastax.oss.driver.api.core.cql.BoundStatement; @@ -40,6 +46,8 @@ import com.datastax.oss.driver.api.core.cql.Row; import com.datastax.oss.driver.api.core.cql.Statement; import com.datastax.oss.driver.api.core.metadata.Node; +import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.internal.core.session.RepreparePayload; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; @@ -48,16 +56,20 @@ import com.datastax.oss.protocol.internal.response.result.Prepared; import com.datastax.oss.protocol.internal.response.result.SetKeyspace; import com.datastax.oss.protocol.internal.util.Bytes; +import io.netty.util.Timer; import java.nio.ByteBuffer; import java.time.Duration; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletionStage; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Test; public class CqlRequestHandlerTest extends CqlRequestHandlerTestBase { @@ -99,6 +111,37 @@ public void should_complete_result_if_first_node_replies_immediately() { } } + @Test + public void should_try_next_node_if_channel_selection_fails() { + RuntimeException failure = new RuntimeException("mock failure"); + try (RequestHandlerTestHarness harness = + RequestHandlerTestHarness.builder() + .withEmptyPool(node1) + .withResponse(node2, defaultFrameOf(singleRow())) + .build()) { + when(harness.getSession().getChannel(eq(node1), anyString(), any(), any())) + .thenThrow(failure); + + CompletionStage result = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, + harness.getSession(), + harness.getContext(), + "test") + .handle(); + + assertThatStage(result) + .isSuccess( + resultSet -> { + assertThat(resultSet.getExecutionInfo().getCoordinator()).isEqualTo(node2); + List> errors = resultSet.getExecutionInfo().getErrors(); + assertThat(errors).hasSize(1); + assertThat(errors.get(0).getKey()).isEqualTo(node1); + assertThat(errors.get(0).getValue()).isSameAs(failure); + }); + } + } + @Test public void should_fail_if_no_node_available() { try (RequestHandlerTestHarness harness = @@ -157,19 +200,182 @@ public void should_fail_if_nodes_unavailable() { } @Test - public void should_cancel_pre_acquired_id_if_request_decoration_fails_before_write() { + public void should_complete_result_and_cleanup_if_immediate_request_setup_fails() { RequestIdGenerator requestIdGenerator = mock(RequestIdGenerator.class); + RequestThrottler throttler = mock(RequestThrottler.class); RuntimeException failure = new RuntimeException("mock failure"); when(requestIdGenerator.getSessionRequestId()).thenReturn("session"); when(requestIdGenerator.getNodeRequestId(any(), eq("session"))).thenReturn("node"); when(requestIdGenerator.getDecoratedStatement(any(), eq("node"))).thenThrow(failure); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(any()); RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder().withRequestIdGenerator(requestIdGenerator); PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); try (RequestHandlerTestHarness harness = harnessBuilder.build()) { - // This only verifies stream-id cleanup; the other failure-path leaks are tracked in #980. + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); + node1Behavior.verifyNoWrite(); + node1Behavior.verifyPreAcquireCancelled(); + assertThat(harness.nextScheduledTimeout().isCancelled()).isTrue(); + verify(throttler).signalError(handler, failure); + } + } + + @Test + public void should_not_release_throttler_if_request_was_not_admitted() { + RequestThrottler throttler = mock(RequestThrottler.class); + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleFailure(failure); + return null; + }) + .when(throttler) + .register(any()); + + try (RequestHandlerTestHarness harness = RequestHandlerTestHarness.builder().build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(harness.nextScheduledTimeout().isCancelled()).isTrue(); + verify(throttler, never()).signalSuccess(any()); + verify(throttler, never()).signalError(any(), any()); + verify(throttler, never()).signalTimeout(any()); + verify(throttler, never()).signalCancel(any()); + } + } + + @Test + public void should_schedule_timeout_before_throttler_registration() { + RequestThrottler throttler = mock(RequestThrottler.class); + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + AtomicReference timeoutSeenDuringRegistration = new AtomicReference<>(); + + try (RequestHandlerTestHarness harness = RequestHandlerTestHarness.builder().build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + doAnswer( + invocation -> { + timeoutSeenDuringRegistration.set(harness.nextScheduledTimeout()); + invocation.getArgument(0, Throttled.class).onThrottleFailure(failure); + return null; + }) + .when(throttler) + .register(any()); + + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + assertThat(timeoutSeenDuringRegistration.get()).isNotNull(); + assertThat(timeoutSeenDuringRegistration.get().isCancelled()).isTrue(); + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); + } + } + + @Test + public void should_defer_queued_timeout_release_until_registration_returns() { + RequestThrottler throttler = mock(RequestThrottler.class); + AtomicBoolean timeoutSignaled = new AtomicBoolean(); + doAnswer( + invocation -> { + timeoutSignaled.set(true); + return null; + }) + .when(throttler) + .signalTimeout(any()); + + try (RequestHandlerTestHarness harness = RequestHandlerTestHarness.builder().build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + doAnswer( + invocation -> { + CapturedTimeout timeout = harness.nextScheduledTimeout(); + timeout.task().run(timeout); + assertThat(timeoutSignaled).isFalse(); + return null; + }) + .when(throttler) + .register(any()); + + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + assertThatStage(handler.handle()) + .isFailed(error -> assertThat(error).isInstanceOf(DriverTimeoutException.class)); + assertThat(timeoutSignaled).isTrue(); + verify(throttler).signalTimeout(handler); + } + } + + @Test + public void should_release_permit_when_registration_throws_after_admission() { + RequestThrottler throttler = mock(RequestThrottler.class); + RuntimeException failure = new RuntimeException("mock failure"); + AtomicReference admittedHandler = new AtomicReference<>(); + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + doAnswer( + invocation -> { + CqlRequestHandler handler = invocation.getArgument(0); + admittedHandler.set(handler); + handler.onThrottleReady(false); + node1Behavior.setWriteSuccess(); + throw failure; + }) + .when(throttler) + .register(any()); + + assertThatThrownBy( + () -> + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, + harness.getSession(), + harness.getContext(), + "test")) + .isSameAs(failure); + + CqlRequestHandler handler = admittedHandler.get(); + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(harness.nextScheduledTimeout().isCancelled()).isTrue(); + node1Behavior.verifyWrite(); + node1Behavior.verifyCancellation(); + verify(throttler).signalError(handler, failure); + } + } + + @Test + public void + should_cancel_timeout_without_releasing_permit_when_registration_fails_before_admission() { + RequestThrottler throttler = mock(RequestThrottler.class); + RuntimeException failure = new RuntimeException("mock failure"); + doAnswer( + invocation -> { + throw failure; + }) + .when(throttler) + .register(any()); + + try (RequestHandlerTestHarness harness = RequestHandlerTestHarness.builder().build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + assertThatThrownBy( () -> new CqlRequestHandler( @@ -179,8 +385,148 @@ public void should_cancel_pre_acquired_id_if_request_decoration_fails_before_wri "test")) .isSameAs(failure); + assertThat(harness.nextScheduledTimeout().isCancelled()).isTrue(); + verify(throttler, never()).signalSuccess(any()); + verify(throttler, never()).signalError(any(), any()); + verify(throttler, never()).signalTimeout(any()); + verify(throttler, never()).signalCancel(any()); + } + } + + @Test + public void should_complete_result_and_cleanup_if_delayed_request_setup_fails() { + RequestIdGenerator requestIdGenerator = mock(RequestIdGenerator.class); + RequestThrottler throttler = mock(RequestThrottler.class); + AtomicReference registeredRequest = new AtomicReference<>(); + RuntimeException failure = new RuntimeException("mock failure"); + when(requestIdGenerator.getSessionRequestId()).thenReturn("session"); + when(requestIdGenerator.getNodeRequestId(any(), eq("session"))).thenReturn("node"); + when(requestIdGenerator.getDecoratedStatement(any(), eq("node"))).thenThrow(failure); + doAnswer( + invocation -> { + registeredRequest.set(invocation.getArgument(0)); + return null; + }) + .when(throttler) + .register(any()); + + RequestHandlerTestHarness.Builder harnessBuilder = + RequestHandlerTestHarness.builder().withRequestIdGenerator(requestIdGenerator); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + registeredRequest.get().onThrottleReady(true); + + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); node1Behavior.verifyNoWrite(); node1Behavior.verifyPreAcquireCancelled(); + assertThat(harness.nextScheduledTimeout().isCancelled()).isTrue(); + verify(throttler).signalError(handler, failure); + } + } + + @Test + public void should_not_register_request_if_timeout_scheduling_fails() { + RequestThrottler throttler = mock(RequestThrottler.class); + Timer timer = mock(Timer.class); + IllegalStateException failure = new IllegalStateException("mock failure"); + when(timer.newTimeout(any(), anyLong(), any())).thenThrow(failure); + + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + harnessBuilder.customBehavior(node1); + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + when(harness.getContext().getNettyOptions().getTimer()).thenReturn(timer); + + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); + verify(throttler, never()).register(any()); + verify(throttler, never()).signalSuccess(any()); + verify(throttler, never()).signalError(any(), any()); + verify(throttler, never()).signalTimeout(any()); + verify(throttler, never()).signalCancel(any()); + } + } + + @Test + public void should_cleanup_if_request_setup_fails_after_write_retry() { + RequestIdGenerator requestIdGenerator = mock(RequestIdGenerator.class); + RequestThrottler throttler = mock(RequestThrottler.class); + RuntimeException writeFailure = new RuntimeException("mock write failure"); + RuntimeException setupFailure = new RuntimeException("mock setup failure"); + when(requestIdGenerator.getSessionRequestId()).thenReturn("session"); + when(requestIdGenerator.getNodeRequestId(any(), eq("session"))).thenReturn("node1", "node2"); + doReturn(UNDEFINED_IDEMPOTENCE_STATEMENT) + .doThrow(setupFailure) + .when(requestIdGenerator) + .getDecoratedStatement(any(), any()); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(any()); + + RequestHandlerTestHarness.Builder harnessBuilder = + RequestHandlerTestHarness.builder().withRequestIdGenerator(requestIdGenerator); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + PoolBehavior node2Behavior = harnessBuilder.customBehavior(node2); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + CompletionStage result = handler.handle(); + + node1Behavior.setWriteFailure(writeFailure); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(setupFailure)); + node2Behavior.verifyNoWrite(); + node2Behavior.verifyPreAcquireCancelled(); + assertThat(harness.nextScheduledTimeout().isCancelled()).isTrue(); + verify(throttler).signalError(handler, setupFailure); + } + } + + @Test + public void should_release_cancelled_request_only_once() { + RequestIdGenerator requestIdGenerator = mock(RequestIdGenerator.class); + RequestThrottler throttler = mock(RequestThrottler.class); + CancellationException failure = new CancellationException("mock cancellation"); + when(requestIdGenerator.getSessionRequestId()).thenReturn("session"); + when(requestIdGenerator.getNodeRequestId(any(), eq("session"))).thenReturn("node"); + when(requestIdGenerator.getDecoratedStatement(any(), eq("node"))).thenThrow(failure); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(any()); + + RequestHandlerTestHarness.Builder harnessBuilder = + RequestHandlerTestHarness.builder().withRequestIdGenerator(requestIdGenerator); + harnessBuilder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + CqlRequestHandler handler = + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, harness.getSession(), harness.getContext(), "test"); + + assertThat(handler.handle().toCompletableFuture()).isCancelled(); + verify(throttler).signalCancel(handler); + verify(throttler, never()).signalError(eq(handler), any()); } } @@ -285,4 +631,54 @@ public void should_reprepare_on_the_fly_if_not_prepared() throws InterruptedExce assertThatStage(resultSetFuture).isSuccess(); } } + + @Test + public void should_release_outer_request_if_reprepare_is_throttled() { + ByteBuffer mockId = Bytes.fromHexString("0xffff"); + PreparedStatement preparedStatement = mock(PreparedStatement.class); + when(preparedStatement.getId()).thenReturn(mockId); + ColumnDefinitions columnDefinitions = mock(ColumnDefinitions.class); + when(columnDefinitions.size()).thenReturn(0); + when(preparedStatement.getResultSetDefinitions()).thenReturn(columnDefinitions); + BoundStatement boundStatement = mock(BoundStatement.class); + when(boundStatement.getPreparedStatement()).thenReturn(preparedStatement); + when(boundStatement.getValues()).thenReturn(Collections.emptyList()); + when(boundStatement.getNowInSeconds()).thenReturn(Statement.NO_NOW_IN_SECONDS); + + RequestThrottler throttler = mock(RequestThrottler.class); + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleFailure(failure); + return null; + }) + .when(throttler) + .register(any()); + + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + ConcurrentMap repreparePayloads = new ConcurrentHashMap<>(); + repreparePayloads.put( + mockId, new RepreparePayload(mockId, "mock query", null, Collections.emptyMap())); + when(harness.getSession().getRepreparePayloads()).thenReturn(repreparePayloads); + + CqlRequestHandler handler = + new CqlRequestHandler(boundStatement, harness.getSession(), harness.getContext(), "test"); + node1Behavior.setWriteSuccess(); + node1Behavior.setResponseSuccess( + defaultFrameOf(new Unprepared("mock message", Bytes.getArray(mockId)))); + + assertThatStage(handler.handle()).isFailed(error -> assertThat(error).isSameAs(failure)); + verify(throttler).signalError(handler, failure); + verify(throttler).signalError(any(), eq(failure)); + } + } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottlerTest.java index 7eb682070cd..f9b4460fb4c 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottlerTest.java @@ -30,6 +30,7 @@ import com.datastax.oss.driver.shaded.guava.common.collect.Lists; import java.util.List; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import org.junit.Before; import org.junit.Test; @@ -154,6 +155,145 @@ public void should_dequeue_when_active_times_out() { should_dequeue_when_active_completes(throttler::signalTimeout); } + @Test + public void should_dequeue_synchronous_failures_without_recursion() { + MockThrottled active = new MockThrottled(); + throttler.register(active); + for (int i = 0; i < 4; i++) { + throttler.register(new MockThrottled()); + } + + AtomicInteger depth = new AtomicInteger(); + AtomicInteger maximumDepth = new AtomicInteger(); + AtomicInteger completed = new AtomicInteger(); + for (int i = 0; i < 10; i++) { + throttler.register( + new Throttled() { + @Override + public void onThrottleReady(boolean wasDelayed) { + int currentDepth = depth.incrementAndGet(); + maximumDepth.accumulateAndGet(currentDepth, Math::max); + try { + completed.incrementAndGet(); + throttler.signalError(this, new RuntimeException("mock failure")); + } finally { + depth.decrementAndGet(); + } + } + + @Override + public void onThrottleFailure(RequestThrottlingException error) {} + }); + } + + throttler.signalSuccess(active); + + assertThat(completed).hasValue(10); + assertThat(maximumDepth).hasValue(1); + assertThat(throttler.getQueue()).isEmpty(); + assertThat(throttler.getConcurrentRequests()).isEqualTo(4); + } + + @Test + public void should_find_reentrant_callback_below_another_throttler() { + ConcurrencyLimitingRequestThrottler other = new ConcurrencyLimitingRequestThrottler(context); + AtomicInteger depth = new AtomicInteger(); + AtomicInteger maximumDepth = new AtomicInteger(); + AtomicInteger completed = new AtomicInteger(); + + Throttled nestedOnFirst = + new Throttled() { + @Override + public void onThrottleReady(boolean wasDelayed) { + int currentDepth = depth.incrementAndGet(); + maximumDepth.accumulateAndGet(currentDepth, Math::max); + completed.incrementAndGet(); + depth.decrementAndGet(); + } + + @Override + public void onThrottleFailure(RequestThrottlingException error) {} + }; + Throttled onOther = + new Throttled() { + @Override + public void onThrottleReady(boolean wasDelayed) { + int currentDepth = depth.incrementAndGet(); + maximumDepth.accumulateAndGet(currentDepth, Math::max); + try { + completed.incrementAndGet(); + throttler.register(nestedOnFirst); + } finally { + depth.decrementAndGet(); + } + } + + @Override + public void onThrottleFailure(RequestThrottlingException error) {} + }; + Throttled onFirst = + new Throttled() { + @Override + public void onThrottleReady(boolean wasDelayed) { + int currentDepth = depth.incrementAndGet(); + maximumDepth.accumulateAndGet(currentDepth, Math::max); + try { + completed.incrementAndGet(); + other.register(onOther); + } finally { + depth.decrementAndGet(); + } + } + + @Override + public void onThrottleFailure(RequestThrottlingException error) {} + }; + + throttler.register(onFirst); + + assertThat(completed).hasValue(3); + assertThat(maximumDepth).hasValue(2); + } + + @Test + public void should_keep_draining_ready_callbacks_after_one_throws() { + MockThrottled active = new MockThrottled(); + throttler.register(active); + for (int i = 0; i < 4; i++) { + throttler.register(new MockThrottled()); + } + + RuntimeException failure = new RuntimeException("mock failure"); + AtomicInteger completed = new AtomicInteger(); + throttler.register( + new Throttled() { + @Override + public void onThrottleReady(boolean wasDelayed) { + throttler.signalError(this, failure); + throw failure; + } + + @Override + public void onThrottleFailure(RequestThrottlingException error) {} + }); + throttler.register( + new Throttled() { + @Override + public void onThrottleReady(boolean wasDelayed) { + completed.incrementAndGet(); + } + + @Override + public void onThrottleFailure(RequestThrottlingException error) {} + }); + + throttler.signalSuccess(active); + + assertThat(completed).hasValue(1); + assertThat(throttler.getQueue()).isEmpty(); + assertThat(throttler.getConcurrentRequests()).isEqualTo(5); + } + private void should_dequeue_when_active_completes(Consumer completeCallback) { // Given MockThrottled first = new MockThrottled();