[server] add support of rebuilding KV from local snapshot

fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java

@@ -40,6 +40,8 @@
 import org.apache.fluss.server.kv.autoinc.AutoIncrementManager;
 import org.apache.fluss.server.kv.autoinc.ZkSequenceGeneratorFactory;
 import org.apache.fluss.server.kv.rowmerger.RowMerger;
+import org.apache.fluss.server.kv.snapshot.CompletedSnapshot;
+import org.apache.fluss.server.kv.snapshot.LocalKvSnapshotUtils;
 import org.apache.fluss.server.log.LogManager;
 import org.apache.fluss.server.log.LogTablet;
 import org.apache.fluss.server.metrics.group.TabletServerMetricGroup;
@@ -341,7 +343,7 @@ public KvTablet getOrCreateKv(
      * @param tableBucket the table bucket
      * @return the tablet directory
      */
-    public File createTabletDir(
+    public File deleteAndCreateTabletDir(
             File dataDir, PhysicalTablePath tablePath, TableBucket tableBucket) {
         File tabletDir = getTabletDir(dataDir, tablePath, tableBucket);
 
@@ -351,6 +353,34 @@ public File createTabletDir(
         return tabletDir;
     }
 
+    /**
+     * Attempts to rebuild a tablet directory from the matching retained local snapshot.
+     *
+     * @return the rebuilt tablet directory, or empty when the local snapshot is unavailable or
+     *     invalid
+     */
+    public Optional<File> restoreKvFromLocalSnapshot(
+            File dataDir,
+            PhysicalTablePath tablePath,
+            TableBucket tableBucket,
+            CompletedSnapshot completedSnapshot) {
+        if (!tableBucket.equals(completedSnapshot.getTableBucket())) {
+            return Optional.empty();
+        }
+
+        File tabletDir = getTabletDir(dataDir, tablePath, tableBucket);
+        if (!tabletDir.isDirectory()) {
+            LOG.debug(
+                    "Skip retained local snapshot recovery because KV tablet directory {} "
+                            + "does not exist or is not a directory.",
+                    tabletDir);
+            return Optional.empty();
+        }
+        return LocalKvSnapshotUtils.restore(tabletDir, completedSnapshot)
+                ? Optional.of(tabletDir)
+                : Optional.empty();
+    }
+
     public Optional<KvTablet> getKv(TableBucket tableBucket) {
         return Optional.ofNullable(currentKvs.get(tableBucket));
     }

fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/LocalKvSnapshotUtils.java

@@ -0,0 +1,208 @@
+/*
+ * 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.server.kv.rocksdb.RocksDBKvBuilder;
+import org.apache.fluss.utils.FileUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.LinkOption;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Utilities for retaining, validating, and restoring local KV snapshot checkpoints. */
+public final class LocalKvSnapshotUtils {
+
+    private static final Logger LOG = LoggerFactory.getLogger(LocalKvSnapshotUtils.class);
+
+    private static final String SNAPSHOT_DIRECTORY_PREFIX = "snap-";
+    private static final String RESTORE_DIRECTORY_PREFIX = ".db-restore-";
+
+    private LocalKvSnapshotUtils() {}
+
+    /** Returns the local checkpoint directory for the given snapshot. */
+    public static File getSnapshotDirectory(File kvTabletDir, long snapshotId) {
+        return new File(kvTabletDir, SNAPSHOT_DIRECTORY_PREFIX + snapshotId);
+    }
+
+    /**
+     * Rebuilds the active RocksDB directory from a matching local checkpoint.
+     *
+     * <p>The checkpoint is accepted only when every file referenced by the committed remote
+     * snapshot metadata exists locally with the expected size and there are no unexpected files.
+     * SST files are hard-linked into a temporary RocksDB directory when possible; mutable metadata
+     * files are copied. The retained checkpoint itself is left untouched so another local restart
+     * can use it before a newer snapshot completes.
+     *
+     * @return whether the active RocksDB directory was rebuilt from the local checkpoint
+     */
+    public static boolean restore(File kvTabletDir, CompletedSnapshot completedSnapshot) {
+        long snapshotId = completedSnapshot.getSnapshotID();
+        File snapshotDirectory = getSnapshotDirectory(kvTabletDir, snapshotId);
+        Map<String, Long> expectedFiles = getExpectedFiles(completedSnapshot);
+        if (!isValid(snapshotDirectory.toPath(), expectedFiles)) {
+            if (snapshotDirectory.exists()) {
+                LOG.warn(
+                        "Retained local KV snapshot {} does not match committed snapshot metadata. "
+                                + "Falling back to remote snapshot recovery.",
+                        snapshotDirectory);
+            }
+            return false;
+        }
+
+        Path restoreDirectory = kvTabletDir.toPath().resolve(RESTORE_DIRECTORY_PREFIX + snapshotId);
+        Path activeDbDirectory = RocksDBKvBuilder.getInstanceRocksDBPath(kvTabletDir).toPath();
+        try {
+            FileUtils.deleteDirectory(restoreDirectory.toFile());
+            Files.createDirectories(restoreDirectory);
+            for (String fileName : expectedFiles.keySet()) {
+                Path source = snapshotDirectory.toPath().resolve(fileName);
+                Path target = restoreDirectory.resolve(fileName);
+                copySnapshotFile(source, target);
+            }
+
+            FileUtils.deleteDirectory(activeDbDirectory.toFile());
+            FileUtils.atomicMoveWithFallback(restoreDirectory, activeDbDirectory);
+            retainOnly(kvTabletDir, snapshotId);
+            return true;
+        } catch (Exception e) {
+            LOG.warn(
+                    "Failed to rebuild local KV directory {} from snapshot {}. "
+                            + "Falling back to remote snapshot recovery.",
+                    activeDbDirectory,
+                    snapshotId,
+                    e);
+            FileUtils.deleteDirectoryQuietly(restoreDirectory.toFile());
+            return false;
+        }
+    }
+
+    static void retainOnly(File kvTabletDir, long snapshotId) {
+        File retainedSnapshot = getSnapshotDirectory(kvTabletDir, snapshotId);
+        for (File snapshotDirectory : FileUtils.listDirectories(kvTabletDir)) {
+            if (isSnapshotDirectory(snapshotDirectory)
+                    && !snapshotDirectory.equals(retainedSnapshot)) {
+                deleteQuietly(snapshotDirectory);
+            }
+        }
+    }
+
+    static void discard(File kvTabletDir, long snapshotId) {
+        deleteQuietly(getSnapshotDirectory(kvTabletDir, snapshotId));
+    }
+
+    private static Map<String, Long> getExpectedFiles(CompletedSnapshot completedSnapshot) {
+        Map<String, Long> expectedFiles = new HashMap<>();
+        KvSnapshotHandle snapshotHandle = completedSnapshot.getKvSnapshotHandle();
+        for (KvFileHandleAndLocalPath file : snapshotHandle.getSharedKvFileHandles()) {
+            addExpectedFile(expectedFiles, file);
+        }
+        for (KvFileHandleAndLocalPath file : snapshotHandle.getPrivateFileHandles()) {
+            addExpectedFile(expectedFiles, file);
+        }
+        return expectedFiles;
+    }
+
+    private static void addExpectedFile(
+            Map<String, Long> expectedFiles, KvFileHandleAndLocalPath file) {
+        String localPath = file.getLocalPath();
+        Long previous = expectedFiles.put(localPath, file.getKvFileHandle().getSize());
+        if (previous != null) {
+            // Duplicate paths cannot describe an unambiguous local checkpoint.
+            expectedFiles.put(localPath, -1L);
+        }
+    }
+
+    private static boolean isValid(Path snapshotDirectory, Map<String, Long> expectedFiles) {
+        if (!Files.isDirectory(snapshotDirectory, LinkOption.NOFOLLOW_LINKS)
+                || expectedFiles.isEmpty()) {
+            return false;
+        }
+
+        try {
+            Path[] actualFiles = FileUtils.listDirectory(snapshotDirectory);
+            if (actualFiles.length != expectedFiles.size()) {
+                return false;
+            }
+
+            for (Map.Entry<String, Long> expectedFile : expectedFiles.entrySet()) {
+                Path relativePath = Paths.get(expectedFile.getKey()).normalize();
+                if (relativePath.isAbsolute()
+                        || relativePath.getNameCount() != 1
+                        || !expectedFile.getKey().equals(relativePath.toString())
+                        || expectedFile.getValue() < 0) {
+                    return false;
+                }
+
+                Path localFile = snapshotDirectory.resolve(relativePath).normalize();
+                if (!localFile.getParent().equals(snapshotDirectory.normalize())
+                        || !Files.isRegularFile(localFile, LinkOption.NOFOLLOW_LINKS)
+                        || Files.size(localFile) != expectedFile.getValue()) {
+                    return false;
+                }
+            }
+
+            for (Path actualFile : actualFiles) {
+                if (!expectedFiles.containsKey(actualFile.getFileName().toString())) {
+                    return false;
+                }
+            }
+            return true;
+        } catch (Exception e) {
+            LOG.warn("Failed to validate local KV snapshot directory {}.", snapshotDirectory, e);
+            return false;
+        }
+    }
+
+    private static void copySnapshotFile(Path source, Path target) throws IOException {
+        if (source.getFileName().toString().endsWith(RocksIncrementalSnapshot.SST_FILE_SUFFIX)) {
+            try {
+                Files.createLink(target, source);
+                return;
+            } catch (UnsupportedOperationException | IOException linkException) {
+                try {
+                    Files.copy(source, target);
+                    return;
+                } catch (IOException copyException) {
+                    copyException.addSuppressed(linkException);
+                    throw copyException;
+                }
+            }
+        }
+        Files.copy(source, target);
+    }
+
+    private static boolean isSnapshotDirectory(File directory) {
+        return directory.getName().startsWith(SNAPSHOT_DIRECTORY_PREFIX);
+    }
+
+    private static void deleteQuietly(File snapshotDirectory) {
+        try {
+            FileUtils.deleteDirectory(snapshotDirectory);
+        } catch (IOException e) {
+            LOG.warn("Could not properly clean local KV snapshot {}.", snapshotDirectory, e);
+        }
+    }
+}

fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/RocksIncrementalSnapshot.java

@@ -113,12 +113,16 @@ public void notifySnapshotComplete(long completedSnapshotId) {
             uploadedSstFiles.keySet().removeIf(snapshotId -> snapshotId < completedSnapshotId);
             lastCompletedSnapshotId = completedSnapshotId;
         }
+        // The local checkpoint is useful after an in-place restart. Keep the committed snapshot and
+        // remove older or uncommitted checkpoints only after the remote commit succeeds.
+        LocalKvSnapshotUtils.retainOnly(instanceBasePath, completedSnapshotId);
     }
 
     public void notifySnapshotAbort(long abortedSnapshotId) {
         synchronized (uploadedSstFiles) {
             uploadedSstFiles.remove(abortedSnapshotId);
         }
+        LocalKvSnapshotUtils.discard(instanceBasePath, abortedSnapshotId);
     }
 
     @Override
@@ -138,7 +142,7 @@ public NativeRocksDBSnapshotResources syncPrepareResources(long snapshotId) thro
     }
 
     private File prepareLocalSnapshotDirectory(long snapshotId) {
-        return new File(instanceBasePath, "snap-" + snapshotId);
+        return LocalKvSnapshotUtils.getSnapshotDirectory(instanceBasePath, snapshotId);
     }
 
     private PreviousSnapshot getPreviousSnapshot(long snapshotId) {
@@ -323,16 +327,9 @@ protected NativeRocksDBSnapshotResources(
 
         @Override
         public void release() {
-            try {
-                if (snapshotDirectory.exists()) {
-                    LOG.trace(
-                            "Running cleanup for local RocksDB backup directory {}.",
-                            snapshotDirectory);
-                    FileUtils.deleteDirectory(snapshotDirectory);
-                }
-            } catch (IOException e) {
-                LOG.warn("Could not properly cleanup local RocksDB backup directory.", e);
-            }
+            // Do not delete the checkpoint when the asynchronous upload finishes. Its lifecycle is
+            // decided only after the snapshot commit result is known: notifySnapshotComplete keeps
+            // the latest committed checkpoint and notifySnapshotAbort removes a failed candidate.
         }
     }
 

fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java

@@ -787,21 +787,9 @@ private Optional<CompletedSnapshot> initKvTablet() {
         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
+        // Prefer a retained local checkpoint only when it exactly matches the latest committed
+        // snapshot metadata. Otherwise, download the committed snapshot from remote storage.
+        // The default recovery offset is 0 when no committed snapshot exists.
         long restoreStartOffset = 0;
         Optional<CompletedSnapshot> optCompletedSnapshot = getLatestSnapshot(tableBucket);
         try {
@@ -814,12 +802,29 @@ private Optional<CompletedSnapshot> initKvTablet() {
                         tableBucket,
                         physicalPath);
                 CompletedSnapshot completedSnapshot = optCompletedSnapshot.get();
-                // always create a new dir for the kv tablet
-                File tabletDir =
-                        kvManager.createTabletDir(
-                                logTablet.getDataDir(), physicalPath, tableBucket);
-                // down the snapshot to target tablet dir
-                downloadKvSnapshots(completedSnapshot, tabletDir.toPath());
+
+                File tabletDir;
+                long start = System.currentTimeMillis();
+                Optional<File> optionalTabletDir =
+                        kvManager.restoreKvFromLocalSnapshot(
+                                logTablet.getDataDir(),
+                                physicalPath,
+                                tableBucket,
+                                completedSnapshot);
+                if (optionalTabletDir.isPresent()) {
+                    tabletDir = optionalTabletDir.get();
+                    LOG.info(
+                            "Rebuilt kv tablet for {} of table {} from retained local snapshot {} that costs {} ms.",
+                            tableBucket,
+                            physicalPath,
+                            completedSnapshot.getSnapshotID(),
+                            System.currentTimeMillis() - start);
+                } else {
+                    tabletDir =
+                            kvManager.deleteAndCreateTabletDir(
+                                    logTablet.getDataDir(), physicalPath, tableBucket);
+                    downloadKvSnapshots(completedSnapshot, tabletDir.toPath());
+                }
 
                 // as we have downloaded kv files into the tablet dir, now, we can load it
                 kvTablet = kvManager.loadKv(tabletDir, schemaGetter, this::onKvFlushComplete);