[SPARK-32915][CORE] Network-layer and shuffle RPC layer changes to support push shuffle blocks

common/network-common/pom.xml

@@ -91,6 +91,10 @@
       <groupId>org.apache.commons</groupId>
       <artifactId>commons-crypto</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.roaringbitmap</groupId>
+      <artifactId>RoaringBitmap</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>

common/network-common/src/main/java/org/apache/spark/network/protocol/Encoders.java

@@ -17,9 +17,11 @@
 
 package org.apache.spark.network.protocol;
 
+import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 
 import io.netty.buffer.ByteBuf;
+import org.roaringbitmap.RoaringBitmap;
 
 /** Provides a canonical set of Encoders for simple types. */
 public class Encoders {
@@ -44,6 +46,40 @@ public static String decode(ByteBuf buf) {
     }
   }
 
+  /** Bitmaps are encoded with their serialization length followed by the serialization bytes. */
+  public static class Bitmaps {
+    public static int encodedLength(RoaringBitmap b) {
+      // Compress the bitmap before serializing it. Note that since BlockTransferMessage
+      // needs to invoke encodedLength first to figure out the length for the ByteBuf, it
+      // guarantees that the bitmap will always be compressed before being serialized.
+      b.trim();
+      b.runOptimize();
+      return b.serializedSizeInBytes();
+    }
+
+    public static void encode(ByteBuf buf, RoaringBitmap b) {
+      int encodedLength = b.serializedSizeInBytes();
+      // RoaringBitmap requires nio ByteBuffer for serde. We expose the netty ByteBuf as a nio
+      // ByteBuffer. Here, we need to explicitly manage the index so we can write into the
+      // ByteBuffer, and the write is reflected in the underneath ByteBuf.
+      b.serialize(buf.nioBuffer(buf.writerIndex(), encodedLength));
+      buf.writerIndex(buf.writerIndex() + encodedLength);
+    }
+
+    public static RoaringBitmap decode(ByteBuf buf) {
+      RoaringBitmap bitmap = new RoaringBitmap();
+      try {
+        bitmap.deserialize(buf.nioBuffer());
+        // RoaringBitmap deserialize does not advance the reader index of the underlying ByteBuf.
+        // Manually update the index here.
+        buf.readerIndex(buf.readerIndex() + bitmap.serializedSizeInBytes());
+      } catch (IOException e) {
+        throw new RuntimeException("Exception while decoding bitmap", e);
+      }
+      return bitmap;
+    }
+  }
+
   /** Byte arrays are encoded with their length followed by bytes. */
   public static class ByteArrays {
     public static int encodedLength(byte[] arr) {
@@ -135,4 +171,31 @@ public static long[] decode(ByteBuf buf) {
       return longs;
     }
   }
+
+  /** Bitmap arrays are encoded with the number of bitmaps followed by per-Bitmap encoding. */
+  public static class BitmapArrays {
+    public static int encodedLength(RoaringBitmap[] bitmaps) {
+      int totalLength = 4;
+      for (RoaringBitmap b : bitmaps) {
+        totalLength += Bitmaps.encodedLength(b);
+      }
+      return totalLength;
+    }
+
+    public static void encode(ByteBuf buf, RoaringBitmap[] bitmaps) {
+      buf.writeInt(bitmaps.length);
+      for (RoaringBitmap b : bitmaps) {
+        Bitmaps.encode(buf, b);
+      }
+    }
+
+    public static RoaringBitmap[] decode(ByteBuf buf) {
+      int numBitmaps = buf.readInt();
+      RoaringBitmap[] bitmaps = new RoaringBitmap[numBitmaps];
+      for (int i = 0; i < bitmaps.length; i ++) {
+        bitmaps[i] = Bitmaps.decode(buf);
+      }
+      return bitmaps;
+    }
+  }
 }

common/network-common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java

@@ -32,6 +32,7 @@
 import org.apache.spark.network.buffer.NioManagedBuffer;
 import org.apache.spark.network.client.*;
 import org.apache.spark.network.protocol.*;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.TransportFrameDecoder;
 
 import static org.apache.spark.network.util.NettyUtils.getRemoteAddress;
