[core] Introduce commit strict mode to avoid data loss in rescale job (apache#5727)

Author: tsreaper

docs/layouts/shortcodes/generated/core_configuration.html

@@ -158,6 +158,12 @@
             <td>Integer</td>
             <td>Maximum number of retries when commit failed.</td>
         </tr>
+        <tr>
+            <td><h5>commit.strict-mode.last-safe-snapshot</h5></td>
+            <td style="word-wrap: break-word;">(none)</td>
+            <td>Long</td>
+            <td>If set, committer will check if there are other commit user's COMPACT / OVERWRITE snapshot, starting from the snapshot after this one. If found, commit will be aborted. If the value of this option is -1, committer will not check for its first commit.</td>
+        </tr>
         <tr>
             <td><h5>commit.timeout</h5></td>
             <td style="word-wrap: break-word;">(none)</td>

paimon-api/src/main/java/org/apache/paimon/CoreOptions.java

@@ -1744,6 +1744,15 @@ public class CoreOptions implements Serializable {
                             "If true, it disables explicit type casting. For ex: it disables converting LONG type to INT type. "
                                     + "Users can enable this option to disable explicit type casting");
 
+    public static final ConfigOption<Long> COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT =
+            ConfigOptions.key("commit.strict-mode.last-safe-snapshot")
+                    .longType()
+                    .noDefaultValue()
+                    .withDescription(
+                            "If set, committer will check if there are other commit user's COMPACT / OVERWRITE snapshot, "
+                                    + "starting from the snapshot after this one. If found, commit will be aborted. "
+                                    + "If the value of this option is -1, committer will not check for its first commit.");
+
     private final Options options;
 
     public CoreOptions(Map<String, String> options) {
@@ -2717,6 +2726,10 @@ public boolean aggregationRemoveRecordOnDelete() {
         return options.get(AGGREGATION_REMOVE_RECORD_ON_DELETE);
     }
 
+    public Optional<Long> commitStrictModeLastSafeSnapshot() {
+        return options.getOptional(COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT);
+    }
+
     /** Specifies the merge engine for table with primary key. */
     public enum MergeEngine implements DescribedEnum {
         DEDUPLICATE("deduplicate", "De-duplicate and keep the last row."),

paimon-core/src/main/java/org/apache/paimon/AbstractFileStore.java

@@ -321,7 +321,8 @@ public FileStoreCommitImpl newCommit(String commitUser, FileStoreTable table) {
                 options.scanManifestParallelism(),
                 createCommitCallbacks(commitUser, table),
                 options.commitMaxRetries(),
-                options.commitTimeout());
+                options.commitTimeout(),
+                options.commitStrictModeLastSafeSnapshot().orElse(null));
     }
 
     @Override

paimon-core/src/main/java/org/apache/paimon/operation/FileStoreCommitImpl.java

@@ -141,6 +141,7 @@ public class FileStoreCommitImpl implements FileStoreCommit {
     private final BucketMode bucketMode;
     private final long commitTimeout;
     private final int commitMaxRetries;
+    @Nullable private Long strictModeLastSafeSnapshot;
     private final InternalRowPartitionComputer partitionComputer;
 
     private boolean ignoreEmptyCommit;
@@ -174,7 +175,8 @@ public FileStoreCommitImpl(
             @Nullable Integer manifestReadParallelism,
             List<CommitCallback> commitCallbacks,
             int commitMaxRetries,
-            long commitTimeout) {
+            long commitTimeout,
+            @Nullable Long strictModeLastSafeSnapshot) {
         this.snapshotCommit = snapshotCommit;
         this.fileIO = fileIO;
         this.schemaManager = schemaManager;
@@ -203,6 +205,7 @@ public FileStoreCommitImpl(
         this.commitCallbacks = commitCallbacks;
         this.commitMaxRetries = commitMaxRetries;
         this.commitTimeout = commitTimeout;
+        this.strictModeLastSafeSnapshot = strictModeLastSafeSnapshot;
         this.partitionComputer =
                 new InternalRowPartitionComputer(
                         options.partitionDefaultName(),
@@ -894,6 +897,28 @@ CommitResult tryCommitOnce(
         long newSnapshotId =
                 latestSnapshot == null ? Snapshot.FIRST_SNAPSHOT_ID : latestSnapshot.id() + 1;
 
+        if (strictModeLastSafeSnapshot != null && strictModeLastSafeSnapshot >= 0) {
+            for (long id = strictModeLastSafeSnapshot + 1; id < newSnapshotId; id++) {
+                Snapshot snapshot = snapshotManager.snapshot(id);
+                if ((snapshot.commitKind() == Snapshot.CommitKind.COMPACT
+                                || snapshot.commitKind() == Snapshot.CommitKind.OVERWRITE)
+                        && !snapshot.commitUser().equals(commitUser)) {
+                    throw new RuntimeException(
+                            String.format(
+                                    "When trying to commit snapshot %d, "
+                                            + "commit user %s has found a %s snapshot (id: %d) by another user %s. "
+                                            + "Giving up committing as %s is set.",
+                                    newSnapshotId,
+                                    commitUser,
+                                    snapshot.commitKind().name(),
+                                    id,
+                                    snapshot.commitUser(),
+                                    CoreOptions.COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT.key()));
+                }
+            }
+            strictModeLastSafeSnapshot = newSnapshotId - 1;
+        }
+
         if (LOG.isDebugEnabled()) {
             LOG.debug("Ready to commit table files to snapshot {}", newSnapshotId);
             for (ManifestEntry entry : deltaFiles) {
@@ -1086,6 +1111,9 @@ CommitResult tryCommitOnce(
                     commitUser,
                     identifier,
                     commitKind.name());
+            if (strictModeLastSafeSnapshot != null) {
+                strictModeLastSafeSnapshot = newSnapshot.id();
+            }
             commitCallbacks.forEach(callback -> callback.call(deltaFiles, indexFiles, newSnapshot));
             return new SuccessResult();
         }

paimon-core/src/test/java/org/apache/paimon/table/sink/TableCommitTest.java

@@ -313,4 +313,72 @@ public void testGiveUpCommitWhenTotalBucketsChanged() throws Exception {
                     .hasMessageContaining("changed from 1 to 2 without overwrite");
         }
     }
+
+    @Test
+    public void testStrictMode() throws Exception {
+        String path = tempDir.toString();
+        RowType rowType =
+                RowType.of(
+                        new DataType[] {DataTypes.INT(), DataTypes.BIGINT()},
+                        new String[] {"k", "v"});
+
+        Options options = new Options();
+        options.set(CoreOptions.PATH, path);
+        options.set(CoreOptions.BUCKET, 1);
+        options.set(CoreOptions.NUM_SORTED_RUNS_COMPACTION_TRIGGER, 10);
+        TableSchema tableSchema =
+                SchemaUtils.forceCommit(
+                        new SchemaManager(LocalFileIO.create(), new Path(path)),
+                        new Schema(
+                                rowType.getFields(),
+                                Collections.emptyList(),
+                                Collections.singletonList("k"),
+                                options.toMap(),
+                                ""));
+        FileStoreTable table =
+                FileStoreTableFactory.create(
+                        LocalFileIO.create(),
+                        new Path(path),
+                        tableSchema,
+                        CatalogEnvironment.empty());
+        String user1 = UUID.randomUUID().toString();
+        TableWriteImpl<?> write1 = table.newWrite(user1);
+        TableCommitImpl commit1 = table.newCommit(user1);
+
+        Map<String, String> newOptions = new HashMap<>();
+        newOptions.put(CoreOptions.COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT.key(), "-1");
+        table = table.copy(newOptions);
+        String user2 = UUID.randomUUID().toString();
+        TableWriteImpl<?> write2 = table.newWrite(user2);
+        TableCommitImpl commit2 = table.newCommit(user2);
+
+        // by default, first commit is not checked
+
+        write1.write(GenericRow.of(0, 0L));
+        write1.compact(BinaryRow.EMPTY_ROW, 0, true);
+        commit1.commit(1, write1.prepareCommit(true, 1));
+
+        write2.write(GenericRow.of(1, 1L));
+        commit2.commit(1, write2.prepareCommit(false, 1));
+
+        // APPEND commit is ignored
+
+        write1.write(GenericRow.of(2, 2L));
+        commit1.commit(2, write1.prepareCommit(false, 2));
+
+        write2.write(GenericRow.of(3, 3L));
+        commit2.commit(2, write2.prepareCommit(false, 2));
+
+        // COMPACT commit should be checked
+
+        write1.write(GenericRow.of(4, 4L));
+        write1.compact(BinaryRow.EMPTY_ROW, 0, true);
+        commit1.commit(3, write1.prepareCommit(true, 3));
+
+        write2.write(GenericRow.of(5, 5L));
+        assertThatThrownBy(() -> commit2.commit(3, write2.prepareCommit(false, 3)))
+                .isInstanceOf(RuntimeException.class)
+                .hasMessageContaining(
+                        "Giving up committing as commit.strict-mode.last-safe-snapshot is set.");
+    }
 }

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/RescaleAction.java

@@ -19,6 +19,7 @@
 package org.apache.paimon.flink.action;
 
 import org.apache.paimon.CoreOptions;
+import org.apache.paimon.Snapshot;
 import org.apache.paimon.flink.sink.FlinkSinkBuilder;
 import org.apache.paimon.flink.source.FlinkSourceBuilder;
 import org.apache.paimon.manifest.ManifestEntry;
@@ -41,6 +42,7 @@
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Optional;
 
 /** Action to rescale one partition of a table. */
 public class RescaleAction extends TableActionBase {
@@ -81,6 +83,21 @@ public void build() throws Exception {
         env.configure(flinkConf);
 
         FileStoreTable fileStoreTable = (FileStoreTable) table;
+        Optional<Snapshot> optionalSnapshot = fileStoreTable.latestSnapshot();
+        if (!optionalSnapshot.isPresent()) {
+            throw new IllegalArgumentException(
+                    "Table " + table.fullName() + " has no snapshot. No need to rescale.");
+        }
+        Snapshot snapshot = optionalSnapshot.get();
+
+        // If someone commits while the rescale job is running, this commit will be lost.
+        // So we use strict mode to make sure nothing is lost.
+        Map<String, String> dynamicOptions = new HashMap<>();
+        dynamicOptions.put(
+                CoreOptions.COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT.key(),
+                String.valueOf(snapshot.id()));
+        fileStoreTable = fileStoreTable.copy(dynamicOptions);
+
         RowType partitionType = fileStoreTable.schema().logicalPartitionType();
         Predicate partitionPredicate =
                 PartitionPredicate.createPartitionPredicate(
@@ -94,7 +111,9 @@ public void build() throws Exception {
                         .env(env)
                         .sourceBounded(true)
                         .sourceParallelism(
-                                scanParallelism == null ? currentBucketNum() : scanParallelism)
+                                scanParallelism == null
+                                        ? currentBucketNum(snapshot)
+                                        : scanParallelism)
                         .predicate(partitionPredicate)
                         .build();
 
@@ -118,14 +137,15 @@ public void build() throws Exception {
     @Override
     public void run() throws Exception {
         build();
-        env.execute("Rescale Postpone Bucket : " + table.fullName());
+        env.execute("Rescale : " + table.fullName());
     }
 
-    private int currentBucketNum() {
+    private int currentBucketNum(Snapshot snapshot) {
         FileStoreTable fileStoreTable = (FileStoreTable) table;
         Iterator<ManifestEntry> it =
                 fileStoreTable
                         .newSnapshotReader()
+                        .withSnapshot(snapshot)
                         .withPartitionFilter(partition)
                         .readFileIterator();
         Preconditions.checkArgument(