Address PR comment

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

Author: rishabhmaurya

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

@@ -34,7 +34,7 @@
  */
 class ClientHeaderMiddleware implements FlightClientMiddleware {
     static final String RAW_HEADER_KEY = "raw-header";
-    static final String REQUEST_ID_KEY = "req-id";
+    static final String CORRELATION_ID_KEY = "correlation-id";
 
     private final HeaderContext context;
     private final Version version;
@@ -60,13 +60,13 @@ class ClientHeaderMiddleware implements FlightClientMiddleware {
     @Override
     public void onHeadersReceived(CallHeaders incomingHeaders) {
         String encodedHeader = incomingHeaders.get(RAW_HEADER_KEY);
-        String reqId = incomingHeaders.get(REQUEST_ID_KEY);
+        String correlationId = incomingHeaders.get(CORRELATION_ID_KEY);
 
         if (encodedHeader == null) {
             throw new StreamException(StreamErrorCode.INVALID_ARGUMENT, "Missing required header: " + RAW_HEADER_KEY);
         }
-        if (reqId == null) {
-            throw new StreamException(StreamErrorCode.INVALID_ARGUMENT, "Missing required header: " + REQUEST_ID_KEY);
+        if (correlationId == null) {
+            throw new StreamException(StreamErrorCode.INVALID_ARGUMENT, "Missing required header: " + CORRELATION_ID_KEY);
         }
 
         try {
@@ -87,12 +87,11 @@ public void onHeadersReceived(CallHeaders incomingHeaders) {
             }
 
             // Store the header in context for later retrieval
-            long requestId = Long.parseLong(reqId);
-            context.setHeader(requestId, header);
+            context.setHeader(Long.parseLong(correlationId), header);
         } catch (IOException e) {
             throw new StreamException(StreamErrorCode.INTERNAL, "Failed to decode header", e);
         } catch (NumberFormatException e) {
-            throw new StreamException(StreamErrorCode.INVALID_ARGUMENT, "Invalid request ID format: " + reqId, e);
+            throw new StreamException(StreamErrorCode.INVALID_ARGUMENT, "Invalid request ID format: " + correlationId, e);
         }
     }
 

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

@@ -47,7 +47,7 @@
  */
 class FlightClientChannel implements TcpChannel {
     private static final Logger logger = LogManager.getLogger(FlightClientChannel.class);
-    private final AtomicLong requestIdGenerator = new AtomicLong();
+    private final AtomicLong correlationIdGenerator = new AtomicLong();
     private final FlightClient client;
     private final DiscoveryNode node;
     private final BoundTransportAddress boundAddress;
@@ -199,7 +199,7 @@ public InetSocketAddress getRemoteAddress() {
     }
 
     @Override
-    public void sendMessage(long reqId, BytesReference reference, ActionListener<Void> listener) {
+    public void sendMessage(long requestId, BytesReference reference, ActionListener<Void> listener) {
         if (!isOpen()) {
             listener.onFailure(new StreamException(StreamErrorCode.UNAVAILABLE, "FlightClientChannel is closed"));
             return;
@@ -214,8 +214,8 @@ public void sendMessage(long reqId, BytesReference reference, ActionListener<Voi
         try {
             // ticket will contain the serialized headers
             Ticket ticket = serializeToTicket(reference);
-            TransportResponseHandler<?> handler = responseHandlers.onResponseReceived(reqId, messageListener);
-            long correlationId = requestIdGenerator.incrementAndGet();
+            TransportResponseHandler<?> handler = responseHandlers.onResponseReceived(requestId, messageListener);
+            long correlationId = correlationIdGenerator.incrementAndGet();
 
             if (callTracker != null) {
                 handler = new MetricsTrackingResponseHandler<>(handler, callTracker);
@@ -243,7 +243,7 @@ public void sendMessage(long reqId, BytesReference reference, ActionListener<Voi
 
     @Override
     public void sendMessage(BytesReference reference, ActionListener<Void> listener) {
-        throw new IllegalStateException("sendMessage must be accompanied with reqId for FlightClientChannel, use the right variant.");
+        throw new IllegalStateException("sendMessage must be accompanied with requestId for FlightClientChannel, use the right variant.");
     }
 
     private void processStreamResponse(FlightTransportResponse<?> streamResponse) {

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

@@ -28,7 +28,7 @@
 import java.io.IOException;
 import java.util.Objects;
 
-import static org.opensearch.arrow.flight.transport.ClientHeaderMiddleware.REQUEST_ID_KEY;
+import static org.opensearch.arrow.flight.transport.ClientHeaderMiddleware.CORRELATION_ID_KEY;
 
 /**
  * Arrow Flight implementation of streaming transport responses.
@@ -42,7 +42,7 @@ class FlightTransportResponse<T extends TransportResponse> implements StreamTran
     private final FlightStream flightStream;
     private final NamedWriteableRegistry namedWriteableRegistry;
     private final HeaderContext headerContext;
-    private final long reqId;
+    private final long correlationId;
     private final FlightTransportConfig config;
 
     private final TransportResponseHandler<T> handler;
@@ -62,21 +62,21 @@ class FlightTransportResponse<T extends TransportResponse> implements StreamTran
      */
     public FlightTransportResponse(
         TransportResponseHandler<T> handler,
-        long reqId,
+        long correlationId,
         FlightClient flightClient,
         HeaderContext headerContext,
         Ticket ticket,
         NamedWriteableRegistry namedWriteableRegistry,
         FlightTransportConfig config
     ) {
         this.handler = handler;
-        this.reqId = reqId;
+        this.correlationId = correlationId;
         this.headerContext = Objects.requireNonNull(headerContext, "headerContext must not be null");
         this.namedWriteableRegistry = namedWriteableRegistry;
         this.config = config;
-        // Initialize Flight stream with request ID header
+        // Initialize Flight stream with correlation ID header
         FlightCallHeaders callHeaders = new FlightCallHeaders();
-        callHeaders.insert(REQUEST_ID_KEY, String.valueOf(reqId));
+        callHeaders.insert(CORRELATION_ID_KEY, String.valueOf(correlationId));
         HeaderCallOption callOptions = new HeaderCallOption(callHeaders);
         this.flightStream = flightClient.getStream(ticket, callOptions);
 
@@ -118,7 +118,7 @@ public T nextResponse() {
 
             if (flightStream.next()) {
                 currentRoot = flightStream.getRoot();
-                currentHeader = headerContext.getHeader(reqId);
+                currentHeader = headerContext.getHeader(correlationId);
                 // Capture the batch size before deserialization
                 currentBatchSize = FlightUtils.calculateVectorSchemaRootSize(currentRoot);
                 return deserializeResponse();
@@ -202,7 +202,7 @@ private synchronized void initializeStreamIfNeeded() {
         try {
             if (flightStream.next()) {
                 currentRoot = flightStream.getRoot();
-                currentHeader = headerContext.getHeader(reqId);
+                currentHeader = headerContext.getHeader(correlationId);
                 // Capture the batch size before deserialization
                 currentBatchSize = FlightUtils.calculateVectorSchemaRootSize(currentRoot);
                 streamInitialized = true;
@@ -212,13 +212,13 @@ private synchronized void initializeStreamIfNeeded() {
         } 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(reqId);
+            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(reqId);
+            currentHeader = headerContext.getHeader(correlationId);
             streamExhausted = true;
             initializationException = new StreamException(StreamErrorCode.INTERNAL, "Stream initialization failed", e);
             logger.warn("Stream initialization failed", e);

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

@@ -15,11 +15,11 @@
 class HeaderContext {
     private final ConcurrentHashMap<Long, Header> headerMap = new ConcurrentHashMap<>();
 
-    void setHeader(long reqId, Header header) {
-        headerMap.put(reqId, header);
+    void setHeader(long correlationId, Header header) {
+        headerMap.put(correlationId, header);
     }
 
-    Header getHeader(long reqId) {
-        return headerMap.remove(reqId);
+    Header getHeader(long correlationId) {
+        return headerMap.remove(correlationId);
     }
 }

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

@@ -17,8 +17,8 @@
 import java.nio.ByteBuffer;
 import java.util.Base64;
 
+import static org.opensearch.arrow.flight.transport.ClientHeaderMiddleware.CORRELATION_ID_KEY;
 import static org.opensearch.arrow.flight.transport.ClientHeaderMiddleware.RAW_HEADER_KEY;
-import static org.opensearch.arrow.flight.transport.ClientHeaderMiddleware.REQUEST_ID_KEY;
 
 /**
  * ServerHeaderMiddleware is created per call to handle the response header
@@ -27,10 +27,10 @@
  */
 class ServerHeaderMiddleware implements FlightServerMiddleware {
     private ByteBuffer headerBuffer;
-    private final String reqId;
+    private final String requestId;
 
-    ServerHeaderMiddleware(String reqId) {
-        this.reqId = reqId;
+    ServerHeaderMiddleware(String requestId) {
+        this.requestId = requestId;
     }
 
     void setHeader(ByteBuffer headerBuffer) {
@@ -44,11 +44,11 @@ public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
             headerBuffer.get(headerBytes);
             String encodedHeader = Base64.getEncoder().encodeToString(headerBytes);
             outgoingHeaders.insert(RAW_HEADER_KEY, encodedHeader);
-            outgoingHeaders.insert(REQUEST_ID_KEY, reqId);
+            outgoingHeaders.insert(CORRELATION_ID_KEY, requestId);
             headerBuffer.rewind();
         } else {
             outgoingHeaders.insert(RAW_HEADER_KEY, "");
-            outgoingHeaders.insert(REQUEST_ID_KEY, reqId);
+            outgoingHeaders.insert(CORRELATION_ID_KEY, requestId);
         }
     }
 
@@ -61,8 +61,8 @@ public void onCallErrored(Throwable err) {}
     public static class Factory implements FlightServerMiddleware.Factory<ServerHeaderMiddleware> {
         @Override
         public ServerHeaderMiddleware onCallStarted(CallInfo callInfo, CallHeaders incomingHeaders, RequestContext context) {
-            String reqId = incomingHeaders.get(REQUEST_ID_KEY);
-            return new ServerHeaderMiddleware(reqId);
+            String requestId = incomingHeaders.get(CORRELATION_ID_KEY);
+            return new ServerHeaderMiddleware(requestId);
         }
     }
 }

server/src/main/java/org/opensearch/transport/TcpChannel.java

@@ -88,11 +88,11 @@ public interface TcpChannel extends CloseableChannel {
      * Sends a tcp message to the channel. The listener will be executed once the send process has been
      * completed.
      *
-     * @param reqId request Id
+     * @param requestId request Id
      * @param reference to send to channel
      * @param listener to execute upon send completion
      */
-    default void sendMessage(long reqId, BytesReference reference, ActionListener<Void> listener) {
+    default void sendMessage(long requestId, BytesReference reference, ActionListener<Void> listener) {
         sendMessage(reference, listener);
     }
 

server/src/main/java/org/opensearch/transport/nativeprotocol/NativeOutboundHandler.java

@@ -180,11 +180,11 @@ public void sendErrorResponse(
         sendMessage(requestId, channel, message, listener);
     }
 
-    private void sendMessage(long reqId, TcpChannel channel, NativeOutboundMessage networkMessage, ActionListener<Void> listener)
+    private void sendMessage(long requestId, TcpChannel channel, NativeOutboundMessage networkMessage, ActionListener<Void> listener)
         throws IOException {
         MessageSerializer serializer = new MessageSerializer(networkMessage, bigArrays);
         OutboundHandler.SendContext sendContext = new OutboundHandler.SendContext(statsTracker, channel, serializer, listener, serializer);
-        handler.sendBytes(reqId, channel, sendContext);
+        handler.sendBytes(requestId, channel, sendContext);
     }
 
     @Override