JAVA-3149: Support request cancellation in request throttler

patch by Lukasz Antoniak; reviewed by Andy Tolbert and Chris Lohfink for JAVA-3149

Author: lukasz-antoniak

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

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

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
+  }
 }

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

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

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

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

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

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<Throttled> completeCallback) {
     // Given

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

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<AsyncResultSet> future = session.executeAsync(QUERY);
+        future.toCompletableFuture().cancel(true);
+      }
+
+      // The next query should be successful, because the previous queries were cancelled
+      session.execute(QUERY);
+    }
+  }
 }