Fix issues at flight transport layer; Add middleware for header management

Signed-off-by: Rishabh Maurya <rishabhmaurya05@gmail.com>

Author: rishabhmaurya

plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/FlightTransportIT.java

@@ -77,7 +77,6 @@ public void setUp() throws Exception {
 
     @LockFeatureFlag(STREAM_TRANSPORT)
     public void testArrowFlightProducer() throws Exception {
-        final SearchRequest searchRequest = new SearchRequest("index");
         ActionFuture<SearchResponse> future = client().prepareStreamSearch("index").execute();
         SearchResponse resp = future.actionGet();
         assertNotNull(resp);

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/stream/ArrowStreamInput.java

@@ -48,10 +48,6 @@ public ArrowStreamInput(VectorSchemaRoot root, NamedWriteableRegistry registry)
 
         for (FieldVector vector : root.getFieldVectors()) {
             String fieldName = vector.getField().getName();
-            // skip the header field
-            if (fieldName.equals("_meta")) {
-                continue;
-            }
             String parentPath = extractParentPath(fieldName);
             vectorsByPath.computeIfAbsent(parentPath, k -> new ArrayList<>()).add(vector);
         }

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/stream/ArrowStreamOutput.java

@@ -29,7 +29,6 @@
 import org.opensearch.core.common.io.stream.Writeable;
 
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -300,16 +299,8 @@ public void writeMap(@Nullable Map<String, Object> map) throws IOException {
         structVector.setValueCount(row + 1);
     }
 
-    public VectorSchemaRoot getUnifiedRoot(ByteBuffer headers) {
+    public VectorSchemaRoot getUnifiedRoot() {
         List<FieldVector> allFields = new ArrayList<>();
-        // TODO: we need a better mechanism to serialize headers; maybe make use of Tcp headers
-        if (headers != null) {
-            Field field = new Field("_meta", new FieldType(true, new ArrowType.Binary(), null, null), null);
-            VarBinaryVector fieldVector = new VarBinaryVector(field, allocator);
-            fieldVector.setSafe(0, headers.array());
-            fieldVector.setValueCount(1);
-            allFields.add(fieldVector);
-        }
         for (VectorSchemaRoot root : roots.values()) {
             allFields.addAll(root.getFieldVectors());
         }

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

@@ -10,6 +10,7 @@
 
 import org.apache.arrow.flight.CallStatus;
 import org.apache.arrow.flight.FlightRuntimeException;
+import org.apache.arrow.flight.FlightServerMiddleware;
 import org.apache.arrow.flight.NoOpFlightProducer;
 import org.apache.arrow.flight.Ticket;
 import org.apache.arrow.memory.BufferAllocator;
@@ -28,8 +29,9 @@ public class ArrowFlightProducer extends NoOpFlightProducer {
     private final BufferAllocator allocator;
     private final InboundPipeline pipeline;
     private static final Logger logger = LogManager.getLogger(ArrowFlightProducer.class);
+    private final FlightServerMiddleware.Key<ServerHeaderMiddleware> middlewareKey;
 
-    public ArrowFlightProducer(FlightTransport flightTransport, BufferAllocator allocator) {
+    public ArrowFlightProducer(FlightTransport flightTransport, BufferAllocator allocator, FlightServerMiddleware.Key<ServerHeaderMiddleware> middlewareKey) {
         final ThreadPool threadPool = flightTransport.getThreadPool();
         final Transport.RequestHandlers requestHandlers = flightTransport.getRequestHandlers();
         this.pipeline = new InboundPipeline(
@@ -41,19 +43,22 @@ public ArrowFlightProducer(FlightTransport flightTransport, BufferAllocator allo
             requestHandlers::getHandler,
             flightTransport::inboundMessage
         );
+        this.middlewareKey = middlewareKey;
         this.allocator = allocator;
     }
 
     @Override
     public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
         try {
-            FlightServerChannel channel = new FlightServerChannel(listener, allocator);
+            FlightServerChannel channel = new FlightServerChannel(listener, allocator, context, context.getMiddleware(middlewareKey));
+            listener.setUseZeroCopy(true);
             BytesArray buf = new BytesArray(ticket.getBytes());
             // nothing changes in inbound logic, so reusing native transport inbound pipeline
             try (ReleasableBytesReference reference = ReleasableBytesReference.wrap(buf)) {
                 pipeline.handleBytes(channel, reference);
             }
         } catch (FlightRuntimeException ex) {
+            logger.error("Unexpected error during stream processing", ex);
             listener.error(ex);
             throw ex;
         } catch (Exception ex) {

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

@@ -0,0 +1,78 @@
+/*
+ * 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.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.CallInfo;
+import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.FlightClientMiddleware;
+import org.opensearch.Version;
+import org.opensearch.core.common.bytes.BytesArray;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.transport.Header;
+import org.opensearch.transport.InboundDecoder;
+import org.opensearch.transport.TransportException;
+import org.opensearch.transport.TransportStatus;
+
+import java.io.IOException;
+import java.util.Base64;
+
+/**
+ * Client middleware for handling Arrow Flight headers. It assumes that one request is sent at a time to {@link FlightClientChannel}
+ */
+public class ClientHeaderMiddleware implements FlightClientMiddleware {
+    private final HeaderContext context;
+    private final Version version;
+
+    ClientHeaderMiddleware(HeaderContext context, Version version) {
+        this.context = context;
+        this.version = version;
+    }
+
+    @Override
+    public void onHeadersReceived(CallHeaders incomingHeaders) {
+        String encodedHeader = incomingHeaders.get("raw-header");
+        byte[] headerBuffer =  Base64.getDecoder().decode(encodedHeader);
+        BytesReference headerRef = new BytesArray(headerBuffer);
+        Header header;
+        try {
+            header = InboundDecoder.readHeader(version, headerRef.length(), headerRef);
+        } catch (IOException e) {
+            throw new TransportException(e);
+        }
+        if (!Version.CURRENT.isCompatible(header.getVersion())) {
+            throw new TransportException("Incompatible version: " + header.getVersion());
+        }
+        if (TransportStatus.isError(header.getStatus())) {
+            throw new TransportException("Received error response");
+        }
+        context.setHeader(header);
+    }
+
+    @Override
+    public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {}
+
+    @Override
+    public void onCallCompleted(CallStatus status) {}
+
+    public static class Factory implements FlightClientMiddleware.Factory {
+        private final Version version;
+        private final HeaderContext context;
+
+        Factory(HeaderContext context, Version version) {
+            this.version = version;
+            this.context = context;
+        }
+
+        @Override
+        public ClientHeaderMiddleware onCallStarted(CallInfo callInfo) {
+            return new ClientHeaderMiddleware(context, version);
+        }
+    }
+}