Fix cancellation races in shared async execution queue

Author: fitekonea

httpclient5-testing/src/test/java/org/apache/hc/client5/http/impl/async/InternalTestHttpAsyncExecRuntime.java

@@ -31,7 +31,6 @@
 import java.util.List;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.hc.client5.http.HttpRoute;
 import org.apache.hc.client5.http.async.AsyncExecRuntime;
@@ -64,7 +63,7 @@ public final class InternalTestHttpAsyncExecRuntime extends InternalHttpAsyncExe
     public InternalTestHttpAsyncExecRuntime(final AsyncClientConnectionManager manager,
                                             final ConnectionInitiator connectionInitiator,
                                             final TlsConfig tlsConfig) {
-        super(LOG, manager, connectionInitiator, null, tlsConfig, -1, new AtomicInteger());
+        super(LOG, manager, connectionInitiator, null, tlsConfig);
         this.cancelled = new AtomicBoolean();
     }
 

httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/H2AsyncClientBuilder.java

@@ -328,12 +328,12 @@ public final H2AsyncClientBuilder disableRequestPriority() {
     }
 
     /**
-     * Sets a hard cap on the number of requests allowed to be queued/in-flight
-     * within the internal async execution pipeline. When the limit is reached,
-     * new submissions fail fast with {@link java.util.concurrent.RejectedExecutionException}.
+     * Sets the maximum number of concurrent request executions within the
+     * internal async execution pipeline. Requests beyond this limit are queued
+     * in FIFO order and dispatched as in-flight executions complete.
      * A value {@code <= 0} means unlimited (default).
      *
-     * @param max maximum number of queued requests; {@code <= 0} to disable the cap
+     * @param max maximum number of concurrent executions; {@code <= 0} to disable the cap
      * @return this builder
      * @since 5.7
      */

httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/HttpAsyncClientBuilder.java

@@ -910,12 +910,12 @@ public final HttpAsyncClientBuilder setTlsRequired(final boolean tlsRequired) {
     }
 
     /**
-     * Sets a hard cap on the number of requests allowed to be queued/in-flight
-     * within the internal async execution pipeline. When the limit is reached,
-     * new submissions fail fast with {@link java.util.concurrent.RejectedExecutionException}.
-     * A value <= 0 means unlimited (default).
+     * Sets the maximum number of concurrent request executions within the
+     * internal async execution pipeline. Requests beyond this limit are queued
+     * in FIFO order and dispatched as in-flight executions complete.
+     * A value {@code <= 0} means unlimited (default).
      *
-     * @param max maximum number of queued requests; <= 0 to disable the cap
+     * @param max maximum number of concurrent executions; {@code <= 0} to disable the cap
      * @return this builder
      * @since 5.7
      */

httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/InternalH2AsyncClient.java

@@ -67,8 +67,8 @@ public final class InternalH2AsyncClient extends InternalAbstractHttpAsyncClient
     private final InternalH2ConnPool connPool;
 
     /**
-     * One shared FIFO queue per client instance (Oleg).
-     * null means "unlimited" / no throttling.
+     * One shared FIFO queue per client instance.
+     * {@code null} means unlimited / no throttling.
      */
     private final SharedRequestExecutionQueue executionQueue;
 

httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/InternalHttpAsyncExecRuntime.java

@@ -383,29 +383,26 @@ public Cancellable execute(
             return startExecution(id, exchangeHandler, context);
         }
 
-        final AtomicReference<Cancellable> activeCancelRef = new AtomicReference<>(Operations.nonCancellable());
+        final ComplexCancellable complexCancellable = new ComplexCancellable();
 
         final Cancellable queued = executionQueue.enqueue(
                 () -> {
                     final AsyncClientExchangeHandler wrapped =
                             new ReleasingAsyncClientExchangeHandler(exchangeHandler, executionQueue::completed);
                     try {
                         final Cancellable cancellable = startExecution(id, wrapped, context);
-                        activeCancelRef.set(cancellable);
+                        complexCancellable.setDependency(cancellable);
                     } catch (final RuntimeException ex) {
                         wrapped.failed(ex);
                     }
                 },
-                () -> {
-                    activeCancelRef.get().cancel();
-                    exchangeHandler.cancel();
-                });
+                exchangeHandler::cancel);
 
         return () -> {
             if (queued.cancel()) {
                 return true;
             }
-            return activeCancelRef.get().cancel();
+            return complexCancellable.cancel();
         };
     }
 

httpclient5/src/main/java/org/apache/hc/client5/http/impl/async/SharedRequestExecutionQueue.java

@@ -30,6 +30,7 @@
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 
 import org.apache.hc.core5.annotation.Internal;
 import org.apache.hc.core5.concurrent.Cancellable;
@@ -83,10 +84,9 @@ private void drain() {
                     if (entry == null) {
                         break;
                     }
-                    if (entry.isCancelled()) {
+                    if (!entry.tryStart()) {
                         continue;
                     }
-                    entry.markStarted();
                     inFlight.incrementAndGet();
                     entry.run();
                 }
@@ -102,36 +102,32 @@ private void drain() {
 
     private static final class Entry implements Cancellable {
 
+        private static final int QUEUED = 0;
+        private static final int STARTED = 1;
+        private static final int CANCELLED = 2;
+
+        private static final AtomicIntegerFieldUpdater<Entry> STATE = AtomicIntegerFieldUpdater.newUpdater(Entry.class, "state");
+
         private final Runnable task;
         private final Runnable onCancel;
-        private final AtomicBoolean started;
-        private final AtomicBoolean cancelled;
+        private volatile int state;
 
         Entry(final Runnable task, final Runnable onCancel) {
             this.task = task;
             this.onCancel = onCancel;
-            this.started = new AtomicBoolean(false);
-            this.cancelled = new AtomicBoolean(false);
         }
 
         void run() {
             task.run();
         }
 
-        void markStarted() {
-            started.set(true);
-        }
-
-        boolean isCancelled() {
-            return cancelled.get();
+        boolean tryStart() {
+            return STATE.compareAndSet(this, QUEUED, STARTED);
         }
 
         @Override
         public boolean cancel() {
-            if (started.get()) {
-                return false;
-            }
-            if (cancelled.compareAndSet(false, true)) {
+            if (STATE.compareAndSet(this, QUEUED, CANCELLED)) {
                 onCancel.run();
                 return true;
             }