[SPARK-27677][CORE] Serve local disk persisted blocks by the external service after releasing executor by dynamic allocation

# What changes were proposed in this pull request?

## Problem statement

An executor which has persisted blocks does not consider to be idle and this way ready to be released by dynamic allocation after the regular timeout `spark.dynamicAllocation.executorIdleTimeout` but there is separate configuration `spark.dynamicAllocation.cachedExecutorIdleTimeout` which defaults to `Integer.MAX_VALUE`. This is because releasing the executor also means losing the persisted blocks (as the metadata for individual blocks called `BlockInfo` are kept in memory) and when the RDD is referenced latter on this lost blocks will be recomputed.
On the other hand keeping the executors too long without any task to work on is also a waste of resources (as executors are reserved for the application by the resource manager).

## Solution

This PR focuses on the first part of SPARK-25888: it extends the external shuffle service with the capability to serve RDD blocks which are persisted on the local disk store by the executors. Moreover when this feature is enabled by setting the `spark.shuffle.service.fetch.rdd.enabled` config to true and a block is reported to be persisted on to disk the external shuffle service instance running on the same host as the executor is also registered (along with the reporting block manager) as a possible location for fetching it.

## Some implementation detail

Some explanation about the decisions made during the development:
- the location list to fetch a block was randomized but the groups (same host, same rack, others) order was kept. In this PR the order of groups are kept and external shuffle service added to the end of the each group.
- `BlockManagerInfo` is not introduced for external shuffle service but only a lightweight solution is taken. A hash map from `BlockId` to `BlockStatus` is introduced. A type alias would make the source more readable but I know it is discouraged. On the other hand a new class wrapping this hash map would introduce unnecessary indirection.
- when this feature is on the cleanup triggered during removing of executors (which is handled in `ExternalShuffleBlockResolver`) is modified to keep the disk persisted RDD blocks. This cleanup is triggered in standalone mode when the `spark.storage.cleanupFilesAfterExecutorExit` config is set.
- the unpersisting of an RDD is extended to use the external shuffle service for disk persisted RDD blocks when the original executor which created the blocks are already released. New block transport messages are introduced to support this: `RemoveBlocks` and `BlocksRemoved`.

# How was this patch tested?

## Unit tests

### ExternalShuffleServiceSuite

Here the complete use case is tested by the "SPARK-25888: using external shuffle service fetching disk persisted blocks" with a tiny difference: here the executor is killed manually, this way the test is a bit faster than waiting for the idle timeout.

### ExternalShuffleBlockHandlerSuite

Tests the fetching of the RDD blocks via the external shuffle service.

### BlockManagerInfoSuite

This a new suite. As the `BlockManagerInfo` behaviour depends very much on whether the external shuffle service enabled or not all the tests are executed with and without it.

### BlockManagerSuite

Tests the sorting of the block locations.

## Manually on YARN

Spark App was:

~~~scala
package com.mycompany

import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkContext, SparkConf}
import org.apache.spark.storage.StorageLevel

object TestAppDiskOnlyLevel {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setAppName("test-app")

    println("Attila: START")
    val sc = new SparkContext(conf)
    val rdd = sc.parallelize(0 until 100, 10)
      .map { i =>
        println(s"Attila: calculate first rdd i=$i")
        Thread.sleep(1000)
        i
      }

    rdd.persist(StorageLevel.DISK_ONLY)
    rdd.count()

    println("Attila: First RDD is processed, waiting for 60 sec")

    Thread.sleep(60 * 1000)

    println("Attila: Num executors must be 0 as executorIdleTimeout is way over")

    val rdd2 = sc.parallelize(0 until 10, 1)
      .map(i => (i, 1))
      .persist(StorageLevel.DISK_ONLY)

    rdd2.count()

    println("Attila: Second RDD with one partition (only one executors must be alive)")

    // reduce runs as user code to detect the empty seq (empty blocks)
    println("Calling collect on the first RDD: " + rdd.collect().reduce(_ + _))

    println("Attila: STOP")
  }
}
~~~

I have submitted with the following configuration:

