AmandeepSingh285 commented on code in PR #3695:
URL: https://github.com/apache/celeborn/pull/3695#discussion_r3272731564
##########
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:
Done
--
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]
