From 5829fa80f9ae6c68e3bef3716875f4640d249e52 Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Wed, 19 Aug 2026 19:49:28 -0400 Subject: [PATCH] fix: contain throttled request setup failures Route synchronous CQL, admin, and continuous-request setup failures through terminal completion, including retry and reprepare callbacks. Cancel reserved stream IDs and scheduled timeouts, and distinguish rejected, admitted, and cancelled requests so throttler permits are returned exactly once. Drain concurrency-throttler readiness callbacks iteratively to prevent recursive failure chains and continue admitting queued work after a callback fails. --- .../ContinuousRequestHandlerBase.java | 40 ++- .../ThrottledAdminRequestHandler.java | 16 +- .../internal/core/cql/CqlRequestHandler.java | 230 +++++++++++----- .../ConcurrencyLimitingRequestThrottler.java | 85 +++++- .../ContinuousCqlRequestHandlerTest.java | 28 +- .../ContinuousGraphRequestHandlerTest.java | 116 ++++++++ .../ThrottledAdminRequestHandlerTest.java | 49 +++- .../core/cql/CqlRequestHandlerRetryTest.java | 28 ++ ...equestHandlerSpeculativeExecutionTest.java | 45 ++++ .../core/cql/CqlRequestHandlerTest.java | 250 +++++++++++++++++- ...ncurrencyLimitingRequestThrottlerTest.java | 140 ++++++++++ 11 files changed, 927 insertions(+), 100 deletions(-) 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..feec5bd6102 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 @@ -283,6 +285,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 +295,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 +383,7 @@ private void sendRequest( } } else if (!chosenCallback.isDone()) { boolean writeSubmitted = false; - Throwable terminalPreWriteFailure = null; + Throwable terminalSetupFailure = null; NodeResponseCallback nodeResponseCallback = null; try { nodeResponseCallback = @@ -392,12 +405,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 +421,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/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java index d14f4ee08c4..94198be547f 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,12 @@ 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); + } else { + // Registration failed before admission, so there is no throttler permit to release. + super.setFinalError(t); + } throw t; } return result; @@ -150,6 +154,7 @@ public CompletionStage start() { @Override public void onThrottleReady(boolean wasDelayed) { + admitted = true; try { if (wasDelayed) { metricUpdater.updateTimer( @@ -163,7 +168,6 @@ public void onThrottleReady(boolean wasDelayed) { } catch (Throwable t) { cancelExternalReservation(); setFinalError(t); - throw t; } } @@ -175,8 +179,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 +199,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/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index ce4b40d6d29..69af9c32697 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,9 @@ public class CqlRequestHandler implements Throttled { private final Optional requestIdGenerator; private final SessionMetricUpdater sessionMetricUpdater; private final DriverExecutionProfile executionProfile; + private final AtomicBoolean admitted = 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 +185,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); @@ -204,35 +212,52 @@ protected CqlRequestHandler( this.timer = context.getNettyOptions().getTimer(); 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()) { + this.scheduledTimeout = scheduleTimeout(timeout); + // Registration can start the request synchronously. Do not leave a timeout behind if the + // request became terminal while the timeout was being installed. + if (result.isDone() && this.scheduledTimeout != null) { + this.scheduledTimeout.cancel(); + } + } } @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.set(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 +280,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( + // If we raced with session shutdown the timer might be closed already; surface a more + // explicit failure through the request's normal terminal path. + setFinalError( + initialStatement, "cannot be started once stopped".equals(e.getMessage()) ? new IllegalStateException("Session is closed") - : e); + : e, + null, + -1); } } return null; @@ -380,27 +408,33 @@ private void sendRequest( } Node node = retriedNode; DriverChannel channel = null; - if (node == null - || (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)); + try { + Token routingToken = getRoutingToken(statement); + if (node == null + || (channel = + session.getChannel( + node, + handlerLogPrefix, + routingToken, + getShardFromTabletMap(statement, node, routingToken))) + == null) { + while (!result.isDone() && (node = queryPlan.poll()) != null) { + channel = + session.getChannel( + node, + handlerLogPrefix, + routingToken, + getShardFromTabletMap(statement, node, routingToken)); + if (channel != null) { + break; + } else { + recordError(node, new NodeUnavailableException(node)); + } } } + } catch (Throwable t) { + handleRequestSetupFailure(statement, t, node, currentExecutionIndex); + return; } if (channel == null) { // We've reached the end of the query plan without finding any node to write to @@ -410,6 +444,7 @@ private void sendRequest( } } else { boolean writeSubmitted = false; + Throwable setupFailure = null; try { Statement finalStatement = statement; String nodeRequestId = @@ -437,11 +472,57 @@ 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, node, currentExecutionIndex); + } + } + } + + private void handleRequestSetupFailure( + Statement statement, Throwable error, @Nullable Node node, int execution) { + if (result.isDone()) { + Loggers.warnWithException( + LOG, + "[{}] Request setup failed after the request had completed", + handlerLogPrefix, + error); + return; + } + if (node != null) { + recordError(node, error); + if (!(requestTracker instanceof NoopRequestTracker)) { + requestTracker.onNodeError( + statement, + error, + System.nanoTime() - startTimeNanos, + executionProfile, + node, + handlerLogPrefix); + } + } + if (activeExecutionsCount.decrementAndGet() == 0) { + setFinalError(statement, error, node, execution); + } else { + Loggers.warnWithException( + LOG, + "[{}] Request setup failed, another execution is still active", + handlerLogPrefix, + error); } } @@ -485,7 +566,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 +689,13 @@ 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) { + terminalError.compareAndSet(null, error); if (error instanceof DriverException) { ((DriverException) error) .setExecutionInfo( @@ -633,16 +719,30 @@ private void setFinalError(Statement statement, Throwable error, Node node, i requestTracker.onError( statement, error, latencyNanos, executionProfile, node, handlerLogPrefix); } + 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 (!throttlerReleased.compareAndSet(false, true)) { + return; + } + if (error instanceof DriverTimeoutException) { + throttler.signalTimeout(this); + } else if (error instanceof CancellationException || !admitted.get()) { + // 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. * @@ -1013,28 +1113,26 @@ private void processErrorResponse(Error errorMessage) { private void processRetryVerdict(RetryVerdict verdict, Throwable error) { LOG.trace("[{}] Processing retry decision {}", logPrefix, verdict); + Statement retryStatement = null; + if (verdict.getRetryDecision() == RetryDecision.RETRY_SAME + || verdict.getRetryDecision() == RetryDecision.RETRY_NEXT) { + try { + retryStatement = verdict.getRetryRequest(statement); + } catch (Throwable t) { + handleRequestSetupFailure(statement, t, node, execution); + return; + } + } switch (verdict.getRetryDecision()) { 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); 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..079eeb3f3fe 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 @@ -63,6 +63,7 @@ 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 +232,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 +243,31 @@ public void should_unwind_execution_if_request_setup_fails_before_write() { } } + @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/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java index 52f5b4a80ef..f1333f0ee26 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 @@ -28,10 +28,12 @@ import static org.mockito.ArgumentMatchers.isNull; 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 +41,7 @@ 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.util.concurrent.Future; import java.time.Duration; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -82,7 +85,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 +106,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,6 +126,48 @@ public void should_complete_result_without_releasing_permit_when_registration_th assertThatStage(handler.result).isFailed(error -> assertThat(error).isSameAs(failure)); } + @Test + public void should_release_permit_when_registration_throws_after_admission() { + RuntimeException failure = new RuntimeException("mock failure"); + ThrottledAdminRequestHandler handler = newHandler(); + Future writeFuture = mock(Future.class); + 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)); + } + + @Test + public void should_release_permit_for_throttling_exception_after_admission() { + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + ThrottledAdminRequestHandler handler = newHandler(); + Future writeFuture = mock(Future.class); + 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() { assertThat(channel.preAcquireId()).isTrue(); assertThat(availableIds.get()).isZero(); 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..d8cc8f2a69c 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 @@ -39,6 +39,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; @@ -68,6 +69,33 @@ 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()) { + 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)); + } + } + @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..ad93cf6f200 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 @@ -19,10 +19,13 @@ 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.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 +34,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 +44,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 +54,19 @@ 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.AtomicReference; import org.junit.Test; public class CqlRequestHandlerTest extends CqlRequestHandlerTestBase { @@ -157,30 +166,200 @@ 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); + RequestThrottlingException failure = new RequestThrottlingException("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()) { + 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()).isNull(); + 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()).isNull(); + verify(throttler, never()).signalError(any(), any()); + } + } + + @Test + public void should_not_propagate_delayed_request_setup_failure() { + 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()) { - // This only verifies stream-id cleanup; the other failure-path leaks are tracked in #980. - assertThatThrownBy( - () -> - new CqlRequestHandler( - UNDEFINED_IDEMPOTENCE_STATEMENT, - harness.getSession(), - harness.getContext(), - "test")) - .isSameAs(failure); + 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_release_admitted_request_if_timeout_scheduling_fails() { + RequestThrottler throttler = mock(RequestThrottler.class); + Timer timer = mock(Timer.class); + IllegalStateException failure = new IllegalStateException("mock failure"); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(any()); + 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).signalError(handler, failure); + } + } + + @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 +464,53 @@ 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); + } + } } 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();