~~~bash
spark-submit --master yarn \
  --conf spark.dynamicAllocation.enabled=true \
  --conf spark.dynamicAllocation.executorIdleTimeout=30 \
  --conf spark.dynamicAllocation.cachedExecutorIdleTimeout=90 \
  --class com.mycompany.TestAppDiskOnlyLevel dyn_alloc_demo-core_2.11-0.1.0-SNAPSHOT-jar-with-dependencies.jar
~~~

Checked the result by filtering for the side effect of the task calculations:

~~~bash
[userserver ~]$ yarn logs -applicationId application_1556299359453_0001 | grep "Attila: calculate" | wc -l
WARNING: YARN_OPTS has been replaced by HADOOP_OPTS. Using value of YARN_OPTS.
19/04/26 10:31:59 INFO client.RMProxy: Connecting to ResourceManager at apiros-1.gce.company.com/172.31.115.165:8032
100
~~~

So it is only 100 task execution and not 200 (which would be the case for re-computation).

Moreover from the submit/launcher log we can see executors really stopped in between (see the new total is 0 before the last line):
~~~
[userserver ~]$ grep "Attila: Num executors must be 0" -B 2 spark-submit.log
19/04/26 10:24:27 INFO cluster.YarnScheduler: Executor 9 on apiros-3.gce.company.com killed by driver.
19/04/26 10:24:27 INFO spark.ExecutorAllocationManager: Existing executor 9 has been removed (new total is 0)
Attila: Num executors must be 0 as executorIdleTimeout is way over
~~~

