[SPARK-38909][CORE][YARN] Encapsulate LevelDB used to store remote/external shuffle state as DB

common/network-common/src/main/java/org/apache/spark/network/shuffledb/DB.java

@@ -0,0 +1,49 @@
+/*
+ * 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.shuffledb;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * The local KV storage used to persist the shuffle state,
+ * the implementations may include leveldb, rocksdb, etc.
+ */
+public interface DB extends Closeable {
+    /**
+     * Set the DB entry for "key" to "value".
+     */
+    void put(byte[] key, byte[] value) throws RuntimeException;
+
+    /**
+     * Get which returns a new byte array storing the value associated
+     * with the specified input key if any.
+     */
+    byte[] get(byte[] key) throws RuntimeException;
+
+    /**
+     * Delete the DB entry (if any) for "key".
+     */
+    void delete(byte[] key) throws RuntimeException;
+
+    /**
+     * Read KV prefixed with `prefix` into a Map from DB.
+     */
+    Map<String, byte[]> readKVToMap(String prefix) throws IOException;
+}

common/network-common/src/main/java/org/apache/spark/network/shuffledb/LevelDB.java

@@ -0,0 +1,69 @@
+/*
+ * 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.shuffledb;
+
+import org.iq80.leveldb.DBIterator;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+
+public class LevelDB implements DB {
+    private final org.iq80.leveldb.DB db;
+
+    public LevelDB(org.iq80.leveldb.DB db) {
+        this.db = db;
+    }
+
+    @Override
+    public void put(byte[] key, byte[] value) throws RuntimeException {
+        db.put(key, value);
+    }
+
+    @Override
+    public byte[] get(byte[] key) throws RuntimeException {
+       return db.get(key);
+    }
+
+    @Override
+    public void delete(byte[] key) throws RuntimeException {
+        db.delete(key);
+    }
+
+    @Override
+    public void close() throws IOException {
+        db.close();
+    }
+
+    @Override
+    public Map<String, byte[]> readKVToMap(String prefix) {
+        Map<String, byte[]> map = new HashMap<>();
+        DBIterator itr = db.iterator();
+        itr.seek(prefix.getBytes(StandardCharsets.UTF_8));
+        while (itr.hasNext()) {
+            Map.Entry<byte[], byte[]> e = itr.next();
+            String key = new String(e.getKey(), StandardCharsets.UTF_8);
+            if (!key.startsWith(prefix)) {
+                break;
+            }
+            map.put(key, e.getValue());
+        }
+        return map;
+    }
+}

common/network-common/src/main/java/org/apache/spark/network/shuffledb/StoreVersion.java

@@ -0,0 +1,54 @@
+/*
+ * 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.shuffledb;
+
+import java.nio.charset.StandardCharsets;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class StoreVersion {
+
+    public static final byte[] KEY = "StoreVersion".getBytes(StandardCharsets.UTF_8);
+
+    public final int major;
+    public final int minor;
+
+    @JsonCreator
+    public StoreVersion(@JsonProperty("major") int major, @JsonProperty("minor") int minor) {
+        this.major = major;
+        this.minor = minor;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        StoreVersion that = (StoreVersion) o;
+
+        return major == that.major && minor == that.minor;
+    }
+
+    @Override
+    public int hashCode() {
+        int result = major;
+        result = 31 * result + minor;
+        return result;
+    }
+}

common/network-common/src/main/java/org/apache/spark/network/util/DBProvider.java

@@ -0,0 +1,51 @@
+/*
+ * 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.util;
+
+import java.io.File;
+import java.io.IOException;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
+
+import org.apache.spark.network.shuffledb.LevelDB;
+import org.apache.spark.network.shuffledb.DB;
+import org.apache.spark.network.shuffledb.StoreVersion;
+
+public class DBProvider {
+    public static DB initDB(File dbFile, StoreVersion version, ObjectMapper mapper)
+        throws IOException {
+        if (dbFile != null) {
+            if (dbFile.getName().endsWith(".ldb")) {
+                org.iq80.leveldb.DB levelDB = LevelDBProvider.initLevelDB(dbFile, version, mapper);
+                return levelDB != null ? new LevelDB(levelDB) : null;
+            } else {
+                return null;
+            }
+        }
+        return null;
+    }
+
+    @VisibleForTesting
+    public static DB initDB(File file) throws IOException {
+        if (file.getName().endsWith(".ldb")) {
+            return new LevelDB(LevelDBProvider.initLevelDB(file));
+        } else {
+            return null;
+        }
+    }
+}

common/network-common/src/main/java/org/apache/spark/network/util/LevelDBProvider.java

@@ -19,18 +19,18 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.nio.charset.StandardCharsets;
 
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
 import org.fusesource.leveldbjni.JniDBFactory;
 import org.fusesource.leveldbjni.internal.NativeDB;
 import org.iq80.leveldb.DB;
 import org.iq80.leveldb.Options;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.spark.network.shuffledb.StoreVersion;
+
 /**
  * LevelDB utility class available in the network package.
  */
