This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 0198bcf53 [CELEBORN-2334] Automatically restore RocksDB in case of
failures
0198bcf53 is described below
commit 0198bcf53a602f585df3a72cc1e5f147dff419a1
Author: AmandeepSingh285 <[email protected]>
AuthorDate: Mon Jun 22 10:48:46 2026 +0800
[CELEBORN-2334] Automatically restore RocksDB in case of failures
### What changes were proposed in this pull request?
The patch re-instantiates RocksDB in case of failures. In the current
implementation, when RocksDB enters a read-only mode due to failures, Celeborn
metadata operations fail and remain blocked until manual intervention or
restart. This pull request adds logic to detect such RocksDB failures and
re-instantiate the RocksDB instance so that metadata operations can recover
automatically and continue functioning without prolonged disruption. RocksDB
can enter a read-only or unusable state [...]
### Why are the changes needed?
Once RocksDB enters a read-only or error state, Celeborn metadata
operations become unavailable because the existing RocksDB instance remains
unusable, which could lead to failures in metadata updates.
### Does this PR resolve a correctness bug?
No.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Unit tests.
Closes #3695 from AmandeepSingh285/auto-recover-rocks-db.
Lead-authored-by: AmandeepSingh285 <[email protected]>
Co-authored-by: amandeeps.28 <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../common/protocol/RocksDBCompressionType.java | 30 +++
.../org/apache/celeborn/common/CelebornConf.scala | 56 ++++
docs/configuration/worker.md | 4 +
.../deploy/worker/shuffledb/DBProvider.java | 11 +-
.../deploy/worker/shuffledb/ManagedRocksDB.java | 107 ++++++++
.../service/deploy/worker/shuffledb/RocksDB.java | 178 ++++++++++++-
.../deploy/worker/shuffledb/RocksDBIterator.java | 25 +-
.../deploy/worker/shuffledb/RocksDBProvider.java | 75 ++++--
.../worker/storage/PartitionFilesSorter.java | 3 +-
.../deploy/worker/storage/StorageManager.scala | 2 +-
.../deploy/worker/shuffledb/DBProviderSuiteJ.java | 6 +-
.../worker/shuffledb/ManagedRocksDBSuiteJ.java | 224 ++++++++++++++++
.../worker/shuffledb/RocksDBRecoverySuiteJ.java | 285 +++++++++++++++++++++
13 files changed, 965 insertions(+), 41 deletions(-)
diff --git
a/common/src/main/java/org/apache/celeborn/common/protocol/RocksDBCompressionType.java
b/common/src/main/java/org/apache/celeborn/common/protocol/RocksDBCompressionType.java
new file mode 100644
index 000000000..49f1005ef
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/protocol/RocksDBCompressionType.java
@@ -0,0 +1,30 @@
+/*
+ * 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.celeborn.common.protocol;
+
+public enum RocksDBCompressionType {
+ NO_COMPRESSION,
+ SNAPPY_COMPRESSION,
+ ZLIB_COMPRESSION,
+ BZLIB2_COMPRESSION,
+ LZ4_COMPRESSION,
+ LZ4HC_COMPRESSION,
+ XPRESS_COMPRESSION,
+ ZSTD_COMPRESSION,
+ DISABLE_COMPRESSION_OPTION;
+}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 668fbcf79..ea6b819fc 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1368,6 +1368,14 @@ class CelebornConf(loadDefaults: Boolean) extends
Cloneable with Logging with Se
def workerGracefulShutdownRecoverPath: String =
get(WORKER_GRACEFUL_SHUTDOWN_RECOVER_PATH)
def workerGracefulShutdownRecoverDbBackend: String =
get(WORKER_GRACEFUL_SHUTDOWN_RECOVER_DB_BACKEND)
+ def workerRecoverDbRocksDBAutoRecoveryEnabled: Boolean =
+ get(WORKER_RECOVERDB_ROCKSDB_AUTORECOVERY_ENABLED)
+ def workerRecoverDbRocksDBCompression: String =
+ get(WORKER_RECOVERDB_ROCKSDB_COMPRESSION)
+ def workerRecoverDbRocksDBBottommostCompression: String =
+ get(WORKER_RECOVERDB_ROCKSDB_BOTTOMMOST_COMPRESSION)
+ def workerRecoverDbRocksDBBloomFilterBitsPerKey: Double =
+ get(WORKER_RECOVERDB_ROCKSDB_BLOOMFILTER_BITSPERKEY)
def workerGracefulShutdownPartitionSorterCloseAwaitTimeMs: Long =
get(WORKER_PARTITION_SORTER_SHUTDOWN_TIMEOUT)
def workerGracefulShutdownFlusherShutdownTimeoutMs: Long =
get(WORKER_FLUSHER_SHUTDOWN_TIMEOUT)
@@ -4533,6 +4541,54 @@ object CelebornConf extends Logging {
.checkValues(Set("LEVELDB", "ROCKSDB"))
.createWithDefault("ROCKSDB")
+ val WORKER_RECOVERDB_ROCKSDB_AUTORECOVERY_ENABLED: ConfigEntry[Boolean] =
+
buildConf("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.autoRecovery.enabled")
+ .categories("worker")
+ .doc("If true, the metadata DB will automatically attempt to recover
from RocksDBException " +
+ "errors during put/get/delete operations. Recovery tries a safe
reopen. " +
+ "If false, RocksDBException errors are propagated directly to the
caller.")
+ .version("0.7.0")
+ .booleanConf
+ .createWithDefault(false)
+
+ private val rocksDBCompressionTypes: Set[String] =
+ RocksDBCompressionType.values().map(_.name()).toSet
+
+ val WORKER_RECOVERDB_ROCKSDB_COMPRESSION: ConfigEntry[String] =
+
buildConf("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.compression")
+ .categories("worker")
+ .doc("Compression type for the recover DB RocksDB upper levels. " +
+ "Must be a valid org.rocksdb.CompressionType value, e.g.
LZ4_COMPRESSION, " +
+ "ZSTD_COMPRESSION, SNAPPY_COMPRESSION, NO_COMPRESSION.")
+ .version("0.7.0")
+ .stringConf
+ .transform(_.toUpperCase(Locale.ROOT))
+ .checkValues(rocksDBCompressionTypes)
+ .createWithDefault(RocksDBCompressionType.LZ4_COMPRESSION.name())
+
+ val WORKER_RECOVERDB_ROCKSDB_BOTTOMMOST_COMPRESSION: ConfigEntry[String] =
+
buildConf("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.bottommostCompression")
+ .categories("worker")
+ .doc("Compression type for the recover DB RocksDB bottommost level. " +
+ "Higher-ratio codecs (e.g. ZSTD_COMPRESSION) trade CPU for disk space
and " +
+ "are typically a good fit since the bottommost level holds the most
data. " +
+ "Must be a valid org.rocksdb.CompressionType value.")
+ .version("0.7.0")
+ .stringConf
+ .transform(_.toUpperCase(Locale.ROOT))
+ .checkValues(rocksDBCompressionTypes)
+ .createWithDefault(RocksDBCompressionType.ZSTD_COMPRESSION.name())
+
+ val WORKER_RECOVERDB_ROCKSDB_BLOOMFILTER_BITSPERKEY: ConfigEntry[Double] =
+
buildConf("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.bloomFilter.bitsPerKey")
+ .categories("worker")
+ .doc("Bits per key for the bloom filter used by the recover DB RocksDB.
" +
+ "Higher values reduce false positives (fewer wasted disk reads) at the
cost " +
+ "of more memory per key.")
+ .version("0.7.0")
+ .doubleConf
+ .createWithDefault(10.0)
+
val WORKER_PARTITION_SORTER_SHUTDOWN_TIMEOUT: ConfigEntry[Long] =
buildConf("celeborn.worker.graceful.shutdown.partitionSorter.shutdownTimeout")
.categories("worker")
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index bb2cec89b..f33a16993 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -105,6 +105,10 @@ license: |
| celeborn.worker.graceful.shutdown.dbDeleteFailurePolicy | IGNORE | false |
Policy for handling DB delete failures during graceful shutdown. THROW: throw
exception, EXIT: trigger graceful shutdown, IGNORE: log error and continue
(default). | 0.7.0 | |
| celeborn.worker.graceful.shutdown.enabled | false | false | When true,
during worker shutdown, the worker will wait for all released slots to be
committed or destroyed. | 0.2.0 | |
| celeborn.worker.graceful.shutdown.partitionSorter.shutdownTimeout | 120s |
false | The wait time of waiting for sorting partition files during worker
graceful shutdown. | 0.2.0 | |
+| celeborn.worker.graceful.shutdown.recoverDb.rocksdb.autoRecovery.enabled |
false | false | If true, the metadata DB will automatically attempt to recover
from RocksDBException errors during put/get/delete operations. Recovery tries a
safe reopen. If false, RocksDBException errors are propagated directly to the
caller. | 0.7.0 | |
+| celeborn.worker.graceful.shutdown.recoverDb.rocksdb.bloomFilter.bitsPerKey |
10.0 | false | Bits per key for the bloom filter used by the recover DB
RocksDB. Higher values reduce false positives (fewer wasted disk reads) at the
cost of more memory per key. | 0.7.0 | |
+| celeborn.worker.graceful.shutdown.recoverDb.rocksdb.bottommostCompression |
ZSTD_COMPRESSION | false | Compression type for the recover DB RocksDB
bottommost level. Higher-ratio codecs (e.g. ZSTD_COMPRESSION) trade CPU for
disk space and are typically a good fit since the bottommost level holds the
most data. Must be a valid org.rocksdb.CompressionType value. | 0.7.0 | |
+| celeborn.worker.graceful.shutdown.recoverDb.rocksdb.compression |
LZ4_COMPRESSION | false | Compression type for the recover DB RocksDB upper
levels. Must be a valid org.rocksdb.CompressionType value, e.g.
LZ4_COMPRESSION, ZSTD_COMPRESSION, SNAPPY_COMPRESSION, NO_COMPRESSION. | 0.7.0
| |
| celeborn.worker.graceful.shutdown.recoverDbBackend | ROCKSDB | false |
Specifies a disk-based store used in local db. ROCKSDB or LEVELDB (deprecated).
| 0.4.0 | |
| celeborn.worker.graceful.shutdown.recoverPath | <tmp>/recover | false
| The path to store DB. | 0.2.0 | |
| celeborn.worker.graceful.shutdown.saveCommittedFileInfo.interval | 5s |
false | Interval for a Celeborn worker to flush committed file infos into DB. |
0.3.1 | |
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
index 8022a62e3..37be632b8 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
@@ -23,6 +23,7 @@ import java.io.IOException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.metrics.source.AbstractSource;
/** Note: code copied from Apache Spark. */
@@ -30,7 +31,11 @@ public class DBProvider {
private static final Logger logger =
LoggerFactory.getLogger(DBProvider.class);
public static DB initDB(
- DBBackend dbBackend, File dbFile, StoreVersion version, AbstractSource
source)
+ DBBackend dbBackend,
+ File dbFile,
+ StoreVersion version,
+ AbstractSource source,
+ CelebornConf conf)
throws IOException {
if (dbFile != null) {
switch (dbBackend) {
@@ -39,8 +44,8 @@ public class DBProvider {
logger.warn("The LEVELDB is deprecated. Please use ROCKSDB
instead.");
return levelDB != null ? new LevelDB(levelDB, source, dbBackend) :
null;
case ROCKSDB:
- org.rocksdb.RocksDB rocksDB = RocksDBProvider.initRockDB(dbFile,
version);
- return rocksDB != null ? new RocksDB(rocksDB, source, dbBackend) :
null;
+ ManagedRocksDB rocksDB = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ return rocksDB != null ? new RocksDB(rocksDB, source, dbBackend,
dbFile, conf) : null;
default:
throw new IllegalArgumentException("Unsupported DBBackend: " +
dbBackend);
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/ManagedRocksDB.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/ManagedRocksDB.java
new file mode 100644
index 000000000..3245f7ead
--- /dev/null
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/ManagedRocksDB.java
@@ -0,0 +1,107 @@
+/*
+ * 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.celeborn.service.deploy.worker.shuffledb;
+
+import java.io.Closeable;
+
+import org.rocksdb.BloomFilter;
+import org.rocksdb.Options;
+import org.rocksdb.RocksDBException;
+import org.rocksdb.RocksIterator;
+import org.rocksdb.WriteOptions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates a {@link org.rocksdb.RocksDB} instance together with the
resources it depends on.
+ * {@link org.rocksdb.RocksDB#close()} releases only the DB handle (and column
family handles); the
+ * {@link Options}, {@link BloomFilter}, and {@link org.rocksdb.Logger} each
own separate off-heap
+ * memory and must be closed explicitly.
+ *
+ * <p>Instances are produced by {@link RocksDBProvider}; callers interact with
the underlying DB
+ * through the forwarding {@code put/get/delete/newIterator} methods and
release everything via a
+ * single {@link #close()} call.
+ */
+public class ManagedRocksDB implements Closeable {
+ private static final Logger logger =
LoggerFactory.getLogger(ManagedRocksDB.class);
+
+ private final Options dbOptions;
+ private final BloomFilter bloomFilter;
+ private final org.rocksdb.Logger rocksDBLogger;
+ private org.rocksdb.RocksDB db;
+
+ ManagedRocksDB(Options dbOptions, BloomFilter bloomFilter,
org.rocksdb.Logger rocksDBLogger) {
+ this.dbOptions = dbOptions;
+ this.bloomFilter = bloomFilter;
+ this.rocksDBLogger = rocksDBLogger;
+ }
+
+ /** Attaches the DB handle */
+ void setDb(org.rocksdb.RocksDB db) {
+ this.db = db;
+ }
+
+ /** Package-private accessor */
+ Options options() {
+ return dbOptions;
+ }
+
+ /** Package-private accessor for {@link RocksDBProvider#checkVersion}. */
+ org.rocksdb.RocksDB db() {
+ return db;
+ }
+
+ public void put(byte[] key, byte[] value) throws RocksDBException {
+ db.put(key, value);
+ }
+
+ public void put(WriteOptions writeOptions, byte[] key, byte[] value) throws
RocksDBException {
+ db.put(writeOptions, key, value);
+ }
+
+ public byte[] get(byte[] key) throws RocksDBException {
+ return db.get(key);
+ }
+
+ public void delete(byte[] key) throws RocksDBException {
+ db.delete(key);
+ }
+
+ public RocksIterator newIterator() {
+ return db.newIterator();
+ }
+
+ @Override
+ public void close() {
+ closeQuietly(db, "RocksDB");
+ closeQuietly(dbOptions, "RocksDB Options");
+ closeQuietly(bloomFilter, "RocksDB BloomFilter");
+ closeQuietly(rocksDBLogger, "RocksDB logger");
+ }
+
+ private static void closeQuietly(AutoCloseable resource, String name) {
+ if (resource == null) {
+ return;
+ }
+ try {
+ resource.close();
+ } catch (Exception e) {
+ logger.warn("Failed to close {}", name, e);
+ }
+ }
+}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
index b45988ca0..02c1b659c 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
@@ -17,63 +17,217 @@
package org.apache.celeborn.service.deploy.worker.shuffledb;
+import java.io.File;
import java.io.IOException;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.metrics.source.AbstractSource;
/**
* RocksDB implementation of the local KV storage used to persist the shuffle
state.
*
+ * <p>This class supports automatic recovery from RocksDB failures when {@code
autoRecoveryEnabled}
+ * is set to {@code true}. When a put/get/delete operation encounters a {@link
RocksDBException},
+ * the DB instance is closed and reopened. If the safe reopen fails, the
exception is propagated.
+ * When {@code autoRecoveryEnabled} is {@code false}, exceptions are
propagated directly without any
+ * recovery attempt.
+ *
+ * <p>Iterators obtained via {@link #iterator()} are invalidated after a
recovery event and will
+ * throw {@link IllegalStateException} on subsequent use.
+ *
* <p>Note: code copied from Apache Spark.
*/
public class RocksDB extends DB {
- private final org.rocksdb.RocksDB db;
+ private static final Logger logger = LoggerFactory.getLogger(RocksDB.class);
+
+ private volatile ManagedRocksDB db;
private final WriteOptions SYNC_WRITE_OPTIONS = new
WriteOptions().setSync(true);
+ private final ReadWriteLock rwLock = new ReentrantReadWriteLock();
+ private final AtomicLong dbGeneration = new AtomicLong(0);
+ private final File dbFile;
+ private final boolean autoRecoveryEnabled;
+ private final CelebornConf conf;
+ private volatile boolean closed = false;
- public RocksDB(org.rocksdb.RocksDB db, AbstractSource source, DBBackend
dbBackend) {
+ public RocksDB(
+ ManagedRocksDB db,
+ AbstractSource source,
+ DBBackend dbBackend,
+ File dbFile,
+ CelebornConf conf) {
super(source, dbBackend);
this.db = db;
+ this.dbFile = dbFile;
+ this.autoRecoveryEnabled =
conf.workerRecoverDbRocksDBAutoRecoveryEnabled();
+ this.conf = conf;
+ }
+
+ /** Attempts to recover the DB by closing and safely reopening it. */
+ private void tryRecoverDBInstance(long failedGeneration) {
+ if (isClosed()) {
+ return;
+ }
+
+ rwLock.writeLock().lock();
+ try {
+ if (dbGeneration.get() != failedGeneration) {
+ logger.info(
+ "Recovery already attempted by another thread (generation {} ->
{}); "
+ + "if DB is still unhealthy, the next operation will retry",
+ failedGeneration,
+ dbGeneration.get());
+ return;
+ }
+
+ if (isClosed()) {
+ return;
+ }
+
+ try {
+ if (db != null) {
+ db.close();
+ }
+ } catch (Exception e) {
+ logger.warn("Failed to close RocksDB instance", e);
+ }
+
+ dbGeneration.incrementAndGet();
+
+ try {
+ db = RocksDBProvider.reopenRocksDB(dbFile, conf);
+ logger.info("RocksDB instance recovered at {}", dbFile);
+ } catch (IOException e) {
+ logger.error("Safe reopen failed for RocksDB at {}. ", dbFile, e);
+ }
+ } finally {
+ rwLock.writeLock().unlock();
+ }
+ }
+
+ private void checkState() {
+ if (isClosed()) {
+ throw new IllegalStateException("DB is closed");
+ }
+ }
+
+ private boolean isClosed() {
+ return closed;
+ }
+
+ @FunctionalInterface
+ interface CheckedSupplier<T> {
+ T get() throws RocksDBException;
+ }
+
+ @FunctionalInterface
+ interface CheckedRunnable {
+ void run() throws RocksDBException;
+ }
+
+ private <T> T withRecovery(CheckedSupplier<T> operation) throws
RocksDBException {
+ checkState();
+ long generation = 0;
+ try {
+ rwLock.readLock().lock();
+ try {
+ if (isClosed()) {
+ throw new IllegalStateException("DB is closed");
+ }
+ generation = dbGeneration.get();
+ return operation.get();
+ } finally {
+ rwLock.readLock().unlock();
+ }
+ } catch (RocksDBException e) {
+ if (autoRecoveryEnabled) {
+ tryRecoverDBInstance(generation);
+ }
+ throw e;
+ }
+ }
+
+ private void runWithRecovery(CheckedRunnable operation) throws
RocksDBException {
+ withRecovery(
+ () -> {
+ operation.run();
+ return null;
+ });
}
@Override
protected void putInternal(byte[] key, byte[] value) throws RocksDBException
{
- db.put(key, value);
+ runWithRecovery(() -> db.put(key, value));
}
@Override
protected void putInternal(byte[] key, byte[] value, boolean sync) throws
RocksDBException {
- if (sync) {
- db.put(SYNC_WRITE_OPTIONS, key, value);
- } else {
- db.put(key, value);
- }
+ runWithRecovery(
+ () -> {
+ if (sync) {
+ db.put(SYNC_WRITE_OPTIONS, key, value);
+ } else {
+ db.put(key, value);
+ }
+ });
}
@Override
protected byte[] getInternal(byte[] key) throws RocksDBException {
- return db.get(key);
+ return withRecovery(() -> db.get(key));
}
@Override
protected void deleteInternal(byte[] key) throws RocksDBException {
- db.delete(key);
+ runWithRecovery(() -> db.delete(key));
}
@Override
protected DBIterator newIterator(MetadataMetrics metrics) {
- return new RocksDBIterator(db.newIterator(), metrics);
+ checkState();
+ rwLock.readLock().lock();
+ try {
+ if (isClosed()) {
+ throw new IllegalStateException("DB is closed");
+ }
+ long generation = dbGeneration.get();
+ return new RocksDBIterator(db.newIterator(), metrics, dbGeneration,
generation);
+ } finally {
+ rwLock.readLock().unlock();
+ }
}
@Override
public void close() throws IOException {
+ rwLock.writeLock().lock();
try {
+ closed = true;
db.close();
} finally {
- // WriteOptions is a native handle; release it even if db.close() throws.
+ rwLock.writeLock().unlock();
SYNC_WRITE_OPTIONS.close();
}
}
+
+ // Visible for testing
+ long getDbGeneration() {
+ return dbGeneration.get();
+ }
+
+ // Visible for testing
+ void forceRecovery() {
+ tryRecoverDBInstance(dbGeneration.get());
+ }
+
+ // Visible for testing
+ void forceRecovery(long atGeneration) {
+ tryRecoverDBInstance(atGeneration);
+ }
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
index cce181dce..c944fb3a8 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
@@ -21,6 +21,7 @@ import java.io.IOException;
import java.util.AbstractMap;
import java.util.Map;
import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicLong;
import org.rocksdb.RocksIterator;
@@ -33,6 +34,8 @@ public class RocksDBIterator implements DBIterator {
private final RocksIterator it;
private final MetadataMetrics metrics;
+ private final AtomicLong dbGeneration;
+ private final long creationGeneration;
private boolean checkedNext;
@@ -40,13 +43,31 @@ public class RocksDBIterator implements DBIterator {
private Map.Entry<byte[], byte[]> next;
- public RocksDBIterator(RocksIterator it, MetadataMetrics metrics) {
+ public RocksDBIterator(
+ RocksIterator it, MetadataMetrics metrics, AtomicLong dbGeneration, long
creationGeneration) {
this.it = it;
this.metrics = metrics;
+ this.dbGeneration = dbGeneration;
+ this.creationGeneration = creationGeneration;
+ }
+
+ // Fail fast on a stale iterator rather than transparently re-acquiring on
the recovered DB:
+ // a mid-iteration recovery means the underlying data was unhealthy enough
to force a reopen,
+ // and silently restarting iteration would hide that signal from the caller
(and risk skipping
+ // or duplicating entries depending on what mutations were in flight).
+ private void checkGeneration() {
+ if (dbGeneration.get() != creationGeneration) {
+ if (!closed) {
+ it.close();
+ closed = true;
+ }
+ throw new IllegalStateException("DB instance was recreated, iterator is
stale");
+ }
}
@Override
public boolean hasNext() {
+ checkGeneration();
if (!checkedNext && !closed) {
next = loadNext();
checkedNext = true;
@@ -63,6 +84,7 @@ public class RocksDBIterator implements DBIterator {
@Override
public Map.Entry<byte[], byte[]> next() {
+ checkGeneration();
if (!hasNext()) {
throw new NoSuchElementException();
}
@@ -83,6 +105,7 @@ public class RocksDBIterator implements DBIterator {
@Override
public void seek(byte[] key) {
+ checkGeneration();
metrics.onRead(
() -> {
it.seek(key);
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
index 03a53ec66..46e962f4d 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java
@@ -33,6 +33,7 @@ import org.rocksdb.Status;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.PbSerDeUtils;
/**
@@ -48,26 +49,55 @@ public class RocksDBProvider {
private static final Logger logger =
LoggerFactory.getLogger(RocksDBProvider.class);
- public static org.rocksdb.RocksDB initRockDB(File dbFile, StoreVersion
version)
+ private static ManagedRocksDB createDBOptions(CelebornConf conf) {
+ BloomFilter fullFilter =
+ new BloomFilter(conf.workerRecoverDbRocksDBBloomFilterBitsPerKey(),
false);
+ BlockBasedTableConfig tableFormatConfig =
+ new BlockBasedTableConfig()
+ .setFilterPolicy(fullFilter)
+ .setEnableIndexCompression(false)
+ .setIndexBlockRestartInterval(8)
+ .setFormatVersion(5);
+
+ Options dbOptions = new Options();
+ RocksDBLogger rocksDBLogger = new RocksDBLogger(dbOptions);
+
+ dbOptions.setCreateIfMissing(false);
+ dbOptions.setBottommostCompressionType(
+
CompressionType.valueOf(conf.workerRecoverDbRocksDBBottommostCompression()));
+
dbOptions.setCompressionType(CompressionType.valueOf(conf.workerRecoverDbRocksDBCompression()));
+ dbOptions.setTableFormatConfig(tableFormatConfig);
+ dbOptions.setLogger(rocksDBLogger);
+
+ return new ManagedRocksDB(dbOptions, fullFilter, rocksDBLogger);
+ }
+
+ /**
+ * Reopen an existing RocksDB without the delete-and-recreate fallback. Use
this for recovery from
+ * transient errors. The returned {@link ManagedRocksDB} owns the DB plus
its open-time native
+ * resources; the caller must close it to release them.
+ */
+ public static ManagedRocksDB reopenRocksDB(File dbFile, CelebornConf conf)
throws IOException {
+ if (dbFile == null || !dbFile.exists()) {
+ throw new IOException("RocksDB path does not exist: " + dbFile);
+ }
+ ManagedRocksDB managedDb = createDBOptions(conf);
+ try {
+ managedDb.setDb(org.rocksdb.RocksDB.open(managedDb.options(),
dbFile.toString()));
+ return managedDb;
+ } catch (RocksDBException e) {
+ managedDb.close();
+ throw new IOException("Failed to reopen RocksDB at " + dbFile, e);
+ }
+ }
+
+ public static ManagedRocksDB initRockDB(File dbFile, StoreVersion version,
CelebornConf conf)
throws IOException {
- org.rocksdb.RocksDB tmpDb = null;
+ ManagedRocksDB managedDb = null;
if (dbFile != null) {
- BloomFilter fullFilter = new BloomFilter(10.0D /*
BloomFilter.DEFAULT_BITS_PER_KEY */, false);
- BlockBasedTableConfig tableFormatConfig =
- new BlockBasedTableConfig()
- .setFilterPolicy(fullFilter)
- .setEnableIndexCompression(false)
- .setIndexBlockRestartInterval(8)
- .setFormatVersion(5);
-
- Options dbOptions = new Options();
- RocksDBLogger rocksDBLogger = new RocksDBLogger(dbOptions);
-
- dbOptions.setCreateIfMissing(false);
- dbOptions.setBottommostCompressionType(CompressionType.ZSTD_COMPRESSION);
- dbOptions.setCompressionType(CompressionType.LZ4_COMPRESSION);
- dbOptions.setTableFormatConfig(tableFormatConfig);
- dbOptions.setLogger(rocksDBLogger);
+ managedDb = createDBOptions(conf);
+ Options dbOptions = managedDb.options();
+ org.rocksdb.RocksDB tmpDb = null;
try {
tmpDb = org.rocksdb.RocksDB.open(dbOptions, dbFile.toString());
@@ -78,6 +108,7 @@ public class RocksDBProvider {
try {
tmpDb = org.rocksdb.RocksDB.open(dbOptions, dbFile.toString());
} catch (RocksDBException dbExc) {
+ managedDb.close();
throw new IOException("Unable to create state store", dbExc);
}
} else {
@@ -102,23 +133,25 @@ public class RocksDBProvider {
try {
tmpDb = org.rocksdb.RocksDB.open(dbOptions, dbFile.toString());
} catch (RocksDBException dbExc) {
+ managedDb.close();
throw new IOException("Unable to create state store", dbExc);
}
}
}
+ managedDb.setDb(tmpDb);
try {
// if there is a version mismatch, we throw an exception, which means
the service
// is unusable
checkVersion(tmpDb, version);
} catch (RocksDBException e) {
- tmpDb.close();
+ managedDb.close();
throw new IOException(e.getMessage(), e);
} catch (IOException ioe) {
- tmpDb.close();
+ managedDb.close();
throw ioe;
}
}
- return tmpDb;
+ return managedDb;
}
private static void createIfMissing(Options dbOptions, File dbFile) {
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
index 5979a8434..f0ac31454 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -124,7 +124,8 @@ public class PartitionFilesSorter extends
ShuffleRecoverHelper {
String recoverySortedFilesFileName =
dbBackend.fileName(RECOVERY_SORTED_FILES_FILE_NAME_PREFIX);
this.recoverFile = new File(recoverPath, recoverySortedFilesFileName);
- this.sortedFilesDb = DBProvider.initDB(dbBackend, recoverFile,
CURRENT_VERSION, source);
+ this.sortedFilesDb =
+ DBProvider.initDB(dbBackend, recoverFile, CURRENT_VERSION, source,
conf);
reloadAndCleanSortedShuffleFiles(this.sortedFilesDb);
} catch (Exception e) {
throw new IllegalStateException(
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 9a2d4a8a7..99f87cd18 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -301,7 +301,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
val dbBackend =
DBBackend.byName(conf.workerGracefulShutdownRecoverDbBackend)
RECOVERY_FILE_NAME = dbBackend.fileName(RECOVERY_FILE_NAME_PREFIX)
val recoverFile = new File(conf.workerGracefulShutdownRecoverPath,
RECOVERY_FILE_NAME)
- this.db = DBProvider.initDB(dbBackend, recoverFile, CURRENT_VERSION,
workerSource)
+ this.db = DBProvider.initDB(dbBackend, recoverFile, CURRENT_VERSION,
workerSource, conf)
reloadAndCleanFileInfos(this.db)
} catch {
case e: Exception =>
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
index d5208e5d1..6af300e4d 100644
---
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
@@ -67,13 +67,15 @@ public class DBProviderSuiteJ {
? Utils.createDirectory(dbDir.getPath(), namePrefix)
: new File(dbDir.getPath(), String.format("%s-%s", namePrefix,
UUID.randomUUID()));
WorkerSource workerSource = new WorkerSource(new CelebornConf());
+ CelebornConf conf = new CelebornConf();
try {
StoreVersion v1 = new StoreVersion(1, 0);
- DBProvider.initDB(dbBackend, dbFile, v1, workerSource).close();
+ DBProvider.initDB(dbBackend, dbFile, v1, workerSource, conf).close();
StoreVersion v2 = new StoreVersion(2, 0);
IOException ioe =
assertThrows(
- IOException.class, () -> DBProvider.initDB(dbBackend, dbFile,
v2, workerSource));
+ IOException.class,
+ () -> DBProvider.initDB(dbBackend, dbFile, v2, workerSource,
conf));
assertTrue(ioe.getMessage().contains("incompatible with current version
StoreVersion[2.0]"));
} finally {
workerSource.destroy();
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/ManagedRocksDBSuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/ManagedRocksDBSuiteJ.java
new file mode 100644
index 000000000..eb1762b16
--- /dev/null
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/ManagedRocksDBSuiteJ.java
@@ -0,0 +1,224 @@
+/*
+ * 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.celeborn.service.deploy.worker.shuffledb;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.rocksdb.BloomFilter;
+import org.rocksdb.CompressionType;
+import org.rocksdb.InfoLogLevel;
+import org.rocksdb.Options;
+import org.rocksdb.RocksIterator;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.JavaUtils;
+
+public class ManagedRocksDBSuiteJ {
+
+ static {
+ // Ensure the native library is loaded for the partial-init test below,
which constructs
+ // RocksDB native resources directly without going through RocksDBProvider.
+ org.rocksdb.RocksDB.loadLibrary();
+ }
+
+ private File dbDir;
+ private File dbFile;
+ private StoreVersion version;
+ private CelebornConf conf;
+
+ @Before
+ public void setUp() throws IOException {
+ dbDir = Files.createTempDirectory("managed-rocksdb-test").toFile();
+ dbFile = new File(dbDir, "test-db");
+ version = new StoreVersion(1, 0);
+ conf = new CelebornConf();
+ }
+
+ @After
+ public void tearDown() throws IOException {
+ JavaUtils.deleteRecursively(dbDir);
+ }
+
+ @Test
+ public void testForwardingPutGetDeleteDelegateToUnderlyingDb() throws
Exception {
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ try {
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+
+ managedDb.put(key, value);
+ assertArrayEquals(value, managedDb.get(key));
+
+ managedDb.delete(key);
+ assertNull(managedDb.get(key));
+ } finally {
+ managedDb.close();
+ }
+ }
+
+ @Test
+ public void testNewIteratorReturnsUsableIterator() throws Exception {
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ try {
+ managedDb.put("a".getBytes(StandardCharsets.UTF_8),
"1".getBytes(StandardCharsets.UTF_8));
+ managedDb.put("b".getBytes(StandardCharsets.UTF_8),
"2".getBytes(StandardCharsets.UTF_8));
+
+ try (RocksIterator iter = managedDb.newIterator()) {
+ int seen = 0;
+ for (iter.seekToFirst(); iter.isValid(); iter.next()) {
+ seen++;
+ }
+ assertTrue("expected at least 2 entries from forwarded iterator, got "
+ seen, seen >= 2);
+ }
+ } finally {
+ managedDb.close();
+ }
+ }
+
+ @Test
+ public void testCloseReleasesUnderlyingNativeHandles() throws Exception {
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ org.rocksdb.RocksDB underlyingDb = managedDb.db();
+ Options underlyingOptions = managedDb.options();
+ assertNotNull(underlyingDb);
+ assertNotNull(underlyingOptions);
+ assertTrue("underlying DB should own native handle pre-close",
underlyingDb.isOwningHandle());
+ assertTrue(
+ "underlying Options should own native handle pre-close",
+ underlyingOptions.isOwningHandle());
+
+ managedDb.close();
+
+ assertFalse(
+ "underlying DB native handle must be released after close",
underlyingDb.isOwningHandle());
+ assertFalse(
+ "underlying Options native handle must be released after close",
+ underlyingOptions.isOwningHandle());
+ }
+
+ @Test
+ public void testCloseIsIdempotent() throws Exception {
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ managedDb.close();
+ // A second close must not throw. The recovery path can race with explicit
close in production,
+ // so this safety is load-bearing.
+ managedDb.close();
+ }
+
+ @Test
+ public void testCloseWithoutSetDbReleasesPartialInitResources() {
+ // Simulates the partial-init state inside RocksDBProvider when
RocksDB.open(...) fails
+ // before setDb is ever called: the four open-time resources are
allocated, but no DB
+ // handle has been attached. close() must still release them and must not
NPE on the
+ // null db reference.
+ Options dbOptions = new Options();
+ BloomFilter bloomFilter = new BloomFilter(10.0D, false);
+ org.rocksdb.Logger rocksDBLogger =
+ new org.rocksdb.Logger(dbOptions.infoLogLevel()) {
+ @Override
+ protected void log(InfoLogLevel infoLogLevel, String logMsg) {}
+ };
+
+ assertTrue(dbOptions.isOwningHandle());
+ assertTrue(bloomFilter.isOwningHandle());
+ assertTrue(rocksDBLogger.isOwningHandle());
+
+ ManagedRocksDB managedDb = new ManagedRocksDB(dbOptions, bloomFilter,
rocksDBLogger);
+ assertNull("db should be unset for the partial-init case", managedDb.db());
+
+ managedDb.close();
+
+ assertFalse(
+ "Options must be released even when db was never attached",
dbOptions.isOwningHandle());
+ assertFalse(
+ "BloomFilter must be released even when db was never attached",
+ bloomFilter.isOwningHandle());
+ assertFalse(
+ "Logger must be released even when db was never attached",
rocksDBLogger.isOwningHandle());
+ }
+
+ @Test
+ public void testDefaultCompressionMatchesPreviousHardcodedValues() throws
Exception {
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ try {
+ assertEquals(CompressionType.LZ4_COMPRESSION,
managedDb.options().compressionType());
+ assertEquals(
+ CompressionType.ZSTD_COMPRESSION,
managedDb.options().bottommostCompressionType());
+ } finally {
+ managedDb.close();
+ }
+ }
+
+ @Test
+ public void testCustomCompressionTypeFromConf() throws Exception {
+
conf.set("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.compression",
"NO_COMPRESSION");
+ conf.set(
+
"celeborn.worker.graceful.shutdown.recoverDb.rocksdb.bottommostCompression",
+ "SNAPPY_COMPRESSION");
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ try {
+ assertEquals(CompressionType.NO_COMPRESSION,
managedDb.options().compressionType());
+ assertEquals(
+ CompressionType.SNAPPY_COMPRESSION,
managedDb.options().bottommostCompressionType());
+ } finally {
+ managedDb.close();
+ }
+ }
+
+ @Test
+ public void testCompressionTypeIsCaseInsensitive() throws Exception {
+
conf.set("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.compression",
"zstd_compression");
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ try {
+ assertEquals(CompressionType.ZSTD_COMPRESSION,
managedDb.options().compressionType());
+ } finally {
+ managedDb.close();
+ }
+ }
+
+ @Test
+ public void testCustomBloomFilterBitsPerKeyOpensCleanly() throws Exception {
+ // Bloom filter bits-per-key is wrapped inside BlockBasedTableConfig and
not directly
+ // readable back via the RocksDB Java API, so the practical assertion is
that the DB
+ // builds and opens without throwing when a non-default value is supplied.
+
conf.set("celeborn.worker.graceful.shutdown.recoverDb.rocksdb.bloomFilter.bitsPerKey",
"15.0");
+ ManagedRocksDB managedDb = RocksDBProvider.initRockDB(dbFile, version,
conf);
+ try {
+ assertNotNull(managedDb.db());
+ byte[] key = "k".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "v".getBytes(StandardCharsets.UTF_8);
+ managedDb.put(key, value);
+ assertArrayEquals(value, managedDb.get(key));
+ } finally {
+ managedDb.close();
+ }
+ }
+}
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBRecoverySuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBRecoverySuiteJ.java
new file mode 100644
index 000000000..670063279
--- /dev/null
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBRecoverySuiteJ.java
@@ -0,0 +1,285 @@
+/*
+ * 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.celeborn.service.deploy.worker.shuffledb;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.rocksdb.RocksDBException;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.JavaUtils;
+import org.apache.celeborn.common.util.ThreadUtils;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
+
+public class RocksDBRecoverySuiteJ {
+
+ private File dbDir;
+ private File dbFile;
+ private CelebornConf defaultConf;
+ private CelebornConf confWithRecovery;
+ private WorkerSource workerSource;
+ private StoreVersion version;
+
+ @Before
+ public void setUp() throws IOException {
+ dbDir = Files.createTempDirectory("rocksdb-recovery-test").toFile();
+ dbFile = new File(dbDir, "test-db");
+ defaultConf = new CelebornConf();
+ confWithRecovery = new CelebornConf();
+ confWithRecovery.set(
+
"celeborn.worker.graceful.shutdown.recoverDb.rocksdb.autoRecovery.enabled",
"true");
+ workerSource = new WorkerSource(defaultConf);
+ version = new StoreVersion(1, 0);
+ }
+
+ @After
+ public void tearDown() throws IOException {
+ workerSource.destroy();
+ JavaUtils.deleteRecursively(dbDir);
+ }
+
+ @Test
+ public void testRecoveryAfterCorruption() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource, defaultConf);
+ assertNotNull(db);
+
+ byte[] key = "test-key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "test-value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+
+ byte[] result = db.get(key);
+ assertNotNull(result);
+ assertEquals("test-value", new String(result, StandardCharsets.UTF_8));
+
+ db.close();
+
+ // Corrupt the DB by overwriting SST files
+ corruptDbFiles(dbFile);
+
+ // Reopen — initRockDB will wipe and recreate since files are corrupt
+ db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version, workerSource,
defaultConf);
+ assertNotNull(db);
+
+ // Data is gone after wipe-and-recreate, but DB is functional
+ byte[] newKey = "new-key".getBytes(StandardCharsets.UTF_8);
+ byte[] newValue = "new-value".getBytes(StandardCharsets.UTF_8);
+ db.put(newKey, newValue);
+
+ result = db.get(newKey);
+ assertNotNull(result);
+ assertEquals("new-value", new String(result, StandardCharsets.UTF_8));
+ db.close();
+ }
+
+ @Test
+ public void testConcurrentRecoveryOnlyReopensOnce() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource, confWithRecovery);
+ assertNotNull(db);
+
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+
+ RocksDB rocksDB = (RocksDB) db;
+ assertEquals(0, rocksDB.getDbGeneration());
+
+ int threadCount = 8;
+ CyclicBarrier barrier = new CyclicBarrier(threadCount);
+ ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+
+ // Trigger recovery directly to verify concurrent recovery deduplication
+ rocksDB.forceRecovery();
+ long genAfterFirstRecovery = rocksDB.getDbGeneration();
+ assertEquals(1, genAfterFirstRecovery);
+
+ // Now launch concurrent threads that all try to trigger recovery at the
same generation.
+ // genAfterFirstRecovery is captured once so every thread calls
forceRecovery with the
+ // same stale-generation value; only one can win the write-lock check and
actually reopen.
+ List<Future<?>> futures = new ArrayList<>();
+ for (int i = 0; i < threadCount; i++) {
+ futures.add(
+ executor.submit(
+ () -> {
+ try {
+ barrier.await();
+ rocksDB.forceRecovery(genAfterFirstRecovery);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }));
+ }
+
+ try {
+ for (Future<?> f : futures) {
+ f.get();
+ }
+ } finally {
+ ThreadUtils.shutdown(executor);
+ }
+
+ // Generation should have incremented exactly once more (all threads saw
the same generation
+ // and only one wins the write lock to perform the actual reopen; the rest
observe the
+ // advanced generation and bail without reopening)
+ assertEquals(genAfterFirstRecovery + 1, rocksDB.getDbGeneration());
+
+ // DB should still be usable
+ db.put(key, value);
+ byte[] result = db.get(key);
+ assertNotNull(result);
+ assertEquals("value", new String(result, StandardCharsets.UTF_8));
+
+ db.close();
+ }
+
+ @Test
+ public void testOperationsAfterCloseDoNotResurrect() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource, confWithRecovery);
+ assertNotNull(db);
+
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+ db.close();
+
+ // All operations after close should throw IllegalStateException
+ DB closedDb = db;
+ assertThrows(IllegalStateException.class, () ->
closedDb.put("k".getBytes(), "v".getBytes()));
+
+ assertThrows(IllegalStateException.class, () ->
closedDb.get("k".getBytes()));
+
+ assertThrows(IllegalStateException.class, () ->
closedDb.delete("k".getBytes()));
+
+ assertThrows(IllegalStateException.class, closedDb::iterator);
+ }
+
+ @Test
+ public void testIteratorInvalidatedAfterRecovery() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource, confWithRecovery);
+ assertNotNull(db);
+
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+
+ // Get an iterator at generation 0
+ DBIterator iter = db.iterator();
+ iter.seek(key);
+ assertTrue(iter.hasNext());
+
+ // Force a recovery so the generation increments
+ RocksDB rocksDB = (RocksDB) db;
+ assertEquals(0, rocksDB.getDbGeneration());
+ rocksDB.forceRecovery();
+ assertEquals(1, rocksDB.getDbGeneration());
+
+ // The stale iterator should throw on hasNext, next, and seek
+ assertThrows(IllegalStateException.class, iter::hasNext);
+ assertThrows(IllegalStateException.class, iter::next);
+ assertThrows(IllegalStateException.class, () -> iter.seek(key));
+
+ // A new iterator should work fine
+ DBIterator newIter = db.iterator();
+ newIter.seek(key);
+ assertTrue(newIter.hasNext());
+ newIter.close();
+
+ db.close();
+ }
+
+ @Test
+ public void testWithRecoveryTriggeredByRocksDBException() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource, confWithRecovery);
+ assertNotNull(db);
+
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+ assertEquals("value", new String(db.get(key), StandardCharsets.UTF_8));
+
+ RocksDB rocksDB = (RocksDB) db;
+ assertEquals(0, rocksDB.getDbGeneration());
+
+ // Swap the internal ManagedRocksDB with a Mockito spy that throws on the
first put.
+ // Recovery will replace this spy with a fresh instance, so subsequent
puts succeed normally.
+ Field dbField = RocksDB.class.getDeclaredField("db");
+ dbField.setAccessible(true);
+ ManagedRocksDB original = (ManagedRocksDB) dbField.get(rocksDB);
+ ManagedRocksDB spied = spy(original);
+ doThrow(new RocksDBException("injected failure"))
+ .doCallRealMethod()
+ .when(spied)
+ .put(any(byte[].class), any(byte[].class));
+ dbField.set(rocksDB, spied);
+
+ RuntimeException thrown =
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ db.put(
+ "k2".getBytes(StandardCharsets.UTF_8),
"v2".getBytes(StandardCharsets.UTF_8)));
+ assertTrue(thrown.getCause() instanceof RocksDBException);
+
+ // Recovery should have incremented the generation
+ assertEquals(1, rocksDB.getDbGeneration());
+
+ // Subsequent operations should succeed on the recovered DB
+ byte[] newKey = "after-recovery".getBytes(StandardCharsets.UTF_8);
+ byte[] newValue = "works".getBytes(StandardCharsets.UTF_8);
+ db.put(newKey, newValue);
+ assertEquals("works", new String(db.get(newKey), StandardCharsets.UTF_8));
+
+ db.close();
+ }
+
+ private void corruptDbFiles(File dir) throws IOException {
+ if (dir.isDirectory()) {
+ File[] files = dir.listFiles();
+ if (files != null) {
+ for (File f : files) {
+ if (f.isFile()
+ && (f.getName().endsWith(".sst")
+ || f.getName().startsWith("MANIFEST-")
+ || f.getName().equals("CURRENT"))) {
+ Files.write(f.toPath(),
"corrupted-data".getBytes(StandardCharsets.UTF_8));
+ }
+ }
+ }
+ }
+ }
+}
