async downstream executor (DE-697)

rashtao · rashtao · commit deba97d1d699 · 2023-10-25T21:45:38.000+02:00
diff --git a/core/src/main/java/com/arangodb/ArangoDB.java b/core/src/main/java/com/arangodb/ArangoDB.java
@@ -340,11 +340,6 @@ public interface ArangoDB extends ArangoSerdeAccessor {
      */
     class Builder extends InternalArangoDBBuilder<Builder> {
 
-        public Builder protocol(final Protocol protocol) {
-            config.setProtocol(protocol);
-            return this;
-        }
-
         /**
          * Returns an instance of {@link ArangoDB}.
          *
diff --git a/core/src/main/java/com/arangodb/internal/ArangoExecutorAsync.java b/core/src/main/java/com/arangodb/internal/ArangoExecutorAsync.java
@@ -27,16 +27,20 @@
 
 import java.lang.reflect.Type;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 /**
  * @author Mark Vollmary
  * @author Michele Rastelli
  */
 public class ArangoExecutorAsync extends ArangoExecutor {
+    private final Executor downstreamExecutor;
 
     public ArangoExecutorAsync(final CommunicationProtocol protocol, final ArangoConfig config) {
         super(protocol, config);
+        downstreamExecutor = config.getAsyncExecutor();
     }
 
     public <T> CompletableFuture<T> execute(final Supplier<InternalRequest> requestSupplier, final Type type) {
@@ -56,18 +60,23 @@ public <T> CompletableFuture<T> execute(
             final ResponseDeserializer<T> responseDeserializer,
             final HostHandle hostHandle) {
 
-        return CompletableFuture.completedFuture(requestSupplier)
+        CompletableFuture<T> cf = CompletableFuture.completedFuture(requestSupplier)
                 .thenApply(Supplier::get)
-                .thenCompose(request -> protocol.executeAsync(interceptRequest(request), hostHandle)
-                        .handle((r, e) -> {
-                            if (e != null) {
-                                throw ArangoDBException.of(e);
-                            } else {
-                                interceptResponse(r);
-                                return responseDeserializer.deserialize(r);
-                            }
-                        })
-                );
+                .thenCompose(request -> protocol.executeAsync(interceptRequest(request), hostHandle))
+                .handle((r, e) -> {
+                    if (e != null) {
+                        throw ArangoDBException.of(e);
+                    } else {
+                        interceptResponse(r);
+                        return responseDeserializer.deserialize(r);
+                    }
+                });
+
+        if (downstreamExecutor != null) {
+            return cf.thenApplyAsync(Function.identity(), downstreamExecutor);
+        } else {
+            return cf;
+        }
     }
 
 }
diff --git a/core/src/main/java/com/arangodb/internal/InternalArangoDBBuilder.java b/core/src/main/java/com/arangodb/internal/InternalArangoDBBuilder.java
@@ -38,6 +38,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.ServiceLoader;
+import java.util.concurrent.Executor;
 
 
 /**
@@ -53,6 +54,12 @@ public T loadProperties(final ArangoConfigProperties properties) {
         return (T) this;
     }
 
+    @SuppressWarnings("unchecked")
+    public T protocol(final Protocol protocol) {
+        config.setProtocol(protocol);
+        return (T) this;
+    }
+
     /**
      * Adds a host to connect to. Multiple hosts can be added to provide fallbacks.
      *
@@ -289,6 +296,19 @@ public T serde(final ArangoSerde serde) {
         return (T) this;
     }
 
+    /**
+     * Sets the downstream async executor that will be used to consume the responses of the async API, that are returned
+     * as {@link java.util.concurrent.CompletableFuture}
+     *
+     * @param executor async downstream executor
+     * @return {@link ArangoDB.Builder}
+     */
+    @SuppressWarnings("unchecked")
+    public T asyncExecutor(final Executor executor) {
+        config.setAsyncExecutor(executor);
+        return (T) this;
+    }
+
     protected ProtocolProvider protocolProvider(Protocol protocol) {
         ServiceLoader<ProtocolProvider> loader = ServiceLoader.load(ProtocolProvider.class);
         for (ProtocolProvider p : loader) {
diff --git a/core/src/main/java/com/arangodb/internal/config/ArangoConfig.java b/core/src/main/java/com/arangodb/internal/config/ArangoConfig.java
@@ -18,6 +18,7 @@
 
 import javax.net.ssl.SSLContext;
 import java.util.*;
+import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
 
 public class ArangoConfig {
@@ -41,6 +42,7 @@ public class ArangoConfig {
     private ArangoSerde userDataSerde;
     private Integer responseQueueTimeSamples;
     private Module protocolModule;
+    private Executor asyncExecutor;
 
     private static final Logger LOG = LoggerFactory.getLogger(ArangoConfig.class);
 
@@ -275,4 +277,12 @@ public void setResponseQueueTimeSamples(Integer responseQueueTimeSamples) {
     public void setProtocolModule(Module m) {
         protocolModule = m;
     }
+
+    public Executor getAsyncExecutor() {
+        return asyncExecutor;
+    }
+
+    public void setAsyncExecutor(Executor asyncExecutor) {
+        this.asyncExecutor = asyncExecutor;
+    }
 }
diff --git a/driver/src/test/java/com/arangodb/ConsumerThreadAsyncTest.java b/driver/src/test/java/com/arangodb/ConsumerThreadAsyncTest.java
@@ -0,0 +1,118 @@
+package com.arangodb;
+
+import com.arangodb.config.ArangoConfigProperties;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class ConsumerThreadAsyncTest extends BaseJunit5 {
+
+    private volatile Thread thread;
+
+    private void setThread() {
+        thread = Thread.currentThread();
+    }
+
+    private void sleep() {
+        try {
+            Thread.sleep(3_000);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @ParameterizedTest
+    @EnumSource(Protocol.class)
+    @Disabled
+    void defaultConsumerThread(Protocol protocol) throws ExecutionException, InterruptedException {
+        ArangoDBAsync adb = new ArangoDB.Builder()
+                .loadProperties(ArangoConfigProperties.fromFile())
+                .protocol(protocol)
+                .build()
+                .async();
+
+        adb.getVersion()
+                .thenAccept(it -> setThread())
+                .get();
+
+        adb.shutdown();
+
+        if (Protocol.VST.equals(protocol)) {
+            assertThat(thread.getName()).startsWith("adb-vst-");
+        } else {
+            assertThat(thread.getName()).startsWith("adb-http-");
+        }
+    }
+
+    @ParameterizedTest
+    @EnumSource(Protocol.class)
+    void customConsumerExecutor(Protocol protocol) throws ExecutionException, InterruptedException {
+        ExecutorService es = Executors.newCachedThreadPool(r -> {
+            Thread t = Executors.defaultThreadFactory().newThread(r);
+            t.setName("custom-" + UUID.randomUUID());
+            return t;
+        });
+        ArangoDBAsync adb = new ArangoDB.Builder()
+                .loadProperties(ArangoConfigProperties.fromFile())
+                .protocol(protocol)
+                .asyncExecutor(es)
+                .build()
+                .async();
+
+        adb.getVersion()
+                .thenAccept(it -> setThread())
+                .get();
+
+        adb.shutdown();
+        es.shutdown();
+        assertThat(thread.getName()).startsWith("custom-");
+    }
+
+    /**
+     * Generates warns from Vert.x BlockedThreadChecker
+     */
+    @ParameterizedTest
+    @EnumSource(Protocol.class)
+    @Disabled
+    void sleepOnDefaultConsumerThread(Protocol protocol) throws ExecutionException, InterruptedException {
+        ArangoDBAsync adb = new ArangoDB.Builder()
+                .loadProperties(ArangoConfigProperties.fromFile())
+                .protocol(protocol)
+                .maxConnections(1)
+                .build()
+                .async();
+
+        adb.getVersion()
+                .thenAccept(it -> sleep())
+                .get();
+
+        adb.shutdown();
+    }
+
+    @ParameterizedTest
+    @EnumSource(Protocol.class)
+    void nestedRequests(Protocol protocol) throws ExecutionException, InterruptedException {
+        ArangoDBAsync adb = new ArangoDB.Builder()
+                .loadProperties(ArangoConfigProperties.fromFile())
+                .protocol(protocol)
+                .maxConnections(1)
+                .build()
+                .async();
+
+        adb.getVersion()
+                .thenCompose(it -> adb.getVersion())
+                .thenCompose(it -> adb.getVersion())
+                .thenCompose(it -> adb.getVersion())
+                .get();
+
+        adb.shutdown();
+    }
+
+}
diff --git a/http/src/main/java/com/arangodb/http/HttpConnection.java b/http/src/main/java/com/arangodb/http/HttpConnection.java
@@ -90,7 +90,7 @@ private static String getUserAgent() {
         timeout = config.getTimeout();
         vertx = Vertx.vertx(new VertxOptions().setPreferNativeTransport(true).setEventLoopPoolSize(1));
         vertx.runOnContext(e -> {
-            Thread.currentThread().setName("adb-eventloop-" + THREAD_COUNT.getAndIncrement());
+            Thread.currentThread().setName("adb-http-" + THREAD_COUNT.getAndIncrement());
             auth = new UsernamePasswordCredentials(
                     config.getUser(), Optional.ofNullable(config.getPassword()).orElse("")
             ).toHttpAuthorization();
diff --git a/vst/src/main/java/com/arangodb/vst/internal/VstConnection.java b/vst/src/main/java/com/arangodb/vst/internal/VstConnection.java
@@ -48,13 +48,15 @@
 import java.util.Date;
 import java.util.Map;
 import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * @author Mark Vollmary
  */
 public abstract class VstConnection<T> implements Connection {
     private static final Logger LOGGER = LoggerFactory.getLogger(VstConnection.class);
+    private static final AtomicInteger THREAD_COUNT = new AtomicInteger();
     private static final byte[] PROTOCOL_HEADER = "VST/1.0\r\n\r\n".getBytes();
     protected final MessageStore messageStore = new MessageStore();
     protected final Integer timeout;
@@ -171,7 +173,12 @@ public synchronized void open() throws IOException {
         }
         sendProtocolHeader();
 
-        executor = Executors.newSingleThreadExecutor();
+        executor = Executors.newSingleThreadExecutor(r -> {
+            Thread t = Executors.defaultThreadFactory().newThread(r);
+            t.setDaemon(true);
+            t.setName("adb-vst-" + THREAD_COUNT.getAndIncrement());
+            return t;
+        });
         executor.submit((Callable<Void>) () -> {
             LOGGER.debug("[" + connectionName + "]: Start Callable");
 
