Copilot commented on code in PR #3695:
URL: https://github.com/apache/celeborn/pull/3695#discussion_r3270908449
##########
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala:
##########
@@ -4520,6 +4522,17 @@ object CelebornConf extends Logging {
.checkValues(Set("LEVELDB", "ROCKSDB"))
.createWithDefault("ROCKSDB")
+ val WORKER_RECOVER_DB_AUTO_RECOVERY: ConfigEntry[Boolean] =
+ buildConf("celeborn.metadata.autoRecovery.enabled")
+ .categories("worker")
+ .doc("If true, the metadata DB will automatically attempt to recover
from RocksDBException " +
+ "errors during put/get/delete operations. Recovery first tries a safe
reopen; if that " +
+ "fails, recreates the DB." +
Review Comment:
`WORKER_RECOVER_DB_AUTO_RECOVERY` doc string concatenation is missing
whitespace between sentences ("recreates the DB." + "If false...") which will
render as "DB.If" in generated docs. Add a space either at the end of the first
string or the beginning of the next one.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBProvider.java:
##########
@@ -48,26 +48,48 @@ public class RocksDBProvider {
private static final Logger logger =
LoggerFactory.getLogger(RocksDBProvider.class);
+ private static Options createDBOptions() {
+ 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);
+
+ return dbOptions;
+ }
+
+ /**
+ * Reopen an existing RocksDB without the delete-and-recreate fallback. Use
this for recovery from
+ * transient errors.
+ */
+ public static org.rocksdb.RocksDB reopenRocksDB(File dbFile) throws
IOException {
+ if (dbFile == null || !dbFile.exists()) {
+ throw new IOException("RocksDB path does not exist: " + dbFile);
+ }
+ Options dbOptions = createDBOptions();
+ try {
+ return org.rocksdb.RocksDB.open(dbOptions, dbFile.toString());
+ } catch (RocksDBException e) {
+ throw new IOException("Failed to reopen RocksDB at " + dbFile, e);
+ }
Review Comment:
`reopenRocksDB` allocates RocksDB native resources (`Options`,
`BloomFilter`, `BlockBasedTableConfig`, logger) but never closes them on
failure, which can leak native memory when recovery repeatedly attempts reopen.
Ensure these handles are closed at least on the exception path (and consider a
longer-term ownership model so they’re closed on success as well).
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java:
##########
@@ -33,20 +34,32 @@ public class RocksDBIterator implements DBIterator {
private final RocksIterator it;
private final MetadataMetrics metrics;
+ private final AtomicLong dbGeneration;
+ private final long creationGeneration;
private boolean checkedNext;
private boolean closed;
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;
+ }
+
+ private void checkGeneration() {
+ if (dbGeneration.get() != creationGeneration) {
+ throw new IllegalStateException("DB instance was recreated, iterator is
stale");
Review Comment:
`checkGeneration()` throws on stale iterators without closing the underlying
`RocksIterator`. Since `RocksIterator` holds native resources, this can leak if
callers hit the exception and skip `close()`. Consider closing `it` (or marking
closed) before throwing when generation mismatches.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java:
##########
@@ -17,63 +17,213 @@
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.metrics.source.AbstractSource;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
/**
* 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. Recovery first attempts a safe
reopen; if that fails, it
+ * falls back to recreating the DB. When {@code autoRecoveryEnabled} is {@code
false} (the default),
+ * 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 org.rocksdb.RocksDB 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 StoreVersion version;
+ private final boolean autoRecoveryEnabled;
+ private volatile boolean closed = false;
- public RocksDB(org.rocksdb.RocksDB db, AbstractSource source, DBBackend
dbBackend) {
+ public RocksDB(
+ org.rocksdb.RocksDB db,
+ AbstractSource source,
+ DBBackend dbBackend,
+ File dbFile,
+ StoreVersion version) {
super(source, dbBackend);
this.db = db;
+ this.dbFile = dbFile;
+ this.version = version;
+ this.autoRecoveryEnabled =
+ (source instanceof WorkerSource) && ((WorkerSource)
source).metadataAutoRecoveryEnabled();
+ }
+
+ private void recreateDBInstance(long failedGeneration) {
+ if (isClosed()) {
+ return;
+ }
+
+ rwLock.writeLock().lock();
+ try {
+ if (dbGeneration.get() != failedGeneration) {
+ logger.info(
+ "RocksDB instance already recovered by another thread (generation
{} -> {})",
+ failedGeneration,
+ dbGeneration.get());
+ return;
+ }
+
+ if (isClosed()) {
+ return;
+ }
+
+ try {
+ if (db != null) {
+ db.close();
+ }
+ } catch (Exception e) {
+ logger.warn("Failed to close RocksDB instance", e);
+ }
+
+ // Phase 1: try safe reopen
+ try {
+ db = RocksDBProvider.reopenRocksDB(dbFile);
+ dbGeneration.incrementAndGet();
+ logger.info("RocksDB instance recovered {}", dbFile);
+ return;
+ } catch (IOException e) {
+ logger.warn("Safe reopen failed for RocksDB at {}", dbFile, e);
+ }
+
+ // Phase 2: recreate
+ try {
+ db = RocksDBProvider.initRockDB(dbFile, version);
+ dbGeneration.incrementAndGet();
+ logger.error("RocksDB {} was recreated.", dbFile);
+ } catch (IOException e) {
+ dbGeneration.incrementAndGet();
+ logger.error("Failed to recreate RocksDB instance 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();
+ rwLock.readLock().lock();
+ boolean unlocked = false;
+ long generation = dbGeneration.get();
+ try {
+ return operation.get();
+ } catch (RocksDBException e) {
Review Comment:
`checkState()` runs before acquiring the read lock in `withRecovery`. A
concurrent `close()` can set `closed=true` and close the native DB after
another thread passes `checkState()` but before it acquires the read lock,
leading to post-close operations throwing `RocksDBException` (or worse) instead
of consistently throwing `IllegalStateException`. Consider checking `closed`
under the lock (e.g., re-check after acquiring the read lock) to make the
closed-state behavior thread-safe.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java:
##########
@@ -17,63 +17,213 @@
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.metrics.source.AbstractSource;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
/**
* 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. Recovery first attempts a safe
reopen; if that fails, it
+ * falls back to recreating the DB. When {@code autoRecoveryEnabled} is {@code
false} (the default),
+ * 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 org.rocksdb.RocksDB 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 StoreVersion version;
+ private final boolean autoRecoveryEnabled;
+ private volatile boolean closed = false;
- public RocksDB(org.rocksdb.RocksDB db, AbstractSource source, DBBackend
dbBackend) {
+ public RocksDB(
+ org.rocksdb.RocksDB db,
+ AbstractSource source,
+ DBBackend dbBackend,
+ File dbFile,
+ StoreVersion version) {
super(source, dbBackend);
this.db = db;
+ this.dbFile = dbFile;
+ this.version = version;
+ this.autoRecoveryEnabled =
+ (source instanceof WorkerSource) && ((WorkerSource)
source).metadataAutoRecoveryEnabled();
+ }
+
+ private void recreateDBInstance(long failedGeneration) {
+ if (isClosed()) {
+ return;
+ }
+
+ rwLock.writeLock().lock();
+ try {
+ if (dbGeneration.get() != failedGeneration) {
+ logger.info(
+ "RocksDB instance already recovered by another thread (generation
{} -> {})",
+ failedGeneration,
+ dbGeneration.get());
+ return;
+ }
+
+ if (isClosed()) {
+ return;
+ }
+
+ try {
+ if (db != null) {
+ db.close();
+ }
+ } catch (Exception e) {
+ logger.warn("Failed to close RocksDB instance", e);
+ }
+
+ // Phase 1: try safe reopen
+ try {
+ db = RocksDBProvider.reopenRocksDB(dbFile);
+ dbGeneration.incrementAndGet();
+ logger.info("RocksDB instance recovered {}", dbFile);
+ return;
+ } catch (IOException e) {
+ logger.warn("Safe reopen failed for RocksDB at {}", dbFile, e);
+ }
+
+ // Phase 2: recreate
+ try {
+ db = RocksDBProvider.initRockDB(dbFile, version);
+ dbGeneration.incrementAndGet();
+ logger.error("RocksDB {} was recreated.", dbFile);
+ } catch (IOException e) {
+ dbGeneration.incrementAndGet();
+ logger.error("Failed to recreate RocksDB instance 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();
+ rwLock.readLock().lock();
+ boolean unlocked = false;
+ long generation = dbGeneration.get();
+ try {
+ return operation.get();
+ } catch (RocksDBException e) {
+ rwLock.readLock().unlock();
+ unlocked = true;
+ if (autoRecoveryEnabled) {
+ recreateDBInstance(generation);
+ }
+ throw e;
+ } finally {
+ if (!unlocked) {
+ rwLock.readLock().unlock();
+ }
+ }
Review Comment:
`withRecovery` recreates the RocksDB instance but always rethrows the
original `RocksDBException` without retrying the failed operation. If the
intent is to make metadata operations recover transparently, consider retrying
the operation once after a successful recovery (or update the config/docs/tests
to explicitly state recovery only applies to subsequent operations).
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java:
##########
@@ -17,63 +17,213 @@
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.metrics.source.AbstractSource;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
/**
* 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. Recovery first attempts a safe
reopen; if that fails, it
+ * falls back to recreating the DB. When {@code autoRecoveryEnabled} is {@code
false} (the default),
+ * 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 org.rocksdb.RocksDB 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 StoreVersion version;
+ private final boolean autoRecoveryEnabled;
+ private volatile boolean closed = false;
- public RocksDB(org.rocksdb.RocksDB db, AbstractSource source, DBBackend
dbBackend) {
+ public RocksDB(
+ org.rocksdb.RocksDB db,
+ AbstractSource source,
+ DBBackend dbBackend,
+ File dbFile,
+ StoreVersion version) {
super(source, dbBackend);
this.db = db;
+ this.dbFile = dbFile;
+ this.version = version;
+ this.autoRecoveryEnabled =
+ (source instanceof WorkerSource) && ((WorkerSource)
source).metadataAutoRecoveryEnabled();
+ }
+
+ private void recreateDBInstance(long failedGeneration) {
+ if (isClosed()) {
+ return;
+ }
+
+ rwLock.writeLock().lock();
+ try {
+ if (dbGeneration.get() != failedGeneration) {
+ logger.info(
+ "RocksDB instance already recovered by another thread (generation
{} -> {})",
+ failedGeneration,
+ dbGeneration.get());
+ return;
+ }
+
+ if (isClosed()) {
+ return;
+ }
+
+ try {
+ if (db != null) {
+ db.close();
+ }
+ } catch (Exception e) {
+ logger.warn("Failed to close RocksDB instance", e);
+ }
+
+ // Phase 1: try safe reopen
+ try {
+ db = RocksDBProvider.reopenRocksDB(dbFile);
+ dbGeneration.incrementAndGet();
+ logger.info("RocksDB instance recovered {}", dbFile);
+ return;
+ } catch (IOException e) {
+ logger.warn("Safe reopen failed for RocksDB at {}", dbFile, e);
+ }
+
+ // Phase 2: recreate
+ try {
+ db = RocksDBProvider.initRockDB(dbFile, version);
+ dbGeneration.incrementAndGet();
+ logger.error("RocksDB {} was recreated.", dbFile);
+ } catch (IOException e) {
+ dbGeneration.incrementAndGet();
+ logger.error("Failed to recreate RocksDB instance at {}. ", dbFile, e);
+ }
Review Comment:
`logger.error("RocksDB {} was recreated")` is used for a normal recovery
path. Logging this at ERROR can create noisy alerts even though the system
self-healed. Consider WARN/INFO for successful recovery and reserve ERROR for
unrecoverable failures.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java:
##########
@@ -17,63 +17,213 @@
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.metrics.source.AbstractSource;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
/**
* 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. Recovery first attempts a safe
reopen; if that fails, it
+ * falls back to recreating the DB. When {@code autoRecoveryEnabled} is {@code
false} (the default),
+ * 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 org.rocksdb.RocksDB 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 StoreVersion version;
+ private final boolean autoRecoveryEnabled;
+ private volatile boolean closed = false;
- public RocksDB(org.rocksdb.RocksDB db, AbstractSource source, DBBackend
dbBackend) {
+ public RocksDB(
+ org.rocksdb.RocksDB db,
+ AbstractSource source,
+ DBBackend dbBackend,
+ File dbFile,
+ StoreVersion version) {
super(source, dbBackend);
this.db = db;
+ this.dbFile = dbFile;
+ this.version = version;
+ this.autoRecoveryEnabled =
+ (source instanceof WorkerSource) && ((WorkerSource)
source).metadataAutoRecoveryEnabled();
+ }
+
+ private void recreateDBInstance(long failedGeneration) {
+ if (isClosed()) {
+ return;
+ }
+
+ rwLock.writeLock().lock();
+ try {
+ if (dbGeneration.get() != failedGeneration) {
+ logger.info(
+ "RocksDB instance already recovered by another thread (generation
{} -> {})",
+ failedGeneration,
+ dbGeneration.get());
+ return;
+ }
+
+ if (isClosed()) {
+ return;
+ }
+
+ try {
+ if (db != null) {
+ db.close();
+ }
+ } catch (Exception e) {
+ logger.warn("Failed to close RocksDB instance", e);
+ }
+
+ // Phase 1: try safe reopen
+ try {
+ db = RocksDBProvider.reopenRocksDB(dbFile);
+ dbGeneration.incrementAndGet();
+ logger.info("RocksDB instance recovered {}", dbFile);
+ return;
+ } catch (IOException e) {
+ logger.warn("Safe reopen failed for RocksDB at {}", dbFile, e);
+ }
+
+ // Phase 2: recreate
+ try {
+ db = RocksDBProvider.initRockDB(dbFile, version);
+ dbGeneration.incrementAndGet();
+ logger.error("RocksDB {} was recreated.", dbFile);
+ } catch (IOException e) {
Review Comment:
Recovery fallback calls `RocksDBProvider.initRockDB`, which (on non-NotFound
open errors) deletes the existing DB files and recreates the store. As
implemented, any `RocksDBException` can therefore trigger destructive
recreation when safe reopen fails. Consider limiting the recreate path to
specific corruption/read-only status codes, or gating destructive recreation
behind a separate explicit config to avoid unexpected data loss.
##########
worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBRecoverySuiteJ.java:
##########
@@ -0,0 +1,181 @@
+/*
+ * 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 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.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.JavaUtils;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
+
+public class RocksDBRecoverySuiteJ {
+
+ private File dbDir;
+ private File dbFile;
+ private WorkerSource workerSource;
+ private WorkerSource workerSourceWithRecovery;
+ private StoreVersion version;
+
+ @Before
+ public void setUp() throws IOException {
+ dbDir = Files.createTempDirectory("rocksdb-recovery-test").toFile();
+ dbFile = new File(dbDir, "test-db");
+ workerSource = new WorkerSource(new CelebornConf());
+
+ CelebornConf confWithRecovery = new CelebornConf();
+ confWithRecovery.set("celeborn.metadata.autoRecovery.enabled", "true");
+ workerSourceWithRecovery = new WorkerSource(confWithRecovery);
+
+ version = new StoreVersion(1, 0);
+ }
+
+ @After
+ public void tearDown() throws IOException {
+ workerSource.destroy();
+ workerSourceWithRecovery.destroy();
+ JavaUtils.deleteRecursively(dbDir);
+ }
+
+ @Test
+ public void testRecoveryAfterCorruption() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource);
+ 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);
+ 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 testConcurrentRecoveryOnlyRecreatesOnce() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource);
+ assertNotNull(db);
+
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+
+ RocksDB rocksDB = (RocksDB) db;
+
+ // Close and reopen to get a working DB, then verify generation tracking
+ db.close();
+ db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version, workerSource);
+ assertNotNull(db);
+ rocksDB = (RocksDB) db;
+
+ assertEquals(0, rocksDB.getDbGeneration());
+
+ db.put(key, value);
+ byte[] result = db.get(key);
+ assertNotNull(result);
+
+ db.close();
+ }
+
+ @Test
+ public void testOperationsAfterCloseDoNotResurrect() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSourceWithRecovery);
+ 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,
workerSourceWithRecovery);
+ 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);
+
+ // Verify iterator works before any recovery
+ assertTrue(iter.hasNext());
+ iter.close();
+
Review Comment:
`testIteratorInvalidatedAfterRecovery` never triggers a recovery event and
therefore does not assert the new stale-iterator behavior. To cover the
iterator invalidation logic, force a recovery (e.g., make an operation throw
`RocksDBException` so generation increments), then assert that the previously
created iterator throws `IllegalStateException` on `hasNext/next/seek`.
##########
worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBRecoverySuiteJ.java:
##########
@@ -0,0 +1,181 @@
+/*
+ * 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 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.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.JavaUtils;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
+
+public class RocksDBRecoverySuiteJ {
+
+ private File dbDir;
+ private File dbFile;
+ private WorkerSource workerSource;
+ private WorkerSource workerSourceWithRecovery;
+ private StoreVersion version;
+
+ @Before
+ public void setUp() throws IOException {
+ dbDir = Files.createTempDirectory("rocksdb-recovery-test").toFile();
+ dbFile = new File(dbDir, "test-db");
+ workerSource = new WorkerSource(new CelebornConf());
+
+ CelebornConf confWithRecovery = new CelebornConf();
+ confWithRecovery.set("celeborn.metadata.autoRecovery.enabled", "true");
+ workerSourceWithRecovery = new WorkerSource(confWithRecovery);
+
+ version = new StoreVersion(1, 0);
+ }
+
+ @After
+ public void tearDown() throws IOException {
+ workerSource.destroy();
+ workerSourceWithRecovery.destroy();
+ JavaUtils.deleteRecursively(dbDir);
+ }
+
+ @Test
+ public void testRecoveryAfterCorruption() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource);
+ 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);
+ 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 testConcurrentRecoveryOnlyRecreatesOnce() throws Exception {
+ DB db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version,
workerSource);
+ assertNotNull(db);
+
+ byte[] key = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] value = "value".getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+
+ RocksDB rocksDB = (RocksDB) db;
+
+ // Close and reopen to get a working DB, then verify generation tracking
+ db.close();
+ db = DBProvider.initDB(DBBackend.ROCKSDB, dbFile, version, workerSource);
+ assertNotNull(db);
+ rocksDB = (RocksDB) db;
+
+ assertEquals(0, rocksDB.getDbGeneration());
+
Review Comment:
`testConcurrentRecoveryOnlyRecreatesOnce` does not exercise recovery or
concurrency (it only closes and reopens a new DB instance and checks
generation==0). To validate the new generation-based single-recreate behavior,
this test should trigger a `RocksDBException` and run concurrent operations
against the *same* `RocksDB` instance, then assert generation increments
exactly once and the DB remains usable afterward.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
