code: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.spark.network.shuffle;
  
  import java.io.DataOutputStream;
  import java.io.File;
  import java.io.FileOutputStream;
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.nio.channels.FileChannel;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.nio.file.Paths;
  import java.util.Arrays;
  import java.util.concurrent.Semaphore;
  
  import com.google.common.base.Throwables;
  import com.google.common.collect.ImmutableMap;
  
  import org.apache.commons.io.FileUtils;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  import org.roaringbitmap.RoaringBitmap;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static org.junit.Assert.*;
  
  import org.apache.spark.network.buffer.FileSegmentManagedBuffer;
  import org.apache.spark.network.client.StreamCallbackWithID;
  import org.apache.spark.network.shuffle.RemoteBlockPushResolver.MergeShuffleFile;
  import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
  import org.apache.spark.network.shuffle.protocol.FinalizeShuffleMerge;
  import org.apache.spark.network.shuffle.protocol.MergeStatuses;
  import org.apache.spark.network.shuffle.protocol.PushBlockStream;
  import org.apache.spark.network.util.MapConfigProvider;
  import org.apache.spark.network.util.TransportConf;
  
  /**
   * Tests for {@link RemoteBlockPushResolver}.
   */
  public class RemoteBlockPushResolverSuite {
  
    private static final Logger log = LoggerFactory.getLogger(RemoteBlockPushResolverSuite.class);
    private final String TEST_APP = "testApp";
    private final String BLOCK_MANAGER_DIR = "blockmgr-193d8401";
  
    private TransportConf conf;
    private RemoteBlockPushResolver pushResolver;
    private Path[] localDirs;
  
    @Before
    public void before() throws IOException {
      localDirs = createLocalDirs(2);
      MapConfigProvider provider = new MapConfigProvider(
        ImmutableMap.of("spark.shuffle.server.minChunkSizeInMergedShuffleFile", "4"));
      conf = new TransportConf("shuffle", provider);
      pushResolver = new RemoteBlockPushResolver(conf);
      registerExecutor(TEST_APP, prepareLocalDirs(localDirs));
    }
  
    @After
    public void after() {
      try {
        for (Path local : localDirs) {
          FileUtils.deleteDirectory(local.toFile());
        }
        removeApplication(TEST_APP);
      } catch (Exception e) {
        // don't fail if clean up doesn't succeed.
        log.debug("Error while tearing down", e);
      }
    }
  
    @Test(expected = RuntimeException.class)
    public void testNoIndexFile() {
      try {
        pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      } catch (Throwable t) {
        assertTrue(t.getMessage().startsWith("Merged shuffle index file"));
        Throwables.propagate(t);
      }
    }
  
    @Test
    public void testBasicBlockMerge() throws IOException {
      PushBlock[] pushBlocks = new PushBlock[] {
        new PushBlock(0, 0, 0, ByteBuffer.wrap(new byte[4])),
        new PushBlock(0, 1, 0, ByteBuffer.wrap(new byte[5]))
      };
      pushBlockHelper(TEST_APP, pushBlocks);
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      validateMergeStatuses(statuses, new int[] {0}, new long[] {9});
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{4, 5}, new int[][]{{0}, {1}});
    }
  
    @Test
    public void testDividingMergedBlocksIntoChunks() throws IOException {
      PushBlock[] pushBlocks = new PushBlock[] {
        new PushBlock(0, 0, 0, ByteBuffer.wrap(new byte[2])),
        new PushBlock(0, 1, 0, ByteBuffer.wrap(new byte[3])),
        new PushBlock(0, 2, 0, ByteBuffer.wrap(new byte[5])),
        new PushBlock(0, 3, 0, ByteBuffer.wrap(new byte[3]))
      };
      pushBlockHelper(TEST_APP, pushBlocks);
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      validateMergeStatuses(statuses, new int[] {0}, new long[] {13});
      MergedBlockMeta meta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, meta, new int[]{5, 5, 3}, new int[][]{{0, 1}, {2}, {3}});
    }
  
    @Test
    public void testFinalizeWithMultipleReducePartitions() throws IOException {
      PushBlock[] pushBlocks = new PushBlock[] {
        new PushBlock(0, 0, 0, ByteBuffer.wrap(new byte[2])),
        new PushBlock(0, 1, 0, ByteBuffer.wrap(new byte[3])),
        new PushBlock(0, 0, 1, ByteBuffer.wrap(new byte[5])),
        new PushBlock(0, 1, 1, ByteBuffer.wrap(new byte[3]))
      };
      pushBlockHelper(TEST_APP, pushBlocks);
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      validateMergeStatuses(statuses, new int[] {0, 1}, new long[] {5, 8});
      MergedBlockMeta meta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, meta, new int[]{5}, new int[][]{{0, 1}});
    }
  
    @Test
    public void testDeferredBufsAreWrittenDuringOnData() throws IOException {
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      // This should be deferred
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[3]));
      // stream 1 now completes
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      stream1.onComplete(stream1.getID());
      // stream 2 has more data and then completes
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[3]));
      stream2.onComplete(stream2.getID());
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{4, 6}, new int[][]{{0}, {1}});
    }
  
    @Test
    public void testDeferredBufsAreWrittenDuringOnComplete() throws IOException {
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      // This should be deferred
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[3]));
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[3]));
      // stream 1 now completes
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      stream1.onComplete(stream1.getID());
      // stream 2 now completes completes
      stream2.onComplete(stream2.getID());
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{4, 6}, new int[][]{{0}, {1}});
    }
  
    @Test
    public void testDuplicateBlocksAreIgnoredWhenPrevStreamHasCompleted() throws IOException {
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      stream1.onComplete(stream1.getID());
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      // This should be ignored
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[2]));
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[2]));
      stream2.onComplete(stream2.getID());
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{4}, new int[][]{{0}});
    }
  
    @Test
    public void testDuplicateBlocksAreIgnoredWhenPrevStreamIsInProgress() throws IOException {
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      // This should be ignored
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[2]));
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[2]));
      // stream 1 now completes
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      stream1.onComplete(stream1.getID());
      // stream 2 now completes completes
      stream2.onComplete(stream2.getID());
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{4}, new int[][]{{0}});
    }
  
    @Test
    public void testFailureAfterData() throws IOException {
      StreamCallbackWithID stream =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[4]));
      stream.onFailure(stream.getID(), new RuntimeException("Forced Failure"));
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      assertEquals("num-chunks", 0, blockMeta.getNumChunks());
    }
  
    @Test
    public void testFailureAfterMultipleDataBlocks() throws IOException {
      StreamCallbackWithID stream =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[2]));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[3]));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[4]));
      stream.onFailure(stream.getID(), new RuntimeException("Forced Failure"));
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      assertEquals("num-chunks", 0, blockMeta.getNumChunks());
    }
  
    @Test
    public void testFailureAfterComplete() throws IOException {
      StreamCallbackWithID stream =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[2]));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[3]));
      stream.onData(stream.getID(), ByteBuffer.wrap(new byte[4]));
      stream.onComplete(stream.getID());
      stream.onFailure(stream.getID(), new RuntimeException("Forced Failure"));
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{9}, new int[][]{{0}});
    }
  
    @Test (expected = RuntimeException.class)
    public void testTooLateArrival() throws IOException {
      ByteBuffer[] blocks = new ByteBuffer[]{
        ByteBuffer.wrap(new byte[4]),
        ByteBuffer.wrap(new byte[5])
      };
      StreamCallbackWithID stream = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      for (ByteBuffer block : blocks) {
        stream.onData(stream.getID(), block);
      }
      stream.onComplete(stream.getID());
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      StreamCallbackWithID stream1 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[4]));
      try {
        stream1.onComplete(stream1.getID());
      } catch (RuntimeException re) {
        assertEquals(
          "Block shufflePush_0_1_0 received after merged shuffle is finalized",
            re.getMessage());
        MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
        validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{9}, new int[][]{{0}});
        throw re;
      }
    }
  
    @Test
    public void testIncompleteStreamsAreOverwritten() throws IOException {
      registerExecutor(TEST_APP, prepareLocalDirs(localDirs));
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[4]));
      // There is a failure
      stream1.onFailure(stream1.getID(), new RuntimeException("forced error"));
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[5]));
      stream2.onComplete(stream2.getID());
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[]{5}, new int[][]{{1}});
    }
  
    @Test (expected = RuntimeException.class)
    public void testCollision() throws IOException {
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      // This should be deferred
      stream2.onData(stream2.getID(), ByteBuffer.wrap(new byte[5]));
      // Since stream2 didn't get any opportunity it will throw couldn't find opportunity error
      try {
        stream2.onComplete(stream2.getID());
      } catch (RuntimeException re) {
        assertEquals(
          "Couldn't find an opportunity to write block shufflePush_0_1_0 to merged shuffle",
          re.getMessage());
        throw re;
      }
    }
  
    @Test (expected = RuntimeException.class)
    public void testFailureInAStreamDoesNotInterfereWithStreamWhichIsWriting() throws IOException {
      StreamCallbackWithID stream1 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      StreamCallbackWithID stream2 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      // There is a failure with stream2
      stream2.onFailure(stream2.getID(), new RuntimeException("forced error"));
      StreamCallbackWithID stream3 =
        pushResolver.receiveBlockDataAsStream(new PushBlockStream(TEST_APP, 0, 2, 0, 0));
      // This should be deferred
      stream3.onData(stream3.getID(), ByteBuffer.wrap(new byte[5]));
      // Since this stream didn't get any opportunity it will throw couldn't find opportunity error
      RuntimeException failedEx = null;
      try {
        stream3.onComplete(stream3.getID());
      } catch (RuntimeException re) {
        assertEquals(
          "Couldn't find an opportunity to write block shufflePush_0_2_0 to merged shuffle",
          re.getMessage());
        failedEx = re;
      }
      // stream 1 now completes
      stream1.onData(stream1.getID(), ByteBuffer.wrap(new byte[2]));
      stream1.onComplete(stream1.getID());
  
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {4}, new int[][] {{0}});
      if (failedEx != null) {
        throw failedEx;
      }
    }
  
    @Test(expected = NullPointerException.class)
    public void testUpdateLocalDirsOnlyOnce() throws IOException {
      String testApp = "updateLocalDirsOnlyOnceTest";
      Path[] activeLocalDirs = createLocalDirs(1);
      registerExecutor(testApp, prepareLocalDirs(activeLocalDirs));
      assertEquals(pushResolver.getMergedBlockDirs(testApp).length, 1);
      assertTrue(pushResolver.getMergedBlockDirs(testApp)[0].contains(
        activeLocalDirs[0].toFile().getPath()));
      // Any later executor register from the same application should not change the active local
      // dirs list
      Path[] updatedLocalDirs = localDirs;
      registerExecutor(testApp, prepareLocalDirs(updatedLocalDirs));
      assertEquals(pushResolver.getMergedBlockDirs(testApp).length, 1);
      assertTrue(pushResolver.getMergedBlockDirs(testApp)[0].contains(
        activeLocalDirs[0].toFile().getPath()));
      removeApplication(testApp);
      try {
        pushResolver.getMergedBlockDirs(testApp);
      } catch (Throwable e) {
        assertTrue(e.getMessage()
          .startsWith("application " + testApp + " is not registered or NM was restarted."));
        Throwables.propagate(e);
      }
    }
  
    @Test
    public void testCleanUpDirectory() throws IOException, InterruptedException {
      String testApp = "cleanUpDirectory";
      Semaphore deleted = new Semaphore(0);
      pushResolver = new RemoteBlockPushResolver(conf) {
        @Override
        void deleteExecutorDirs(Path[] dirs) {
          super.deleteExecutorDirs(dirs);
          deleted.release();
        }
      };
      Path[] activeDirs = createLocalDirs(1);
      registerExecutor(testApp, prepareLocalDirs(activeDirs));
      PushBlock[] pushBlocks = new PushBlock[] {
        new PushBlock(0, 0, 0, ByteBuffer.wrap(new byte[4]))};
      pushBlockHelper(testApp, pushBlocks);
      pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(testApp, 0));
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(testApp, 0, 0);
      validateChunks(testApp, 0, 0, blockMeta, new int[]{4}, new int[][]{{0}});
      String[] mergeDirs = pushResolver.getMergedBlockDirs(testApp);
      pushResolver.applicationRemoved(testApp,  true);
      // Since the cleanup happen in a different thread, check few times to see if the merge dirs gets
      // deleted.
      deleted.acquire();
      for (String mergeDir : mergeDirs) {
        Assert.assertFalse(Files.exists(Paths.get(mergeDir)));
      }
    }
  
    @Test
    public void testRecoverIndexFileAfterIOExceptions() throws IOException {
      useTestFiles(true, false);
      RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[4]));
      callback1.onComplete(callback1.getID());
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback1.getPartitionInfo();
      // Close the index stream so it throws IOException
      TestMergeShuffleFile testIndexFile = (TestMergeShuffleFile) partitionInfo.getIndexFile();
      testIndexFile.close();
      StreamCallbackWithID callback2 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      callback2.onData(callback2.getID(), ByteBuffer.wrap(new byte[5]));
      // This will complete without any IOExceptions because number of IOExceptions are less than
      // the threshold but the update to index file will be unsuccessful.
      callback2.onComplete(callback2.getID());
      assertEquals("index position", 16, testIndexFile.getPos());
      // Restore the index stream so it can write successfully again.
      testIndexFile.restore();
      StreamCallbackWithID callback3 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 2, 0, 0));
      callback3.onData(callback3.getID(), ByteBuffer.wrap(new byte[2]));
      callback3.onComplete(callback3.getID());
      assertEquals("index position", 24, testIndexFile.getPos());
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      validateMergeStatuses(statuses, new int[] {0}, new long[] {11});
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {4, 7}, new int[][] {{0}, {1, 2}});
    }
  
    @Test
    public void testRecoverIndexFileAfterIOExceptionsInFinalize() throws IOException {
      useTestFiles(true, false);
      RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[4]));
      callback1.onComplete(callback1.getID());
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback1.getPartitionInfo();
      // Close the index stream so it throws IOException
      TestMergeShuffleFile testIndexFile = (TestMergeShuffleFile) partitionInfo.getIndexFile();
      testIndexFile.close();
      StreamCallbackWithID callback2 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      callback2.onData(callback2.getID(), ByteBuffer.wrap(new byte[5]));
      // This will complete without any IOExceptions because number of IOExceptions are less than
      // the threshold but the update to index file will be unsuccessful.
      callback2.onComplete(callback2.getID());
      assertEquals("index position", 16, testIndexFile.getPos());
      // The last update to index was unsuccessful however any further updates will be successful.
      // Restore the index stream so it can write successfully again.
      testIndexFile.restore();
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      assertEquals("index position", 24, testIndexFile.getPos());
      validateMergeStatuses(statuses, new int[] {0}, new long[] {9});
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {4, 5}, new int[][] {{0}, {1}});
    }
  
    @Test
    public void testRecoverMetaFileAfterIOExceptions() throws IOException {
      useTestFiles(false, true);
      RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[4]));
      callback1.onComplete(callback1.getID());
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback1.getPartitionInfo();
      // Close the meta stream so it throws IOException
      TestMergeShuffleFile testMetaFile = (TestMergeShuffleFile) partitionInfo.getMetaFile();
      long metaPosBeforeClose = testMetaFile.getPos();
      testMetaFile.close();
      StreamCallbackWithID callback2 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      callback2.onData(callback2.getID(), ByteBuffer.wrap(new byte[5]));
      // This will complete without any IOExceptions because number of IOExceptions are less than
      // the threshold but the update to index and meta file will be unsuccessful.
      callback2.onComplete(callback2.getID());
      assertEquals("index position", 16, partitionInfo.getIndexFile().getPos());
      assertEquals("meta position", metaPosBeforeClose, testMetaFile.getPos());
      // Restore the meta stream so it can write successfully again.
      testMetaFile.restore();
      StreamCallbackWithID callback3 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 2, 0, 0));
      callback3.onData(callback3.getID(), ByteBuffer.wrap(new byte[2]));
      callback3.onComplete(callback3.getID());
      assertEquals("index position", 24, partitionInfo.getIndexFile().getPos());
      assertTrue("meta position", testMetaFile.getPos() > metaPosBeforeClose);
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      validateMergeStatuses(statuses, new int[] {0}, new long[] {11});
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {4, 7}, new int[][] {{0}, {1, 2}});
    }
  
    @Test
    public void testRecoverMetaFileAfterIOExceptionsInFinalize() throws IOException {
      useTestFiles(false, true);
      RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[4]));
      callback1.onComplete(callback1.getID());
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback1.getPartitionInfo();
      // Close the meta stream so it throws IOException
      TestMergeShuffleFile testMetaFile = (TestMergeShuffleFile) partitionInfo.getMetaFile();
      long metaPosBeforeClose = testMetaFile.getPos();
      testMetaFile.close();
      StreamCallbackWithID callback2 = pushResolver.receiveBlockDataAsStream(
        new PushBlockStream(TEST_APP, 0, 1, 0, 0));
      callback2.onData(callback2.getID(), ByteBuffer.wrap(new byte[5]));
      // This will complete without any IOExceptions because number of IOExceptions are less than
      // the threshold but the update to index and meta file will be unsuccessful.
      callback2.onComplete(callback2.getID());
      MergeShuffleFile indexFile = partitionInfo.getIndexFile();
      assertEquals("index position", 16, indexFile.getPos());
      assertEquals("meta position", metaPosBeforeClose, testMetaFile.getPos());
      // Restore the meta stream so it can write successfully again.
      testMetaFile.restore();
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      assertEquals("index position", 24, indexFile.getPos());
      assertTrue("meta position", testMetaFile.getPos() > metaPosBeforeClose);
      validateMergeStatuses(statuses, new int[] {0}, new long[] {9});
      MergedBlockMeta blockMeta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 0);
      validateChunks(TEST_APP, 0, 0, blockMeta, new int[] {4, 5}, new int[][] {{0}, {1}});
    }
  
    @Test (expected = RuntimeException.class)
    public void testIOExceptionsExceededThreshold() throws IOException {
      RemoteBlockPushResolver.PushBlockStreamCallback callback =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback.getPartitionInfo();
      callback.onData(callback.getID(), ByteBuffer.wrap(new byte[4]));
      callback.onComplete(callback.getID());
      // Close the data stream so it throws continuous IOException
      partitionInfo.getDataChannel().close();
      for (int i = 1; i < 5; i++) {
        RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
          (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
            new PushBlockStream(TEST_APP, 0, i, 0, 0));
        try {
          callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[2]));
        } catch (IOException ioe) {
          // this will throw IOException so the client can retry.
          callback1.onFailure(callback1.getID(), ioe);
        }
      }
      assertEquals(4, partitionInfo.getNumIOExceptions());
      // After 4 IOException, the server will respond with IOExceptions exceeded threshold
      try {
        RemoteBlockPushResolver.PushBlockStreamCallback callback2 =
          (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
            new PushBlockStream(TEST_APP, 0, 5, 0, 0));
        callback2.onData(callback.getID(), ByteBuffer.wrap(new byte[1]));
      } catch (Throwable t) {
        assertEquals("IOExceptions exceeded the threshold when merging shufflePush_0_5_0",
          t.getMessage());
        throw t;
      }
    }
  
    @Test (expected = RuntimeException.class)
    public void testIOExceptionsDuringMetaUpdateIncreasesExceptionCount() throws IOException {
      useTestFiles(true, false);
      RemoteBlockPushResolver.PushBlockStreamCallback callback =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback.getPartitionInfo();
      callback.onData(callback.getID(), ByteBuffer.wrap(new byte[4]));
      callback.onComplete(callback.getID());
      TestMergeShuffleFile testIndexFile = (TestMergeShuffleFile) partitionInfo.getIndexFile();
      testIndexFile.close();
      for (int i = 1; i < 5; i++) {
        RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
          (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
            new PushBlockStream(TEST_APP, 0, i, 0, 0));
        callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[5]));
        // This will complete without any exceptions but the exception count is increased.
        callback1.onComplete(callback1.getID());
      }
      assertEquals(4, partitionInfo.getNumIOExceptions());
      // After 4 IOException, the server will respond with IOExceptions exceeded threshold for any
      // new request for this partition.
      try {
        RemoteBlockPushResolver.PushBlockStreamCallback callback2 =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 5, 0, 0));
        callback2.onData(callback2.getID(), ByteBuffer.wrap(new byte[4]));
        callback2.onComplete(callback2.getID());
      } catch (Throwable t) {
        assertEquals("IOExceptions exceeded the threshold when merging shufflePush_0_5_0",
          t.getMessage());
        throw t;
      }
    }
  
    @Test (expected = RuntimeException.class)
    public void testRequestForAbortedShufflePartitionThrowsException() {
      try {
        testIOExceptionsDuringMetaUpdateIncreasesExceptionCount();
      } catch (Throwable t) {
        // No more blocks can be merged to this partition.
      }
      try {
        pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 10, 0, 0));
      } catch (Throwable t) {
        assertEquals("IOExceptions exceeded the threshold when merging shufflePush_0_10_0",
          t.getMessage());
        throw t;
      }
    }
  
    @Test (expected = RuntimeException.class)
    public void testPendingBlockIsAbortedImmediately() throws IOException {
      useTestFiles(true, false);
      RemoteBlockPushResolver.PushBlockStreamCallback callback =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback.getPartitionInfo();
      TestMergeShuffleFile testIndexFile = (TestMergeShuffleFile) partitionInfo.getIndexFile();
      testIndexFile.close();
      for (int i = 1; i < 6; i++) {
        RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
          (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
            new PushBlockStream(TEST_APP, 0, i, 0, 0));
        try {
          callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[5]));
          // This will complete without any exceptions but the exception count is increased.
          callback1.onComplete(callback1.getID());
        } catch (Throwable t) {
          callback1.onFailure(callback1.getID(), t);
        }
      }
      assertEquals(5, partitionInfo.getNumIOExceptions());
      // The server will respond with IOExceptions exceeded threshold for any additional attempts
      // to write.
      try {
        callback.onData(callback.getID(), ByteBuffer.wrap(new byte[4]));
      } catch (Throwable t) {
        assertEquals("IOExceptions exceeded the threshold when merging shufflePush_0_0_0",
          t.getMessage());
        throw t;
      }
    }
  
    @Test (expected = RuntimeException.class)
    public void testWritingPendingBufsIsAbortedImmediatelyDuringComplete() throws IOException {
      useTestFiles(true, false);
      RemoteBlockPushResolver.PushBlockStreamCallback callback =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 0, 0, 0));
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback.getPartitionInfo();
      TestMergeShuffleFile testIndexFile = (TestMergeShuffleFile) partitionInfo.getIndexFile();
      testIndexFile.close();
      for (int i = 1; i < 5; i++) {
        RemoteBlockPushResolver.PushBlockStreamCallback callback1 =
          (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
            new PushBlockStream(TEST_APP, 0, i, 0, 0));
        try {
          callback1.onData(callback1.getID(), ByteBuffer.wrap(new byte[5]));
          // This will complete without any exceptions but the exception count is increased.
          callback1.onComplete(callback1.getID());
        } catch (Throwable t) {
          callback1.onFailure(callback1.getID(), t);
        }
      }
      assertEquals(4, partitionInfo.getNumIOExceptions());
      RemoteBlockPushResolver.PushBlockStreamCallback callback2 =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 5, 0, 0));
      callback2.onData(callback2.getID(), ByteBuffer.wrap(new byte[5]));
      // This is deferred
      callback.onData(callback.getID(), ByteBuffer.wrap(new byte[4]));
      // Callback2 completes which will throw another exception.
      try {
        callback2.onComplete(callback2.getID());
      } catch (Throwable t) {
        callback2.onFailure(callback2.getID(), t);
      }
      assertEquals(5, partitionInfo.getNumIOExceptions());
      // Restore index file so that any further writes to it are successful and any exceptions are
      // due to IOExceptions exceeding threshold.
      testIndexFile.restore();
      try {
        callback.onComplete(callback.getID());
      } catch (Throwable t) {
        assertEquals("IOExceptions exceeded the threshold when merging shufflePush_0_0_0",
          t.getMessage());
        throw t;
      }
    }
  
    @Test
    public void testFailureWhileTruncatingFiles() throws IOException {
      useTestFiles(true, false);
      PushBlock[] pushBlocks = new PushBlock[] {
        new PushBlock(0, 0, 0, ByteBuffer.wrap(new byte[2])),
        new PushBlock(0, 1, 0, ByteBuffer.wrap(new byte[3])),
        new PushBlock(0, 0, 1, ByteBuffer.wrap(new byte[5])),
        new PushBlock(0, 1, 1, ByteBuffer.wrap(new byte[3]))
      };
      pushBlockHelper(TEST_APP, pushBlocks);
      RemoteBlockPushResolver.PushBlockStreamCallback callback =
        (RemoteBlockPushResolver.PushBlockStreamCallback) pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(TEST_APP, 0, 2, 0, 0));
      callback.onData(callback.getID(), ByteBuffer.wrap(new byte[2]));
      callback.onComplete(callback.getID());
      RemoteBlockPushResolver.AppShufflePartitionInfo partitionInfo = callback.getPartitionInfo();
      TestMergeShuffleFile testIndexFile = (TestMergeShuffleFile) partitionInfo.getIndexFile();
      // Close the index file so truncate throws IOException
      testIndexFile.close();
      MergeStatuses statuses = pushResolver.finalizeShuffleMerge(
        new FinalizeShuffleMerge(TEST_APP, 0));
      validateMergeStatuses(statuses, new int[] {1}, new long[] {8});
      MergedBlockMeta meta = pushResolver.getMergedBlockMeta(TEST_APP, 0, 1);
      validateChunks(TEST_APP, 0, 1, meta, new int[]{5, 3}, new int[][]{{0},{1}});
    }
  
    private void useTestFiles(boolean useTestIndexFile, boolean useTestMetaFile) throws IOException {
      pushResolver = new RemoteBlockPushResolver(conf) {
        @Override
        AppShufflePartitionInfo newAppShufflePartitionInfo(AppShuffleId appShuffleId, int reduceId,
          File dataFile, File indexFile, File metaFile) throws IOException {
          MergeShuffleFile mergedIndexFile = useTestIndexFile ? new TestMergeShuffleFile(indexFile)
            : new MergeShuffleFile(indexFile);
          MergeShuffleFile mergedMetaFile = useTestMetaFile ? new TestMergeShuffleFile(metaFile) :
            new MergeShuffleFile(metaFile);
          return new AppShufflePartitionInfo(appShuffleId, reduceId, dataFile, mergedIndexFile,
            mergedMetaFile);
        }
      };
      registerExecutor(TEST_APP, prepareLocalDirs(localDirs));
    }
  
    private Path[] createLocalDirs(int numLocalDirs) throws IOException {
      Path[] localDirs = new Path[numLocalDirs];
      for (int i = 0; i < localDirs.length; i++) {
        localDirs[i] = Files.createTempDirectory("shuffleMerge");
        localDirs[i].toFile().deleteOnExit();
      }
      return localDirs;
    }
  
    private void registerExecutor(String appId, String[] localDirs) throws IOException {
      ExecutorShuffleInfo shuffleInfo = new ExecutorShuffleInfo(localDirs, 1, "mergedShuffle");
      pushResolver.registerExecutor(appId, shuffleInfo);
    }
  
    private String[] prepareLocalDirs(Path[] localDirs) throws IOException {
      String[] blockMgrDirs = new String[localDirs.length];
      for (int i = 0; i< localDirs.length; i++) {
        Files.createDirectories(localDirs[i].resolve(
          RemoteBlockPushResolver.MERGE_MANAGER_DIR + File.separator + "00"));
        blockMgrDirs[i] = localDirs[i].toFile().getPath() + File.separator + BLOCK_MANAGER_DIR;
      }
      return blockMgrDirs;
    }
  
    private void removeApplication(String appId) {
      // PushResolver cleans up the local dirs in a different thread which can conflict with the test
      // data of other tests, since they are using the same Application Id.
      pushResolver.applicationRemoved(appId,  false);
    }
  
    private void validateMergeStatuses(
        MergeStatuses mergeStatuses,
        int[] expectedReduceIds,
        long[] expectedSizes) {
      assertArrayEquals(expectedReduceIds, mergeStatuses.reduceIds);
      assertArrayEquals(expectedSizes, mergeStatuses.sizes);
    }
  
    private void validateChunks(
        String appId,
        int shuffleId,
        int reduceId,
        MergedBlockMeta meta,
        int[] expectedSizes,
        int[][] expectedMapsPerChunk) throws IOException {
      assertEquals("num chunks", expectedSizes.length, meta.getNumChunks());
      RoaringBitmap[] bitmaps = meta.readChunkBitmaps();
      assertEquals("num of bitmaps", meta.getNumChunks(), bitmaps.length);
      for (int i = 0; i < meta.getNumChunks(); i++) {
        RoaringBitmap chunkBitmap = bitmaps[i];
        Arrays.stream(expectedMapsPerChunk[i]).forEach(x -> assertTrue(chunkBitmap.contains(x)));
      }
      for (int i = 0; i < meta.getNumChunks(); i++) {
        FileSegmentManagedBuffer mb =
          (FileSegmentManagedBuffer) pushResolver.getMergedBlockData(appId, shuffleId, reduceId, i);
        assertEquals(expectedSizes[i], mb.getLength());
      }
    }
  
    private void pushBlockHelper(
        String appId,
        PushBlock[] blocks) throws IOException {
      for (int i = 0; i < blocks.length; i++) {
        StreamCallbackWithID stream = pushResolver.receiveBlockDataAsStream(
          new PushBlockStream(appId, blocks[i].shuffleId, blocks[i].mapIndex, blocks[i].reduceId, 0));
        stream.onData(stream.getID(), blocks[i].buffer);
        stream.onComplete(stream.getID());
      }
    }
  
    private static class PushBlock {
      private final int shuffleId;
      private final int mapIndex;
      private final int reduceId;
      private final ByteBuffer buffer;
      PushBlock(int shuffleId, int mapIndex, int reduceId, ByteBuffer buffer) {
        this.shuffleId = shuffleId;
        this.mapIndex = mapIndex;
        this.reduceId = reduceId;
        this.buffer = buffer;
      }
    }
  
    private static class TestMergeShuffleFile extends MergeShuffleFile {
      private DataOutputStream activeDos;
      private File file;
      private FileChannel channel;
  
      private TestMergeShuffleFile(File file) throws IOException {
        super(null, null);
        this.file = file;
        FileOutputStream fos = new FileOutputStream(file);
        channel = fos.getChannel();
        activeDos = new DataOutputStream(fos);
      }
  
      @Override
      DataOutputStream getDos() {
        return activeDos;
      }
  
      @Override
      FileChannel getChannel() {
        return channel;
      }
  
      @Override
      void close() throws IOException {
        activeDos.close();
      }
  
      void restore() throws IOException {
        FileOutputStream fos = new FileOutputStream(file, true);
        channel = fos.getChannel();
        activeDos = new DataOutputStream(fos);
      }
    }
  }
  