You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@fluss.apache.org by "swuferhong (via GitHub)" <gi...@apache.org> on 2025/12/15 13:32:41 UTC

[PR] [kv] Introduce the API and implement the KV snapshot consumer [fluss]

swuferhong opened a new pull request, #2179:
URL: https://github.com/apache/fluss/pull/2179

   <!--
   *Thank you very much for contributing to Fluss - we are happy that you want to help us improve Fluss. To help the community review your contribution in the best possible way, please go through the checklist below, which will get the contribution into a shape in which it can be best reviewed.*
   
   ## Contribution Checklist
   
     - Make sure that the pull request corresponds to a [GitHub issue](https://github.com/apache/fluss/issues). Exceptions are made for typos in JavaDoc or documentation files, which need no issue.
   
     - Name the pull request in the format "[component] Title of the pull request", where *[component]* should be replaced by the name of the component being changed. Typically, this corresponds to the component label assigned to the issue (e.g., [kv], [log], [client], [flink]). Skip *[component]* if you are unsure about which is the best component.
   
     - Fill out the template below to describe the changes contributed by the pull request. That will give reviewers the context they need to do the review.
   
     - Make sure that the change passes the automated tests, i.e., `mvn clean verify` passes.
   
     - Each pull request should address only one issue, not mix up code from multiple issues.
   
   
   **(The sections below can be removed for hotfixes or typos)**
   -->
   
   ### Purpose
   
   <!-- Linking this pull request to the issue -->
   Linked issue: close #2172 
   
   <!-- What is the purpose of the change -->
   
   Introduce the API and implement the KV snapshot consumer
   
   ### Brief change log
   
   <!-- Please describe the changes made in this pull request and explain how they address the issue -->
   
   ### Tests
   
   <!-- List UT and IT cases to verify this change -->
   
   ### API and Format
   
   <!-- Does this change affect API or storage format -->
   
   ### Documentation
   
   <!-- Does this change introduce a new feature -->
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "wuchong (via GitHub)" <gi...@apache.org>.
wuchong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2779335910


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),

Review Comment:
   Could you create an issue to improve this in the future?
   
   trace by: https://github.com/apache/fluss/issues/2602



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "platinumhamburg (via GitHub)" <gi...@apache.org>.
platinumhamburg commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2654917690


