test for shuffle service w/ NM restarts

Author: squito

network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandler.java

@@ -48,7 +48,8 @@
 public class ExternalShuffleBlockHandler extends RpcHandler {
   private final Logger logger = LoggerFactory.getLogger(ExternalShuffleBlockHandler.class);
 
-  private final ExternalShuffleBlockResolver blockManager;
+  @VisibleForTesting
+  final ExternalShuffleBlockResolver blockManager;
   private final OneForOneStreamManager streamManager;
 
   public ExternalShuffleBlockHandler(TransportConf conf, File registeredExecutorFile)

network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolver.java

@@ -50,14 +50,16 @@ public class ExternalShuffleBlockResolver {
   private static final Logger logger = LoggerFactory.getLogger(ExternalShuffleBlockResolver.class);
 
   // Map containing all registered executors' metadata.
-  private final ConcurrentMap<AppExecId, ExecutorShuffleInfo> executors;
+  @VisibleForTesting
+  final ConcurrentMap<AppExecId, ExecutorShuffleInfo> executors;
 
   // Single-threaded Java executor used to perform expensive recursive directory deletion.
   private final Executor directoryCleaner;
 
   private final TransportConf conf;
 
-  private final File registeredExecutorFile;
+  @VisibleForTesting
+  final File registeredExecutorFile;
 
   public ExternalShuffleBlockResolver(TransportConf conf, File registeredExecutorFile)
       throws IOException, ClassNotFoundException {
@@ -233,11 +235,12 @@ static File getFile(String[] localDirs, int subDirsPerLocalDir, String filename)
   }
 
   /** Simply encodes an executor's full ID, which is appId + execId. */
-  private static class AppExecId implements Serializable {
+  @VisibleForTesting
+  static class AppExecId implements Serializable {
     final String appId;
     final String execId;
 
-    private AppExecId(String appId, String execId) {
+    AppExecId(String appId, String execId) {
       this.appId = appId;
       this.execId = execId;
     }

network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -22,6 +22,7 @@
 import java.nio.ByteBuffer;
 import java.util.List;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ContainerId;
@@ -79,9 +80,11 @@ public class YarnShuffleService extends AuxiliaryService {
   private TransportServer shuffleServer = null;
 
   // Handles registering executors and opening shuffle blocks
-  private ExternalShuffleBlockHandler blockHandler;
+  @VisibleForTesting
+  ExternalShuffleBlockHandler blockHandler;
 
-  private File registeredExecutorFile;
+  @VisibleForTesting
+  File registeredExecutorFile;
 
   public YarnShuffleService() {
     super("spark_shuffle");

yarn/src/test/scala/org/apache/spark/network/shuffle/TestUtil.scala

@@ -0,0 +1,47 @@
+/*
+ * 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.File
+
+import org.apache.spark.network.shuffle.ExternalShuffleBlockResolver.AppExecId
+import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo
+
+/**
+ * just a cheat to get package-visible members in tests
+ */
+object TestUtil {
+
+  def getBlockResolver(handler: ExternalShuffleBlockHandler): ExternalShuffleBlockResolver = {
+    handler.blockManager
+  }
+  
+  def getExecutorInfo(
+      appId: String,
+      execId: String,
+      resolver: ExternalShuffleBlockResolver
+  ): Option[ExecutorShuffleInfo] = {
+    val id = new AppExecId(appId, execId)
+    Option(resolver.executors.get(id))
+  }
+
+  def registeredExecutorFile(resolver: ExternalShuffleBlockResolver): File = {
+    resolver.registeredExecutorFile
+  }
+
+
+}

yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceSuite.scala

@@ -0,0 +1,90 @@
+/*
+ * 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.yarn
+
+import java.io.File
+
+import org.apache.commons.io.FileUtils
+import org.apache.hadoop.yarn.api.records.ApplicationId
+import org.apache.hadoop.yarn.conf.YarnConfiguration
+import org.apache.hadoop.yarn.server.api.ApplicationInitializationContext
+import org.scalatest.Matchers
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.network.shuffle.TestUtil
+import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo
+
+class YarnShuffleServiceSuite extends SparkFunSuite with Matchers {
+  private[yarn] var yarnConfig: YarnConfiguration = new YarnConfiguration
+
+  {
+    yarnConfig.set(YarnConfiguration.NM_AUX_SERVICES, "spark_shuffle");
+    yarnConfig.set(YarnConfiguration.NM_AUX_SERVICE_FMT.format("spark_shuffle"),
+      "org.apache.spark.network.yarn.YarnShuffleService");
+
+    yarnConfig.get("yarn.nodemanager.local-dirs").split(",").foreach { dir =>
+      println("making dir " + dir)
+      val d = new File(dir)
+      if (d.exists()) {
+        FileUtils.deleteDirectory(d)
+      }
+      FileUtils.forceMkdir(d)
+    }
+  }
+
+  test("executor state kept across NM restart") {
+    val service: YarnShuffleService = new YarnShuffleService
+    service.init(yarnConfig)
+    val appId = ApplicationId.newInstance(0, 0)
+    val appData: ApplicationInitializationContext =
+      new ApplicationInitializationContext("user", appId, null)
+    service.initializeApplication(appData)
+
+    val execStateFile = service.registeredExecutorFile
+    execStateFile should not be (null)
+    execStateFile.exists() should be (false)
+
+    val blockHandler = service.blockHandler
+    val blockResolver = TestUtil.getBlockResolver(blockHandler)
+    TestUtil.registeredExecutorFile(blockResolver) should be (execStateFile)
+
+    val shuffleInfo = new ExecutorShuffleInfo(Array("/foo", "/bar"), 3, "sort")
+    blockResolver.registerExecutor(appId.toString, "exec-1", shuffleInfo)
+    val executor = TestUtil.getExecutorInfo(appId.toString, "exec-1", blockResolver)
+    executor should be (Some(shuffleInfo))
+
+    execStateFile.exists() should be (true)
+
+    // now we pretend the shuffle service goes down, and comes back up
+    service.stop()
+
+    val s2: YarnShuffleService = new YarnShuffleService
+    s2.init(yarnConfig)
+    service.registeredExecutorFile should be (execStateFile)
+
+    val handler2 = service.blockHandler
+    val resolver2 = TestUtil.getBlockResolver(handler2)
+
+    // until we initial the application, don't know about any executors
+
+//    TestUtil.getExecutorInfo(appId.toString, "exec-1", blockResolver) should be (None)
+
+    s2.initializeApplication(appData)
+    val ex2 = TestUtil.getExecutorInfo(appId.toString, "exec-1", resolver2)
+    ex2 should be (Some(shuffleInfo))
+  }
+}