change to kv snapshot lease

Author: swuferhong

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

@@ -18,10 +18,10 @@
 package org.apache.fluss.client.admin;
 
 import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.client.metadata.AcquireKvSnapshotLeaseResult;
 import org.apache.fluss.client.metadata.KvSnapshotMetadata;
 import org.apache.fluss.client.metadata.KvSnapshots;
 import org.apache.fluss.client.metadata.LakeSnapshot;
-import org.apache.fluss.client.metadata.RegisterKvSnapshotResult;
 import org.apache.fluss.cluster.ServerNode;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.cluster.AlterConfig;
@@ -34,7 +34,7 @@
 import org.apache.fluss.exception.InvalidPartitionException;
 import org.apache.fluss.exception.InvalidReplicationFactorException;
 import org.apache.fluss.exception.InvalidTableException;
-import org.apache.fluss.exception.KvSnapshotConsumerNotExistException;
+import org.apache.fluss.exception.KvSnapshotLeaseNotExistException;
 import org.apache.fluss.exception.KvSnapshotNotExistException;
 import org.apache.fluss.exception.LakeTableSnapshotNotExistException;
 import org.apache.fluss.exception.NonPrimaryKeyTableException;
@@ -405,60 +405,73 @@ CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
     /**
-     * Register the kv snapshot consumer of the given tableBucket set asynchronously. After
-     * registered, the kv snapshot will not be deleted until unregistered or the consumer expired.
+     * Acquires a lease for specific KV snapshots of the given tableBuckets asynchronously.
      *
-     * <p>For the details bucket to consume, you can call {@link #getLatestKvSnapshots(TablePath)}
-     * or {@link #getLatestKvSnapshots(TablePath, String)} first.
+     * <p>Once acquired, the specified KV snapshots will be protected from garbage collection for
+     * the duration of the {@code leaseDuration}. The client must call {@link #renewKvSnapshotLease}
+     * periodically to keep the lease alive, or call {@link #releaseKvSnapshotLease} to release the
+     * lock early when reading is finished.
      *
-     * <p>The following exceptions can be anticipated when calling {@code get()} on returned future.
+     * <p>If the lease expires (no renew received within duration), the server is free to delete the
+     * snapshot files.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on returned future:
      *
      * <ul>
      *   <li>{@link TableNotExistException} if the table does not exist.
-     *   <li>{@link PartitionNotExistException} if the partition does not exist
+     *   <li>{@link PartitionNotExistException} if the partition does not exist.
      * </ul>
      *
-     * @param consumerId the consumer id.
-     * @param consumeBuckets the tableBuckets to consume, a map from TableBucket to kvSnapshotId.
-     * @return the result of registering kv snapshot consumer.
+     * @param leaseId The unique ID for this lease session (usually a UUID generated by client).
+     * @param snapshotIds The snapshots to consume, a map from TableBucket to kvSnapshotId.
+     * @param leaseDuration The duration (in milliseconds) for which the snapshots should be kept.
+     * @return The result of the acquire operation, containing any buckets that failed to be locked.
      */
-    CompletableFuture<RegisterKvSnapshotResult> registerKvSnapshotConsumer(
-            String consumerId, Map<TableBucket, Long> consumeBuckets);
+    CompletableFuture<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration);
 
     /**
-     * Unregister the kv snapshot consumer of the given tableBucket set asynchronously.
+     * Renews the lease for the given leaseId asynchronously.
+     *
+     * <p>This acts as a heartbeat. It resets the expiration time for ALL tableBuckets currently
+     * associated with this {@code leaseId}.
      *
-     * <p>As the following situation happens, unregister operation will be partially ignored or
-     * completely ignored
+     * <p>The following exceptions can be anticipated:
      *
      * <ul>
-     *   <li>If the table/partition not exists, the unregister operation will be ignored for this
-     *       table/partition, others will success.
-     *   <li>If the partition bucket/kvSnapshotId not exists, the unregister operation will be
-     *       ignored for this bucket/kvSnapshotId, others will success.
-     *   <li>If the consumerId not exists, the unregister operation will be ignored.
+     *   <li>{@link KvSnapshotLeaseNotExistException} if the leaseId does not exist or has already
+     *       expired.
      * </ul>
      *
-     * @param consumerId the consumer id.
-     * @param bucketsToUnregister the tableBucket to unregister.
+     * @param leaseId The lease id to renew.
+     * @param leaseDuration The new duration (in milliseconds) from now.
      */
