review feedback

Author: squito

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

@@ -57,6 +57,10 @@
 public class ExternalShuffleBlockResolver {
   private static final Logger logger = LoggerFactory.getLogger(ExternalShuffleBlockResolver.class);
 
+  private static final ObjectMapper mapper = new ObjectMapper();
+  private static final String APP_KEY_PREFIX = "AppExecShuffleInfo";
+  private static final StoreVersion CURRENT_VERSION = new StoreVersion(1, 0);
+
   // Map containing all registered executors' metadata.
   @VisibleForTesting
   final ConcurrentMap<AppExecId, ExecutorShuffleInfo> executors;
@@ -339,8 +343,6 @@ public String toString() {
     }
   }
 
-  static ObjectMapper mapper = new ObjectMapper();
-
   private static byte[] dbAppExecKey(AppExecId appExecId) throws IOException {
     // we stick a common prefix on all the keys so we can find them in the DB
     String appExecJson = mapper.writeValueAsString(appExecId);
@@ -355,8 +357,6 @@ private static AppExecId parseDbAppExecKey(String s) throws IOException {
     return parsed;
   }
 
-  private static final String APP_KEY_PREFIX = "AppExecShuffleInfo";
-
   @VisibleForTesting
   static ConcurrentMap<AppExecId, ExecutorShuffleInfo> reloadRegisteredExecutors(DB db)
       throws IOException {
@@ -367,11 +367,11 @@ static ConcurrentMap<AppExecId, ExecutorShuffleInfo> reloadRegisteredExecutors(D
       while (itr.hasNext()) {
         Map.Entry<byte[], byte[]> e = itr.next();
         String key = new String(e.getKey(), Charsets.UTF_8);
-        if (!key.startsWith(APP_KEY_PREFIX))
+        if (!key.startsWith(APP_KEY_PREFIX)) {
           break;
+        }
         AppExecId id = parseDbAppExecKey(key);
-        ExecutorShuffleInfo shuffleInfo =
-          mapper.readValue(new String(e.getValue(), Charsets.UTF_8), ExecutorShuffleInfo.class);
+        ExecutorShuffleInfo shuffleInfo = mapper.readValue(e.getValue(), ExecutorShuffleInfo.class);
         registeredExecutors.put(id, shuffleInfo);
       }
     }
@@ -387,18 +387,17 @@ public void log(String message) {
     }
   }
 
-  private static final StoreVersion CURRENT_VERSION = new StoreVersion(1,0);
+  /**
+   * Simple major.minor versioning scheme.  Any incompatible changes should be across major
+   * versions.  Minor version differences are allowed -- meaning we should be able to read
+   * dbs that are either earlier *or* later on the minor version.
+   */
   private static void checkVersion(DB db) throws IOException {
     byte[] bytes = db.get(StoreVersion.KEY);
     if (bytes == null) {
       storeVersion(db);
-    } else if (bytes.length != 8) {
-      throw new IOException("unexpected version format");
     } else {
-      DataInputStream in = new DataInputStream(new ByteArrayInputStream(bytes));
-      int major = in.readInt();
-      int minor = in.readInt();
-      StoreVersion version = new StoreVersion(major, minor);
+      StoreVersion version = mapper.readValue(bytes, StoreVersion.class);
       if (version.major != CURRENT_VERSION.major) {
         throw new IOException("cannot read state DB with version " + version + ", incompatible " +
           "with current version " + CURRENT_VERSION);
@@ -408,23 +407,18 @@ private static void checkVersion(DB db) throws IOException {
   }
 
   private static void storeVersion(DB db) throws IOException {
-    ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
-    DataOutputStream out = new DataOutputStream(bytesOut);
-    out.writeInt(CURRENT_VERSION.major);
-    out.writeInt(CURRENT_VERSION.minor);
-    out.close();
-    db.put(StoreVersion.KEY, bytesOut.toByteArray());
+    db.put(StoreVersion.KEY, mapper.writeValueAsBytes(CURRENT_VERSION));
   }
 
 
-  private static class StoreVersion {
+  public static class StoreVersion {
 
     final static byte[] KEY = "StoreVersion".getBytes(Charsets.UTF_8);
 
-    final int major;
-    final int minor;
+    public final int major;
+    public final int minor;
 
-    StoreVersion(int major, int minor) {
+    @JsonCreator public StoreVersion(@JsonProperty("major") int major, @JsonProperty("minor") int minor) {
       this.major = major;
       this.minor = minor;
     }
@@ -436,10 +430,7 @@ public boolean equals(Object o) {
 
       StoreVersion that = (StoreVersion) o;
 
-      if (major != that.major) return false;
-      if (minor != that.minor) return false;
-
-      return true;
+      return major == that.major && minor == that.minor;
     }
 
     @Override
@@ -450,6 +441,4 @@ public int hashCode() {
     }
   }
 
-
-
 }

network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java

@@ -21,6 +21,7 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.io.CharStreams;
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
 import org.apache.spark.network.util.SystemPropertyConfigProvider;
@@ -131,27 +132,25 @@ public void testHashShuffleBlocks() throws IOException {
 
   @Test
   public void jsonSerializationOfExecutorRegistration() throws IOException {
+    ObjectMapper mapper = new ObjectMapper();
     AppExecId appId = new AppExecId("foo", "bar");
-    String appIdJson = ExternalShuffleBlockResolver.mapper.writeValueAsString(appId);
-    AppExecId parsedAppId =
-      ExternalShuffleBlockResolver.mapper.readValue(appIdJson, AppExecId.class);
+    String appIdJson = mapper.writeValueAsString(appId);
+    AppExecId parsedAppId = mapper.readValue(appIdJson, AppExecId.class);
     assertEquals(parsedAppId, appId);
 
     ExecutorShuffleInfo shuffleInfo =
       new ExecutorShuffleInfo(new String[]{"/bippy", "/flippy"}, 7, "hash");
-    String shuffleJson = ExternalShuffleBlockResolver.mapper.writeValueAsString(shuffleInfo);
+    String shuffleJson = mapper.writeValueAsString(shuffleInfo);
     ExecutorShuffleInfo parsedShuffleInfo =
-      ExternalShuffleBlockResolver.mapper.readValue(shuffleJson, ExecutorShuffleInfo.class);
+      mapper.readValue(shuffleJson, ExecutorShuffleInfo.class);
     assertEquals(parsedShuffleInfo, shuffleInfo);
 
     // Intentionally keep these hard-coded strings in here, to check backwards-compatability.
     // its not legacy yet, but keeping this here in case anybody changes it
     String legacyAppIdJson = "{\"appId\":\"foo\", \"execId\":\"bar\"}";
-    assertEquals(appId,
-      ExternalShuffleBlockResolver.mapper.readValue(legacyAppIdJson, AppExecId.class));
+    assertEquals(appId, mapper.readValue(legacyAppIdJson, AppExecId.class));
     String legacyShuffleJson = "{\"localDirs\": [\"/bippy\", \"/flippy\"], " +
       "\"subDirsPerLocalDir\": 7, \"shuffleManager\": \"hash\"}";
-    assertEquals(shuffleInfo,
-      ExternalShuffleBlockResolver.mapper.readValue(legacyShuffleJson, ExecutorShuffleInfo.class));
+    assertEquals(shuffleInfo, mapper.readValue(legacyShuffleJson, ExecutorShuffleInfo.class));
   }
 }

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

@@ -16,14 +16,14 @@
  */
 package org.apache.spark.network.yarn
 
-import java.io.{DataOutputStream, FileOutputStream, PrintWriter, File}
+import java.io.{DataOutputStream, File, FileOutputStream}
 
 import scala.annotation.tailrec
 
 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.{ApplicationTerminationContext, ApplicationInitializationContext}
+import org.apache.hadoop.yarn.server.api.{ApplicationInitializationContext, ApplicationTerminationContext}
 import org.scalatest.{BeforeAndAfterEach, Matchers}
 
 import org.apache.spark.SparkFunSuite
@@ -33,11 +33,10 @@ import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo
 class YarnShuffleServiceSuite extends SparkFunSuite with Matchers with BeforeAndAfterEach {
   private[yarn] var yarnConfig: YarnConfiguration = new YarnConfiguration
 
-
   override def beforeEach(): Unit = {
     yarnConfig.set(YarnConfiguration.NM_AUX_SERVICES, "spark_shuffle")
     yarnConfig.set(YarnConfiguration.NM_AUX_SERVICE_FMT.format("spark_shuffle"),
-      classOf[YarnShuffleService].getCanonicalName);
+      classOf[YarnShuffleService].getCanonicalName)
 
     yarnConfig.get("yarn.nodemanager.local-dirs").split(",").foreach { dir =>
       val d = new File(dir)