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();