Cap total number of concurrent requests per HTTP/2 connection

Author: arturobernalg

httpcore5-h2/src/main/java/org/apache/hc/core5/http2/impl/nio/bootstrap/H2MultiplexingRequester.java

@@ -32,7 +32,11 @@
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.hc.core5.annotation.Internal;
 import org.apache.hc.core5.concurrent.Cancellable;
@@ -43,7 +47,6 @@
 import org.apache.hc.core5.function.Callback;
 import org.apache.hc.core5.function.Decorator;
 import org.apache.hc.core5.function.Resolver;
-import org.apache.hc.core5.http.ConnectionClosedException;
 import org.apache.hc.core5.http.EntityDetails;
 import org.apache.hc.core5.http.Header;
 import org.apache.hc.core5.http.HttpException;
@@ -86,6 +89,8 @@
 public class H2MultiplexingRequester extends AsyncRequester {
 
     private final H2ConnPool connPool;
+    private final ConcurrentMap<IOSession, AtomicInteger> inFlightPerSession;
+    private final int maxRequestsPerConnection;
 
     /**
      * Use {@link H2MultiplexingRequesterBootstrap} to create instances of this class.
@@ -100,11 +105,44 @@ public H2MultiplexingRequester(
             final Resolver<HttpHost, InetSocketAddress> addressResolver,
             final TlsStrategy tlsStrategy,
             final IOReactorMetricsListener threadPoolListener,
-            final IOWorkerSelector workerSelector) {
+            final IOWorkerSelector workerSelector,
+            final int maxRequestsPerConnection) {
         super(eventHandlerFactory, ioReactorConfig, ioSessionDecorator, exceptionCallback, sessionListener,
                 ShutdownCommand.GRACEFUL_IMMEDIATE_CALLBACK, DefaultAddressResolver.INSTANCE,
                 threadPoolListener, workerSelector);
         this.connPool = new H2ConnPool(this, addressResolver, tlsStrategy);
+        this.inFlightPerSession = new ConcurrentHashMap<>();
+        this.maxRequestsPerConnection = maxRequestsPerConnection;
+    }
+
+    private boolean tryAcquireSlot(final IOSession ioSession, final int max) {
+        if (max <= 0) {
+            return true;
+        }
+        final AtomicInteger counter = inFlightPerSession.computeIfAbsent(ioSession, s -> new AtomicInteger(0));
+        for (;;) {
+            final int q = counter.get();
+            if (q >= max) {
+                return false;
+            }
+            if (counter.compareAndSet(q, q + 1)) {
+                return true;
+            }
+        }
+    }
+
+    private void releaseSlot(final IOSession ioSession, final int max) {
+        if (max <= 0) {
+            return;
+        }
+        final AtomicInteger counter = inFlightPerSession.get(ioSession);
+        if (counter == null) {
+            return;
+        }
+        final int q = counter.decrementAndGet();
+        if (q <= 0) {
+            inFlightPerSession.remove(ioSession, counter);
+        }
     }
 
     public void closeIdle(final TimeValue idleTime) {
@@ -182,15 +220,29 @@ private void execute(
                 if (request.getAuthority() == null) {
                     request.setAuthority(new URIAuthority(host));
                 }
+                if (request.getScheme() == null) {
+                    request.setScheme(host.getSchemeName());
+                }
                 connPool.getSession(host, timeout, new FutureCallback<IOSession>() {
 
                     @Override
                     public void completed(final IOSession ioSession) {
+                        if (!tryAcquireSlot(ioSession, maxRequestsPerConnection)) {
+                            exchangeHandler.failed(new RejectedExecutionException(
+                                    "Maximum number of concurrent requests per connection reached (max=" + maxRequestsPerConnection + ")"));
+                            exchangeHandler.releaseResources();
+                            return;
+                        }
+
+                        final AsyncClientExchangeHandler actual = maxRequestsPerConnection > 0
+                                ? new ReleasingAsyncClientExchangeHandler(exchangeHandler, () -> releaseSlot(ioSession, maxRequestsPerConnection))
+                                : exchangeHandler;
+
                         final AsyncClientExchangeHandler handlerProxy = new AsyncClientExchangeHandler() {
 
                             @Override
                             public void releaseResources() {
-                                exchangeHandler.releaseResources();
+                                actual.releaseResources();
                             }
 
                             @Override
@@ -199,67 +251,70 @@ public void produceRequest(final RequestChannel channel, final HttpContext httpC
                             }
 
                             @Override
-                            public int available() {
-                                return exchangeHandler.available();
+                            public void consumeResponse(
+                                    final HttpResponse response, final EntityDetails entityDetails, final HttpContext httpContext) throws HttpException, IOException {
+                                actual.consumeResponse(response, entityDetails, httpContext);
                             }
 
                             @Override
-                            public void produce(final DataStreamChannel channel) throws IOException {
-                                exchangeHandler.produce(channel);
+                            public void consumeInformation(final HttpResponse response, final HttpContext httpContext) throws HttpException, IOException {
+                                actual.consumeInformation(response, httpContext);
                             }
 
                             @Override
-                            public void consumeInformation(final HttpResponse response, final HttpContext httpContext) throws HttpException, IOException {
-                                exchangeHandler.consumeInformation(response, httpContext);
+                            public int available() {
+                                return actual.available();
                             }
 
                             @Override
-                            public void consumeResponse(
-                                    final HttpResponse response, final EntityDetails entityDetails, final HttpContext httpContext) throws HttpException, IOException {
-                                exchangeHandler.consumeResponse(response, entityDetails, httpContext);
+                            public void produce(final DataStreamChannel channel) throws IOException {
+                                actual.produce(channel);
                             }
 
                             @Override
                             public void updateCapacity(final CapacityChannel capacityChannel) throws IOException {
-                                exchangeHandler.updateCapacity(capacityChannel);
+                                actual.updateCapacity(capacityChannel);
                             }
 
                             @Override
                             public void consume(final ByteBuffer src) throws IOException {
-                                exchangeHandler.consume(src);
+                                actual.consume(src);
                             }
 
                             @Override
                             public void streamEnd(final List<? extends Header> trailers) throws HttpException, IOException {
-                                exchangeHandler.streamEnd(trailers);
+                                actual.streamEnd(trailers);
                             }
 
                             @Override
                             public void cancel() {
-                                exchangeHandler.cancel();
+                                actual.cancel();
                             }
 
                             @Override
                             public void failed(final Exception cause) {
-                                exchangeHandler.failed(cause);
+                                actual.failed(cause);
                             }
 
                         };
                         final Timeout socketTimeout = ioSession.getSocketTimeout();
-                        ioSession.enqueue(new RequestExecutionCommand(
-                                        handlerProxy,
-                                        pushHandlerFactory,
-                                        context,
-                                        streamControl -> {
-                                            cancellableDependency.setDependency(streamControl);
-                                            if (socketTimeout != null) {
-                                                streamControl.setTimeout(socketTimeout);
-                                            }
-                                        }),
-                                Command.Priority.NORMAL);
-                        if (!ioSession.isOpen()) {
-                            exchangeHandler.failed(new ConnectionClosedException());
+                        try {
+                            ioSession.enqueue(new RequestExecutionCommand(
+                                            handlerProxy,
+                                            pushHandlerFactory,
+                                            context,
+                                            streamControl -> {
+                                                cancellableDependency.setDependency(streamControl);
+                                                if (socketTimeout != null) {
+                                                    streamControl.setTimeout(socketTimeout);
+                                                }
+                                            }),
+                                    Command.Priority.NORMAL);
+                        } catch (final RuntimeException ex) {
+                            actual.failed(ex);
+                            actual.releaseResources();
                         }
+
                     }
 
                     @Override

httpcore5-h2/src/main/java/org/apache/hc/core5/http2/impl/nio/bootstrap/H2MultiplexingRequesterBootstrap.java

@@ -29,6 +29,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import org.apache.hc.core5.annotation.Experimental;
 import org.apache.hc.core5.function.Callback;
 import org.apache.hc.core5.function.Decorator;
 import org.apache.hc.core5.function.Supplier;
@@ -76,6 +77,8 @@ public class H2MultiplexingRequesterBootstrap {
 
     private IOReactorMetricsListener threadPoolListener;
 
+    private int maxRequestsPerConnection;
+
     private H2MultiplexingRequesterBootstrap() {
         this.routeEntries = new ArrayList<>();
     }
@@ -180,6 +183,21 @@ public final H2MultiplexingRequesterBootstrap setIOReactorMetricsListener(final
         return this;
     }
 
+    /**
+     * Sets a hard cap on the number of requests allowed to be queued / in-flight per connection.
+     * When the limit is reached, new submissions fail fast with {@link java.util.concurrent.RejectedExecutionException}.
+     * A value {@code <= 0} means unlimited (default).
+     *
+     * @param max maximum number of requests per connection; {@code <= 0} to disable the cap
+     * @return this instance.
+     * @since 5.5
+     */
+    @Experimental
+    public final H2MultiplexingRequesterBootstrap setMaxRequestsPerConnection(final int max) {
+        this.maxRequestsPerConnection = max;
+        return this;
+    }
+
     /**
      * Sets {@link H2StreamListener} instance.
      *
@@ -274,7 +292,8 @@ public H2MultiplexingRequester create() {
                 DefaultAddressResolver.INSTANCE,
                 tlsStrategy != null ? tlsStrategy : new H2ClientTlsStrategy(),
                 threadPoolListener,
-                null);
+                null,
+                maxRequestsPerConnection);
     }
 
 }

httpcore5-h2/src/main/java/org/apache/hc/core5/http2/impl/nio/bootstrap/ReleasingAsyncClientExchangeHandler.java

@@ -0,0 +1,118 @@
+/*
+ * ====================================================================
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ * ====================================================================
+ *
+ * This software consists of voluntary contributions made by many
+ * individuals on behalf of the Apache Software Foundation.  For more
+ * information on the Apache Software Foundation, please see
+ * <http://www.apache.org/>.
+ *
+ */
+package org.apache.hc.core5.http2.impl.nio.bootstrap;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.apache.hc.core5.http.EntityDetails;
+import org.apache.hc.core5.http.Header;
+import org.apache.hc.core5.http.HttpException;
+import org.apache.hc.core5.http.HttpResponse;
+import org.apache.hc.core5.http.nio.AsyncClientExchangeHandler;
+import org.apache.hc.core5.http.nio.CapacityChannel;
+import org.apache.hc.core5.http.nio.DataStreamChannel;
+import org.apache.hc.core5.http.nio.RequestChannel;
+import org.apache.hc.core5.http.protocol.HttpContext;
+
+final class ReleasingAsyncClientExchangeHandler implements AsyncClientExchangeHandler {
+
+    private final AsyncClientExchangeHandler exchangeHandler;
+    private final Runnable onRelease;
+    private final AtomicBoolean released;
+
+    ReleasingAsyncClientExchangeHandler(final AsyncClientExchangeHandler exchangeHandler, final Runnable onRelease) {
+        this.exchangeHandler = exchangeHandler;
+        this.onRelease = onRelease;
+        this.released = new AtomicBoolean(false);
+    }
+
+    @Override
+    public void produceRequest(final RequestChannel channel, final HttpContext context) throws HttpException, IOException {
+        exchangeHandler.produceRequest(channel, context);
+    }
+
+    @Override
+    public void consumeResponse(final HttpResponse response, final EntityDetails entityDetails, final HttpContext context)
+            throws HttpException, IOException {
+        exchangeHandler.consumeResponse(response, entityDetails, context);
+    }
+
+    @Override
+    public void consumeInformation(final HttpResponse response, final HttpContext context) throws HttpException, IOException {
+        exchangeHandler.consumeInformation(response, context);
+    }
+
+    @Override
+    public int available() {
+        return exchangeHandler.available();
+    }
+
+    @Override
+    public void produce(final DataStreamChannel channel) throws IOException {
+        exchangeHandler.produce(channel);
+    }
+
+    @Override
+    public void updateCapacity(final CapacityChannel capacityChannel) throws IOException {
+        exchangeHandler.updateCapacity(capacityChannel);
+    }
+
+    @Override
+    public void consume(final ByteBuffer src) throws IOException {
+        exchangeHandler.consume(src);
+    }
+
+    @Override
+    public void streamEnd(final List<? extends Header> trailers) throws HttpException, IOException {
+        exchangeHandler.streamEnd(trailers);
+    }
+
+    @Override
+    public void failed(final Exception cause) {
+        exchangeHandler.failed(cause);
+    }
+
+    @Override
+    public void cancel() {
+        exchangeHandler.cancel();
+    }
+
+    @Override
+    public void releaseResources() {
+        try {
+            exchangeHandler.releaseResources();
+        } finally {
+            if (released.compareAndSet(false, true)) {
+                onRelease.run();
+            }
+        }
+    }
+
+}