[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,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.shuffledb;
+
+import java.io.Closeable;
+
+/**
+ * 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);
+
+    /**
+     * Get which returns a new byte array storing the value associated
+     * with the specified input key if any.
+     */
+    byte[] get(byte[] key);
+
+    /**
+     * Delete the DB entry (if any) for "key".
+     */
+    void delete(byte[] key);
+
+    /**
+     * Return an iterator over the contents of the DB.
+     */
+    DBIterator iterator();
+}

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

@@ -0,0 +1,42 @@
+/*
+ * 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.util.Locale;
+
+/**
+ * The enum `DBBackend` use to specify a disk-based store used in shuffle service local db.
+ * Only LEVELDB is supported now.
+ */
+public enum DBBackend {
+    LEVELDB(".ldb");
+
+    private final String fileSuffix;
+
+    DBBackend(String fileSuffix) {
+      this.fileSuffix = fileSuffix;
+    }
+
+    public String fileName(String prefix) {
+      return prefix + fileSuffix;
+    }
+
+    public static DBBackend byName(String value) {
+      return DBBackend.valueOf(value.toUpperCase(Locale.ROOT));
+    }
+}

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

@@ -0,0 +1,35 @@
+/*
+ * 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.util.Iterator;
+import java.util.Map;
+
+public interface DBIterator extends Iterator<Map.Entry<byte[], byte[]>>, Closeable {
+
+    /**
+     * Position at the first entry in the source whose `key` is at target.
+     */
+    void seek(byte[] key);
+
+    default void remove() {
+        throw new UnsupportedOperationException();
+    }
+
+}

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

@@ -0,0 +1,53 @@
+/*
+ * 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;
+
+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) {
+        db.put(key, value);
+    }
+
+    @Override
+    public byte[] get(byte[] key) {
+       return db.get(key);
+    }
+
+    @Override
+    public void delete(byte[] key) {
+        db.delete(key);
+    }
+
+    @Override
+    public void close() throws IOException {
+        db.close();
+    }
+
+    @Override
+    public DBIterator iterator() {
+        return new LevelDBIterator(db.iterator());
+    }
+}

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

@@ -0,0 +1,93 @@
+/*
+ * 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 com.google.common.base.Throwables;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.NoSuchElementException;
+
+public class LevelDBIterator implements DBIterator {
+
+    private final org.iq80.leveldb.DBIterator it;
+
+    private boolean checkedNext;
+
+    private boolean closed;
+
+    private Map.Entry<byte[], byte[]> next;
+
+    public LevelDBIterator(org.iq80.leveldb.DBIterator 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);
+    }
+
+    @Override
+    public void remove() {
+      throw new UnsupportedOperationException();
+    }
+
+    private Map.Entry<byte[], byte[]> loadNext() {
+        boolean hasNext = it.hasNext();
+        if (!hasNext) {
+            return null;
+        }
+        return it.next();
+    }
+}

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

@@ -0,0 +1,57 @@
+/*
+ * 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;
+
+/**
+ * Used to identify the version of data stored in local shuffle state DB.
+ */
+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,61 @@
+/*
+ * 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.DBBackend;
+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(
+        DBBackend dbBackend,
+        File dbFile,
+        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;
+          default:
+            throw new IllegalArgumentException("Unsupported DBBackend: " + dbBackend);
+        }
+      }
+      return null;
+    }
+
+    @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);
+        }
+      }
+      return null;
+    }
+}