-    CompletableFuture<Void> unregisterKvSnapshotConsumer(
-            String consumerId, Set<TableBucket> bucketsToUnregister);
+    CompletableFuture<Void> renewKvSnapshotLease(String leaseId, long leaseDuration);
 
     /**
-     * Clear the kv snapshot consumer asynchronously. After cleared, all the registered kv snapshot
-     * consume info for this consumer will be cleared.
+     * Releases the lease for specific tableBuckets asynchronously.
      *
-     * <p>The following exceptions can be anticipated when calling {@code get()} on returned future.
+     * <p>This is typically called when a client finishes reading a specific bucket (or a batch of
+     * buckets) but is still reading others under the same leaseId.
      *
-     * <ul>
-     *   <li>{@link TableNotExistException} if the table does not exist.
-     *   <li>{@link KvSnapshotConsumerNotExistException} if the kv snapshot consumer does not exist.
-     * </ul>
+     * <p>If {@code bucketsToRelease} contains all buckets under this leaseId, the lease itself will
+     * be removed.
+     *
+     * @param leaseId The lease id.
+     * @param bucketsToRelease The specific tableBuckets to release.
+     */
+    CompletableFuture<Void> releaseKvSnapshotLease(
+            String leaseId, Set<TableBucket> bucketsToRelease);
+
+    /**
+     * Drops the entire lease asynchronously.
+     *
+     * <p>All snapshots locked under this {@code leaseId} will be released immediately. This is
+     * equivalent to calling {@link #releaseKvSnapshotLease} with all held buckets.
      *
-     * @param consumerId the consumer id.
+     * @param leaseId The lease id to drop.
      */
