Reduce lock held duration in ConcurrencyLimitingRequestThrottler

It might take some (small) time for callback handling when the
throttler request proceeds to submission.

Before this change, the throttler proceed request will happen while
holding the lock, preventing other tasks from proceeding when there is
spare capacity and even preventing tasks from enqueuing until the
callback completes.

By tracking the expected outcome, we can perform the callback outside
of the lock. This means that request registration and submission can
proceed even when a long callback is being processed.

patch by Jason Koch; Reviewed by Andy Tolbert and Chris Lohfink for CASSANDRA-19922

Author: jasonk000

core/src/main/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottler.java

@@ -25,6 +25,7 @@
 import com.datastax.oss.driver.api.core.session.throttling.Throttled;
 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.locks.ReentrantLock;
@@ -87,6 +88,8 @@ public ConcurrencyLimitingRequestThrottler(DriverContext context) {
 
   @Override
   public void register(@NonNull Throttled request) {
+    boolean notifyReadyRequired = false;
+
     lock.lock();
     try {
       if (closed) {
@@ -96,7 +99,7 @@ public void register(@NonNull Throttled request) {
         // We have capacity for one more concurrent request
         LOG.trace("[{}] Starting newly registered request", logPrefix);
         concurrentRequests += 1;
-        request.onThrottleReady(false);
+        notifyReadyRequired = true;
       } else if (queue.size() < maxQueueSize) {
         LOG.trace("[{}] Enqueuing request", logPrefix);
         queue.add(request);
@@ -112,16 +115,26 @@ public void register(@NonNull Throttled request) {
     } finally {
       lock.unlock();
     }
+
+    // no need to hold the lock while allowing the task to progress
+    if (notifyReadyRequired) {
+      request.onThrottleReady(false);
+    }
   }
 
   @Override
   public void signalSuccess(@NonNull Throttled request) {
+    Throttled nextRequest = null;
     lock.lock();
     try {
-      onRequestDone();
+      nextRequest = onRequestDoneAndDequeNext();
     } finally {
       lock.unlock();
     }
+
+    if (nextRequest != null) {
+      nextRequest.onThrottleReady(true);
+    }
   }
 
   @Override
@@ -131,48 +144,62 @@ public void signalError(@NonNull Throttled request, @NonNull Throwable error) {
 
   @Override
   public void signalTimeout(@NonNull Throttled request) {
+    Throttled nextRequest = null;
     lock.lock();
     try {
       if (!closed) {
         if (queue.remove(request)) { // The request timed out before it was active
           LOG.trace("[{}] Removing timed out request from the queue", logPrefix);
         } else {
-          onRequestDone();
+          nextRequest = onRequestDoneAndDequeNext();
         }
       }
     } finally {
       lock.unlock();
     }
+
+    if (nextRequest != null) {
+      nextRequest.onThrottleReady(true);
+    }
   }
 
   @Override
   public void signalCancel(@NonNull Throttled request) {
+    Throttled nextRequest = null;
     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();
+          nextRequest = onRequestDoneAndDequeNext();
         }
       }
     } finally {
       lock.unlock();
     }
+
+    if (nextRequest != null) {
+      nextRequest.onThrottleReady(true);
+    }
   }
 
   @SuppressWarnings("GuardedBy") // this method is only called with the lock held
-  private void onRequestDone() {
+  @Nullable
+  private Throttled onRequestDoneAndDequeNext() {
     assert lock.isHeldByCurrentThread();
     if (!closed) {
       if (queue.isEmpty()) {
         concurrentRequests -= 1;
       } else {
         LOG.trace("[{}] Starting dequeued request", logPrefix);
-        queue.poll().onThrottleReady(true);
         // don't touch concurrentRequests since we finished one but started another
+        return queue.poll();
       }
     }
+
+    // no next task was dequeued
+    return null;
   }
 
   @Override

core/src/test/java/com/datastax/oss/driver/internal/core/session/throttling/ConcurrencyLimitingRequestThrottlerTest.java

@@ -29,6 +29,7 @@
 import com.datastax.oss.driver.api.core.session.throttling.Throttled;
 import com.datastax.oss.driver.shaded.guava.common.collect.Lists;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.function.Consumer;
 import org.junit.Before;
 import org.junit.Test;
@@ -67,7 +68,7 @@ public void should_start_immediately_when_under_capacity() {
     throttler.register(request);
 
     // Then
-    assertThatStage(request.started).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
+    assertThatStage(request.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
     assertThat(throttler.getConcurrentRequests()).isEqualTo(1);
     assertThat(throttler.getQueue()).isEmpty();
   }
@@ -98,7 +99,7 @@ private void should_allow_new_request_when_active_one_completes(
     // Given
     MockThrottled first = new MockThrottled();
     throttler.register(first);
-    assertThatStage(first.started).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
+    assertThatStage(first.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
     for (int i = 0; i < 4; i++) { // fill to capacity
       throttler.register(new MockThrottled());
     }
@@ -113,7 +114,7 @@ private void should_allow_new_request_when_active_one_completes(
     throttler.register(incoming);
 
     // Then
-    assertThatStage(incoming.started).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
+    assertThatStage(incoming.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
     assertThat(throttler.getConcurrentRequests()).isEqualTo(5);
     assertThat(throttler.getQueue()).isEmpty();
   }
@@ -132,7 +133,7 @@ public void should_enqueue_when_over_capacity() {
     throttler.register(incoming);
 
     // Then
-    assertThatStage(incoming.started).isNotDone();
+    assertThatStage(incoming.ended).isNotDone();
     assertThat(throttler.getConcurrentRequests()).isEqualTo(5);
     assertThat(throttler.getQueue()).containsExactly(incoming);
   }
@@ -157,20 +158,20 @@ private void should_dequeue_when_active_completes(Consumer<Throttled> completeCa
     // Given
     MockThrottled first = new MockThrottled();
     throttler.register(first);
-    assertThatStage(first.started).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
+    assertThatStage(first.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
     for (int i = 0; i < 4; i++) {
       throttler.register(new MockThrottled());
     }
 
     MockThrottled incoming = new MockThrottled();
     throttler.register(incoming);
-    assertThatStage(incoming.started).isNotDone();
+    assertThatStage(incoming.ended).isNotDone();
 
     // When
     completeCallback.accept(first);
 
     // Then
-    assertThatStage(incoming.started).isSuccess(wasDelayed -> assertThat(wasDelayed).isTrue());
+    assertThatStage(incoming.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isTrue());
     assertThat(throttler.getConcurrentRequests()).isEqualTo(5);
     assertThat(throttler.getQueue()).isEmpty();
   }
@@ -189,7 +190,7 @@ public void should_reject_when_queue_is_full() {
     throttler.register(incoming);
 
     // Then
-    assertThatStage(incoming.started)
+    assertThatStage(incoming.ended)
         .isFailed(error -> assertThat(error).isInstanceOf(RequestThrottlingException.class));
   }
 
@@ -208,7 +209,7 @@ public void should_remove_timed_out_request_from_queue() {
     throttler.signalTimeout(queued1);
 
     // Then
-    assertThatStage(queued2.started).isNotDone();
+    assertThatStage(queued2.ended).isNotDone();
     assertThat(throttler.getConcurrentRequests()).isEqualTo(5);
     assertThat(throttler.getQueue()).hasSize(1);
   }
@@ -223,7 +224,7 @@ public void should_reject_enqueued_when_closing() {
     for (int i = 0; i < 10; i++) {
       MockThrottled request = new MockThrottled();
       throttler.register(request);
-      assertThatStage(request.started).isNotDone();
+      assertThatStage(request.ended).isNotDone();
       enqueued.add(request);
     }
 
@@ -232,7 +233,7 @@ public void should_reject_enqueued_when_closing() {
 
     // Then
     for (MockThrottled request : enqueued) {
-      assertThatStage(request.started)
+      assertThatStage(request.ended)
           .isFailed(error -> assertThat(error).isInstanceOf(RequestThrottlingException.class));
     }
 
@@ -241,7 +242,125 @@ public void should_reject_enqueued_when_closing() {
     throttler.register(request);
 
     // Then
-    assertThatStage(request.started)
+    assertThatStage(request.ended)
         .isFailed(error -> assertThat(error).isInstanceOf(RequestThrottlingException.class));
   }
+
+  @Test
+  public void should_run_throttle_callbacks_concurrently() throws InterruptedException {
+    // Given
+
+    // a task is enqueued, which when in onThrottleReady, will stall latch countDown()ed
+    // register() should automatically start onThrottleReady on same thread
+
+    // start a parallel thread
+    CountDownLatch firstRelease = new CountDownLatch(1);
+    MockThrottled first = new MockThrottled(firstRelease);
+    Runnable r =
+        () -> {
+          throttler.register(first);
+          first.ended.toCompletableFuture().thenRun(() -> throttler.signalSuccess(first));
+        };
+    Thread t = new Thread(r);
+    t.start();
+
+    // wait for the registration threads to reach await state
+    assertThatStage(first.started).isSuccess();
+    assertThatStage(first.ended).isNotDone();
+
+    // When
+    // we concurrently submit a second shorter task
+    MockThrottled second = new MockThrottled();
+    // (on a second thread, so that we can join and force a timeout in case
+    // registration is delayed)
+    Thread t2 = new Thread(() -> throttler.register(second));
+    t2.start();
+    t2.join(1_000);
+
+    // Then
+    // registration will trigger callback, should complete ~immediately
+    assertThatStage(second.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
+    // first should still be unfinished
+    assertThatStage(first.started).isDone();
+    assertThatStage(first.ended).isNotDone();
+    // now finish, and verify
+    firstRelease.countDown();
+    assertThatStage(first.ended).isSuccess(wasDelayed -> assertThat(wasDelayed).isFalse());
+
+    t.join(1_000);
+  }
+
+  @Test
+  public void should_enqueue_tasks_quickly_when_callbacks_blocked() throws InterruptedException {
+    // Given
+
+    // Multiple tasks are registered, up to the limit, and proceed into their
+    // callback
+
+    // start five parallel threads
+    final int THREADS = 5;
+    Thread[] threads = new Thread[THREADS];
+    CountDownLatch[] latches = new CountDownLatch[THREADS];
+    MockThrottled[] throttled = new MockThrottled[THREADS];
+    for (int i = 0; i < threads.length; i++) {
+      latches[i] = new CountDownLatch(1);
+      final MockThrottled itThrottled = new MockThrottled(latches[i]);
+      throttled[i] = itThrottled;
+      threads[i] =
+          new Thread(
+              () -> {
+                throttler.register(itThrottled);
+                itThrottled
+                    .ended
+                    .toCompletableFuture()
+                    .thenRun(() -> throttler.signalSuccess(itThrottled));
+              });
+      threads[i].start();
+    }
+
+    // wait for the registration threads to be launched
+    // they are all waiting now
+    for (int i = 0; i < throttled.length; i++) {
+      assertThatStage(throttled[i].started).isSuccess();
+      assertThatStage(throttled[i].ended).isNotDone();
+    }
+
+    // When
+    // we concurrently submit another task
+    MockThrottled last = new MockThrottled();
+    throttler.register(last);
+
+    // Then
+    // registration will enqueue the callback, and it should not
+    // take any time to proceed (ie: we should not be blocked)
+    // and there should be an element in the queue
+    assertThatStage(last.started).isNotDone();
+    assertThatStage(last.ended).isNotDone();
+    assertThat(throttler.getQueue()).containsExactly(last);
+
+    // we still have not released, so old throttled threads should be waiting
+    for (int i = 0; i < throttled.length; i++) {
+      assertThatStage(throttled[i].started).isDone();
+      assertThatStage(throttled[i].ended).isNotDone();
+    }
+
+    // now let us release ..
+    for (int i = 0; i < latches.length; i++) {
+      latches[i].countDown();
+    }
+
+    // .. and check everything finished up OK
+    for (int i = 0; i < latches.length; i++) {
+      assertThatStage(throttled[i].started).isSuccess();
+      assertThatStage(throttled[i].ended).isSuccess();
+    }
+
+    // for good measure, we will also wait for the enqueued to complete
+    assertThatStage(last.started).isSuccess();
+    assertThatStage(last.ended).isSuccess();
+
+    for (int i = 0; i < threads.length; i++) {
+      threads[i].join(1_000);
+    }
+  }
 }