@@ -181,6 +182,19 @@ public void onFailure(Throwable e) {
   private void processStreamUpload(final UploadStream req) {
     assert (req.body() == null);
     try {
+      // Retain the original metadata buffer, since it will be used during the invocation of
+      // this method. Will be released later. The metadata will be sent back to the client
+      // in the response in case the client needs the original metadata to handle the response
+      // in a callback.
+      req.meta.retain();
+      // Make a copy of the original metadata buffer. In benchmark, we noticed that
+      // we cannot respond the original metadata buffer back to the client, otherwise
+      // in cases where multiple concurrent shuffles are present, a wrong metadata might
+      // be sent back to client. This is related to the eager release of the metadata buffer,
+      // i.e., we always release the original buffer by the time the invocation of this
+      // method ends, instead of by the time we respond it to the client. This is necessary,
+      // otherwise we start seeing memory issues very quickly in benchmarks.
+      ByteBuffer meta = cloneBuffer(req.meta.nioByteBuffer());
       RpcResponseCallback callback = new RpcResponseCallback() {
         @Override
         public void onSuccess(ByteBuffer response) {
@@ -189,13 +203,17 @@ public void onSuccess(ByteBuffer response) {
 
         @Override
         public void onFailure(Throwable e) {
-          respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
+          // Piggyback request metadata as part of the exception error String, so we can
+          // respond the metadata upon a failure without changing the existing protocol.
+          respond(new RpcFailure(req.requestId,
+              JavaUtils.encodeHeaderIntoErrorString(meta.duplicate(), e)));
+          req.meta.release();
         }
       };
       TransportFrameDecoder frameDecoder = (TransportFrameDecoder)
           channel.pipeline().get(TransportFrameDecoder.HANDLER_NAME);
-      ByteBuffer meta = req.meta.nioByteBuffer();
-      StreamCallbackWithID streamHandler = rpcHandler.receiveStream(reverseClient, meta, callback);
+      StreamCallbackWithID streamHandler =
+          rpcHandler.receiveStream(reverseClient, meta.duplicate(), callback);
       if (streamHandler == null) {
         throw new NullPointerException("rpcHandler returned a null streamHandler");
       }
@@ -209,12 +227,17 @@ public void onData(String streamId, ByteBuffer buf) throws IOException {
         public void onComplete(String streamId) throws IOException {
            try {
              streamHandler.onComplete(streamId);
-             callback.onSuccess(ByteBuffer.allocate(0));
+             callback.onSuccess(meta.duplicate());
            } catch (Exception ex) {
              IOException ioExc = new IOException("Failure post-processing complete stream;" +
                " failing this rpc and leaving channel active", ex);
+             // req.meta will be released once inside callback.onFailure. Retain it one more
+             // time to be released in the finally block.
+             req.meta.retain();
              callback.onFailure(ioExc);
              streamHandler.onFailure(streamId, ioExc);
+           } finally {
+             req.meta.release();
            }
         }
 
@@ -238,12 +261,28 @@ public String getID() {
       }
     } catch (Exception e) {
       logger.error("Error while invoking RpcHandler#receive() on RPC id " + req.requestId, e);
-      respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
+      try {
+        // It's OK to respond the original metadata buffer here, because this is still inside
+        // the invocation of this method.
+        respond(new RpcFailure(req.requestId,
+            JavaUtils.encodeHeaderIntoErrorString(req.meta.nioByteBuffer(), e)));
+      } catch (IOException ioe) {
+        // No exception will be thrown here. req.meta.nioByteBuffer will not throw IOException
+        // because it's a NettyManagedBuffer. This try-catch block is to make compiler happy.
+        logger.error("Error in handling failure while invoking RpcHandler#receive() on RPC id {}",
+            req.requestId, e);
+        assert false : "Unexpected IOException in handling failure while invoking "
+            + "RpcHandler#receive()";
+      } finally {
+        req.meta.release();
+      }
       // We choose to totally fail the channel, rather than trying to recover as we do in other
       // cases.  We don't know how many bytes of the stream the client has already sent for the
       // stream, it's not worth trying to recover.
       channel.pipeline().fireExceptionCaught(e);
     } finally {
+      // Make sure we always release the original metadata buffer by the time we exit the
+      // invocation of this method. Otherwise, we see memory issues fairly quickly in benchmarks.
       req.meta.release();
     }
   }
@@ -258,6 +297,16 @@ private void processOneWayMessage(OneWayMessage req) {
     }
   }
 
+  /**
+   * Make a full copy of a nio ByteBuffer.
+   */
+  private static ByteBuffer cloneBuffer(ByteBuffer buf) {
+    ByteBuffer clone = ByteBuffer.allocate(buf.capacity());
+    clone.put(buf.duplicate());
+    clone.flip();
+    return clone;
+  }
+
   /**
    * Responds to a single message with some Encodable object. If a failure occurs while sending,
    * it will be logged and the channel closed.

common/network-common/src/main/java/org/apache/spark/network/util/JavaUtils.java

@@ -27,6 +27,7 @@
 import java.util.regex.Pattern;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
 import io.netty.buffer.Unpooled;
 import org.apache.commons.lang3.SystemUtils;
@@ -80,6 +81,17 @@ public static String bytesToString(ByteBuffer b) {
     return Unpooled.wrappedBuffer(b).toString(StandardCharsets.UTF_8);
   }
 
+  /**
+   * Encode the header ByteBuffer into the error string to be returned via RpcFailure.
+   * Use ISO_8859_1 encoding instead of UTF_8. UTF_8 will change the byte content
+   * for bytes larger than 127. This would render incorrect result when encoding
+   * decoding the index inside the PushBlockStream message.
+   */
+  public static String encodeHeaderIntoErrorString(ByteBuffer headerBuffer, Throwable e) {
+    String encodedHeader = StandardCharsets.ISO_8859_1.decode(headerBuffer).toString();
+    return encodedHeader + Throwables.getStackTraceAsString(e);
+  }
+
   /**
    * Delete a file or directory and its contents recursively.
    * Don't follow directories if they are symlinks.

common/network-common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java

@@ -19,6 +19,7 @@
 
 import java.io.*;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Semaphore;
@@ -168,6 +169,7 @@ public static void tearDown() {
   static class RpcResult {
     public Set<String> successMessages;
     public Set<String> errorMessages;
+    public Set<String> successResponses;
   }
 
   private RpcResult sendRPC(String ... commands) throws Exception {
@@ -210,6 +212,7 @@ private RpcResult sendRpcWithStream(String... streams) throws Exception {
     RpcResult res = new RpcResult();
     res.successMessages = Collections.synchronizedSet(new HashSet<>());
     res.errorMessages = Collections.synchronizedSet(new HashSet<>());
+    res.successResponses = Collections.synchronizedSet(new HashSet<>());
 
     for (String stream : streams) {
       int idx = stream.lastIndexOf('/');
@@ -247,6 +250,7 @@ private static class RpcStreamCallback implements RpcResponseCallback {
     @Override
     public void onSuccess(ByteBuffer message) {
       res.successMessages.add(streamId);
+      res.successResponses.add(JavaUtils.bytesToString(message));
       sem.release();
     }
 
@@ -326,6 +330,7 @@ public void sendRpcWithStreamOneAtATime() throws Exception {
       RpcResult res = sendRpcWithStream(stream);
       assertTrue("there were error messages!" + res.errorMessages, res.errorMessages.isEmpty());
       assertEquals(Sets.newHashSet(stream), res.successMessages);
+      assertEquals(Sets.newHashSet(stream), res.successResponses);
     }
   }
 
@@ -336,30 +341,54 @@ public void sendRpcWithStreamConcurrently() throws Exception {
       streams[i] = StreamTestHelper.STREAMS[i % StreamTestHelper.STREAMS.length];
     }
     RpcResult res = sendRpcWithStream(streams);
-    assertEquals(Sets.newHashSet(StreamTestHelper.STREAMS), res.successMessages);
+    Set<String> streamSet = Sets.newHashSet(StreamTestHelper.STREAMS);
+    assertEquals(streamSet, res.successMessages);
+    assertEquals(streamSet, res.successResponses);
     assertTrue(res.errorMessages.isEmpty());
   }
 
   @Test
   public void sendRpcWithStreamFailures() throws Exception {
     // when there is a failure reading stream data, we don't try to keep the channel usable,
     // just send back a decent error msg.
+    String failStream = "fail/exception-ondata/smallBuffer";
     RpcResult exceptionInCallbackResult =
-        sendRpcWithStream("fail/exception-ondata/smallBuffer", "smallBuffer");
+        sendRpcWithStream(failStream, "smallBuffer");
     assertErrorAndClosed(exceptionInCallbackResult, "Destination failed while reading stream");
+    assertDecodedErrorsContain(exceptionInCallbackResult.errorMessages, failStream);
 
+    failStream = "fail/null/smallBuffer";
     RpcResult nullStreamHandler =
-        sendRpcWithStream("fail/null/smallBuffer", "smallBuffer");
+        sendRpcWithStream(failStream, "smallBuffer");
     assertErrorAndClosed(exceptionInCallbackResult, "Destination failed while reading stream");
+    assertDecodedErrorsContain(nullStreamHandler.errorMessages, failStream);
 
     // OTOH, if there is a failure during onComplete, the channel should still be fine
+    failStream = "fail/exception-oncomplete/smallBuffer";
     RpcResult exceptionInOnComplete =
-        sendRpcWithStream("fail/exception-oncomplete/smallBuffer", "smallBuffer");
+        sendRpcWithStream(failStream, "smallBuffer");
     assertErrorsContain(exceptionInOnComplete.errorMessages,
         Sets.newHashSet("Failure post-processing"));
+    assertDecodedErrorsContain(exceptionInOnComplete.errorMessages, failStream);
     assertEquals(Sets.newHashSet("smallBuffer"), exceptionInOnComplete.successMessages);
   }
 
+  private void assertDecodedErrorsContain(Set<String> errors, String contain) {
+    Set<String> decodedErrors = Sets.newHashSet();
+    for (String error : errors) {
+      ByteBuffer rawBuffer = ByteBuffer.wrap(error.getBytes(StandardCharsets.ISO_8859_1));
+      decodedErrors.add(JavaUtils.bytesToString(rawBuffer));
+    }
+    boolean foundMatch = false;
+    for (String error : decodedErrors) {
+      if (error.contains(contain)) {
+        foundMatch = true;
+        break;
+      }
+    }
+    assertTrue("Could not find decoded error containing " + contain, foundMatch);
+  }
+
   private void assertErrorsContain(Set<String> errors, Set<String> contains) {
     assertEquals("Expected " + contains.size() + " errors, got " + errors.size() + "errors: " +
         errors, contains.size(), errors.size());

common/network-shuffle/pom.xml

@@ -57,6 +57,10 @@
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.roaringbitmap</groupId>
+      <artifactId>RoaringBitmap</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>
@@ -93,6 +97,11 @@
       <artifactId>mockito-core</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>commons-io</groupId>
+      <artifactId>commons-io</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>