[core] Make LocalKvDb flush failure-safe

paimon-common/src/main/java/org/apache/paimon/lookup/sort/db/LocalKvDb.java

@@ -471,12 +471,10 @@ public void flush() throws IOException {
 
     private void flushMemTable() throws IOException {
         TreeMap<MemorySlice, byte[]> snapshot = memTable;
-        memTable = new TreeMap<>(keyComparator);
-        memTableSize = 0;
-
         SstFileMetadata metadata = writeMemTableToSst(snapshot);
-
         levels.addLevelZeroFile(metadata);
+        memTable = new TreeMap<>(keyComparator);
+        memTableSize = 0;
 
         LOG.info(
                 "Flushed MemTable to L0 SST file: {}, entries: {}",
@@ -610,12 +608,14 @@ private byte[] lookupInFile(File file, byte[] key) throws IOException {
     private SstFileMetadata writeMemTableToSst(TreeMap<MemorySlice, byte[]> data)
             throws IOException {
         File sstFile = newSstFile();
-        SortLookupStoreWriter writer =
-                storeFactory.createWriter(sstFile, bloomFilterBuilderFactory.apply(data.size()));
+        SortLookupStoreWriter writer = null;
         MemorySlice minKey = null;
         MemorySlice maxKey = null;
         long tombstoneCount = 0;
         try {
+            writer =
+                    storeFactory.createWriter(
+                            sstFile, bloomFilterBuilderFactory.apply(data.size()));
             for (Map.Entry<MemorySlice, byte[]> entry : data.entrySet()) {
                 writer.put(entry.getKey().copyBytes(), entry.getValue());
                 if (minKey == null) {
@@ -626,10 +626,20 @@ private SstFileMetadata writeMemTableToSst(TreeMap<MemorySlice, byte[]> data)
                     tombstoneCount++;
                 }
             }
-        } finally {
             writer.close();
+            writer = null;
+            return new SstFileMetadata(sstFile, minKey, maxKey, tombstoneCount, 0);
+        } catch (IOException | RuntimeException e) {
+            if (writer != null) {
+                try {
+                    writer.close();
+                } catch (IOException suppressed) {
+                    e.addSuppressed(suppressed);
+                }
+            }
+            deleteFileQuietly(sstFile);
+            throw e;
         }
-        return new SstFileMetadata(sstFile, minKey, maxKey, tombstoneCount, 0);
     }
 
     private File newSstFile() {

paimon-common/src/test/java/org/apache/paimon/lookup/sort/db/LocalKvDbTest.java

@@ -30,6 +30,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.RandomAccessFile;
+import java.nio.file.Files;
 import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -391,6 +392,31 @@ public void testCloseFlushesMemTable() throws IOException {
         Assertions.assertEquals(1, db.getSstFileCount());
     }
 
+    @Test
+    public void testFlushFailureKeepsMemTableForRetry() throws IOException {
+        File dbDir = new File(tempDir.toFile(), "flush-failure-db");
+        try (LocalKvDb db =
+                LocalKvDb.builder(dbDir)
+                        .memTableFlushThreshold(1024 * 1024)
+                        .blockSize(256)
+                        .compressOptions(new CompressOptions("none", 1))
+                        .build()) {
+            putString(db, "key", "value");
+
+            Files.delete(dbDir.toPath());
+            Files.createFile(dbDir.toPath());
+            Assertions.assertThrows(IOException.class, db::flush);
+            Assertions.assertEquals("value", getString(db, "key"));
+            Assertions.assertTrue(db.getMemTableSize() > 0);
+
+            Files.delete(dbDir.toPath());
+            Files.createDirectories(dbDir.toPath());
+            db.flush();
+            Assertions.assertEquals("value", getString(db, "key"));
+            Assertions.assertEquals(0, db.getMemTableSize());
+        }
+    }
+
     @Test
     public void testClosedDbThrowsException() throws IOException {
         LocalKvDb db = createDb();