applying review comments of Imran v2.0

Author: attilapiros

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/CleanupNonShuffleServiceServedFilesSuite.java

@@ -177,7 +177,7 @@ private static void assertCleanedUp(TestShuffleDataContext dataContext) {
   }
 
   private static TestShuffleDataContext initDataContext(boolean withFilesToKeep)
-    throws IOException {
+      throws IOException {
     TestShuffleDataContext dataContext = new TestShuffleDataContext(10, 5);
     dataContext.create();
     if (withFilesToKeep) {
@@ -193,7 +193,7 @@ private static void createFilesToKeep(TestShuffleDataContext dataContext) throws
     dataContext.insertSortShuffleData(rand.nextInt(1000), rand.nextInt(1000), new byte[][] {
         "ABC".getBytes(StandardCharsets.UTF_8),
         "DEF".getBytes(StandardCharsets.UTF_8)});
-    dataContext.insertCachedRddData();
+    dataContext.insertCachedRddData(12, 34, new byte[] { 42 });
   }
 
   private static void createRemovableTestFiles(TestShuffleDataContext dataContext)

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java

@@ -17,10 +17,12 @@
 
 package org.apache.spark.network.shuffle;
 
+import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
@@ -31,6 +33,8 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
+import org.apache.spark.network.buffer.FileSegmentManagedBuffer;
+import org.apache.spark.network.server.OneForOneStreamManager;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -52,6 +56,11 @@ public class ExternalShuffleIntegrationSuite {
   private static final String APP_ID = "app-id";
   private static final String SORT_MANAGER = "org.apache.spark.shuffle.sort.SortShuffleManager";
 
+  private static final int RDD_ID = 1;
+  private static final int SPLIT_INDEX_VALID_BLOCK = 0;
+  private static final int SPLIT_INDEX_MISSING_FILE = 1;
+  private static final int SPLIT_INDEX_CORRUPT_LENGTH = 2;
+
   // Executor 0 is sort-based
   static TestShuffleDataContext dataContext0;
 
@@ -60,6 +69,8 @@ public class ExternalShuffleIntegrationSuite {
   static TransportConf conf;
   static TransportContext transportContext;
 
+  static byte[] exec0RddBlock = new byte[123];
+
   static byte[][] exec0Blocks = new byte[][] {
     new byte[123],
     new byte[12345],
@@ -81,13 +92,36 @@ public static void beforeAll() throws IOException {
     for (byte[] block: exec1Blocks) {
       rand.nextBytes(block);
     }
+    rand.nextBytes(exec0RddBlock);
 
     dataContext0 = new TestShuffleDataContext(2, 5);
     dataContext0.create();
     dataContext0.insertSortShuffleData(0, 0, exec0Blocks);
-
-    conf = new TransportConf("shuffle", MapConfigProvider.EMPTY);
-    handler = new ExternalShuffleBlockHandler(conf, null);
+    dataContext0.insertCachedRddData( RDD_ID, SPLIT_INDEX_VALID_BLOCK, exec0RddBlock);
+
+    HashMap<String, String> config = new HashMap<>();
+    config.put("spark.shuffle.io.maxRetries", "0");
+    conf = new TransportConf("shuffle", new MapConfigProvider(config));
+    handler = new ExternalShuffleBlockHandler(
+      new OneForOneStreamManager(),
+      new ExternalShuffleBlockResolver(conf, null) {
+        @Override
+        public ManagedBuffer getRddBlockData(String appId, String execId, int rddId, int splitIdx) {
+          ManagedBuffer res;
+          if (rddId == RDD_ID) {
+            switch (splitIdx) {
+              case SPLIT_INDEX_CORRUPT_LENGTH:
+                res = new FileSegmentManagedBuffer(conf, new File("missing.file"), 0, 12);
+                break;
+              default:
+                res = super.getRddBlockData(appId, execId, rddId, splitIdx);
+            }
+          } else {
+            res = super.getRddBlockData(appId, execId, rddId, splitIdx);
+          }
+          return res;
+        }
+      });
     transportContext = new TransportContext(conf, handler);
     server = transportContext.createServer();
   }
@@ -199,9 +233,38 @@ public void testRegisterInvalidExecutor() throws Exception {
   @Test
   public void testFetchWrongBlockId() throws Exception {
     registerExecutor("exec-1", dataContext0.createExecutorInfo(SORT_MANAGER));
-    FetchResult execFetch = fetchBlocks("exec-1", new String[] { "rdd_1_0_0" });
+    FetchResult execFetch = fetchBlocks("exec-1", new String[] { "broadcast_1" });
+    assertTrue(execFetch.successBlocks.isEmpty());
+    assertEquals(Sets.newHashSet("broadcast_1"), execFetch.failedBlocks);
+  }
+
+  @Test
+  public void testFetchValidRddBlock() throws Exception {
+    registerExecutor("exec-1", dataContext0.createExecutorInfo(SORT_MANAGER));
+    String validBlockId = "rdd_" + RDD_ID +"_" + SPLIT_INDEX_VALID_BLOCK;
+    FetchResult execFetch = fetchBlocks("exec-1", new String[] { validBlockId });
+    assertTrue(execFetch.failedBlocks.isEmpty());
+    assertEquals(Sets.newHashSet(validBlockId), execFetch.successBlocks);
+    assertBuffersEqual(new NioManagedBuffer(ByteBuffer.wrap(exec0RddBlock)),
+      execFetch.buffers.get(0));
+  }
+
+  @Test
+  public void testFetchDeletedRddBlock() throws Exception {
+    registerExecutor("exec-1", dataContext0.createExecutorInfo(SORT_MANAGER));
+    String missingBlockId = "rdd_" + RDD_ID +"_" + SPLIT_INDEX_MISSING_FILE;
+    FetchResult execFetch = fetchBlocks("exec-1", new String[] { missingBlockId });
+    assertTrue(execFetch.successBlocks.isEmpty());
+    assertEquals(Sets.newHashSet(missingBlockId), execFetch.failedBlocks);
+  }
+
+  @Test
+  public void testFetchCorruptRddBlock() throws Exception {
+    registerExecutor("exec-1", dataContext0.createExecutorInfo(SORT_MANAGER));
+    String corruptBlockId = "rdd_" + RDD_ID +"_" + SPLIT_INDEX_CORRUPT_LENGTH;
+    FetchResult execFetch = fetchBlocks("exec-1", new String[] { corruptBlockId });
     assertTrue(execFetch.successBlocks.isEmpty());
-    assertEquals(Sets.newHashSet("rdd_1_0_0"), execFetch.failedBlocks);
+    assertEquals(Sets.newHashSet(corruptBlockId), execFetch.failedBlocks);
   }
 
   @Test

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/TestShuffleDataContext.java

@@ -111,17 +111,20 @@ public void insertTempShuffleData() throws IOException {
     insertFile(filename);
   }
 
-  public void insertCachedRddData() throws IOException {
-    String filename = "rdd_12_34";
-    insertFile(filename);
+  public void insertCachedRddData(int rddId, int splitId, byte[] block) throws IOException {
+    String blockId = "rdd_" + rddId + "_" + splitId;
+    insertFile(blockId, block);
   }
-
   private void insertFile(String filename) throws IOException {
+    insertFile(filename, new byte[] { 42 });
+  }
+
+  private void insertFile(String filename, byte[] block) throws IOException {
     OutputStream dataStream = null;
     try {
       dataStream = new FileOutputStream(
         ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, filename));
-      dataStream.write(42);
+      dataStream.write(block);
     } finally {
       Closeables.close(dataStream, false);
     }

core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala

@@ -225,7 +225,7 @@ class BlockManagerMaster(
   /**
    * Find out if the executor has cached blocks which are not available via the external shuffle
    * service.
-   * This method does not consider broadcast blocks, since they are not reported the master.
+   * This method does not consider broadcast blocks, since they are not reported to the master.
    */
   def hasExclusiveCachedBlocks(executorId: String): Boolean = {
     driverEndpoint.askSync[Boolean](HasExclusiveCachedBlocks(executorId))

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

@@ -357,6 +357,8 @@ class BlockManagerMasterEndpoint(
   }
 
   private def externalShuffleServiceIdOnHost(blockManagerId: BlockManagerId): BlockManagerId = {
+    // we need to keep the executor ID of the original executor to let the shuffle service know
+    // which local directories should be used to look for the file
     BlockManagerId(blockManagerId.executorId, blockManagerId.host, externalShuffleServicePort)
   }