[Full spark submit log](https://github.com/attilapiros/spark/files/3122465/spark-submit.log)

I have done a test also after changing the `DISK_ONLY` storage level to `MEMORY_ONLY` for the first RDD. After this change during the 60sec waiting no executor was removed.

Closes #24499 from attilapiros/SPARK-25888-final.

Authored-by: “attilapiros” <[email protected]>
Signed-off-by: Marcelo Vanzin <[email protected]>

Author: attilapiros

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

@@ -91,6 +91,9 @@ protected void channelRead0(
     try {
       streamManager.checkAuthorization(client, msg.streamChunkId.streamId);
       buf = streamManager.getChunk(msg.streamChunkId.streamId, msg.streamChunkId.chunkIndex);
+      if (buf == null) {
+        throw new IllegalStateException("Chunk was not found");
+      }
     } catch (Exception e) {
       logger.error(String.format("Error opening block %s for request from %s",
         msg.streamChunkId, getRemoteAddress(channel)), e);

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

@@ -125,7 +125,10 @@ public void connectionTerminated(Channel channel) {
 
         // Release all remaining buffers.
         while (state.buffers.hasNext()) {
-          state.buffers.next().release();
+          ManagedBuffer buffer = state.buffers.next();
+          if (buffer != null) {
+            buffer.release();
+          }
         }
       }
     }

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

@@ -23,6 +23,7 @@
 
 import io.netty.channel.Channel;
 import org.apache.spark.network.server.ChunkFetchRequestHandler;
+import org.junit.Assert;
 import org.junit.Test;
 
 import static org.mockito.Mockito.*;
@@ -45,9 +46,8 @@ public void handleChunkFetchRequest() throws Exception {
     Channel channel = mock(Channel.class);
     ChannelHandlerContext context = mock(ChannelHandlerContext.class);
     when(context.channel())
-      .thenAnswer(invocationOnMock0 -> {
-        return channel;
-      });
+      .thenAnswer(invocationOnMock0 -> channel);
+
     List<Pair<Object, ExtendedChannelPromise>> responseAndPromisePairs =
       new ArrayList<>();
     when(channel.writeAndFlush(any()))
@@ -62,6 +62,7 @@ public void handleChunkFetchRequest() throws Exception {
     List<ManagedBuffer> managedBuffers = new ArrayList<>();
     managedBuffers.add(new TestManagedBuffer(10));
     managedBuffers.add(new TestManagedBuffer(20));
+    managedBuffers.add(null);
     managedBuffers.add(new TestManagedBuffer(30));
     managedBuffers.add(new TestManagedBuffer(40));
     long streamId = streamManager.registerStream("test-app", managedBuffers.iterator(), channel);
@@ -71,31 +72,40 @@ public void handleChunkFetchRequest() throws Exception {
 
     RequestMessage request0 = new ChunkFetchRequest(new StreamChunkId(streamId, 0));
     requestHandler.channelRead(context, request0);
-    assert responseAndPromisePairs.size() == 1;
-    assert responseAndPromisePairs.get(0).getLeft() instanceof ChunkFetchSuccess;
-    assert ((ChunkFetchSuccess) (responseAndPromisePairs.get(0).getLeft())).body() ==
-      managedBuffers.get(0);
+    Assert.assertEquals(1, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(0).getLeft() instanceof ChunkFetchSuccess);
+    Assert.assertEquals(managedBuffers.get(0),
+      ((ChunkFetchSuccess) (responseAndPromisePairs.get(0).getLeft())).body());
 
     RequestMessage request1 = new ChunkFetchRequest(new StreamChunkId(streamId, 1));
     requestHandler.channelRead(context, request1);
-    assert responseAndPromisePairs.size() == 2;
-    assert responseAndPromisePairs.get(1).getLeft() instanceof ChunkFetchSuccess;
-    assert ((ChunkFetchSuccess) (responseAndPromisePairs.get(1).getLeft())).body() ==
-      managedBuffers.get(1);
+    Assert.assertEquals(2, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(1).getLeft() instanceof ChunkFetchSuccess);
+    Assert.assertEquals(managedBuffers.get(1),
+      ((ChunkFetchSuccess) (responseAndPromisePairs.get(1).getLeft())).body());
 
     // Finish flushing the response for request0.
     responseAndPromisePairs.get(0).getRight().finish(true);
 
     RequestMessage request2 = new ChunkFetchRequest(new StreamChunkId(streamId, 2));
     requestHandler.channelRead(context, request2);
-    assert responseAndPromisePairs.size() == 3;
-    assert responseAndPromisePairs.get(2).getLeft() instanceof ChunkFetchSuccess;
-    assert ((ChunkFetchSuccess) (responseAndPromisePairs.get(2).getLeft())).body() ==
-      managedBuffers.get(2);
+    Assert.assertEquals(3, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(2).getLeft() instanceof ChunkFetchFailure);
+    ChunkFetchFailure chunkFetchFailure =
+        ((ChunkFetchFailure) (responseAndPromisePairs.get(2).getLeft()));
+    Assert.assertEquals("java.lang.IllegalStateException: Chunk was not found",
+        chunkFetchFailure.errorString.split("\\r?\\n")[0]);
 
     RequestMessage request3 = new ChunkFetchRequest(new StreamChunkId(streamId, 3));
     requestHandler.channelRead(context, request3);
+    Assert.assertEquals(4, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(3).getLeft() instanceof ChunkFetchSuccess);
+    Assert.assertEquals(managedBuffers.get(3),
+      ((ChunkFetchSuccess) (responseAndPromisePairs.get(3).getLeft())).body());
+
+    RequestMessage request4 = new ChunkFetchRequest(new StreamChunkId(streamId, 4));
+    requestHandler.channelRead(context, request4);
     verify(channel, times(1)).close();
-    assert responseAndPromisePairs.size() == 3;
+    Assert.assertEquals(4, responseAndPromisePairs.size());
   }
 }

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

@@ -21,6 +21,7 @@
 import java.util.List;
 
 import io.netty.channel.Channel;
+import org.junit.Assert;
 import org.junit.Test;
 
 import static org.mockito.Mockito.*;
@@ -38,7 +39,7 @@
 public class TransportRequestHandlerSuite {
 
   @Test
-  public void handleStreamRequest() throws Exception {
+  public void handleStreamRequest() {
     RpcHandler rpcHandler = new NoOpRpcHandler();
     OneForOneStreamManager streamManager = (OneForOneStreamManager) (rpcHandler.getStreamManager());
     Channel channel = mock(Channel.class);
@@ -56,48 +57,56 @@ public void handleStreamRequest() throws Exception {
     List<ManagedBuffer> managedBuffers = new ArrayList<>();
     managedBuffers.add(new TestManagedBuffer(10));
     managedBuffers.add(new TestManagedBuffer(20));
+    managedBuffers.add(null);
     managedBuffers.add(new TestManagedBuffer(30));
     managedBuffers.add(new TestManagedBuffer(40));
     long streamId = streamManager.registerStream("test-app", managedBuffers.iterator(), channel);
 
-    assert streamManager.numStreamStates() == 1;
+    Assert.assertEquals(1, streamManager.numStreamStates());
 
     TransportClient reverseClient = mock(TransportClient.class);
     TransportRequestHandler requestHandler = new TransportRequestHandler(channel, reverseClient,
       rpcHandler, 2L);
 
     RequestMessage request0 = new StreamRequest(String.format("%d_%d", streamId, 0));
     requestHandler.handle(request0);
-    assert responseAndPromisePairs.size() == 1;
-    assert responseAndPromisePairs.get(0).getLeft() instanceof StreamResponse;
-    assert ((StreamResponse) (responseAndPromisePairs.get(0).getLeft())).body() ==
-      managedBuffers.get(0);
+    Assert.assertEquals(1, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(0).getLeft() instanceof StreamResponse);
+    Assert.assertEquals(managedBuffers.get(0),
+      ((StreamResponse) (responseAndPromisePairs.get(0).getLeft())).body());
 
     RequestMessage request1 = new StreamRequest(String.format("%d_%d", streamId, 1));
     requestHandler.handle(request1);
-    assert responseAndPromisePairs.size() == 2;
-    assert responseAndPromisePairs.get(1).getLeft() instanceof StreamResponse;
-    assert ((StreamResponse) (responseAndPromisePairs.get(1).getLeft())).body() ==
-      managedBuffers.get(1);
+    Assert.assertEquals(2, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(1).getLeft() instanceof StreamResponse);
+    Assert.assertEquals(managedBuffers.get(1),
+      ((StreamResponse) (responseAndPromisePairs.get(1).getLeft())).body());
 
     // Finish flushing the response for request0.
     responseAndPromisePairs.get(0).getRight().finish(true);
 
-    RequestMessage request2 = new StreamRequest(String.format("%d_%d", streamId, 2));
+    StreamRequest request2 = new StreamRequest(String.format("%d_%d", streamId, 2));
     requestHandler.handle(request2);
-    assert responseAndPromisePairs.size() == 3;
-    assert responseAndPromisePairs.get(2).getLeft() instanceof StreamResponse;
-    assert ((StreamResponse) (responseAndPromisePairs.get(2).getLeft())).body() ==
-      managedBuffers.get(2);
+    Assert.assertEquals(3, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(2).getLeft() instanceof StreamFailure);
+    Assert.assertEquals(String.format("Stream '%s' was not found.", request2.streamId),
+        ((StreamFailure) (responseAndPromisePairs.get(2).getLeft())).error);
 
-    // Request3 will trigger the close of channel, because the number of max chunks being
-    // transferred is 2;
     RequestMessage request3 = new StreamRequest(String.format("%d_%d", streamId, 3));
     requestHandler.handle(request3);
+    Assert.assertEquals(4, responseAndPromisePairs.size());
+    Assert.assertTrue(responseAndPromisePairs.get(3).getLeft() instanceof StreamResponse);
+    Assert.assertEquals(managedBuffers.get(3),
+      ((StreamResponse) (responseAndPromisePairs.get(3).getLeft())).body());
+
+    // Request4 will trigger the close of channel, because the number of max chunks being
+    // transferred is 2;
+    RequestMessage request4 = new StreamRequest(String.format("%d_%d", streamId, 4));
+    requestHandler.handle(request4);
     verify(channel, times(1)).close();
-    assert responseAndPromisePairs.size() == 3;
+    Assert.assertEquals(4, responseAndPromisePairs.size());
 
     streamManager.connectionTerminated(channel);
-    assert streamManager.numStreamStates() == 0;
+    Assert.assertEquals(0, streamManager.numStreamStates());
   }
 }

common/network-common/src/test/java/org/apache/spark/network/server/OneForOneStreamManagerSuite.java

@@ -21,6 +21,8 @@
 import java.util.List;
 
 import io.netty.channel.Channel;
+import org.junit.After;
+import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
 
@@ -29,23 +31,69 @@
 
 public class OneForOneStreamManagerSuite {
 
+  List<ManagedBuffer> managedBuffersToRelease = new ArrayList<>();
+
+  @After
+  public void tearDown() {
+    managedBuffersToRelease.forEach(managedBuffer -> managedBuffer.release());
+    managedBuffersToRelease.clear();
+  }
+
+  private ManagedBuffer getChunk(OneForOneStreamManager manager, long streamId, int chunkIndex) {
+    ManagedBuffer chunk = manager.getChunk(streamId, chunkIndex);
+    if (chunk != null) {
+      managedBuffersToRelease.add(chunk);
+    }
+    return chunk;
+  }
+
   @Test
-  public void managedBuffersAreFeedWhenConnectionIsClosed() throws Exception {
+  public void testMissingChunk() {
     OneForOneStreamManager manager = new OneForOneStreamManager();
     List<ManagedBuffer> buffers = new ArrayList<>();
     TestManagedBuffer buffer1 = Mockito.spy(new TestManagedBuffer(10));
     TestManagedBuffer buffer2 = Mockito.spy(new TestManagedBuffer(20));
+    TestManagedBuffer buffer3 = Mockito.spy(new TestManagedBuffer(20));
+
     buffers.add(buffer1);
+    // the nulls here are to simulate a file which goes missing before being read,
+    // just as a defensive measure
+    buffers.add(null);
     buffers.add(buffer2);
+    buffers.add(null);
+    buffers.add(buffer3);
 
     Channel dummyChannel = Mockito.mock(Channel.class, Mockito.RETURNS_SMART_NULLS);
-    manager.registerStream("appId", buffers.iterator(), dummyChannel);
-    assert manager.numStreamStates() == 1;
+    long streamId = manager.registerStream("appId", buffers.iterator(), dummyChannel);
+    Assert.assertEquals(1, manager.numStreamStates());
+    Assert.assertNotNull(getChunk(manager, streamId, 0));
+    Assert.assertNull(getChunk(manager, streamId, 1));
+    Assert.assertNotNull(getChunk(manager, streamId, 2));
+    manager.connectionTerminated(dummyChannel);
+
+    // loaded buffers are not released yet as in production a MangedBuffer returned by getChunk()
+    // would only be released by Netty after it is written to the network
+    Mockito.verify(buffer1, Mockito.never()).release();
+    Mockito.verify(buffer2, Mockito.never()).release();
+    Mockito.verify(buffer3, Mockito.times(1)).release();
+  }
 
+  @Test
+  public void managedBuffersAreFreedWhenConnectionIsClosed() {
+    OneForOneStreamManager manager = new OneForOneStreamManager();
+    List<ManagedBuffer> buffers = new ArrayList<>();
+    TestManagedBuffer buffer1 = Mockito.spy(new TestManagedBuffer(10));
+    TestManagedBuffer buffer2 = Mockito.spy(new TestManagedBuffer(20));
+    buffers.add(buffer1);
+    buffers.add(buffer2);
+
+    Channel dummyChannel = Mockito.mock(Channel.class, Mockito.RETURNS_SMART_NULLS);
+    manager.registerStream("appId", buffers.iterator(), dummyChannel);
+    Assert.assertEquals(1, manager.numStreamStates());
     manager.connectionTerminated(dummyChannel);
 
     Mockito.verify(buffer1, Mockito.times(1)).release();
     Mockito.verify(buffer2, Mockito.times(1)).release();
-    assert manager.numStreamStates() == 0;
+    Assert.assertEquals(0, manager.numStreamStates());
   }
 }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/Constants.java

@@ -0,0 +1,24 @@
+/*
+ * 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;
+
+public class Constants {
+
+  public static final String SHUFFLE_SERVICE_FETCH_RDD_ENABLED =
+    "spark.shuffle.service.fetch.rdd.enabled";
+}