[Issue #736] pipeline execution of hash partitioned join (cont.)

pixels-common/src/main/java/io/pixelsdb/pixels/common/physical/StreamPath.java

@@ -0,0 +1,37 @@
+package io.pixelsdb.pixels.common.physical;
+
+import static java.util.Objects.requireNonNull;
+
+public class StreamPath
+{
+    private String host;
+    private int port;
+    public boolean valid = false;
+
+    public StreamPath(String path)
+    {
+        requireNonNull(path);
+        if (path.contains(":///"))
+        {
+            path = path.substring(path.indexOf(":///") + 4);
+        }
+        int colon = path.indexOf(':');
+        if (colon > 0)
+        {
+            host = path.substring(0, colon);
+            port = Integer.parseInt(path.substring(colon + 1));
+            this.valid = true;
+        }
+    }
+
+    public String getHostName()
+    {
+        return host;
+    }
+
+    public int getPort()
+    {
+        return port;
+    }
+
+}

pixels-common/src/main/java/io/pixelsdb/pixels/common/utils/Constants.java

@@ -32,6 +32,7 @@ public final class Constants
     public static final int S3_BUFFER_SIZE = 8 * 1024 * 1024;
     public static final int REDIS_BUFFER_SIZE = 8 * 1024 * 1024;
     public static final int GCS_BUFFER_SIZE = 8 * 1024 * 1024;
+    public static final int STREAM_BUFFER_SIZE = 8 * 1024 * 1024;
 
     public static final int MIN_REPEAT = 3;
     public static final int MAX_SCOPE = 512;

pixels-common/src/main/java/io/pixelsdb/pixels/common/utils/HttpServer.java

@@ -40,16 +40,20 @@
 public final class HttpServer
 {
     final HttpServerInitializer initializer;
+    private EventLoopGroup bossGroup;
+    private EventLoopGroup workerGroup;
+    private Channel channel;
 
     public HttpServer(HttpServerHandler handler) throws CertificateException, SSLException
     {
         this.initializer = new HttpServerInitializer(HttpServerUtil.buildSslContext(), handler);
+        handler.setServerCloser(this::close);
     }
 
     public void serve(int PORT) throws InterruptedException
     {
-        EventLoopGroup bossGroup = new NioEventLoopGroup(1);
-        EventLoopGroup workerGroup = new NioEventLoopGroup();
+        bossGroup = new NioEventLoopGroup(1);
+        workerGroup = new NioEventLoopGroup();
         try
         {
             ServerBootstrap b = new ServerBootstrap();
@@ -59,13 +63,28 @@ public void serve(int PORT) throws InterruptedException
                     .handler(new LoggingHandler(LogLevel.DEBUG))
                     .childHandler(this.initializer);
 
-            Channel ch = b.bind(PORT).sync().channel();
-
-            ch.closeFuture().sync();
+            channel = b.bind(PORT).sync().channel();
+            channel.closeFuture().sync();
         } finally
         {
             bossGroup.shutdownGracefully();
             workerGroup.shutdownGracefully();
         }
     }
+
+    public void close()
+    {
+        if (channel != null)
+        {
+            channel.close();
+        }
+        if (bossGroup != null)
+        {
+            bossGroup.shutdownGracefully();
+        }
+        if (workerGroup != null)
+        {
+            workerGroup.shutdownGracefully();
+        }
+    }
 }

pixels-common/src/main/java/io/pixelsdb/pixels/common/utils/HttpServerHandler.java

@@ -43,6 +43,7 @@
 @ChannelHandler.Sharable
 public class HttpServerHandler extends SimpleChannelInboundHandler<HttpObject>
 {
+    protected Runnable serverCloser;
 
     @Override
     public void channelReadComplete(ChannelHandlerContext ctx)
@@ -100,4 +101,8 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause)
         ChannelFuture f = ctx.writeAndFlush(response);
         f.addListener(ChannelFutureListener.CLOSE);
     }
+
+    public void setServerCloser(Runnable serverCloser) {
+        this.serverCloser = serverCloser;
+    }
 }

pixels-core/src/main/java/io/pixelsdb/pixels/core/PixelsReaderStreamImpl.java

