Async flush of batches at server for Stream Transport

Signed-off-by: Rishabh Maurya <[email protected]>

Author: rishabhmaurya

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/bootstrap/ServerConfig.java

@@ -30,6 +30,7 @@
 import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.channel.socket.nio.NioServerSocketChannel;
 import io.netty.channel.socket.nio.NioSocketChannel;
+import io.netty.util.NettyRuntime;
 
 /**
  * Configuration class for OpenSearch Flight server settings.
@@ -87,6 +88,13 @@ public ServerConfig() {}
         Setting.Property.NodeScope
     );
 
+    static final Setting<Integer> FLIGHT_EVENT_LOOP_THREADS = Setting.intSetting(
+        "flight.event_loop.threads",
+        Math.max(1, NettyRuntime.availableProcessors() * 2),
+        1,
+        Setting.Property.NodeScope
+    );
+
     static final Setting<Boolean> ARROW_SSL_ENABLE = Setting.boolSetting(
         "flight.ssl.enable",
         false, // TODO: get default from security enabled
@@ -112,6 +120,7 @@ public ServerConfig() {}
     private static int threadPoolMin;
     private static int threadPoolMax;
     private static TimeValue keepAlive;
+    private static int eventLoopThreads;
 
     /**
      * Initializes the server configuration with the provided settings.
@@ -134,6 +143,7 @@ public static void init(Settings settings) {
         threadPoolMin = FLIGHT_THREAD_POOL_MIN_SIZE.get(settings);
         threadPoolMax = FLIGHT_THREAD_POOL_MAX_SIZE.get(settings);
         keepAlive = FLIGHT_THREAD_POOL_KEEP_ALIVE.get(settings);
+        eventLoopThreads = FLIGHT_EVENT_LOOP_THREADS.get(settings);
     }
 
     /**
@@ -172,6 +182,15 @@ public static ScalingExecutorBuilder getClientExecutorBuilder() {
         return new ScalingExecutorBuilder(FLIGHT_CLIENT_THREAD_POOL_NAME, threadPoolMin, threadPoolMax, keepAlive);
     }
 
+    /**
+     * Gets the configured number of event loop threads.
+     *
+     * @return The number of event loop threads
+     */
+    public static int getEventLoopThreads() {
+        return eventLoopThreads;
+    }
+
     /**
      * Returns a list of all settings managed by this configuration class.
      *
@@ -184,7 +203,8 @@ public static List<Setting<?>> getSettings() {
                 ARROW_ENABLE_NULL_CHECK_FOR_GET,
                 ARROW_ENABLE_DEBUG_ALLOCATOR,
                 ARROW_ENABLE_UNSAFE_MEMORY_ACCESS,
-                ARROW_SSL_ENABLE
+                ARROW_SSL_ENABLE,
+                FLIGHT_EVENT_LOOP_THREADS
             )
         );
     }

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ArrowFlightProducer.java

@@ -61,7 +61,13 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
         // https://github.com/apache/arrow/issues/38668
         executor.execute(() -> {
             FlightCallTracker callTracker = statsCollector.createServerCallTracker();
-            FlightServerChannel channel = new FlightServerChannel(listener, allocator, middleware, callTracker);
+            FlightServerChannel channel = new FlightServerChannel(
+                listener,
+                allocator,
+                middleware,
+                callTracker,
+                flightTransport.getNextFlightExecutor()
+            );
             try {
                 BytesArray buf = new BytesArray(ticket.getBytes());
                 callTracker.recordRequestBytes(buf.ramBytesUsed());

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/BatchTask.java

@@ -0,0 +1,135 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.arrow.flight.transport;
+
+import org.opensearch.Version;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.transport.TransportResponse;
+import org.opensearch.transport.TcpChannel;
+
+import java.util.Set;
+
+/**
+ * Represents a batch processing task for Flight streaming responses.
+ * @opensearch.internal
+ */
+class BatchTask implements AutoCloseable {
+    final Version nodeVersion;
+    final Set<String> features;
+    final TcpChannel channel;
+    final FlightTransportChannel transportChannel;
+    final long requestId;
+    final String action;
+    final TransportResponse response;
+    final boolean compress;
+    final boolean isHandshake;
+    final boolean isComplete;
+    final boolean isError;
+    final Exception error;
+    final ThreadContext.StoredContext storedContext;
+
+    BatchTask(
+        Version nodeVersion,
+        Set<String> features,
+        TcpChannel channel,
+        FlightTransportChannel transportChannel,
+        long requestId,
+        String action,
+        TransportResponse response,
+        boolean compress,
+        boolean isHandshake,
+        boolean isComplete,
+        ThreadContext.StoredContext storedContext
+    ) {
+        this(
+            nodeVersion,
+            features,
+            channel,
+            transportChannel,
+            requestId,
+            action,
+            response,
+            compress,
+            isHandshake,
+            isComplete,
+            false,
+            null,
+            storedContext
+        );
+    }
+
+    BatchTask(
+        Version nodeVersion,
+        Set<String> features,
+        TcpChannel channel,
+        FlightTransportChannel transportChannel,
+        long requestId,
+        String action,
+        boolean compress,
+        boolean isHandshake,
+        Exception error,
+        ThreadContext.StoredContext storedContext
+    ) {
+        this(
+            nodeVersion,
+            features,
+            channel,
+            transportChannel,
+            requestId,
+            action,
+            null,
+            compress,
+            isHandshake,
+            false,
+            true,
+            error,
+            storedContext
+        );
+    }
+
+    private BatchTask(
+        Version nodeVersion,
+        Set<String> features,
+        TcpChannel channel,
+        FlightTransportChannel transportChannel,
+        long requestId,
+        String action,
+        TransportResponse response,
+        boolean compress,
+        boolean isHandshake,
+        boolean isComplete,
+        boolean isError,
+        Exception error,
+        ThreadContext.StoredContext storedContext
+    ) {
+        this.nodeVersion = nodeVersion;
+        this.features = features;
+        this.channel = channel;
+        this.transportChannel = transportChannel;
+        this.requestId = requestId;
+        this.action = action;
+        this.response = response;
+        this.compress = compress;
+        this.isHandshake = isHandshake;
+        this.isComplete = isComplete;
+        this.isError = isError;
+        this.error = error;
+        this.storedContext = storedContext;
+    }
+
+    @Override
+    public void close() {
+        if (storedContext != null) {
+            storedContext.close();
+        }
+        if ((isComplete || isError) && transportChannel != null) {
+            transportChannel.releaseChannel(isError);
+        }
+    }
+}