##########
fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java:
##########
@@ -408,6 +411,56 @@ CompletableFuture<Void> dropPartition(
     CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
+    /**
+     * Acquires a lease for specific KV snapshots of the given tableBuckets asynchronously.
+     *
+     * <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
+     * #releaseKvSnapshotLease} to release the lock early when reading is finished.
+     *
+     * <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.
+     * </ul>
+     *
+     * @param leaseId The unique ID for this lease session (usually a UUID generated by client).
+     * @param snapshotIds The snapshots to lease, 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<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration);
+
+    /**
+     * Releases the lease for specific tableBuckets asynchronously.
+     *
+     * <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.
+     *
+     * <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 leaseId The lease id to drop.
+     */
+    CompletableFuture<Void> dropKvSnapshotLease(String leaseId);

Review Comment:
   ditto



##########
fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java:
##########
@@ -408,6 +411,56 @@ CompletableFuture<Void> dropPartition(
     CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
+    /**
+     * Acquires a lease for specific KV snapshots of the given tableBuckets asynchronously.
+     *
+     * <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
+     * #releaseKvSnapshotLease} to release the lock early when reading is finished.
+     *
+     * <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.
+     * </ul>
+     *
+     * @param leaseId The unique ID for this lease session (usually a UUID generated by client).
+     * @param snapshotIds The snapshots to lease, 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<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration);
+
+    /**
+     * Releases the lease for specific tableBuckets asynchronously.
+     *
+     * <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.
+     *
+     * <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.

Review Comment:
   Prefer 'releases' over 'drops'—more precise and idiomatic for leases.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")

Review Comment:
   I'm not entirely sure why the refCountLock is needed—it doesn't seem to be protecting multiple atomic states. Could using just a ConcurrentHashMap along with Atomic types lead to consistency issues?
   
   



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java:
##########
@@ -0,0 +1,36 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+
+/** Procedure to drop kv snapshot lease. */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {

Review Comment:
   Ditto—“drop” isn’t the most appropriate verb to use with “lease.”
   
   



##########
fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java:
##########
@@ -53,6 +53,10 @@ public class MetricNames {
     public static final String KV_ALL_SNAPSHOT_SIZE = "allKvSnapshotSize";
     public static final String SERVER_PHYSICAL_STORAGE_REMOTE_KV_SIZE = "remoteKvSize";
 
+    // for kv snapshot lease.
+    public static final String KV_SNAPSHOT_LEASE_COUNT = "kvSnapshotLeaseCount";

Review Comment:
   I think it should be implemented at the table level so we can quickly identify the affected business when lease leaks occur.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();

Review Comment:
   Instead of getLeasesList() returning a List<String>, it could directly return a List<KvSnapshotLease>.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerde.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.core.JsonGenerator;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
+import org.apache.fluss.utils.json.JsonDeserializer;
+import org.apache.fluss.utils.json.JsonSerializer;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Json serializer and deserializer for {@link KvSnapshotLeaseMetadata}. */
+public class KvSnapshotLeaseMetadataJsonSerde
+        implements JsonSerializer<KvSnapshotLeaseMetadata>,
+                JsonDeserializer<KvSnapshotLeaseMetadata> {
+
+    public static final KvSnapshotLeaseMetadataJsonSerde INSTANCE =
+            new KvSnapshotLeaseMetadataJsonSerde();
+
+    private static final String VERSION_KEY = "version";
+    private static final String EXPIRATION_TIME = "expiration_time";
+    private static final String TABLES = "tables";
+    private static final String TABLE_ID = "table_id";
+    private static final String KV_SNAPSHOT_PATH = "lease_metadata_path";
+
+    private static final int VERSION = 1;
+
+    @Override
+    public void serialize(KvSnapshotLeaseMetadata kvSnapshotLeaseMetadata, JsonGenerator generator)

Review Comment:
   I noticed we're using Jackson's very low-level APIs, which makes the serialization/deserialization logic quite complex. What's the reason behind this approach?
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;

Review Comment:
   The usage of KvSnapshotLeaseMetadataHelper involves both metadata queries and updates. The suffix "Helper" typically implies a stateless utility class, which doesn't align well with the actual role of KvSnapshotLeaseMetadataHelper. It might be more appropriate to rename it to KvSnapshotLeaseMetadataManager.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);

Review Comment:
   This code can be optimized to call get only once.



##########
fluss-common/src/main/java/org/apache/fluss/metadata/KvSnapshotLeaseForBucket.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.metadata;
+
+import java.util.Objects;
+
+/** An entity for kv snapshot lease for bucket. */
+public class KvSnapshotLeaseForBucket {

Review Comment:
   A lease inherently implies a time-bound reservation—without a TTL, it’s just a static reference. Consider adding a ttl (e.g., in milliseconds) to this class to properly reflect lease semantics.
   
   



##########
fluss-server/src/test/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManagerTest.java:
##########
@@ -0,0 +1,481 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.PhysicalTablePath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.server.metrics.group.TestingMetricGroups;
+import org.apache.fluss.server.zk.NOPErrorHandler;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.ZooKeeperExtension;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService;
+import org.apache.fluss.utils.clock.ManualClock;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.fluss.record.TestData.DATA1_PARTITIONED_TABLE_DESCRIPTOR;
+import static org.apache.fluss.record.TestData.DATA1_TABLE_ID_PK;
+import static org.apache.fluss.record.TestData.DATA1_TABLE_INFO_PK;
+import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH_PK;
+import static org.apache.fluss.server.coordinator.CoordinatorTestUtils.createServers;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Test for {@link KvSnapshotLeaseManager}. */
+public class KvSnapshotLeaseManagerTest {
+
+    @RegisterExtension
+    public static final AllCallbackWrapper<ZooKeeperExtension> ZOO_KEEPER_EXTENSION_WRAPPER =
+            new AllCallbackWrapper<>(new ZooKeeperExtension());
+
+    private static final TablePath PARTITION_TABLE_PATH =
+            new TablePath("test_db_1", "test_partition_table");
+    private static final long PARTITION_TABLE_ID = 150008L;
+    private static final TableInfo PARTITION_TABLE_INFO =

Review Comment:
   Can we reuse TestData, or should we add the test data into TestData?
   
   



##########
fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java:
##########
@@ -53,6 +53,10 @@ public class MetricNames {
     public static final String KV_ALL_SNAPSHOT_SIZE = "allKvSnapshotSize";
     public static final String SERVER_PHYSICAL_STORAGE_REMOTE_KV_SIZE = "remoteKvSize";
 
+    // for kv snapshot lease.
+    public static final String KV_SNAPSHOT_LEASE_COUNT = "kvSnapshotLeaseCount";
+    public static final String LEASED_KV_SNAPSHOT_COUNT = "leasedKvSnapshotCount";

Review Comment:
   Please add documentation to monitor-metrics.md.
   
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLeaseJsonSerde.java:
##########
@@ -0,0 +1,116 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.core.JsonGenerator;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
+import org.apache.fluss.utils.json.JsonDeserializer;
+import org.apache.fluss.utils.json.JsonSerdeUtils;
+import org.apache.fluss.utils.json.JsonSerializer;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Json serializer and deserializer for {@link KvSnapshotTableLease}. */
+public class KvSnapshotTableLeaseJsonSerde
+        implements JsonSerializer<KvSnapshotTableLease>, JsonDeserializer<KvSnapshotTableLease> {
+
+    public static final KvSnapshotTableLeaseJsonSerde INSTANCE =
+            new KvSnapshotTableLeaseJsonSerde();
+
+    private static final String VERSION_KEY = "version";
+    private static final String TABLE_ID = "table_id";
+    private static final String PARTITION_SNAPSHOTS = "partition_snapshots";
+    private static final String PARTITION_ID = "partition_id";
+    private static final String BUCKET_SNAPSHOTS = "bucket_snapshots";
+
+    private static final int VERSION = 1;
+
+    @Override
+    public void serialize(KvSnapshotTableLease kvSnapshotTableLease, JsonGenerator generator)

Review Comment:
   ditto



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());

Review Comment:
   I noticed that the management state for individual leases is scattered across multiple maps, such as leaseLocks and kvSnapshotLeaseMap. Perhaps all lease-related state—including locks—should be encapsulated into a single cohesive class. This would effectively reduce both code complexity and the complexity of managing concurrent access safely.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(
+                                "kv snapshot lease '"
+                                        + leaseId
+                                        + "' has been acquired. The lease expiration time is "
+                                        + kvSnapshotLease.getExpirationTime());
+                    } else {
+                        kvSnapshotLease = kvSnapshotLeaseMap.get(leaseId);
+                        kvSnapshotLease.setExpirationTime(newLeaseDuration);
+                    }
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataHelper.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(
+                                                tableBucket, originalSnapshotId));
+                            }
+                            incrementRefCount(bucket);
+                        }
+                    }
+
+                    if (update) {
+                        metadataHelper.updateLease(leaseId, kvSnapshotLease);
+                    } else {
+                        metadataHelper.registerLease(leaseId, kvSnapshotLease);
+                    }
+
+                    return unavailableSnapshots;
+                });
+    }
+
+    public void release(String leaseId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotLease lease = kvSnapshotLeaseMap.get(leaseId);
+                    if (lease == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = lease.releaseBucket(bucket);
+                            if (snapshotId != -1L) {
+                                leasedBucketCount.decrementAndGet();
+                                decrementRefCount(new KvSnapshotLeaseForBucket(bucket, snapshotId));
+                            }
+                        }
+                    }
+
+                    if (lease.isEmpty()) {
+                        drop(leaseId);
+                    } else {
+                        metadataHelper.updateLease(leaseId, lease);
+                    }
+                });
+    }
+
+    /**
+     * Drop kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @return true if clear success, false if lease not exist
+     */
+    public boolean drop(String leaseId) throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseMap.remove(leaseId);
+                            if (kvSnapshotLease == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotLease);
+                            metadataHelper.deleteLease(leaseId);
+
+                            LOG.info("kv snapshot lease '" + leaseId + "' has been dropped.");
+                            return true;
+                        });
+
+        leaseLocks.remove(leaseId);
+        return exist;
+    }
+
+    private void initializeRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+
+                    incrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        incrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                    }
+                }
+            }
+        }
+    }
+
+    private void clearRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+                    decrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                    leasedBucketCount.decrementAndGet();
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        decrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                        leasedBucketCount.decrementAndGet();
+                    }
+                }
+            }
+        }
+    }
+
+    private void incrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () ->
+                        refCount.computeIfAbsent(
+                                        kvSnapshotLeaseForBucket, k -> new AtomicInteger(0))
+                                .incrementAndGet());
+    }
+
+    private void decrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger atomicInteger = refCount.get(kvSnapshotLeaseForBucket);
+                    if (atomicInteger != null) {
+                        int decrementAndGet = atomicInteger.decrementAndGet();
+                        if (decrementAndGet <= 0) {
+                            refCount.remove(kvSnapshotLeaseForBucket);
+                        }
+                    }
+                });
+    }
+
+    private void expireLeases() {
+        long currentTime = clock.milliseconds();
+        List<String> expiredLeases =
+                kvSnapshotLeaseMap.entrySet().stream()
+                        .filter(entry -> entry.getValue().getExpirationTime() < currentTime)
+                        .map(Map.Entry::getKey)
+                        .collect(Collectors.toList());
+        for (String leaseId : expiredLeases) {
+            try {
+                drop(leaseId);
+            } catch (Exception e) {
+                LOG.error("Failed to clear kv snapshot lease {}", leaseId, e);
+            }
+        }
+    }
+
+    private void registerMetrics(CoordinatorMetricGroup coordinatorMetricGroup) {
+        coordinatorMetricGroup.gauge(MetricNames.KV_SNAPSHOT_LEASE_COUNT, this::getLeaseCount);
+        // TODO register as table or bucket level.
+        coordinatorMetricGroup.gauge(
+                MetricNames.LEASED_KV_SNAPSHOT_COUNT, this::getLeasedBucketCount);
+    }
+
+    @VisibleForTesting
+    int getLeaseCount() {
+        return kvSnapshotLeaseMap.size();
+    }
+
+    @VisibleForTesting
+    int getLeasedBucketCount() {
+        return leasedBucketCount.get();
+    }
+
+    @VisibleForTesting
+    int getRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    if (!refCount.containsKey(kvSnapshotLeaseForBucket)) {
+                        return 0;
+                    } else {
+                        return refCount.get(kvSnapshotLeaseForBucket).get();

Review Comment:
   This will result in two get calls.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(
+                                "kv snapshot lease '"
+                                        + leaseId
+                                        + "' has been acquired. The lease expiration time is "
+                                        + kvSnapshotLease.getExpirationTime());
+                    } else {
+                        kvSnapshotLease = kvSnapshotLeaseMap.get(leaseId);
+                        kvSnapshotLease.setExpirationTime(newLeaseDuration);
+                    }
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataHelper.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(

Review Comment:
   If the number of buckets reaches hundreds of thousands, would managing buckets via KvSnapshotLeaseForBucket become a memory-expensive model?
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(
+                                "kv snapshot lease '"
+                                        + leaseId
+                                        + "' has been acquired. The lease expiration time is "
+                                        + kvSnapshotLease.getExpirationTime());
+                    } else {
+                        kvSnapshotLease = kvSnapshotLeaseMap.get(leaseId);
+                        kvSnapshotLease.setExpirationTime(newLeaseDuration);
+                    }
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataHelper.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(
+                                                tableBucket, originalSnapshotId));
+                            }
+                            incrementRefCount(bucket);
+                        }
+                    }
+
+                    if (update) {
+                        metadataHelper.updateLease(leaseId, kvSnapshotLease);
+                    } else {
+                        metadataHelper.registerLease(leaseId, kvSnapshotLease);
+                    }
+
+                    return unavailableSnapshots;
+                });
+    }
+
+    public void release(String leaseId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotLease lease = kvSnapshotLeaseMap.get(leaseId);
+                    if (lease == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = lease.releaseBucket(bucket);
+                            if (snapshotId != -1L) {
+                                leasedBucketCount.decrementAndGet();
+                                decrementRefCount(new KvSnapshotLeaseForBucket(bucket, snapshotId));
+                            }
+                        }
+                    }
+
+                    if (lease.isEmpty()) {
+                        drop(leaseId);
+                    } else {
+                        metadataHelper.updateLease(leaseId, lease);
+                    }
+                });
+    }
+
+    /**
+     * Drop kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @return true if clear success, false if lease not exist
+     */
+    public boolean drop(String leaseId) throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseMap.remove(leaseId);
+                            if (kvSnapshotLease == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotLease);
+                            metadataHelper.deleteLease(leaseId);
+
+                            LOG.info("kv snapshot lease '" + leaseId + "' has been dropped.");
+                            return true;
+                        });
+
+        leaseLocks.remove(leaseId);
+        return exist;
+    }
+
+    private void initializeRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+
+                    incrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        incrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                    }
+                }
+            }
+        }
+    }
+
+    private void clearRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+                    decrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                    leasedBucketCount.decrementAndGet();
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        decrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                        leasedBucketCount.decrementAndGet();
+                    }
+                }
+            }
+        }
+    }
+
+    private void incrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () ->
+                        refCount.computeIfAbsent(
+                                        kvSnapshotLeaseForBucket, k -> new AtomicInteger(0))
+                                .incrementAndGet());
+    }
+
+    private void decrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger atomicInteger = refCount.get(kvSnapshotLeaseForBucket);
+                    if (atomicInteger != null) {
+                        int decrementAndGet = atomicInteger.decrementAndGet();
+                        if (decrementAndGet <= 0) {
+                            refCount.remove(kvSnapshotLeaseForBucket);
+                        }
+                    }
+                });
+    }
+
+    private void expireLeases() {
+        long currentTime = clock.milliseconds();
+        List<String> expiredLeases =
+                kvSnapshotLeaseMap.entrySet().stream()
+                        .filter(entry -> entry.getValue().getExpirationTime() < currentTime)
+                        .map(Map.Entry::getKey)
+                        .collect(Collectors.toList());
+        for (String leaseId : expiredLeases) {

Review Comment:
   The logic here is somewhat inefficient—it first creates a list and then iterates over it, resulting in multiple loops and unnecessary copying.
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778797005


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java:
##########
@@ -77,13 +88,53 @@ public FlinkSourceReader(
                 recordEmitter,
                 context.getConfiguration(),
                 context);
+        this.finishedKvSnapshotConsumeBuckets = new HashSet<>();
     }
 
     @Override
     protected void onSplitFinished(Map<String, SourceSplitState> map) {
         // do nothing
     }
 
+    @Override
+    public List<SourceSplitBase> snapshotState(long checkpointId) {
+        Set<TableBucket> bucketsFinishedConsumeKvSnapshot = new HashSet<>();
+
+        // do not modify this state.
+        List<SourceSplitBase> sourceSplitBases = super.snapshotState(checkpointId);
+        for (SourceSplitBase sourceSplitBase : sourceSplitBases) {
+            TableBucket tableBucket = sourceSplitBase.getTableBucket();
+            if (finishedKvSnapshotConsumeBuckets.contains(tableBucket)) {
+                continue;
+            }
+
+            if (sourceSplitBase.isHybridSnapshotLogSplit()) {
+                HybridSnapshotLogSplit hybridSnapshotLogSplit =
+                        sourceSplitBase.asHybridSnapshotLogSplit();
+                if (hybridSnapshotLogSplit.isSnapshotFinished()) {
+                    bucketsFinishedConsumeKvSnapshot.add(tableBucket);
+                }
+            }
+        }
+
+        // report finished kv snapshot consume event.
+        if (!bucketsFinishedConsumeKvSnapshot.isEmpty()) {
+            LOG.info(
+                    "reader has finished kv snapshot read for bucket: {}, checkpoint id {}",
+                    bucketsFinishedConsumeKvSnapshot,
+                    checkpointId);
+
+            context.sendSourceEventToCoordinator(

Review Comment:
   Link issue: https://github.com/apache/fluss/issues/2597



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778756027


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSourceBuilder.java:
##########
@@ -71,6 +72,8 @@ public class FlussSourceBuilder<OUT> {
     private Long scanPartitionDiscoveryIntervalMs;
     private OffsetsInitializer offsetsInitializer;
     private FlussDeserializationSchema<OUT> deserializationSchema;
+    private String kvSnapshotLeaseId;
+    private long kvSnapshotLeaseDurationMs;

Review Comment:
   DataStream API support setting lease id be support in another pr? it always use  `LeaseContext.DEFAULT` now.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "platinumhamburg (via GitHub)" <gi...@apache.org>.
platinumhamburg commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2654962686


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(
+                                "kv snapshot lease '"
+                                        + leaseId
+                                        + "' has been acquired. The lease expiration time is "
+                                        + kvSnapshotLease.getExpirationTime());
+                    } else {
+                        kvSnapshotLease = kvSnapshotLeaseMap.get(leaseId);
+                        kvSnapshotLease.setExpirationTime(newLeaseDuration);
+                    }
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataHelper.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(

Review Comment:
   If the number of buckets reaches hundreds of thousands, would managing buckets via KvSnapshotLeaseForBucket become a memory-expensive model?
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "wuchong (via GitHub)" <gi...@apache.org>.
wuchong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778842331


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSourceBuilder.java:
##########
@@ -71,6 +72,8 @@ public class FlussSourceBuilder<OUT> {
     private Long scanPartitionDiscoveryIntervalMs;
     private OffsetsInitializer offsetsInitializer;
     private FlussDeserializationSchema<OUT> deserializationSchema;
+    private String kvSnapshotLeaseId;
+    private long kvSnapshotLeaseDurationMs;

Review Comment:
   OK



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778670479


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java:
##########
@@ -2035,6 +2073,53 @@ private ControlledShutdownResponse tryProcessControlledShutdown(
         return response;
     }
 
+    private AcquireKvSnapshotLeaseResponse tryProcessAcquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseEvent event) throws Exception {
+        AcquireKvSnapshotLeaseResponse response = new AcquireKvSnapshotLeaseResponse();
+        Map<TableBucket, Long> unavailableSnapshots =
+                kvSnapshotLeaseManager.acquireLease(
+                        event.getLeaseId(),
+                        event.getLeaseDuration(),
+                        event.getTableIdToLeasedBucket());
+
+        Map<Long, List<PbKvSnapshotLeaseForBucket>> pbFailedTables = new HashMap<>();
+        for (Map.Entry<TableBucket, Long> entry : unavailableSnapshots.entrySet()) {
+            TableBucket tb = entry.getKey();
+            Long snapshotId = entry.getValue();
+            PbKvSnapshotLeaseForBucket pbBucket =
+                    new PbKvSnapshotLeaseForBucket().setBucketId(tb.getBucket());
+            if (tb.getPartitionId() != null) {
+                pbBucket.setPartitionId(tb.getPartitionId());
+            }
+            pbBucket.setSnapshotId(snapshotId);
+            pbFailedTables.computeIfAbsent(tb.getTableId(), k -> new ArrayList<>()).add(pbBucket);
+        }
+
+        for (Map.Entry<Long, List<PbKvSnapshotLeaseForBucket>> entry : pbFailedTables.entrySet()) {
+            response.addTablesLeaseRe()
+                    .setTableId(entry.getKey())
+                    .addAllBucketsReqs(entry.getValue());
+        }
+        return response;
+    }
+
+    private ReleaseKvSnapshotLeaseResponse tryProcessReleaseKvSnapshotLease(
+            ReleaseKvSnapshotLeaseEvent event) throws Exception {
+        ReleaseKvSnapshotLeaseResponse response = new ReleaseKvSnapshotLeaseResponse();
+        Map<Long, List<TableBucket>> tableIdToReleasedBucket = event.getTableIdToReleasedBucket();
+        if (tableIdToReleasedBucket.isEmpty()) {
+            // release all
+            boolean exist = kvSnapshotLeaseManager.releaseAll(event.getLeaseId());
+            if (!exist) {
+                throw new KvSnapshotLeaseNotExistException(

Review Comment:
   I change not to throw exception.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "platinumhamburg (via GitHub)" <gi...@apache.org>.
platinumhamburg commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2684178818


##########
fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java:
##########
@@ -412,6 +415,56 @@ CompletableFuture<Void> dropPartition(
     CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
+    /**
+     * Acquires a lease for specific KV snapshots of the given tableBuckets asynchronously.
+     *
+     * <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
+     * #releaseKvSnapshotLease} to release the lock early when reading is finished.
+     *
+     * <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.
+     * </ul>
+     *
+     * @param leaseId The unique ID for this lease session (usually a UUID generated by client).
+     * @param snapshotIds The snapshots to lease, 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<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration);
+
+    /**
+     * Releases the lease for specific tableBuckets asynchronously.
+     *
+     * <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.
+     *
+     * <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);
+
+    /**
+     * Releases 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 leaseId The lease id to release.
+     */
+    CompletableFuture<Void> dropKvSnapshotLease(String leaseId);

Review Comment:
   rename to releaseAllSnapshotLeases()



##########
fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java:
##########
@@ -387,6 +392,33 @@ public CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
                 .thenApply(ClientRpcMessageUtils::toKvSnapshotMetadata);
     }
 
+    @Override
+    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");
+        }
+
+        return gateway.acquireKvSnapshotLease(
+                        makeAcquireKvSnapshotLeaseRequest(leaseId, snapshotIds, leaseDuration))
+                .thenApply(ClientRpcMessageUtils::toAcquireKvSnapshotLeaseResult);
+    }
+
+    @Override
+    public CompletableFuture<Void> releaseKvSnapshotLease(
+            String leaseId, Set<TableBucket> bucketsToRelease) {
+        return gateway.releaseKvSnapshotLease(
+                        makeReleaseKvSnapshotLeaseRequest(leaseId, bucketsToRelease))
+                .thenApply(r -> null);
+    }
+
+    @Override
+    public CompletableFuture<Void> dropKvSnapshotLease(String leaseId) {

Review Comment:
   rename to releaseAllKvSnapshotLeases()



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataHelper.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);

Review Comment:
   In high-availability mode, can this initialization model ensure distributed concurrency safety?
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);

Review Comment:
   It's very odd for an atomic type to be used as a constant.
   
   



##########
fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManagerTest.java:
##########
@@ -0,0 +1,146 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.server.zk.NOPErrorHandler;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.ZooKeeperExtension;
+import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.utils.IOUtils;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.ByteArrayOutputStream;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Test for {@link KvSnapshotLeaseMetadataManager}. */
+public class KvSnapshotLeaseMetadataManagerTest {
+
+    @RegisterExtension
+    public static final AllCallbackWrapper<ZooKeeperExtension> ZOO_KEEPER_EXTENSION_WRAPPER =
+            new AllCallbackWrapper<>(new ZooKeeperExtension());
+
+    protected static ZooKeeperClient zookeeperClient;
+    private @TempDir Path tempDir;
+    private KvSnapshotLeaseMetadataManager metadataHelper;
+
+    @BeforeAll
+    static void beforeAll() {
+        zookeeperClient =
+                ZOO_KEEPER_EXTENSION_WRAPPER
+                        .getCustomExtension()
+                        .getZooKeeperClient(NOPErrorHandler.INSTANCE);
+    }
+
+    @BeforeEach
+    void beforeEach() {
+        metadataHelper = new KvSnapshotLeaseMetadataManager(zookeeperClient, tempDir.toString());
+    }
+
+    @AfterEach
+    void afterEach() {
+        ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupRoot();
+    }
+
+    @Test
+    void testGetLeasesList() throws Exception {
+        List<String> leasesList = metadataHelper.getLeasesList();
+        assertThat(leasesList).isEmpty();
+
+        metadataHelper.registerLease("leaseId1", new KvSnapshotLease(1000L));
+
+        Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+        tableIdToTableLease.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L, -1L}));
+        metadataHelper.registerLease("leaseId2", new KvSnapshotLease(2000L, tableIdToTableLease));
+        leasesList = metadataHelper.getLeasesList();
+        assertThat(leasesList).containsExactlyInAnyOrder("leaseId1", "leaseId2");
+    }
+
+    @Test
+    void testRegisterAndUpdateLease() throws Exception {
+        Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+        tableIdToTableLease.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L, -1L}));
+
+        Map<Long, Long[]> partitionSnapshots = new HashMap<>();
+        partitionSnapshots.put(1000L, new Long[] {111L, 122L});
+        partitionSnapshots.put(1001L, new Long[] {122L, -1L});
+        tableIdToTableLease.put(2L, new KvSnapshotTableLease(2L, partitionSnapshots));
+
+        KvSnapshotLease expectedLease = new KvSnapshotLease(1000L, tableIdToTableLease);
+        metadataHelper.registerLease("leaseId1", expectedLease);

Review Comment:
   The metaHelper in the code can be uniformly renamed to leaseMetadataManager.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLease.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.utils.MapUtils;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
+
+/** The entity of kv snapshot lease. */
+@NotThreadSafe
+public class KvSnapshotLease {
+    private long expirationTime;
+
+    /** A map from table id to kv snapshot lease for one table. */
+    private final Map<Long, KvSnapshotTableLease> tableIdToTableLease;
+
+    public KvSnapshotLease(long leaseDuration) {
+        this(leaseDuration, MapUtils.newConcurrentHashMap());
+    }
+
+    public KvSnapshotLease(
+            long expirationTime, Map<Long, KvSnapshotTableLease> tableIdToTableLease) {

Review Comment:
   expirationTime is a more precise variable and parameter name than duration. However, the codebase extensively mixes both duration and expirationTime; please standardize them.
   
   



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java:
##########
@@ -0,0 +1,36 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+
+/** Procedure to release kv snapshot lease. */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {
+
+    @ProcedureHint(
+            argument = {
+                @ArgumentHint(name = "leaseId", type = @DataTypeHint("STRING")),
+            })
+    public String[] call(ProcedureContext context, String leaseId) throws Exception {

Review Comment:
   Add a description in the procedures.md document.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLease.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.utils.MapUtils;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
+
+/** The entity of kv snapshot lease. */
+@NotThreadSafe
+public class KvSnapshotLease {
+    private long expirationTime;
+
+    /** A map from table id to kv snapshot lease for one table. */
+    private final Map<Long, KvSnapshotTableLease> tableIdToTableLease;
+
+    public KvSnapshotLease(long leaseDuration) {
+        this(leaseDuration, MapUtils.newConcurrentHashMap());
+    }
+
+    public KvSnapshotLease(
+            long expirationTime, Map<Long, KvSnapshotTableLease> tableIdToTableLease) {

Review Comment:
   See `FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_DURATION` also.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataHelper.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        refCount.getOrDefault(kvSnapshotLeaseForBucket, new AtomicInteger(0)).get()
+                                <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(
+                                "kv snapshot lease '"
+                                        + leaseId
+                                        + "' has been acquired. The lease expiration time is "
+                                        + kvSnapshotLease.getExpirationTime());
+                    } else {
+                        kvSnapshotLease = kvSnapshotLeaseMap.get(leaseId);
+                        kvSnapshotLease.setExpirationTime(newLeaseDuration);
+                    }
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataHelper.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(
+                                                tableBucket, originalSnapshotId));
+                            }
+                            incrementRefCount(bucket);
+                        }
+                    }
+
+                    if (update) {
+                        metadataHelper.updateLease(leaseId, kvSnapshotLease);
+                    } else {
+                        metadataHelper.registerLease(leaseId, kvSnapshotLease);
+                    }
+
+                    return unavailableSnapshots;
+                });
+    }
+
+    public void release(String leaseId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotLease lease = kvSnapshotLeaseMap.get(leaseId);
+                    if (lease == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = lease.releaseBucket(bucket);
+                            if (snapshotId != -1L) {
+                                leasedBucketCount.decrementAndGet();
+                                decrementRefCount(new KvSnapshotLeaseForBucket(bucket, snapshotId));
+                            }
+                        }
+                    }
+
+                    if (lease.isEmpty()) {
+                        drop(leaseId);
+                    } else {
+                        metadataHelper.updateLease(leaseId, lease);
+                    }
+                });
+    }
+
+    /**
+     * Drop kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @return true if clear success, false if lease not exist
+     */
+    public boolean drop(String leaseId) throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseMap.remove(leaseId);
+                            if (kvSnapshotLease == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotLease);
+                            metadataHelper.deleteLease(leaseId);
+
+                            LOG.info("kv snapshot lease '" + leaseId + "' has been dropped.");
+                            return true;
+                        });
+
+        leaseLocks.remove(leaseId);
+        return exist;
+    }
+
+    private void initializeRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+
+                    incrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        incrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                    }
+                }
+            }
+        }
+    }
+
+    private void clearRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+                    decrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                    leasedBucketCount.decrementAndGet();
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        decrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                        leasedBucketCount.decrementAndGet();
+                    }
+                }
+            }
+        }
+    }
+
+    private void incrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () ->
+                        refCount.computeIfAbsent(
+                                        kvSnapshotLeaseForBucket, k -> new AtomicInteger(0))
+                                .incrementAndGet());
+    }
+
+    private void decrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger atomicInteger = refCount.get(kvSnapshotLeaseForBucket);
+                    if (atomicInteger != null) {
+                        int decrementAndGet = atomicInteger.decrementAndGet();
+                        if (decrementAndGet <= 0) {
+                            refCount.remove(kvSnapshotLeaseForBucket);
+                        }
+                    }
+                });
+    }
+
+    private void expireLeases() {
+        long currentTime = clock.milliseconds();
+        kvSnapshotLeaseMap.entrySet().stream()
+                .filter(entry -> entry.getValue().getExpirationTime() < currentTime)
+                .map(Map.Entry::getKey)
+                .forEach(
+                        leaseId -> {
+                            try {
+                                drop(leaseId);

Review Comment:
   Although ConcurrentHashMap is used, calling drop to modify the map during iteration may lead to inconsistent behavior.
   Recommendation: first collect the lease IDs that need to be removed, then delete them in bulk.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -132,6 +140,13 @@ public class FlinkSourceEnumerator
     private final OffsetsInitializer startingOffsetsInitializer;
     private final OffsetsInitializer stoppingOffsetsInitializer;
 
+    private final LeaseContext leaseContext;
+
+    private final ReadWriteLock lock = new ReentrantReadWriteLock();
+    /** checkpointId -> tableBuckets who finished consume kv snapshots. */
+    @GuardedBy("lock")
+    private final TreeMap<Long, Set<TableBucket>> consumedKvSnapshotMap = new TreeMap<>();

Review Comment:
   Is there a concurrency safety issue here? Why is a read-write lock needed?
   
   



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java:
##########
@@ -0,0 +1,36 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+
+/** Procedure to release kv snapshot lease. */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {

Review Comment:
   It's recommended to improve the comments and include usage examples.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;

Review Comment:
   You've annotated KvSnapshotLeaseManager with @ThreadSafe, which implies that it may be used in a multi-threaded environment and makes no guarantee that it shares a thread with CoordinatorEventProcessor. Under this assumption, directly accessing CoordinatorContext—which is marked as @NotThreadSafe—introduces a thread-safety model violation.
   
   



##########
website/docs/maintenance/configuration.md:
##########
@@ -157,7 +158,7 @@ during the Fluss cluster working.
 | kv.rocksdb.use-bloom-filter                       | Boolean    | true                          | If true, every newly created SST file will contain a Bloom filter. It is enabled by default.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
 | kv.rocksdb.bloom-filter.bits-per-key              | Double     | 10.0                          | Bits per key that bloom filter will use, this only take effect when bloom filter is used. The default value is 10.0.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
 | kv.rocksdb.bloom-filter.block-based-mode          | Boolean    | false                         | If true, RocksDB will use block-based filter instead of full filter, this only take effect when bloom filter is used. The default value is `false`.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
-| kv.rocksdb.shared-rate-limiter-bytes-per-sec              | MemorySize | Long.MAX_VALUE                | The bytes per second rate limit for RocksDB flush and compaction operations shared across all RocksDB instances on the TabletServer. The rate limiter is always enabled. The default value is Long.MAX_VALUE (effectively unlimited). Set to a lower value (e.g., 100MB) to limit the rate. This configuration can be updated dynamically without server restart. See [Updating Configs](operations/updating-configs.md) for more details.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          |
+| kv.rocksdb.shared-rate-limiter-bytes-per-sec      | MemorySize | Long.MAX_VALUE                | The bytes per second rate limit for RocksDB flush and compaction operations shared across all RocksDB instances on the TabletServer. The rate limiter is always enabled. The default value is Long.MAX_VALUE (effectively unlimited). Set to a lower value (e.g., 100MB) to limit the rate. This configuration can be updated dynamically without server restart. See [Updating Configs](operations/updating-configs.md) for more details.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |

Review Comment:
   Clean up unrelated changes.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;

Review Comment:
   Rename to metadataManager



##########
website/docs/maintenance/observability/monitor-metrics.md:
##########
@@ -325,6 +325,16 @@ Some metrics might not be exposed when using other JVM implementations (e.g. IBM
       <td>The total number of partitions in this cluster.</td>
       <td>Gauge</td>
     </tr>
+    <tr>
+      <td>kvSnapshotLeaseCount</td>
+      <td>The total number of kv snapshots in this cluster.</td>

Review Comment:
   The metric name does not match its description.
   
   



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +543,96 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease util all buckets
+        // acquire success. The reason is that getLatestKvSnapshots and acquireKvSnapshotLease
+        // are not atomic operations, the latest kv snapshot obtained via get may become outdated by
+        // the time it is passed to acquire. Therefore, this logic must implement a retry
+        // mechanism: the unavailable tableBuckets in the AcquiredKvSnapshotLeaseResult returned by
+        // acquireKvSnapshotLease must be retried repeatedly until all buckets are successfully
+        // acquired.
+        try {
+            Set<TableBucket> remainingTableBuckets;
+            do {
+                KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+                remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+                tableId = kvSnapshots.getTableId();
+                partitionId = kvSnapshots.getPartitionId();
+
+                Set<TableBucket> ignoreBuckets = new HashSet<>();
+                Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+                for (TableBucket tb : remainingTableBuckets) {
+                    int bucket = tb.getBucket();
+                    OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                    OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                    if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                        bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                    } else {
+                        ignoreBuckets.add(tb);
+                    }
+
+                    snapshotIds.put(
+                            bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                    logOffsets.put(
+                            bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+                }
+
+                if (!ignoreBuckets.isEmpty()) {
+                    remainingTableBuckets.removeAll(ignoreBuckets);
+                }
+
+                if (!bucketsToLease.isEmpty()) {
+                    String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                    LOG.info(
+                            "Try to acquire kv snapshot lease {} for table {}",
+                            kvSnapshotLeaseId,
+                            tablePath);
+                    remainingTableBuckets =
+                            flussAdmin
+                                    .acquireKvSnapshotLease(
+                                            kvSnapshotLeaseId,
+                                            bucketsToLease,
+                                            leaseContext.getKvSnapshotLeaseDurationMs())

Review Comment:
   acquireKvSnapshotLease() may throw NullPointerException when [null == easeContext.getKvSnapshotLeaseDurationMs()]



##########
website/docs/engine-flink/options.md:
##########
@@ -91,20 +91,22 @@ See more details about [ALTER TABLE ... SET](engine-flink/ddl.md#set-properties)
 
 ## Read Options
 
-| Option                                              | Type       | Default                                         | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
-|-----------------------------------------------------|------------|-------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| scan.startup.mode                                   | Enum       | full                                            | The scan startup mode enables you to specify the starting point for data consumption. Fluss currently supports the following `scan.startup.mode` options: `full` (default), earliest, latest, timestamp. See the [Start Reading Position](engine-flink/reads.md#start-reading-position) for more details.                                                                                                                                                                                                                                          |
-| scan.startup.timestamp                              | Long       | (None)                                          | The timestamp to start reading the data from. This option is only valid when `scan.startup.mode` is set to `timestamp`. The format is 'milli-second-since-epoch' or `yyyy-MM-dd HH:mm:ss`, like `1678883047356` or `2023-12-09 23:09:12`.                                                                                                                                                                                                                                                                                                          |
-| scan.partition.discovery.interval                   | Duration   | 1min                                            | The time interval for the Fluss source to discover the new partitions for partitioned table while scanning. A non-positive value disables the partition discovery. The default value is 1 minute. Currently, since Fluss Admin#listPartitions(TablePath tablePath) requires a large number of requests to ZooKeeper in server, this option cannot be set too small, as a small value would cause frequent requests and increase server load. In the future, once list partitions is optimized, the default value of this parameter can be reduced. |
-| client.scanner.log.check-crc                        | Boolean    | true                                            | Automatically check the CRC3 of the read records for LogScanner. This ensures no on-the-wire or on-disk corruption to the messages occurred. This check adds some overhead, so it may be disabled in cases seeking extreme performance.                                                                                                                                                                                                                                                                                                            |
-| client.scanner.log.max-poll-records                 | Integer    | 500                                             | The maximum number of records returned in a single call to poll() for LogScanner. Note that this config doesn't impact the underlying fetching behavior. The Scanner will cache the records from each fetch request and returns them incrementally from each poll.                                                                                                                                                                                                                                                                                 |
-| client.scanner.log.fetch.max-bytes                  | MemorySize | 16mb                                            | The maximum amount of data the server should return for a fetch request from client. Records are fetched in batches, and if the first record batch in the first non-empty bucket of the fetch is larger than this value, the record batch will still be returned to ensure that the fetch can make progress. As such, this is not a absolute maximum.                                                                                                                                                                                              |
-| client.scanner.log.fetch.max-bytes-for-bucket       | MemorySize | 1mb                                             | The maximum amount of data the server should return for a table bucket in fetch request fom client. Records are fetched in batches, and the max bytes size is config by this option.                                                                                                                                                                                                                                                                                                                                                               |
-| client.scanner.log.fetch.min-bytes                  | MemorySize | 1b                                              | The minimum bytes expected for each fetch log request from client to response. If not enough bytes, wait up to client.scanner.log.fetch-wait-max-time time to return.                                                                                                                                                                                                                                                                                                                                                                              |
-| client.scanner.log.fetch.wait-max-time              | Duration   | 500ms                                           | The maximum time to wait for enough bytes to be available for a fetch log request from client to response.                                                                                                                                                                                                                                                                                                                                                                                                                                         |
-| client.scanner.io.tmpdir                            | String     | System.getProperty("java.io.tmpdir") + "/fluss" | Local directory that is used by client for storing the data files (like kv snapshot, log segment files) to read temporarily                                                                                                                                                                                                                                                                                                                                                                                                                        |
-| client.scanner.remote-log.prefetch-num              | Integer    | 4                                               | The number of remote log segments to keep in local temp file for LogScanner, which download from remote storage. The default setting is 4.                                                                                                                                                                                                                                                                                                                                                                                                         |
-| client.remote-file.download-thread-num              | Integer    | 3                                               | The number of threads the client uses to download remote files.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+| Option                                        | Type       | Default                                         | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
+|-----------------------------------------------|------------|-------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| scan.startup.mode                             | Enum       | full                                            | The scan startup mode enables you to specify the starting point for data consumption. Fluss currently supports the following `scan.startup.mode` options: `full` (default), earliest, latest, timestamp. See the [Start Reading Position](engine-flink/reads.md#start-reading-position) for more details.                                                                                                                                                                                                                                          |
+| scan.startup.timestamp                        | Long       | (None)                                          | The timestamp to start reading the data from. This option is only valid when `scan.startup.mode` is set to `timestamp`. The format is 'milli-second-since-epoch' or `yyyy-MM-dd HH:mm:ss`, like `1678883047356` or `2023-12-09 23:09:12`.                                                                                                                                                                                                                                                                                                          |
+| scan.partition.discovery.interval             | Duration   | 1min                                            | The time interval for the Fluss source to discover the new partitions for partitioned table while scanning. A non-positive value disables the partition discovery. The default value is 1 minute. Currently, since Fluss Admin#listPartitions(TablePath tablePath) requires a large number of requests to ZooKeeper in server, this option cannot be set too small, as a small value would cause frequent requests and increase server load. In the future, once list partitions is optimized, the default value of this parameter can be reduced. |
+| scan.kv.snapshot.lease.id                     | String     | UUID                                            | The lease id to lease kv snapshots. If set, the acquired kv snapshots will not be deleted until the consumer finished consuming all the snapshots or the lease duration time is reached. If not set, an UUID will be set.                                                                                                                                                                                                                                                                                                                          |
+| scan.kv.snapshot.lease.duration               | Duration   | 1day                                            | The time period how long to wait before expiring the kv snapshot lease to avoid kv snapshot blocking to delete.                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+| client.scanner.log.check-crc                  | Boolean    | true                                            | Automatically check the CRC3 of the read records for LogScanner. This ensures no on-the-wire or on-disk corruption to the messages occurred. This check adds some overhead, so it may be disabled in cases seeking extreme performance.                                                                                                                                                                                                                                                                                                            |
+| client.scanner.log.max-poll-records           | Integer    | 500                                             | The maximum number of records returned in a single call to poll() for LogScanner. Note that this config doesn't impact the underlying fetching behavior. The Scanner will cache the records from each fetch request and returns them incrementally from each poll.                                                                                                                                                                                                                                                                                 |
+| client.scanner.log.fetch.max-bytes            | MemorySize | 16mb                                            | The maximum amount of data the server should return for a fetch request from client. Records are fetched in batches, and if the first record batch in the first non-empty bucket of the fetch is larger than this value, the record batch will still be returned to ensure that the fetch can make progress. As such, this is not a absolute maximum.                                                                                                                                                                                              |
+| client.scanner.log.fetch.max-bytes-for-bucket | MemorySize | 1mb                                             | The maximum amount of data the server should return for a table bucket in fetch request fom client. Records are fetched in batches, and the max bytes size is config by this option.                                                                                                                                                                                                                                                                                                                                                               |
+| client.scanner.log.fetch.min-bytes            | MemorySize | 1b                                              | The minimum bytes expected for each fetch log request from client to response. If not enough bytes, wait up to client.scanner.log.fetch-wait-max-time time to return.                                                                                                                                                                                                                                                                                                                                                                              |
+| client.scanner.log.fetch.wait-max-time        | Duration   | 500ms                                           | The maximum time to wait for enough bytes to be available for a fetch log request from client to response.                                                                                                                                                                                                                                                                                                                                                                                                                                         |
+| client.scanner.io.tmpdir                      | String     | System.getProperty("java.io.tmpdir") + "/fluss" | Local directory that is used by client for storing the data files (like kv snapshot, log segment files) to read temporarily                                                                                                                                                                                                                                                                                                                                                                                                                        |
+| client.scanner.remote-log.prefetch-num        | Integer    | 4                                               | The number of remote log segments to keep in local temp file for LogScanner, which download from remote storage. The default setting is 4.                                                                                                                                                                                                                                                                                                                                                                                                         |

Review Comment:
   Clean up unrelated changes.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataHelper.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        refCount.getOrDefault(kvSnapshotLeaseForBucket, new AtomicInteger(0)).get()
+                                <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(
+                                "kv snapshot lease '"
+                                        + leaseId
+                                        + "' has been acquired. The lease expiration time is "
+                                        + kvSnapshotLease.getExpirationTime());
+                    } else {
+                        kvSnapshotLease = kvSnapshotLeaseMap.get(leaseId);

Review Comment:
   Using containsKey() followed by get() is a bad pattern, as it entails querying the map multiple times. Moreover, this approach is overly complex and can be simplified using computeIfAbsent().
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataHelper.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        refCount.getOrDefault(kvSnapshotLeaseForBucket, new AtomicInteger(0)).get()

Review Comment:
   new AtomicInteger() here is inefficient.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/LeaseContext.java:
##########
@@ -0,0 +1,79 @@
+/*
+ * 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.flink.source.reader;
+
+import javax.annotation.Nullable;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+/** Context for lease. */
+public class LeaseContext implements Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    // kv snapshot lease id. null for log table.
+    private final @Nullable String kvSnapshotLeaseId;
+
+    // kv snapshot lease duration. null for log table.
+    private final @Nullable Long kvSnapshotLeaseDurationMs;
+
+    public LeaseContext(
+            @Nullable String kvSnapshotLeaseId, @Nullable Long kvSnapshotLeaseDurationMs) {
+        this.kvSnapshotLeaseId = kvSnapshotLeaseId;
+        this.kvSnapshotLeaseDurationMs = kvSnapshotLeaseDurationMs;
+    }
+
+    public @Nullable String getKvSnapshotLeaseId() {
+        return kvSnapshotLeaseId;
+    }
+
+    public @Nullable Long getKvSnapshotLeaseDurationMs() {

Review Comment:
   Why allow nullable?



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java:
##########
@@ -0,0 +1,36 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+
+/** Procedure to drop kv snapshot lease. */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {

Review Comment:
   I recommend to update the API name. Although it differs slightly from the FIP, it should be fine with an explanatory comment.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java:
##########
@@ -0,0 +1,226 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.types.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * The manager to handle {@link KvSnapshotLease} to register/update/delete metadata from zk and
+ * remote fs.
+ */
+public class KvSnapshotLeaseMetadataManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseMetadataManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final String remoteDataDir;
+
+    public KvSnapshotLeaseMetadataManager(ZooKeeperClient zkClient, String remoteDataDir) {
+        this.zkClient = zkClient;
+        this.remoteDataDir = remoteDataDir;
+    }
+
+    public List<String> getLeasesList() throws Exception {
+        return zkClient.getKvSnapshotLeasesList();
+    }
+
+    /**
+     * Register a new kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param lease the kv snapshot lease.
+     */
+    public void registerLease(String leaseId, KvSnapshotLease lease) throws Exception {
+        Map<Long, FsPath> tableIdToRemoteMetadataFsPath = generateMetadataFile(leaseId, lease);
+
+        // generate remote fsPath of metadata.
+        KvSnapshotLeaseMetadata leaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        lease.getExpirationTime(), tableIdToRemoteMetadataFsPath);
+
+        // register kv snapshot metadata to zk.
+        try {
+            zkClient.registerKvSnapshotLeaseMetadata(leaseId, leaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to register kv snapshot lease metadata to zk.", e);
+            leaseMetadata.discard();
+            throw e;
+        }
+    }
+
+    /**
+     * Update a kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param kvSnapshotLease the kv snapshot lease.
+     */
+    public void updateLease(String leaseId, KvSnapshotLease kvSnapshotLease) throws Exception {
+        // TODO change this to incremental update to avoid create too many remote metadata files.
+
+        Optional<KvSnapshotLeaseMetadata> originalLeaseMetadata =
+                zkClient.getKvSnapshotLeaseMetadata(leaseId);
+
+        Map<Long, FsPath> tableIdToNewRemoteMetadataFsPath =
+                generateMetadataFile(leaseId, kvSnapshotLease);
+
+        // generate  new kv snapshot lease metadata.
+        KvSnapshotLeaseMetadata newLeaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        kvSnapshotLease.getExpirationTime(), tableIdToNewRemoteMetadataFsPath);
+        // register new snapshot metadata to zk.
+        try {
+            zkClient.updateKvSnapshotLeaseMetadata(leaseId, newLeaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to update kv snapshot lease metadata to zk.", e);
+            newLeaseMetadata.discard();

Review Comment:
   Perhaps we should have a mechanism to scan and clean up orphaned lease files during initialization.



##########
fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java:
##########
@@ -200,14 +208,21 @@ interface SubsumeAction {
         void subsume(CompletedSnapshot snapshot) throws Exception;
     }
 
-    private static boolean canSubsume(CompletedSnapshot next, CompletedSnapshot latest) {
+    /** A function to determine whether a snapshot can be subsumed. */
+    @FunctionalInterface
+    public interface CanSubsume {

Review Comment:
   “CanSubsume” as a class name isn’t ideal—it reads more like a predicate method (e.g., canSubsume()) than a noun that represents a type or entity. In object-oriented design, class names should typically be nouns that clearly convey what the class is or represents. (e.g., SubsumptionChecker).



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +543,96 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease util all buckets
+        // acquire success. The reason is that getLatestKvSnapshots and acquireKvSnapshotLease
+        // are not atomic operations, the latest kv snapshot obtained via get may become outdated by
+        // the time it is passed to acquire. Therefore, this logic must implement a retry
+        // mechanism: the unavailable tableBuckets in the AcquiredKvSnapshotLeaseResult returned by
+        // acquireKvSnapshotLease must be retried repeatedly until all buckets are successfully
+        // acquired.
+        try {
+            Set<TableBucket> remainingTableBuckets;
+            do {
+                KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+                remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+                tableId = kvSnapshots.getTableId();
+                partitionId = kvSnapshots.getPartitionId();
+
+                Set<TableBucket> ignoreBuckets = new HashSet<>();
+                Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+                for (TableBucket tb : remainingTableBuckets) {
+                    int bucket = tb.getBucket();
+                    OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                    OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                    if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                        bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                    } else {
+                        ignoreBuckets.add(tb);
+                    }
+
+                    snapshotIds.put(
+                            bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                    logOffsets.put(
+                            bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+                }
+
+                if (!ignoreBuckets.isEmpty()) {
+                    remainingTableBuckets.removeAll(ignoreBuckets);
+                }
+
+                if (!bucketsToLease.isEmpty()) {
+                    String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                    LOG.info(
+                            "Try to acquire kv snapshot lease {} for table {}",
+                            kvSnapshotLeaseId,
+                            tablePath);
+                    remainingTableBuckets =
+                            flussAdmin
+                                    .acquireKvSnapshotLease(
+                                            kvSnapshotLeaseId,
+                                            bucketsToLease,
+                                            leaseContext.getKvSnapshotLeaseDurationMs())
+                                    .get()
+                                    .getUnavailableTableBucketSet();
+                    if (!remainingTableBuckets.isEmpty()) {
+                        LOG.info(
+                                "Failed to acquire kv snapshot lease for table {}: {}. Retry to re-acquire",
+                                tablePath,
+                                remainingTableBuckets);
+                    }
+                }
+            } while (!remainingTableBuckets.isEmpty());

Review Comment:
   Do not use infinite retry loops, and add a delay between requests to avoid retry storms.
   
   



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java:
##########
@@ -51,6 +57,11 @@ public class FlinkSourceReader<OUT>
         extends SingleThreadMultiplexSourceReaderBaseAdapter<
                 RecordAndPos, OUT, SourceSplitBase, SourceSplitState> {
 
+    private static final Logger LOG = LoggerFactory.getLogger(FlinkSourceReader.class);
+
+    /** the tableBuckets ignore to send FinishedKvSnapshotConsumeEvent as it already sending. */
+    private final Set<TableBucket> ignoreBuckets;

Review Comment:
   The variable name does not clearly indicate its purpose or usage.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();

Review Comment:
   The current state management design using Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount is inefficient. I strongly recommend using Map<TableBucket, ConcurrentSet<String /* leaseId */>> instead, for two reasons:
   ① Every method that accesses this map in the code must instantiate a new KvSnapshotLeaseForBucket object, creating unnecessary GC pressure.
   ② Maintaining the correctness of refCount is extremely difficult, which severely degrades code readability and leads to code rot.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLease.java:
##########
@@ -0,0 +1,188 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.utils.MapUtils;
+
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+
+/** The lease of kv snapshot for a table. */
+@NotThreadSafe
+public class KvSnapshotTableLease {
+    private final long tableId;
+    private final @Nullable Long[] bucketSnapshots;
+    private final Map<Long, Long[]> partitionSnapshots;
+
+    public KvSnapshotTableLease(long tableId) {
+        this(tableId, null, MapUtils.newConcurrentHashMap());
+    }
+
+    public KvSnapshotTableLease(long tableId, Long[] bucketSnapshots) {
+        this(tableId, bucketSnapshots, Collections.emptyMap());
+    }
+
+    public KvSnapshotTableLease(long tableId, Map<Long, Long[]> partitionSnapshots) {
+        this(tableId, null, partitionSnapshots);
+    }
+
+    public KvSnapshotTableLease(
+            long tableId, @Nullable Long[] bucketSnapshots, Map<Long, Long[]> partitionSnapshots) {
+        this.tableId = tableId;
+        this.bucketSnapshots = bucketSnapshots;
+        this.partitionSnapshots = partitionSnapshots;
+    }
+
+    public long getTableId() {
+        return tableId;
+    }
+
+    public @Nullable Long[] getBucketSnapshots() {
+        return bucketSnapshots;
+    }
+
+    public @Nullable Long[] getBucketSnapshots(long partitionId) {
+        return partitionSnapshots.get(partitionId);
+    }
+
+    public Map<Long, Long[]> getPartitionSnapshots() {
+        return partitionSnapshots;
+    }
+
+    public void addPartitionSnapshots(long partitionId, Long[] snapshots) {
+        if (bucketSnapshots != null) {
+            throw new IllegalStateException("This is an none partition table lease.");
+        }
+        partitionSnapshots.put(partitionId, snapshots);
+    }
+
+    public int getLeasedSnapshotCount() {
+        int count = 0;
+        if (bucketSnapshots != null) {
+            for (Long snapshot : bucketSnapshots) {
+                if (snapshot != -1L) {
+                    count++;
+                }
+            }
+        } else {
+            for (Long[] snapshots : partitionSnapshots.values()) {
+                for (Long snapshot : snapshots) {
+                    if (snapshot != -1L) {
+                        count++;
+                    }
+                }
+            }
+        }
+        return count;
+    }
+
+    @Override
+    public String toString() {
+        String partitionSnapshotsStr = formatLongArrayMap(partitionSnapshots);
+        return "KvSnapshotTableLease{"
+                + "tableId="
+                + tableId
+                + ", bucketSnapshots="
+                + Arrays.toString(bucketSnapshots)
+                + ", partitionSnapshots="
+                + partitionSnapshotsStr
+                + '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        KvSnapshotTableLease that = (KvSnapshotTableLease) o;
+        return tableId == that.tableId
+                && Arrays.equals(bucketSnapshots, that.bucketSnapshots)
+                && deepEqualsMapOfArrays(partitionSnapshots, that.partitionSnapshots);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = Objects.hash(tableId);
+        result = 31 * result + Arrays.hashCode(bucketSnapshots);
+        result = 31 * result + deepHashCodeMapOfArrays(partitionSnapshots);
+        return result;
+    }
+
+    private static String formatLongArrayMap(Map<Long, Long[]> map) {
+        if (map == null) {
+            return "null";
+        }
+        StringBuilder sb = new StringBuilder("{");
+        boolean first = true;
+        for (Map.Entry<Long, Long[]> entry : map.entrySet()) {
+            if (!first) {
+                sb.append(", ");
+            }
+            sb.append(entry.getKey()).append("=").append(Arrays.toString(entry.getValue()));
+            first = false;
+        }
+        sb.append("}");
+        return sb.toString();
+    }
+
+    private static boolean deepEqualsMapOfArrays(Map<Long, Long[]> map1, Map<Long, Long[]> map2) {
+        if (map1 == map2) {
+            return true;
+        }
+        if (map1 == null || map2 == null || map1.size() != map2.size()) {
+            return false;
+        }
+
+        for (Map.Entry<Long, Long[]> entry : map1.entrySet()) {
+            Long key = entry.getKey();
+            Long[] value1 = entry.getValue();
+            Long[] value2 = map2.get(key);
+
+            if (value2 == null) {
+                return false;
+            }
+
+            if (!Arrays.equals(value1, value2)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    private static int deepHashCodeMapOfArrays(Map<Long, Long[]> map) {
+        if (map == null) {
+            return 0;
+        }
+        int hash = 0;
+        for (Map.Entry<Long, Long[]> entry : map.entrySet()) {
+            Long key = entry.getKey();
+            Long[] value = entry.getValue();
+            // Combine key hash and array content hash
+            hash += Objects.hashCode(key) ^ Arrays.hashCode(value);

Review Comment:
   hash = 31 * hash + (Objects.hashCode(key) ^ Arrays.hashCode(value));



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataHelper.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);

Review Comment:
   What I mean is, after a standby node is initialized, can it stay consistent with the primary node? And if a failover occurs, can it correctly handle the business logic?
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataHelper.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        refCount.getOrDefault(kvSnapshotLeaseForBucket, new AtomicInteger(0)).get()
+                                <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    KvSnapshotLease kvSnapshotLease;
+                    long newLeaseDuration = clock.milliseconds() + leaseDuration;
+                    if (!update) {
+                        // set the expiration time as: current time + leaseDuration
+                        kvSnapshotLease = new KvSnapshotLease(newLeaseDuration);
+                        kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+                        LOG.info(

Review Comment:
   Don't use string concatenation to print logs.
   
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,443 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
+
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
+    private final CoordinatorContext coordinatorContext;

Review Comment:
   Although in practice we only access coordinatorContext within acquireLease, and acquireLease is exclusively invoked by coordinatorProcessor, meaning there is actually no thread-safety issue, this thread-safety guarantee is overly implicit and obscure—indicating a bad code smell. Other maintainers of the codebase could easily misunderstand or inadvertently modify the code, thereby breaking the implicit safety assumptions. While refactoring and fixing this properly might be costly, we can at least add comments to explicitly document the thread-safety constraints here.
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2703138735


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +543,96 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease util all buckets
+        // acquire success. The reason is that getLatestKvSnapshots and acquireKvSnapshotLease
+        // are not atomic operations, the latest kv snapshot obtained via get may become outdated by
+        // the time it is passed to acquire. Therefore, this logic must implement a retry
+        // mechanism: the unavailable tableBuckets in the AcquiredKvSnapshotLeaseResult returned by
+        // acquireKvSnapshotLease must be retried repeatedly until all buckets are successfully
+        // acquired.
+        try {
+            Set<TableBucket> remainingTableBuckets;
+            do {
+                KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+                remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+                tableId = kvSnapshots.getTableId();
+                partitionId = kvSnapshots.getPartitionId();
+
+                Set<TableBucket> ignoreBuckets = new HashSet<>();
+                Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+                for (TableBucket tb : remainingTableBuckets) {
+                    int bucket = tb.getBucket();
+                    OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                    OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                    if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                        bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                    } else {
+                        ignoreBuckets.add(tb);
+                    }
+
+                    snapshotIds.put(
+                            bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                    logOffsets.put(
+                            bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+                }
+
+                if (!ignoreBuckets.isEmpty()) {
+                    remainingTableBuckets.removeAll(ignoreBuckets);
+                }
+
+                if (!bucketsToLease.isEmpty()) {
+                    String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                    LOG.info(
+                            "Try to acquire kv snapshot lease {} for table {}",
+                            kvSnapshotLeaseId,
+                            tablePath);
+                    remainingTableBuckets =
+                            flussAdmin
+                                    .acquireKvSnapshotLease(
+                                            kvSnapshotLeaseId,
+                                            bucketsToLease,
+                                            leaseContext.getKvSnapshotLeaseDurationMs())
+                                    .get()
+                                    .getUnavailableTableBucketSet();
+                    if (!remainingTableBuckets.isEmpty()) {
+                        LOG.info(
+                                "Failed to acquire kv snapshot lease for table {}: {}. Retry to re-acquire",
+                                tablePath,
+                                remainingTableBuckets);
+                    }
+                }
+            } while (!remainingTableBuckets.isEmpty());

Review Comment:
   I will remove this loop first as we choose to reserve more than one kv snapshot.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778878702


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataManager.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger count = refCount.get(kvSnapshotLeaseForBucket);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    // set the expiration time as: current time + leaseDuration
+                    long newExpirationTime = clock.milliseconds() + leaseDuration;
+                    KvSnapshotLease kvSnapshotLease =
+                            kvSnapshotLeaseMap.compute(
+                                    leaseId,
+                                    (key, existingLease) -> {
+                                        if (existingLease == null) {
+                                            LOG.info(
+                                                    "kv snapshot lease '{}' has been acquired. The lease expiration "
+                                                            + "time is {}",
+                                                    leaseId,
+                                                    newExpirationTime);
+                                            return new KvSnapshotLease(newExpirationTime);
+                                        } else {
+                                            existingLease.setExpirationTime(newExpirationTime);
+                                            return existingLease;
+                                        }
+                                    });
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataManager.isSnapshotExists(tableBucket, kvSnapshotId);

Review Comment:
   Link issue: https://github.com/apache/fluss/issues/2600



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2780266369


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,427 @@
+/*
+ * 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.server.coordinator.lease;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableBucketSnapshot;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final long leaseExpirationCheckInterval;
+
+    private final ReadWriteLock managerLock = new ReentrantReadWriteLock();
+
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("managerLock")
+    private final ConcurrentHashMap<String, KvSnapshotLeaseHandler> kvSnapshotLeaseMap =
+            MapUtils.newConcurrentHashMap();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("managerLock")
+    private final Map<TableBucketSnapshot, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            long leaseExpirationCheckInterval,
+            ZooKeeperClient zkClient,
+            String remoteDataDir,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                leaseExpirationCheckInterval,
+                zkClient,
+                remoteDataDir,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            long leaseExpirationCheckInterval,
+            ZooKeeperClient zkClient,
+            String remoteDataDir,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = new KvSnapshotLeaseMetadataManager(zkClient, remoteDataDir);
+        this.leaseExpirationCheckInterval = leaseExpirationCheckInterval;
+        this.scheduledExecutor = scheduledExecutor;
+        this.clock = clock;
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+
+        List<String> leasesList = new ArrayList<>();
+        try {
+            leasesList = metadataManager.getLeasesList();
+        } catch (Exception e) {
+            LOG.error("Failed to get leases list from zookeeper.", e);
+        }
+
+        for (String leaseId : leasesList) {
+            Optional<KvSnapshotLeaseHandler> kvSnapshotLeaseOpt = Optional.empty();
+            try {
+                kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            } catch (Exception e) {
+                LOG.error("Failed to get kv snapshot lease from zookeeper.", e);
+            }
+
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLeaseHandler kvSnapshotLeasehandle = kvSnapshotLeaseOpt.get();
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLeasehandle);
+
+                initializeRefCount(kvSnapshotLeasehandle);
+
+                leasedBucketCount.addAndGet(kvSnapshotLeasehandle.getLeasedSnapshotCount());
+            }
+        }
+
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases, 0L, leaseExpirationCheckInterval, TimeUnit.MILLISECONDS);
+    }
+
+    public boolean snapshotLeaseNotExist(TableBucketSnapshot tableBucketSnapshot) {
+        return inReadLock(
+                managerLock,
+                () -> {
+                    AtomicInteger count = refCount.get(tableBucketSnapshot);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(

Review Comment:
   trace by: https://github.com/apache/fluss/issues/2603



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2780970813


##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java:
##########
@@ -352,6 +354,46 @@ void testPkTableReadMixSnapshotAndLog() throws Exception {
         assertResultsIgnoreOrder(rowIter, expectedRows, true);
     }
 
+    @Test
+    void testPkTableReadWithKvSnapshotLease() throws Exception {
+        tEnv.executeSql(
+                "create table pk_table_with_kv_snapshot_lease (a int not null primary key not enforced, b varchar)");
+        TablePath tablePath = TablePath.of(DEFAULT_DB, "pk_table_with_kv_snapshot_lease");
+
+        List<InternalRow> rows = Arrays.asList(row(1, "v1"), row(2, "v2"), row(3, "v3"));
+
+        // write records
+        writeRows(conn, tablePath, rows, false);
+
+        FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshot(tablePath);
+
+        // enable checkpoint to make sure the kv snapshot lease will be cleared.
+        execEnv.enableCheckpointing(100);

Review Comment:
   Trace by: https://github.com/apache/fluss/issues/2605



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "wuchong (via GitHub)" <gi...@apache.org>.
wuchong merged PR #2179:
URL: https://github.com/apache/fluss/pull/2179


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2703003966


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java:
##########
@@ -0,0 +1,36 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+
+/** Procedure to drop kv snapshot lease. */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {

Review Comment:
   i will change to `releaseAllKvSnapshotLease(String leaseId)`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2703107337


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/LeaseContext.java:
##########
@@ -0,0 +1,79 @@
+/*
+ * 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.flink.source.reader;
+
+import javax.annotation.Nullable;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+/** Context for lease. */
+public class LeaseContext implements Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    // kv snapshot lease id. null for log table.
+    private final @Nullable String kvSnapshotLeaseId;
+
+    // kv snapshot lease duration. null for log table.
+    private final @Nullable Long kvSnapshotLeaseDurationMs;
+
+    public LeaseContext(
+            @Nullable String kvSnapshotLeaseId, @Nullable Long kvSnapshotLeaseDurationMs) {
+        this.kvSnapshotLeaseId = kvSnapshotLeaseId;
+        this.kvSnapshotLeaseDurationMs = kvSnapshotLeaseDurationMs;
+    }
+
+    public @Nullable String getKvSnapshotLeaseId() {
+        return kvSnapshotLeaseId;
+    }
+
+    public @Nullable Long getKvSnapshotLeaseDurationMs() {

Review Comment:
   This is mainly for compatibility when restoring from old state. Perhaps we can provide a default value. I will change to set default value.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "platinumhamburg (via GitHub)" <gi...@apache.org>.
platinumhamburg commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2654927818


##########
fluss-common/src/main/java/org/apache/fluss/metadata/KvSnapshotLeaseForBucket.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.metadata;
+
+import java.util.Objects;
+
+/** An entity for kv snapshot lease for bucket. */
+public class KvSnapshotLeaseForBucket {

Review Comment:
   A lease inherently implies a time-bound reservation—without a TTL, it’s just a static reference. Consider adding a ttl (e.g., in milliseconds) to this class to properly reflect lease semantics.
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2702398182


##########
website/docs/engine-flink/options.md:
##########
@@ -91,20 +91,22 @@ See more details about [ALTER TABLE ... SET](engine-flink/ddl.md#set-properties)
 
 ## Read Options
 
-| Option                                              | Type       | Default                                         | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
-|-----------------------------------------------------|------------|-------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| scan.startup.mode                                   | Enum       | full                                            | The scan startup mode enables you to specify the starting point for data consumption. Fluss currently supports the following `scan.startup.mode` options: `full` (default), earliest, latest, timestamp. See the [Start Reading Position](engine-flink/reads.md#start-reading-position) for more details.                                                                                                                                                                                                                                          |
-| scan.startup.timestamp                              | Long       | (None)                                          | The timestamp to start reading the data from. This option is only valid when `scan.startup.mode` is set to `timestamp`. The format is 'milli-second-since-epoch' or `yyyy-MM-dd HH:mm:ss`, like `1678883047356` or `2023-12-09 23:09:12`.                                                                                                                                                                                                                                                                                                          |
-| scan.partition.discovery.interval                   | Duration   | 1min                                            | The time interval for the Fluss source to discover the new partitions for partitioned table while scanning. A non-positive value disables the partition discovery. The default value is 1 minute. Currently, since Fluss Admin#listPartitions(TablePath tablePath) requires a large number of requests to ZooKeeper in server, this option cannot be set too small, as a small value would cause frequent requests and increase server load. In the future, once list partitions is optimized, the default value of this parameter can be reduced. |
-| client.scanner.log.check-crc                        | Boolean    | true                                            | Automatically check the CRC3 of the read records for LogScanner. This ensures no on-the-wire or on-disk corruption to the messages occurred. This check adds some overhead, so it may be disabled in cases seeking extreme performance.                                                                                                                                                                                                                                                                                                            |
-| client.scanner.log.max-poll-records                 | Integer    | 500                                             | The maximum number of records returned in a single call to poll() for LogScanner. Note that this config doesn't impact the underlying fetching behavior. The Scanner will cache the records from each fetch request and returns them incrementally from each poll.                                                                                                                                                                                                                                                                                 |
-| client.scanner.log.fetch.max-bytes                  | MemorySize | 16mb                                            | The maximum amount of data the server should return for a fetch request from client. Records are fetched in batches, and if the first record batch in the first non-empty bucket of the fetch is larger than this value, the record batch will still be returned to ensure that the fetch can make progress. As such, this is not a absolute maximum.                                                                                                                                                                                              |
-| client.scanner.log.fetch.max-bytes-for-bucket       | MemorySize | 1mb                                             | The maximum amount of data the server should return for a table bucket in fetch request fom client. Records are fetched in batches, and the max bytes size is config by this option.                                                                                                                                                                                                                                                                                                                                                               |
-| client.scanner.log.fetch.min-bytes                  | MemorySize | 1b                                              | The minimum bytes expected for each fetch log request from client to response. If not enough bytes, wait up to client.scanner.log.fetch-wait-max-time time to return.                                                                                                                                                                                                                                                                                                                                                                              |
-| client.scanner.log.fetch.wait-max-time              | Duration   | 500ms                                           | The maximum time to wait for enough bytes to be available for a fetch log request from client to response.                                                                                                                                                                                                                                                                                                                                                                                                                                         |
-| client.scanner.io.tmpdir                            | String     | System.getProperty("java.io.tmpdir") + "/fluss" | Local directory that is used by client for storing the data files (like kv snapshot, log segment files) to read temporarily                                                                                                                                                                                                                                                                                                                                                                                                                        |
-| client.scanner.remote-log.prefetch-num              | Integer    | 4                                               | The number of remote log segments to keep in local temp file for LogScanner, which download from remote storage. The default setting is 4.                                                                                                                                                                                                                                                                                                                                                                                                         |
-| client.remote-file.download-thread-num              | Integer    | 3                                               | The number of threads the client uses to download remote files.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+| Option                                        | Type       | Default                                         | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
+|-----------------------------------------------|------------|-------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| scan.startup.mode                             | Enum       | full                                            | The scan startup mode enables you to specify the starting point for data consumption. Fluss currently supports the following `scan.startup.mode` options: `full` (default), earliest, latest, timestamp. See the [Start Reading Position](engine-flink/reads.md#start-reading-position) for more details.                                                                                                                                                                                                                                          |
+| scan.startup.timestamp                        | Long       | (None)                                          | The timestamp to start reading the data from. This option is only valid when `scan.startup.mode` is set to `timestamp`. The format is 'milli-second-since-epoch' or `yyyy-MM-dd HH:mm:ss`, like `1678883047356` or `2023-12-09 23:09:12`.                                                                                                                                                                                                                                                                                                          |
+| scan.partition.discovery.interval             | Duration   | 1min                                            | The time interval for the Fluss source to discover the new partitions for partitioned table while scanning. A non-positive value disables the partition discovery. The default value is 1 minute. Currently, since Fluss Admin#listPartitions(TablePath tablePath) requires a large number of requests to ZooKeeper in server, this option cannot be set too small, as a small value would cause frequent requests and increase server load. In the future, once list partitions is optimized, the default value of this parameter can be reduced. |
+| scan.kv.snapshot.lease.id                     | String     | UUID                                            | The lease id to lease kv snapshots. If set, the acquired kv snapshots will not be deleted until the consumer finished consuming all the snapshots or the lease duration time is reached. If not set, an UUID will be set.                                                                                                                                                                                                                                                                                                                          |
+| scan.kv.snapshot.lease.duration               | Duration   | 1day                                            | The time period how long to wait before expiring the kv snapshot lease to avoid kv snapshot blocking to delete.                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+| client.scanner.log.check-crc                  | Boolean    | true                                            | Automatically check the CRC3 of the read records for LogScanner. This ensures no on-the-wire or on-disk corruption to the messages occurred. This check adds some overhead, so it may be disabled in cases seeking extreme performance.                                                                                                                                                                                                                                                                                                            |
+| client.scanner.log.max-poll-records           | Integer    | 500                                             | The maximum number of records returned in a single call to poll() for LogScanner. Note that this config doesn't impact the underlying fetching behavior. The Scanner will cache the records from each fetch request and returns them incrementally from each poll.                                                                                                                                                                                                                                                                                 |
+| client.scanner.log.fetch.max-bytes            | MemorySize | 16mb                                            | The maximum amount of data the server should return for a fetch request from client. Records are fetched in batches, and if the first record batch in the first non-empty bucket of the fetch is larger than this value, the record batch will still be returned to ensure that the fetch can make progress. As such, this is not a absolute maximum.                                                                                                                                                                                              |
+| client.scanner.log.fetch.max-bytes-for-bucket | MemorySize | 1mb                                             | The maximum amount of data the server should return for a table bucket in fetch request fom client. Records are fetched in batches, and the max bytes size is config by this option.                                                                                                                                                                                                                                                                                                                                                               |
+| client.scanner.log.fetch.min-bytes            | MemorySize | 1b                                              | The minimum bytes expected for each fetch log request from client to response. If not enough bytes, wait up to client.scanner.log.fetch-wait-max-time time to return.                                                                                                                                                                                                                                                                                                                                                                              |
+| client.scanner.log.fetch.wait-max-time        | Duration   | 500ms                                           | The maximum time to wait for enough bytes to be available for a fetch log request from client to response.                                                                                                                                                                                                                                                                                                                                                                                                                                         |
+| client.scanner.io.tmpdir                      | String     | System.getProperty("java.io.tmpdir") + "/fluss" | Local directory that is used by client for storing the data files (like kv snapshot, log segment files) to read temporarily                                                                                                                                                                                                                                                                                                                                                                                                                        |
+| client.scanner.remote-log.prefetch-num        | Integer    | 4                                               | The number of remote log segments to keep in local temp file for LogScanner, which download from remote storage. The default setting is 4.                                                                                                                                                                                                                                                                                                                                                                                                         |

Review Comment:
   This is format by idea check style.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2703100891


##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java:
##########
@@ -0,0 +1,226 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.types.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * The manager to handle {@link KvSnapshotLease} to register/update/delete metadata from zk and
+ * remote fs.
+ */
+public class KvSnapshotLeaseMetadataManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseMetadataManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final String remoteDataDir;
+
+    public KvSnapshotLeaseMetadataManager(ZooKeeperClient zkClient, String remoteDataDir) {
+        this.zkClient = zkClient;
+        this.remoteDataDir = remoteDataDir;
+    }
+
+    public List<String> getLeasesList() throws Exception {
+        return zkClient.getKvSnapshotLeasesList();
+    }
+
+    /**
+     * Register a new kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param lease the kv snapshot lease.
+     */
+    public void registerLease(String leaseId, KvSnapshotLease lease) throws Exception {
+        Map<Long, FsPath> tableIdToRemoteMetadataFsPath = generateMetadataFile(leaseId, lease);
+
+        // generate remote fsPath of metadata.
+        KvSnapshotLeaseMetadata leaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        lease.getExpirationTime(), tableIdToRemoteMetadataFsPath);
+
+        // register kv snapshot metadata to zk.
+        try {
+            zkClient.registerKvSnapshotLeaseMetadata(leaseId, leaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to register kv snapshot lease metadata to zk.", e);
+            leaseMetadata.discard();
+            throw e;
+        }
+    }
+
+    /**
+     * Update a kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param kvSnapshotLease the kv snapshot lease.
+     */
+    public void updateLease(String leaseId, KvSnapshotLease kvSnapshotLease) throws Exception {
+        // TODO change this to incremental update to avoid create too many remote metadata files.
+
+        Optional<KvSnapshotLeaseMetadata> originalLeaseMetadata =
+                zkClient.getKvSnapshotLeaseMetadata(leaseId);
+
+        Map<Long, FsPath> tableIdToNewRemoteMetadataFsPath =
+                generateMetadataFile(leaseId, kvSnapshotLease);
+
+        // generate  new kv snapshot lease metadata.
+        KvSnapshotLeaseMetadata newLeaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        kvSnapshotLease.getExpirationTime(), tableIdToNewRemoteMetadataFsPath);
+        // register new snapshot metadata to zk.
+        try {
+            zkClient.updateKvSnapshotLeaseMetadata(leaseId, newLeaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to update kv snapshot lease metadata to zk.", e);
+            newLeaseMetadata.discard();

Review Comment:
   This is a common issue also an complex issue that we can track separately in a dedicated issue.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2702397633


##########
website/docs/maintenance/configuration.md:
##########
@@ -157,7 +158,7 @@ during the Fluss cluster working.
 | kv.rocksdb.use-bloom-filter                       | Boolean    | true                          | If true, every newly created SST file will contain a Bloom filter. It is enabled by default.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
 | kv.rocksdb.bloom-filter.bits-per-key              | Double     | 10.0                          | Bits per key that bloom filter will use, this only take effect when bloom filter is used. The default value is 10.0.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
 | kv.rocksdb.bloom-filter.block-based-mode          | Boolean    | false                         | If true, RocksDB will use block-based filter instead of full filter, this only take effect when bloom filter is used. The default value is `false`.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
-| kv.rocksdb.shared-rate-limiter-bytes-per-sec              | MemorySize | Long.MAX_VALUE                | The bytes per second rate limit for RocksDB flush and compaction operations shared across all RocksDB instances on the TabletServer. The rate limiter is always enabled. The default value is Long.MAX_VALUE (effectively unlimited). Set to a lower value (e.g., 100MB) to limit the rate. This configuration can be updated dynamically without server restart. See [Updating Configs](operations/updating-configs.md) for more details.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          |
+| kv.rocksdb.shared-rate-limiter-bytes-per-sec      | MemorySize | Long.MAX_VALUE                | The bytes per second rate limit for RocksDB flush and compaction operations shared across all RocksDB instances on the TabletServer. The rate limiter is always enabled. The default value is Long.MAX_VALUE (effectively unlimited). Set to a lower value (e.g., 100MB) to limit the rate. This configuration can be updated dynamically without server restart. See [Updating Configs](operations/updating-configs.md) for more details.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |

Review Comment:
   This is format by idea check style.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "wuchong (via GitHub)" <gi...@apache.org>.
wuchong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2779161277


##########
fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java:
##########
@@ -1060,6 +1062,119 @@ void testRebalanceDuringConcurrentTableCreation() throws Exception {
         }
     }
 
+    // ------------------------------------------------------------------------
+    //  KV Snapshot Lease Authorization Tests
+    // ------------------------------------------------------------------------
+
+    @Test
+    void testAcquireKvSnapshotLease() throws Exception {
+        TableInfo tableInfo = rootAdmin.getTableInfo(DATA1_TABLE_PATH_PK).get();
+        long tableId = tableInfo.getTableId();
+        FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId);
+
+        KvSnapshotLease kvSnapshotLease =
+                guestAdmin.createKvSnapshotLease(
+                        "test-acquire-lease", Duration.ofDays(1).toMillis());
+        Map<TableBucket, Long> snapshotIds = new HashMap<>();
+        snapshotIds.put(new TableBucket(tableId, 0), 0L);
+
+        // test acquireKvSnapshotLease without READ permission on table resource
+        assertThatThrownBy(() -> kvSnapshotLease.acquireSnapshots(snapshotIds).get())
+                .rootCause()
+                .isInstanceOf(AuthorizationException.class)
+                .hasMessageContaining(
+                        String.format(
+                                "Principal %s have no authorization to operate READ on resource Resource{type=TABLE, name='%s'}",
+                                guestPrincipal, DATA1_TABLE_PATH_PK));
+
+        // add READ permission to guest user on table resource
+        List<AclBinding> aclBindings =
+                Collections.singletonList(
+                        new AclBinding(
+                                Resource.table(DATA1_TABLE_PATH_PK),
+                                new AccessControlEntry(
+                                        guestPrincipal, "*", READ, PermissionType.ALLOW)));
+        rootAdmin.createAcls(aclBindings).all().get();
+        FLUSS_CLUSTER_EXTENSION.waitUntilAuthenticationSync(aclBindings, true);
+
+        // test acquireKvSnapshotLease with READ permission should succeed
+        // (no AuthorizationException should be thrown)
+        kvSnapshotLease.acquireSnapshots(snapshotIds).get();
+
+        // cleanup: drop the lease using root admin
+        rootAdmin

Review Comment:
   use `guestAdmin`



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +540,78 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // Get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : kvSnapshots.getTableBuckets()) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!bucketsToLease.isEmpty()) {
+                String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                LOG.info(
+                        "Try to acquire kv snapshot lease {} for table {}",
+                        kvSnapshotLeaseId,
+                        PhysicalTablePath.of(tablePath, partitionName));
+                long kvSnapshotLeaseDurationMs = leaseContext.getKvSnapshotLeaseDurationMs();
+                checkNotNull(kvSnapshotLeaseDurationMs, "kv snapshot lease duration is null.");
+                Set<TableBucket> unavailableTableBucketSet =
+                        flussAdmin
+                                .createKvSnapshotLease(kvSnapshotLeaseId, kvSnapshotLeaseDurationMs)
+                                .acquireSnapshots(bucketsToLease)
+                                .get()
+                                .getUnavailableTableBucketSet();
+                if (!unavailableTableBucketSet.isEmpty()) {
+                    LOG.info(

Review Comment:
   use `LOG.warn` or `LOG.error`



##########
fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java:
##########
@@ -135,55 +154,101 @@ CompletedSnapshot addSnapshotAndSubsumeOldestOne(
 
                     // Remove completed snapshot from queue and snapshotStateHandleStore, not
                     // discard.
-                    Optional<CompletedSnapshot> subsume =
-                            subsume(
-                                    completedSnapshots,
-                                    maxNumberOfSnapshotsToRetain,
-                                    completedSnapshot -> {
-                                        remove(
-                                                completedSnapshot.getTableBucket(),
-                                                completedSnapshot.getSnapshotID());
-                                        snapshotsCleaner.addSubsumedSnapshot(completedSnapshot);
-                                    });
-
+                    subsume(
+                            completedSnapshots,
+                            maxNumberOfSnapshotsToRetain,
+                            completedSnapshot -> {
+                                remove(
+                                        completedSnapshot.getTableBucket(),
+                                        completedSnapshot.getSnapshotID());
+                                snapshotsCleaner.addSubsumedSnapshot(completedSnapshot);
+                            },
+                            subsumptionChecker);
+
+                    // Move leased snapshots that should have been subsumed but couldn't
+                    // (protected by a lease) from completedSnapshots to stillInUseSnapshots.
+                    // This ensures the effective lowestSnapshotID is computed from retained
+                    // (non-leased) snapshots only, allowing SST files from non-leased
+                    // subsumed snapshots to be cleaned up properly.
+                    CompletedSnapshot latest = completedSnapshots.peekLast();
+                    Iterator<CompletedSnapshot> leaseIt = completedSnapshots.iterator();
+                    while (leaseIt.hasNext()) {
+                        CompletedSnapshot next = leaseIt.next();
+                        if (next != latest
+                                && !subsumptionChecker.canSubsume(
+                                        new TableBucketSnapshot(
+                                                next.getTableBucket(), next.getSnapshotID()))) {
+                            leaseIt.remove();
+                            stillInUseSnapshots.put(next.getSnapshotID(), next);
+                            LOG.debug(
+                                    "Moved leased snapshot {} to stillInUseSnapshots",
+                                    next.getSnapshotID());
+                        }
+                    }
+
+                    // Check if any previously still-in-use snapshots can now be released
+                    // (lease expired).
+                    Iterator<Map.Entry<Long, CompletedSnapshot>> stillInUseIter =
+                            stillInUseSnapshots.entrySet().iterator();
+                    while (stillInUseIter.hasNext()) {
+                        Map.Entry<Long, CompletedSnapshot> entry = stillInUseIter.next();
+                        CompletedSnapshot s = entry.getValue();
+                        if (subsumptionChecker.canSubsume(
+                                new TableBucketSnapshot(s.getTableBucket(), s.getSnapshotID()))) {
+                            stillInUseIter.remove();
+                            try {
+                                remove(s.getTableBucket(), s.getSnapshotID());
+                            } catch (Exception e) {
+                                LOG.warn(
+                                        "Failed to remove released snapshot {} from store",
+                                        s.getSnapshotID(),
+                                        e);
+                            }
+                            snapshotsCleaner.addSubsumedSnapshot(s);
+                            LOG.debug(
+                                    "Released snapshot {} from stillInUseSnapshots (lease expired)",
+                                    s.getSnapshotID());
+                        }
+                    }
+
+                    // SST file cleanup: compute effective lowest from retained (non-leased)
+                    // snapshots only, and protect files referenced by still-in-use snapshots.
+                    Set<Long> stillInUseIds = new HashSet<>(stillInUseSnapshots.keySet());
                     findLowest(completedSnapshots)
                             .ifPresent(
-                                    id -> {
-                                        // unregister the unused kv file, which will then cause the
-                                        // kv file
-                                        // deletion
-                                        sharedKvFileRegistry.unregisterUnusedKvFile(id);
-                                        snapshotsCleaner.cleanSubsumedSnapshots(
-                                                id,
-                                                Collections.emptySet(),
-                                                postCleanup,
-                                                ioExecutor);
-                                    });
-                    return subsume.orElse(null);
+                                    id ->
+                                            sharedKvFileRegistry.unregisterUnusedKvFile(
+                                                    id, stillInUseIds));
+
+                    // Snapshot metadata/private files cleanup: use the latest snapshot
+                    // ID + 1 so subsumed snapshots can be cleaned even when a lower
+                    // snapshot has a lease. This is safe because
+                    // KvSnapshotHandle.discard() only deletes private files and
+                    // metadata, not shared SST files registered in SharedKvFileRegistry.
+                    snapshotsCleaner.cleanSubsumedSnapshots(
+                            snapshot.getSnapshotID() + 1, stillInUseIds, postCleanup, ioExecutor);
+                    return null;

Review Comment:
   The return value seems not used, change the signature of the method to return `void`?



##########
fluss-client/src/main/java/org/apache/fluss/client/admin/KvSnapshotLeaseImpl.java:
##########
@@ -0,0 +1,90 @@
+/*
+ * 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.client.metadata.AcquireKvSnapshotLeaseResult;
+import org.apache.fluss.client.utils.ClientRpcMessageUtils;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.rpc.gateway.AdminGateway;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAcquireKvSnapshotLeaseRequest;
+
+/** The default implementation of KvSnapshotLease. */
+public class KvSnapshotLeaseImpl implements KvSnapshotLease {
+    private final String leaseId;
+    private final long leaseDurationMs;
+    private final AdminGateway gateway;
+
+    public KvSnapshotLeaseImpl(String leaseId, long leaseDurationMs, AdminGateway gateway) {
+        this.leaseId = leaseId;
+        this.leaseDurationMs = leaseDurationMs;
+        this.gateway = gateway;
+    }
+
+    @Override
+    public String leaseId() {
+        return leaseId;
+    }
+
+    @Override
+    public long leaseDurationMs() {
+        return leaseDurationMs;
+    }
+
+    @Override
+    public CompletableFuture<AcquireKvSnapshotLeaseResult> acquireSnapshots(
+            Map<TableBucket, Long> snapshotIds) {
+        if (snapshotIds.isEmpty()) {
+            throw new IllegalArgumentException(
+                    "The snapshotIds to acquire kv snapshot lease is empty");
+        }
+
+        return gateway.acquireKvSnapshotLease(
+                        makeAcquireKvSnapshotLeaseRequest(leaseId, snapshotIds, leaseDurationMs))
+                .thenApply(ClientRpcMessageUtils::toAcquireKvSnapshotLeaseResult);
+    }
+
+    @Override
+    public CompletableFuture<Void> renew() {
+        AcquireKvSnapshotLeaseRequest request =
+                new AcquireKvSnapshotLeaseRequest()
+                        .setLeaseId(leaseId)
+                        .setLeaseDurationMs(leaseDurationMs);
+        return gateway.acquireKvSnapshotLease(request).thenApply(r -> null);
+    }
+
+    @Override
+    public CompletableFuture<Void> releaseSnapshots(Set<TableBucket> bucketsToRelease) {
+        ReleaseKvSnapshotLeaseRequest request =
+                new ReleaseKvSnapshotLeaseRequest().setLeaseId(leaseId);

Review Comment:
   The `bucketsToRelease` is not passed. 
   
   Use `makeReleaseKvSnapshotLeaseRequest(leaseId, bucketsToRelease)` instead.



##########
fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java:
##########
@@ -1060,6 +1062,119 @@ void testRebalanceDuringConcurrentTableCreation() throws Exception {
         }
     }
 
+    // ------------------------------------------------------------------------
+    //  KV Snapshot Lease Authorization Tests
+    // ------------------------------------------------------------------------
+
+    @Test
+    void testAcquireKvSnapshotLease() throws Exception {
+        TableInfo tableInfo = rootAdmin.getTableInfo(DATA1_TABLE_PATH_PK).get();
+        long tableId = tableInfo.getTableId();
+        FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId);
+
+        KvSnapshotLease kvSnapshotLease =
+                guestAdmin.createKvSnapshotLease(
+                        "test-acquire-lease", Duration.ofDays(1).toMillis());
+        Map<TableBucket, Long> snapshotIds = new HashMap<>();
+        snapshotIds.put(new TableBucket(tableId, 0), 0L);
+
+        // test acquireKvSnapshotLease without READ permission on table resource
+        assertThatThrownBy(() -> kvSnapshotLease.acquireSnapshots(snapshotIds).get())
+                .rootCause()
+                .isInstanceOf(AuthorizationException.class)
+                .hasMessageContaining(
+                        String.format(
+                                "Principal %s have no authorization to operate READ on resource Resource{type=TABLE, name='%s'}",
+                                guestPrincipal, DATA1_TABLE_PATH_PK));
+
+        // add READ permission to guest user on table resource
+        List<AclBinding> aclBindings =
+                Collections.singletonList(
+                        new AclBinding(
+                                Resource.table(DATA1_TABLE_PATH_PK),
+                                new AccessControlEntry(
+                                        guestPrincipal, "*", READ, PermissionType.ALLOW)));
+        rootAdmin.createAcls(aclBindings).all().get();
+        FLUSS_CLUSTER_EXTENSION.waitUntilAuthenticationSync(aclBindings, true);
+
+        // test acquireKvSnapshotLease with READ permission should succeed
+        // (no AuthorizationException should be thrown)
+        kvSnapshotLease.acquireSnapshots(snapshotIds).get();
+
+        // cleanup: drop the lease using root admin
+        rootAdmin
+                .createKvSnapshotLease("test-acquire-lease", Duration.ofDays(1).toMillis())
+                .dropLease()
+                .get();
+    }
+
+    @Test
+    void testReleaseKvSnapshotLease() throws Exception {
+        TableInfo tableInfo = rootAdmin.getTableInfo(DATA1_TABLE_PATH_PK).get();
+        long tableId = tableInfo.getTableId();
+        FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId);
+
+        AdminGateway guestGateway = ((FlussAdmin) guestAdmin).getAdminGateway();
+        ReleaseKvSnapshotLeaseRequest request =
+                ClientRpcMessageUtils.makeReleaseKvSnapshotLeaseRequest(
+                        "test-release-lease", Collections.singleton(new TableBucket(tableId, 0)));
+
+        // test releaseKvSnapshotLease without READ permission on table resource
+        assertThatThrownBy(() -> guestGateway.releaseKvSnapshotLease(request).get())
+                .rootCause()
+                .isInstanceOf(AuthorizationException.class)
+                .hasMessageContaining(
+                        String.format(
+                                "Principal %s have no authorization to operate READ on resource Resource{type=TABLE, name='%s'}",
+                                guestPrincipal, DATA1_TABLE_PATH_PK));
+
+        // add READ permission to guest user on table resource
+        List<AclBinding> aclBindings =
+                Collections.singletonList(
+                        new AclBinding(
+                                Resource.table(DATA1_TABLE_PATH_PK),
+                                new AccessControlEntry(
+                                        guestPrincipal, "*", READ, PermissionType.ALLOW)));
+        rootAdmin.createAcls(aclBindings).all().get();
+        FLUSS_CLUSTER_EXTENSION.waitUntilAuthenticationSync(aclBindings, true);
+
+        // test releaseKvSnapshotLease with READ permission should succeed
+        // (the lease doesn't exist, but no AuthorizationException should be thrown)
+        guestGateway.releaseKvSnapshotLease(request).get();
+    }
+
+    @Test
+    void testDropKvSnapshotLease() throws Exception {
+        KvSnapshotLease kvSnapshotLease =
+                guestAdmin.createKvSnapshotLease("test-drop-lease", Duration.ofDays(1).toMillis());
+
+        // test dropKvSnapshotLease without WRITE permission on cluster resource
+        assertThatThrownBy(() -> kvSnapshotLease.dropLease().get())
+                .rootCause()
+                .isInstanceOf(AuthorizationException.class)
+                .hasMessageContaining(
+                        String.format(
+                                "Principal %s have no authorization to operate WRITE on resource Resource{type=CLUSTER, name='fluss-cluster'}",
+                                guestPrincipal));
+
+        // add WRITE permission to guest user on cluster resource

Review Comment:
   should use `READ` permission. Note that the kvSnapshotLease is used for read kv snapshots, so most users only have read permission to this table, and we should make this common case work. 



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/FlussSourceEnumeratorStateSerializer.java:
##########
@@ -260,4 +287,17 @@ private List<SourceSplitBase> deserializeRemainingHybridLakeFlussSplits(
             return null;
         }
     }
+
+    private void serializeLeaseContext(final DataOutputSerializer out, SourceEnumeratorState state)

Review Comment:
   rename method to `serializeLeaseId()`.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -923,15 +1000,72 @@ public void addReader(int subtaskId) {
     public SourceEnumeratorState snapshotState(long checkpointId) {
         final SourceEnumeratorState enumeratorState =
                 new SourceEnumeratorState(
-                        assignedTableBuckets, assignedPartitions, pendingHybridLakeFlussSplits);
+                        assignedTableBuckets,
+                        assignedPartitions,
+                        pendingHybridLakeFlussSplits,
+                        leaseContext.getKvSnapshotLeaseId());
         LOG.debug("Source Checkpoint is {}", enumeratorState);
         return enumeratorState;
     }
 
+    @Override
+    public void notifyCheckpointComplete(long checkpointId) throws Exception {
+        // lower than this checkpoint id.
+        Set<TableBucket> consumedKvSnapshots = getAndRemoveConsumedBucketsUpTo(checkpointId);
+
+        LOG.info(
+                "kv snapshot has already consumed and try to release kv snapshot lease for: {}, checkpoint id: {}",
+                consumedKvSnapshots,
+                checkpointId);
+
+        // send request to fluss to unregister the kv snapshot lease.
+        try {
+            flussAdmin
+                    .createKvSnapshotLease(
+                            leaseContext.getKvSnapshotLeaseId(),
+                            leaseContext.getKvSnapshotLeaseDurationMs())
+                    .releaseSnapshots(consumedKvSnapshots)
+                    .get();
+        } catch (Exception e) {
+            LOG.error("Failed to release kv snapshot lease. These snapshot need to re-enqueue", e);
+            // use the current checkpoint id to re-enqueue the buckets
+            consumedKvSnapshots.forEach(
+                    tableBucket -> addConsumedBucket(checkpointId, tableBucket));
+        }
+    }
+
+    /** Add bucket who has been consumed kv snapshot to the consumedKvSnapshotMap. */
+    public void addConsumedBucket(long checkpointId, TableBucket tableBucket) {
+        consumedKvSnapshotMap.computeIfAbsent(checkpointId, k -> new HashSet<>()).add(tableBucket);
+    }
+
+    /** Get and remove the buckets who have been consumed kv snapshot up to the checkpoint id. */
+    public Set<TableBucket> getAndRemoveConsumedBucketsUpTo(long checkpointId) {
+        NavigableMap<Long, Set<TableBucket>> toRemove =
+                consumedKvSnapshotMap.headMap(checkpointId, false);
+        Set<TableBucket> result = new HashSet<>();
+        for (Set<TableBucket> snapshots : toRemove.values()) {
+            result.addAll(snapshots);
+        }
+        toRemove.clear();
+        return result;
+    }
+
     @Override
     public void close() throws IOException {
         try {
             closed = true;
+
+            if (!isStreaming) {

Review Comment:
   Change the condition to check it's pk table and read from snapshot to avoid unnecessary IO. 
   
   ```java
   if (!isStreaming
                       && hasPrimaryKey
                       && startingOffsetsInitializer instanceof SnapshotOffsetsInitializer)
   ```



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java:
##########
@@ -856,6 +872,95 @@ public CompletableFuture<ControlledShutdownResponse> controlledShutdown(
         return response;
     }
 
+    @Override
+    public CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseRequest request) {
+        // Authorization: require WRITE permission on all tables in the request
+        if (authorizer != null) {
+            for (PbKvSnapshotLeaseForTable kvSnapshotLeaseForTable :
+                    request.getSnapshotsToLeasesList()) {
+                long tableId = kvSnapshotLeaseForTable.getTableId();
+                authorizeTable(OperationType.READ, tableId);
+            }
+        }
+
+        String leaseId = request.getLeaseId();
+        return CompletableFuture.supplyAsync(
+                () -> {
+                    try {
+                        return makeAcquireKvSnapshotLeaseResponse(
+                                kvSnapshotLeaseManager.acquireLease(
+                                        leaseId,
+                                        request.getLeaseDurationMs(),
+                                        getAcquireKvSnapshotLeaseData(request)));
+                    } catch (ApiException e) {
+                        // Re-throw ApiExceptions as-is to preserve exception type for client
+                        throw e;
+                    } catch (Exception e) {
+                        throw new UnknownServerException(
+                                "Failed to acquire kv snapshot lease for" + leaseId, e);
+                    }
+                },
+                ioExecutor);
+    }
+
+    @Override
+    public CompletableFuture<ReleaseKvSnapshotLeaseResponse> releaseKvSnapshotLease(
+            ReleaseKvSnapshotLeaseRequest request) {
+        // Authorization: require WRITE permission on all tables in the request
+        if (authorizer != null) {
+            for (PbTableBucket tableBucket : request.getBucketsToReleasesList()) {
+                long tableId = tableBucket.getTableId();
+                authorizeTable(OperationType.READ, tableId);
+            }
+        }
+
+        String leaseId = request.getLeaseId();
+        return CompletableFuture.supplyAsync(
+                () -> {
+                    try {
+                        ReleaseKvSnapshotLeaseResponse response =
+                                new ReleaseKvSnapshotLeaseResponse();
+                        kvSnapshotLeaseManager.release(
+                                leaseId, getReleaseKvSnapshotLeaseData(request));
+                        return response;
+                    } catch (ApiException e) {
+                        // Re-throw ApiExceptions as-is to preserve exception type for client
+                        throw e;
+                    } catch (Exception e) {
+                        throw new UnknownServerException(
+                                "Failed to release kv snapshot lease for" + leaseId, e);
+                    }
+                },
+                ioExecutor);
+    }
+
+    @Override
+    public CompletableFuture<DropKvSnapshotLeaseResponse> dropKvSnapshotLease(
+            DropKvSnapshotLeaseRequest request) {
+        // Authorization: require WRITE permission on the cluster
+        if (authorizer != null) {
+            authorizer.authorize(currentSession(), OperationType.WRITE, Resource.cluster());

Review Comment:
   We should use **`READ` permission on the table** instead. Requiring **admin-level cluster permissions** just to read a primary key table is unreasonable and violates the principle of least privilege. Table-level `READ` access should be sufficient for such operations.



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -378,6 +378,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration_ms = 2;
+  repeated PbKvSnapshotLeaseForTable snapshots_to_lease = 3;
+}
+
+message AcquireKvSnapshotLeaseResponse {
+  repeated PbKvSnapshotLeaseForTable unavailable_snapshots = 3;

Review Comment:
   ```suggestion
     repeated PbKvSnapshotLeaseForTable unavailable_snapshots = 1;
   ```



##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java:
##########
@@ -113,7 +114,8 @@ void testPkTableNoSnapshotSplits() throws Throwable {
                             DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
                             streaming,
                             null,
-                            null);
+                            null,
+                            new LeaseContext("kv_snapshot_lease", Duration.ofDays(1).toMillis()));

Review Comment:
   Update all the `LeaseContext` parameter in this class to use `LeaseContext.DEFAULT`



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,427 @@
+/*
+ * 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.server.coordinator.lease;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableBucketSnapshot;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final long leaseExpirationCheckInterval;
+
+    private final ReadWriteLock managerLock = new ReentrantReadWriteLock();
+
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("managerLock")
+    private final ConcurrentHashMap<String, KvSnapshotLeaseHandler> kvSnapshotLeaseMap =
+            MapUtils.newConcurrentHashMap();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("managerLock")
+    private final Map<TableBucketSnapshot, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            long leaseExpirationCheckInterval,
+            ZooKeeperClient zkClient,
+            String remoteDataDir,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                leaseExpirationCheckInterval,
+                zkClient,
+                remoteDataDir,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            long leaseExpirationCheckInterval,
+            ZooKeeperClient zkClient,
+            String remoteDataDir,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = new KvSnapshotLeaseMetadataManager(zkClient, remoteDataDir);
+        this.leaseExpirationCheckInterval = leaseExpirationCheckInterval;
+        this.scheduledExecutor = scheduledExecutor;
+        this.clock = clock;
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+
+        List<String> leasesList = new ArrayList<>();
+        try {
+            leasesList = metadataManager.getLeasesList();
+        } catch (Exception e) {
+            LOG.error("Failed to get leases list from zookeeper.", e);
+        }
+
+        for (String leaseId : leasesList) {
+            Optional<KvSnapshotLeaseHandler> kvSnapshotLeaseOpt = Optional.empty();
+            try {
+                kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            } catch (Exception e) {
+                LOG.error("Failed to get kv snapshot lease from zookeeper.", e);
+            }
+
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLeaseHandler kvSnapshotLeasehandle = kvSnapshotLeaseOpt.get();
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLeasehandle);
+
+                initializeRefCount(kvSnapshotLeasehandle);
+
+                leasedBucketCount.addAndGet(kvSnapshotLeasehandle.getLeasedSnapshotCount());
+            }
+        }
+
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases, 0L, leaseExpirationCheckInterval, TimeUnit.MILLISECONDS);
+    }
+
+    public boolean snapshotLeaseNotExist(TableBucketSnapshot tableBucketSnapshot) {
+        return inReadLock(
+                managerLock,
+                () -> {
+                    AtomicInteger count = refCount.get(tableBucketSnapshot);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(

Review Comment:
   It seems this comment https://github.com/apache/fluss/pull/2179#discussion_r2777746406 is not addressed. I post it here again. I think we can create an issue to track this and fix in the future. 
   
   ------
   
   The current use of `leaseLocks` feels like a hack. If thread safety is required for `KvSnapshotLease`, then the class itself should be made inherently thread-safe, which would also significantly simplify the surrounding code.
   
   To achieve atomicity and thread safety, we should refactor the existing `KvSnapshotLease#acquireBucket` method into a single, cohesive operation:
   
   ```java
   KvSnapshotLease#acquireSnapshots(
       long leaseDurationMs,
       Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket,
       ConcurrentHashMap<KvSnapshotLeaseForBucket, AtomicInteger> refCount
   )
   ```
   
   This method should:
   - Perform all necessary validations and state updates **atomically**.
   - Update reference counts internally as part of the acquisition logic.
   
   Once this method completes successfully, the ZooKeeper metadata update should be performed **outside the lock**, but with **version-based consistency** to ensure correctness in concurrent scenarios. Specifically:
   - Use conditional ZK writes based on the current version of the lease node.
   - Follow the pattern established in `ProducerOffsetsStore#xxxVersion()` methods, which use versioned ZK operations to avoid race conditions during metadata updates.
   
   This approach eliminates the need for external locking, improves code clarity, and ensures both memory and ZK state remain consistent under concurrency.



##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceBatchITCase.java:
##########
@@ -82,6 +85,55 @@ void after() {
         tEnv.executeSql(String.format("drop database %s cascade", DEFAULT_DB));
     }
 
+    @Test
+    void testPkTableBatchReadWithKvSnapshotLease() throws Exception {

Review Comment:
   Actually, this doesn't read kv snapshot, because the limit is pushed down. I checked the code again, and realized that we dont' allow read kv snapshot when batch mode. So I think we can remove this test (the `close()` logic in source enumerator as well). 



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java:
##########
@@ -196,6 +208,18 @@ protected void startServices() throws Exception {
                     Executors.newFixedThreadPool(
                             conf.get(ConfigOptions.SERVER_IO_POOL_SIZE),
                             new ExecutorThreadFactory("coordinator-io"));
+
+            // Initialize and start the kv snapshot lease manager
+            this.kvSnapshotLeaseManager =

Review Comment:
   close `kvSnapshotLeaseManager` (the thread pool in it) when shutdown the server.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +540,78 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // Get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : kvSnapshots.getTableBuckets()) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!bucketsToLease.isEmpty()) {
+                String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                LOG.info(
+                        "Try to acquire kv snapshot lease {} for table {}",
+                        kvSnapshotLeaseId,
+                        PhysicalTablePath.of(tablePath, partitionName));
+                long kvSnapshotLeaseDurationMs = leaseContext.getKvSnapshotLeaseDurationMs();
+                checkNotNull(kvSnapshotLeaseDurationMs, "kv snapshot lease duration is null.");

Review Comment:
   not need to check not null for a primitive type.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -243,6 +259,8 @@ public FlinkSourceEnumerator(
                 streaming ? new NoStoppingOffsetsInitializer() : OffsetsInitializer.latest();
         this.lakeSource = lakeSource;
         this.workerExecutor = workerExecutor;
+        this.leaseContext = leaseContext;
+        this.isStreaming = streaming;

Review Comment:
   We already have `streaming` as the member variable, no need to introduce a new one. 



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseHandler.java:
##########
@@ -0,0 +1,233 @@
+/*
+ * 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.server.coordinator.lease;
+
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
+
+/** handler of kv snapshot lease. */
+@NotThreadSafe
+public class KvSnapshotLeaseHandler {
+    private long expirationTime;
+
+    /** A map from table id to kv snapshot lease for one table. */
+    private final Map<Long, KvSnapshotTableLease> tableIdToTableLease;
+
+    public KvSnapshotLeaseHandler(long expirationTime) {
+        this(expirationTime, MapUtils.newConcurrentHashMap());
+    }
+
+    public KvSnapshotLeaseHandler(
+            long expirationTime, Map<Long, KvSnapshotTableLease> tableIdToTableLease) {
+        this.expirationTime = expirationTime;
+        this.tableIdToTableLease = tableIdToTableLease;
+    }
+
+    public void setExpirationTime(long expirationTime) {
+        this.expirationTime = expirationTime;
+    }
+
+    public long getExpirationTime() {
+        return expirationTime;
+    }
+
+    public Map<Long, KvSnapshotTableLease> getTableIdToTableLease() {
+        return tableIdToTableLease;
+    }
+
+    /**
+     * Acquire a bucket to the lease id. If the bucket array already exists but is too small to
+     * accommodate the given bucket id, the array will be dynamically expanded to {@code bucketId +
+     * 1}.
+     *
+     * @param tableBucket table bucket
+     * @param snapshotId snapshot id
+     * @return the original registered snapshotId. if -1 means the bucket is new registered
+     */
+    public long acquireBucket(TableBucket tableBucket, long snapshotId) {
+        Long[] bucketSnapshot;
+        Long partitionId = tableBucket.getPartitionId();
+        long tableId = tableBucket.getTableId();
+        int bucketId = tableBucket.getBucket();
+        if (partitionId == null) {
+            // For none-partitioned table.
+            KvSnapshotTableLease tableLease =
+                    tableIdToTableLease.computeIfAbsent(
+                            tableId,
+                            k -> {
+                                Long[] array = new Long[bucketId + 1];
+                                Arrays.fill(array, -1L);
+                                return new KvSnapshotTableLease(tableId, array);
+                            });
+            bucketSnapshot = tableLease.getBucketSnapshots();
+            // Dynamically expand the array if the bucket id exceeds the current array size.
+            // This can happen when new buckets are added to an existing table.
+            if (bucketSnapshot != null && bucketId >= bucketSnapshot.length) {
+                bucketSnapshot = expandArray(bucketSnapshot, bucketId + 1);

Review Comment:
   Currently, we iterate over requested buckets sequentially. If the bucket IDs range from 0 to 10,000, this causes the internal `int[]` array to be resized (expanded) up to 10,000 times—an extremely inefficient process. Worse, this happens while holding locks, which amplifies performance degradation and increases contention.
   
   My earlier suggestion was to refactor `KvSnapshotLeaseHandler` to use a **batched API** that accepts the full `Map<Long, List<TableBucketSnapshot>> tableIdToLeaseBucket` upfront, enabling more efficient memory allocation and processing.
   
   If you plan to implement that refactoring later, **at a minimum**, this PR should include an initialization step that pre-sizes the snapshot `int[]` array based on the **maximum bucket ID** in the request. This avoids repeated array reallocations and significantly reduces lock-holding time during lease acquisition.



##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java:
##########
@@ -352,6 +354,46 @@ void testPkTableReadMixSnapshotAndLog() throws Exception {
         assertResultsIgnoreOrder(rowIter, expectedRows, true);
     }
 
+    @Test
+    void testPkTableReadWithKvSnapshotLease() throws Exception {
+        tEnv.executeSql(
+                "create table pk_table_with_kv_snapshot_lease (a int not null primary key not enforced, b varchar)");
+        TablePath tablePath = TablePath.of(DEFAULT_DB, "pk_table_with_kv_snapshot_lease");
+
+        List<InternalRow> rows = Arrays.asList(row(1, "v1"), row(2, "v2"), row(3, "v3"));
+
+        // write records
+        writeRows(conn, tablePath, rows, false);
+
+        FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshot(tablePath);
+
+        // enable checkpoint to make sure the kv snapshot lease will be cleared.
+        execEnv.enableCheckpointing(100);

Review Comment:
   I think we should add another test case covering scenarios where **checkpointing is disabled** (e.g., debug mode or jobs without checkpointing enabled). In such cases, leases should still be released promptly when the job finishes or is canceled.
   
   One possible approach is to drop the lease in `FlinkSourceEnumerator#close()` under the following conditions:
   - A checkpoint has already been successfully completed (i.e., `snapshotState(cp)` was called), **and**
   - The enumerator was **not restored** from a prior state (i.e., it wasn’t created via `FlinkSource#restoreEnumerator()`).
   
   This would mean: if **no checkpoint ever completed** and the job terminates (either normally or via cancellation), there’s no need to retain the lease—so we can safely release it immediately in `close()`.
   
   That said, this is an optimization that can be addressed later. I’m fine with creating a follow-up issue to track this improvement.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java:
##########
@@ -856,6 +872,95 @@ public CompletableFuture<ControlledShutdownResponse> controlledShutdown(
         return response;
     }
 
+    @Override
+    public CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseRequest request) {
+        // Authorization: require WRITE permission on all tables in the request

Review Comment:
   ```suggestion
           // Authorization: require READ permission on all tables in the request
   ```



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/FlussSourceEnumeratorStateSerializer.java:
##########
@@ -260,4 +287,17 @@ private List<SourceSplitBase> deserializeRemainingHybridLakeFlussSplits(
             return null;
         }
     }
+
+    private void serializeLeaseContext(final DataOutputSerializer out, SourceEnumeratorState state)
+            throws IOException {
+        String leaseId = state.getLeaseId();
+        out.writeUTF(leaseId);
+    }
+
+    private LeaseContext deserializeLeaseContext(final DataInputDeserializer in)

Review Comment:
   rename method to `String deserializeLeaseId()` as there is only lease id is deserialized. 



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java:
##########
@@ -856,6 +872,95 @@ public CompletableFuture<ControlledShutdownResponse> controlledShutdown(
         return response;
     }
 
+    @Override
+    public CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseRequest request) {
+        // Authorization: require WRITE permission on all tables in the request
+        if (authorizer != null) {
+            for (PbKvSnapshotLeaseForTable kvSnapshotLeaseForTable :
+                    request.getSnapshotsToLeasesList()) {
+                long tableId = kvSnapshotLeaseForTable.getTableId();
+                authorizeTable(OperationType.READ, tableId);
+            }
+        }
+
+        String leaseId = request.getLeaseId();
+        return CompletableFuture.supplyAsync(
+                () -> {
+                    try {
+                        return makeAcquireKvSnapshotLeaseResponse(
+                                kvSnapshotLeaseManager.acquireLease(
+                                        leaseId,
+                                        request.getLeaseDurationMs(),
+                                        getAcquireKvSnapshotLeaseData(request)));
+                    } catch (ApiException e) {
+                        // Re-throw ApiExceptions as-is to preserve exception type for client
+                        throw e;
+                    } catch (Exception e) {
+                        throw new UnknownServerException(
+                                "Failed to acquire kv snapshot lease for" + leaseId, e);
+                    }
+                },
+                ioExecutor);
+    }
+
+    @Override
+    public CompletableFuture<ReleaseKvSnapshotLeaseResponse> releaseKvSnapshotLease(
+            ReleaseKvSnapshotLeaseRequest request) {
+        // Authorization: require WRITE permission on all tables in the request

Review Comment:
   ```suggestion
           // Authorization: require READ permission on all tables in the request
   ```



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),

Review Comment:
   Could you create an issue to improve this in the future?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2702399391


##########
fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java:
##########
@@ -200,14 +208,21 @@ interface SubsumeAction {
         void subsume(CompletedSnapshot snapshot) throws Exception;
     }
 
-    private static boolean canSubsume(CompletedSnapshot next, CompletedSnapshot latest) {
+    /** A function to determine whether a snapshot can be subsumed. */
+    @FunctionalInterface
+    public interface CanSubsume {

Review Comment:
   Rename to `SubsumptionChecker `.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "wuchong (via GitHub)" <gi...@apache.org>.
wuchong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2767987165


##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -1129,4 +1150,15 @@ message PbBucketOffset {
 message PbProducerTableOffsets {
   required int64 table_id = 1;
   repeated PbBucketOffset bucket_offsets = 2;
+}
+
+message PbKvSnapshotLeaseForTable {
+  required int64 table_id = 1;
+  repeated PbKvSnapshotLeaseForBucket buckets_req = 2;

Review Comment:
   `PbKvSnapshotLeaseForTable` is also used as response, `buckets_req` is not good for this,  rename to `bucket_snapshots`?



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java:
##########
@@ -2035,6 +2073,53 @@ private ControlledShutdownResponse tryProcessControlledShutdown(
         return response;
     }
 
+    private AcquireKvSnapshotLeaseResponse tryProcessAcquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseEvent event) throws Exception {
+        AcquireKvSnapshotLeaseResponse response = new AcquireKvSnapshotLeaseResponse();
+        Map<TableBucket, Long> unavailableSnapshots =
+                kvSnapshotLeaseManager.acquireLease(
+                        event.getLeaseId(),
+                        event.getLeaseDuration(),
+                        event.getTableIdToLeasedBucket());
+
+        Map<Long, List<PbKvSnapshotLeaseForBucket>> pbFailedTables = new HashMap<>();
+        for (Map.Entry<TableBucket, Long> entry : unavailableSnapshots.entrySet()) {
+            TableBucket tb = entry.getKey();
+            Long snapshotId = entry.getValue();
+            PbKvSnapshotLeaseForBucket pbBucket =
+                    new PbKvSnapshotLeaseForBucket().setBucketId(tb.getBucket());
+            if (tb.getPartitionId() != null) {
+                pbBucket.setPartitionId(tb.getPartitionId());
+            }
+            pbBucket.setSnapshotId(snapshotId);
+            pbFailedTables.computeIfAbsent(tb.getTableId(), k -> new ArrayList<>()).add(pbBucket);
+        }
+
+        for (Map.Entry<Long, List<PbKvSnapshotLeaseForBucket>> entry : pbFailedTables.entrySet()) {
+            response.addTablesLeaseRe()
+                    .setTableId(entry.getKey())
+                    .addAllBucketsReqs(entry.getValue());
+        }
+        return response;

Review Comment:
   Move this to `ServerRpcMessageUtils`



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -375,6 +372,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;
+  repeated PbKvSnapshotLeaseForTable table_lease_req = 3;

Review Comment:
   rename to `all_snapshots_to_lease`?



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataManager.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger count = refCount.get(kvSnapshotLeaseForBucket);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    // set the expiration time as: current time + leaseDuration
+                    long newExpirationTime = clock.milliseconds() + leaseDuration;
+                    KvSnapshotLease kvSnapshotLease =
+                            kvSnapshotLeaseMap.compute(
+                                    leaseId,
+                                    (key, existingLease) -> {
+                                        if (existingLease == null) {
+                                            LOG.info(
+                                                    "kv snapshot lease '{}' has been acquired. The lease expiration "
+                                                            + "time is {}",
+                                                    leaseId,
+                                                    newExpirationTime);
+                                            return new KvSnapshotLease(newExpirationTime);
+                                        } else {
+                                            existingLease.setExpirationTime(newExpirationTime);
+                                            return existingLease;
+                                        }
+                                    });
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataManager.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(
+                                                tableBucket, originalSnapshotId));
+                            }
+                            incrementRefCount(bucket);
+                        }
+                    }
+
+                    if (update) {
+                        metadataManager.updateLease(leaseId, kvSnapshotLease);
+                    } else {
+                        metadataManager.registerLease(leaseId, kvSnapshotLease);
+                    }
+
+                    return unavailableSnapshots;
+                });
+    }
+
+    public void release(String leaseId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotLease lease = kvSnapshotLeaseMap.get(leaseId);
+                    if (lease == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = lease.releaseBucket(bucket);
+                            if (snapshotId != -1L) {
+                                leasedBucketCount.decrementAndGet();
+                                decrementRefCount(new KvSnapshotLeaseForBucket(bucket, snapshotId));
+                            }
+                        }
+                    }
+
+                    if (lease.isEmpty()) {
+                        releaseAll(leaseId);
+                    } else {
+                        metadataManager.updateLease(leaseId, lease);
+                    }
+                });
+    }
+
+    /**
+     * Release kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @return true if clear success, false if lease not exist
+     */
+    public boolean releaseAll(String leaseId) throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseMap.remove(leaseId);
+                            if (kvSnapshotLease == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotLease);
+                            metadataManager.deleteLease(leaseId);
+
+                            LOG.info(
+                                    "kv snapshots of lease '"
+                                            + leaseId
+                                            + "' has been all released.");
+                            return true;
+                        });
+
+        leaseLocks.remove(leaseId);
+        return exist;
+    }
+
+    private void initializeRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+
+                    incrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        incrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                    }
+                }
+            }
+        }
+    }
+
+    private void clearRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+                    decrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                    leasedBucketCount.decrementAndGet();
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        decrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                        leasedBucketCount.decrementAndGet();
+                    }
+                }
+            }
+        }
+    }
+
+    private void incrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () ->
+                        refCount.computeIfAbsent(
+                                        kvSnapshotLeaseForBucket, k -> new AtomicInteger(0))
+                                .incrementAndGet());
+    }
+
+    private void decrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger atomicInteger = refCount.get(kvSnapshotLeaseForBucket);
+                    if (atomicInteger != null) {
+                        int decrementAndGet = atomicInteger.decrementAndGet();
+                        if (decrementAndGet <= 0) {
+                            refCount.remove(kvSnapshotLeaseForBucket);
+                        }
+                    }
+                });

Review Comment:
   We can remove the writeLock if we refactor this into using the atomic API `computeIfPresent`.
   
   ```java
   refCount.computeIfPresent(
                               kvSnapshotLeaseForBucket,
                               (k, v) -> {
                                   int newCount = v.decrementAndGet();
                                   return newCount <= 0 ? null : v;
                               });
   ```



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +537,88 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+            Set<TableBucket> remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Set<TableBucket> ignoreBuckets = new HashSet<>();
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : remainingTableBuckets) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                } else {
+                    ignoreBuckets.add(tb);
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!ignoreBuckets.isEmpty()) {
+                remainingTableBuckets.removeAll(ignoreBuckets);

Review Comment:
   I don't see the usage of `remainingTableBuckets`, it is overriden by the following return value of `flussAdmin.acquireKvSnapshotLease`. 



##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java:
##########
@@ -352,6 +354,51 @@ void testPkTableReadMixSnapshotAndLog() throws Exception {
         assertResultsIgnoreOrder(rowIter, expectedRows, true);
     }
 
+    @Test
+    void testPkTableReadWithKvSnapshotLease() throws Exception {
+        tEnv.executeSql(
+                "create table pk_table_with_kv_snapshot_lease (a int not null primary key not enforced, b varchar)");
+        TablePath tablePath = TablePath.of(DEFAULT_DB, "pk_table_with_kv_snapshot_lease");
+
+        List<InternalRow> rows = Arrays.asList(row(1, "v1"), row(2, "v2"), row(3, "v3"));
+
+        // write records
+        writeRows(conn, tablePath, rows, false);
+
+        FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshot(tablePath);
+
+        // enable checkpoint to make sure the kv snapshot lease will be cleared.
+        execEnv.enableCheckpointing(100);
+
+        List<String> expectedRows = Arrays.asList("+I[1, v1]", "+I[2, v2]", "+I[3, v3]");
+        org.apache.flink.util.CloseableIterator<Row> rowIter =
+                tEnv.executeSql(
+                                "select * from pk_table_with_kv_snapshot_lease "
+                                        + "/*+ OPTIONS('scan.kv.snapshot.lease.id' = 'test-lease-10001') */")

Review Comment:
   Can we add a test without the dynamic option or modify this? I think that is the common way that users will use and we should cover that in tests. We can assert lease is dropped by using `zkClient.getKvSnapshotLeasesList()` method.



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -378,6 +378,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;

Review Comment:
   `lease_duration` -> `lease_duration_ms`



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,453 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataHelper,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataHelper metadataHelper,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataHelper = metadataHelper;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> leaseIds = metadataHelper.getLeasesList();
+        for (String leaseId : leaseIds) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(kvSnapshotLeaseForBucket)
+                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());

Review Comment:
   I agree with @platinumhamburg—the current use of `leaseLocks` feels like a hack. If thread safety is required for `KvSnapshotLease`, then the class itself should be made inherently thread-safe, which would also significantly simplify the surrounding code.
   
   To achieve atomicity and thread safety, we should refactor the existing `KvSnapshotLease#acquireBucket` method into a single, cohesive operation:
   
   ```java
   KvSnapshotLease#acquireSnapshots(
       long leaseDurationMs,
       Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket,
       ConcurrentHashMap<KvSnapshotLeaseForBucket, AtomicInteger> refCount
   )
   ```
   
   This method should:
   - Perform all necessary validations and state updates **atomically**.
   - Update reference counts internally as part of the acquisition logic.
   
   Once this method completes successfully, the ZooKeeper metadata update should be performed **outside the lock**, but with **version-based consistency** to ensure correctness in concurrent scenarios. Specifically:
   - Use conditional ZK writes based on the current version of the lease node.
   - Follow the pattern established in `ProducerOffsetsStore#xxxVersion()` methods, which use versioned ZK operations to avoid race conditions during metadata updates.
   
   This approach eliminates the need for external locking, improves code clarity, and ensures both memory and ZK state remain consistent under concurrency.



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -375,6 +372,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;
+  repeated PbKvSnapshotLeaseForTable table_lease_req = 3;
+}
+
+message AcquireKvSnapshotLeaseResponse {
+  repeated PbKvSnapshotLeaseForTable tables_lease_res = 3;

Review Comment:
   rename to `unavailable_snapshots_to_lease`?



##########
fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java:
##########
@@ -428,6 +430,57 @@ CompletableFuture<Void> dropPartition(
     CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
+    /**
+     * Acquires a lease for specific KV snapshots of the given tableBuckets asynchronously.
+     *
+     * <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
+     * #releaseKvSnapshotLease} to release the lock early when reading is finished.
+     *
+     * <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.
+     * </ul>
+     *
+     * @param leaseId The unique ID for this lease session (usually a UUID generated by client).
+     * @param snapshotIds The snapshots to lease, 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<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration);
+
+    /**
+     * Releases the lease for specific tableBuckets asynchronously.
+     *
+     * <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.
+     *
+     * <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);
+
+    /**
+     * Releases the entire lease asynchronously for all leased snapshots for all table buckets of
+     * specified leaseId.
+     *
+     * <p>All snapshots locked under this {@code leaseId} will be released immediately. This is
+     * equivalent to calling {@link #releaseKvSnapshotLease} with all held buckets.
+     *
+     * @param leaseId The lease id to release.
+     */
+    CompletableFuture<Void> releaseAllKvSnapshotLease(String leaseId);

Review Comment:
   I have concerns about the current naming of `releaseKvSnapshotLease` and `releaseAllKvSnapshotLease`. The names are misleading:
   
   - `releaseKvSnapshotLease` sounds like it releases a lease, but it actually releases **snapshots** associated with that lease.
   - `releaseAllKvSnapshotLease` implies releasing *all* leases, yet it only releases snapshots for **a single lease**.
   
   This confusion stems from a deeper design issue: all snapshot lease–related methods are crammed into the already large `Admin` interface. To disambiguate method names, we’re forced to prefix everything with `KvSnapshotLease`, which leads to awkward and inconsistent semantics.
   
   To improve clarity and maintainability, I propose extracting all snapshot lease functionality into a dedicated, well-defined interface:
   
   ```java
   interface Admin {
       /**
        * Creates a new KV snapshot lease with the given ID and duration.
        *
        * @param leaseId the unique identifier for the lease
        * @param leaseDurationMs the lease duration in milliseconds
        * @return a {@link KvSnapshotLease} instance representing the created lease
        */
       KvSnapshotLease createKvSnapshotLease(String leaseId, long leaseDurationMs);
   }
   
   /**
    * Represents a lease for managing KV snapshots. A lease allows acquiring,
    * renewing, releasing, and dropping snapshot references under a time-bound agreement.
    */
   @PublicEvolving
   public interface KvSnapshotLease {
   
       /** Returns the unique identifier of this lease. */
       String leaseId();
   
       /** Returns the lease duration in milliseconds. */
       long leaseDurationMs();
   
       /**
        * Acquires snapshots for the specified table buckets under this lease.
        *
        * @param snapshotIds mapping from table buckets to desired snapshot IDs
        * @return a future containing the result of the acquisition
        */
       CompletableFuture<AcquireKvSnapshotLeaseResult> acquireSnapshots(
               Map<TableBucket, Long> snapshotIds);
   
       /**
        * Renews the lease, extending its validity period.
        *
        * @return a future that completes when the renewal is acknowledged
        */
       CompletableFuture<Void> renew();
   
       /**
        * Releases snapshots for the specified table buckets, freeing associated resources.
        *
        * @param bucketsToRelease the set of table buckets whose snapshots should be released
        * @return a future that completes when the release is processed
        */
       CompletableFuture<Void> releaseSnapshots(Set<TableBucket> bucketsToRelease);
   
       /**
        * Drops the entire lease, releasing all held snapshots and invalidating the lease.
        *
        * @return a future that completes when the lease is fully dropped
        */
       CompletableFuture<Void> dropLease();
   }
   ```
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java:
##########
@@ -2035,6 +2073,53 @@ private ControlledShutdownResponse tryProcessControlledShutdown(
         return response;
     }
 
+    private AcquireKvSnapshotLeaseResponse tryProcessAcquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseEvent event) throws Exception {
+        AcquireKvSnapshotLeaseResponse response = new AcquireKvSnapshotLeaseResponse();
+        Map<TableBucket, Long> unavailableSnapshots =
+                kvSnapshotLeaseManager.acquireLease(
+                        event.getLeaseId(),
+                        event.getLeaseDuration(),
+                        event.getTableIdToLeasedBucket());
+
+        Map<Long, List<PbKvSnapshotLeaseForBucket>> pbFailedTables = new HashMap<>();
+        for (Map.Entry<TableBucket, Long> entry : unavailableSnapshots.entrySet()) {
+            TableBucket tb = entry.getKey();
+            Long snapshotId = entry.getValue();
+            PbKvSnapshotLeaseForBucket pbBucket =
+                    new PbKvSnapshotLeaseForBucket().setBucketId(tb.getBucket());
+            if (tb.getPartitionId() != null) {
+                pbBucket.setPartitionId(tb.getPartitionId());
+            }
+            pbBucket.setSnapshotId(snapshotId);
+            pbFailedTables.computeIfAbsent(tb.getTableId(), k -> new ArrayList<>()).add(pbBucket);
+        }
+
+        for (Map.Entry<Long, List<PbKvSnapshotLeaseForBucket>> entry : pbFailedTables.entrySet()) {
+            response.addTablesLeaseRe()
+                    .setTableId(entry.getKey())
+                    .addAllBucketsReqs(entry.getValue());
+        }
+        return response;
+    }
+
+    private ReleaseKvSnapshotLeaseResponse tryProcessReleaseKvSnapshotLease(
+            ReleaseKvSnapshotLeaseEvent event) throws Exception {
+        ReleaseKvSnapshotLeaseResponse response = new ReleaseKvSnapshotLeaseResponse();
+        Map<Long, List<TableBucket>> tableIdToReleasedBucket = event.getTableIdToReleasedBucket();
+        if (tableIdToReleasedBucket.isEmpty()) {
+            // release all
+            boolean exist = kvSnapshotLeaseManager.releaseAll(event.getLeaseId());
+            if (!exist) {
+                throw new KvSnapshotLeaseNotExistException(

Review Comment:
   Should this throw an error?  Dropping a lease should be idempotent and reentrant; otherwise, client-side handling becomes very cumbersome.  
   
   Additionally, if we're going to throw a `LeaseNotExistException` here, why doesn't the `release(leaseId, bucket)` method below throw the same exception?   Fundamentally, these two operations are equivalent in this regard.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSourceBuilder.java:
##########
@@ -300,6 +313,9 @@ public FlussSource<OUT> build() {
                 offsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
                 deserializationSchema,
-                true);
+                true,
+                hasPrimaryKey
+                        ? new LeaseContext(kvSnapshotLeaseId, kvSnapshotLeaseDurationMs)
+                        : new LeaseContext(null, null));

Review Comment:
   ditto. we don't need to check primary key table.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLease.java:
##########
@@ -0,0 +1,198 @@
+/*
+ * 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.server.zk.data.lease;

Review Comment:
   The purpose of the package `org.apache.fluss.server.zk.data` is to contain **only** POJO classes that represent ZNode data entities, along with their corresponding JSON serialization/deserialization logic. These classes should be pure data carriers and **must not include any business logic**.
   
   However, `KvSnapshotLease` does not belong here because:
   - It is **not a ZNode data entity** (i.e., it is not serialized into or stored in ZooKeeper).
   - It encapsulates substantial business logic, such as `acquireBucket()` and `releaseBucket()`.
   
   Therefore, it should **not** reside under `org.apache.fluss.server.zk.data.lease`.
   
   Instead, we propose moving it to `org.apache.fluss.server.coordinator.lease`, alongside related components like `KvSnapshotLeaseManager` and `KvSnapshotLeaseMetadataManager`.
   
   Additionally, to avoid confusion with a potential public-facing `KvSnapshotLease` interface (e.g., one exposed via the `Admin` API), we suggest renaming this internal class to **`KvSnapshotLeaseHandle`**, clearly signaling its role as an internal coordination handle rather than a user-facing abstraction.



##########
fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java:
##########
@@ -178,7 +186,7 @@ private static Optional<CompletedSnapshot> subsume(
         Iterator<CompletedSnapshot> iterator = snapshots.iterator();
         while (snapshots.size() > numRetain && iterator.hasNext()) {
             CompletedSnapshot next = iterator.next();
-            if (canSubsume(next, latest)) {
+            if (canSubsume(next, latest, subsumptionChecker)) {

Review Comment:
   
   I believe the current snapshot/SST expiration mechanism is insufficient.  
   
   Currently, we rely on the `subsumptionChecker` to retain leased snapshots while cleaning up newer ones. For example, given snapshots `[1, 2, 3, 4, 5]`, if we lease snapshot `1` and configure retention to keep only **2** snapshots, the `subsume()` operation will retain `[1, 5]`.
   
   However, in `CompletedSnapshotStore#addSnapshotAndSubsumeOldestOne`, the cleanup logic calls `findLowest(completedSnapshots)`, which returns `1` as the `lowestSnapshotID`. This ID is then used to determine which SST files to delete—**but since `1` is still leased, no SST files are cleaned up at all**. As a result, **all files remain until the lease expires (e.g., after 24 hours)**, even though most are no longer needed.
   
   ### ✅ Expected Behavior
   When snapshot `1` is leased:
   - The system should retain the **newest N snapshots** (e.g., `[4, 5]` for N=2) for normal subsumption.
   - Snapshot `1` should be placed in a separate **`stillInUse` set**, indicating it’s protected by a lease but not part of the standard retention window.
   - During cleanup (`sharedKvFileRegistry.unregisterUnusedKvFile` and `snapshotsCleaner.cleanSubsumedSnapshots`), the **effective lowest snapshot ID** should be `4` (not `1`), **and** any snapshot in `stillInUse` (the `1`) must be excluded from deletion.
   
   ### 🔧 Proposed Solution
   
   1. **Track Leased Snapshots Separately**  
      Introduce a `Map<Long, CompletedSnapshot> stillInUseSnapshots` in `CompletedSnapshotStore`.  
      - When `subsume()` is called, if a snapshot is leased, **do not remove it** from the store or delete its ZK metadata.  
      - Instead, add it to `stillInUseSnapshots`.
   
   2. **Update Cleanup Logic**  
      Both `sharedKvFileRegistry.unregisterUnusedKvFile()` and `snapshotsCleaner.cleanSubsumedSnapshots()` must:
      - Compute the effective `lowestSnapshotID` **from the retained (non-leased) snapshots only**.
      - Skip deletion of any file referenced by snapshots in `stillInUseSnapshots`.
   
   3. **Enhance `SharedKvEntry` Metadata**  
      Currently, `SharedKvEntry` only tracks `lastUsedSnapshotID`. We should add `firstUsedSnapshotID` to support accurate cleanup, so it track the min and max snapshot range of this kv file. 
      Then, in `SharedKvFileRegistry#unregisterUnusedKvFile()`, a KV file can only be deleted if **no snapshot in `stillInUseSnapshots` falls within `[min, max]`** of that file’s referencing range.
   
   **Please also add tests for the kv file cleanup when lease involved, like in `CompletedSnapshotStoreTest`.**
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,

Review Comment:
   This isn't an interface—it's essentially an inner class—so there's no need to expose it in the constructor parameters. Instead, just accept the atomic parameters directly: `ZooKeeperClient zkClient` and `String remoteDataDir` and contruct the inner class in the constructor. 



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java:
##########
@@ -142,6 +143,15 @@ public DynamicTableSource createDynamicTableSource(Context context) {
                         .get(FlinkConnectorOptions.SCAN_PARTITION_DISCOVERY_INTERVAL)
                         .toMillis();
 
+        LeaseContext leaseContext =
+                primaryKeyIndexes.length > 0
+                        ? new LeaseContext(
+                                tableOptions.get(FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_ID),
+                                tableOptions
+                                        .get(FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_DURATION)
+                                        .toMillis())
+                        : new LeaseContext(null, null);

Review Comment:
   We can simplify this into one line `LeaseContext.fromConf(tableOptions)` by providing a `fromConf` static method on `LeaseContext`. No matter it is a PK table or LOG table, it's fine to have default values of lease context. 



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -378,6 +378,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;
+  repeated PbKvSnapshotLeaseForTable table_lease_req = 3;
+}
+
+message AcquireKvSnapshotLeaseResponse {
+  repeated PbKvSnapshotLeaseForTable tables_lease_res = 3;

Review Comment:
   rename the `tables_lease_res` to `unavailable_snapshots`. This `tables_lease_res` is very confusing, I thought it is the snapshot that acquired, but it is opposite. 



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -923,11 +1007,55 @@ public void addReader(int subtaskId) {
     public SourceEnumeratorState snapshotState(long checkpointId) {
         final SourceEnumeratorState enumeratorState =
                 new SourceEnumeratorState(
-                        assignedTableBuckets, assignedPartitions, pendingHybridLakeFlussSplits);
+                        assignedTableBuckets,
+                        assignedPartitions,
+                        pendingHybridLakeFlussSplits,
+                        leaseContext);
         LOG.debug("Source Checkpoint is {}", enumeratorState);
         return enumeratorState;
     }
 
+    @Override
+    public void notifyCheckpointComplete(long checkpointId) throws Exception {
+        // lower than this checkpoint id.
+        Set<TableBucket> consumedKvSnapshots = getAndRemoveConsumedBucketsUpTo(checkpointId);
+
+        LOG.info(
+                "kv snapshot has already consumed and try to release kv snapshot lease for: {}, checkpoint id: {}",
+                consumedKvSnapshots,
+                checkpointId);
+
+        // send request to fluss to unregister the kv snapshot lease.
+        try {
+            flussAdmin
+                    .releaseKvSnapshotLease(

Review Comment:
   I don't see `drop_lease` or `release_all`. What happens in batch mode? There is no checkpoints, so the `notifyCheckpointComplete` will not be invoked. I suggest to drop lease in `close()` method, if it works in batch mode and has registered lease. 



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;

Review Comment:
   ```suggestion
       private final ConcurrentHashMap<String, KvSnapshotLease> kvSnapshotLeaseMap =
               MapUtils.newConcurrentHashMap();
   ```
   
   Mark the field type `ConcurrentHashMap` and assign it inline to make the semantic clear. 



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java:
##########
@@ -76,6 +77,23 @@ public class FlinkConnectorOptions {
                             "A list of host/port pairs to use for establishing the initial connection to the Fluss cluster. "
                                     + "The list should be in the form host1:port1,host2:port2,....");
 
+    public static final ConfigOption<String> SCAN_KV_SNAPSHOT_LEASE_ID =
+            ConfigOptions.key("scan.kv.snapshot.lease.id")
+                    .stringType()
+                    .defaultValue(String.valueOf(UUID.randomUUID()))
+                    .withDescription(
+                            "The lease id to ping kv snapshots. If set, the acquired kv snapshots will not be deleted "
+                                    + "until the consumer finished consuming all the snapshots or the lease duration time "
+                                    + "is reached. If not set, an UUID will be set.");

Review Comment:
   Change to this (and also the doc):
   
   The lease ID used to protect acquired KV snapshots from deletion. If specified, the snapshots will be retained until either the consumer finishes processing all of them or the lease duration expires.  By default, this value is set to a randomly generated UUID string if not explicitly provided.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +537,88 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+            Set<TableBucket> remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Set<TableBucket> ignoreBuckets = new HashSet<>();
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : remainingTableBuckets) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                } else {
+                    ignoreBuckets.add(tb);
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!ignoreBuckets.isEmpty()) {
+                remainingTableBuckets.removeAll(ignoreBuckets);
+            }
+
+            if (!bucketsToLease.isEmpty()) {
+                String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                LOG.info(
+                        "Try to acquire kv snapshot lease {} for table {}",
+                        kvSnapshotLeaseId,
+                        tablePath);
+                Long kvSnapshotLeaseDurationMs = leaseContext.getKvSnapshotLeaseDurationMs();
+                checkNotNull(kvSnapshotLeaseDurationMs, "kv snapshot lease duration is null.");
+                remainingTableBuckets =
+                        flussAdmin
+                                .acquireKvSnapshotLease(
+                                        kvSnapshotLeaseId,
+                                        bucketsToLease,
+                                        kvSnapshotLeaseDurationMs)
+                                .get()
+                                .getUnavailableTableBucketSet();
+                if (!remainingTableBuckets.isEmpty()) {
+                    LOG.info(
+                            "Failed to acquire kv snapshot lease for table {}: {}. Retry to re-acquire",

Review Comment:
   I don't see the `Retry to re-acquire` logic. 



##########
fluss-common/src/main/java/org/apache/fluss/metadata/KvSnapshotLeaseForBucket.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.metadata;
+
+import java.util.Objects;
+
+/** An entity for kv snapshot lease for bucket. */
+public class KvSnapshotLeaseForBucket {
+    private final TableBucket tableBucket;
+    private final long kvSnapshotId;
+
+    public KvSnapshotLeaseForBucket(TableBucket tableBucket, long kvSnapshotId) {
+        this.tableBucket = tableBucket;
+        this.kvSnapshotId = kvSnapshotId;
+    }
+
+    public TableBucket getTableBucket() {
+        return tableBucket;
+    }
+
+    public long getKvSnapshotId() {

Review Comment:
   rename to `getSnapshotId()`.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/FlussSourceEnumeratorStateSerializer.java:
##########
@@ -260,4 +287,39 @@ private List<SourceSplitBase> deserializeRemainingHybridLakeFlussSplits(
             return null;
         }
     }
+
+    private void serializeLeaseContext(final DataOutputSerializer out, SourceEnumeratorState state)
+            throws IOException {
+        LeaseContext leaseContext = state.getLeaseContext();
+        String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+        if (kvSnapshotLeaseId != null) {
+            out.writeBoolean(true);
+            out.writeUTF(kvSnapshotLeaseId);
+        } else {
+            out.writeBoolean(false);
+        }
+
+        Long kvSnapshotLeaseDurationMs = leaseContext.getKvSnapshotLeaseDurationMs();
+        if (kvSnapshotLeaseDurationMs != null) {
+            out.writeBoolean(true);
+            out.writeLong(kvSnapshotLeaseDurationMs);

Review Comment:
   We don't need to serialize the duration config into state. It's not a state. Actually, we shoulld support the reconfigure for the duration config, setting it into state will block use to reconfigure it after job restarts. 



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;

Review Comment:
   only have lease expiration check interval as member variable is enought, remove the conf.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataManager.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger count = refCount.get(kvSnapshotLeaseForBucket);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    // set the expiration time as: current time + leaseDuration
+                    long newExpirationTime = clock.milliseconds() + leaseDuration;
+                    KvSnapshotLease kvSnapshotLease =
+                            kvSnapshotLeaseMap.compute(
+                                    leaseId,
+                                    (key, existingLease) -> {
+                                        if (existingLease == null) {
+                                            LOG.info(
+                                                    "kv snapshot lease '{}' has been acquired. The lease expiration "
+                                                            + "time is {}",
+                                                    leaseId,
+                                                    newExpirationTime);
+                                            return new KvSnapshotLease(newExpirationTime);
+                                        } else {
+                                            existingLease.setExpirationTime(newExpirationTime);
+                                            return existingLease;
+                                        }
+                                    });
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();

Review Comment:
   Actually, we don't need to access the `coordinatorContext` for the `numBuckets`. We can grow the snapshot array size in lease if the new acquired bucket id is larger than the previous array size. 



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java:
##########
@@ -0,0 +1,226 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.types.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * The manager to handle {@link KvSnapshotLease} to register/update/delete metadata from zk and
+ * remote fs.
+ */
+public class KvSnapshotLeaseMetadataManager {

Review Comment:
   ditto. it's not an entity class for znode. should be moved to `org.apache.fluss.server.coordinator.lease`.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataManager.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger count = refCount.get(kvSnapshotLeaseForBucket);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    // set the expiration time as: current time + leaseDuration
+                    long newExpirationTime = clock.milliseconds() + leaseDuration;
+                    KvSnapshotLease kvSnapshotLease =
+                            kvSnapshotLeaseMap.compute(
+                                    leaseId,
+                                    (key, existingLease) -> {
+                                        if (existingLease == null) {
+                                            LOG.info(
+                                                    "kv snapshot lease '{}' has been acquired. The lease expiration "
+                                                            + "time is {}",
+                                                    leaseId,
+                                                    newExpirationTime);
+                                            return new KvSnapshotLease(newExpirationTime);
+                                        } else {
+                                            existingLease.setExpirationTime(newExpirationTime);
+                                            return existingLease;
+                                        }
+                                    });
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataManager.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(
+                                                tableBucket, originalSnapshotId));
+                            }
+                            incrementRefCount(bucket);
+                        }
+                    }
+
+                    if (update) {
+                        metadataManager.updateLease(leaseId, kvSnapshotLease);
+                    } else {
+                        metadataManager.registerLease(leaseId, kvSnapshotLease);
+                    }
+
+                    return unavailableSnapshots;
+                });
+    }
+
+    public void release(String leaseId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotLease lease = kvSnapshotLeaseMap.get(leaseId);
+                    if (lease == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = lease.releaseBucket(bucket);
+                            if (snapshotId != -1L) {
+                                leasedBucketCount.decrementAndGet();
+                                decrementRefCount(new KvSnapshotLeaseForBucket(bucket, snapshotId));
+                            }
+                        }
+                    }
+
+                    if (lease.isEmpty()) {
+                        releaseAll(leaseId);
+                    } else {
+                        metadataManager.updateLease(leaseId, lease);
+                    }
+                });
+    }
+
+    /**
+     * Release kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @return true if clear success, false if lease not exist
+     */
+    public boolean releaseAll(String leaseId) throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseMap.remove(leaseId);
+                            if (kvSnapshotLease == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotLease);
+                            metadataManager.deleteLease(leaseId);
+
+                            LOG.info(
+                                    "kv snapshots of lease '"
+                                            + leaseId
+                                            + "' has been all released.");
+                            return true;
+                        });
+
+        leaseLocks.remove(leaseId);
+        return exist;
+    }
+
+    private void initializeRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+
+                    incrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        incrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                    }
+                }
+            }
+        }
+    }
+
+    private void clearRefCount(KvSnapshotLease lease) {
+        for (Map.Entry<Long, KvSnapshotTableLease> tableEntry :
+                lease.getTableIdToTableLease().entrySet()) {
+            long tableId = tableEntry.getKey();
+            KvSnapshotTableLease tableLease = tableEntry.getValue();
+            if (tableLease.getBucketSnapshots() != null) {
+                Long[] snapshots = tableLease.getBucketSnapshots();
+                for (int i = 0; i < snapshots.length; i++) {
+                    if (snapshots[i] == -1L) {
+                        continue;
+                    }
+                    decrementRefCount(
+                            new KvSnapshotLeaseForBucket(
+                                    new TableBucket(tableId, i), snapshots[i]));
+                    leasedBucketCount.decrementAndGet();
+                }
+            } else {
+                Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+                for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+                    Long partitionId = entry.getKey();
+                    Long[] snapshots = entry.getValue();
+                    for (int i = 0; i < snapshots.length; i++) {
+                        if (snapshots[i] == -1L) {
+                            continue;
+                        }
+
+                        decrementRefCount(
+                                new KvSnapshotLeaseForBucket(
+                                        new TableBucket(tableId, partitionId, i), snapshots[i]));
+                        leasedBucketCount.decrementAndGet();
+                    }
+                }
+            }
+        }
+    }
+
+    private void incrementRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        inWriteLock(
+                refCountLock,
+                () ->
+                        refCount.computeIfAbsent(
+                                        kvSnapshotLeaseForBucket, k -> new AtomicInteger(0))
+                                .incrementAndGet());
+    }

Review Comment:
   We can remove the `inWriteLock` because `refCount` is already a `ConcurrentHashMap`, which is thread-safe by design. Eliminating the lock will also improve performance.
   
   The current `inWriteLock` is highly contended: methods like `initializeRefCount` frequently invoke `incrementRefCount` and `decrementRefCount`, leading to excessive lock acquisition and release overhead. This should be optimized.
   
   Since `ConcurrentHashMap` provides sufficient concurrency guarantees for reference counting, we don’t need an additional lock. Alternatively, if a non-concurrent map were used, a read-write lock might be justified—but in this case, sticking with `ConcurrentHashMap` alone is simpler and more efficient.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),

Review Comment:
   We should avoid start one after another threads. Could you reuse the scheduler with `ProducerOffsetsManager`? The `KvSnapshotLeaseManager` should be moved into `CoordinatorService`, so it should be easy to share the thread pool.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/SourceEnumeratorState.java:
##########
@@ -41,13 +42,18 @@ public class SourceEnumeratorState {
     // lake snapshot
     @Nullable private final List<SourceSplitBase> remainingHybridLakeFlussSplits;
 
+    // lease context for restore.
+    private final LeaseContext leaseContext;
+
     public SourceEnumeratorState(
             Set<TableBucket> assignedBuckets,
             Map<Long, String> assignedPartitions,
-            @Nullable List<SourceSplitBase> remainingHybridLakeFlussSplits) {
+            @Nullable List<SourceSplitBase> remainingHybridLakeFlussSplits,
+            @Nullable LeaseContext leaseContext) {

Review Comment:
   Just serialize enumerator state with lease id



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +537,88 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+            Set<TableBucket> remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Set<TableBucket> ignoreBuckets = new HashSet<>();
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : remainingTableBuckets) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                } else {
+                    ignoreBuckets.add(tb);
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!ignoreBuckets.isEmpty()) {
+                remainingTableBuckets.removeAll(ignoreBuckets);
+            }
+
+            if (!bucketsToLease.isEmpty()) {
+                String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                LOG.info(
+                        "Try to acquire kv snapshot lease {} for table {}",
+                        kvSnapshotLeaseId,
+                        tablePath);

Review Comment:
   nit: we should print the nullable partition name as well in the log. 



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataManager.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger count = refCount.get(kvSnapshotLeaseForBucket);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    // set the expiration time as: current time + leaseDuration
+                    long newExpirationTime = clock.milliseconds() + leaseDuration;
+                    KvSnapshotLease kvSnapshotLease =
+                            kvSnapshotLeaseMap.compute(
+                                    leaseId,
+                                    (key, existingLease) -> {
+                                        if (existingLease == null) {
+                                            LOG.info(
+                                                    "kv snapshot lease '{}' has been acquired. The lease expiration "
+                                                            + "time is {}",
+                                                    leaseId,
+                                                    newExpirationTime);
+                                            return new KvSnapshotLease(newExpirationTime);
+                                        } else {
+                                            existingLease.setExpirationTime(newExpirationTime);
+                                            return existingLease;
+                                        }
+                                    });
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataManager.isSnapshotExists(tableBucket, kvSnapshotId);

Review Comment:
   **This is a VERY HEAVY operation, if there is 100K bucket, retain 10 snapshots, this will involve 1 million ZK IO operations! AND IN THE LOCK!**
   
   This will soon become a bottleneck of the system and the Flink job. 
   
   I think a better solution is to check the availability (and also pin snapshots) via `CompletedSnapshotStoreManager`. But I think this will involve a lot of code refactoring. I remembered we have discussed this, we have increased the retained snapshot number to `2` (this is the reason we increased it), so it's rarely to happen there is unavailable snapshots when we leasing it. And we can implement the unavailabe check in the future task. 
   
   So could you remove the zk snapshot check here, and create an issue for this, and add a TODO comment about the missing check with the issue id.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java:
##########
@@ -0,0 +1,226 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.types.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * The manager to handle {@link KvSnapshotLease} to register/update/delete metadata from zk and
+ * remote fs.
+ */
+public class KvSnapshotLeaseMetadataManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseMetadataManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final String remoteDataDir;
+
+    public KvSnapshotLeaseMetadataManager(ZooKeeperClient zkClient, String remoteDataDir) {
+        this.zkClient = zkClient;
+        this.remoteDataDir = remoteDataDir;
+    }
+
+    public List<String> getLeasesList() throws Exception {
+        return zkClient.getKvSnapshotLeasesList();
+    }
+
+    /**
+     * Register a new kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param lease the kv snapshot lease.
+     */
+    public void registerLease(String leaseId, KvSnapshotLease lease) throws Exception {
+        Map<Long, FsPath> tableIdToRemoteMetadataFsPath = generateMetadataFile(leaseId, lease);
+
+        // generate remote fsPath of metadata.
+        KvSnapshotLeaseMetadata leaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        lease.getExpirationTime(), tableIdToRemoteMetadataFsPath);
+
+        // register kv snapshot metadata to zk.
+        try {
+            zkClient.registerKvSnapshotLeaseMetadata(leaseId, leaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to register kv snapshot lease metadata to zk.", e);
+            leaseMetadata.discard();

Review Comment:
   This looks very dangerous, will this remove previous success metadata files if there is exception happen this time? Could you learn from `ProducerOffsetsStore.tryStoreOffsets` to cleanup only the `createdFiles`?



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {

Review Comment:
   move `initialize()` into  `start()`, and schedule the cleanup task after the `initialize()`. It's strange we have 2 start up methods. Besides



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +537,88 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+            Set<TableBucket> remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Set<TableBucket> ignoreBuckets = new HashSet<>();
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : remainingTableBuckets) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                } else {
+                    ignoreBuckets.add(tb);
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!ignoreBuckets.isEmpty()) {
+                remainingTableBuckets.removeAll(ignoreBuckets);
+            }
+
+            if (!bucketsToLease.isEmpty()) {
+                String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();
+                LOG.info(
+                        "Try to acquire kv snapshot lease {} for table {}",
+                        kvSnapshotLeaseId,
+                        tablePath);
+                Long kvSnapshotLeaseDurationMs = leaseContext.getKvSnapshotLeaseDurationMs();
+                checkNotNull(kvSnapshotLeaseDurationMs, "kv snapshot lease duration is null.");
+                remainingTableBuckets =

Review Comment:
   Please use a separate variable `unavailableSnapshots` for the return value of `acquireKvSnapshotLease`. Overwrite the existing `remainingTableBuckets` is very confusing and no sense. 



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLease.java:
##########
@@ -0,0 +1,188 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.utils.MapUtils;
+
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+
+/** The lease of kv snapshot for a table. */
+@NotThreadSafe
+public class KvSnapshotTableLease {
+    private final long tableId;
+    private final @Nullable Long[] bucketSnapshots;
+    private final Map<Long, Long[]> partitionSnapshots;
+
+    public KvSnapshotTableLease(long tableId) {
+        this(tableId, null, MapUtils.newConcurrentHashMap());

Review Comment:
   Mark `partitionSnapshots` as `ConcurrnetHashMap`, if it needs to be thread-safe. I see many places assigns different hash maps to the same variable. 



##########
fluss-common/src/main/java/org/apache/fluss/metadata/KvSnapshotLeaseForBucket.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.metadata;
+
+import java.util.Objects;
+
+/** An entity for kv snapshot lease for bucket. */
+public class KvSnapshotLeaseForBucket {

Review Comment:
   rename the class to `TableBucketSnapshot`, so that it can be reused not only for lease.



##########
fluss-server/src/test/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManagerTest.java:
##########
@@ -0,0 +1,470 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.PhysicalTablePath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.metrics.group.TestingMetricGroups;
+import org.apache.fluss.server.zk.NOPErrorHandler;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.ZooKeeperExtension;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.testutils.common.ManuallyTriggeredScheduledExecutorService;
+import org.apache.fluss.utils.clock.ManualClock;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.fluss.record.TestData.DATA1_TABLE_ID_PK;
+import static org.apache.fluss.record.TestData.DATA1_TABLE_INFO_PK;
+import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH_PK;
+import static org.apache.fluss.record.TestData.PARTITION_TABLE_ID;
+import static org.apache.fluss.record.TestData.PARTITION_TABLE_INFO;
+import static org.apache.fluss.record.TestData.PARTITION_TABLE_PATH;
+import static org.apache.fluss.server.coordinator.CoordinatorTestUtils.createServers;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Test for {@link KvSnapshotLeaseManager}. */
+public class KvSnapshotLeaseManagerTest {
+
+    @RegisterExtension
+    public static final AllCallbackWrapper<ZooKeeperExtension> ZOO_KEEPER_EXTENSION_WRAPPER =
+            new AllCallbackWrapper<>(new ZooKeeperExtension());
+
+    private static final long PARTITION_ID_1 = 19001L;
+    private static final PhysicalTablePath PARTITION_TABLE_PATH_1 =
+            PhysicalTablePath.of(PARTITION_TABLE_PATH, "2024");
+
+    private static final long PARTITION_ID_2 = 19002L;
+    private static final PhysicalTablePath PARTITION_TABLE_PATH_2 =
+            PhysicalTablePath.of(PARTITION_TABLE_PATH, "2025");
+
+    private static final int NUM_BUCKETS = DATA1_TABLE_INFO_PK.getNumBuckets();
+    private static final TableBucket t0b0 = new TableBucket(DATA1_TABLE_ID_PK, 0);
+    private static final TableBucket t0b1 = new TableBucket(DATA1_TABLE_ID_PK, 1);
+    private static final TableBucket t1p0b0 =
+            new TableBucket(PARTITION_TABLE_ID, PARTITION_ID_1, 0);
+    private static final TableBucket t1p0b1 =
+            new TableBucket(PARTITION_TABLE_ID, PARTITION_ID_1, 1);
+    private static final TableBucket t1p1b0 =
+            new TableBucket(PARTITION_TABLE_ID, PARTITION_ID_2, 0);
+
+    protected static ZooKeeperClient zookeeperClient;
+
+    private CoordinatorContext coordinatorContext;
+    private ManualClock manualClock;
+    private ManuallyTriggeredScheduledExecutorService clearLeaseScheduler;
+    private KvSnapshotLeaseManager kvSnapshotLeaseManager;
+    private KvSnapshotLeaseMetadataManager metadataManager;
+
+    private @TempDir Path tempDir;
+
+    @BeforeAll
+    static void beforeAll() {
+        zookeeperClient =
+                ZOO_KEEPER_EXTENSION_WRAPPER
+                        .getCustomExtension()
+                        .getZooKeeperClient(NOPErrorHandler.INSTANCE);
+    }
+
+    @BeforeEach
+    void beforeEach() throws Exception {
+        initCoordinatorContext();
+        Configuration conf = new Configuration();
+        // set a huge expiration check interval to avoid expiration check.
+        conf.set(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL, Duration.ofDays(7));
+        manualClock = new ManualClock(System.currentTimeMillis());
+        clearLeaseScheduler = new ManuallyTriggeredScheduledExecutorService();
+        metadataManager = new KvSnapshotLeaseMetadataManager(zookeeperClient, tempDir.toString());
+        kvSnapshotLeaseManager =
+                new KvSnapshotLeaseManager(
+                        conf,
+                        metadataManager,
+                        coordinatorContext,
+                        clearLeaseScheduler,
+                        manualClock,
+                        TestingMetricGroups.COORDINATOR_METRICS);
+        kvSnapshotLeaseManager.start();
+        initialZookeeper();
+    }
+
+    @AfterEach
+    void afterEach() {
+        ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupRoot();
+    }
+
+    private static void initialZookeeper() throws Exception {
+        List<TableBucket> tableBuckets = Arrays.asList(t0b0, t0b1, t1p0b0, t1p0b1, t1p1b0);
+        for (TableBucket tb : tableBuckets) {
+            zookeeperClient.registerTableBucketSnapshot(
+                    tb, new BucketSnapshot(0L, 0L, "test-path"));
+        }
+    }
+
+    @Test
+    void testInitialize() throws Exception {
+        assertThat(
+                        snapshotLeaseNotExists(
+                                Arrays.asList(
+                                        new KvSnapshotLeaseForBucket(t0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p1b0, 0L))))
+                .isTrue();
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+
+        // test initialize from zookeeper when coordinator is started.
+        KvSnapshotLease kvSnapshotLease = new KvSnapshotLease(1000L);
+        acquire(kvSnapshotLease, new KvSnapshotLeaseForBucket(t0b0, 0L));
+        acquire(kvSnapshotLease, new KvSnapshotLeaseForBucket(t0b1, 0L));
+        acquire(kvSnapshotLease, new KvSnapshotLeaseForBucket(t1p0b0, 0L));
+        acquire(kvSnapshotLease, new KvSnapshotLeaseForBucket(t1p0b1, 0L));
+        acquire(kvSnapshotLease, new KvSnapshotLeaseForBucket(t1p1b0, 0L));
+        metadataManager.registerLease("lease1", kvSnapshotLease);
+
+        kvSnapshotLeaseManager.initialize();
+
+        assertThat(
+                        snapshotLeaseExists(
+                                Arrays.asList(
+                                        new KvSnapshotLeaseForBucket(t0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p1b0, 0L))))
+                .isTrue();
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(5);
+
+        // check detail content.
+        Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+        tableIdToTableLease.put(
+                DATA1_TABLE_ID_PK,
+                new KvSnapshotTableLease(DATA1_TABLE_ID_PK, new Long[] {0L, 0L, -1L}));
+        KvSnapshotTableLease leaseForPartitionTable = new KvSnapshotTableLease(PARTITION_TABLE_ID);
+        leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_1, new Long[] {0L, 0L, -1L});
+        leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_2, new Long[] {0L, -1L, -1L});
+        tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+        KvSnapshotLease expectedLease = new KvSnapshotLease(1000L, tableIdToTableLease);
+        assertThat(kvSnapshotLeaseManager.getKvSnapshotLease("lease1")).isEqualTo(expectedLease);
+        assertThat(metadataManager.getLease("lease1")).hasValue(expectedLease);
+    }
+
+    @Test
+    void testAcquireAndRelease() throws Exception {
+        Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToRegisterBucket = initRegisterBuckets();
+        acquire("lease1", tableIdToRegisterBucket);
+
+        // first register snapshot to zk.
+        zookeeperClient.registerTableBucketSnapshot(
+                t1p0b0, new BucketSnapshot(1L, 10L, "test-path"));
+        zookeeperClient.registerTableBucketSnapshot(
+                t1p0b1, new BucketSnapshot(1L, 10L, "test-path"));
+
+        tableIdToRegisterBucket = new HashMap<>();
+        tableIdToRegisterBucket.put(
+                PARTITION_TABLE_ID,
+                Arrays.asList(
+                        new KvSnapshotLeaseForBucket(t1p0b0, 1L),
+                        new KvSnapshotLeaseForBucket(t1p0b1, 1L)));
+        acquire("lease2", tableIdToRegisterBucket);
+
+        assertThat(
+                        snapshotLeaseExists(
+                                Arrays.asList(
+                                        new KvSnapshotLeaseForBucket(t0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b0, 0),
+                                        new KvSnapshotLeaseForBucket(t1p0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p1b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b0, 1L),
+                                        new KvSnapshotLeaseForBucket(t1p0b1, 1L))))
+                .isTrue();
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(7);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(2);
+
+        // update lease register.
+        tableIdToRegisterBucket = new HashMap<>();
+        zookeeperClient.registerTableBucketSnapshot(t0b0, new BucketSnapshot(1L, 10L, "test-path"));
+        tableIdToRegisterBucket.put(
+                DATA1_TABLE_ID_PK,
+                Collections.singletonList(new KvSnapshotLeaseForBucket(t0b0, 1L)));
+        acquire("lease1", tableIdToRegisterBucket);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(7);
+
+        // new insert.
+        tableIdToRegisterBucket = new HashMap<>();
+        TableBucket newTableBucket = new TableBucket(DATA1_TABLE_ID_PK, 2);
+
+        zookeeperClient.registerTableBucketSnapshot(
+                newTableBucket, new BucketSnapshot(1L, 10L, "test-path"));
+        tableIdToRegisterBucket.put(
+                DATA1_TABLE_ID_PK,
+                Collections.singletonList(new KvSnapshotLeaseForBucket(newTableBucket, 1L)));
+        acquire("lease1", tableIdToRegisterBucket);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(8);
+
+        // release
+        Map<Long, List<TableBucket>> tableIdToUnregisterBucket = new HashMap<>();
+        tableIdToUnregisterBucket.put(DATA1_TABLE_ID_PK, Collections.singletonList(newTableBucket));
+        release("lease1", tableIdToUnregisterBucket);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(7);
+
+        // release a non-exist bucket.
+        tableIdToUnregisterBucket = new HashMap<>();
+        tableIdToUnregisterBucket.put(
+                DATA1_TABLE_ID_PK,
+                Collections.singletonList(new TableBucket(DATA1_TABLE_ID_PK, PARTITION_ID_1, 2)));
+        release("lease1", tableIdToUnregisterBucket);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(7);
+
+        // check detail content for lease1.
+        Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+        tableIdToTableLease.put(
+                DATA1_TABLE_ID_PK,
+                new KvSnapshotTableLease(DATA1_TABLE_ID_PK, new Long[] {1L, 0L, -1L}));
+        KvSnapshotTableLease leaseForPartitionTable = new KvSnapshotTableLease(PARTITION_TABLE_ID);
+        leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_1, new Long[] {0L, 0L, -1L});
+        leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_2, new Long[] {0L, -1L, -1L});
+        tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+        KvSnapshotLease expectedLease =
+                new KvSnapshotLease(manualClock.milliseconds() + 1000L, tableIdToTableLease);
+        assertThat(kvSnapshotLeaseManager.getKvSnapshotLease("lease1")).isEqualTo(expectedLease);
+        assertThat(metadataManager.getLease("lease1")).hasValue(expectedLease);
+
+        // check detail content for lease2.
+        tableIdToTableLease = new HashMap<>();
+        leaseForPartitionTable = new KvSnapshotTableLease(PARTITION_TABLE_ID);
+        leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_1, new Long[] {1L, 1L, -1L});
+        tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+        KvSnapshotLease expectedLease2 =
+                new KvSnapshotLease(manualClock.milliseconds() + 1000L, tableIdToTableLease);
+        assertThat(kvSnapshotLeaseManager.getKvSnapshotLease("lease2")).isEqualTo(expectedLease2);
+        assertThat(metadataManager.getLease("lease2")).hasValue(expectedLease2);
+    }
+
+    @Test
+    void testUnregisterAll() throws Exception {
+        Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToRegisterBucket = initRegisterBuckets();
+        acquire("lease1", tableIdToRegisterBucket);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(5);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(1);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isPresent();
+
+        Map<Long, List<TableBucket>> tableIdToUnregisterBucket = new HashMap<>();
+        tableIdToUnregisterBucket.put(DATA1_TABLE_ID_PK, Arrays.asList(t0b0, t0b1));
+        tableIdToUnregisterBucket.put(PARTITION_TABLE_ID, Arrays.asList(t1p0b0, t1p0b1, t1p1b0));
+
+        // unregister all will clear this lease.
+        release("lease1", tableIdToUnregisterBucket);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(0);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isNotPresent();
+    }
+
+    @Test
+    void testClear() throws Exception {
+        Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToRegisterBucket = initRegisterBuckets();
+        acquire("lease1", tableIdToRegisterBucket);
+
+        // first register snapshot to zk.
+        zookeeperClient.registerTableBucketSnapshot(
+                t1p0b0, new BucketSnapshot(1L, 10L, "test-path"));
+        zookeeperClient.registerTableBucketSnapshot(
+                t1p0b1, new BucketSnapshot(1L, 10L, "test-path"));
+        tableIdToRegisterBucket = new HashMap<>();
+        tableIdToRegisterBucket.put(
+                PARTITION_TABLE_ID,
+                Arrays.asList(
+                        new KvSnapshotLeaseForBucket(t0b0, 0L), // same ref.
+                        new KvSnapshotLeaseForBucket(t1p0b0, 1L),
+                        new KvSnapshotLeaseForBucket(t1p0b1, 1L)));
+        acquire("lease2", tableIdToRegisterBucket);
+
+        assertThat(
+                        snapshotLeaseExists(
+                                Arrays.asList(
+                                        new KvSnapshotLeaseForBucket(t0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b1, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p1b0, 0L),
+                                        new KvSnapshotLeaseForBucket(t1p0b0, 1L),
+                                        new KvSnapshotLeaseForBucket(t1p0b1, 1L))))
+                .isTrue();
+        assertThat(kvSnapshotLeaseManager.getRefCount(new KvSnapshotLeaseForBucket(t0b0, 0L)))
+                .isEqualTo(2);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(8);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(2);
+
+        kvSnapshotLeaseManager.releaseAll("lease1");
+        assertThat(kvSnapshotLeaseManager.getRefCount(new KvSnapshotLeaseForBucket(t0b0, 0L)))
+                .isEqualTo(1);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(3);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(1);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isEmpty();
+
+        kvSnapshotLeaseManager.releaseAll("lease2");
+        assertThat(kvSnapshotLeaseManager.getRefCount(new KvSnapshotLeaseForBucket(t0b0, 0L)))
+                .isEqualTo(0);
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(0);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease2")).isEmpty();
+
+        assertThat(kvSnapshotLeaseManager.releaseAll("non-exist")).isFalse();
+    }
+
+    @Test
+    void testExpireLeases() throws Exception {
+        // test lease expire by expire thread.
+        Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket = initRegisterBuckets();
+
+        // expire after 1000ms.
+        kvSnapshotLeaseManager.acquireLease("lease1", 1000L, tableIdToLeaseBucket);
+
+        tableIdToLeaseBucket = new HashMap<>();
+        zookeeperClient.registerTableBucketSnapshot(
+                t1p0b0, new BucketSnapshot(1L, 10L, "test-path"));
+        zookeeperClient.registerTableBucketSnapshot(
+                t1p0b1, new BucketSnapshot(1L, 10L, "test-path"));
+        tableIdToLeaseBucket.put(
+                PARTITION_TABLE_ID,
+                Arrays.asList(
+                        new KvSnapshotLeaseForBucket(t0b0, 0L), // same ref.
+                        new KvSnapshotLeaseForBucket(t1p0b0, 1L),
+                        new KvSnapshotLeaseForBucket(t1p0b1, 1L)));
+        // expire after 2000ms.
+        kvSnapshotLeaseManager.acquireLease("lease2", 2000L, tableIdToLeaseBucket);
+
+        clearLeaseScheduler.triggerPeriodicScheduledTasks();
+        // no lease expire.
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(8);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(2);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isPresent();
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease2")).isPresent();
+
+        manualClock.advanceTime(1005L, TimeUnit.MILLISECONDS);
+        clearLeaseScheduler.triggerPeriodicScheduledTasks();
+        // lease1 expire.
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(3);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(1);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isNotPresent();
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease2")).isPresent();
+
+        manualClock.advanceTime(1005L, TimeUnit.MILLISECONDS);
+        clearLeaseScheduler.triggerPeriodicScheduledTasks();
+        // lease2 expire.
+        assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+        assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(0);
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isNotPresent();
+        assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease2")).isNotPresent();
+    }
+
+    @Test
+    void registerWithNotExistSnapshotId() throws Exception {
+        Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToRegisterBucket = new HashMap<>();
+        tableIdToRegisterBucket.put(
+                DATA1_TABLE_ID_PK,
+                Arrays.asList(
+                        new KvSnapshotLeaseForBucket(t0b0, 1000L),
+                        new KvSnapshotLeaseForBucket(t0b1, 1000L)));
+        assertThat(
+                        kvSnapshotLeaseManager
+                                .acquireLease("lease1", 1000L, tableIdToRegisterBucket)
+                                .keySet())
+                .contains(t0b0);

Review Comment:
   ```suggestion
                   .containsExactlyInAnyOrder(t0b0, t0b1);
   ```



##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlussSourceITCase.java:
##########
@@ -151,6 +156,8 @@ public void testRowDataPKTableSource() throws Exception {
                         .setTable(pkTableName)
                         .setStartingOffsets(OffsetsInitializer.earliest())
                         .setScanPartitionDiscoveryIntervalMs(1000L)
+                        .setKvSnapshotLeaseId("test-lease-sdafa3")
+                        .setKvSnapshotLeaseDurationMs(Duration.ofDays(1).toMillis())

Review Comment:
   remove this. all existing DataStream API jobs should work without modification, otherwise, it's an in-compatible.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java:
##########
@@ -531,22 +537,88 @@ private List<SourceSplitBase> initPrimaryKeyTablePartitionSplits(
         List<SourceSplitBase> splits = new ArrayList<>();
         for (Partition partition : newPartitions) {
             String partitionName = partition.getPartitionName();
-            // get the table snapshot info
-            final KvSnapshots kvSnapshots;
-            try {
-                kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
-            } catch (Exception e) {
-                throw new FlinkRuntimeException(
-                        String.format(
-                                "Failed to get table snapshot for table %s and partition %s",
-                                tablePath, partitionName),
-                        ExceptionUtils.stripCompletionException(e));
-            }
-            splits.addAll(getSnapshotAndLogSplits(kvSnapshots, partitionName));
+            splits.addAll(
+                    getSnapshotAndLogSplits(
+                            getLatestKvSnapshotsAndRegister(partitionName), partitionName));
         }
         return splits;
     }
 
+    private KvSnapshots getLatestKvSnapshotsAndRegister(@Nullable String partitionName) {
+        long tableId;
+        Long partitionId;
+        Map<Integer, Long> snapshotIds = new HashMap<>();
+        Map<Integer, Long> logOffsets = new HashMap<>();
+
+        // retry to get the latest kv snapshots and acquire kvSnapshot lease.
+        try {
+            KvSnapshots kvSnapshots = getLatestKvSnapshots(partitionName);
+            Set<TableBucket> remainingTableBuckets = new HashSet<>(kvSnapshots.getTableBuckets());
+
+            tableId = kvSnapshots.getTableId();
+            partitionId = kvSnapshots.getPartitionId();
+
+            Set<TableBucket> ignoreBuckets = new HashSet<>();
+            Map<TableBucket, Long> bucketsToLease = new HashMap<>();
+            for (TableBucket tb : remainingTableBuckets) {
+                int bucket = tb.getBucket();
+                OptionalLong snapshotIdOpt = kvSnapshots.getSnapshotId(bucket);
+                OptionalLong logOffsetOpt = kvSnapshots.getLogOffset(bucket);
+                if (snapshotIdOpt.isPresent() && !ignoreTableBucket(tb)) {
+                    bucketsToLease.put(tb, snapshotIdOpt.getAsLong());
+                } else {
+                    ignoreBuckets.add(tb);
+                }
+
+                snapshotIds.put(
+                        bucket, snapshotIdOpt.isPresent() ? snapshotIdOpt.getAsLong() : null);
+                logOffsets.put(bucket, logOffsetOpt.isPresent() ? logOffsetOpt.getAsLong() : null);
+            }
+
+            if (!ignoreBuckets.isEmpty()) {
+                remainingTableBuckets.removeAll(ignoreBuckets);
+            }
+
+            if (!bucketsToLease.isEmpty()) {
+                String kvSnapshotLeaseId = leaseContext.getKvSnapshotLeaseId();

Review Comment:
   The `kvSnapshotLeaseId` can be `null` when a job is recovered from an older version, since the enumerator state did not include a lease ID previously. This currently triggers exceptions because the code assumes the lease ID is always non-null.
   
   A clean solution is to ensure that both `kvSnapshotLeaseId` and `kvSnapshotLeaseDurationMs` in `LeaseContext` are never null by providing sensible defaults. Specifically:
   
   - Introduce a static factory method `LeaseContext.default()` that generates a random UUID for the lease ID and sets the lease duration to **1 day**.
   - This aligns with the default behavior of the `SCAN_KV_SNAPSHOT_LEASE_ID` configuration option, which also uses a random UUID when unspecified.
   
   With this approach, during deserialization of legacy enumerator state (where `kvSnapshotLeaseId` is missing or null), we can safely use `LeaseContext.default()` to maintain consistent semantics across versions.



##########
website/docs/engine-flink/procedures.md:
##########
@@ -487,4 +487,38 @@ CALL sys.cancel_rebalance();
 
 -- Cancel a specific rebalance operation by ID
 CALL sys.cancel_rebalance('rebalance-12345');
+```
+
+## kv snapshot lease
+
+Fluss provides procedures to manage KV snapshot leases, allowing you to release leased kv snapshots.
+
+### release_all_kv_snapshot_lease
+
+Release all KV snapshots leased under a specified leaseId. This is typically used for handle the scenario of lease
+remnants. After a normal job completion, the registered lease is not released, requiring a manual trigger of the
+procedure for cleanup.

Review Comment:
   `After a normal job completion, the registered lease is not released` This sounds very terrible. It sounds by design, the lease will not be released, and requires a manual release. 
   
   I think this should be "After an abnormal job termination (e.g., crash or forced cancellation), the registered lease may not be released automatically and could require manual cleanup."
   



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/ProcedureManager.java:
##########
@@ -77,7 +77,9 @@ private enum ProcedureEnum {
         REMOVE_SERVER_TAG("sys.remove_server_tag", RemoveServerTagProcedure.class),
         REBALANCE("sys.rebalance", RebalanceProcedure.class),
         CANCEL_REBALANCE("sys.cancel_rebalance", CancelRebalanceProcedure.class),
-        LIST_REBALANCE_PROGRESS("sys.list_rebalance", ListRebalanceProcessProcedure.class);
+        LIST_REBALANCE_PROGRESS("sys.list_rebalance", ListRebalanceProcessProcedure.class),
+        RELEASE_ALL_KV_SNAPSHOT_LEASE(
+                "sys.release_all_kv_snapshot_lease", ReleaseAllKvSnapshotLeaseProcedure.class);

Review Comment:
   `release_all_kv_snapshot_lease` -> `drop_kv_snapshot_lease`



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadata.java:
##########
@@ -0,0 +1,94 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Objects;
+
+import static org.apache.fluss.metrics.registry.MetricRegistry.LOG;
+
+/** The zkNode data of kv snapshot lease. */
+public class KvSnapshotLeaseMetadata {
+    private final long expirationTime;
+    private final Map<Long, FsPath> tableIdToRemoteMetadataFilePath;
+
+    public KvSnapshotLeaseMetadata(
+            long expirationTime, Map<Long, FsPath> tableIdToRemoteMetadataFilePath) {
+        this.expirationTime = expirationTime;
+        this.tableIdToRemoteMetadataFilePath = tableIdToRemoteMetadataFilePath;
+    }
+
+    public long getExpirationTime() {
+        return expirationTime;
+    }
+
+    public Map<Long, FsPath> getTableIdToRemoteMetadataFilePath() {
+        return tableIdToRemoteMetadataFilePath;
+    }
+
+    public void discard() {
+        // delete all remote metadata file.
+        tableIdToRemoteMetadataFilePath.values().forEach(this::delete);
+    }
+
+    private void delete(FsPath fsPath) {

Review Comment:
   `KvSnapshotLeaseMetadata` is a zk entity class (POJO class), we should move the clean logic out of the class to keep the entity clean. 



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java:
##########
@@ -0,0 +1,226 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.types.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * The manager to handle {@link KvSnapshotLease} to register/update/delete metadata from zk and
+ * remote fs.
+ */
+public class KvSnapshotLeaseMetadataManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseMetadataManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final String remoteDataDir;
+
+    public KvSnapshotLeaseMetadataManager(ZooKeeperClient zkClient, String remoteDataDir) {
+        this.zkClient = zkClient;
+        this.remoteDataDir = remoteDataDir;
+    }
+
+    public List<String> getLeasesList() throws Exception {
+        return zkClient.getKvSnapshotLeasesList();
+    }
+
+    /**
+     * Register a new kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param lease the kv snapshot lease.
+     */
+    public void registerLease(String leaseId, KvSnapshotLease lease) throws Exception {
+        Map<Long, FsPath> tableIdToRemoteMetadataFsPath = generateMetadataFile(leaseId, lease);
+
+        // generate remote fsPath of metadata.
+        KvSnapshotLeaseMetadata leaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        lease.getExpirationTime(), tableIdToRemoteMetadataFsPath);
+
+        // register kv snapshot metadata to zk.
+        try {
+            zkClient.registerKvSnapshotLeaseMetadata(leaseId, leaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to register kv snapshot lease metadata to zk.", e);
+            leaseMetadata.discard();
+            throw e;
+        }
+    }
+
+    /**
+     * Update a kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param kvSnapshotLease the kv snapshot lease.
+     */
+    public void updateLease(String leaseId, KvSnapshotLease kvSnapshotLease) throws Exception {
+        // TODO change this to incremental update to avoid create too many remote metadata files.
+
+        Optional<KvSnapshotLeaseMetadata> originalLeaseMetadata =
+                zkClient.getKvSnapshotLeaseMetadata(leaseId);
+
+        Map<Long, FsPath> tableIdToNewRemoteMetadataFsPath =
+                generateMetadataFile(leaseId, kvSnapshotLease);
+
+        // generate  new kv snapshot lease metadata.
+        KvSnapshotLeaseMetadata newLeaseMetadata =
+                new KvSnapshotLeaseMetadata(
+                        kvSnapshotLease.getExpirationTime(), tableIdToNewRemoteMetadataFsPath);
+        // register new snapshot metadata to zk.
+        try {
+            zkClient.updateKvSnapshotLeaseMetadata(leaseId, newLeaseMetadata);
+        } catch (Exception e) {
+            LOG.warn("Failed to update kv snapshot lease metadata to zk.", e);
+            newLeaseMetadata.discard();

Review Comment:
   ditto



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -378,6 +378,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;
+  repeated PbKvSnapshotLeaseForTable table_lease_req = 3;
+}
+
+message AcquireKvSnapshotLeaseResponse {
+  repeated PbKvSnapshotLeaseForTable tables_lease_res = 3;
+}
+
+message ReleaseKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  repeated PbTable release_tables = 2;

Review Comment:
   We should add `optional bool release_all` to indicate it is a drop lease operation (or use the `DropKvSnapshotLeaseRequest` for dropping lease), rather than relying on empty `release_tables`. Usually, empty set means there is nothing to release, not release all...
   
   
   Besides, could you reuse the `PbTableBucket` structure to avoid introducing new type? And rename the field name to `buckets_to_release`.
   



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.metadataManager = metadataManager;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.kvSnapshotLeaseMap = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        LOG.info("kv snapshot lease manager has been started.");
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireLeases,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        for (String leaseId : metadataManager.getLeasesList()) {
+            Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataManager.getLease(leaseId);
+            if (kvSnapshotLeaseOpt.isPresent()) {
+                KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
+                this.leaseLocks.put(leaseId, new ReentrantReadWriteLock());
+                this.kvSnapshotLeaseMap.put(leaseId, kvSnapshotLease);
+
+                initializeRefCount(kvSnapshotLease);
+
+                leasedBucketCount.addAndGet(kvSnapshotLease.getLeasedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
+        return inReadLock(
+                refCountLock,
+                () -> {
+                    AtomicInteger count = refCount.get(kvSnapshotLeaseForBucket);
+                    return count == null || count.get() <= 0;
+                });
+    }
+
+    /**
+     * Acquire kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @param leaseDuration the lease duration
+     * @param tableIdToLeaseBucket the table id to lease bucket
+     * @return the map of unavailable snapshots that failed to be leased
+     */
+    public Map<TableBucket, Long> acquireLease(
+            String leaseId,
+            long leaseDuration,
+            Map<Long, List<KvSnapshotLeaseForBucket>> tableIdToLeaseBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.computeIfAbsent(leaseId, k -> new ReentrantReadWriteLock());
+        return inWriteLock(
+                lock,
+                () -> {
+                    // To record the unavailable snapshots such as the kv snapshotId to lease not
+                    // exists.
+                    Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+
+                    boolean update = kvSnapshotLeaseMap.containsKey(leaseId);
+                    // set the expiration time as: current time + leaseDuration
+                    long newExpirationTime = clock.milliseconds() + leaseDuration;
+                    KvSnapshotLease kvSnapshotLease =
+                            kvSnapshotLeaseMap.compute(
+                                    leaseId,
+                                    (key, existingLease) -> {
+                                        if (existingLease == null) {
+                                            LOG.info(
+                                                    "kv snapshot lease '{}' has been acquired. The lease expiration "
+                                                            + "time is {}",
+                                                    leaseId,
+                                                    newExpirationTime);
+                                            return new KvSnapshotLease(newExpirationTime);
+                                        } else {
+                                            existingLease.setExpirationTime(newExpirationTime);
+                                            return existingLease;
+                                        }
+                                    });
+
+                    for (Map.Entry<Long, List<KvSnapshotLeaseForBucket>> entry :
+                            tableIdToLeaseBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<KvSnapshotLeaseForBucket> buckets = entry.getValue();
+                        for (KvSnapshotLeaseForBucket bucket : buckets) {
+
+                            TableBucket tableBucket = bucket.getTableBucket();
+                            long kvSnapshotId = bucket.getKvSnapshotId();
+                            try {
+                                boolean snapshotExists =
+                                        metadataManager.isSnapshotExists(tableBucket, kvSnapshotId);
+                                if (!snapshotExists) {
+                                    unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                    continue;
+                                }
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Failed to check snapshotExists for tableBucket when acquire kv "
+                                                + "snapshot kvSnapshotLease {}.",
+                                        tableBucket,
+                                        e);
+                                unavailableSnapshots.put(tableBucket, kvSnapshotId);
+                                continue;
+                            }
+
+                            long originalSnapshotId =
+                                    kvSnapshotLease.acquireBucket(
+                                            tableBucket, kvSnapshotId, numBuckets);
+                            if (originalSnapshotId == -1L) {
+                                leasedBucketCount.incrementAndGet();
+                            } else {
+                                // clear the original ref.
+                                decrementRefCount(
+                                        new KvSnapshotLeaseForBucket(
+                                                tableBucket, originalSnapshotId));
+                            }
+                            incrementRefCount(bucket);
+                        }
+                    }
+
+                    if (update) {
+                        metadataManager.updateLease(leaseId, kvSnapshotLease);
+                    } else {
+                        metadataManager.registerLease(leaseId, kvSnapshotLease);
+                    }
+
+                    return unavailableSnapshots;
+                });
+    }
+
+    public void release(String leaseId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotLease lease = kvSnapshotLeaseMap.get(leaseId);
+                    if (lease == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = lease.releaseBucket(bucket);
+                            if (snapshotId != -1L) {
+                                leasedBucketCount.decrementAndGet();
+                                decrementRefCount(new KvSnapshotLeaseForBucket(bucket, snapshotId));
+                            }
+                        }
+                    }
+
+                    if (lease.isEmpty()) {
+                        releaseAll(leaseId);
+                    } else {
+                        metadataManager.updateLease(leaseId, lease);
+                    }
+                });
+    }
+
+    /**
+     * Release kv snapshot lease.
+     *
+     * @param leaseId the lease id
+     * @return true if clear success, false if lease not exist
+     */
+    public boolean releaseAll(String leaseId) throws Exception {
+        ReadWriteLock lock = leaseLocks.get(leaseId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseMap.remove(leaseId);
+                            if (kvSnapshotLease == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotLease);
+                            metadataManager.deleteLease(leaseId);
+
+                            LOG.info(
+                                    "kv snapshots of lease '"
+                                            + leaseId
+                                            + "' has been all released.");

Review Comment:
   nit: use `{}` in logging



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -378,6 +378,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;
+  repeated PbKvSnapshotLeaseForTable table_lease_req = 3;

Review Comment:
   rename to `snapshots_to_lease` or `snapshots_to_acquire`



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java:
##########
@@ -0,0 +1,226 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.BucketSnapshot;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.types.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * The manager to handle {@link KvSnapshotLease} to register/update/delete metadata from zk and
+ * remote fs.
+ */
+public class KvSnapshotLeaseMetadataManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseMetadataManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final String remoteDataDir;
+
+    public KvSnapshotLeaseMetadataManager(ZooKeeperClient zkClient, String remoteDataDir) {
+        this.zkClient = zkClient;
+        this.remoteDataDir = remoteDataDir;
+    }
+
+    public List<String> getLeasesList() throws Exception {
+        return zkClient.getKvSnapshotLeasesList();
+    }
+
+    /**
+     * Register a new kv snapshot lease to zk and remote fs.
+     *
+     * @param leaseId the lease id.
+     * @param lease the kv snapshot lease.
+     */
+    public void registerLease(String leaseId, KvSnapshotLease lease) throws Exception {
+        Map<Long, FsPath> tableIdToRemoteMetadataFsPath = generateMetadataFile(leaseId, lease);

Review Comment:
   It would be better to wrap the `generateMetadataFile` call into the following try-catch block, as it involves multiple remote I/O operations and has a high likelihood of throwing an `IOException`. And we should fallback and cleanup created files.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java:
##########
@@ -77,13 +88,53 @@ public FlinkSourceReader(
                 recordEmitter,
                 context.getConfiguration(),
                 context);
+        this.finishedKvSnapshotConsumeBuckets = new HashSet<>();
     }
 
     @Override
     protected void onSplitFinished(Map<String, SourceSplitState> map) {
         // do nothing
     }
 
+    @Override
+    public List<SourceSplitBase> snapshotState(long checkpointId) {
+        Set<TableBucket> bucketsFinishedConsumeKvSnapshot = new HashSet<>();
+
+        // do not modify this state.
+        List<SourceSplitBase> sourceSplitBases = super.snapshotState(checkpointId);
+        for (SourceSplitBase sourceSplitBase : sourceSplitBases) {
+            TableBucket tableBucket = sourceSplitBase.getTableBucket();
+            if (finishedKvSnapshotConsumeBuckets.contains(tableBucket)) {
+                continue;
+            }
+
+            if (sourceSplitBase.isHybridSnapshotLogSplit()) {
+                HybridSnapshotLogSplit hybridSnapshotLogSplit =
+                        sourceSplitBase.asHybridSnapshotLogSplit();
+                if (hybridSnapshotLogSplit.isSnapshotFinished()) {
+                    bucketsFinishedConsumeKvSnapshot.add(tableBucket);
+                }
+            }
+        }
+
+        // report finished kv snapshot consume event.
+        if (!bucketsFinishedConsumeKvSnapshot.isEmpty()) {
+            LOG.info(
+                    "reader has finished kv snapshot read for bucket: {}, checkpoint id {}",
+                    bucketsFinishedConsumeKvSnapshot,
+                    checkpointId);
+
+            context.sendSourceEventToCoordinator(

Review Comment:
   This introduces external IO operation here which I have concerns. But I don't have other solutions yet. Could you please add `TODO` comments before this code to outline the IO problem and future improvements?



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSourceBuilder.java:
##########
@@ -71,6 +72,8 @@ public class FlussSourceBuilder<OUT> {
     private Long scanPartitionDiscoveryIntervalMs;
     private OffsetsInitializer offsetsInitializer;
     private FlussDeserializationSchema<OUT> deserializationSchema;
+    private String kvSnapshotLeaseId;
+    private long kvSnapshotLeaseDurationMs;

Review Comment:
   They should have default values (use the config option defaults), otherwise, DataStream API jobs will fail without setting these 2 parameters. (backward-incompatible)



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java:
##########
@@ -856,6 +864,35 @@ public CompletableFuture<ControlledShutdownResponse> controlledShutdown(
         return response;
     }
 
+    @Override
+    public CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseRequest request) {
+        CompletableFuture<AcquireKvSnapshotLeaseResponse> response = new CompletableFuture<>();
+        eventManagerSupplier
+                .get()
+                .put(
+                        new AcquireKvSnapshotLeaseEvent(
+                                request.getLeaseId(),
+                                request.getLeaseDuration(),
+                                getAcquireKvSnapshotLeaseData(request),
+                                response));
+        return response;
+    }
+
+    @Override
+    public CompletableFuture<ReleaseKvSnapshotLeaseResponse> releaseKvSnapshotLease(
+            ReleaseKvSnapshotLeaseRequest request) {
+        CompletableFuture<ReleaseKvSnapshotLeaseResponse> response = new CompletableFuture<>();

Review Comment:
   ditto. check READ permission. and add tests



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java:
##########
@@ -633,6 +659,18 @@ public void process(CoordinatorEvent event) {
             completeFromCallable(
                     listRebalanceProgressEvent.getRespCallback(),
                     () -> processListRebalanceProgress(listRebalanceProgressEvent));
+        } else if (event instanceof AcquireKvSnapshotLeaseEvent) {
+            AcquireKvSnapshotLeaseEvent acquireKvSnapshotLeaseEvent =
+                    (AcquireKvSnapshotLeaseEvent) event;
+            completeFromCallable(
+                    acquireKvSnapshotLeaseEvent.getRespCallback(),
+                    () -> tryProcessAcquireKvSnapshotLease(acquireKvSnapshotLeaseEvent));

Review Comment:
   The lease requests don't need to access `CoordinatorContext` and should be processed in coordinator queue. The lease operation involves heavy lock operation and IO operations, I have concenrs it will make the coordinator queue much worse. We should handle the lease requests directly in `CoordinatorService`.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java:
##########
@@ -856,6 +864,35 @@ public CompletableFuture<ControlledShutdownResponse> controlledShutdown(
         return response;
     }
 
+    @Override
+    public CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+            AcquireKvSnapshotLeaseRequest request) {
+        CompletableFuture<AcquireKvSnapshotLeaseResponse> response = new CompletableFuture<>();

Review Comment:
   We should add permission check. We should throw auth exception if the session user doesn't have READ permission on the table id. 



##########
fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java:
##########
@@ -427,7 +434,12 @@ void testDiscoverPartitionsPeriodically(boolean isPrimaryKeyTable) throws Throwa
                                 streaming,
                                 null,
                                 null,
-                                workExecutor)) {
+                                workExecutor,
+                                isPrimaryKeyTable
+                                        ? new LeaseContext(
+                                                "kv_snapshot_lease1", Duration.ofDays(1).toMillis())
+                                        : new LeaseContext(null, null))) {

Review Comment:
   This conditional processing looks too verbose and make the code unreadable. As I suggested before, we can use `LeaseContext.default()` directly here. Having default values for the lease id and duration has no cost, but can improve a lot of code. 



##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -1138,4 +1171,15 @@ message PbDatabaseSummary {
   required string database_name = 1;
   required int64 created_time = 2;
   required int32 table_count = 3;
+}
+
+message PbKvSnapshotLeaseForTable {
+  required int64 table_id = 1;
+  repeated PbKvSnapshotLeaseForBucket buckets_req = 2;

Review Comment:
   rename to `bucket_snapshots`. This is also used in response, so having `req` in the field name is not a good idea.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot consumer [fluss]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2621949912


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());
+                this.consumers.put(consumer, kvSnapshotConsumer);

Review Comment:
   This field access (publicly accessible via [this expression](1)) is not protected by any monitor, but the class is annotated as @ThreadSafe.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());

Review Comment:
   This field access (publicly accessible via [this expression](1)) is not protected by any monitor, but the class is annotated as @ThreadSafe.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());
+                this.consumers.put(consumer, kvSnapshotConsumer);
+
+                initializeRefCount(kvSnapshotConsumer);
+
+                consumedBucketCount.getAndAdd(kvSnapshotConsumer.getConsumedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotConsumerNotExist(ConsumeKvSnapshotForBucket consumeKvSnapshotForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(consumeKvSnapshotForBucket)
+                                || refCount.get(consumeKvSnapshotForBucket).get() <= 0);
+    }
+
+    public void register(
+            String consumerId,
+            long expirationTime,
+            Map<Long, List<ConsumeKvSnapshotForBucket>> tableIdToRegisterBucket)
+            throws Exception {
+        ReadWriteLock lock =
+                consumerLocks.computeIfAbsent(consumerId, k -> new ReentrantReadWriteLock());
+        inWriteLock(
+                lock,
+                () -> {
+                    boolean update = consumers.containsKey(consumerId);
+                    KvSnapshotConsumer consumer;
+                    if (!update) {
+                        // set the expiration time as: current time + expirationTime
+                        consumer = new KvSnapshotConsumer(clock.milliseconds() + expirationTime);
+                        consumers.put(consumerId, consumer);
+                        LOG.info("kv snapshot consumer '" + consumerId + "' has been registered.");
+                    } else {
+                        consumer = consumers.get(consumerId);
+                    }
+
+                    for (Map.Entry<Long, List<ConsumeKvSnapshotForBucket>> entry :
+                            tableIdToRegisterBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<ConsumeKvSnapshotForBucket> buckets = entry.getValue();
+                        for (ConsumeKvSnapshotForBucket bucket : buckets) {
+                            boolean isUpdate =
+                                    consumer.registerBucket(
+                                            bucket.getTableBucket(),
+                                            bucket.getKvSnapshotId(),
+                                            numBuckets);
+                            if (!isUpdate) {
+                                consumedBucketCount.getAndIncrement();
+                                inWriteLock(
+                                        refCountLock,
+                                        () -> {
+                                            refCount.computeIfAbsent(
+                                                            bucket, k -> new AtomicInteger(0))
+                                                    .getAndIncrement();
+                                        });
+                            }
+                        }
+                    }
+
+                    if (update) {
+                        zkClient.updateKvSnapshotConsumer(consumerId, consumer);
+                    } else {
+                        zkClient.registerKvSnapshotConsumer(consumerId, consumer);
+                    }
+                });
+    }
+
+    public void unregister(
+            String consumerId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = consumerLocks.get(consumerId);

Review Comment:
   This field access (publicly accessible via [this expression](1)) is not protected by any monitor, but the class is annotated as @ThreadSafe.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());
+                this.consumers.put(consumer, kvSnapshotConsumer);
+
+                initializeRefCount(kvSnapshotConsumer);
+
+                consumedBucketCount.getAndAdd(kvSnapshotConsumer.getConsumedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotConsumerNotExist(ConsumeKvSnapshotForBucket consumeKvSnapshotForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(consumeKvSnapshotForBucket)
+                                || refCount.get(consumeKvSnapshotForBucket).get() <= 0);
+    }
+
+    public void register(
+            String consumerId,
+            long expirationTime,
+            Map<Long, List<ConsumeKvSnapshotForBucket>> tableIdToRegisterBucket)
+            throws Exception {
+        ReadWriteLock lock =
+                consumerLocks.computeIfAbsent(consumerId, k -> new ReentrantReadWriteLock());

Review Comment:
   This field access (publicly accessible via [this expression](1)) is not protected by any monitor, but the class is annotated as @ThreadSafe.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/KvSnapshotConsumerJsonSerde.java:
##########
@@ -0,0 +1,160 @@
+/*
+ * 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.server.zk.data;
+
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.core.JsonGenerator;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.fluss.utils.json.JsonDeserializer;
+import org.apache.fluss.utils.json.JsonSerializer;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/** Json serializer and deserializer for {@link KvSnapshotConsumer}. */
+public class KvSnapshotConsumerJsonSerde
+        implements JsonSerializer<KvSnapshotConsumer>, JsonDeserializer<KvSnapshotConsumer> {
+
+    public static final KvSnapshotConsumerJsonSerde INSTANCE = new KvSnapshotConsumerJsonSerde();
+
+    private static final String VERSION_KEY = "version";
+    private static final String EXPIRATION_TIME = "expiration_time";
+    private static final String TABLES = "tables";
+
+    private static final int VERSION = 1;
+
+    @Override
+    public void serialize(KvSnapshotConsumer kvSnapshotConsumer, JsonGenerator generator)
+            throws IOException {
+        generator.writeStartObject();
+        generator.writeNumberField(VERSION_KEY, VERSION);
+        generator.writeNumberField(EXPIRATION_TIME, kvSnapshotConsumer.getExpirationTime());
+
+        generator.writeFieldName(TABLES);
+        generator.writeStartObject();
+
+        // for none-partitioned table.
+        for (Map.Entry<Long, Long[]> entry :
+                kvSnapshotConsumer.getTableIdToSnapshots().entrySet()) {
+            generator.writeFieldName(entry.getKey().toString());
+            writeLongArray(generator, entry.getValue());
+        }
+
+        // for partitioned table.
+        Map<Long, Long[]> partitionToSortedSnapshotIds =
+                kvSnapshotConsumer.getPartitionIdToSnapshots();
+        for (Map.Entry<Long, Set<Long>> entry :
+                kvSnapshotConsumer.getTableIdToPartitions().entrySet()) {
+            Long tableId = entry.getKey();
+
+            generator.writeFieldName(tableId.toString());
+            generator.writeStartObject();
+
+            for (Long partitionId : entry.getValue()) {
+                Long[] snapshotIds = partitionToSortedSnapshotIds.get(partitionId);
+                if (snapshotIds != null) {
+                    generator.writeFieldName(partitionId.toString());
+                    writeLongArray(generator, snapshotIds);
+                }
+            }
+
+            generator.writeEndObject();
+        }
+        generator.writeEndObject(); // tables
+
+        generator.writeEndObject(); // root
+    }
+
+    @Override
+    public KvSnapshotConsumer deserialize(JsonNode node) {
+        long expirationTime = node.get(EXPIRATION_TIME).asLong();
+        JsonNode tablesNode = node.get(TABLES);
+
+        Map<Long, Long[]> tableIdToSnapshots = new HashMap<>();
+        Map<Long, Set<Long>> tableIdToPartitions = new HashMap<>();
+        Map<Long, Long[]> partitionIdToSnapshots = new HashMap<>();
+
+        ObjectNode tablesObj = (ObjectNode) tablesNode;
+        Iterator<Map.Entry<String, JsonNode>> tableFields = tablesObj.fields();
+        while (tableFields.hasNext()) {
+            Map.Entry<String, JsonNode> tableEntry = tableFields.next();
+            Long tableId = Long.parseLong(tableEntry.getKey());
+            JsonNode tableValue = tableEntry.getValue();
+
+            if (tableValue.isArray()) {
+                // Non-partitioned table, like: "1": [1, -1, 1, 2]
+                List<Long> snapshotIds = new ArrayList<>();
+                for (JsonNode elem : tableValue) {
+                    snapshotIds.add(elem.asLong());
+                }
+                tableIdToSnapshots.put(tableId, snapshotIds.toArray(new Long[0]));
+            } else if (tableValue.isObject()) {
+                // Partitioned table, like: "2": { "1001": [...], "1002": [...] }
+                Set<Long> partitions = new HashSet<>();
+                ObjectNode partitionObj = (ObjectNode) tableValue;
+                Iterator<Map.Entry<String, JsonNode>> partFields = partitionObj.fields();
+                while (partFields.hasNext()) {
+                    Map.Entry<String, JsonNode> partEntry = partFields.next();
+                    Long partitionId = Long.parseLong(partEntry.getKey());

Review Comment:
   The variable 'partitionId' is only assigned values of primitive type and is never 'null', but it is declared with the boxed type 'Long'.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());
+                this.consumers.put(consumer, kvSnapshotConsumer);
+
+                initializeRefCount(kvSnapshotConsumer);
+
+                consumedBucketCount.getAndAdd(kvSnapshotConsumer.getConsumedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotConsumerNotExist(ConsumeKvSnapshotForBucket consumeKvSnapshotForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(consumeKvSnapshotForBucket)
+                                || refCount.get(consumeKvSnapshotForBucket).get() <= 0);
+    }
+
+    public void register(
+            String consumerId,
+            long expirationTime,
+            Map<Long, List<ConsumeKvSnapshotForBucket>> tableIdToRegisterBucket)
+            throws Exception {
+        ReadWriteLock lock =
+                consumerLocks.computeIfAbsent(consumerId, k -> new ReentrantReadWriteLock());
+        inWriteLock(
+                lock,
+                () -> {
+                    boolean update = consumers.containsKey(consumerId);
+                    KvSnapshotConsumer consumer;
+                    if (!update) {
+                        // set the expiration time as: current time + expirationTime
+                        consumer = new KvSnapshotConsumer(clock.milliseconds() + expirationTime);
+                        consumers.put(consumerId, consumer);
+                        LOG.info("kv snapshot consumer '" + consumerId + "' has been registered.");
+                    } else {
+                        consumer = consumers.get(consumerId);
+                    }
+
+                    for (Map.Entry<Long, List<ConsumeKvSnapshotForBucket>> entry :
+                            tableIdToRegisterBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<ConsumeKvSnapshotForBucket> buckets = entry.getValue();
+                        for (ConsumeKvSnapshotForBucket bucket : buckets) {
+                            boolean isUpdate =
+                                    consumer.registerBucket(
+                                            bucket.getTableBucket(),
+                                            bucket.getKvSnapshotId(),
+                                            numBuckets);
+                            if (!isUpdate) {
+                                consumedBucketCount.getAndIncrement();
+                                inWriteLock(
+                                        refCountLock,
+                                        () -> {
+                                            refCount.computeIfAbsent(
+                                                            bucket, k -> new AtomicInteger(0))
+                                                    .getAndIncrement();
+                                        });
+                            }
+                        }
+                    }
+
+                    if (update) {
+                        zkClient.updateKvSnapshotConsumer(consumerId, consumer);
+                    } else {
+                        zkClient.registerKvSnapshotConsumer(consumerId, consumer);
+                    }
+                });
+    }
+
+    public void unregister(
+            String consumerId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = consumerLocks.get(consumerId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotConsumer consumer = consumers.get(consumerId);
+
+                    if (consumer == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = consumer.unregisterBucket(bucket);
+                            if (snapshotId != -1L) {
+                                consumedBucketCount.getAndDecrement();
+                                inWriteLock(
+                                        refCountLock,
+                                        () -> {
+                                            refCount.get(
+                                                            new ConsumeKvSnapshotForBucket(
+                                                                    bucket, snapshotId))
+                                                    .getAndDecrement();
+                                        });
+                            }
+                        }
+                    }
+
+                    if (consumer.isEmpty()) {
+                        clear(consumerId);
+                    } else {
+                        zkClient.updateKvSnapshotConsumer(consumerId, consumer);
+                    }
+                });
+    }
+
+    /**
+     * Clear kv snapshot consumer.
+     *
+     * @param consumerId the consumer id
+     * @return true if clear success, false if consumer not exist
+     */
+    public boolean clear(String consumerId) throws Exception {
+        ReadWriteLock lock = consumerLocks.get(consumerId);

Review Comment:
   This field access (publicly accessible via [this expression](1)) is not protected by any monitor, but the class is annotated as @ThreadSafe.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/KvSnapshotConsumer.java:
##########
@@ -0,0 +1,311 @@
+/*
+ * 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.server.zk.data;
+
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.utils.MapUtils;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+/** The zkNode data of kv snapshot consumer. */
+@NotThreadSafe
+public class KvSnapshotConsumer {
+    private final long expirationTime;
+
+    /**
+     * mapping from table id to sorted snapshot ids for none-partitioned table, the value is a list
+     * of consumed snapshot ids sorted by bucket id.
+     */
+    private final Map<Long, Long[]> tableIdToSnapshots;
+
+    /** tableId to partition ids for partitioned table. */
+    private final Map<Long, Set<Long>> tableIdToPartitions;
+
+    /**
+     * mapping from partition id to sorted snapshot ids for partitioned table, the value is a list
+     * of consumed snapshot ids sorted by bucket id.
+     */
+    private final Map<Long, Long[]> partitionIdToSnapshots;
+
+    public KvSnapshotConsumer(long expirationTime) {
+        this(
+                expirationTime,
+                MapUtils.newConcurrentHashMap(),
+                MapUtils.newConcurrentHashMap(),
+                MapUtils.newConcurrentHashMap());
+    }
+
+    public KvSnapshotConsumer(
+            long expirationTime,
+            Map<Long, Long[]> tableIdToSnapshots,
+            Map<Long, Set<Long>> tableIdToPartitions,
+            Map<Long, Long[]> partitionIdToSnapshots) {
+        this.expirationTime = expirationTime;
+        this.tableIdToSnapshots = tableIdToSnapshots;
+        this.tableIdToPartitions = tableIdToPartitions;
+        this.partitionIdToSnapshots = partitionIdToSnapshots;
+    }
+
+    public long getExpirationTime() {
+        return expirationTime;
+    }
+
+    public Map<Long, Long[]> getTableIdToSnapshots() {
+        return tableIdToSnapshots;
+    }
+
+    public Map<Long, Set<Long>> getTableIdToPartitions() {
+        return tableIdToPartitions;
+    }
+
+    public Map<Long, Long[]> getPartitionIdToSnapshots() {
+        return partitionIdToSnapshots;
+    }
+
+    /**
+     * Register a bucket to the consumer.
+     *
+     * @param tableBucket table bucket
+     * @param snapshotId snapshot id
+     * @param bucketNum bucket number of this table or partition
+     * @return true if this operation is update, false if this operation is insert
+     */
+    public boolean registerBucket(TableBucket tableBucket, long snapshotId, int bucketNum) {
+        Long[] bucketIndex;
+        Long partitionId = tableBucket.getPartitionId();
+        long tableId = tableBucket.getTableId();
+        int bucketId = tableBucket.getBucket();
+        if (partitionId == null) {
+            // For none-partitioned table.
+            bucketIndex =
+                    tableIdToSnapshots.computeIfAbsent(
+                            tableId,
+                            k -> {
+                                Long[] array = new Long[bucketNum];
+                                Arrays.fill(array, -1L);
+                                return array;
+                            });
+        } else {
+            // For partitioned table.
+
+            // first add partition to table.
+            Set<Long> partitions =
+                    tableIdToPartitions.computeIfAbsent(tableId, k -> new HashSet<>());
+            partitions.add(partitionId);
+
+            // then add bucket to partition.
+            bucketIndex =
+                    partitionIdToSnapshots.computeIfAbsent(
+                            partitionId,
+                            k -> {
+                                Long[] array = new Long[bucketNum];
+                                Arrays.fill(array, -1L);
+                                return array;
+                            });
+        }
+
+        if (bucketIndex.length != bucketNum) {
+            throw new IllegalArgumentException(
+                    "The input bucket number is not equal to the bucket number of the table.");
+        }
+        boolean isUpdate = bucketIndex[bucketId] != -1L;
+        bucketIndex[bucketId] = snapshotId;
+        return isUpdate;
+    }
+
+    /**
+     * Unregister a bucket from the consumer.
+     *
+     * @param tableBucket table bucket
+     * @return the snapshot id of the unregistered bucket
+     */
+    public long unregisterBucket(TableBucket tableBucket) {
+        Long[] bucketIndex;
+        long tableId = tableBucket.getTableId();
+        Long partitionId = tableBucket.getPartitionId();
+        int bucketId = tableBucket.getBucket();
+        if (partitionId == null) {
+            // For none-partitioned table.
+            bucketIndex = tableIdToSnapshots.get(tableId);
+        } else {
+            // For partitioned table.
+            bucketIndex = partitionIdToSnapshots.get(partitionId);
+        }
+
+        Long snapshotId = -1L;
+        if (bucketIndex != null) {
+            snapshotId = bucketIndex[bucketId];
+            bucketIndex[bucketId] = -1L;
+
+            boolean needRemove = true;
+            for (Long bucket : bucketIndex) {
+                if (bucket != -1L) {
+                    needRemove = false;
+                    break;
+                }
+            }
+
+            if (needRemove) {
+                if (partitionId == null) {
+                    tableIdToSnapshots.remove(tableId);
+                } else {
+                    tableIdToPartitions.remove(tableId);
+                    partitionIdToSnapshots.remove(partitionId);
+                }
+            }
+        }
+        return snapshotId;
+    }
+
+    public boolean isEmpty() {
+        return tableIdToSnapshots.isEmpty()
+                && tableIdToPartitions.isEmpty()
+                && partitionIdToSnapshots.isEmpty();
+    }
+
+    public int getConsumedSnapshotCount() {
+        int count = 0;
+        for (Long[] buckets : tableIdToSnapshots.values()) {
+            for (Long bucket : buckets) {
+                if (bucket != -1L) {
+                    count++;
+                }
+            }
+        }
+
+        for (Long[] buckets : partitionIdToSnapshots.values()) {
+            for (Long bucket : buckets) {
+                if (bucket != -1L) {
+                    count++;
+                }
+            }
+        }
+        return count;
+    }
+
+    @Override
+    public String toString() {
+        String tableSnapshotsStr = formatLongArrayMap(tableIdToSnapshots);
+        String partitionSnapshotsStr = formatLongArrayMap(partitionIdToSnapshots);
+
+        return "KvSnapshotConsumer{"
+                + "expirationTime="
+                + expirationTime
+                + ", tableIdToSnapshots="
+                + tableSnapshotsStr
+                + ", tableIdToPartitions="
+                + tableIdToPartitions
+                + ", partitionIdToSnapshots="
+                + partitionSnapshotsStr
+                + '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {

Review Comment:
   This 'equals()' method does not check argument type.



##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/event/FinishedKvSnapshotConsumeEvent.java:
##########
@@ -0,0 +1,74 @@
+/*
+ * 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.flink.source.event;
+
+import org.apache.fluss.metadata.TableBucket;
+
+import org.apache.flink.api.connector.source.SourceEvent;
+
+import java.util.Objects;
+import java.util.Set;
+
+/** SourceEvent used to represent a Fluss table bucket has complete consume kv snapshot. */
+public class FinishedKvSnapshotConsumeEvent implements SourceEvent {

Review Comment:
   Class FinishedKvSnapshotConsumeEvent overrides [hashCode](1) but not equals.



##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/KvSnapshotConsumerJsonSerde.java:
##########
@@ -0,0 +1,160 @@
+/*
+ * 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.server.zk.data;
+
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.core.JsonGenerator;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.fluss.utils.json.JsonDeserializer;
+import org.apache.fluss.utils.json.JsonSerializer;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/** Json serializer and deserializer for {@link KvSnapshotConsumer}. */
+public class KvSnapshotConsumerJsonSerde
+        implements JsonSerializer<KvSnapshotConsumer>, JsonDeserializer<KvSnapshotConsumer> {
+
+    public static final KvSnapshotConsumerJsonSerde INSTANCE = new KvSnapshotConsumerJsonSerde();
+
+    private static final String VERSION_KEY = "version";
+    private static final String EXPIRATION_TIME = "expiration_time";
+    private static final String TABLES = "tables";
+
+    private static final int VERSION = 1;
+
+    @Override
+    public void serialize(KvSnapshotConsumer kvSnapshotConsumer, JsonGenerator generator)
+            throws IOException {
+        generator.writeStartObject();
+        generator.writeNumberField(VERSION_KEY, VERSION);
+        generator.writeNumberField(EXPIRATION_TIME, kvSnapshotConsumer.getExpirationTime());
+
+        generator.writeFieldName(TABLES);
+        generator.writeStartObject();
+
+        // for none-partitioned table.
+        for (Map.Entry<Long, Long[]> entry :
+                kvSnapshotConsumer.getTableIdToSnapshots().entrySet()) {
+            generator.writeFieldName(entry.getKey().toString());
+            writeLongArray(generator, entry.getValue());
+        }
+
+        // for partitioned table.
+        Map<Long, Long[]> partitionToSortedSnapshotIds =
+                kvSnapshotConsumer.getPartitionIdToSnapshots();
+        for (Map.Entry<Long, Set<Long>> entry :
+                kvSnapshotConsumer.getTableIdToPartitions().entrySet()) {
+            Long tableId = entry.getKey();
+
+            generator.writeFieldName(tableId.toString());
+            generator.writeStartObject();
+
+            for (Long partitionId : entry.getValue()) {
+                Long[] snapshotIds = partitionToSortedSnapshotIds.get(partitionId);
+                if (snapshotIds != null) {
+                    generator.writeFieldName(partitionId.toString());
+                    writeLongArray(generator, snapshotIds);
+                }
+            }
+
+            generator.writeEndObject();
+        }
+        generator.writeEndObject(); // tables
+
+        generator.writeEndObject(); // root
+    }
+
+    @Override
+    public KvSnapshotConsumer deserialize(JsonNode node) {
+        long expirationTime = node.get(EXPIRATION_TIME).asLong();
+        JsonNode tablesNode = node.get(TABLES);
+
+        Map<Long, Long[]> tableIdToSnapshots = new HashMap<>();
+        Map<Long, Set<Long>> tableIdToPartitions = new HashMap<>();
+        Map<Long, Long[]> partitionIdToSnapshots = new HashMap<>();
+
+        ObjectNode tablesObj = (ObjectNode) tablesNode;
+        Iterator<Map.Entry<String, JsonNode>> tableFields = tablesObj.fields();
+        while (tableFields.hasNext()) {
+            Map.Entry<String, JsonNode> tableEntry = tableFields.next();
+            Long tableId = Long.parseLong(tableEntry.getKey());

Review Comment:
   The variable 'tableId' is only assigned values of primitive type and is never 'null', but it is declared with the boxed type 'Long'.



##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());
+                this.consumers.put(consumer, kvSnapshotConsumer);
+
+                initializeRefCount(kvSnapshotConsumer);
+
+                consumedBucketCount.getAndAdd(kvSnapshotConsumer.getConsumedSnapshotCount());
+            }
+        }
+    }
+
+    public boolean snapshotConsumerNotExist(ConsumeKvSnapshotForBucket consumeKvSnapshotForBucket) {
+        return inReadLock(
+                refCountLock,
+                () ->
+                        !refCount.containsKey(consumeKvSnapshotForBucket)
+                                || refCount.get(consumeKvSnapshotForBucket).get() <= 0);
+    }
+
+    public void register(
+            String consumerId,
+            long expirationTime,
+            Map<Long, List<ConsumeKvSnapshotForBucket>> tableIdToRegisterBucket)
+            throws Exception {
+        ReadWriteLock lock =
+                consumerLocks.computeIfAbsent(consumerId, k -> new ReentrantReadWriteLock());
+        inWriteLock(
+                lock,
+                () -> {
+                    boolean update = consumers.containsKey(consumerId);
+                    KvSnapshotConsumer consumer;
+                    if (!update) {
+                        // set the expiration time as: current time + expirationTime
+                        consumer = new KvSnapshotConsumer(clock.milliseconds() + expirationTime);
+                        consumers.put(consumerId, consumer);
+                        LOG.info("kv snapshot consumer '" + consumerId + "' has been registered.");
+                    } else {
+                        consumer = consumers.get(consumerId);
+                    }
+
+                    for (Map.Entry<Long, List<ConsumeKvSnapshotForBucket>> entry :
+                            tableIdToRegisterBucket.entrySet()) {
+                        Long tableId = entry.getKey();
+                        TableInfo tableInfo = coordinatorContext.getTableInfoById(tableId);
+                        int numBuckets = tableInfo.getNumBuckets();
+                        List<ConsumeKvSnapshotForBucket> buckets = entry.getValue();
+                        for (ConsumeKvSnapshotForBucket bucket : buckets) {
+                            boolean isUpdate =
+                                    consumer.registerBucket(
+                                            bucket.getTableBucket(),
+                                            bucket.getKvSnapshotId(),
+                                            numBuckets);
+                            if (!isUpdate) {
+                                consumedBucketCount.getAndIncrement();
+                                inWriteLock(
+                                        refCountLock,
+                                        () -> {
+                                            refCount.computeIfAbsent(
+                                                            bucket, k -> new AtomicInteger(0))
+                                                    .getAndIncrement();
+                                        });
+                            }
+                        }
+                    }
+
+                    if (update) {
+                        zkClient.updateKvSnapshotConsumer(consumerId, consumer);
+                    } else {
+                        zkClient.registerKvSnapshotConsumer(consumerId, consumer);
+                    }
+                });
+    }
+
+    public void unregister(
+            String consumerId, Map<Long, List<TableBucket>> tableIdToUnregisterBucket)
+            throws Exception {
+        ReadWriteLock lock = consumerLocks.get(consumerId);
+        if (lock == null) {
+            return;
+        }
+
+        inWriteLock(
+                lock,
+                () -> {
+                    KvSnapshotConsumer consumer = consumers.get(consumerId);
+
+                    if (consumer == null) {
+                        return;
+                    }
+
+                    for (Map.Entry<Long, List<TableBucket>> entry :
+                            tableIdToUnregisterBucket.entrySet()) {
+                        List<TableBucket> buckets = entry.getValue();
+                        for (TableBucket bucket : buckets) {
+                            long snapshotId = consumer.unregisterBucket(bucket);
+                            if (snapshotId != -1L) {
+                                consumedBucketCount.getAndDecrement();
+                                inWriteLock(
+                                        refCountLock,
+                                        () -> {
+                                            refCount.get(
+                                                            new ConsumeKvSnapshotForBucket(
+                                                                    bucket, snapshotId))
+                                                    .getAndDecrement();
+                                        });
+                            }
+                        }
+                    }
+
+                    if (consumer.isEmpty()) {
+                        clear(consumerId);
+                    } else {
+                        zkClient.updateKvSnapshotConsumer(consumerId, consumer);
+                    }
+                });
+    }
+
+    /**
+     * Clear kv snapshot consumer.
+     *
+     * @param consumerId the consumer id
+     * @return true if clear success, false if consumer not exist
+     */
+    public boolean clear(String consumerId) throws Exception {
+        ReadWriteLock lock = consumerLocks.get(consumerId);
+        if (lock == null) {
+            return false;
+        }
+
+        boolean exist =
+                inWriteLock(
+                        lock,
+                        () -> {
+                            KvSnapshotConsumer kvSnapshotConsumer = consumers.remove(consumerId);
+                            if (kvSnapshotConsumer == null) {
+                                return false;
+                            }
+
+                            clearRefCount(kvSnapshotConsumer);
+                            zkClient.deleteKvSnapshotConsumer(consumerId);
+
+                            LOG.info("kv snapshot consumer '" + consumerId + "' has been cleared.");
+                            return true;
+                        });
+
+        consumerLocks.remove(consumerId);

Review Comment:
   This field access (publicly accessible via [this expression](1)) is not protected by any monitor, but the class is annotated as @ThreadSafe.



##########
fluss-common/src/main/java/org/apache/fluss/metadata/ConsumeKvSnapshotForBucket.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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.metadata;
+
+import java.util.Objects;
+
+/** An entity for consume kv snapshot for bucket. */
+public class ConsumeKvSnapshotForBucket {
+    private final TableBucket tableBucket;
+    private final long kvSnapshotId;
+
+    public ConsumeKvSnapshotForBucket(TableBucket tableBucket, long kvSnapshotId) {
+        this.tableBucket = tableBucket;
+        this.kvSnapshotId = kvSnapshotId;
+    }
+
+    public TableBucket getTableBucket() {
+        return tableBucket;
+    }
+
+    public long getKvSnapshotId() {
+        return kvSnapshotId;
+    }
+
+    @Override
+    public String toString() {
+        return "ConsumeKvSnapshotForBucket{"
+                + "tableBucket="
+                + tableBucket
+                + ", kvSnapshotId="
+                + kvSnapshotId
+                + '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {

Review Comment:
   This 'equals()' method does not check argument type.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2703144815


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/LeaseContext.java:
##########
@@ -0,0 +1,79 @@
+/*
+ * 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.flink.source.reader;
+
+import javax.annotation.Nullable;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+/** Context for lease. */
+public class LeaseContext implements Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    // kv snapshot lease id. null for log table.
+    private final @Nullable String kvSnapshotLeaseId;
+
+    // kv snapshot lease duration. null for log table.
+    private final @Nullable Long kvSnapshotLeaseDurationMs;
+
+    public LeaseContext(
+            @Nullable String kvSnapshotLeaseId, @Nullable Long kvSnapshotLeaseDurationMs) {
+        this.kvSnapshotLeaseId = kvSnapshotLeaseId;
+        this.kvSnapshotLeaseDurationMs = kvSnapshotLeaseDurationMs;
+    }
+
+    public @Nullable String getKvSnapshotLeaseId() {
+        return kvSnapshotLeaseId;
+    }
+
+    public @Nullable Long getKvSnapshotLeaseDurationMs() {

Review Comment:
   nullable for logTable. I will reserve this.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2658790889


##########
fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java:
##########
@@ -53,6 +53,10 @@ public class MetricNames {
     public static final String KV_ALL_SNAPSHOT_SIZE = "allKvSnapshotSize";
     public static final String SERVER_PHYSICAL_STORAGE_REMOTE_KV_SIZE = "remoteKvSize";
 
+    // for kv snapshot lease.
+    public static final String KV_SNAPSHOT_LEASE_COUNT = "kvSnapshotLeaseCount";

Review Comment:
   Sure, I was thinking the same thing. I'll add a TODO for now, and this will be implemented in the next PR since this PR is already too large. Trace by: https://github.com/apache/fluss/issues/2297



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot consumer [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2621986860


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotConsumerManager.java:
##########
@@ -0,0 +1,425 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.ConsumeKvSnapshotForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.KvSnapshotConsumer;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot consumer register/unregister/clear. */
+@ThreadSafe
+public class KvSnapshotConsumerManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotConsumerManager.class);
+
+    private final ZooKeeperClient zkClient;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> consumerLocks = MapUtils.newConcurrentHashMap();
+    /** Consumer id to Consumer. */
+    @GuardedBy("consumerLocks")
+    private final Map<String, KvSnapshotConsumer> consumers;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * ConsumeKvSnapshotForBucket to the ref count, which means this table bucket + snapshotId has
+     * been consumed by how many consumers.
+     */
+    @GuardedBy("refCountLock")
+    private final Map<ConsumeKvSnapshotForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger consumedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                zkClient,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-consumer-cleaner")),
+                clock,
+                coordinatorMetricGroup);
+    }
+
+    @VisibleForTesting
+    public KvSnapshotConsumerManager(
+            Configuration conf,
+            ZooKeeperClient zkClient,
+            CoordinatorContext coordinatorContext,
+            ScheduledExecutorService scheduledExecutor,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this.zkClient = zkClient;
+        this.conf = conf;
+        this.scheduledExecutor = scheduledExecutor;
+        this.coordinatorContext = coordinatorContext;
+        this.clock = clock;
+        this.consumers = MapUtils.newConcurrentHashMap();
+
+        registerMetrics(coordinatorMetricGroup);
+    }
+
+    public void start() {
+        scheduledExecutor.scheduleWithFixedDelay(
+                this::expireConsumers,
+                0L,
+                conf.get(ConfigOptions.KV_SNAPSHOT_CONSUMER_EXPIRATION_CHECK_INTERVAL).toMillis(),
+                TimeUnit.MILLISECONDS);
+    }
+
+    public void initialize() throws Exception {
+        List<String> consumers = zkClient.getKvSnapshotConsumerList();
+        for (String consumer : consumers) {
+            Optional<KvSnapshotConsumer> kvSnapshotConsumerOpt =
+                    zkClient.getKvSnapshotConsumer(consumer);
+            if (kvSnapshotConsumerOpt.isPresent()) {
+                KvSnapshotConsumer kvSnapshotConsumer = kvSnapshotConsumerOpt.get();
+                this.consumerLocks.put(consumer, new ReentrantReadWriteLock());

Review Comment:
   Hi, copilot. I think this variable is thread safe as I already announce it as ConcurrentHashMap.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2658795116


##########
fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerde.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.server.zk.data.lease;
+
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.core.JsonGenerator;
+import org.apache.fluss.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
+import org.apache.fluss.utils.json.JsonDeserializer;
+import org.apache.fluss.utils.json.JsonSerializer;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Json serializer and deserializer for {@link KvSnapshotLeaseMetadata}. */
+public class KvSnapshotLeaseMetadataJsonSerde
+        implements JsonSerializer<KvSnapshotLeaseMetadata>,
+                JsonDeserializer<KvSnapshotLeaseMetadata> {
+
+    public static final KvSnapshotLeaseMetadataJsonSerde INSTANCE =
+            new KvSnapshotLeaseMetadataJsonSerde();
+
+    private static final String VERSION_KEY = "version";
+    private static final String EXPIRATION_TIME = "expiration_time";
+    private static final String TABLES = "tables";
+    private static final String TABLE_ID = "table_id";
+    private static final String KV_SNAPSHOT_PATH = "lease_metadata_path";
+
+    private static final int VERSION = 1;
+
+    @Override
+    public void serialize(KvSnapshotLeaseMetadata kvSnapshotLeaseMetadata, JsonGenerator generator)

Review Comment:
   This implementation is indeed quite old—it was originally based on Flink's serialization. We could consider refactoring it, but perhaps not in this PR?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2658792249


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java:
##########
@@ -0,0 +1,36 @@
+/*
+ * 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.flink.procedure;
+
+import org.apache.flink.table.annotation.ArgumentHint;
+import org.apache.flink.table.annotation.DataTypeHint;
+import org.apache.flink.table.annotation.ProcedureHint;
+import org.apache.flink.table.procedure.ProcedureContext;
+
+/** Procedure to drop kv snapshot lease. */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {

Review Comment:
   Do you intend to rename this API to `release`, or just update the comment? According to the FIP, this API is indeed named `dropKvSnapshotLease`, and its behavior genuinely performs a "drop" operation.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2702409837


##########
fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java:
##########
@@ -51,6 +57,11 @@ public class FlinkSourceReader<OUT>
         extends SingleThreadMultiplexSourceReaderBaseAdapter<
                 RecordAndPos, OUT, SourceSplitBase, SourceSplitState> {
 
+    private static final Logger LOG = LoggerFactory.getLogger(FlinkSourceReader.class);
+
+    /** the tableBuckets ignore to send FinishedKvSnapshotConsumeEvent as it already sending. */
+    private final Set<TableBucket> ignoreBuckets;

Review Comment:
   change to `finishedKvSnapshotConsumeBuckets`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2703007776


##########
fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java:
##########
@@ -412,6 +415,56 @@ CompletableFuture<Void> dropPartition(
     CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
+    /**
+     * Acquires a lease for specific KV snapshots of the given tableBuckets asynchronously.
+     *
+     * <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
+     * #releaseKvSnapshotLease} to release the lock early when reading is finished.
+     *
+     * <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.
+     * </ul>
+     *
+     * @param leaseId The unique ID for this lease session (usually a UUID generated by client).
+     * @param snapshotIds The snapshots to lease, 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<AcquireKvSnapshotLeaseResult> acquireKvSnapshotLease(
+            String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration);
+
+    /**
+     * Releases the lease for specific tableBuckets asynchronously.
+     *
+     * <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.
+     *
+     * <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);
+
+    /**
+     * Releases 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 leaseId The lease id to release.
+     */
+    CompletableFuture<Void> dropKvSnapshotLease(String leaseId);

Review Comment:
   Ok, I will also remove `DropKvSnapshotLeaseRequest/DropKvSnapshotLeaseResponse`. Instead the `releaseAllKvSnapshotLeases()` and `releaseKvSnapshotLease` will use `ReleaseKvSnapshotLeaseRequest/ReleaseKvSnapshotLeaseResponse` together.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778938199


##########
fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java:
##########
@@ -0,0 +1,459 @@
+/*
+ * 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.server.coordinator;
+
+import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.KvSnapshotLeaseForBucket;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.utils.MapUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
+
+import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock;
+import static org.apache.fluss.utils.concurrent.LockUtils.inWriteLock;
+
+/** A manager to manage kv snapshot lease acquire, renew, release and drop. */
+@ThreadSafe
+public class KvSnapshotLeaseManager {
+    private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+
+    private final KvSnapshotLeaseMetadataManager metadataManager;
+    private final CoordinatorContext coordinatorContext;
+    private final Clock clock;
+    private final ScheduledExecutorService scheduledExecutor;
+    private final Configuration conf;
+
+    private final Map<String, ReadWriteLock> leaseLocks = MapUtils.newConcurrentHashMap();
+    /** lease id to kv snapshot lease. */
+    @GuardedBy("leaseLocks")
+    private final Map<String, KvSnapshotLease> kvSnapshotLeaseMap;
+
+    private final ReadWriteLock refCountLock = new ReentrantReadWriteLock();
+
+    /**
+     * KvSnapshotLeaseForBucket to the ref count, which means this table bucket + snapshotId has
+     * been leased by how many lease id.
+     */
+    private final Map<KvSnapshotLeaseForBucket, AtomicInteger> refCount =
+            MapUtils.newConcurrentHashMap();
+
+    /** For metrics. */
+    private final AtomicInteger leasedBucketCount = new AtomicInteger(0);
+
+    public KvSnapshotLeaseManager(
+            Configuration conf,
+            KvSnapshotLeaseMetadataManager metadataManager,
+            CoordinatorContext coordinatorContext,
+            Clock clock,
+            CoordinatorMetricGroup coordinatorMetricGroup) {
+        this(
+                conf,
+                metadataManager,
+                coordinatorContext,
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("kv-snapshot-lease-cleaner")),

Review Comment:
   I haven't changed this yet because `KvSnapshotLeaseManager` is managed within `CoordinatorServer` (`CoordinatorProcessor` still needs it). Therefore, it needs to be passed here. However, the `ioExecutor` in `CoordinatorServer` is not a `ScheduledExecutorService`, so I'm not sure whether we need to abstract out a separate `ScheduledExecutorService`? if need, i can new a shared thread pool.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [kv] Support kv snapshot lease [fluss]

Posted by "swuferhong (via GitHub)" <gi...@apache.org>.
swuferhong commented on code in PR #2179:
URL: https://github.com/apache/fluss/pull/2179#discussion_r2778429053


##########
fluss-rpc/src/main/proto/FlussApi.proto:
##########
@@ -378,6 +378,29 @@ message GetKvSnapshotMetadataResponse {
   repeated PbRemotePathAndLocalFile snapshot_files = 2;
 }
 
+message AcquireKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  required int64 lease_duration = 2;
+  repeated PbKvSnapshotLeaseForTable table_lease_req = 3;
+}
+
+message AcquireKvSnapshotLeaseResponse {
+  repeated PbKvSnapshotLeaseForTable tables_lease_res = 3;
+}
+
+message ReleaseKvSnapshotLeaseRequest {
+  required string lease_id = 1;
+  repeated PbTable release_tables = 2;

Review Comment:
   I will introduce dropKvSnapshotLease API instead of reusing this api.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@fluss.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org