@@ -84,6 +84,7 @@ public class PixelsReaderStreamImpl implements PixelsReader
     // In partitioned mode, we use byteBufBlockingMap to map hash value to corresponding ByteBuf
     private final BlockingMap<Integer, ByteBuf> byteBufBlockingMap;
     private final boolean partitioned;
+    private final int httpPort;
     private final AtomicReference<Integer> numPartitionsReceived = new AtomicReference<>(0);
     private final List<PixelsRecordReaderStreamImpl> recordReaders;
 
@@ -98,7 +99,7 @@ public class PixelsReaderStreamImpl implements PixelsReader
 
     public PixelsReaderStreamImpl(String endpoint) throws Exception
     {
-        this(endpoint, false, -1);
+        this(endpoint, false, -2);
     }
 
     public PixelsReaderStreamImpl(int port) throws Exception
@@ -113,7 +114,7 @@ public PixelsReaderStreamImpl(String endpoint, boolean partitioned, int numParti
         this.streamHeader = null;
         URI uri = new URI(endpoint);
         String IP = uri.getHost();
-        int httpPort = uri.getPort();
+        this.httpPort = uri.getPort();
         logger.debug("In Pixels stream reader constructor, IP: " + IP + ", port: " + httpPort +
                 ", partitioned: " + partitioned + ", numPartitions: " + numPartitions);
         if (!Objects.equals(IP, "127.0.0.1") && !Objects.equals(IP, "localhost"))
@@ -151,13 +152,6 @@ public void channelRead0(ChannelHandlerContext ctx, HttpObject msg)
                         ", partition ID header: " + req.headers().get("X-Partition-Id") +
                         ", HTTP request object body total length: " + req.content().readableBytes());
 
-                // schema packet: only 1 packet expected, so close the connection immediately
-                // partitioned mode: close the connection if all partitions received
-                // else (non-partitioned mode, data packet): close connection if empty packet received
-                boolean needCloseParentChannel = partitionId == PixelsWriterStreamImpl.PARTITION_ID_SCHEMA_WRITER ||
-                        (partitioned && numPartitionsReceived.get() == numPartitions) ||
-                        (Objects.equals(req.headers().get(CONNECTION), CLOSE.toString()) &&
-                                req.content().readableBytes() == 0);
                 ByteBuf byteBuf = req.content();
                 try
                 {
@@ -178,27 +172,29 @@ public void channelRead0(ChannelHandlerContext ctx, HttpObject msg)
                         catch (IOException e)
                         {
                             logger.error("Invalid stream header values: ", e);
-                            sendResponseAndClose(ctx, req, BAD_REQUEST, needCloseParentChannel);
+                            sendResponseAndClose(ctx, req, BAD_REQUEST, false);
                             return;
                         }
                     }
                     else if (partitioned)
                     {
                         // In partitioned mode, every packet brings a streamHeader to prevent errors from possibly
-                        // out-of-order packet arrivals, so we need to parse it, but do not need the return value
-                        // (except for the first incoming packet processed above).
+                        //  out-of-order packet arrivals, so we need to parse it, but do not need the return value
+                        //  (except for the first incoming packet processed above).
+                        // XXX: Now we have each worker pass the schema in a separate packet, so this is no longer
+                        //  necessary. We can remove this block of code in PixelsWriterStreamImpl.
                         parseStreamHeader(byteBuf);
                     }
                 }
                 catch (InvalidProtocolBufferException | IndexOutOfBoundsException e)
                 {
                     logger.error("Malformed or corrupted stream header", e);
-                    sendResponseAndClose(ctx, req, BAD_REQUEST, needCloseParentChannel);
+                    sendResponseAndClose(ctx, req, BAD_REQUEST, false);
                     return;
                 }
 
                 // We only need to put the byteBuf into the blocking queue to pass it to the recordReader, if the
