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