[kv] Add Producer Offset Snapshot for Exactly-Once semantics (#2434)

Author: platinumhamburg

fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java

@@ -71,6 +71,7 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -604,4 +605,65 @@ CompletableFuture<Optional<RebalanceProgress>> listRebalanceProgress(
      *     NoRebalanceInProgressException} will be thrown.
      */
     CompletableFuture<Void> cancelRebalance(@Nullable String rebalanceId);
+
+    // ==================================================================================
+    // Producer Offset Management APIs (for Exactly-Once Semantics)
+    // ==================================================================================
+
+    /**
+     * Register producer offset snapshot.
+     *
+     * <p>This method provides atomic "check and register" semantics:
+     *
+     * <ul>
+     *   <li>If snapshot does not exist: create new snapshot and return {@link
+     *       RegisterResult#CREATED}
+     *   <li>If snapshot already exists: do NOT overwrite and return {@link
+     *       RegisterResult#ALREADY_EXISTS}
+     * </ul>
+     *
+     * <p>The atomicity is guaranteed by the server implementation. This enables the caller to
+     * determine whether undo recovery is needed based on the return value.
+     *
+     * <p>The snapshot will be automatically cleaned up after the configured TTL expires.
+     *
+     * <p>This API is typically used by Flink Operator Coordinator at job startup to register the
+     * initial offset snapshot before any data is written.
+     *
+     * @param producerId the ID of the producer (typically Flink job ID)
+     * @param offsets map of TableBucket to offset for all tables
+     * @return a CompletableFuture containing the registration result indicating whether the
+     *     snapshot was newly created or already existed
+     * @since 0.9
+     */
+    CompletableFuture<RegisterResult> registerProducerOffsets(
+            String producerId, Map<TableBucket, Long> offsets);
+
+    /**
+     * Get producer offset snapshot.
+     *
+     * <p>This method retrieves the registered offset snapshot for a producer. Returns null if no
+     * snapshot exists for the given producer ID.
+     *
+     * <p>This API is typically used by Flink Operator Coordinator at job startup to check if a
+     * previous snapshot exists (indicating a failover before first checkpoint).
+     *
+     * @param producerId the ID of the producer
+     * @return a CompletableFuture containing the producer offsets, or null if not found
+     * @since 0.9
+     */
+    CompletableFuture<ProducerOffsetsResult> getProducerOffsets(String producerId);
+
+    /**
+     * Delete producer offset snapshot.
+     *
+     * <p>This method deletes the registered offset snapshot for a producer. This is typically
+     * called after the first checkpoint completes successfully, as the checkpoint state will be
+     * used for recovery instead of the initial snapshot.
+     *
+     * @param producerId the ID of the producer
+     * @return a CompletableFuture that completes when deletion succeeds
+     * @since 0.9
+     */
+    CompletableFuture<Void> deleteProducerOffsets(String producerId);
 }

fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java

@@ -57,6 +57,7 @@
 import org.apache.fluss.rpc.messages.CreateTableRequest;
 import org.apache.fluss.rpc.messages.DatabaseExistsRequest;
 import org.apache.fluss.rpc.messages.DatabaseExistsResponse;
+import org.apache.fluss.rpc.messages.DeleteProducerOffsetsRequest;
 import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest;
 import org.apache.fluss.rpc.messages.DropAclsRequest;
 import org.apache.fluss.rpc.messages.DropDatabaseRequest;
@@ -65,6 +66,7 @@
 import org.apache.fluss.rpc.messages.GetKvSnapshotMetadataRequest;
 import org.apache.fluss.rpc.messages.GetLatestKvSnapshotsRequest;
 import org.apache.fluss.rpc.messages.GetLatestLakeSnapshotRequest;
+import org.apache.fluss.rpc.messages.GetProducerOffsetsRequest;
 import org.apache.fluss.rpc.messages.GetTableInfoRequest;
 import org.apache.fluss.rpc.messages.GetTableSchemaRequest;
 import org.apache.fluss.rpc.messages.ListAclsRequest;
@@ -105,6 +107,7 @@
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeDropPartitionRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeListOffsetsRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makePbPartitionSpec;
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeRegisterProducerOffsetsRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.toConfigEntries;
 import static org.apache.fluss.client.utils.MetadataUtils.sendMetadataRequestAndRebuildCluster;
 import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toAclBindings;
@@ -588,6 +591,49 @@ public CompletableFuture<Void> cancelRebalance(@Nullable String rebalanceId) {
         return gateway.cancelRebalance(request).thenApply(r -> null);
     }
 