-                //  client is a data writer rather than a schema writer. In the latter case,
+                //  packet is a data packet rather than a schema packet. Because in the latter case,
                 //  the schema packet has been processed when parsing the stream header above.
                 if (partitionId != PixelsWriterStreamImpl.PARTITION_ID_SCHEMA_WRITER)
                 {
@@ -209,7 +205,7 @@ else if (partitioned)
                         if (partitionId < 0 || partitionId >= numPartitions)
                         {
                             logger.warn("Client sent invalid partitionId value: " + partitionId);
-                            sendResponseAndClose(ctx, req, BAD_REQUEST, needCloseParentChannel);
+                            sendResponseAndClose(ctx, req, BAD_REQUEST, false);
                             return;
                         }
                         byteBufBlockingMap.put(partitionId, byteBuf);
@@ -222,6 +218,13 @@ else if (partitioned)
                     }
                 }
 
+                // schema reader: only 1 packet expected, so close the connection immediately
+                // partitioned mode: close the connection if all partitions received
+                // else (non-partitioned mode, data packet): close connection if empty packet received
+                boolean needCloseParentChannel = numPartitions == PixelsWriterStreamImpl.PARTITION_ID_SCHEMA_WRITER ||
+                        (partitioned && numPartitionsReceived.get() == numPartitions) ||
+                        (numPartitions == -1 && Objects.equals(req.headers().get(CONNECTION), CLOSE.toString()) &&
+                                req.content().readableBytes() == 0);
                 sendResponseAndClose(ctx, req, HttpResponseStatus.OK, needCloseParentChannel);
             }
 
@@ -437,15 +440,7 @@ public int getRowGroupNum()
     @Override
     public boolean isPartitioned()
     {
-        try
-        {
-            streamHeaderLatch.await();
-        }
-        catch (InterruptedException e)
-        {
-            logger.error("Interrupted while waiting for stream header", e);
-        }
-        return this.streamHeader.hasPartitioned() && this.streamHeader.getPartitioned();
+        return partitioned;
     }
 
     /**
@@ -523,6 +518,7 @@ public PixelsProto.Footer getFooter()
     public void close()
             throws IOException
     {
+        logger.debug("Closing PixelsReaderStreamImpl");
         new Thread(() -> {
             // Conditions for closing:
             // 1. streamHeaderLatch.await() to ensure that the stream header has been received
@@ -539,11 +535,11 @@ public void close()
 
             try
             {
-                if (!this.httpServerFuture.isDone()) this.httpServerFuture.get(5, TimeUnit.SECONDS);
+                if (!this.httpServerFuture.isDone()) this.httpServerFuture.get(300, TimeUnit.SECONDS);
             }
             catch (TimeoutException e)
             {
-                logger.warn("In close(), HTTP server did not shut down in 5 seconds, doing forceful shutdown");
+                logger.warn("In close(), HTTP server on port " + httpPort + " did not shut down in 300 seconds, doing forceful shutdown");
                 this.httpServerFuture.cancel(true);
             }
             catch (InterruptedException | ExecutionException e)

pixels-core/src/main/java/io/pixelsdb/pixels/core/PixelsWriterStreamImpl.java

@@ -512,6 +512,7 @@ public void addRowBatch(VectorizedRowBatch rowBatch, int hashValue) throws IOExc
         currHashValue = hashValue;
         hashValueIsSet = true;
         curRowGroupDataLength = 0;
+        if (rowBatch == null) return;
         curRowGroupNumOfRows += rowBatch.size;
         writeColumnVectors(rowBatch.cols, rowBatch.size);
     }
@@ -550,8 +551,9 @@ public void close()
     {
         try
         {
-            if (curRowGroupNumOfRows != 0)
+            if (partitioned || curRowGroupNumOfRows != 0)
             {
+                // In partitioned mode, even an empty row group has to be sent to the server.
                 writeRowGroup();
             }
             // If the outgoing stream is empty (addRowBatch() and thus writeRowGroup() never called), we artificially
@@ -609,7 +611,9 @@ else if (isFirstRowGroup)
 
     private void writeRowGroup() throws IOException
     {
-        if (isFirstRowGroup || partitioned)
+        // XXX: Now that we have each worker pass the schema in a separate packet in partitioned mode, it is no longer
+        //  necessary to add a stream header to every packet. We can modify this block of code.
+        if (isFirstRowGroup || partitioned)  // if (isFirstRowGroup)
         {
             writeStreamHeader();
             isFirstRowGroup = false;
@@ -769,7 +773,8 @@ private void writeRowGroup() throws IOException
             uri = URI.create(fileNameToUri(fileName));
         }
         String reqUri = partitioned ? uris.get(currHashValue).toString() : uri.toString();
-        logger.debug("Sending row group with length: " + byteBuf.writerIndex() + " to endpoint: " + reqUri);
+        logger.debug("Sending row group to endpoint: " + reqUri + ", length: " + byteBuf.writerIndex()
+                + ", partitionId: " + partitionId);
         Request req = httpClient.preparePost(reqUri)
                 .setBody(byteBuf.nioBuffer())
                 .addHeader("X-Partition-Id", String.valueOf(partitionId))
@@ -786,8 +791,8 @@ private void writeRowGroup() throws IOException
         try
         {
             outstandingHTTPRequestSemaphore.acquire();
-            int maxAttempts = 30000;
-            long backoffMillis = 10;
+            int maxAttempts = 3000;
+            long backoffMillis = 100;
             int attempt = 0;
             boolean success = false;
 

pixels-core/src/main/java/io/pixelsdb/pixels/core/reader/PixelsRecordReaderStreamImpl.java

@@ -37,6 +37,8 @@
 import java.util.*;
 import java.util.concurrent.BlockingQueue;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 /**
  * PixelsRecordReaderStreamImpl is the variant of {@link PixelsRecordReaderImpl} for streaming mode.
  * <p>
@@ -153,6 +155,12 @@ public PixelsRecordReaderStreamImpl(boolean partitioned,
      */
     public void lateInitialization(PixelsStreamProto.StreamHeader streamHeader) throws IOException
     {
+        if (this.streamHeader != null)
+        {
+            checkArgument(this.streamHeader == streamHeader,
+                    "streamHeader used for lateInitialization() is not the same as the one in the RecordReader.");
+            return;
+        }
         this.streamHeader = streamHeader;
         checkBeforeRead();
     }