-    CompletableFuture<Void> clearKvSnapshotConsumer(String consumerId);
+    CompletableFuture<Void> dropKvSnapshotLease(String leaseId);
 
     /**
      * Get table lake snapshot info of the given table asynchronously.

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

@@ -17,15 +17,14 @@
 
 package org.apache.fluss.client.admin;
 
+import org.apache.fluss.client.metadata.AcquireKvSnapshotLeaseResult;
 import org.apache.fluss.client.metadata.KvSnapshotMetadata;
 import org.apache.fluss.client.metadata.KvSnapshots;
 import org.apache.fluss.client.metadata.LakeSnapshot;
 import org.apache.fluss.client.metadata.MetadataUpdater;
-import org.apache.fluss.client.metadata.RegisterKvSnapshotResult;
 import org.apache.fluss.client.utils.ClientRpcMessageUtils;
 import org.apache.fluss.cluster.Cluster;
 import org.apache.fluss.cluster.ServerNode;
-import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.config.cluster.AlterConfig;
 import org.apache.fluss.config.cluster.ConfigEntry;
@@ -49,7 +48,6 @@
 import org.apache.fluss.rpc.gateway.TabletServerGateway;
 import org.apache.fluss.rpc.messages.AlterClusterConfigsRequest;
 import org.apache.fluss.rpc.messages.AlterTableRequest;
-import org.apache.fluss.rpc.messages.ClearKvSnapshotConsumerRequest;
 import org.apache.fluss.rpc.messages.CreateAclsRequest;
 import org.apache.fluss.rpc.messages.CreateDatabaseRequest;
 import org.apache.fluss.rpc.messages.CreateTableRequest;
@@ -58,6 +56,7 @@
 import org.apache.fluss.rpc.messages.DescribeClusterConfigsRequest;
 import org.apache.fluss.rpc.messages.DropAclsRequest;
 import org.apache.fluss.rpc.messages.DropDatabaseRequest;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseRequest;
 import org.apache.fluss.rpc.messages.DropTableRequest;
 import org.apache.fluss.rpc.messages.GetDatabaseInfoRequest;
 import org.apache.fluss.rpc.messages.GetKvSnapshotMetadataRequest;
@@ -76,6 +75,7 @@
 import org.apache.fluss.rpc.messages.PbListOffsetsRespForBucket;
 import org.apache.fluss.rpc.messages.PbPartitionSpec;
 import org.apache.fluss.rpc.messages.PbTablePath;
+import org.apache.fluss.rpc.messages.RenewKvSnapshotLeaseRequest;
 import org.apache.fluss.rpc.messages.TableExistsRequest;
 import org.apache.fluss.rpc.messages.TableExistsResponse;
 import org.apache.fluss.rpc.protocol.ApiError;
@@ -94,13 +94,13 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAcquireKvSnapshotLeaseRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAlterTableRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeCreatePartitionRequest;
 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.makeRegisterKvSnapshotConsumerRequest;
-import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeUnregisterKvSnapshotConsumerRequest;
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeReleaseKvSnapshotLeaseRequest;
 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;
@@ -388,35 +388,39 @@ public CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
     }
 
     @Override
-    public CompletableFuture<RegisterKvSnapshotResult> registerKvSnapshotConsumer(
-            String consumerId, Map<TableBucket, Long> consumeBuckets) {
-        if (consumeBuckets.isEmpty()) {
-            throw new IllegalArgumentException("consumeBuckets is empty");
+    public CompletableFuture<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration) {
+        if (snapshotIds.isEmpty()) {
+            throw new IllegalArgumentException(
+                    "The snapshotIds to acquire kv snapshot lease is empty");
         }
 
-        long expirationTime =
-                clientConfig
-                        .get(ConfigOptions.CLIENT_SCANNER_KV_SNAPSHOT_CONSUMER_EXPIRATION_TIME)
-                        .toMillis();
-        return gateway.registerKvSnapshotConsumer(
-                        makeRegisterKvSnapshotConsumerRequest(
-                                consumerId, consumeBuckets, expirationTime))
-                .thenApply(ClientRpcMessageUtils::toRegisterKvSnapshotResult);
+        return gateway.acquireKvSnapshotLease(
+                        makeAcquireKvSnapshotLeaseRequest(leaseId, snapshotIds, leaseDuration))
+                .thenApply(ClientRpcMessageUtils::toAcquireKvSnapshotLeaseResult);
     }
 
     @Override
-    public CompletableFuture<Void> unregisterKvSnapshotConsumer(
-            String consumerId, Set<TableBucket> bucketsToUnregister) {
-        return gateway.unregisterKvSnapshotConsumer(
-                        makeUnregisterKvSnapshotConsumerRequest(consumerId, bucketsToUnregister))
+    public CompletableFuture<Void> renewKvSnapshotLease(String leaseId, long leaseDuration) {
+        RenewKvSnapshotLeaseRequest request =
+                new RenewKvSnapshotLeaseRequest()
+                        .setLeaseId(leaseId)
+                        .setLeaseDuration(leaseDuration);
+        return gateway.renewKvSnapshotLease(request).thenApply(r -> null);
+    }
+
+    @Override
+    public CompletableFuture<Void> releaseKvSnapshotLease(
+            String leaseId, Set<TableBucket> bucketsToRelease) {
+        return gateway.releaseKvSnapshotLease(
+                        makeReleaseKvSnapshotLeaseRequest(leaseId, bucketsToRelease))
                 .thenApply(r -> null);
     }
 
     @Override
-    public CompletableFuture<Void> clearKvSnapshotConsumer(String consumerId) {
-        ClearKvSnapshotConsumerRequest request =
-                new ClearKvSnapshotConsumerRequest().setConsumerId(consumerId);
-        return gateway.clearKvSnapshotConsumer(request).thenApply(r -> null);
+    public CompletableFuture<Void> dropKvSnapshotLease(String leaseId) {
+        DropKvSnapshotLeaseRequest request = new DropKvSnapshotLeaseRequest().setLeaseId(leaseId);
+        return gateway.dropKvSnapshotLease(request).thenApply(r -> null);
     }
 
     @Override

…t/metadata/RegisterKvSnapshotResult.java …tadata/AcquireKvSnapshotLeaseResult.javafluss-client/src/main/java/org/apache/fluss/client/metadata/RegisterKvSnapshotResult.java renamed to fluss-client/src/main/java/org/apache/fluss/client/metadata/AcquireKvSnapshotLeaseResult.java fluss-client/src/main/java/org/apache/fluss/client/metadata/RegisterKvSnapshotResult.java renamed to fluss-client/src/main/java/org/apache/fluss/client/metadata/AcquireKvSnapshotLeaseResult.java

@@ -20,27 +20,36 @@
 import org.apache.fluss.annotation.PublicEvolving;
 import org.apache.fluss.metadata.TableBucket;
 
+import java.util.Map;
 import java.util.Set;
 
 /**
- * A class to represent the result of registering kv snapshot. It contains:
+ * A class to represent the result of acquire kv snapshot lease. It contains:
  *
  * <ul>
- *   <li>An set of failed tableBuckets. Such as the specify snapshotId is not exist for this table
+ *   <li>A map of unavailable snapshots. Such as the specify snapshotId is not exist for this table
  *       bucket.
  * </ul>
  *
  * @since 0.9
  */
 @PublicEvolving
-public class RegisterKvSnapshotResult {
-    private final Set<TableBucket> failedTableBucketSet;
+public class AcquireKvSnapshotLeaseResult {
+    private final Map<TableBucket, Long> unavailableSnapshots;
 
-    public RegisterKvSnapshotResult(Set<TableBucket> failedTableBucketSet) {
-        this.failedTableBucketSet = failedTableBucketSet;
+    public AcquireKvSnapshotLeaseResult(Map<TableBucket, Long> unavailableSnapshots) {
+        this.unavailableSnapshots = unavailableSnapshots;
     }
 
-    public Set<TableBucket> getFailedTableBucketSet() {
-        return failedTableBucketSet;
+    /**
+     * Returns the set of buckets that could not be locked (e.g., snapshot ID doesn't exist or has
+     * already been GC'ed).
+     */
+    public Map<TableBucket, Long> getUnavailableSnapshots() {
+        return unavailableSnapshots;
+    }
+
+    public Set<TableBucket> getUnavailableTableBucketSet() {
+        return unavailableSnapshots.keySet();
     }
 }