[server] fix KV snapshot restoring fallback

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManager.java

@@ -203,8 +203,7 @@ private CompletedSnapshotStore createCompletedSnapshotStore(
                 retrievedSnapshots.add(
                         checkNotNull(snapshotStateHandle.retrieveCompleteSnapshot()));
             } catch (Exception e) {
-                if (e.getMessage()
-                        .contains(CompletedSnapshot.SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE)) {
+                if (CompletedSnapshot.isSnapshotDataNotExists(e)) {
                     LOG.error(
                             "Metadata not found for snapshot {} of table bucket {}, maybe snapshot already removed or broken.",
                             snapshotStateHandle.getSnapshotId(),

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

@@ -88,6 +88,9 @@ public class CompletedSnapshot {
 
     public static final String SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE = "No such file or directory";
 
+    private static final String HADOOP_SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE =
+            "File does not exist";
+
     public CompletedSnapshot(
             TableBucket tableBucket,
             long snapshotID,
@@ -163,9 +166,8 @@ public void registerSharedKvFilesAfterRestored(SharedKvFileRegistry sharedKvFile
     }
 
     public CompletableFuture<Void> discardAsync(Executor ioExecutor) {
-        // it'll discard the snapshot files for kv, it'll always discard
-        // the private files; for shared files, only if they're not be registered in
-        // SharedKvRegistry, can the files be deleted.
+        // Always discard private files. Shared files are discarded directly only when the handle
+        // represents a newly created snapshot that has not been registered.
         CompletableFuture<Void> discardKvFuture =
                 FutureUtils.runAsync(kvSnapshotHandle::discard, ioExecutor);
 
@@ -195,6 +197,24 @@ public static FsPath getMetadataFilePath(FsPath snapshotLocation) {
         return new FsPath(snapshotLocation, SNAPSHOT_METADATA_FILE_NAME);
     }
 
+    /**
+     * Returns whether the throwable or any of its causes indicates that snapshot data no longer
+     * exists in remote storage.
+     */
+    public static boolean isSnapshotDataNotExists(Throwable throwable) {
+        Throwable cause = throwable;
+        while (cause != null) {
+            String message = cause.getMessage();
+            if (message != null
+                    && (message.contains(SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE)
+                            || message.contains(HADOOP_SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE))) {
+                return true;
+            }
+            cause = cause.getCause();
+        }
+        return false;
+    }
+
     private void disposeMetadata() throws IOException {
         FsPath metadataFilePath = getMetadataFilePath();
         FileSystem fileSystem = metadataFilePath.getFileSystem();

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

@@ -193,7 +193,7 @@ public CompletedSnapshot deserialize(JsonNode node) {
 
         // construct CompletedSnapshot
         KvSnapshotHandle kvSnapshotHandle =
-                new KvSnapshotHandle(sharedFileHandles, privateFileHandles, incrementalSize);
+                KvSnapshotHandle.restore(sharedFileHandles, privateFileHandles, incrementalSize);
 
         Long rowCount = null;
         if (node.has(ROW_COUNT)) {

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

@@ -27,7 +27,6 @@
 import java.util.stream.Collectors;
 
 import static org.apache.fluss.utils.Preconditions.checkNotNull;
-import static org.apache.fluss.utils.Preconditions.checkState;
 
 /**
  * A handle to the snapshot of a kv tablet. It contains the share file handles and the private file
@@ -46,22 +45,36 @@ public class KvSnapshotHandle {
     private final long incrementalSize;
 
     /**
-     * Once the shared states are registered, it is the {@link SharedKvFileRegistry}'s
-     * responsibility to cleanup those shared states. But in the cases where the state handle is
-     * discarded before performing the registration, the handle should delete all the shared states
-     * created by it.
-     *
-     * <p>This variable is not null iff the handles was registered.
+     * Whether {@link #discard()} should directly delete the shared files. This is true for newly
+     * created handles until registration and false for restored handles.
      */
-    private transient SharedKvFileRegistry sharedKvFileRegistry;
+    private boolean ownsSharedFiles;
 
-    public KvSnapshotHandle(
+    KvSnapshotHandle(
             List<KvFileHandleAndLocalPath> sharedFileHandles,
             List<KvFileHandleAndLocalPath> privateFileHandles,
-            long incrementalSize) {
+            long incrementalSize,
+            boolean ownsSharedFiles) {
         this.sharedFileHandles = sharedFileHandles;
         this.privateFileHandles = privateFileHandles;
         this.incrementalSize = incrementalSize;
+        this.ownsSharedFiles = ownsSharedFiles;
+    }
+
+    /** Creates a handle that directly discards its shared files until they are registered. */
+    public static KvSnapshotHandle create(
+            List<KvFileHandleAndLocalPath> sharedFileHandles,
+            List<KvFileHandleAndLocalPath> privateFileHandles,
+            long incrementalSize) {
+        return new KvSnapshotHandle(sharedFileHandles, privateFileHandles, incrementalSize, true);
+    }
+
+    /** Restores a handle that only references already-persisted shared files. */
+    static KvSnapshotHandle restore(
+            List<KvFileHandleAndLocalPath> sharedFileHandles,
+            List<KvFileHandleAndLocalPath> privateFileHandles,
+            long incrementalSize) {
+        return new KvSnapshotHandle(sharedFileHandles, privateFileHandles, incrementalSize, false);
     }
 
     public List<KvFileHandleAndLocalPath> getSharedKvFileHandles() {
@@ -98,9 +111,6 @@ public long getSnapshotSize() {
     }
 
     public void discard() {
-        SharedKvFileRegistry registry = this.sharedKvFileRegistry;
-        final boolean isRegistered = (registry != null);
-
         try {
             SnapshotUtil.bestEffortDiscardAllKvFiles(
                     privateFileHandles.stream()
@@ -110,7 +120,7 @@ public void discard() {
             LOG.warn("Could not properly discard misc file states.", e);
         }
 
-        if (!isRegistered) {
+        if (ownsSharedFiles) {
             try {
                 SnapshotUtil.bestEffortDiscardAllKvFiles(
                         sharedFileHandles.stream()
@@ -123,11 +133,8 @@ public void discard() {
     }
 
     public void registerKvFileHandles(SharedKvFileRegistry registry, long snapshotID) {
-        checkState(
-                sharedKvFileRegistry != registry,
-                "The kv file handle has already registered its shared kv files to the given registry.");
-
-        sharedKvFileRegistry = checkNotNull(registry);
+        checkNotNull(registry);
+        ownsSharedFiles = false;
 
         for (KvFileHandleAndLocalPath handleAndLocalPath : sharedFileHandles) {
             registry.registerReference(

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

@@ -221,7 +221,7 @@ public SnapshotResult get(CloseableRegistry snapshotCloseableRegistry) throws Ex
                 completed = true;
                 // We make the 'sstFiles' as the 'shared' in KvSnapshotHandle,
                 final KvSnapshotHandle kvSnapshotHandle =
-                        new KvSnapshotHandle(sstFiles, miscFiles, snapshotIncrementalSize);
+                        KvSnapshotHandle.create(sstFiles, miscFiles, snapshotIncrementalSize);
                 return new SnapshotResult(
                         kvSnapshotHandle, snapshotLocation.getSnapshotDirectory(), tabletState);
             } finally {

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

@@ -854,7 +854,7 @@ private void downloadKvSnapshots(CompletedSnapshot completedSnapshot, Path kvTab
         try {
             kvSnapshotDataDownloader.transferAllDataToDirectory(downloadSpec, closeableRegistry);
         } catch (Exception e) {
-            if (e.getMessage().contains(CompletedSnapshot.SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE)) {
+            if (CompletedSnapshot.isSnapshotDataNotExists(e)) {
                 try {
                     snapshotContext.handleSnapshotBroken(completedSnapshot);
                 } catch (Exception t) {

fluss-server/src/test/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManagerTest.java

@@ -186,6 +186,20 @@ void testRemoveCompletedSnapshotStoreFromManager() throws Exception {
 
     @Test
     void testMetadataInconsistencyWithMetadataNotExistsException() throws Exception {
+        verifyMissingSnapshotMetadataIsCleanedUp(
+                new FileNotFoundException(
+                        CompletedSnapshot.SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE));
+    }
+
+    @Test
+    void testMetadataInconsistencyWithHadoopFileDoesNotExistException() throws Exception {
+        verifyMissingSnapshotMetadataIsCleanedUp(
+                new IOException(
+                        "Failed to open snapshot metadata.",
+                        new IOException("File does not exist: /remote/snapshot/_METADATA")));
+    }
+
+    private void verifyMissingSnapshotMetadataIsCleanedUp(IOException exception) throws Exception {
         // setup test data with mixed valid and invalid snapshots
         TableBucket tableBucket = new TableBucket(1, 1);
         CompletedSnapshot validSnapshot =
@@ -196,7 +210,7 @@ void testMetadataInconsistencyWithMetadataNotExistsException() throws Exception
         CompletedSnapshot invalidSnapshot =
                 KvTestUtils.mockCompletedSnapshot(tempDir, tableBucket, 2L);
         TestingCompletedSnapshotHandle invalidSnapshotHandle =
-                new TestingCompletedSnapshotHandleWithFileNotFound(invalidSnapshot);
+                new TestingBrokenCompletedSnapshotHandle(invalidSnapshot, exception);
 
         // create CompletedSnapshotHandleStore with real implementations
         CompletedSnapshotHandleStore completedSnapshotHandleStore =
@@ -220,6 +234,7 @@ void testMetadataInconsistencyWithMetadataNotExistsException() throws Exception
                         DATA1_TABLE_PATH, tableBucket);
         assertThat(completedSnapshotStore.getAllSnapshots()).hasSize(1);
         assertThat(completedSnapshotStore.getAllSnapshots().get(0).getSnapshotID()).isEqualTo(1L);
+        assertThat(completedSnapshotHandleStore.get(tableBucket, 2L)).isEmpty();
     }
 
     private CompletedSnapshotStoreManager createCompletedSnapshotStoreManager(
@@ -281,17 +296,20 @@ private Set<TableBucket> createTableBuckets(int tableNum, int bucketNum) {
      * A test-specific implementation of CompletedSnapshotHandle that throws FileNotFoundException
      * with the specific error message expected by CompletedSnapshotStoreManager.
      */
-    private static class TestingCompletedSnapshotHandleWithFileNotFound
+    private static class TestingBrokenCompletedSnapshotHandle
             extends TestingCompletedSnapshotHandle {
 
-        public TestingCompletedSnapshotHandleWithFileNotFound(CompletedSnapshot snapshot) {
+        private final IOException exception;
+
+        public TestingBrokenCompletedSnapshotHandle(
+                CompletedSnapshot snapshot, IOException exception) {
             super(snapshot, false);
+            this.exception = exception;
         }
 
         @Override
         public CompletedSnapshot retrieveCompleteSnapshot() throws IOException {
-            throw new FileNotFoundException(
-                    CompletedSnapshot.SNAPSHOT_DATA_NOT_EXISTS_ERROR_MESSAGE);
+            throw exception;
         }
     }
 

fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorServiceOrphanRpcsITCase.java

@@ -221,7 +221,7 @@ private CompletedSnapshot createMinimalSnapshot(TableBucket tb, long snapshotId)
                 tb,
                 snapshotId,
                 new FsPath(snapDir.resolve("_metadata").toUri().toString()),
-                new KvSnapshotHandle(Collections.emptyList(), Collections.emptyList(), 0L),
+                KvSnapshotHandle.create(Collections.emptyList(), Collections.emptyList(), 0L),
                 snapshotId,
                 null,
                 null);

fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotJsonSerdeTest.java

@@ -56,7 +56,7 @@ protected CompletedSnapshot[] createObjects() {
                         new TableBucket(1, 1),
                         1,
                         new FsPath("oss://bucket/snapshot"),
-                        new KvSnapshotHandle(sharedFileHandles, privateFileHandles, 5),
+                        KvSnapshotHandle.create(sharedFileHandles, privateFileHandles, 5),
                         10,
                         null,
                         null);
@@ -65,7 +65,7 @@ protected CompletedSnapshot[] createObjects() {
                         new TableBucket(1, 10L, 1),
                         1,
                         new FsPath("oss://bucket/snapshot"),
-                        new KvSnapshotHandle(sharedFileHandles, privateFileHandles, 5),
+                        KvSnapshotHandle.create(sharedFileHandles, privateFileHandles, 5),
                         10,
                         1234L,
                         Collections.singletonList(new AutoIncIDRange(2, 10000, 20000)));

fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStoreTest.java

@@ -746,7 +746,7 @@ private CompletedSnapshot getSnapshot(long id) {
                 tableBucket,
                 id,
                 new FsPath(tempDir.toString(), "test_snapshot"),
-                new KvSnapshotHandle(Collections.emptyList(), Collections.emptyList(), 0));
+                KvSnapshotHandle.create(Collections.emptyList(), Collections.emptyList(), 0));
     }
 
     private CompletedSnapshot getSnapshotWithSharedFiles(
@@ -756,7 +756,7 @@ private CompletedSnapshot getSnapshotWithSharedFiles(
                 tableBucket,
                 id,
                 new FsPath(tempDir.toString(), "snapshot_" + id),
-                new KvSnapshotHandle(sharedFiles, Collections.emptyList(), 0));
+                KvSnapshotHandle.create(sharedFiles, Collections.emptyList(), 0));
     }
 
     private void testSnapshotRetention(
@@ -804,7 +804,7 @@ private List<Tuple2<CompletedSnapshotHandle, String>> createSnapshotHandles(
                             new TableBucket(1, 1),
                             i,
                             new FsPath("test_snapshot"),
-                            new KvSnapshotHandle(
+                            KvSnapshotHandle.create(
                                     Collections.emptyList(), Collections.emptyList(), -1));
             final CompletedSnapshotHandle snapshotStateHandle =
                     new TestingCompletedSnapshotHandle(

fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotTest.java

@@ -27,6 +27,7 @@
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
@@ -105,8 +106,32 @@ void testKvSnapshotSize(@TempDir Path tempDir) throws Exception {
         assertThat(snapshot.getKvSnapshotHandle().getSnapshotSize()).isEqualTo(400L);
     }
 
+    @Test
+    void testDiscardRestoredSnapshotKeepsSharedFiles(@TempDir Path tempDir) throws Exception {
+        TableBucket tableBucket = new TableBucket(1, 1);
+        Path localFileDir = makeDir(tempDir, "local");
+        Path snapshotBaseLocation = makeDir(tempDir, "snapshot");
+        Path shareDir = makeDir(snapshotBaseLocation, "share");
+        Path snapshotPath = makeDir(snapshotBaseLocation, "snapshot-1");
+        KvSnapshotHandle kvSnapshotHandle =
+                makeSnapshotHandle(localFileDir, snapshotPath, shareDir, 100);
+        CompletedSnapshot snapshot =
+                new CompletedSnapshot(
+                        tableBucket,
+                        1,
+                        FsPath.fromLocalFile(snapshotPath.toFile()),
+                        kvSnapshotHandle);
+        Files.createFile(snapshotPath.resolve("_METADATA"));
+
+        CompletedSnapshot restoredSnapshot =
+                CompletedSnapshotJsonSerde.fromJson(CompletedSnapshotJsonSerde.toJson(snapshot));
+        restoredSnapshot.discardAsync(Executors.directExecutor()).get();
+
+        checkCompletedSnapshotCleanUp(snapshotPath, restoredSnapshot.getKvSnapshotHandle(), false);
+    }
+
     private void checkCompletedSnapshotCleanUp(
-            Path snapshotPath, KvSnapshotHandle kvSnapshotHandle, boolean isShareFileShouldDelete) {
+            Path snapshotPath, KvSnapshotHandle kvSnapshotHandle, boolean shouldDeleteSharedFiles) {
         // private should be deleted, but the local file should still remain
         for (KvFileHandleAndLocalPath kvFileHandleAndLocalPath :
                 kvSnapshotHandle.getPrivateFileHandles()) {
@@ -118,8 +143,8 @@ private void checkCompletedSnapshotCleanUp(
         // check the share files is as expected, and the local file should still remain
         for (KvFileHandleAndLocalPath kvFileHandleAndLocalPath :
                 kvSnapshotHandle.getSharedKvFileHandles()) {
-            // share files should also be deleted, but the local file should still remain
-            if (isShareFileShouldDelete) {
+            // shared files should also be deleted, but the local file should still remain
+            if (shouldDeleteSharedFiles) {
                 assertThat(new File(kvFileHandleAndLocalPath.getKvFileHandle().getFilePath()))
                         .doesNotExist();
             } else {
@@ -168,7 +193,7 @@ KvSnapshotHandle makeSnapshotHandle(
                             new KvFileHandle(privateFile.getPath(), privateFile.length()),
                             localFile.getPath()));
         }
-        return new KvSnapshotHandle(sharedFileHandles, privateFileHandles, 10);
+        return KvSnapshotHandle.create(sharedFileHandles, privateFileHandles, 10);
     }
 
     private Path makeDir(Path basePath, String dirName) {

fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvSnapshotDataDownloaderTest.java

@@ -162,7 +162,8 @@ private KvSnapshotDownloadSpec createDownloadRequestForContent(
                             handles.get(i), String.format("private-%d-%d", remoteHandleId, i)));
         }
 
-        KvSnapshotHandle kvSnapshotHandle = new KvSnapshotHandle(sharedStates, privateStates, -1);
+        KvSnapshotHandle kvSnapshotHandle =
+                KvSnapshotHandle.create(sharedStates, privateStates, -1);
 
         return new KvSnapshotDownloadSpec(kvSnapshotHandle, dstPath);
     }

fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/SnapshotsCleanerTest.java

@@ -118,7 +118,7 @@ private static class TestKvSnapshotHandle extends KvSnapshotHandle {
         public TestKvSnapshotHandle(
                 List<KvFileHandleAndLocalPath> sharedFileHandles,
                 List<KvFileHandleAndLocalPath> privateFileHandles) {
-            super(sharedFileHandles, privateFileHandles, -1);
+            super(sharedFileHandles, privateFileHandles, -1, true);
         }
 
         @Override