@@ -426,6 +434,7 @@ public int prepareBatch(int batchSize)
      */
     private VectorizedRowBatch createEmptyEOFRowBatch(int size)
     {
+        logger.debug("In createEmptyEOFRowBatch(), size = " + size);
         TypeDescription resultSchema = TypeDescription.createSchema(new ArrayList<>());
         VectorizedRowBatch resultRowBatch = resultSchema.createRowBatch(0);
         resultRowBatch.projectionSize = 0;
@@ -494,6 +503,14 @@ public VectorizedRowBatch readBatch(int batchSize, boolean reuse)
         }
 
         int rgRowCount = (int) curRowGroupStreamFooter.getNumberOfRows();
+        if (rgRowCount == 0)
+        {
+            // Empty row group, mark the current row group as unreadable.
+            curRowGroupByteBuf.readerIndex(curRowGroupByteBuf.readerIndex() + curRowGroupByteBuf.readableBytes());
+            curRGIdx++;
+            return resultSchema.createRowBatch(0, resultColumnsEncoded);
+        }
+
         int curBatchSize;
         ColumnVector[] columnVectors = resultRowBatch.cols;
 
@@ -702,6 +719,7 @@ private void acquireNewRowGroup(boolean reuse) throws IOException
         else
         // incoming byteBuf unreadable, must be end of stream
         {
+            logger.debug("In acquireNewRowGroup(), end of file");
             // checkValid = false; // Issue #105: to reject continuous read.
             if (reuse && resultRowBatch != null)
                 // XXX: Before we implement necessary checks, the close() below might be called before our readBatch()

pixels-core/src/main/java/io/pixelsdb/pixels/core/utils/BlockingMap.java

@@ -51,7 +51,7 @@ public void put(K key, V value)
 
     public V get(K key) throws InterruptedException
     {
-        V ret = getQueue(key).poll(60, TimeUnit.SECONDS);
+        V ret = getQueue(key).poll(300, TimeUnit.SECONDS);
         if (ret == null)
         {
             throw new RuntimeException("BlockingMap.get() timed out");