[SPARK-38888][BUILD][CORE][YARN][DOCS] Add RocksDB support for shuffle service state store

common/network-common/pom.xml

@@ -82,6 +82,10 @@
       <artifactId>leveldbjni-all</artifactId>
       <version>1.8</version>
     </dependency>
+    <dependency>
+      <groupId>org.rocksdb</groupId>
+      <artifactId>rocksdbjni</artifactId>
+    </dependency>
 
     <dependency>
       <groupId>com.fasterxml.jackson.core</groupId>

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

@@ -21,10 +21,10 @@
 
 /**
  * The enum `DBBackend` use to specify a disk-based store used in shuffle service local db.
- * Only LEVELDB is supported now.
+ * Support the use of LevelDB and RocksDB.
  */
 public enum DBBackend {
-    LEVELDB(".ldb");
+    LEVELDB(".ldb"), ROCKSDB(".rdb");
 
     private final String fileSuffix;
 

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

@@ -0,0 +1,71 @@
+/*
+ * 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.IOException;
+
+import com.google.common.base.Throwables;
+import org.rocksdb.RocksDBException;
+
+/**
+ * RocksDB implementation of the local KV storage used to persist the shuffle state.
+ */
+public class RocksDB implements DB {
+    private final org.rocksdb.RocksDB db;
+
+    public RocksDB(org.rocksdb.RocksDB db) {
+      this.db = db;
+    }
+
+    @Override
+    public void put(byte[] key, byte[] value) {
+      try {
+        db.put(key, value);
+      } catch (RocksDBException e) {
+        throw Throwables.propagate(e);
+      }
+    }
+
+    @Override
+    public byte[] get(byte[] key) {
+      try {
+        return db.get(key);
+      } catch (RocksDBException e) {
+        throw Throwables.propagate(e);
+      }
+    }
+
+    @Override
+    public void delete(byte[] key) {
+      try {
+        db.delete(key);
+      } catch (RocksDBException e) {
+        throw Throwables.propagate(e);
+      }
+    }
+
+    @Override
+    public DBIterator iterator() {
+      return new RocksDBIterator(db.newIterator());
+    }
+
+    @Override
+    public void close() throws IOException {
+      db.close();
+    }
+}

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

@@ -0,0 +1,95 @@
+/*
+ * 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.IOException;
+import java.util.AbstractMap;
+import java.util.Map;
+import java.util.NoSuchElementException;
+
+import com.google.common.base.Throwables;
+import org.rocksdb.RocksIterator;
+
+/**
+ * RocksDB implementation of `DBIterator`.
+ */
+public class RocksDBIterator implements DBIterator {
+
+    private final RocksIterator it;
+
+    private boolean checkedNext;
+
+    private boolean closed;
+
+    private Map.Entry<byte[], byte[]> next;
+
+    public RocksDBIterator(RocksIterator it) {
+      this.it = it;
+    }
+
+    @Override
+    public boolean hasNext() {
+      if (!checkedNext && !closed) {
+        next = loadNext();
+        checkedNext = true;
+      }
+      if (!closed && next == null) {
+        try {
+          close();
+        } catch (IOException ioe) {
+          throw Throwables.propagate(ioe);
+        }
+      }
+      return next != null;
+    }
+
+    @Override
+    public Map.Entry<byte[], byte[]> next() {
+      if (!hasNext()) {
+        throw new NoSuchElementException();
+      }
+      checkedNext = false;
+      Map.Entry<byte[], byte[]> ret = next;
+      next = null;
+      return ret;
+    }
+
+    @Override
+    public void close() throws IOException {
+      if (!closed) {
+        it.close();
+        closed = true;
+        next = null;
+      }
+    }
+
+    @Override
+    public void seek(byte[] key) {
+      it.seek(key);
+    }
+
+    private Map.Entry<byte[], byte[]> loadNext() {
+      if (it.isValid()) {
+        Map.Entry<byte[], byte[]> nextEntry =
+          new AbstractMap.SimpleEntry<>(it.key(), it.value());
+        it.next();
+        return nextEntry;
+      }
+      return null;
+    }
+}

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

@@ -22,9 +22,10 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 
+import org.apache.spark.network.shuffledb.DB;
 import org.apache.spark.network.shuffledb.DBBackend;
 import org.apache.spark.network.shuffledb.LevelDB;
-import org.apache.spark.network.shuffledb.DB;
+import org.apache.spark.network.shuffledb.RocksDB;
 import org.apache.spark.network.shuffledb.StoreVersion;
 
 public class DBProvider {
@@ -34,11 +35,13 @@ public static DB initDB(
         StoreVersion version,
         ObjectMapper mapper) throws IOException {
       if (dbFile != null) {
-        // TODO: SPARK-38888, add rocksdb implementation.
         switch (dbBackend) {
           case LEVELDB:
             org.iq80.leveldb.DB levelDB = LevelDBProvider.initLevelDB(dbFile, version, mapper);
             return levelDB != null ? new LevelDB(levelDB) : null;
+          case ROCKSDB:
+            org.rocksdb.RocksDB rocksDB = RocksDBProvider.initRockDB(dbFile, version, mapper);
+            return rocksDB != null ? new RocksDB(rocksDB) : null;
           default:
             throw new IllegalArgumentException("Unsupported DBBackend: " + dbBackend);
         }
@@ -49,11 +52,11 @@ public static DB initDB(
     @VisibleForTesting
     public static DB initDB(DBBackend dbBackend, File file) throws IOException {
       if (file != null) {
-        // TODO: SPARK-38888, add rocksdb implementation.
         switch (dbBackend) {
           case LEVELDB: return new LevelDB(LevelDBProvider.initLevelDB(file));
-        default:
-          throw new IllegalArgumentException("Unsupported DBBackend: " + dbBackend);
+          case ROCKSDB: return new RocksDB(RocksDBProvider.initRocksDB(file));
+          default:
+            throw new IllegalArgumentException("Unsupported DBBackend: " + dbBackend);
         }
       }
       return null;