From 7e378436b4fd4a5e956593ead09197074eec69dd Mon Sep 17 00:00:00 2001 From: Rishabh Maurya Date: Sun, 4 Jan 2026 13:52:46 -0800 Subject: [PATCH 1/4] Stream transport refactor * make FlightClientChannel sendMessage async and use of virtual threads * fix for headers when no batches are sent * set version from headers for incoming response deserialization Signed-off-by: Rishabh Maurya --- CHANGELOG.md | 1 + .../arrow/flight/bootstrap/ServerConfig.java | 3 +- .../flight/transport/FlightClientChannel.java | 72 ++--- .../transport/FlightOutboundHandler.java | 2 +- .../flight/transport/FlightServerChannel.java | 36 +-- .../flight/transport/FlightStreamPlugin.java | 3 +- .../flight/transport/FlightTransport.java | 7 +- .../transport/FlightTransportResponse.java | 255 ++++++------------ .../flight/transport/VectorStreamOutput.java | 51 ++-- .../ArrowStreamSerializationTests.java | 3 +- .../transport/FlightClientChannelTests.java | 5 +- .../transport/FlightTransportTestBase.java | 6 +- 12 files changed, 205 insertions(+), 239 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 75c21e6e6756b..51e725413c01f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -51,6 +51,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix Netty deprecation warnings in transport-reactor-netty4 module ([20429](https://github.com/opensearch-project/OpenSearch/pull/20429)) - Fix stats aggregation returning zero results with `size:0`. ([20427](https://github.com/opensearch-project/OpenSearch/pull/20427)) - Remove child level directory on refresh for CompositeIndexWriter ([#20326](https://github.com/opensearch-project/OpenSearch/pull/20326)) +- Fixes and refactoring in stream transport to make it more robust ([#20359](https://github.com/opensearch-project/OpenSearch/pull/20359)) ### Dependencies - Bump `com.google.auth:google-auth-library-oauth2-http` from 1.38.0 to 1.41.0 ([#20183](https://github.com/opensearch-project/OpenSearch/pull/20183)) diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/bootstrap/ServerConfig.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/bootstrap/ServerConfig.java index 2aadae9542845..2e8b687e6c7a6 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/bootstrap/ServerConfig.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/bootstrap/ServerConfig.java @@ -204,7 +204,8 @@ public static List> getSettings() { ARROW_ENABLE_DEBUG_ALLOCATOR, ARROW_ENABLE_UNSAFE_MEMORY_ACCESS, ARROW_SSL_ENABLE, - FLIGHT_EVENT_LOOP_THREADS + FLIGHT_EVENT_LOOP_THREADS, + FLIGHT_THREAD_POOL_MIN_SIZE ) ); } diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java index c20ace2d41b74..25b7738a65c8a 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java @@ -16,7 +16,6 @@ import org.opensearch.arrow.flight.stats.FlightCallTracker; import org.opensearch.arrow.flight.stats.FlightStatsCollector; import org.opensearch.cluster.node.DiscoveryNode; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; @@ -45,6 +44,7 @@ */ class FlightClientChannel implements TcpChannel { private static final Logger logger = LogManager.getLogger(FlightClientChannel.class); + private static final AtomicLong GLOBAL_CHANNEL_COUNTER = new AtomicLong(); private final AtomicLong correlationIdGenerator = new AtomicLong(); private final FlightClient client; private final DiscoveryNode node; @@ -112,6 +112,11 @@ public FlightClientChannel( this.closeListeners = new CopyOnWriteArrayList<>(); this.stats = new ChannelStats(); this.isClosed = false; + // Initialize with timestamp + global counter to ensure uniqueness with multiple channels + // Upper bits: timestamp, lower 20 bits: channel ID + long channelId = GLOBAL_CHANNEL_COUNTER.incrementAndGet() & 0xFFFFF; // 20 bits for channel ID + long initialValue = (System.currentTimeMillis() << 20) | channelId; + this.correlationIdGenerator.set(initialValue); if (statsCollector != null) { statsCollector.incrementClientChannelsActive(); } @@ -229,7 +234,8 @@ public void sendMessage(long requestId, BytesReference reference, ActionListener config ); - processStreamResponse(streamResponse); + // Open stream and prefetch first batch, invoke handler when ready + openStreamAndInvokeHandler(streamResponse); listener.onResponse(null); } catch (Exception e) { if (callTracker != null) { @@ -244,39 +250,45 @@ public void sendMessage(BytesReference reference, ActionListener listener) throw new IllegalStateException("sendMessage must be accompanied with requestId for FlightClientChannel, use the right variant."); } - private void processStreamResponse(FlightTransportResponse streamResponse) { - try { - executeWithThreadContext(streamResponse); - } catch (Exception e) { - handleStreamException(streamResponse, e); - } - } - @SuppressWarnings({ "unchecked", "rawtypes" }) - private void executeWithThreadContext(FlightTransportResponse streamResponse) { - final ThreadContext threadContext = threadPool.getThreadContext(); - final String executor = streamResponse.getHandler().executor(); + private void openStreamAndInvokeHandler(FlightTransportResponse streamResponse) { + TransportResponseHandler handler = streamResponse.getHandler(); + String executor = handler.executor(); + if (ThreadPool.Names.SAME.equals(executor)) { - executeHandler(threadContext, streamResponse); - } else { - threadPool.executor(executor).execute(() -> executeHandler(threadContext, streamResponse)); + logger.warn("Stream transport handler using SAME executor, which may cause blocking behavior"); } - } - @SuppressWarnings({ "unchecked", "rawtypes" }) - private void executeHandler(ThreadContext threadContext, FlightTransportResponse streamResponse) { - try (ThreadContext.StoredContext ignored = threadContext.stashContext()) { - Header header = streamResponse.getHeader(); - if (header == null) { - throw new StreamException(StreamErrorCode.INTERNAL, "Header is null"); + var threadContext = threadPool.getThreadContext(); + CompletableFuture
future = new CompletableFuture<>(); + streamResponse.openAndPrefetchAsync(future); + + future.whenComplete((header, error) -> { + if (error != null) { + handleStreamException(streamResponse, error instanceof Exception ? (Exception) error : new Exception(error)); + return; } - TransportResponseHandler handler = streamResponse.getHandler(); - threadContext.setHeaders(header.getHeaders()); - handler.handleStreamResponse(streamResponse); - } catch (Exception e) { - cleanupStreamResponse(streamResponse); - throw e; - } + + Runnable task = () -> { + try (var ignored = threadContext.stashContext()) { + if (header == null) { + cleanupStreamResponse(streamResponse); + throw new StreamException(StreamErrorCode.INTERNAL, "Header is null"); + } + threadContext.setHeaders(header.getHeaders()); + handler.handleStreamResponse(streamResponse); + } catch (Exception e) { + cleanupStreamResponse(streamResponse); + throw e; + } + }; + + if (ThreadPool.Names.SAME.equals(executor)) { + task.run(); + } else { + threadPool.executor(executor).execute(task); + } + }); } private void cleanupStreamResponse(StreamTransportResponse streamResponse) { diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java index b340b98492f01..a4a76d518bad2 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightOutboundHandler.java @@ -215,7 +215,7 @@ private void processCompleteTask(BatchTask task) { } try { - flightChannel.completeStream(); + flightChannel.completeStream(getHeaderBuffer(task.requestId(), task.nodeVersion(), task.features())); messageListener.onResponseSent(task.requestId(), task.action(), TransportResponse.Empty.INSTANCE); } catch (Exception e) { messageListener.onResponseSent(task.requestId(), task.action(), e); diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java index cef4ed0d99c92..ad14291437ad8 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java @@ -28,7 +28,6 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.Optional; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; @@ -49,7 +48,7 @@ class FlightServerChannel implements TcpChannel { private final InetSocketAddress remoteAddress; private final List> closeListeners = Collections.synchronizedList(new ArrayList<>()); private final ServerHeaderMiddleware middleware; - private volatile Optional root = Optional.empty(); + private volatile VectorSchemaRoot root = null; private final FlightCallTracker callTracker; private volatile boolean cancelled = false; private final ExecutorService executor; @@ -63,13 +62,10 @@ public FlightServerChannel( ) { this.serverStreamListener = serverStreamListener; this.serverStreamListener.setUseZeroCopy(true); - this.serverStreamListener.setOnCancelHandler(new Runnable() { - @Override - public void run() { - cancelled = true; - callTracker.recordCallEnd(StreamErrorCode.CANCELLED.name()); - close(); - } + this.serverStreamListener.setOnCancelHandler(() -> { + cancelled = true; + callTracker.recordCallEnd(StreamErrorCode.CANCELLED.name()); + close(); }); this.allocator = allocator; this.middleware = middleware; @@ -83,7 +79,7 @@ public BufferAllocator getAllocator() { return allocator; } - Optional getRoot() { + VectorSchemaRoot getRoot() { return root; } @@ -108,12 +104,12 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { } long batchStartTime = System.nanoTime(); // Only set for the first batch - if (root.isEmpty()) { + if (root == null) { middleware.setHeader(header); - root = Optional.of(output.getRoot()); - serverStreamListener.start(root.get()); + root = output.getRoot(); + serverStreamListener.start(root); } else { - root = Optional.of(output.getRoot()); + root = output.getRoot(); // placeholder to clear and fill the root with data for the next batch } @@ -121,7 +117,7 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { // its transmitted at transport; we close them all at complete stream. TODO: optimize this behaviour serverStreamListener.putNext(); if (callTracker != null) { - long rootSize = FlightUtils.calculateVectorSchemaRootSize(root.get()); + long rootSize = FlightUtils.calculateVectorSchemaRootSize(root); callTracker.recordBatchSent(rootSize, System.nanoTime() - batchStartTime); } } @@ -130,11 +126,15 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { * Completes the streaming response and closes all pending roots. * */ - public void completeStream() { + public void completeStream(ByteBuffer header) { try { if (!open.get()) { throw new IllegalStateException("FlightServerChannel already closed."); } + if (root == null) { + // Set header if no batches were sent + middleware.setHeader(header); + } serverStreamListener.completed(); } finally { callTracker.recordCallEnd(StreamErrorCode.OK.name()); @@ -210,7 +210,9 @@ public void close() { return; } open.set(false); - root.ifPresent(VectorSchemaRoot::close); + if (root != null) { + root.close(); + } notifyCloseListeners(); } diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightStreamPlugin.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightStreamPlugin.java index 10550d9730a40..dd7cb34a9db34 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightStreamPlugin.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightStreamPlugin.java @@ -364,7 +364,8 @@ public List> getSettings() { ServerComponents.SETTING_FLIGHT_PORTS, ServerComponents.SETTING_FLIGHT_HOST, ServerComponents.SETTING_FLIGHT_BIND_HOST, - ServerComponents.SETTING_FLIGHT_PUBLISH_HOST + ServerComponents.SETTING_FLIGHT_PUBLISH_HOST, + ServerComponents.SETTING_FLIGHT_PUBLISH_PORT ) ) { { diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java index 6ad73468dbe19..d42a93b425cb9 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java @@ -64,6 +64,7 @@ import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -132,7 +133,7 @@ public FlightTransport( this.sslContextProvider = sslContextProvider; this.statsCollector = statsCollector; this.bossEventLoopGroup = createEventLoopGroup("os-grpc-boss-ELG", 1); - this.workerEventLoopGroup = createEventLoopGroup("os-grpc-worker-ELG", Runtime.getRuntime().availableProcessors() * 2); + this.workerEventLoopGroup = createEventLoopGroup("os-grpc-worker-ELG", Runtime.getRuntime().availableProcessors()); this.serverExecutor = threadPool.executor(ServerConfig.GRPC_EXECUTOR_THREAD_POOL_NAME); this.clientExecutor = threadPool.executor(ServerConfig.FLIGHT_CLIENT_THREAD_POOL_NAME); this.threadPool = threadPool; @@ -409,7 +410,9 @@ protected InboundHandler createInboundHandler( } private EventLoopGroup createEventLoopGroup(String name, int threads) { - return new NioEventLoopGroup(threads); + AtomicInteger threadCounter = new AtomicInteger(0); + ThreadFactory threadFactory = r -> new Thread(r, name + "-" + threadCounter.incrementAndGet()); + return new NioEventLoopGroup(threads, threadFactory); } private void gracefullyShutdownELG(EventLoopGroup group, String name) { diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java index 048c1875b5679..011ac0e877e1d 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java @@ -27,40 +27,34 @@ import java.io.IOException; import java.util.Objects; +import java.util.concurrent.CompletableFuture; import static org.opensearch.arrow.flight.transport.ClientHeaderMiddleware.CORRELATION_ID_KEY; /** - * Arrow Flight implementation of streaming transport responses. - * - *

Handles streaming responses from Arrow Flight servers with lazy batch processing. - * Headers are extracted when first accessed, and responses are deserialized on demand. + * Streaming transport response implementation using Arrow Flight. + * Manages Flight stream lifecycle with lazy initialization and prefetching support. */ class FlightTransportResponse implements StreamTransportResponse { private static final Logger logger = LogManager.getLogger(FlightTransportResponse.class); - private final FlightStream flightStream; + private final FlightClient flightClient; + private final Ticket ticket; + private final FlightCallHeaders callHeaders; private final NamedWriteableRegistry namedWriteableRegistry; private final HeaderContext headerContext; - private final long correlationId; + private final TransportResponseHandler handler; private final FlightTransportConfig config; + private final long correlationId; - private final TransportResponseHandler handler; - private boolean isClosed; - - // Stream state - private VectorSchemaRoot currentRoot; - private Header currentHeader; - private boolean streamInitialized = false; - private boolean streamExhausted = false; - private boolean firstResponseConsumed = false; - private StreamException initializationException; - private long currentBatchSize; - - /** - * Creates a new Flight transport response. - */ - public FlightTransportResponse( + private volatile FlightStream flightStream; + private volatile long currentBatchSize; + private volatile boolean firstBatchConsumed; + private volatile boolean closed; + private volatile boolean prefetchStarted; + private volatile Header initialHeader; + + FlightTransportResponse( TransportResponseHandler handler, long correlationId, FlightClient flightClient, @@ -69,94 +63,94 @@ public FlightTransportResponse( NamedWriteableRegistry namedWriteableRegistry, FlightTransportConfig config ) { - this.handler = handler; + this.handler = Objects.requireNonNull(handler); this.correlationId = correlationId; - this.headerContext = Objects.requireNonNull(headerContext, "headerContext must not be null"); - this.namedWriteableRegistry = namedWriteableRegistry; - this.config = config; - // Initialize Flight stream with correlation ID header - FlightCallHeaders callHeaders = new FlightCallHeaders(); - callHeaders.insert(CORRELATION_ID_KEY, String.valueOf(correlationId)); - HeaderCallOption callOptions = new HeaderCallOption(callHeaders); - this.flightStream = flightClient.getStream(ticket, callOptions); - - this.isClosed = false; + this.flightClient = Objects.requireNonNull(flightClient); + this.headerContext = Objects.requireNonNull(headerContext); + this.ticket = Objects.requireNonNull(ticket); + this.namedWriteableRegistry = Objects.requireNonNull(namedWriteableRegistry); + this.config = Objects.requireNonNull(config); + this.callHeaders = new FlightCallHeaders(); + this.callHeaders.insert(CORRELATION_ID_KEY, String.valueOf(correlationId)); } - /** - * Gets the header for the current batch. - * If no batch has been fetched yet, fetches the first batch to extract headers. - */ - public Header getHeader() { - ensureOpen(); - initializeStreamIfNeeded(); - return currentHeader; + void openAndPrefetchAsync(CompletableFuture

future) { + if (prefetchStarted) return; + + synchronized (this) { + if (prefetchStarted) return; + if (closed) { + future.completeExceptionally(new StreamException(StreamErrorCode.UNAVAILABLE, "Stream is closed")); + return; + } + + prefetchStarted = true; + + Thread.ofVirtual().start(() -> { + try { + long start = System.nanoTime(); + flightStream = flightClient.getStream(ticket, new HeaderCallOption(callHeaders)); + if ((System.nanoTime() - start) / 1_000_000 > 10) { + logger.debug( + "FlightClient.getStream() for correlationId: {} took {}ms", + correlationId, + (System.nanoTime() - start) / 1_000_000 + ); + } + + flightStream.next(); + initialHeader = headerContext.getHeader(correlationId); + future.complete(initialHeader); + } catch (FlightRuntimeException e) { + future.completeExceptionally(FlightErrorMapper.fromFlightException(e)); + } catch (Exception e) { + future.completeExceptionally(new StreamException(StreamErrorCode.INTERNAL, "Stream open/prefetch failed", e)); + } + }); + } + } + + TransportResponseHandler getHandler() { + return handler; } - /** - * Gets the next response from the stream. - */ @Override public T nextResponse() { - ensureOpen(); - initializeStreamIfNeeded(); - - if (streamExhausted) { - if (initializationException != null) { - throw initializationException; - } - return null; - } + if (closed) throw new StreamException(StreamErrorCode.UNAVAILABLE, "Stream is closed"); + if (flightStream == null) throw new IllegalStateException("openAndPrefetch() must be called first"); long startTime = System.currentTimeMillis(); try { - if (!firstResponseConsumed) { - // First call - use the batch we already fetched during initialization - firstResponseConsumed = true; - return deserializeResponse(); - } - - if (flightStream.next()) { - currentRoot = flightStream.getRoot(); - currentHeader = headerContext.getHeader(correlationId); - // Capture the batch size before deserialization - currentBatchSize = FlightUtils.calculateVectorSchemaRootSize(currentRoot); - return deserializeResponse(); - } else { - streamExhausted = true; - return null; + boolean hasNext = firstBatchConsumed ? flightStream.next() : (firstBatchConsumed = true); + if (!hasNext) return null; + + VectorSchemaRoot root = flightStream.getRoot(); + currentBatchSize = FlightUtils.calculateVectorSchemaRootSize(root); + try (VectorStreamInput input = new VectorStreamInput(root, namedWriteableRegistry)) { + input.setVersion(initialHeader.getVersion()); + return handler.read(input); } } catch (FlightRuntimeException e) { - streamExhausted = true; throw FlightErrorMapper.fromFlightException(e); - } catch (Exception e) { - streamExhausted = true; - throw new StreamException(StreamErrorCode.INTERNAL, "Failed to fetch next batch", e); + } catch (IOException e) { + throw new StreamException(StreamErrorCode.INTERNAL, "Failed to deserialize batch", e); } finally { - logSlowOperation(startTime); + long took = System.currentTimeMillis() - startTime; + if (took > config.getSlowLogThreshold().millis()) { + logger.warn("Flight stream next() took [{}ms], exceeding threshold [{}ms]", took, config.getSlowLogThreshold().millis()); + } } } - /** - * Gets the size of the current batch in bytes. - * - * @return the size in bytes, or 0 if no batch is available - */ - public long getCurrentBatchSize() { + long getCurrentBatchSize() { return currentBatchSize; } - /** - * Cancels the Flight stream. - */ @Override public void cancel(String reason, Throwable cause) { - if (isClosed) { - return; - } + if (closed) return; try { - flightStream.cancel(reason, cause); - logger.debug("Cancelled flight stream: {}", reason); + if (flightStream != null) flightStream.cancel(reason, cause); } catch (Exception e) { logger.warn("Error cancelling flight stream", e); } finally { @@ -164,88 +158,17 @@ public void cancel(String reason, Throwable cause) { } } - /** - * Closes the Flight stream and releases resources. - */ @Override public void close() { - if (isClosed) { - return; - } - try { - if (currentRoot != null) { - currentRoot.close(); - currentRoot = null; + if (closed) return; + closed = true; + + if (flightStream != null) { + try { + flightStream.close(); + } catch (IllegalStateException ignore) {} catch (Exception e) { + throw new StreamException(StreamErrorCode.INTERNAL, "Error closing flight stream", e); } - flightStream.close(); - } catch (IllegalStateException ignore) { - // this is fine if the allocator is already closed - } catch (Exception e) { - throw new StreamException(StreamErrorCode.INTERNAL, "Error while closing flight stream", e); - } finally { - isClosed = true; - } - } - - public TransportResponseHandler getHandler() { - return handler; - } - - /** - * Initializes the stream by fetching the first batch to extract headers. - */ - private synchronized void initializeStreamIfNeeded() { - if (streamInitialized || streamExhausted) { - return; - } - long startTime = System.currentTimeMillis(); - try { - if (flightStream.next()) { - currentRoot = flightStream.getRoot(); - currentHeader = headerContext.getHeader(correlationId); - // Capture the batch size before deserialization - currentBatchSize = FlightUtils.calculateVectorSchemaRootSize(currentRoot); - streamInitialized = true; - } else { - streamExhausted = true; - } - } catch (FlightRuntimeException e) { - // TODO maybe add a check - handshake and validate if node is connected - // Try to get headers even if stream failed - currentHeader = headerContext.getHeader(correlationId); - streamExhausted = true; - initializationException = FlightErrorMapper.fromFlightException(e); - logger.warn("Stream initialization failed", e); - } catch (Exception e) { - // Try to get headers even if stream failed - currentHeader = headerContext.getHeader(correlationId); - streamExhausted = true; - initializationException = new StreamException(StreamErrorCode.INTERNAL, "Stream initialization failed", e); - logger.warn("Stream initialization failed", e); - } finally { - logSlowOperation(startTime); - } - } - - private T deserializeResponse() { - try (VectorStreamInput input = new VectorStreamInput(currentRoot, namedWriteableRegistry)) { - return handler.read(input); - } catch (IOException e) { - throw new StreamException(StreamErrorCode.INTERNAL, "Failed to deserialize response", e); - } - } - - private void ensureOpen() { - if (isClosed) { - throw new StreamException(StreamErrorCode.UNAVAILABLE, "Stream is closed"); - } - } - - private void logSlowOperation(long startTime) { - long took = System.currentTimeMillis() - startTime; - long thresholdMs = config.getSlowLogThreshold().millis(); - if (took > thresholdMs) { - logger.warn("Flight stream next() took [{}ms], exceeding threshold [{}ms]", took, thresholdMs); } } } diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java index 09e9e4a54c6c9..9755c4c77dbda 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/VectorStreamOutput.java @@ -18,45 +18,60 @@ import java.io.IOException; import java.util.List; -import java.util.Optional; class VectorStreamOutput extends StreamOutput { private int row = 0; private final VarBinaryVector vector; - private Optional root = Optional.empty(); + private VectorSchemaRoot root; + private final byte[] tempBuffer = new byte[8192]; + private int tempBufferPos = 0; - public VectorStreamOutput(BufferAllocator allocator, Optional root) { - if (root.isPresent()) { - vector = (VarBinaryVector) root.get().getVector(0); + public VectorStreamOutput(BufferAllocator allocator, VectorSchemaRoot root) { + if (root != null) { + vector = (VarBinaryVector) root.getVector(0); this.root = root; } else { Field field = new Field("0", new FieldType(true, new ArrowType.Binary(), null, null), null); vector = (VarBinaryVector) field.createVector(allocator); + // Pre-allocate with reasonable capacity to avoid repeated allocations + vector.setInitialCapacity(16); + vector.allocateNew(); } - vector.allocateNew(); } @Override - public void writeByte(byte b) throws IOException { - vector.setInitialCapacity(row + 1); - vector.setSafe(row++, new byte[] { b }); + public void writeByte(byte b) { + // Buffer small writes to reduce vector operations + if (tempBufferPos >= tempBuffer.length) { + flushTempBuffer(); + } + tempBuffer[tempBufferPos++] = b; } @Override - public void writeBytes(byte[] b, int offset, int length) throws IOException { - vector.setInitialCapacity(row + 1); + public void writeBytes(byte[] b, int offset, int length) { if (length == 0) { return; } if (b.length < (offset + length)) { throw new IllegalArgumentException("Illegal offset " + offset + "/length " + length + " for byte[] of length " + b.length); } + if (tempBufferPos > 0) { + flushTempBuffer(); + } vector.setSafe(row++, b, offset, length); } + private void flushTempBuffer() { + if (tempBufferPos > 0) { + vector.setSafe(row++, tempBuffer, 0, tempBufferPos); + tempBufferPos = 0; + } + } + @Override - public void flush() throws IOException { + public void flush() { } @@ -67,17 +82,19 @@ public void close() throws IOException { } @Override - public void reset() throws IOException { + public void reset() { row = 0; + tempBufferPos = 0; vector.clear(); } public VectorSchemaRoot getRoot() { + flushTempBuffer(); vector.setValueCount(row); - if (!root.isPresent()) { - root = Optional.of(new VectorSchemaRoot(List.of(vector))); + if (root == null) { + root = new VectorSchemaRoot(List.of(vector)); } - root.get().setRowCount(row); - return root.get(); + root.setRowCount(row); + return root; } } diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java index 843ddbcc1e385..16024496216da 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/ArrowStreamSerializationTests.java @@ -23,7 +23,6 @@ import java.io.IOException; import java.util.Arrays; import java.util.Collections; -import java.util.Optional; public class ArrowStreamSerializationTests extends OpenSearchTestCase { private NamedWriteableRegistry registry; @@ -51,7 +50,7 @@ public void tearDown() throws Exception { public void testInternalAggregationSerializationDeserialization() throws IOException { StringTerms original = createTestStringTerms(); - try (VectorStreamOutput output = new VectorStreamOutput(allocator, Optional.empty())) { + try (VectorStreamOutput output = new VectorStreamOutput(allocator, null)) { output.writeNamedWriteable(original); VectorSchemaRoot unifiedRoot = output.getRoot(); diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightClientChannelTests.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightClientChannelTests.java index ebdafdf0225bd..2d513c9bd9cf9 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightClientChannelTests.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightClientChannelTests.java @@ -555,7 +555,10 @@ public void handleStreamResponse(StreamTransportResponse streamRes } @Override - public void handleException(TransportException exp) {} + public void handleException(TransportException exp) { + handlerException.set(exp); + handlerLatch.countDown(); + } @Override public String executor() { diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java index a9a4d19f7e9a1..12a05d94ee64d 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java @@ -17,6 +17,7 @@ import org.opensearch.common.network.NetworkService; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.PageCacheRecycler; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; @@ -41,6 +42,7 @@ import java.util.Collections; import java.util.concurrent.atomic.AtomicInteger; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; @@ -103,7 +105,9 @@ public void setUp() throws Exception { ); flightTransport.start(); TransportService transportService = mock(TransportService.class); - when(transportService.getTaskManager()).thenReturn(mock(TaskManager.class)); + TaskManager taskManager = mock(TaskManager.class); + when(taskManager.taskExecutionStarted(any())).thenReturn(mock(ThreadContext.StoredContext.class)); + when(transportService.getTaskManager()).thenReturn(taskManager); streamTransportService = spy( new StreamTransportService( settings, From 537d14300f487e8fa43934bbfb0b6b0bc538c4f8 Mon Sep 17 00:00:00 2001 From: Rishabh Maurya Date: Fri, 23 Jan 2026 16:25:28 -0800 Subject: [PATCH 2/4] more debug logs and address PR comments Signed-off-by: Rishabh Maurya --- .../flight/transport/FlightServerChannel.java | 36 +++++++++++++++++-- .../flight/transport/FlightTransport.java | 5 +-- .../transport/FlightTransportResponse.java | 14 ++++---- .../transport/ServerHeaderMiddleware.java | 4 +++ 4 files changed, 47 insertions(+), 12 deletions(-) diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java index ad14291437ad8..69b222bbeacd2 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java @@ -15,6 +15,7 @@ import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.OpenSearchException; import org.opensearch.arrow.flight.stats.FlightCallTracker; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.bytes.BytesReference; @@ -30,6 +31,7 @@ import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import static org.opensearch.arrow.flight.transport.FlightErrorMapper.mapFromCallStatus; @@ -52,6 +54,8 @@ class FlightServerChannel implements TcpChannel { private final FlightCallTracker callTracker; private volatile boolean cancelled = false; private final ExecutorService executor; + private final long correlationId; + private AtomicInteger batchNumber = new AtomicInteger(0); public FlightServerChannel( ServerStreamListener serverStreamListener, @@ -60,6 +64,8 @@ public FlightServerChannel( FlightCallTracker callTracker, ExecutorService executor ) { + this.correlationId = Long.parseLong(middleware.getCorrelationId()); + logger.debug("Creating FlightServerChannel for correlation ID: {}", correlationId); this.serverStreamListener = serverStreamListener; this.serverStreamListener.setUseZeroCopy(true); this.serverStreamListener.setOnCancelHandler(() -> { @@ -102,6 +108,7 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { if (!open.get()) { throw new IllegalStateException("FlightServerChannel already closed."); } + batchNumber.incrementAndGet(); long batchStartTime = System.nanoTime(); // Only set for the first batch if (root == null) { @@ -112,13 +119,30 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { root = output.getRoot(); // placeholder to clear and fill the root with data for the next batch } - + logger.debug("Sending batch #{} for correlation ID: {}", batchNumber, correlationId); // we do not want to close the root right after putNext() call as we do not know the status of it whether // its transmitted at transport; we close them all at complete stream. TODO: optimize this behaviour serverStreamListener.putNext(); + long putNextTime = (System.nanoTime() - batchStartTime) / 1_000_000; if (callTracker != null) { long rootSize = FlightUtils.calculateVectorSchemaRootSize(root); callTracker.recordBatchSent(rootSize, System.nanoTime() - batchStartTime); + logger.debug( + "Batch #{} sent for correlation ID: {} in {}ms, size: {} bytes, putNext: {}ms", + batchNumber, + correlationId, + putNextTime / 1_000_000, + rootSize, + putNextTime + ); + } else { + logger.debug( + "Batch #{} sent for correlation ID: {} in {}ms, bytes, putNext: {}ms", + batchNumber, + correlationId, + putNextTime / 1_000_000, + putNextTime + ); } } @@ -134,6 +158,9 @@ public void completeStream(ByteBuffer header) { if (root == null) { // Set header if no batches were sent middleware.setHeader(header); + logger.debug("Completing empty stream for correlation ID: {}", correlationId); + } else { + logger.debug("Completing stream for correlation ID: {} after {} batches", correlationId, batchNumber); } serverStreamListener.completed(); } finally { @@ -160,8 +187,13 @@ public void sendError(ByteBuffer header, Exception error) { .toRuntimeException(); } middleware.setHeader(header); + if (error instanceof OpenSearchException) { + logger.debug("Error in Flight stream: {}", error.getMessage()); + } else { + logger.error("Unexpected error in Flight stream", error); + } + logger.debug("Sending error for correlation ID: {} after {} batches: {}", correlationId, batchNumber, error.getMessage()); serverStreamListener.error(flightExc); - logger.debug(error); } finally { StreamErrorCode errorCode = flightExc != null ? mapFromCallStatus(flightExc) : StreamErrorCode.UNKNOWN; callTracker.recordCallEnd(errorCode.name()); diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java index d42a93b425cb9..ac3dd44ff8731 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransport.java @@ -71,7 +71,8 @@ import java.util.stream.Collectors; import io.netty.channel.EventLoopGroup; -import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.MultiThreadIoEventLoopGroup; +import io.netty.channel.nio.NioIoHandler; import static org.opensearch.arrow.flight.bootstrap.ServerComponents.SETTING_FLIGHT_BIND_HOST; import static org.opensearch.arrow.flight.bootstrap.ServerComponents.SETTING_FLIGHT_PORTS; @@ -412,7 +413,7 @@ protected InboundHandler createInboundHandler( private EventLoopGroup createEventLoopGroup(String name, int threads) { AtomicInteger threadCounter = new AtomicInteger(0); ThreadFactory threadFactory = r -> new Thread(r, name + "-" + threadCounter.incrementAndGet()); - return new NioEventLoopGroup(threads, threadFactory); + return new MultiThreadIoEventLoopGroup(threads, threadFactory, NioIoHandler.newFactory()); } private void gracefullyShutdownELG(EventLoopGroup group, String name) { diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java index 011ac0e877e1d..1047faab274b3 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightTransportResponse.java @@ -90,15 +90,12 @@ void openAndPrefetchAsync(CompletableFuture
future) { try { long start = System.nanoTime(); flightStream = flightClient.getStream(ticket, new HeaderCallOption(callHeaders)); - if ((System.nanoTime() - start) / 1_000_000 > 10) { - logger.debug( - "FlightClient.getStream() for correlationId: {} took {}ms", - correlationId, - (System.nanoTime() - start) / 1_000_000 - ); - } - + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + logger.debug("FlightClient.getStream() for correlationId: {} took {}ms", correlationId, elapsedMs); + start = System.nanoTime(); flightStream.next(); + elapsedMs = (System.nanoTime() - start) / 1_000_000; + logger.debug("First FlightClient.next() for correlationId: {} took {}ms", correlationId, elapsedMs); initialHeader = headerContext.getHeader(correlationId); future.complete(initialHeader); } catch (FlightRuntimeException e) { @@ -139,6 +136,7 @@ public T nextResponse() { if (took > config.getSlowLogThreshold().millis()) { logger.warn("Flight stream next() took [{}ms], exceeding threshold [{}ms]", took, config.getSlowLogThreshold().millis()); } + logger.debug("FlightClient.next() for correlationId: {} took {}ms", correlationId, took); } } diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ServerHeaderMiddleware.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ServerHeaderMiddleware.java index b8e96b16c9c35..3d8f46ae04370 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ServerHeaderMiddleware.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/ServerHeaderMiddleware.java @@ -37,6 +37,10 @@ void setHeader(ByteBuffer headerBuffer) { this.headerBuffer = headerBuffer; } + String getCorrelationId() { + return requestId; + } + @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { if (headerBuffer != null) { From b25be8be5296d7d670af75ada016ccd020d44fa3 Mon Sep 17 00:00:00 2001 From: Rishabh Maurya Date: Tue, 27 Jan 2026 13:03:38 -0800 Subject: [PATCH 3/4] Address PR comments Signed-off-by: Rishabh Maurya --- .../arrow/flight/transport/FlightClientChannel.java | 2 +- .../arrow/flight/transport/FlightServerChannel.java | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java index 25b7738a65c8a..2ad49645ef58d 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java @@ -256,7 +256,7 @@ private void openStreamAndInvokeHandler(FlightTransportResponse streamRespons String executor = handler.executor(); if (ThreadPool.Names.SAME.equals(executor)) { - logger.warn("Stream transport handler using SAME executor, which may cause blocking behavior"); + logger.debug("Stream transport handler using SAME executor, which may cause blocking behavior"); } var threadContext = threadPool.getThreadContext(); diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java index 69b222bbeacd2..cd05d53c1a380 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java @@ -131,7 +131,7 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { "Batch #{} sent for correlation ID: {} in {}ms, size: {} bytes, putNext: {}ms", batchNumber, correlationId, - putNextTime / 1_000_000, + putNextTime, rootSize, putNextTime ); @@ -140,7 +140,7 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { "Batch #{} sent for correlation ID: {} in {}ms, bytes, putNext: {}ms", batchNumber, correlationId, - putNextTime / 1_000_000, + putNextTime, putNextTime ); } From 35edeec1279ca97f025797b0d9bae8a2e908e7fe Mon Sep 17 00:00:00 2001 From: Rishabh Maurya Date: Tue, 27 Jan 2026 16:47:54 -0800 Subject: [PATCH 4/4] Address PR comments Signed-off-by: Rishabh Maurya --- .../arrow/flight/transport/FlightClientChannel.java | 3 +-- .../arrow/flight/transport/FlightServerChannel.java | 13 +++---------- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java index 2ad49645ef58d..d9816adbe14f4 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightClientChannel.java @@ -272,8 +272,7 @@ private void openStreamAndInvokeHandler(FlightTransportResponse streamRespons Runnable task = () -> { try (var ignored = threadContext.stashContext()) { if (header == null) { - cleanupStreamResponse(streamResponse); - throw new StreamException(StreamErrorCode.INTERNAL, "Header is null"); + handleStreamException(streamResponse, new StreamException(StreamErrorCode.INTERNAL, "Header is null")); } threadContext.setHeaders(header.getHeaders()); handler.handleStreamResponse(streamResponse); diff --git a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java index cd05d53c1a380..a923c7843d987 100644 --- a/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java +++ b/plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/transport/FlightServerChannel.java @@ -55,7 +55,7 @@ class FlightServerChannel implements TcpChannel { private volatile boolean cancelled = false; private final ExecutorService executor; private final long correlationId; - private AtomicInteger batchNumber = new AtomicInteger(0); + private final AtomicInteger batchNumber = new AtomicInteger(0); public FlightServerChannel( ServerStreamListener serverStreamListener, @@ -128,21 +128,14 @@ public void sendBatch(ByteBuffer header, VectorStreamOutput output) { long rootSize = FlightUtils.calculateVectorSchemaRootSize(root); callTracker.recordBatchSent(rootSize, System.nanoTime() - batchStartTime); logger.debug( - "Batch #{} sent for correlation ID: {} in {}ms, size: {} bytes, putNext: {}ms", + "Batch #{} sent for correlation ID: {}, size: {} bytes, putNext: {}ms", batchNumber, correlationId, - putNextTime, rootSize, putNextTime ); } else { - logger.debug( - "Batch #{} sent for correlation ID: {} in {}ms, bytes, putNext: {}ms", - batchNumber, - correlationId, - putNextTime, - putNextTime - ); + logger.debug("Batch #{} sent for correlation ID: {}, putNext: {}ms", batchNumber, correlationId, putNextTime); } }