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 9a7be344721..0453022cb6a 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 @@ -410,6 +410,7 @@ public void cancel() { cancelScheduledTasks(null); cancelGlobalTimeout(); + throttler.signalCancel(this); } private void cancelGlobalTimeout() { diff --git a/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java b/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java index 702da69b855..5c9ceb00df2 100644 --- a/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java +++ b/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java @@ -153,6 +153,7 @@ public class GraphRequestHandler implements Throttled { try { if (t instanceof CancellationException) { cancelScheduledTasks(); + context.getRequestThrottler().signalCancel(this); } } catch (Throwable t2) { Loggers.warnWithException(LOG, "[{}] Uncaught exception", logPrefix, t2); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/throttling/RequestThrottler.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/throttling/RequestThrottler.java index cb55fac336c..7e2b41ebbdb 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/throttling/RequestThrottler.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/session/throttling/RequestThrottler.java @@ -56,4 +56,12 @@ public interface RequestThrottler extends Closeable { * perform time-based eviction on pending requests. */ void signalTimeout(@NonNull Throttled request); + + /** + * Signals that a request has been cancelled. This indicates to the throttler that another request + * might be started. + */ + default void signalCancel(@NonNull Throttled request) { + // no-op for backward compatibility purposes + } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java index 8fe1adb20b1..1ee1f303ab2 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java @@ -124,6 +124,7 @@ protected CqlPrepareHandler( try { if (t instanceof CancellationException) { cancelTimeout(); + context.getRequestThrottler().signalCancel(this); } } catch (Throwable t2) { Loggers.warnWithException(LOG, "[{}] Uncaught exception", logPrefix, t2); 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 a1c6b0e5466..0808bdce63f 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 @@ -152,6 +152,7 @@ protected CqlRequestHandler( try { if (t instanceof CancellationException) { cancelScheduledTasks(); + context.getRequestThrottler().signalCancel(this); } } catch (Throwable t2) { Loggers.warnWithException(LOG, "[{}] Uncaught exception", logPrefix, t2); 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 e8f27467c6f..438bed0953b 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 @@ -145,6 +145,22 @@ public void signalTimeout(@NonNull Throttled request) { } } + @Override + public void signalCancel(@NonNull Throttled request) { + lock.lock(); + try { + if (!closed) { + if (queue.remove(request)) { // The request has been cancelled before it was active + LOG.trace("[{}] Removing cancelled request from the queue", logPrefix); + } else { + onRequestDone(); + } + } + } finally { + lock.unlock(); + } + } + @SuppressWarnings("GuardedBy") // this method is only called with the lock held private void onRequestDone() { assert lock.isHeldByCurrentThread(); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/PassThroughRequestThrottler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/PassThroughRequestThrottler.java index 714c712a4e8..2210e4b26f1 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/PassThroughRequestThrottler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/PassThroughRequestThrottler.java @@ -69,6 +69,11 @@ public void signalTimeout(@NonNull Throttled request) { // nothing to do } + @Override + public void signalCancel(@NonNull Throttled request) { + // nothing to do + } + @Override public void close() throws IOException { // nothing to do diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottler.java index 6536804ffee..03a693dc0fe 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottler.java @@ -198,6 +198,18 @@ public void signalTimeout(@NonNull Throttled request) { } } + @Override + public void signalCancel(@NonNull Throttled request) { + lock.lock(); + try { + if (!closed && queue.remove(request)) { // The request has been cancelled before it was active + LOG.trace("[{}] Removing cancelled request from the queue", logPrefix); + } + } finally { + lock.unlock(); + } + } + @Override public void close() { lock.lock(); 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 b587ac3daa2..c01b26c1e9f 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 @@ -88,6 +88,11 @@ public void should_allow_new_request_when_active_one_times_out() { should_allow_new_request_when_active_one_completes(throttler::signalTimeout); } + @Test + public void should_allow_new_request_when_active_one_canceled() { + should_allow_new_request_when_active_one_completes(throttler::signalCancel); + } + private void should_allow_new_request_when_active_one_completes( Consumer completeCallback) { // Given diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottlerTest.java index 7336fb447b6..0e0fe7c1c65 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/RateLimitingRequestThrottlerTest.java @@ -25,6 +25,7 @@ import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import com.datastax.oss.driver.api.core.config.DriverConfig; import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; import com.datastax.oss.driver.internal.core.context.InternalDriverContext; import com.datastax.oss.driver.internal.core.context.NettyOptions; import com.datastax.oss.driver.internal.core.util.concurrent.ScheduledTaskCapturingEventLoop; @@ -33,6 +34,7 @@ import java.time.Duration; import java.util.List; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -164,6 +166,15 @@ public void should_reject_when_queue_is_full() { @Test public void should_remove_timed_out_request_from_queue() { + testRemoveInvalidEventFromQueue(throttler::signalTimeout); + } + + @Test + public void should_remove_cancel_request_from_queue() { + testRemoveInvalidEventFromQueue(throttler::signalCancel); + } + + private void testRemoveInvalidEventFromQueue(Consumer completeCallback) { // Given for (int i = 0; i < 5; i++) { throttler.register(new MockThrottled()); @@ -174,7 +185,7 @@ public void should_remove_timed_out_request_from_queue() { throttler.register(queued2); // When - throttler.signalTimeout(queued1); + completeCallback.accept(queued1); // Then assertThatStage(queued2.started).isNotDone(); diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/throttling/ThrottlingIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/throttling/ThrottlingIT.java index a6e7295eb09..6fa1a37355b 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/throttling/ThrottlingIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/throttling/ThrottlingIT.java @@ -24,13 +24,16 @@ import com.datastax.oss.driver.api.core.RequestThrottlingException; import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import com.datastax.oss.driver.api.core.config.DriverConfigLoader; +import com.datastax.oss.driver.api.core.cql.AsyncResultSet; import com.datastax.oss.driver.api.testinfra.session.SessionUtils; import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule; import com.datastax.oss.driver.categories.ParallelizableTests; import com.datastax.oss.driver.internal.core.session.throttling.ConcurrencyLimitingRequestThrottler; import com.datastax.oss.simulacron.common.cluster.ClusterSpec; import com.datastax.oss.simulacron.common.stubbing.PrimeDsl; +import java.util.concurrent.CompletionStage; import java.util.concurrent.TimeUnit; +import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -39,21 +42,20 @@ public class ThrottlingIT { private static final String QUERY = "select * from foo"; + private static final int maxConcurrentRequests = 10; + private static final int maxQueueSize = 10; @Rule public SimulacronRule simulacron = new SimulacronRule(ClusterSpec.builder().withNodes(1)); - @Test - public void should_reject_request_when_throttling_by_concurrency() { + private DriverConfigLoader loader = null; + @Before + public void setUp() { // Add a delay so that requests don't complete during the test simulacron .cluster() .prime(PrimeDsl.when(QUERY).then(PrimeDsl.noRows()).delay(5, TimeUnit.SECONDS)); - - int maxConcurrentRequests = 10; - int maxQueueSize = 10; - - DriverConfigLoader loader = + loader = SessionUtils.configLoaderBuilder() .withClass( DefaultDriverOption.REQUEST_THROTTLER_CLASS, @@ -63,7 +65,10 @@ public void should_reject_request_when_throttling_by_concurrency() { maxConcurrentRequests) .withInt(DefaultDriverOption.REQUEST_THROTTLER_MAX_QUEUE_SIZE, maxQueueSize) .build(); + } + @Test + public void should_reject_request_when_throttling_by_concurrency() { try (CqlSession session = SessionUtils.newSession(simulacron, loader)) { // Saturate the session and fill the queue @@ -81,4 +86,19 @@ public void should_reject_request_when_throttling_by_concurrency() { + "(concurrent requests: 10, queue size: 10)"); } } + + @Test + public void should_propagate_cancel_to_throttler() { + try (CqlSession session = SessionUtils.newSession(simulacron, loader)) { + + // Try to saturate the session and fill the queue + for (int i = 0; i < maxConcurrentRequests + maxQueueSize; i++) { + CompletionStage future = session.executeAsync(QUERY); + future.toCompletableFuture().cancel(true); + } + + // The next query should be successful, because the previous queries were cancelled + session.execute(QUERY); + } + } }