@@ -85,6 +85,14 @@ public static DB initLevelDB(File dbFile, StoreVersion version, ObjectMapper map
     return tmpDb;
   }
 
+  @VisibleForTesting
+  static DB initLevelDB(File file) throws IOException {
+    Options options = new Options();
+    options.createIfMissing(true);
+    JniDBFactory factory = new JniDBFactory();
+    return factory.open(file, options);
+  }
+
   private static class LevelDBLogger implements org.iq80.leveldb.Logger {
     private static final Logger LOG = LoggerFactory.getLogger(LevelDBLogger.class);
 
@@ -118,35 +126,4 @@ public static void storeVersion(DB db, StoreVersion version, ObjectMapper mapper
       throws IOException {
     db.put(StoreVersion.KEY, mapper.writeValueAsBytes(version));
   }
-
-  public static class StoreVersion {
-
-    static final byte[] KEY = "StoreVersion".getBytes(StandardCharsets.UTF_8);
-
-    public final int major;
-    public final int minor;
-
-    @JsonCreator
-    public StoreVersion(@JsonProperty("major") int major, @JsonProperty("minor") int minor) {
-      this.major = major;
-      this.minor = minor;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-      if (this == o) return true;
-      if (o == null || getClass() != o.getClass()) return false;
-
-      StoreVersion that = (StoreVersion) o;
-
-      return major == that.major && minor == that.minor;
-    }
-
-    @Override
-    public int hashCode() {
-      int result = major;
-      result = 31 * result + minor;
-      return result;
-    }
-  }
 }

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

@@ -38,8 +38,6 @@
 import com.google.common.cache.LoadingCache;
 import com.google.common.cache.Weigher;
 import com.google.common.collect.Maps;
-import org.iq80.leveldb.DB;
-import org.iq80.leveldb.DBIterator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -48,8 +46,9 @@
 import org.apache.spark.network.shuffle.checksum.Cause;
 import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper;
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
-import org.apache.spark.network.util.LevelDBProvider;
-import org.apache.spark.network.util.LevelDBProvider.StoreVersion;
+import org.apache.spark.network.shuffledb.DB;
+import org.apache.spark.network.shuffledb.StoreVersion;
+import org.apache.spark.network.util.DBProvider;
 import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.NettyUtils;
 import org.apache.spark.network.util.TransportConf;
@@ -124,7 +123,7 @@ public ShuffleIndexInformation load(String filePath) throws IOException {
       .weigher((Weigher<String, ShuffleIndexInformation>)
         (filePath, indexInfo) -> indexInfo.getRetainedMemorySize())
       .build(indexCacheLoader);
-    db = LevelDBProvider.initLevelDB(this.registeredExecutorFile, CURRENT_VERSION, mapper);
+    db = DBProvider.initDB(this.registeredExecutorFile, CURRENT_VERSION, mapper);
     if (db != null) {
       executors = reloadRegisteredExecutors(db);
     } else {
@@ -457,15 +456,8 @@ static ConcurrentMap<AppExecId, ExecutorShuffleInfo> reloadRegisteredExecutors(D
       throws IOException {
     ConcurrentMap<AppExecId, ExecutorShuffleInfo> registeredExecutors = Maps.newConcurrentMap();
     if (db != null) {
-      DBIterator itr = db.iterator();
-      itr.seek(APP_KEY_PREFIX.getBytes(StandardCharsets.UTF_8));
-      while (itr.hasNext()) {
-        Map.Entry<byte[], byte[]> e = itr.next();
-        String key = new String(e.getKey(), StandardCharsets.UTF_8);
-        if (!key.startsWith(APP_KEY_PREFIX)) {
-          break;
-        }
-        AppExecId id = parseDbAppExecKey(key);
+      for (Map.Entry<String, byte[]> e : db.readKVToMap(APP_KEY_PREFIX).entrySet()) {
+        AppExecId id = parseDbAppExecKey(e.getKey());
         logger.info("Reloading registered executors: " +  id.toString());
         ExecutorShuffleInfo shuffleInfo = mapper.readValue(e.getValue(), ExecutorShuffleInfo.class);
         registeredExecutors.put(id, shuffleInfo);