address's Yang Wang's comments

swuferhong · swuferhong · commit d59571dfb771 · 2026-01-12T11:44:03.000+08:00
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java
@@ -72,8 +72,8 @@
 
 import java.util.Collection;
 import java.util.List;
-import java.util.Optional;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
@@ -456,12 +456,12 @@ CompletableFuture<Void> releaseKvSnapshotLease(
             String leaseId, Set<TableBucket> bucketsToRelease);
 
     /**
-     * Drops the entire lease asynchronously.
+     * 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 drop.
+     * @param leaseId The lease id to release.
      */
     CompletableFuture<Void> dropKvSnapshotLease(String leaseId);
 
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/batch/KvSnapshotBatchScannerITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/batch/KvSnapshotBatchScannerITCase.java
@@ -37,7 +37,7 @@
 import org.apache.fluss.row.encode.KeyEncoder;
 import org.apache.fluss.server.zk.ZooKeeperClient;
 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.KvSnapshotLeaseMetadataManager;
 import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
 import org.apache.fluss.types.DataTypes;
 
@@ -205,8 +205,8 @@ public void testKvSnapshotLease() throws Exception {
 
         ZooKeeperClient zkClient = FLUSS_CLUSTER_EXTENSION.getZooKeeperClient();
         String remoteDataDir = FLUSS_CLUSTER_EXTENSION.getRemoteDataDir();
-        KvSnapshotLeaseMetadataHelper metadataHelper =
-                new KvSnapshotLeaseMetadataHelper(zkClient, remoteDataDir);
+        KvSnapshotLeaseMetadataManager metadataHelper =
+                new KvSnapshotLeaseMetadataManager(zkClient, remoteDataDir);
 
         assertThat(zkClient.getKvSnapshotLeasesList()).isEmpty();
 
@@ -346,7 +346,7 @@ private void waitUntilAllSnapshotFinished(Set<TableBucket> tableBuckets, long sn
     }
 
     private void checkKvSnapshotLeaseEquals(
-            KvSnapshotLeaseMetadataHelper metadataHelper,
+            KvSnapshotLeaseMetadataManager metadataHelper,
             String leaseId,
             long tableId,
             Long[] expectedBucketIndex)
diff --git a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java b/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
@@ -54,6 +54,7 @@ public class MetricNames {
     public static final String SERVER_PHYSICAL_STORAGE_REMOTE_KV_SIZE = "remoteKvSize";
 
     // for kv snapshot lease.
+    // TODO implemented it at the table level. Trace by: https://github.com/apache/fluss/issues/2297
     public static final String KV_SNAPSHOT_LEASE_COUNT = "kvSnapshotLeaseCount";
     public static final String LEASED_KV_SNAPSHOT_COUNT = "leasedKvSnapshotCount";
 
diff --git a/fluss-common/src/test/java/org/apache/fluss/record/TestData.java b/fluss-common/src/test/java/org/apache/fluss/record/TestData.java
@@ -97,6 +97,10 @@ public final class TestData {
                     currentMillis);
 
     // for log table / partition table
+    public static final TablePath PARTITION_TABLE_PATH =
+            new TablePath("test_db_1", "test_partition_table");
+    public static final long PARTITION_TABLE_ID = 150008L;
+
     public static final TableDescriptor DATA1_PARTITIONED_TABLE_DESCRIPTOR =
             TableDescriptor.builder()
                     .schema(DATA1_SCHEMA)
@@ -107,6 +111,16 @@ public final class TestData {
                             ConfigOptions.TABLE_AUTO_PARTITION_TIME_UNIT,
                             AutoPartitionTimeUnit.YEAR)
                     .build();
+
+    public static final TableInfo PARTITION_TABLE_INFO =
+            TableInfo.of(
+                    PARTITION_TABLE_PATH,
+                    PARTITION_TABLE_ID,
+                    1,
+                    DATA1_PARTITIONED_TABLE_DESCRIPTOR,
+                    System.currentTimeMillis(),
+                    System.currentTimeMillis());
+
     public static final PhysicalTablePath DATA1_PHYSICAL_TABLE_PATH_PA_2024 =
             PhysicalTablePath.of(DATA1_TABLE_PATH, "2024");
 
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java
@@ -22,7 +22,7 @@
 import org.apache.flink.table.annotation.ProcedureHint;
 import org.apache.flink.table.procedure.ProcedureContext;
 
-/** Procedure to drop kv snapshot lease. */
+/** Procedure to release kv snapshot lease. */
 public class DropKvSnapshotLeaseProcedure extends ProcedureBase {
 
     @ProcedureHint(
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/FlussSourceEnumeratorStateSerializer.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/FlussSourceEnumeratorStateSerializer.java
@@ -199,7 +199,8 @@ private SourceEnumeratorState deserializeV0(byte[] serialized) throws IOExceptio
         return new SourceEnumeratorState(
                 assignBucketAndPartitions.f0,
                 assignBucketAndPartitions.f1,
-                remainingHybridLakeFlussSplits, leaseContext);
+                remainingHybridLakeFlussSplits,
+                leaseContext);
     }
 
     private SourceEnumeratorState deserializeV1(byte[] serialized) throws IOException {
@@ -213,14 +214,14 @@ private SourceEnumeratorState deserializeV1(byte[] serialized) throws IOExceptio
         // this logic no longer depends on the lakeSource flag. This unconditional
         // deserialization is the intended behavior change compared to VERSION_0.
 
-
         // deserialize lease context
         LeaseContext leaseContext = deserializeLeaseContext(in);
 
         return new SourceEnumeratorState(
                 assignBucketAndPartitions.f0,
                 assignBucketAndPartitions.f1,
-                remainingHybridLakeFlussSplits, leaseContext);
+                remainingHybridLakeFlussSplits,
+                leaseContext);
     }
 
     private Tuple2<Set<TableBucket>, Map<Long, String>> deserializeAssignBucketAndPartitions(
diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/procedure/FlinkProcedureITCase.java
@@ -17,14 +17,11 @@
 
 package org.apache.fluss.flink.procedure;
 
-import org.apache.fluss.client.Connection;
-import org.apache.fluss.client.ConnectionFactory;
-import org.apache.fluss.client.admin.Admin;
-import org.apache.fluss.cluster.rebalance.RebalanceProgress;
 import org.apache.fluss.client.Connection;
 import org.apache.fluss.client.ConnectionFactory;
 import org.apache.fluss.client.admin.Admin;
 import org.apache.fluss.client.metadata.KvSnapshots;
+import org.apache.fluss.cluster.rebalance.RebalanceProgress;
 import org.apache.fluss.cluster.rebalance.ServerTag;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.Configuration;
@@ -34,7 +31,6 @@
 import org.apache.fluss.metadata.DataLakeFormat;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.row.InternalRow;
-import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.server.testutils.FlussClusterExtension;
 import org.apache.fluss.server.zk.ZooKeeperClient;
 import org.apache.fluss.server.zk.data.ServerTags;
@@ -46,7 +42,6 @@
 import org.apache.flink.util.CollectionUtil;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
-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;
@@ -64,10 +59,10 @@
 import static org.apache.fluss.cluster.rebalance.ServerTag.PERMANENT_OFFLINE;
 import static org.apache.fluss.flink.source.testutils.FlinkRowAssertionsUtils.assertResultsIgnoreOrder;
 import static org.apache.fluss.flink.utils.FlinkTestBase.writeRows;
-import static org.apache.fluss.testutils.DataTestUtils.row;
-import static org.apache.fluss.testutils.common.CommonTestUtils.waitUntil;
 import static org.apache.fluss.server.testutils.FlussClusterExtension.BUILTIN_DATABASE;
+import static org.apache.fluss.testutils.DataTestUtils.row;
 import static org.apache.fluss.testutils.common.CommonTestUtils.retry;
+import static org.apache.fluss.testutils.common.CommonTestUtils.waitUntil;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/state/SourceEnumeratorStateSerializerTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/state/SourceEnumeratorStateSerializerTest.java
@@ -27,6 +27,7 @@
 import org.apache.fluss.lake.source.TestingLakeSplit;
 import org.apache.fluss.metadata.TableBucket;
 
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
@@ -117,7 +118,8 @@ void testV0Compatibility() throws Exception {
         assignedPartitions.put(1L, "partition1");
         assignedPartitions.put(2L, "partition2");
         SourceEnumeratorState sourceEnumeratorState =
-                new SourceEnumeratorState(assignedBuckets, assignedPartitions, null);
+                new SourceEnumeratorState(
+                        assignedBuckets, assignedPartitions, null, new LeaseContext(null, null));
         byte[] serialized = serializer.serializeV0(sourceEnumeratorState);
 
         // then deserialize
@@ -134,7 +136,10 @@ void testV0Compatibility() throws Exception {
         remainingHybridLakeFlussSplits.add(logSplit);
         sourceEnumeratorState =
                 new SourceEnumeratorState(
-                        assignedBuckets, assignedPartitions, remainingHybridLakeFlussSplits);
+                        assignedBuckets,
+                        assignedPartitions,
+                        remainingHybridLakeFlussSplits,
+                        new LeaseContext(null, null));
 
         serialized = serializer.serializeV0(sourceEnumeratorState);
 
@@ -155,7 +160,8 @@ void testInconsistentLakeSourceSerde() throws Exception {
         assignedPartitions.put(1L, "partition1");
         assignedPartitions.put(2L, "partition2");
         SourceEnumeratorState sourceEnumeratorState =
-                new SourceEnumeratorState(assignedBuckets, assignedPartitions, null);
+                new SourceEnumeratorState(
+                        assignedBuckets, assignedPartitions, null, new LeaseContext(null, null));
         byte[] serialized = serializer.serialize(sourceEnumeratorState);
 
         // test deserialize with nonnull lake source
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
@@ -60,8 +60,8 @@
 import org.apache.fluss.rpc.messages.ListRebalanceProgressResponse;
 import org.apache.fluss.rpc.messages.PbCommitLakeTableSnapshotRespForTable;
 import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForBucket;
-import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseResponse;
 import org.apache.fluss.rpc.messages.RebalanceResponse;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseResponse;
 import org.apache.fluss.rpc.messages.RemoveServerTagResponse;
 import org.apache.fluss.rpc.protocol.ApiError;
 import org.apache.fluss.server.coordinator.event.AccessContextEvent;
@@ -89,8 +89,8 @@
 import org.apache.fluss.server.coordinator.event.NotifyKvSnapshotOffsetEvent;
 import org.apache.fluss.server.coordinator.event.NotifyLakeTableOffsetEvent;
 import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent;
-import org.apache.fluss.server.coordinator.event.ReleaseKvSnapshotLeaseEvent;
 import org.apache.fluss.server.coordinator.event.RebalanceEvent;
+import org.apache.fluss.server.coordinator.event.ReleaseKvSnapshotLeaseEvent;
 import org.apache.fluss.server.coordinator.event.RemoveServerTagEvent;
 import org.apache.fluss.server.coordinator.event.SchemaChangeEvent;
 import org.apache.fluss.server.coordinator.event.watcher.TableChangeWatcher;
@@ -125,7 +125,7 @@
 import org.apache.fluss.server.zk.data.lake.LakeTable;
 import org.apache.fluss.server.zk.data.lake.LakeTableHelper;
 import org.apache.fluss.server.zk.data.lake.LakeTableSnapshot;
-import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataHelper;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataManager;
 import org.apache.fluss.utils.clock.Clock;
 import org.apache.fluss.utils.types.Tuple2;
 
@@ -244,7 +244,7 @@ public CoordinatorEventProcessor(
         this.kvSnapshotLeaseManager =
                 new KvSnapshotLeaseManager(
                         conf,
-                        new KvSnapshotLeaseMetadataHelper(zooKeeperClient, remoteDataDir),
+                        new KvSnapshotLeaseMetadataManager(zooKeeperClient, remoteDataDir),
                         coordinatorContext,
                         clock,
                         coordinatorMetricGroup);
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java
@@ -128,9 +128,9 @@
 import org.apache.fluss.server.coordinator.event.ControlledShutdownEvent;
 import org.apache.fluss.server.coordinator.event.DropKvSnapshotLeaseEvent;
 import org.apache.fluss.server.coordinator.event.EventManager;
-import org.apache.fluss.server.coordinator.event.ReleaseKvSnapshotLeaseEvent;
 import org.apache.fluss.server.coordinator.event.ListRebalanceProgressEvent;
 import org.apache.fluss.server.coordinator.event.RebalanceEvent;
+import org.apache.fluss.server.coordinator.event.ReleaseKvSnapshotLeaseEvent;
 import org.apache.fluss.server.coordinator.event.RemoveServerTagEvent;
 import org.apache.fluss.server.coordinator.rebalance.goal.Goal;
 import org.apache.fluss.server.entity.CommitKvSnapshotData;
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManager.java
@@ -26,7 +26,7 @@
 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.KvSnapshotLeaseMetadataManager;
 import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
 import org.apache.fluss.utils.MapUtils;
 import org.apache.fluss.utils.clock.Clock;
@@ -48,7 +48,6 @@
 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;
@@ -57,8 +56,9 @@
 @ThreadSafe
 public class KvSnapshotLeaseManager {
     private static final Logger LOG = LoggerFactory.getLogger(KvSnapshotLeaseManager.class);
+    private static final AtomicInteger ZERO = new AtomicInteger(0);
 
-    private final KvSnapshotLeaseMetadataHelper metadataHelper;
+    private final KvSnapshotLeaseMetadataManager metadataHelper;
     private final CoordinatorContext coordinatorContext;
     private final Clock clock;
     private final ScheduledExecutorService scheduledExecutor;
@@ -75,7 +75,6 @@ public class KvSnapshotLeaseManager {
      * 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();
 
@@ -84,7 +83,7 @@ public class KvSnapshotLeaseManager {
 
     public KvSnapshotLeaseManager(
             Configuration conf,
-            KvSnapshotLeaseMetadataHelper metadataHelper,
+            KvSnapshotLeaseMetadataManager metadataHelper,
             CoordinatorContext coordinatorContext,
             Clock clock,
             CoordinatorMetricGroup coordinatorMetricGroup) {
@@ -101,7 +100,7 @@ public KvSnapshotLeaseManager(
     @VisibleForTesting
     public KvSnapshotLeaseManager(
             Configuration conf,
-            KvSnapshotLeaseMetadataHelper metadataHelper,
+            KvSnapshotLeaseMetadataManager metadataHelper,
             CoordinatorContext coordinatorContext,
             ScheduledExecutorService scheduledExecutor,
             Clock clock,
@@ -126,8 +125,7 @@ public void start() {
     }
 
     public void initialize() throws Exception {
-        List<String> leaseIds = metadataHelper.getLeasesList();
-        for (String leaseId : leaseIds) {
+        for (String leaseId : metadataHelper.getLeasesList()) {
             Optional<KvSnapshotLease> kvSnapshotLeaseOpt = metadataHelper.getLease(leaseId);
             if (kvSnapshotLeaseOpt.isPresent()) {
                 KvSnapshotLease kvSnapshotLease = kvSnapshotLeaseOpt.get();
@@ -145,8 +143,8 @@ public boolean snapshotLeaseNotExist(KvSnapshotLeaseForBucket kvSnapshotLeaseFor
         return inReadLock(
                 refCountLock,
                 () ->
-                        !refCount.containsKey(kvSnapshotLeaseForBucket)
-                                || refCount.get(kvSnapshotLeaseForBucket).get() <= 0);
+                        refCount.getOrDefault(kvSnapshotLeaseForBucket, new AtomicInteger(0)).get()
+                                <= 0);
     }
 
     /**
@@ -402,18 +400,17 @@ private void decrementRefCount(KvSnapshotLeaseForBucket 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);
-            }
-        }
+        kvSnapshotLeaseMap.entrySet().stream()
+                .filter(entry -> entry.getValue().getExpirationTime() < currentTime)
+                .map(Map.Entry::getKey)
+                .forEach(
+                        leaseId -> {
+                            try {
+                                drop(leaseId);
+                            } catch (Exception e) {
+                                LOG.error("Failed to clear kv snapshot lease {}", leaseId, e);
+                            }
+                        });
     }
 
     private void registerMetrics(CoordinatorMetricGroup coordinatorMetricGroup) {
@@ -436,14 +433,7 @@ int getLeasedBucketCount() {
     @VisibleForTesting
     int getRefCount(KvSnapshotLeaseForBucket kvSnapshotLeaseForBucket) {
         return inReadLock(
-                refCountLock,
-                () -> {
-                    if (!refCount.containsKey(kvSnapshotLeaseForBucket)) {
-                        return 0;
-                    } else {
-                        return refCount.get(kvSnapshotLeaseForBucket).get();
-                    }
-                });
+                refCountLock, () -> refCount.getOrDefault(kvSnapshotLeaseForBucket, ZERO).get());
     }
 
     @VisibleForTesting
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/utils/ServerRpcMessageUtils.java b/fluss-server/src/main/java/org/apache/fluss/server/utils/ServerRpcMessageUtils.java
@@ -150,8 +150,8 @@
 import org.apache.fluss.rpc.messages.ProduceLogResponse;
 import org.apache.fluss.rpc.messages.PutKvRequest;
 import org.apache.fluss.rpc.messages.PutKvResponse;
-import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
 import org.apache.fluss.rpc.messages.RebalanceResponse;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
 import org.apache.fluss.rpc.messages.StopReplicaRequest;
 import org.apache.fluss.rpc.messages.StopReplicaResponse;
 import org.apache.fluss.rpc.messages.UpdateMetadataRequest;
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManager.java
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/KvSnapshotLeaseManagerTest.java
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataManagerTest.java
diff --git a/website/docs/maintenance/observability/monitor-metrics.md b/website/docs/maintenance/observability/monitor-metrics.md
