AmandeepSingh285 commented on code in PR #3695:
URL: https://github.com/apache/celeborn/pull/3695#discussion_r3281474061
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java:
##########
@@ -17,63 +17,209 @@
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 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.metadataAutoRecoveryEnabled();
+ }
+
+ /** 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(
+ "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);
+ }
+
+ try {
+ db = RocksDBProvider.reopenRocksDB(dbFile);
+ dbGeneration.incrementAndGet();
+ logger.info("RocksDB instance recovered at {}", dbFile);
+ } catch (IOException e) {
+ logger.error("Safe reopen failed for RocksDB at {}. ", dbFile, e);
+ }
Review Comment:
Good point updated
--
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]
