platinumhamburg commented on code in PR #2835: URL: https://github.com/apache/fluss/pull/2835#discussion_r3339665498
########## fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/KvSnapshotManager.java: ########## @@ -0,0 +1,1080 @@ +/* + * 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.fluss.server.kv.snapshot; + +import org.apache.fluss.annotation.VisibleForTesting; +import org.apache.fluss.fs.FileSystemSafetyNet; +import org.apache.fluss.fs.FsPath; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.kv.rocksdb.RocksDBKvBuilder; +import org.apache.fluss.utils.CloseableRegistry; +import org.apache.fluss.utils.FileUtils; +import org.apache.fluss.utils.IOUtils; +import org.apache.fluss.utils.MathUtils; +import org.apache.fluss.utils.clock.Clock; +import org.apache.fluss.utils.concurrent.FutureUtils; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +import java.io.Closeable; +import java.io.File; +import java.io.IOException; +import java.nio.file.AtomicMoveNotSupportedException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.LongSupplier; +import java.util.stream.Collectors; + +import static org.apache.fluss.server.kv.snapshot.RocksIncrementalSnapshot.SST_FILE_SUFFIX; +import static org.apache.fluss.utils.Preconditions.checkNotNull; + +/* This file is based on source code of Apache Flink Project (https://flink.apache.org/), licensed by the Apache + * Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE file distributed with this work for + * additional information regarding copyright ownership. */ + +/** + * For a leader replica of PrimaryKey Table, it is a stateless snapshot manager which will trigger + * upload kv snapshot periodically. It'll use a {@link ScheduledExecutorService} to schedule the + * snapshot initialization and a {@link ExecutorService} to complete async phase of snapshot. + * + * <p>For a standby replica of PrimaryKey Table, it will trigger by the + * NotifyKvSnapshotOffsetRequest to incremental download sst files for remote to keep the data up to + * the latest kv snapshot. + * + * <p>For a follower replica of PrimaryKey Table, it will do nothing. + */ +public class KvSnapshotManager implements Closeable { + + private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotManager.class); + + /** Number of consecutive snapshot failures. */ + private final AtomicInteger numberOfConsecutiveFailures; + + /** Whether upload snapshot is started. */ + private volatile boolean isLeader = false; + + /** Whether the replica is a standby replica. */ + private volatile boolean isStandby = false; + + /** Whether the standby replica is initializing. */ + private volatile boolean standbyInitializing = false; Review Comment: Perhaps it would be simpler to use an enum to represent these orthogonal states? ########## fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvSnapshotManagerTest.java: ########## @@ -0,0 +1,1092 @@ +/* + * 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.fluss.server.kv.snapshot; + +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.fs.FsPath; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.kv.KvSnapshotResource; +import org.apache.fluss.server.kv.rocksdb.RocksDBKvBuilder; +import org.apache.fluss.server.zk.NOPErrorHandler; +import org.apache.fluss.server.zk.ZooKeeperClient; +import org.apache.fluss.server.zk.ZooKeeperExtension; +import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService; +import org.apache.fluss.utils.CloseableRegistry; +import org.apache.fluss.utils.clock.ManualClock; +import org.apache.fluss.utils.function.FunctionWithException; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.apache.fluss.shaded.guava32.com.google.common.collect.Iterators.getOnlyElement; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Test for {@link KvSnapshotManager} . */ +class KvSnapshotManagerTest { + @RegisterExtension + public static final AllCallbackWrapper<ZooKeeperExtension> ZOO_KEEPER_EXTENSION_WRAPPER = + new AllCallbackWrapper<>(new ZooKeeperExtension()); + + private static final long periodicMaterializeDelay = 10_000L; + private static ZooKeeperClient zkClient; + private final TableBucket tableBucket = new TableBucket(1, 1); + private ManuallyTriggeredScheduledExecutorService scheduledExecutorService; + private ManuallyTriggeredScheduledExecutorService asyncSnapshotExecutorService; + private KvSnapshotResource kvSnapshotResource; + private KvSnapshotManager kvSnapshotManager; + private DefaultSnapshotContext snapshotContext; + private Configuration conf; + private ManualClock manualClock; + private @TempDir File tmpKvDir; + + @BeforeAll + static void baseBeforeAll() { + zkClient = + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .getZooKeeperClient(NOPErrorHandler.INSTANCE); + } + + @BeforeEach + void before() { + conf = new Configuration(); + conf.set(ConfigOptions.KV_SNAPSHOT_INTERVAL, Duration.ofMillis(periodicMaterializeDelay)); + scheduledExecutorService = new ManuallyTriggeredScheduledExecutorService(); + asyncSnapshotExecutorService = new ManuallyTriggeredScheduledExecutorService(); + ExecutorService dataTransferThreadPool = Executors.newFixedThreadPool(1); + kvSnapshotResource = + new KvSnapshotResource( + scheduledExecutorService, + new KvSnapshotDataUploader(dataTransferThreadPool), + new KvSnapshotDataDownloader(dataTransferThreadPool), + asyncSnapshotExecutorService); + snapshotContext = + DefaultSnapshotContext.create( + zkClient, + new TestingCompletedKvSnapshotCommitter(), + kvSnapshotResource, + conf); + manualClock = new ManualClock(System.currentTimeMillis()); + } + + @AfterEach + void close() { + if (kvSnapshotManager != null) { + kvSnapshotManager.close(); + } + } + + @Test + void testInitialDelay() { + kvSnapshotManager = createSnapshotManager(true); + startPeriodicUploadSnapshot(NopUploadSnapshotTarget.INSTANCE); + checkOnlyOneScheduledTasks(); + } + + @Test + void testInitWithNonPositiveSnapshotInterval() { + conf.set(ConfigOptions.KV_SNAPSHOT_INTERVAL, Duration.ofMillis(0)); + snapshotContext = + DefaultSnapshotContext.create( + zkClient, + new TestingCompletedKvSnapshotCommitter(), + kvSnapshotResource, + conf); + kvSnapshotManager = createSnapshotManager(snapshotContext); + startPeriodicUploadSnapshot(NopUploadSnapshotTarget.INSTANCE); + // periodic snapshot is disabled when periodicMaterializeDelay is not positive + Assertions.assertEquals(0, scheduledExecutorService.getAllScheduledTasks().size()); Review Comment: Use AssertJ instead. ########## fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java: ########## @@ -746,57 +1001,36 @@ private void mayFlushKv(long newHighWatermark) { */ private Optional<CompletedSnapshot> initKvTablet() { checkNotNull(kvManager); + checkNotNull(kvSnapshotManager); long startTime = clock.milliseconds(); LOG.info("Start to init kv tablet for {} of table {}.", tableBucket, physicalPath); - // todo: we may need to handle the following cases: - // case1: no kv files in local, restore from remote snapshot; and apply - // the log; - // case2: kv files in local - // - if no remote snapshot, restore from local and apply the log known to the local - // files. - // - have snapshot, if the known offset to the local files is much less than(maybe - // some value configured) - // the remote snapshot; restore from remote snapshot; - - // currently for simplicity, we'll always download the snapshot files and restore from - // the snapshots as kv files won't exist in our current implementation for - // when replica become follower, we'll always delete the kv files. - - // get the offset from which, we should restore from. default is 0 - long restoreStartOffset = 0; - Optional<CompletedSnapshot> optCompletedSnapshot = getLatestSnapshot(tableBucket); + Optional<CompletedSnapshot> optCompletedSnapshot; try { + optCompletedSnapshot = kvSnapshotManager.downloadLatestSnapshot(); + long restoreStartOffset = 0; Long rowCount; AutoIncIDRange autoIncIDRange; if (optCompletedSnapshot.isPresent()) { LOG.info( - "Use snapshot {} to restore kv tablet for {} of table {}.", - optCompletedSnapshot.get(), + "Init kv tablet for download latest snapshot {} of {} finish, cost {} ms.", Review Comment: The log semantics are a bit strange. A more reasonable description would be Init KvTablet for TableBucket xxx with Snapshot xx. ########## fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java: ########## @@ -311,39 +311,50 @@ public Optional<KvTablet> getKv(TableBucket tableBucket) { return Optional.ofNullable(currentKvs.get(tableBucket)); } - public void dropKv(TableBucket tableBucket) { - KvTablet dropKvTablet = + public void closeOrDropKv(TableBucket tableBucket, boolean needDrop) { Review Comment: Using the name closeThenDropKv or simply closeKv might be more accurate, as the current name looks like a choice between two actions. ########## fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java: ########## @@ -711,24 +940,50 @@ private void createKv() { e); } } - // start periodic kv snapshot - startPeriodicKvSnapshot(snapshotUsed.orElse(null)); + // start periodic upload kv snapshot + startPeriodicUploadKvSnapshot(snapshotUsed.orElse(null)); Review Comment: The Snapshot action consists of three sub-actions: checkpoint + upload + commit, so theoretically the method name should probably not be modified. ########## fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvSnapshotManagerTest.java: ########## @@ -0,0 +1,1092 @@ +/* + * 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.fluss.server.kv.snapshot; + +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.fs.FsPath; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.kv.KvSnapshotResource; +import org.apache.fluss.server.kv.rocksdb.RocksDBKvBuilder; +import org.apache.fluss.server.zk.NOPErrorHandler; +import org.apache.fluss.server.zk.ZooKeeperClient; +import org.apache.fluss.server.zk.ZooKeeperExtension; +import org.apache.fluss.testutils.common.AllCallbackWrapper; +import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService; +import org.apache.fluss.utils.CloseableRegistry; +import org.apache.fluss.utils.clock.ManualClock; +import org.apache.fluss.utils.function.FunctionWithException; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.apache.fluss.shaded.guava32.com.google.common.collect.Iterators.getOnlyElement; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Test for {@link KvSnapshotManager} . */ +class KvSnapshotManagerTest { + @RegisterExtension + public static final AllCallbackWrapper<ZooKeeperExtension> ZOO_KEEPER_EXTENSION_WRAPPER = + new AllCallbackWrapper<>(new ZooKeeperExtension()); + + private static final long periodicMaterializeDelay = 10_000L; + private static ZooKeeperClient zkClient; + private final TableBucket tableBucket = new TableBucket(1, 1); + private ManuallyTriggeredScheduledExecutorService scheduledExecutorService; + private ManuallyTriggeredScheduledExecutorService asyncSnapshotExecutorService; + private KvSnapshotResource kvSnapshotResource; + private KvSnapshotManager kvSnapshotManager; + private DefaultSnapshotContext snapshotContext; + private Configuration conf; + private ManualClock manualClock; + private @TempDir File tmpKvDir; + + @BeforeAll + static void baseBeforeAll() { + zkClient = + ZOO_KEEPER_EXTENSION_WRAPPER + .getCustomExtension() + .getZooKeeperClient(NOPErrorHandler.INSTANCE); + } + + @BeforeEach + void before() { Review Comment: The design of before(0 seems unreasonable, as it causes a lot of resource leaks, repeated initialization, and waste. ########## fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java: ########## @@ -684,6 +818,101 @@ public void updateTieredLogLocalSegments(int tieredLogLocalSegments) { tieredLogLocalSegments); } + /** + * Asynchronously download the latest snapshot for standby replica. + * + * <p>This method submits the snapshot download task to an async thread pool to avoid blocking + * the makeFollower operation. The download can be retried later via + * NotifyKvSnapshotOffsetRequest if it fails. + * + * <p>The returned {@link CompletableFuture} is stored in {@link #standbyDownloadFuture} so that + * any subsequent role transition (standby -> leader / follower) can wait for the in-flight + * download to finish before mutating the same tablet directory. See {@link + * #awaitStandbyDownload()}. + */ + private void becomeStandbyAsync() { + checkNotNull(snapshotContext); + // Wait for any previous standby download to drain before submitting a new one. This + // prevents two downloaders from concurrently writing into the same db/ directory when + // the replica is repeatedly notified to become standby (e.g. across leader epochs). + awaitStandbyDownload(); + synchronized (this) { Review Comment: `becomeStandbyAsync()` calls `awaitStandbyDownload()` **outside** `synchronized(this)`, then assigns `standbyDownloadFuture` **inside** `synchronized(this)`. Between these two steps, `submitStandbyDownload()` (called from `ReplicaManager.notifyKvSnapshotOffset`, which runs **outside** `replicaStateChangeLock`) can slip in, chain a new download future, and start it on the executor. This is reachable through a natural protocol sequence during leader failover: the old leader's last committed snapshot triggers `NotifyKvSnapshotOffsetRequest`, and the leader re-election triggers `NotifyLeaderAndIsrRequest` (re-assigning the same standby). Both RPCs arrive at the TabletServer independently and are handled by different RPC threads. `notifyKvSnapshotOffset` releases `replicaStateChangeLock` before calling `submitStandbyDownload`, so it does not block `makeFollower`. ``` Thread C (notifyKvSnapshotOffset, replicaStateChangeLock already released): │ submitStandbyDownload(X): │ synchronized(this): ← slips into the gap │ F1 = sDF.handleAsync(downloadSnapshot(X)) │ sDF = F1 ← F1 starts executing │ Thread B (makeFollower, holds replicaStateChangeLock + leaderIsrUpdateLock(W)): │ onBecomeNewFollower(): │ awaitStandbyDownload() ← may see old sDF (before F1), returns │ kvSnapshotManager.becomeStandby() ← bumps epoch to E1 │ becomeStandbyAsync(): │ awaitStandbyDownload() ← reads sDF; if F1 was created after the read, misses it │ synchronized(this): │ sDF = runAsync(downloadLatest) → F3 ← overwrites F1 ``` **Result:** F1 is orphaned — no one awaits or tracks it. F1 and F3 execute concurrently, both writing to the same `db/` directory. `downloadEpoch` does not help because both capture the same epoch value (bumped once by `becomeStandby()` before `becomeStandbyAsync()`). `standbyInitializing` is also unreliable since F1 may have passed the check before F3 sets the flag. **Suggestion:** Make `becomeStandbyAsync()` use the same chaining pattern as `submitStandbyDownload()`, so the new download is serialized after any in-flight one: ```java private void becomeStandbyAsync() { synchronized (this) { standbyDownloadFuture = standbyDownloadFuture.handleAsync( (v, prevEx) -> { try { checkNotNull(kvSnapshotManager); kvSnapshotManager.downloadLatestSnapshot(); } catch (Exception e) { LOG.warn("Failed to download snapshot ...", e); } return null; }, snapshotContext.getAsyncOperationsThreadPool()); } } ``` This eliminates the gap between await and assignment, and ensures all downloads — whether from `becomeStandbyAsync` or `submitStandbyDownload` — are serialized through the same future chain. ########## fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/KvSnapshotManager.java: ########## @@ -0,0 +1,1080 @@ +/* + * 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.fluss.server.kv.snapshot; + +import org.apache.fluss.annotation.VisibleForTesting; +import org.apache.fluss.fs.FileSystemSafetyNet; +import org.apache.fluss.fs.FsPath; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.server.kv.rocksdb.RocksDBKvBuilder; +import org.apache.fluss.utils.CloseableRegistry; +import org.apache.fluss.utils.FileUtils; +import org.apache.fluss.utils.IOUtils; +import org.apache.fluss.utils.MathUtils; +import org.apache.fluss.utils.clock.Clock; +import org.apache.fluss.utils.concurrent.FutureUtils; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +import java.io.Closeable; +import java.io.File; +import java.io.IOException; +import java.nio.file.AtomicMoveNotSupportedException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.LongSupplier; +import java.util.stream.Collectors; + +import static org.apache.fluss.server.kv.snapshot.RocksIncrementalSnapshot.SST_FILE_SUFFIX; +import static org.apache.fluss.utils.Preconditions.checkNotNull; + +/* This file is based on source code of Apache Flink Project (https://flink.apache.org/), licensed by the Apache + * Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE file distributed with this work for + * additional information regarding copyright ownership. */ + +/** + * For a leader replica of PrimaryKey Table, it is a stateless snapshot manager which will trigger + * upload kv snapshot periodically. It'll use a {@link ScheduledExecutorService} to schedule the + * snapshot initialization and a {@link ExecutorService} to complete async phase of snapshot. + * + * <p>For a standby replica of PrimaryKey Table, it will trigger by the + * NotifyKvSnapshotOffsetRequest to incremental download sst files for remote to keep the data up to + * the latest kv snapshot. + * + * <p>For a follower replica of PrimaryKey Table, it will do nothing. + */ +public class KvSnapshotManager implements Closeable { + + private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotManager.class); + + /** Number of consecutive snapshot failures. */ + private final AtomicInteger numberOfConsecutiveFailures; + + /** Whether upload snapshot is started. */ + private volatile boolean isLeader = false; + + /** Whether the replica is a standby replica. */ + private volatile boolean isStandby = false; + + /** Whether the standby replica is initializing. */ + private volatile boolean standbyInitializing = false; + + /** + * A supplier to get the current snapshot delay. This allows dynamic reconfiguration of the + * snapshot interval at runtime. + */ + private final LongSupplier snapshotIntervalSupplier; + + private final long initialDelay; + /** The table bucket that the snapshot manager is for. */ + private final TableBucket tableBucket; + + private final File tabletDir; + + private final SnapshotContext snapshotContext; + private final Clock clock; + + /** The target on which the snapshot will be done. */ + private @Nullable UploadSnapshotTarget uploadSnapshotTarget; + + /** + * The scheduled snapshot task. + * + * <p>Since all reads and writes of {@code scheduledTask} are protected by synchronized, the + * volatile modifier is not necessary here. + */ + private ScheduledFuture<?> scheduledTask = null; + + /** The sst files downloaded for standby replicas. */ + private @Nullable Set<Path> downloadedSstFiles; + + private @Nullable Set<Path> downloadedMiscFiles; + private long standbySnapshotSize; + + /** + * Highest snapshot id that has been successfully downloaded by either {@link + * #downloadSnapshot(long)} or {@link #downloadLatestSnapshot()}. Used to make the notification + * path idempotent so that retries / duplicate notifications from coordinator (e.g. across + * coordinator failover) do not redo verify+move on the same SST set. + */ + private long lastDownloadedSnapshotId = -1L; + + /** + * Monotonically increasing epoch bumped on every role transition. Downloads capture the epoch + * at their start and periodically compare against the current value; a mismatch indicates a + * concurrent role change and the download aborts cooperatively. This prevents a timed-out + * standby download from continuing to write into the {@code db/} directory while the new role's + * code path (e.g. leader's loadKv or follower's deleteDirectory) operates on the same files. + */ + private final AtomicInteger downloadEpoch = new AtomicInteger(0); + + protected KvSnapshotManager( + TableBucket tableBucket, File tabletDir, SnapshotContext snapshotContext, Clock clock) { + this.tableBucket = tableBucket; + this.tabletDir = tabletDir; + this.snapshotContext = snapshotContext; + this.numberOfConsecutiveFailures = new AtomicInteger(0); + this.snapshotIntervalSupplier = snapshotContext::getSnapshotIntervalMs; + long periodicSnapshotDelay = snapshotIntervalSupplier.getAsLong(); + this.initialDelay = + periodicSnapshotDelay > 0 + ? MathUtils.murmurHash(tableBucket.hashCode()) % periodicSnapshotDelay + : 0; + this.clock = clock; + this.uploadSnapshotTarget = null; + this.downloadedSstFiles = null; + this.downloadedMiscFiles = null; + this.standbySnapshotSize = 0; + } + + public static KvSnapshotManager create( + TableBucket tableBucket, File tabletDir, SnapshotContext snapshotContext, Clock clock) { + return new KvSnapshotManager(tableBucket, tabletDir, snapshotContext, clock); + } + + public void becomeLeader() { + downloadEpoch.incrementAndGet(); + isLeader = true; + isStandby = false; + // Clear standby download cache when leaving standby role + clearStandbyDownloadCache(); + ensureDbDirectoryExists(); + } + + public void becomeFollower() { + downloadEpoch.incrementAndGet(); + isLeader = false; + isStandby = false; + // Clear standby download cache when leaving standby role + clearStandbyDownloadCache(); + } + + public void becomeStandby() { + downloadEpoch.incrementAndGet(); + isLeader = false; + isStandby = true; + // Clear standby download cache when new added to standby role + clearStandbyDownloadCache(); + ensureDbDirectoryExists(); + } + + @VisibleForTesting + public @Nullable Set<Path> getDownloadedSstFiles() { + return downloadedSstFiles; + } + + @VisibleForTesting + public @Nullable Set<Path> getDownloadedMiscFiles() { + return downloadedMiscFiles; + } + + /** + * Clear the standby download cache. + * + * <p>This method should be called when a replica leaves the standby role (becomes a regular + * follower or leader). It clears the cached state of downloaded SST files, misc files, and + * snapshot size. This ensures that if the replica becomes standby again later, it will perform + * a fresh download based on the actual local files, rather than reusing stale cache that + * references deleted files. + */ + private void clearStandbyDownloadCache() { + downloadedSstFiles = null; + downloadedMiscFiles = null; + standbySnapshotSize = 0; + // Reset idempotency cursor: when the replica leaves standby, the next standby + // promotion may need to re-download a snapshot whose id is <= the previous one + // (e.g. coordinator failover replays an older snapshot id). Forgetting the cursor + // here ensures we don't accidentally short-circuit a legitimate fresh download. + lastDownloadedSnapshotId = -1L; + LOG.info( + "Cleared standby download cache for table bucket {}, will reload from local files on next standby promotion", + tableBucket); + } + + private Path getKvDbPath() { + return tabletDir.toPath().resolve(RocksDBKvBuilder.DB_INSTANCE_DIR_STRING); + } + + /** Ensure the RocksDB data directory exists. */ + private void ensureDbDirectoryExists() { + Path kvDbPath = getKvDbPath(); + if (!kvDbPath.toFile().exists()) { + kvDbPath.toFile().mkdirs(); + } + } + + /** + * The guardedExecutor is an executor that uses to trigger upload snapshot. + * + * <p>It's expected to be passed with a guarded executor to prevent any concurrent modification + * to KvTablet during trigger snapshotting. + */ + public void startPeriodicUploadSnapshot( + Executor guardedExecutor, UploadSnapshotTarget uploadSnapshotTarget) { + this.uploadSnapshotTarget = uploadSnapshotTarget; + + if (snapshotIntervalSupplier.getAsLong() > 0) { + // Cancel any stale scheduled task from a previous epoch to prevent + // triggerAllNonPeriodicTasks() from entering an infinite loop when the stale + // task finds no new data and synchronously re-schedules itself. + synchronized (this) { + if (scheduledTask != null && !scheduledTask.isDone()) { + scheduledTask.cancel(true); + } + } + long delay = Math.max(initialDelay, 1); + LOG.info("TableBucket {} starts periodic snapshot", tableBucket); + scheduleNextSnapshot(delay, guardedExecutor); + } + } + + public void downloadSnapshot(long snapshotId) throws Exception { + if (standbyInitializing || !isStandby) { + LOG.info( + "Skip downloading snapshot {} for bucket {} because replica is no " + + "longer standby or is initializing. is standby: {}, is initializing: {}", + snapshotId, + tableBucket, + isStandby, + standbyInitializing); + return; + } + // Idempotent guard: skip duplicate notifications for the same (or older) snapshot id. + if (snapshotId <= lastDownloadedSnapshotId) { + LOG.debug( + "Skip downloading snapshot {} for bucket {} because a snapshot with id >= {} " + + "has already been downloaded.", + snapshotId, + tableBucket, + lastDownloadedSnapshotId); + return; + } + CompletedSnapshot completedSnapshot = + snapshotContext.getCompletedSnapshotProvider(tableBucket, snapshotId); + if (completedSnapshot == null) { + LOG.warn( + "Snapshot {} not found for bucket {}, skip downloading.", + snapshotId, + tableBucket); + return; + } + int epochAtStart = downloadEpoch.get(); + long startTime = clock.milliseconds(); + incrementalDownloadSnapshot(completedSnapshot, epochAtStart); + standbySnapshotSize = completedSnapshot.getSnapshotSize(); + lastDownloadedSnapshotId = snapshotId; + LOG.info( + "Standby download snapshot {} for bucket {} completed in {} ms, snapshot size: {}.", + snapshotId, + tableBucket, + clock.milliseconds() - startTime, + standbySnapshotSize); + } + + /** + * download the latest snapshot. + * + * <p>For a standby replica, it will download the latest snapshot to keep the data up to the + * latest kv snapshot. + * + * @return the latest snapshot + */ + public Optional<CompletedSnapshot> downloadLatestSnapshot() throws Exception { + // standbyInitializing is used to prevent concurrent download via + // downloadSnapshot(snapshotId). + standbyInitializing = true; + try { + // Note: no isStandby check here. This method is called from both: + // 1. initKvTablet() during leader initialization - isStandby is already false + // 2. becomeStandbyAsync() during standby initialization - isStandby is true + // The isStandby guard is only needed in downloadSnapshot(snapshotId) which is + // called from the notification path exclusively for standby replicas. + int epochAtStart = downloadEpoch.get(); + Optional<CompletedSnapshot> latestSnapshot = getLatestSnapshot(); + if (latestSnapshot.isPresent()) { + CompletedSnapshot completedSnapshot = latestSnapshot.get(); + long startTime = clock.milliseconds(); + incrementalDownloadSnapshot(completedSnapshot, epochAtStart); + standbySnapshotSize = completedSnapshot.getSnapshotSize(); + lastDownloadedSnapshotId = + Math.max(lastDownloadedSnapshotId, completedSnapshot.getSnapshotID()); + LOG.info( + "Downloaded latest snapshot {} for bucket {} in {} ms, " + + "snapshot log offset: {}, snapshot size: {}.", + completedSnapshot.getSnapshotID(), + tableBucket, + clock.milliseconds() - startTime, + completedSnapshot.getLogOffset(), + standbySnapshotSize); + } else { + LOG.info( + "No snapshot found for bucket {}, will init kv from scratch.", tableBucket); + } + + return latestSnapshot; + } finally { + standbyInitializing = false; + } + } + + private void incrementalDownloadSnapshot(CompletedSnapshot completedSnapshot, int expectedEpoch) + throws Exception { + if (downloadedSstFiles == null || downloadedMiscFiles == null) { + // first try to load all ready exists sst files. + downloadedSstFiles = new HashSet<>(); + downloadedMiscFiles = new HashSet<>(); + loadKvLocalFiles(downloadedSstFiles, downloadedMiscFiles); + } + + checkDownloadNotCancelled(expectedEpoch); + + Set<Path> sstFilesToDelete = new HashSet<>(); + KvSnapshotHandle incrementalKvSnapshotHandle = + getIncrementalKvSnapshotHandle( + completedSnapshot, downloadedSstFiles, sstFilesToDelete); + CompletedSnapshot incrementalSnapshot = + completedSnapshot.getIncrementalSnapshot(incrementalKvSnapshotHandle); + + // Use atomic download: download to temp dir first, then atomically move to final location + atomicDownloadSnapshot(incrementalSnapshot, sstFilesToDelete, expectedEpoch); + + // Update cached file sets with absolute paths (resolved against the db directory) + // to stay consistent with loadKvLocalFiles() which returns absolute paths. + Path kvDbPath = getKvDbPath(); + KvSnapshotHandle kvSnapshotHandle = completedSnapshot.getKvSnapshotHandle(); + downloadedSstFiles = + kvSnapshotHandle.getSharedKvFileHandles().stream() + .map(handler -> kvDbPath.resolve(handler.getLocalPath())) + .collect(Collectors.toSet()); + downloadedMiscFiles = + kvSnapshotHandle.getPrivateFileHandles().stream() + .map(handler -> kvDbPath.resolve(handler.getLocalPath())) + .collect(Collectors.toSet()); + } + + /** + * Check whether the download epoch has changed since the download started. If so, a role + * transition has occurred and this download should abort to avoid corrupting the db/ directory + * that the new role is about to use. + */ + private void checkDownloadNotCancelled(int expectedEpoch) throws IOException { + if (downloadEpoch.get() != expectedEpoch) { + throw new IOException( + "Download cancelled due to role transition for bucket " + tableBucket); + } + } + + public long getSnapshotSize() { + if (uploadSnapshotTarget != null) { + return uploadSnapshotTarget.getSnapshotSize(); + } else { + return 0L; + } + } + + public long getStandbySnapshotSize() { + return standbySnapshotSize; + } + + public File getTabletDir() { + return tabletDir; + } + + public Optional<CompletedSnapshot> getLatestSnapshot() { + try { + return Optional.ofNullable( + snapshotContext.getLatestCompletedSnapshotProvider().apply(tableBucket)); + } catch (Exception e) { + LOG.warn("Get latest completed snapshot for {} failed.", tableBucket, e); + } + return Optional.empty(); + } + + // schedule thread and asyncOperationsThreadPool can access this method + private synchronized void scheduleNextSnapshot(long delay, Executor guardedExecutor) { + ScheduledExecutorService snapshotScheduler = snapshotContext.getSnapshotScheduler(); + if (isLeader && !snapshotScheduler.isShutdown()) { + LOG.debug( + "TableBucket {} schedules the next snapshot in {} seconds", + tableBucket, + delay / 1000); + scheduledTask = + snapshotScheduler.schedule( + () -> triggerUploadSnapshot(guardedExecutor), + delay, + TimeUnit.MILLISECONDS); + } + } + + @VisibleForTesting + public long currentSnapshotId() { + return uploadSnapshotTarget.currentSnapshotId(); + } + + /** Trigger upload local snapshot to remote storage. */ + public void triggerUploadSnapshot(Executor guardedExecutor) { + // todo: consider shrink the scope + // of using guardedExecutor + guardedExecutor.execute( + () -> { + if (isLeader) { + LOG.debug("TableBucket {} triggers snapshot.", tableBucket); + long triggerTime = System.currentTimeMillis(); + + Optional<SnapshotRunnable> snapshotRunnableOptional; + try { + checkNotNull(uploadSnapshotTarget); + snapshotRunnableOptional = uploadSnapshotTarget.initSnapshot(); + } catch (Exception e) { + LOG.error("Fail to init snapshot during triggering snapshot.", e); + return; + } + if (snapshotRunnableOptional.isPresent()) { + SnapshotRunnable runnable = snapshotRunnableOptional.get(); + snapshotContext + .getAsyncOperationsThreadPool() + .execute( + () -> + asyncUploadSnapshotPhase( + triggerTime, + runnable.getSnapshotId(), + runnable.getCoordinatorEpoch(), + runnable.getBucketLeaderEpoch(), + runnable.getSnapshotLocation(), + runnable.getSnapshotRunnable(), + guardedExecutor)); + } else { + scheduleNextSnapshot(guardedExecutor); + LOG.debug( + "TableBucket {} has no data updates since last snapshot, " + + "skip this one and schedule the next one in {} seconds", + tableBucket, + snapshotIntervalSupplier.getAsLong() / 1000); + } + } + }); + } + + private void asyncUploadSnapshotPhase( + long triggerTime, + long snapshotId, + int coordinatorEpoch, + int bucketLeaderEpoch, + SnapshotLocation snapshotLocation, + RunnableFuture<SnapshotResult> snapshotedRunnableFuture, + Executor guardedExecutor) { + uploadSnapshot(snapshotedRunnableFuture) + .whenComplete( + (snapshotResult, throwable) -> { + // if succeed + if (throwable == null) { + numberOfConsecutiveFailures.set(0); + + try { + checkNotNull(uploadSnapshotTarget); + uploadSnapshotTarget.handleSnapshotResult( + snapshotId, + coordinatorEpoch, + bucketLeaderEpoch, + snapshotLocation, + snapshotResult); + LOG.info( + "TableBucket {} snapshot {} finished successfully, cost {} ms.", + tableBucket, + snapshotId, + System.currentTimeMillis() - triggerTime); + } catch (Throwable t) { + LOG.warn( + "Fail to handle snapshot result during snapshot of TableBucket {}", + tableBucket, + t); + } + scheduleNextSnapshot(guardedExecutor); + } else { + // if failed + notifyFailureOrCancellation( + snapshotId, snapshotLocation, throwable); + int retryTime = numberOfConsecutiveFailures.incrementAndGet(); + LOG.info( + "TableBucket {} asynchronous part of snapshot is not completed for the {} time.", + tableBucket, + retryTime, + throwable); + + scheduleNextSnapshot(guardedExecutor); + } + }); + } + + private void notifyFailureOrCancellation( + long snapshot, SnapshotLocation snapshotLocation, Throwable cause) { + LOG.warn("TableBucket {} snapshot {} failed.", tableBucket, snapshot, cause); + checkNotNull(uploadSnapshotTarget); + uploadSnapshotTarget.handleSnapshotFailure(snapshot, snapshotLocation, cause); + } + + private CompletableFuture<SnapshotResult> uploadSnapshot( + RunnableFuture<SnapshotResult> snapshotedRunnableFuture) { + + FileSystemSafetyNet.initializeSafetyNetForThread(); + CompletableFuture<SnapshotResult> result = new CompletableFuture<>(); + try { + FutureUtils.runIfNotDoneAndGet(snapshotedRunnableFuture); + + LOG.debug("TableBucket {} finishes asynchronous part of snapshot.", tableBucket); + + result.complete(snapshotedRunnableFuture.get()); + } catch (Exception e) { + result.completeExceptionally(e); + discardFailedUploads(snapshotedRunnableFuture); + } finally { + FileSystemSafetyNet.closeSafetyNetAndGuardedResourcesForThread(); + } + + return result; + } + + private void discardFailedUploads(RunnableFuture<SnapshotResult> snapshotedRunnableFuture) { + LOG.info("TableBucket {} cleanup asynchronous runnable for snapshot.", tableBucket); + + if (snapshotedRunnableFuture != null) { + // snapshot has started + if (!snapshotedRunnableFuture.cancel(true)) { + try { + SnapshotResult snapshotResult = snapshotedRunnableFuture.get(); + if (snapshotResult != null) { + snapshotResult.getKvSnapshotHandle().discard(); + FsPath remoteSnapshotPath = snapshotResult.getSnapshotPath(); + remoteSnapshotPath.getFileSystem().delete(remoteSnapshotPath, true); + } + } catch (Exception ex) { + LOG.debug( + "TableBucket {} cancelled execution of snapshot future runnable. Cancellation produced the following exception, which is expected and can be ignored.", + tableBucket, + ex); + } + } + } + } + + private void scheduleNextSnapshot(Executor guardedExecutor) { + scheduleNextSnapshot(snapshotIntervalSupplier.getAsLong(), guardedExecutor); + } + + private void loadKvLocalFiles(Set<Path> downloadedSstFiles, Set<Path> downloadedMiscFiles) + throws Exception { + if (tabletDir.exists()) { + Path kvDbPath = getKvDbPath(); + Path[] files = FileUtils.listDirectory(kvDbPath); + for (Path filePath : files) { + final String fileName = filePath.getFileName().toString(); + if (fileName.endsWith(SST_FILE_SUFFIX)) { + downloadedSstFiles.add(filePath); + } else { + downloadedMiscFiles.add(filePath); + } + } + LOG.info( + "Loaded {} local sst files and {} misc files from {} for bucket {}.", + downloadedSstFiles.size(), + downloadedMiscFiles.size(), + kvDbPath, + tableBucket); + } else { + LOG.info("No local tablet directory found for bucket {}, starting fresh.", tableBucket); + } + } + + private KvSnapshotHandle getIncrementalKvSnapshotHandle( + CompletedSnapshot completedSnapshot, + Set<Path> downloadedSstFiles, + Set<Path> sstFilesToDelete) { + // get downloaded sst files name to path. + Map<String, Path> downloadedSstFilesMap = new HashMap<>(); + for (Path sstPath : downloadedSstFiles) { + downloadedSstFilesMap.put(sstPath.getFileName().toString(), sstPath); + } + + KvSnapshotHandle completedSnapshotHandler = completedSnapshot.getKvSnapshotHandle(); + List<KvFileHandleAndLocalPath> sstFileHandles = + completedSnapshotHandler.getSharedKvFileHandles(); + List<KvFileHandleAndLocalPath> privateFileHandles = + completedSnapshotHandler.getPrivateFileHandles(); + + List<KvFileHandleAndLocalPath> incrementalSstFileHandles = new ArrayList<>(); + Set<String> downloadedSstFileNames = downloadedSstFilesMap.keySet(); + for (KvFileHandleAndLocalPath sstFileHandle : sstFileHandles) { + String remoteName = sstFileHandle.getLocalPath(); + if (!downloadedSstFileNames.contains(remoteName)) { + // File doesn't exist locally, needs download. + incrementalSstFileHandles.add(sstFileHandle); + } else { + // File exists locally with the same name. Verify the size matches to detect + // files from different RocksDB instances that happen to share the same file + // number (e.g., after leader re-election to a different server). + Path localPath = downloadedSstFilesMap.get(remoteName); + long remoteSize = sstFileHandle.getKvFileHandle().getSize(); + try { + long localSize = Files.size(localPath); + if (localSize != remoteSize) { + LOG.info( + "Local SST file {} has size {} but remote expects size {}, " + + "will re-download for bucket {}", + remoteName, + localSize, + remoteSize, + tableBucket); + incrementalSstFileHandles.add(sstFileHandle); + sstFilesToDelete.add(localPath); + } + } catch (IOException e) { + LOG.warn( + "Failed to read size of local SST file {}, will re-download for bucket {}", + localPath, + tableBucket, + e); + incrementalSstFileHandles.add(sstFileHandle); + sstFilesToDelete.add(localPath); + } + } + } + + Set<String> newSstFileNames = + completedSnapshotHandler.getSharedKvFileHandles().stream() + .map(KvFileHandleAndLocalPath::getLocalPath) + .collect(Collectors.toSet()); + for (String sstFileName : downloadedSstFileNames) { + if (!newSstFileNames.contains(sstFileName)) { + sstFilesToDelete.add(downloadedSstFilesMap.get(sstFileName)); + } + } + + long totalSnapshotSize = + sstFileHandles.stream().mapToLong(f -> f.getKvFileHandle().getSize()).sum() + + privateFileHandles.stream() + .mapToLong(f -> f.getKvFileHandle().getSize()) + .sum(); + long incrementalSnapshotSize = + incrementalSstFileHandles.stream() + .mapToLong(f -> f.getKvFileHandle().getSize()) + .sum() + + privateFileHandles.stream() + .mapToLong(f -> f.getKvFileHandle().getSize()) + .sum(); + int reusedSstFiles = sstFileHandles.size() - incrementalSstFileHandles.size(); + LOG.info( + "Incremental snapshot for bucket {}: remote has {} sst files (total size: {}), " + + "{} reused locally, {} to download, {} to delete, " + + "incremental download size: {}", + tableBucket, + sstFileHandles.size(), + totalSnapshotSize, + reusedSstFiles, + incrementalSstFileHandles.size(), + sstFilesToDelete.size(), + incrementalSnapshotSize); + return new KvSnapshotHandle( + incrementalSstFileHandles, privateFileHandles, incrementalSnapshotSize); + } + + /** + * Atomically download snapshot files to ensure consistency. + * + * <p>This method implements atomic snapshot download by: + * + * <ol> + * <li>Downloading all files to a temporary directory + * <li>Verifying the download completeness + * <li>Deleting obsolete files from the final directory + * <li>Atomically moving new files from temp to final directory + * <li>Cleaning up the temp directory + * </ol> + * + * <p>If any step fails before step 3 (deleting obsolete files), the final directory remains in + * its original consistent state. Note that if a crash occurs between step 3 (deleting obsolete + * files) and step 4 (moving new files), the final directory may be in an incomplete state. This + * is acceptable for standby replicas because they can re-download the full snapshot on restart. + * + * @param incrementalSnapshot the incremental snapshot to download + * @param sstFilesToDelete SST files that should be deleted from the final directory + * @throws IOException if download or file operations fail + */ + private void atomicDownloadSnapshot( + CompletedSnapshot incrementalSnapshot, Set<Path> sstFilesToDelete, int expectedEpoch) + throws IOException { + Path kvTabletDir = tabletDir.toPath(); + Path kvDbPath = getKvDbPath(); + Path tempDownloadDir = + kvTabletDir.resolve(".tmp_snapshot_" + incrementalSnapshot.getSnapshotID()); + + boolean downloadSuccessful = false; + CloseableRegistry closeableRegistry = new CloseableRegistry(); + try { + // Step 1: Create temporary download directory + Files.createDirectories(tempDownloadDir); + LOG.debug( + "Created temporary snapshot download directory {} for bucket {}", + tempDownloadDir, + tableBucket); + + // Step 2: Download all snapshot files to temporary directory + KvSnapshotDownloadSpec downloadSpec = + new KvSnapshotDownloadSpec( + incrementalSnapshot.getKvSnapshotHandle(), tempDownloadDir); + long start = clock.milliseconds(); + LOG.info( + "Start to download kv snapshot {} to temporary directory {}.", + incrementalSnapshot, + tempDownloadDir); + + KvSnapshotDataDownloader kvSnapshotDataDownloader = + snapshotContext.getSnapshotDataDownloader(); + try { + kvSnapshotDataDownloader.transferAllDataToDirectory( + downloadSpec, closeableRegistry); + } catch (Exception e) { + String msg = e.getMessage(); + if (msg != null + && msg.contains(CompletedSnapshot.SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE)) { + try { + snapshotContext.handleSnapshotBroken(incrementalSnapshot); + } catch (Exception t) { + LOG.error("Handle broken snapshot {} failed.", incrementalSnapshot, t); + } + } + throw new IOException("Fail to download kv snapshot to temporary directory.", e); + } + + long downloadTime = clock.milliseconds() - start; + LOG.info( + "Downloaded incremental kv snapshot {} to temp directory in {} ms for bucket {}.", + incrementalSnapshot.getSnapshotID(), + downloadTime, + tableBucket); + + // Step 3: Verify download completeness + verifySnapshotCompleteness(incrementalSnapshot, tempDownloadDir); + + downloadSuccessful = true; + + // Check for role transition before mutating the final directory. + checkDownloadNotCancelled(expectedEpoch); + + // Step 4: Delete obsolete SST files from final directory + for (Path sstFileToDelete : sstFilesToDelete) { + try { + FileUtils.deleteFileOrDirectory(sstFileToDelete.toFile()); + LOG.debug( + "Deleted obsolete SST file {} for bucket {}", + sstFileToDelete, + tableBucket); + } catch (IOException e) { + LOG.warn( + "Failed to delete obsolete SST file {} for bucket {}", + sstFileToDelete, + tableBucket, + e); + // Continue deletion even if one file fails + } + } + + // Step 5: Delete obsolete misc files from final directory + // Use local reference to avoid NPE if clearStandbyDownloadCache() runs concurrently + Set<Path> miscFilesToClean = downloadedMiscFiles; + if (miscFilesToClean != null) { + for (Path miscFileToDelete : miscFilesToClean) { + try { + FileUtils.deleteFileOrDirectory(miscFileToDelete.toFile()); + LOG.debug( + "Deleted obsolete misc file {} for bucket {}", + miscFileToDelete, + tableBucket); + } catch (IOException e) { + LOG.warn( + "Failed to delete obsolete misc file {} for bucket {}", + miscFileToDelete, + tableBucket, + e); + // Continue deletion even if one file fails + } + } + } + + // Check again before moving files into the final directory. + checkDownloadNotCancelled(expectedEpoch); + + // Step 6: Atomically move downloaded files from temp to final directory + moveSnapshotFilesToFinalDirectory(tempDownloadDir, kvDbPath); + + long totalTime = clock.milliseconds() - start; + LOG.info( + "Applied kv snapshot {} for bucket {} in {} ms (download: {} ms, verify+move: {} ms).", + incrementalSnapshot.getSnapshotID(), + tableBucket, + totalTime, + downloadTime, + totalTime - downloadTime); + + } finally { + // Step 7: Clean up closeable registry + IOUtils.closeQuietly(closeableRegistry); + + // Step 8: Clean up temporary directory + if (tempDownloadDir.toFile().exists()) { + try { + FileUtils.deleteDirectory(tempDownloadDir.toFile()); + LOG.debug( + "Cleaned up temporary snapshot directory {} for bucket {} (download {})", + tempDownloadDir, + tableBucket, + downloadSuccessful ? "succeeded" : "failed"); + } catch (IOException e) { + LOG.warn( + "Failed to clean up temporary snapshot directory {} for bucket {}", + tempDownloadDir, + tableBucket, + e); + } + } + } + } + + /** + * Verify that all expected snapshot files have been downloaded to the temporary directory. + * + * @param snapshot the snapshot being verified + * @param tempDir the temporary directory containing downloaded files + * @throws IOException if verification fails + */ + private void verifySnapshotCompleteness(CompletedSnapshot snapshot, Path tempDir) + throws IOException { + KvSnapshotHandle handle = snapshot.getKvSnapshotHandle(); + List<KvFileHandleAndLocalPath> allFiles = new ArrayList<>(); + allFiles.addAll(handle.getSharedKvFileHandles()); + allFiles.addAll(handle.getPrivateFileHandles()); + + for (KvFileHandleAndLocalPath fileHandle : allFiles) { + String fileName = Paths.get(fileHandle.getLocalPath()).getFileName().toString(); + Path expectedFile = tempDir.resolve(fileName); + + if (!Files.exists(expectedFile)) { + throw new IOException( + String.format( + "Snapshot verification failed for bucket %s: expected file %s not found in temp directory %s", + tableBucket, fileName, tempDir)); + } + + long expectedSize = fileHandle.getKvFileHandle().getSize(); + long actualSize = Files.size(expectedFile); + if (expectedSize != actualSize) { + throw new IOException( + String.format( + "Snapshot verification failed for bucket %s: file %s size mismatch (expected: %d, actual: %d)", + tableBucket, fileName, expectedSize, actualSize)); + } + } + + LOG.info( + "Verified completeness of snapshot {} for bucket {}: {} files, total size {} bytes", + snapshot.getSnapshotID(), + tableBucket, + allFiles.size(), + snapshot.getSnapshotSize()); + } + + /** + * Move snapshot files from temporary directory to final RocksDB directory. + * + * <p>This method attempts atomic moves when possible (same filesystem), falling back to + * copy-then-delete if atomic move is not supported. + * + * @param tempDir the temporary directory containing downloaded files + * @param finalDir the final RocksDB db directory + * @throws IOException if file move operations fail + */ + private void moveSnapshotFilesToFinalDirectory(Path tempDir, Path finalDir) throws IOException { + File[] files = tempDir.toFile().listFiles(); + if (files == null || files.length == 0) { + LOG.debug( + "No files to move from temp directory {} to final directory {} for bucket {}", + tempDir, + finalDir, + tableBucket); + return; + } + + int movedCount = 0; + for (File file : files) { Review Comment: There are two `checkDownloadNotCancelled(expectedEpoch)` calls before this loop (before deleting obsolete files and before moving), but none **inside** the loop itself. If `awaitStandbyDownload()` times out and the caller bumps `downloadEpoch`, the download thread will not notice until after the entire move loop completes. In practice this race window is extremely narrow (the move loop is just filesystem metadata ops, milliseconds), but since `checkDownloadNotCancelled` is already used elsewhere as a cooperative cancellation mechanism, adding one inside the loop is zero-cost and closes the theoretical gap: ```java for (File file : files) { checkDownloadNotCancelled(expectedEpoch); Path sourcePath = file.toPath(); Path targetPath = finalDir.resolve(file.getName()); ... ``` -- 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]