+    // ==================================================================================
+    // Producer Offset Management APIs (for Exactly-Once Semantics)
+    // ==================================================================================
+
+    @Override
+    public CompletableFuture<RegisterResult> registerProducerOffsets(
+            String producerId, Map<TableBucket, Long> offsets) {
+        checkNotNull(producerId, "producerId must not be null");
+        checkNotNull(offsets, "offsets must not be null");
+
+        return gateway.registerProducerOffsets(
+                        makeRegisterProducerOffsetsRequest(producerId, offsets))
+                .thenApply(
+                        response -> {
+                            int code =
+                                    response.hasResult()
+                                            ? response.getResult()
+                                            : RegisterResult.CREATED.getCode();
+                            return RegisterResult.fromCode(code);
+                        });
+    }
+
+    @Override
+    public CompletableFuture<ProducerOffsetsResult> getProducerOffsets(String producerId) {
+        checkNotNull(producerId, "producerId must not be null");
+
+        GetProducerOffsetsRequest request = new GetProducerOffsetsRequest();
+        request.setProducerId(producerId);
+
+        return gateway.getProducerOffsets(request)
+                .thenApply(ClientRpcMessageUtils::toProducerOffsetsResult);
+    }
+
+    @Override
+    public CompletableFuture<Void> deleteProducerOffsets(String producerId) {
+        checkNotNull(producerId, "producerId must not be null");
+
+        DeleteProducerOffsetsRequest request = new DeleteProducerOffsetsRequest();
+        request.setProducerId(producerId);
+
+        return gateway.deleteProducerOffsets(request).thenApply(r -> null);
+    }
+
     @Override
     public void close() {
         // nothing to do yet

fluss-client/src/main/java/org/apache/fluss/client/admin/ProducerOffsetsResult.java

@@ -0,0 +1,99 @@
+/*
+ * 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.client.admin;
+
+import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.metadata.TableBucket;
+
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * Result containing producer offset snapshot data.
+ *
+ * <p>This class holds the offset snapshot for a producer, which is used for undo recovery when a
+ * Flink job fails over before completing its first checkpoint.
+ *
+ * <p>The snapshot contains bucket offsets organized by table ID, allowing the Flink Operator
+ * Coordinator to coordinate undo recovery across all subtasks.
+ *
+ * @since 0.9
+ */
+@PublicEvolving
+public class ProducerOffsetsResult {
+
+    private final String producerId;
+    private final Map<Long, Map<TableBucket, Long>> tableOffsets;
+    private final long expirationTime;
+
+    /**
+     * Creates a new ProducerOffsetsResult.
+     *
+     * @param producerId the producer ID (typically Flink job ID)
+     * @param tableOffsets map of table ID to bucket offsets
+     * @param expirationTime the expiration timestamp in milliseconds
+     */
+    public ProducerOffsetsResult(
+            String producerId,
+            Map<Long, Map<TableBucket, Long>> tableOffsets,
+            long expirationTime) {
+        this.producerId = producerId;
+        this.tableOffsets = Collections.unmodifiableMap(tableOffsets);
+        this.expirationTime = expirationTime;
+    }
+
+    /**
+     * Get the producer ID.
+     *
+     * @return the producer ID
+     */
+    public String getProducerId() {
+        return producerId;
+    }
+
+    /**
+     * Get the offset snapshot for all tables.
+     *
+     * @return unmodifiable map of table ID to bucket offsets
+     */
+    public Map<Long, Map<TableBucket, Long>> getTableOffsets() {
+        return tableOffsets;
+    }
+
+    /**
+     * Get the expiration timestamp.
+     *
+     * @return the expiration timestamp in milliseconds since epoch
+     */
+    public long getExpirationTime() {
+        return expirationTime;
+    }
+
+    @Override
+    public String toString() {
+        return "ProducerOffsetsResult{"
+                + "producerId='"
+                + producerId
+                + '\''
+                + ", tableCount="
+                + tableOffsets.size()
+                + ", expirationTime="
+                + expirationTime
+                + '}';
+    }
+}

fluss-client/src/main/java/org/apache/fluss/client/admin/RegisterResult.java

@@ -0,0 +1,69 @@
+/*
+ * 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.client.admin;
+
+import org.apache.fluss.annotation.PublicEvolving;
+
+/**
+ * Result of producer offset registration.
+ *
+ * <p>This enum indicates whether a producer offset snapshot was newly created or already existed
+ * when calling registerProducerOffsets API.
+ *
+ * @since 0.9
+ */
+@PublicEvolving
+public enum RegisterResult {
+    /**
+     * Snapshot was newly created.
+     *
+     * <p>This indicates a first startup scenario where no previous snapshot existed. The caller
+     * does not need to perform undo recovery.
+     */
+    CREATED(0),
+
+    /**
+     * Snapshot already existed and was not overwritten.
+     *
+     * <p>This indicates a failover scenario where a previous snapshot exists. The caller should
+     * perform undo recovery using the existing snapshot offsets.
+     */
+    ALREADY_EXISTS(1);
+
+    /** The code used in RPC messages. */
+    private final int code;
+
+    RegisterResult(int code) {
+        this.code = code;
+    }
+
+    /** Returns the code used in RPC messages. */
+    public int getCode() {
+        return code;
+    }
+
+    /** Returns the RegisterResult for the given code. */
+    public static RegisterResult fromCode(int code) {
+        for (RegisterResult result : values()) {
+            if (result.code == code) {
+                return result;
+            }
+        }
+        throw new IllegalArgumentException("Unknown RegisterResult code: " + code);
+    }
+}