Asynchronous FlightServerChannel to avoid thread contention per request

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/FlightOutboundHandler.java

@@ -20,6 +20,7 @@
 import org.opensearch.Version;
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.transport.TransportResponse;
 import org.opensearch.threadpool.ThreadPool;
@@ -89,85 +90,195 @@ public void sendResponse(
         );
     }
 
-    /** This needs to be synchronized for the cases when multiple batches are written concurrently,
-     * as VectorSchemaRoot is shared across batches **/
-    public synchronized void sendResponseBatch(
+    @Override
+    public void sendErrorResponse(
+        Version nodeVersion,
+        Set<String> features,
+        TcpChannel channel,
+        long requestId,
+        String action,
+        Exception error
+    ) throws IOException {
+        throw new UnsupportedOperationException(
+            "sendResponse() is not supported for streaming requests in FlightOutboundHandler; use sendResponseBatch()"
+        );
+    }
+
+    public void sendResponseBatch(
         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
     ) throws IOException {
-        // TODO add support for compression
+        ThreadContext.StoredContext storedContext = threadPool.getThreadContext().stashContext();
+        BatchTask task = new BatchTask(
+            nodeVersion,
+            features,
+            channel,
+            transportChannel,
+            requestId,
+            action,
+            response,
+            compress,
+            isHandshake,
+            false,
+            false,
+            null,
+            storedContext
+        );
+
         if (!(channel instanceof FlightServerChannel flightChannel)) {
-            throw new IllegalStateException("Expected FlightServerChannel, got " + channel.getClass().getName());
+            messageListener.onResponseSent(requestId, action, new IllegalStateException("Expected FlightServerChannel"));
+            return;
+        }
+
+        flightChannel.getExecutor().execute(() -> {
+            try (BatchTask ignored = task) {
+                processBatchTask(task);
+            } catch (Exception e) {
+                messageListener.onResponseSent(requestId, action, e);
+            }
+        });
+    }
+
+    private void processBatchTask(BatchTask task) {
+        task.storedContext().restore();
+        if (!(task.channel() instanceof FlightServerChannel flightChannel)) {
+            Exception error = new IllegalStateException("Expected FlightServerChannel, got " + task.channel().getClass().getName());
+            messageListener.onResponseSent(task.requestId(), task.action(), error);
+            return;
         }
+
         try {
             try (VectorStreamOutput out = new VectorStreamOutput(flightChannel.getAllocator(), flightChannel.getRoot())) {
-                response.writeTo(out);
-                flightChannel.sendBatch(getHeaderBuffer(requestId, nodeVersion, features), out);
-                messageListener.onResponseSent(requestId, action, response);
+                task.response().writeTo(out);
+                flightChannel.sendBatch(getHeaderBuffer(task.requestId(), task.nodeVersion(), task.features()), out);
+                messageListener.onResponseSent(task.requestId(), task.action(), task.response());
             }
-        } catch (StreamException e) {
-            messageListener.onResponseSent(requestId, action, e);
-            // Let StreamException propagate as is - it will be converted to FlightRuntimeException at a higher level
-            throw e;
         } catch (FlightRuntimeException e) {
-            messageListener.onResponseSent(requestId, action, e);
-            throw FlightErrorMapper.fromFlightException(e);
+            messageListener.onResponseSent(task.requestId(), task.action(), FlightErrorMapper.fromFlightException(e));
         } catch (Exception e) {
-            messageListener.onResponseSent(requestId, action, e);
-            throw e;
+            messageListener.onResponseSent(task.requestId(), task.action(), e);
         }
     }
 
     public void completeStream(
         final Version nodeVersion,
         final Set<String> features,
         final TcpChannel channel,
+        final FlightTransportChannel transportChannel,
         final long requestId,
         final String action
     ) {
+        ThreadContext.StoredContext storedContext = threadPool.getThreadContext().stashContext();
+        BatchTask completeTask = new BatchTask(
+            nodeVersion,
+            features,
+            channel,
+            transportChannel,
+            requestId,
+            action,
+            TransportResponse.Empty.INSTANCE,
+            false,
+            false,
+            true,
+            false,
+            null,
+            storedContext
+        );
+
         if (!(channel instanceof FlightServerChannel flightChannel)) {
-            throw new IllegalStateException("Expected FlightServerChannel, got " + channel.getClass().getName());
+            messageListener.onResponseSent(requestId, action, new IllegalStateException("Expected FlightServerChannel"));
+            return;
         }
+
+        flightChannel.getExecutor().execute(() -> {
+            try (BatchTask ignored = completeTask) {
+                processCompleteTask(completeTask);
+            } catch (Exception e) {
+                messageListener.onResponseSent(requestId, action, e);
+            }
+        });
+    }
+
+    private void processCompleteTask(BatchTask task) {
+        task.storedContext().restore();
+        if (!(task.channel() instanceof FlightServerChannel flightChannel)) {
+            Exception error = new IllegalStateException("Expected FlightServerChannel, got " + task.channel().getClass().getName());
+            messageListener.onResponseSent(task.requestId(), task.action(), error);
+            return;
+        }
+
         try {
             flightChannel.completeStream();
-            messageListener.onResponseSent(requestId, action, TransportResponse.Empty.INSTANCE);
-        } catch (FlightRuntimeException e) {
-            messageListener.onResponseSent(requestId, action, e);
-            throw FlightErrorMapper.fromFlightException(e);
+            messageListener.onResponseSent(task.requestId(), task.action(), TransportResponse.Empty.INSTANCE);
         } catch (Exception e) {
-            messageListener.onResponseSent(requestId, action, e);
-            throw e;
+            messageListener.onResponseSent(task.requestId(), task.action(), e);
         }
     }
 
-    @Override
     public void sendErrorResponse(
         final Version nodeVersion,
         final Set<String> features,
         final TcpChannel channel,
+        final FlightTransportChannel transportChannel,
         final long requestId,
         final String action,
         final Exception error
-    ) throws IOException {
-        if (!(channel instanceof FlightServerChannel flightServerChannel)) {
-            throw new IllegalStateException("Expected FlightServerChannel, got " + channel.getClass().getName());
+    ) {
+        ThreadContext.StoredContext storedContext = threadPool.getThreadContext().stashContext();
+        BatchTask errorTask = new BatchTask(
+            nodeVersion,
+            features,
+            channel,
+            transportChannel,
+            requestId,
+            action,
+            null,
+            false,
+            false,
+            false,
+            true,
+            error,
+            storedContext
+        );
+
+        if (!(channel instanceof FlightServerChannel flightChannel)) {
+            messageListener.onResponseSent(requestId, action, new IllegalStateException("Expected FlightServerChannel"));
+            return;
+        }
+
+        flightChannel.getExecutor().execute(() -> {
+            try (BatchTask ignored = errorTask) {
+                processErrorTask(errorTask);
+            } catch (Exception e) {
+                messageListener.onResponseSent(requestId, action, e);
+            }
+        });
+    }
+
+    private void processErrorTask(BatchTask task) {
+        task.storedContext().restore();
+        if (!(task.channel() instanceof FlightServerChannel flightServerChannel)) {
+            Exception error = new IllegalStateException("Expected FlightServerChannel, got " + task.channel().getClass().getName());
+            messageListener.onResponseSent(task.requestId(), task.action(), error);
+            return;
         }
+
         try {
-            Exception flightError = error;
-            if (error instanceof StreamException) {
-                flightError = FlightErrorMapper.toFlightException((StreamException) error);
+            Exception flightError = task.error();
+            if (task.error() instanceof StreamException) {
+                flightError = FlightErrorMapper.toFlightException((StreamException) task.error());
             }
-            flightServerChannel.sendError(getHeaderBuffer(requestId, version, features), flightError);
-            messageListener.onResponseSent(requestId, action, error);
+            flightServerChannel.sendError(getHeaderBuffer(task.requestId(), task.nodeVersion(), task.features()), flightError);
+            messageListener.onResponseSent(task.requestId(), task.action(), task.error());
         } catch (Exception e) {
-            messageListener.onResponseSent(requestId, action, e);
-            throw e;
+            messageListener.onResponseSent(task.requestId(), task.action(), e);
         }
     }
 
@@ -197,4 +308,19 @@ private ByteBuffer getHeaderBuffer(long requestId, Version nodeVersion, Set<Stri
             return ByteBuffer.wrap(headerBytes.toBytesRef().bytes);
         }
     }
+
+    record 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) implements AutoCloseable {
+
+        @Override
+        public void close() {
+            if (storedContext != null) {
+                storedContext.close();
+            }
+            if ((isComplete || isError) && transportChannel != null) {
+                transportChannel.releaseChannel(isError);
+            }
+        }
+    }
 }