Copilot commented on code in PR #3695: URL: https://github.com/apache/celeborn/pull/3695#discussion_r3280973307
########## worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBRecoverySuiteJ.java: ########## @@ -0,0 +1,228 @@ +/* + * 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 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.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 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.metadata.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); + + // Force a recovery to simulate a RocksDBException scenario + rocksDB.forceRecovery(); + long genAfterFirstRecovery = rocksDB.getDbGeneration(); + assertEquals(1, genAfterFirstRecovery); + + // Now launch concurrent threads that all try to trigger recovery at the same generation + List<Future<?>> futures = new ArrayList<>(); + for (int i = 0; i < threadCount; i++) { + futures.add( + executor.submit( + () -> { + try { + barrier.await(); + rocksDB.forceRecovery(); + } catch (Exception e) { + throw new RuntimeException(e); + } + })); + } + + for (Future<?> f : futures) { + f.get(); + } + executor.shutdown(); + Review Comment: The ExecutorService in this test is only shutdown at the end of the happy path. If any future throws (e.g., barrier await breaks or forceRecovery throws), the test will exit early and leave the fixed thread pool running, which can hang the test JVM. Wrap the submit/get section in a try/finally and always call shutdownNow/shutdown + awaitTermination in the finally block. ########## 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: In tryRecoverDBInstance(), the code closes the current DB before attempting reopen, but only increments dbGeneration when reopen succeeds. If reopen fails, the DB handle has still been closed, yet dbGeneration stays the same, so existing iterators won’t be invalidated by the generation check and may continue using native resources that were just closed. Consider bumping dbGeneration as soon as the DB instance is closed for recovery (or otherwise marking the instance unusable) so stale iterators fail fast even when reopen fails, and to avoid continuing to operate against a closed handle. -- 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]
