You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@fluss.apache.org by ja...@apache.org on 2026/02/09 09:42:01 UTC
(fluss) branch main updated: [kv] Support kv snapshot lease (#2179)
This is an automated email from the ASF dual-hosted git repository.
jark pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss.git
The following commit(s) were added to refs/heads/main by this push:
new 0cee38343 [kv] Support kv snapshot lease (#2179)
0cee38343 is described below
commit 0cee383439323edc5fcc323f32b8c2b49dbea6dd
Author: yunhong <33...@qq.com>
AuthorDate: Mon Feb 9 17:41:57 2026 +0800
[kv] Support kv snapshot lease (#2179)
---
.../java/org/apache/fluss/client/admin/Admin.java | 9 +
.../org/apache/fluss/client/admin/FlussAdmin.java | 5 +
.../apache/fluss/client/admin/KvSnapshotLease.java | 73 +++
.../fluss/client/admin/KvSnapshotLeaseImpl.java | 90 +++
.../metadata/AcquireKvSnapshotLeaseResult.java | 55 ++
.../apache/fluss/client/metadata/KvSnapshots.java | 8 +
.../fluss/client/utils/ClientRpcMessageUtils.java | 77 +++
.../client/admin/CustomFlussClusterITCase.java | 125 +++++
.../client/metadata/TestingClientSchemaGetter.java | 6 +-
.../security/acl/FlussAuthorizationITCase.java | 132 +++++
.../batch/KvSnapshotBatchScannerITCase.java | 114 ++++
.../client/table/scanner/log/LogFetcherTest.java | 5 +-
.../org/apache/fluss/config/ConfigOptions.java | 10 +-
.../apache/fluss/metadata/TableBucketSnapshot.java | 66 +++
.../java/org/apache/fluss/metrics/MetricNames.java | 5 +
.../java/org/apache/fluss/utils/FlussPaths.java | 37 ++
.../java/org/apache/fluss/record/TestData.java | 14 +
.../apache/fluss/flink/FlinkConnectorOptions.java | 19 +
.../fluss/flink/catalog/FlinkTableFactory.java | 7 +-
.../procedure/DropKvSnapshotLeaseProcedure.java | 50 ++
.../fluss/flink/procedure/ProcedureManager.java | 3 +-
.../org/apache/fluss/flink/source/FlinkSource.java | 20 +-
.../fluss/flink/source/FlinkTableSource.java | 13 +-
.../org/apache/fluss/flink/source/FlussSource.java | 4 +-
.../source/enumerator/FlinkSourceEnumerator.java | 189 ++++++-
.../event/FinishedKvSnapshotConsumeEvent.java | 74 +++
.../flink/source/reader/FlinkSourceReader.java | 53 ++
.../fluss/flink/source/reader/LeaseContext.java | 91 +++
.../FlussSourceEnumeratorStateSerializer.java | 45 +-
.../flink/source/state/SourceEnumeratorState.java | 16 +-
.../flink/procedure/FlinkProcedureITCase.java | 43 +-
.../fluss/flink/source/FlinkTableSourceITCase.java | 49 ++
.../enumerator/FlinkSourceEnumeratorTest.java | 29 +-
.../state/SourceEnumeratorStateSerializerTest.java | 23 +-
.../org/apache/fluss/rpc/gateway/AdminGateway.java | 18 +
.../org/apache/fluss/rpc/protocol/ApiKeys.java | 5 +-
fluss-rpc/src/main/proto/FlussApi.proto | 34 ++
.../coordinator/CompletedSnapshotStoreManager.java | 15 +-
.../coordinator/CoordinatorEventProcessor.java | 9 +-
.../server/coordinator/CoordinatorServer.java | 35 +-
.../server/coordinator/CoordinatorService.java | 121 +++-
.../coordinator/lease/KvSnapshotLeaseHandler.java | 233 ++++++++
.../coordinator/lease/KvSnapshotLeaseManager.java | 489 ++++++++++++++++
.../lease/KvSnapshotLeaseMetadataManager.java | 281 ++++++++++
.../server/kv/snapshot/CompletedSnapshotStore.java | 128 +++--
.../server/kv/snapshot/SharedKvFileRegistry.java | 57 +-
.../fluss/server/utils/ServerRpcMessageUtils.java | 68 +++
.../apache/fluss/server/zk/ZooKeeperClient.java | 33 ++
.../org/apache/fluss/server/zk/data/ZkData.java | 36 ++
.../zk/data/lease/KvSnapshotLeaseMetadata.java | 71 +++
.../lease/KvSnapshotLeaseMetadataJsonSerde.java | 84 +++
.../server/zk/data/lease/KvSnapshotTableLease.java | 190 +++++++
.../data/lease/KvSnapshotTableLeaseJsonSerde.java | 116 ++++
.../CompletedSnapshotStoreManagerTest.java | 6 +-
.../coordinator/CoordinatorEventProcessorTest.java | 18 +-
.../server/coordinator/TestCoordinatorGateway.java | 24 +
.../lease/KvSnapshotLeaseHandlerTest.java | 287 ++++++++++
.../lease/KvSnapshotLeaseManagerTest.java | 612 +++++++++++++++++++++
.../lease/KvSnapshotLeaseMetadataManagerTest.java | 281 ++++++++++
.../rebalance/RebalanceManagerTest.java | 18 +-
.../statemachine/TableBucketStateMachineTest.java | 18 +-
.../kv/snapshot/CompletedSnapshotStoreTest.java | 250 ++++++++-
.../kv/snapshot/KvTabletSnapshotTargetTest.java | 3 +-
.../kv/snapshot/SharedKvFileRegistryTest.java | 90 +++
.../ZooKeeperCompletedSnapshotStoreTest.java | 3 +-
.../server/testutils/FlussClusterExtension.java | 2 +-
.../fluss/server/zk/ZooKeeperClientTest.java | 32 ++
.../KvSnapshotLeaseMetadataJsonSerdeTest.java | 57 ++
.../lease/KvSnapshotTableLeaseJsonSerdeTest.java | 54 ++
fluss-test-coverage/pom.xml | 3 +
website/docs/_configs/_partial_config.mdx | 16 +-
website/docs/engine-flink/options.md | 30 +-
website/docs/engine-flink/procedures.md | 34 ++
website/docs/maintenance/configuration.md | 3 +-
.../maintenance/observability/monitor-metrics.md | 12 +-
75 files changed, 5291 insertions(+), 144 deletions(-)
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
index b6e9ea1d6..f2e75ae3e 100644
--- 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
@@ -428,6 +428,15 @@ public interface Admin extends AutoCloseable {
CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
TableBucket bucket, long snapshotId);
+ /**
+ * 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);
+
/**
* Get table lake snapshot info of the given table asynchronously.
*
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
index 9df4bf248..a905823c4 100644
--- a/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
+++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java
@@ -400,6 +400,11 @@ public class FlussAdmin implements Admin {
.thenApply(ClientRpcMessageUtils::toKvSnapshotMetadata);
}
+ @Override
+ public KvSnapshotLease createKvSnapshotLease(String leaseId, long leaseDurationMs) {
+ return new KvSnapshotLeaseImpl(leaseId, leaseDurationMs, gateway);
+ }
+
@Override
public CompletableFuture<LakeSnapshot> getLatestLakeSnapshot(TablePath tablePath) {
GetLatestLakeSnapshotRequest request = new GetLatestLakeSnapshotRequest();
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/KvSnapshotLease.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/KvSnapshotLease.java
new file mode 100644
index 000000000..e8a654bc0
--- /dev/null
+++ b/fluss-client/src/main/java/org/apache/fluss/client/admin/KvSnapshotLease.java
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.fluss.client.admin;
+
+import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.client.metadata.AcquireKvSnapshotLeaseResult;
+import org.apache.fluss.metadata.TableBucket;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Represents a lease for managing KV snapshots. A lease allows acquiring, renewing, releasing, and
+ * dropping snapshot references under a time-bound agreement.
+ *
+ * @since 0.9
+ */
+@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();
+}
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/admin/KvSnapshotLeaseImpl.java b/fluss-client/src/main/java/org/apache/fluss/client/admin/KvSnapshotLeaseImpl.java
new file mode 100644
index 000000000..9d8900c06
--- /dev/null
+++ b/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 java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAcquireKvSnapshotLeaseRequest;
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeReleaseKvSnapshotLeaseRequest;
+
+/** 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) {
+ return gateway.releaseKvSnapshotLease(
+ makeReleaseKvSnapshotLeaseRequest(leaseId, bucketsToRelease))
+ .thenApply(r -> null);
+ }
+
+ @Override
+ public CompletableFuture<Void> dropLease() {
+ DropKvSnapshotLeaseRequest request = new DropKvSnapshotLeaseRequest().setLeaseId(leaseId);
+ return gateway.dropKvSnapshotLease(request).thenApply(r -> null);
+ }
+}
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/metadata/AcquireKvSnapshotLeaseResult.java b/fluss-client/src/main/java/org/apache/fluss/client/metadata/AcquireKvSnapshotLeaseResult.java
new file mode 100644
index 000000000..3852acd78
--- /dev/null
+++ b/fluss-client/src/main/java/org/apache/fluss/client/metadata/AcquireKvSnapshotLeaseResult.java
@@ -0,0 +1,55 @@
+/*
+ * 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.metadata;
+
+import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.metadata.TableBucket;
+
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * A class to represent the result of acquire kv snapshot lease. It contains:
+ *
+ * <ul>
+ * <li>A map of unavailable snapshots. Such as the specify snapshotId is not exist for this table
+ * bucket.
+ * </ul>
+ *
+ * @since 0.9
+ */
+@PublicEvolving
+public class AcquireKvSnapshotLeaseResult {
+ private final Map<TableBucket, Long> unavailableSnapshots;
+
+ public AcquireKvSnapshotLeaseResult(Map<TableBucket, Long> unavailableSnapshots) {
+ this.unavailableSnapshots = unavailableSnapshots;
+ }
+
+ /**
+ * Returns the set of buckets that could not be locked (e.g., snapshot ID doesn't exist or has
+ * already been GC'ed).
+ */
+ public Map<TableBucket, Long> getUnavailableSnapshots() {
+ return unavailableSnapshots;
+ }
+
+ public Set<TableBucket> getUnavailableTableBucketSet() {
+ return unavailableSnapshots.keySet();
+ }
+}
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/metadata/KvSnapshots.java b/fluss-client/src/main/java/org/apache/fluss/client/metadata/KvSnapshots.java
index 9f01ed8f1..91428dcb5 100644
--- a/fluss-client/src/main/java/org/apache/fluss/client/metadata/KvSnapshots.java
+++ b/fluss-client/src/main/java/org/apache/fluss/client/metadata/KvSnapshots.java
@@ -18,12 +18,14 @@
package org.apache.fluss.client.metadata;
import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.metadata.TableBucket;
import javax.annotation.Nullable;
import java.util.Map;
import java.util.OptionalLong;
import java.util.Set;
+import java.util.stream.Collectors;
/**
* A class representing the kv snapshots of a table or a partition. It contains multiple snapshots
@@ -71,6 +73,12 @@ public class KvSnapshots {
return snapshotIds.keySet();
}
+ public Set<TableBucket> getTableBuckets() {
+ return snapshotIds.keySet().stream()
+ .map(bucketId -> new TableBucket(tableId, partitionId, bucketId))
+ .collect(Collectors.toSet());
+ }
+
/**
* Get the latest snapshot id for this kv tablet (bucket), or empty if there are no snapshots.
*/
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java
index 1748e0e43..79f8ef4fb 100644
--- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java
+++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java
@@ -21,6 +21,7 @@ import org.apache.fluss.client.admin.OffsetSpec;
import org.apache.fluss.client.admin.ProducerOffsetsResult;
import org.apache.fluss.client.lookup.LookupBatch;
import org.apache.fluss.client.lookup.PrefixLookupBatch;
+import org.apache.fluss.client.metadata.AcquireKvSnapshotLeaseResult;
import org.apache.fluss.client.metadata.KvSnapshotMetadata;
import org.apache.fluss.client.metadata.KvSnapshots;
import org.apache.fluss.client.metadata.LakeSnapshot;
@@ -43,6 +44,8 @@ import org.apache.fluss.metadata.PhysicalTablePath;
import org.apache.fluss.metadata.TableBucket;
import org.apache.fluss.metadata.TableChange;
import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.AlterTableRequest;
import org.apache.fluss.rpc.messages.CreatePartitionRequest;
import org.apache.fluss.rpc.messages.DropPartitionRequest;
@@ -65,6 +68,8 @@ import org.apache.fluss.rpc.messages.PbDescribeConfig;
import org.apache.fluss.rpc.messages.PbDropColumn;
import org.apache.fluss.rpc.messages.PbKeyValue;
import org.apache.fluss.rpc.messages.PbKvSnapshot;
+import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForBucket;
+import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForTable;
import org.apache.fluss.rpc.messages.PbLakeSnapshotForBucket;
import org.apache.fluss.rpc.messages.PbLookupReqForBucket;
import org.apache.fluss.rpc.messages.PbModifyColumn;
@@ -78,10 +83,12 @@ import org.apache.fluss.rpc.messages.PbRebalanceProgressForBucket;
import org.apache.fluss.rpc.messages.PbRebalanceProgressForTable;
import org.apache.fluss.rpc.messages.PbRemotePathAndLocalFile;
import org.apache.fluss.rpc.messages.PbRenameColumn;
+import org.apache.fluss.rpc.messages.PbTableBucket;
import org.apache.fluss.rpc.messages.PrefixLookupRequest;
import org.apache.fluss.rpc.messages.ProduceLogRequest;
import org.apache.fluss.rpc.messages.PutKvRequest;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsRequest;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
import org.apache.fluss.rpc.protocol.MergeMode;
import org.apache.fluss.utils.json.DataTypeJsonSerde;
import org.apache.fluss.utils.json.JsonSerdeUtils;
@@ -411,6 +418,76 @@ public class ClientRpcMessageUtils {
return request;
}
+ public static AcquireKvSnapshotLeaseRequest makeAcquireKvSnapshotLeaseRequest(
+ String leaseId, Map<TableBucket, Long> snapshotIds, long leaseDuration) {
+ AcquireKvSnapshotLeaseRequest request = new AcquireKvSnapshotLeaseRequest();
+ request.setLeaseId(leaseId).setLeaseDurationMs(leaseDuration);
+
+ Map<Long, List<PbKvSnapshotLeaseForBucket>> pbLeaseForTables = new HashMap<>();
+ for (Map.Entry<TableBucket, Long> entry : snapshotIds.entrySet()) {
+ TableBucket tableBucket = entry.getKey();
+ Long snapshotId = entry.getValue();
+ PbKvSnapshotLeaseForBucket pbLeaseForBucket =
+ new PbKvSnapshotLeaseForBucket()
+ .setBucketId(tableBucket.getBucket())
+ .setSnapshotId(snapshotId);
+ if (tableBucket.getPartitionId() != null) {
+ pbLeaseForBucket.setPartitionId(tableBucket.getPartitionId());
+ }
+ pbLeaseForTables
+ .computeIfAbsent(tableBucket.getTableId(), k -> new ArrayList<>())
+ .add(pbLeaseForBucket);
+ }
+
+ for (Map.Entry<Long, List<PbKvSnapshotLeaseForBucket>> entry :
+ pbLeaseForTables.entrySet()) {
+ request.addSnapshotsToLease()
+ .setTableId(entry.getKey())
+ .addAllBucketSnapshots(entry.getValue());
+ }
+ return request;
+ }
+
+ public static AcquireKvSnapshotLeaseResult toAcquireKvSnapshotLeaseResult(
+ AcquireKvSnapshotLeaseResponse response) {
+ Map<TableBucket, Long> unavailableSnapshots = new HashMap<>();
+ for (PbKvSnapshotLeaseForTable unavailableSnapshot :
+ response.getUnavailableSnapshotsList()) {
+ long tableId = unavailableSnapshot.getTableId();
+ for (PbKvSnapshotLeaseForBucket leaseForBucket :
+ unavailableSnapshot.getBucketSnapshotsList()) {
+ TableBucket tableBucket =
+ new TableBucket(
+ tableId,
+ leaseForBucket.hasPartitionId()
+ ? leaseForBucket.getPartitionId()
+ : null,
+ leaseForBucket.getBucketId());
+ unavailableSnapshots.put(tableBucket, leaseForBucket.getSnapshotId());
+ }
+ }
+ return new AcquireKvSnapshotLeaseResult(unavailableSnapshots);
+ }
+
+ public static ReleaseKvSnapshotLeaseRequest makeReleaseKvSnapshotLeaseRequest(
+ String leaseId, Set<TableBucket> bucketsToRelease) {
+ ReleaseKvSnapshotLeaseRequest request = new ReleaseKvSnapshotLeaseRequest();
+ request.setLeaseId(leaseId);
+
+ List<PbTableBucket> pbTableBuckets = new ArrayList<>();
+ for (TableBucket tb : bucketsToRelease) {
+ PbTableBucket pbBucket =
+ new PbTableBucket().setTableId(tb.getTableId()).setBucketId(tb.getBucket());
+ if (tb.getPartitionId() != null) {
+ pbBucket.setPartitionId(tb.getPartitionId());
+ }
+ pbTableBuckets.add(pbBucket);
+ }
+
+ request.addAllBucketsToReleases(pbTableBuckets);
+ return request;
+ }
+
public static Optional<RebalanceProgress> toRebalanceProgress(
ListRebalanceProgressResponse response) {
if (!response.hasRebalanceId()) {
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/admin/CustomFlussClusterITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/admin/CustomFlussClusterITCase.java
index 28b725595..c5595c5d4 100644
--- a/fluss-client/src/test/java/org/apache/fluss/client/admin/CustomFlussClusterITCase.java
+++ b/fluss-client/src/test/java/org/apache/fluss/client/admin/CustomFlussClusterITCase.java
@@ -32,23 +32,34 @@ import org.apache.fluss.config.MemorySize;
import org.apache.fluss.metadata.DataLakeFormat;
import org.apache.fluss.metadata.DatabaseDescriptor;
import org.apache.fluss.metadata.MergeEngineType;
+import org.apache.fluss.metadata.TableBucket;
import org.apache.fluss.metadata.TableDescriptor;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.record.ChangeType;
import org.apache.fluss.row.InternalRow;
import org.apache.fluss.server.testutils.FlussClusterExtension;
+import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.types.RowType;
import org.junit.jupiter.api.Test;
import java.time.Duration;
import java.util.ArrayList;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.apache.fluss.record.LogRecordBatchFormat.V0_RECORD_BATCH_HEADER_SIZE;
import static org.apache.fluss.record.TestData.DATA1_SCHEMA_PK;
+import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR_PK;
import static org.apache.fluss.testutils.DataTestUtils.row;
import static org.apache.fluss.testutils.InternalRowAssert.assertThatRow;
+import static org.apache.fluss.testutils.common.CommonTestUtils.retry;
import static org.assertj.core.api.Assertions.assertThat;
/** IT case for tests that require manual cluster management. */
@@ -182,6 +193,120 @@ class CustomFlussClusterITCase {
}
}
+ @Test
+ void testConcurrentKvSnapshotLeaseOperations() throws Exception {
+ Configuration conf = initConfig();
+ // Use a short snapshot interval for stress testing
+ conf.set(ConfigOptions.KV_SNAPSHOT_INTERVAL, Duration.ofSeconds(1));
+ final FlussClusterExtension flussClusterExtension =
+ FlussClusterExtension.builder()
+ .setNumOfTabletServers(3)
+ .setClusterConf(conf)
+ .build();
+ flussClusterExtension.start();
+
+ TablePath tablePath = TablePath.of("test_stress_db", "test_concurrent_kv_snapshot_lease");
+
+ try (Connection connection =
+ ConnectionFactory.createConnection(
+ flussClusterExtension.getClientConfig());
+ Admin admin = connection.getAdmin()) {
+ admin.createDatabase(tablePath.getDatabaseName(), DatabaseDescriptor.EMPTY, false)
+ .get();
+ admin.createTable(tablePath, DATA1_TABLE_DESCRIPTOR_PK, false).get();
+
+ long tableId = admin.getTableInfo(tablePath).get().getTableId();
+ int bucketNum = 3;
+
+ // Write initial data
+ try (Table table = connection.getTable(tablePath)) {
+ UpsertWriter upsertWriter = table.newUpsert().createWriter();
+ for (int i = 0; i < 20; i++) {
+ upsertWriter.upsert(row(i, "value_" + i));
+ }
+ upsertWriter.flush();
+ }
+
+ // Trigger and wait for snapshot
+ flussClusterExtension.triggerAndWaitSnapshot(tablePath);
+
+ // Prepare snapshot IDs for lease operations
+ Map<TableBucket, Long> bucketSnapshotIds = new HashMap<>();
+ for (int bucket = 0; bucket < bucketNum; bucket++) {
+ bucketSnapshotIds.put(new TableBucket(tableId, bucket), 0L);
+ }
+
+ int numLeases = 5;
+ int numConcurrentOps = 3;
+ ExecutorService executor = Executors.newFixedThreadPool(numLeases * numConcurrentOps);
+ AtomicInteger successCount = new AtomicInteger(0);
+ AtomicInteger errorCount = new AtomicInteger(0);
+ List<CompletableFuture<Void>> futures = new ArrayList<>();
+
+ // Concurrently acquire, renew, and drop leases
+ for (int i = 0; i < numLeases; i++) {
+ final String leaseId = "stress-lease-" + i;
+ final long leaseDuration = Duration.ofDays(1).toMillis();
+
+ CompletableFuture<Void> future =
+ CompletableFuture.runAsync(
+ () -> {
+ try {
+ KvSnapshotLease lease =
+ admin.createKvSnapshotLease(leaseId, leaseDuration);
+
+ // Acquire snapshots
+ lease.acquireSnapshots(bucketSnapshotIds).get();
+
+ // Renew the lease multiple times
+ for (int r = 0; r < 3; r++) {
+ lease.renew().get();
+ }
+
+ // Drop the lease
+ lease.dropLease().get();
+
+ successCount.incrementAndGet();
+ } catch (Exception e) {
+ errorCount.incrementAndGet();
+ }
+ },
+ executor);
+ futures.add(future);
+ }
+
+ // Wait for all concurrent operations to complete
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+ .get(2, TimeUnit.MINUTES);
+
+ executor.shutdown();
+ assertThat(executor.awaitTermination(30, TimeUnit.SECONDS)).isTrue();
+
+ // Verify all operations succeeded
+ assertThat(successCount.get()).isEqualTo(numLeases);
+ assertThat(errorCount.get()).isEqualTo(0);
+
+ // Verify all leases are cleaned up
+ ZooKeeperClient zkClient = flussClusterExtension.getZooKeeperClient();
+ retry(
+ Duration.ofMinutes(1),
+ () -> assertThat(zkClient.getKvSnapshotLeasesList()).isEmpty());
+
+ // Verify data is still consistent after concurrent operations
+ try (Table table = connection.getTable(tablePath)) {
+ Lookuper lookuper = table.newLookup().createLookuper();
+ for (int i = 0; i < 20; i++) {
+ InternalRow gotRow = lookuper.lookup(row(i)).get().getSingletonRow();
+ assertThatRow(gotRow)
+ .withSchema(DATA1_SCHEMA_PK.getRowType())
+ .isEqualTo(row(i, "value_" + i));
+ }
+ }
+ } finally {
+ flussClusterExtension.close();
+ }
+ }
+
protected static Configuration initConfig() {
Configuration conf = new Configuration();
conf.setInt(ConfigOptions.DEFAULT_REPLICATION_FACTOR, 3);
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingClientSchemaGetter.java b/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingClientSchemaGetter.java
index a983eea11..ffb2c54dc 100644
--- a/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingClientSchemaGetter.java
+++ b/fluss-client/src/test/java/org/apache/fluss/client/metadata/TestingClientSchemaGetter.java
@@ -32,13 +32,13 @@ public class TestingClientSchemaGetter extends ClientSchemaGetter {
public TestingClientSchemaGetter(
TablePath tablePath,
SchemaInfo latestSchemaInfo,
- TestingMetadataUpdater metadataUpdater) {
+ TestingMetadataUpdater metadataUpdater,
+ Configuration conf) {
super(
tablePath,
latestSchemaInfo,
new FlussAdmin(
- RpcClient.create(
- new Configuration(), TestingClientMetricGroup.newInstance(), false),
+ RpcClient.create(conf, TestingClientMetricGroup.newInstance(), false),
metadataUpdater));
}
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java
index fdb5e3015..55aa2a126 100644
--- a/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java
+++ b/fluss-client/src/test/java/org/apache/fluss/client/security/acl/FlussAuthorizationITCase.java
@@ -22,6 +22,7 @@ import org.apache.fluss.client.ConnectionFactory;
import org.apache.fluss.client.FlussConnection;
import org.apache.fluss.client.admin.Admin;
import org.apache.fluss.client.admin.FlussAdmin;
+import org.apache.fluss.client.admin.KvSnapshotLease;
import org.apache.fluss.client.table.Table;
import org.apache.fluss.client.table.scanner.batch.BatchScanner;
import org.apache.fluss.client.table.writer.AppendWriter;
@@ -55,6 +56,7 @@ import org.apache.fluss.rpc.messages.GetKvSnapshotMetadataRequest;
import org.apache.fluss.rpc.messages.InitWriterRequest;
import org.apache.fluss.rpc.messages.InitWriterResponse;
import org.apache.fluss.rpc.messages.MetadataRequest;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
import org.apache.fluss.rpc.metrics.TestingClientMetricGroup;
import org.apache.fluss.security.acl.AccessControlEntry;
import org.apache.fluss.security.acl.AccessControlEntryFilter;
@@ -1060,6 +1062,136 @@ public class FlussAuthorizationITCase {
}
}
+ // ------------------------------------------------------------------------
+ // 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
+ guestAdmin
+ .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 {
+ TableInfo tableInfo = rootAdmin.getTableInfo(DATA1_TABLE_PATH_PK).get();
+ long tableId = tableInfo.getTableId();
+ FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId);
+
+ KvSnapshotLease kvSnapshotLease =
+ guestAdmin.createKvSnapshotLease("test-drop-lease", Duration.ofDays(1).toMillis());
+ Map<TableBucket, Long> snapshotIds = new HashMap<>();
+ snapshotIds.put(new TableBucket(tableId, 0), 0L);
+
+ // this there is no lease exists, so no AuthorizationException should be thrown
+ kvSnapshotLease.dropLease().get();
+
+ // 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();
+
+ // remove the READ permission to guest user on table resource
+ rootAdmin.dropAcls(Collections.singletonList(AclBindingFilter.ANY)).all().get();
+
+ // After removing the READ permission, AuthorizationException should be thrown
+ assertThatThrownBy(() -> kvSnapshotLease.dropLease().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
+ rootAdmin.createAcls(aclBindings).all().get();
+ FLUSS_CLUSTER_EXTENSION.waitUntilAuthenticationSync(aclBindings, true);
+
+ // test dropKvSnapshotLease with READ permission should succeed
+ // (no AuthorizationException should be thrown)
+ kvSnapshotLease.dropLease().get();
+ }
+
// ------------------------------------------------------------------------
// Producer Offsets Authorization Tests
// ------------------------------------------------------------------------
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
index 4127608fc..24415f075 100644
--- 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
@@ -19,6 +19,7 @@ package org.apache.fluss.client.table.scanner.batch;
import org.apache.fluss.bucketing.BucketingFunction;
import org.apache.fluss.client.admin.ClientToServerITCaseBase;
+import org.apache.fluss.client.admin.KvSnapshotLease;
import org.apache.fluss.client.metadata.KvSnapshots;
import org.apache.fluss.client.table.Table;
import org.apache.fluss.client.table.scanner.RemoteFileDownloader;
@@ -34,17 +35,23 @@ import org.apache.fluss.row.InternalRow;
import org.apache.fluss.row.ProjectedRow;
import org.apache.fluss.row.encode.CompactedKeyEncoder;
import org.apache.fluss.row.encode.KeyEncoder;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseHandler;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseMetadataManager;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
import org.apache.fluss.types.DataTypes;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import static org.apache.fluss.record.TestData.DATA1_ROW_TYPE;
import static org.apache.fluss.testutils.DataTestUtils.compactedRow;
@@ -180,6 +187,97 @@ class KvSnapshotBatchScannerITCase extends ClientToServerITCaseBase {
testSnapshotRead(tablePath, expectedRowByBuckets);
}
+ @Test
+ public void testKvSnapshotLease() throws Exception {
+ TablePath tablePath = TablePath.of(DEFAULT_DB, "test-kv-snapshot-lease");
+ long tableId = createTable(tablePath, DEFAULT_TABLE_DESCRIPTOR, true);
+
+ String kvSnapshotLeaseId1 = "test-lease";
+ String kvSnapshotLeaseId2 = "test-lease2";
+
+ // scan the snapshot
+ Map<TableBucket, List<InternalRow>> expectedRowByBuckets = putRows(tableId, tablePath, 10);
+
+ // wait snapshot finish
+ FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshots(expectedRowByBuckets.keySet());
+
+ ZooKeeperClient zkClient = FLUSS_CLUSTER_EXTENSION.getZooKeeperClient();
+ String remoteDataDir = FLUSS_CLUSTER_EXTENSION.getRemoteDataDir();
+ KvSnapshotLeaseMetadataManager metadataManager =
+ new KvSnapshotLeaseMetadataManager(zkClient, remoteDataDir);
+
+ assertThat(zkClient.getKvSnapshotLeasesList()).isEmpty();
+
+ // test register kv snapshot lease for snapshot 0.
+ Map<TableBucket, Long> consumeBuckets = new HashMap<>();
+ KvSnapshots kvSnapshots = admin.getLatestKvSnapshots(tablePath).get();
+ for (int bucketId : kvSnapshots.getBucketIds()) {
+ TableBucket tableBucket = new TableBucket(kvSnapshots.getTableId(), bucketId);
+ consumeBuckets.put(tableBucket, kvSnapshots.getSnapshotId(bucketId).getAsLong());
+ }
+
+ KvSnapshotLease kvSnapshotLease1 =
+ admin.createKvSnapshotLease(kvSnapshotLeaseId1, Duration.ofDays(1).toMillis());
+ kvSnapshotLease1.acquireSnapshots(consumeBuckets).get();
+ checkKvSnapshotLeaseEquals(
+ metadataManager, kvSnapshotLeaseId1, tableId, new Long[] {0L, 0L, 0L});
+
+ expectedRowByBuckets = putRows(tableId, tablePath, 10);
+ // wait snapshot2 finish
+ FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshots(expectedRowByBuckets.keySet());
+
+ // test register kv snapshot lease for snapshot 1.
+ consumeBuckets = new HashMap<>();
+ kvSnapshots = admin.getLatestKvSnapshots(tablePath).get();
+ for (int bucketId : kvSnapshots.getBucketIds()) {
+ TableBucket tableBucket = new TableBucket(kvSnapshots.getTableId(), bucketId);
+ consumeBuckets.put(tableBucket, kvSnapshots.getSnapshotId(bucketId).getAsLong());
+ }
+
+ KvSnapshotLease kvSnapshotLease2 =
+ admin.createKvSnapshotLease(kvSnapshotLeaseId2, Duration.ofDays(1).toMillis());
+ kvSnapshotLease2.acquireSnapshots(consumeBuckets).get();
+ checkKvSnapshotLeaseEquals(
+ metadataManager, kvSnapshotLeaseId2, tableId, new Long[] {1L, 1L, 1L});
+ // check even snapshot1 is generated, snapshot0 also retained as lease exists.
+ for (TableBucket tb : expectedRowByBuckets.keySet()) {
+ assertThat(zkClient.getTableBucketSnapshot(tb, 0L).isPresent()).isTrue();
+ assertThat(zkClient.getTableBucketSnapshot(tb, 1L).isPresent()).isTrue();
+ }
+
+ expectedRowByBuckets = putRows(tableId, tablePath, 10);
+ // wait snapshot3 finish
+ FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshots(expectedRowByBuckets.keySet());
+
+ // release lease1.
+ kvSnapshotLease1.releaseSnapshots(Collections.singleton(new TableBucket(tableId, 0))).get();
+ checkKvSnapshotLeaseEquals(
+ metadataManager, kvSnapshotLeaseId1, tableId, new Long[] {-1L, 0L, 0L});
+
+ // release lease2.
+ kvSnapshotLease2.releaseSnapshots(consumeBuckets.keySet()).get();
+ assertThat(zkClient.getKvSnapshotLeasesList()).doesNotContain(kvSnapshotLeaseId2);
+
+ // release all kv snapshot lease of lease1
+ kvSnapshotLease1.dropLease().get();
+ assertThat(zkClient.getKvSnapshotLeasesList()).isEmpty();
+
+ expectedRowByBuckets = putRows(tableId, tablePath, 10);
+ // wait snapshot2 finish
+ FLUSS_CLUSTER_EXTENSION.triggerAndWaitSnapshots(expectedRowByBuckets.keySet());
+ // as all leases are dropped, and new snapshot is generated, all old snapshot are
+ // cleared.
+ for (TableBucket tb : expectedRowByBuckets.keySet()) {
+ assertThat(zkClient.getTableBucketSnapshot(tb, 0L).isPresent()).isFalse();
+ assertThat(zkClient.getTableBucketSnapshot(tb, 1L).isPresent()).isFalse();
+ }
+
+ // drop no exist lease, no exception.
+ admin.createKvSnapshotLease("no-exist-lease", Duration.ofDays(1).toMillis())
+ .dropLease()
+ .get();
+ }
+
private Map<TableBucket, List<InternalRow>> putRows(
long tableId, TablePath tablePath, int rowNumber) throws Exception {
List<InternalRow> rows = new ArrayList<>();
@@ -238,4 +336,20 @@ class KvSnapshotBatchScannerITCase extends ClientToServerITCaseBase {
byte[] key = bucketKeyEncoder.encodeKey(row);
return function.bucketing(key, DEFAULT_BUCKET_NUM);
}
+
+ private void checkKvSnapshotLeaseEquals(
+ KvSnapshotLeaseMetadataManager metadataManager,
+ String leaseId,
+ long tableId,
+ Long[] expectedBucketIndex)
+ throws Exception {
+ assertThat(metadataManager.getLeasesList()).contains(leaseId);
+ Optional<KvSnapshotLeaseHandler> leaseOpt = metadataManager.getLease(leaseId);
+ assertThat(leaseOpt).isPresent();
+ KvSnapshotLeaseHandler actualLease = leaseOpt.get();
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = actualLease.getTableIdToTableLease();
+ KvSnapshotTableLease tableLease = tableIdToTableLease.get(tableId);
+ assertThat(tableLease).isNotNull();
+ assertThat(tableLease.getBucketSnapshots()).isEqualTo(expectedBucketIndex);
+ }
}
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherTest.java b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherTest.java
index f06f88614..a01809c58 100644
--- a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherTest.java
+++ b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogFetcherTest.java
@@ -70,7 +70,10 @@ public class LogFetcherTest {
metadataUpdater = initializeMetadataUpdater();
ClientSchemaGetter clientSchemaGetter =
new TestingClientSchemaGetter(
- DATA1_TABLE_PATH, new SchemaInfo(DATA1_SCHEMA, 0), metadataUpdater);
+ DATA1_TABLE_PATH,
+ new SchemaInfo(DATA1_SCHEMA, 0),
+ metadataUpdater,
+ new Configuration());
LogScannerStatus logScannerStatus = initializeLogScannerStatus();
logFetcher =
new LogFetcher(
diff --git a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java
index f12153922..08742ebc2 100644
--- a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java
+++ b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java
@@ -1552,9 +1552,17 @@ public class ConfigOptions {
public static final ConfigOption<Integer> KV_MAX_RETAINED_SNAPSHOTS =
key("kv.snapshot.num-retained")
.intType()
- .defaultValue(1)
+ .defaultValue(2)
.withDescription("The maximum number of completed snapshots to retain.");
+ public static final ConfigOption<Duration> KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL =
+ key("kv.snapshot.lease.expiration-check-interval")
+ .durationType()
+ .defaultValue(Duration.ofMinutes(10))
+ .withDescription(
+ "The interval to check the expiration of kv snapshot lease. "
+ + "The default setting is 10 minutes.");
+
public static final ConfigOption<Integer> KV_MAX_BACKGROUND_THREADS =
key("kv.rocksdb.thread.num")
.intType()
diff --git a/fluss-common/src/main/java/org/apache/fluss/metadata/TableBucketSnapshot.java b/fluss-common/src/main/java/org/apache/fluss/metadata/TableBucketSnapshot.java
new file mode 100644
index 000000000..2b9b47096
--- /dev/null
+++ b/fluss-common/src/main/java/org/apache/fluss/metadata/TableBucketSnapshot.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 of table bucket. */
+public class TableBucketSnapshot {
+ private final TableBucket tableBucket;
+ private final long snapshotId;
+
+ public TableBucketSnapshot(TableBucket tableBucket, long snapshotId) {
+ this.tableBucket = tableBucket;
+ this.snapshotId = snapshotId;
+ }
+
+ public TableBucket getTableBucket() {
+ return tableBucket;
+ }
+
+ public long getSnapshotId() {
+ return snapshotId;
+ }
+
+ @Override
+ public String toString() {
+ return "TableBucketSnapshot{"
+ + "tableBucket="
+ + tableBucket
+ + ", snapshotId="
+ + snapshotId
+ + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ TableBucketSnapshot that = (TableBucketSnapshot) o;
+ return snapshotId == that.snapshotId && Objects.equals(tableBucket, that.tableBucket);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(tableBucket, snapshotId);
+ }
+}
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
index 7f4ae5a03..e1a5e28a6 100644
--- 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,11 @@ 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.
+ // 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";
+
// --------------------------------------------------------------------------------------------
// metrics for tablet server
// --------------------------------------------------------------------------------------------
diff --git a/fluss-common/src/main/java/org/apache/fluss/utils/FlussPaths.java b/fluss-common/src/main/java/org/apache/fluss/utils/FlussPaths.java
index c2e796964..9a0659f18 100644
--- a/fluss-common/src/main/java/org/apache/fluss/utils/FlussPaths.java
+++ b/fluss-common/src/main/java/org/apache/fluss/utils/FlussPaths.java
@@ -99,6 +99,8 @@ public class FlussPaths {
/** Suffix of a producer offsets file. */
private static final String PRODUCER_OFFSETS_FILE_SUFFIX = ".offsets";
+ private static final String REMOTE_LEASE_DIR_NAME = "lease";
+
// ----------------------------------------------------------------------------------------
// LOG/KV Tablet Paths
// ----------------------------------------------------------------------------------------
@@ -728,6 +730,41 @@ public class FlussPaths {
UUID.randomUUID()));
}
+ /**
+ * Returns the remote directory path for storing kv snapshot lease files.
+ *
+ * <p>The path contract:
+ *
+ * <pre>
+ * {$remote.data.dir}/lease/kv-snapshot/{leaseId}/{tableId}/
+ * </pre>
+ */
+ private static FsPath remoteKvSnapshotLeaseDir(
+ String remoteDataDir, String leaseId, long tableId) {
+ return new FsPath(
+ String.format(
+ "%s/%s/kv-snapshot/%s/%d",
+ remoteDataDir, REMOTE_LEASE_DIR_NAME, leaseId, tableId));
+ }
+
+ /**
+ * Returns the remote file path for storing kv snapshot lease files.
+ *
+ * <p>The path contract:
+ *
+ * <pre>
+ * {$remoteKvSnapshotLeaseDir}/{uuid}.metadata
+ * </pre>
+ */
+ public static FsPath remoteKvSnapshotLeaseFile(
+ String remoteDataDir, String leaseId, long tableId) {
+ return new FsPath(
+ String.format(
+ "%s/%s.metadata",
+ remoteKvSnapshotLeaseDir(remoteDataDir, leaseId, tableId),
+ UUID.randomUUID()));
+ }
+
/**
* Returns the remote directory path for storing kv snapshot shared files (SST files with UUID
* prefix).
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
index b67e6443e..28ab9ecec 100644
--- 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/FlinkConnectorOptions.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java
index e00f131f1..c30ee5e89 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/FlinkConnectorOptions.java
@@ -31,6 +31,7 @@ import org.apache.flink.table.catalog.IntervalFreshness;
import java.time.Duration;
import java.util.Arrays;
import java.util.List;
+import java.util.UUID;
import static org.apache.flink.configuration.description.TextElement.text;
@@ -76,6 +77,24 @@ 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 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.");
+
+ public static final ConfigOption<Duration> SCAN_KV_SNAPSHOT_LEASE_DURATION =
+ ConfigOptions.key("scan.kv.snapshot.lease.duration")
+ .durationType()
+ .defaultValue(Duration.ofDays(1))
+ .withDescription(
+ "The time period how long to wait before expiring the kv snapshot lease to "
+ + "avoid kv snapshot blocking to delete.");
+
// --------------------------------------------------------------------------------------------
// Lookup specific options
// --------------------------------------------------------------------------------------------
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java
index 50223b520..e4b52ab4b 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java
@@ -27,6 +27,7 @@ import org.apache.fluss.flink.sink.shuffle.DistributionMode;
import org.apache.fluss.flink.source.BinlogFlinkTableSource;
import org.apache.fluss.flink.source.ChangelogFlinkTableSource;
import org.apache.fluss.flink.source.FlinkTableSource;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.utils.FlinkConnectorOptionsUtils;
import org.apache.fluss.metadata.TablePath;
@@ -142,6 +143,7 @@ public class FlinkTableFactory implements DynamicTableSourceFactory, DynamicTabl
.get(FlinkConnectorOptions.SCAN_PARTITION_DISCOVERY_INTERVAL)
.toMillis();
+ LeaseContext leaseContext = LeaseContext.fromConf(tableOptions);
return new FlinkTableSource(
toFlussTablePath(context.getObjectIdentifier()),
toFlussClientConfig(
@@ -158,7 +160,8 @@ public class FlinkTableFactory implements DynamicTableSourceFactory, DynamicTabl
partitionDiscoveryIntervalMs,
tableOptions.get(toFlinkOption(ConfigOptions.TABLE_DATALAKE_ENABLED)),
tableOptions.get(toFlinkOption(ConfigOptions.TABLE_MERGE_ENGINE)),
- context.getCatalogTable().getOptions());
+ context.getCatalogTable().getOptions(),
+ leaseContext);
}
@Override
@@ -214,6 +217,8 @@ public class FlinkTableFactory implements DynamicTableSourceFactory, DynamicTabl
FlinkConnectorOptions.SCAN_STARTUP_MODE,
FlinkConnectorOptions.SCAN_STARTUP_TIMESTAMP,
FlinkConnectorOptions.SCAN_PARTITION_DISCOVERY_INTERVAL,
+ FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_ID,
+ FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_DURATION,
FlinkConnectorOptions.LOOKUP_ASYNC,
FlinkConnectorOptions.LOOKUP_INSERT_IF_NOT_EXISTS,
FlinkConnectorOptions.SINK_IGNORE_DELETE,
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
new file mode 100644
index 000000000..a55255cf4
--- /dev/null
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/DropKvSnapshotLeaseProcedure.java
@@ -0,0 +1,50 @@
+/*
+ * 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.fluss.client.admin.KvSnapshotLease;
+
+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;
+
+import java.time.Duration;
+
+/**
+ * Procedure to drop all kv snapshots leased of specified leaseId. See {@link
+ * KvSnapshotLease#dropLease()} for more details.
+ *
+ * <p>Usage examples:
+ *
+ * <pre>
+ * -- Drop kv snapshots leased of specified leaseId
+ * CALL sys.drop_kv_snapshot_lease('test-lease-id');
+ * </pre>
+ */
+public class DropKvSnapshotLeaseProcedure extends ProcedureBase {
+
+ @ProcedureHint(
+ argument = {
+ @ArgumentHint(name = "leaseId", type = @DataTypeHint("STRING")),
+ })
+ public String[] call(ProcedureContext context, String leaseId) throws Exception {
+ admin.createKvSnapshotLease(leaseId, Duration.ofDays(1).toMillis()).dropLease().get();
+ return new String[] {"success"};
+ }
+}
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/ProcedureManager.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/ProcedureManager.java
index b97b3cd8e..fc2632e85 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/ProcedureManager.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/procedure/ProcedureManager.java
@@ -77,7 +77,8 @@ public class ProcedureManager {
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),
+ DROP_KV_SNAPSHOT_LEASE("sys.drop_kv_snapshot_lease", DropKvSnapshotLeaseProcedure.class);
private final String path;
private final Class<? extends ProcedureBase> procedureClass;
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkSource.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkSource.java
index 768fa7dda..fba394edf 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkSource.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkSource.java
@@ -25,6 +25,7 @@ import org.apache.fluss.flink.source.emitter.FlinkRecordEmitter;
import org.apache.fluss.flink.source.enumerator.FlinkSourceEnumerator;
import org.apache.fluss.flink.source.metrics.FlinkSourceReaderMetrics;
import org.apache.fluss.flink.source.reader.FlinkSourceReader;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.source.reader.RecordAndPos;
import org.apache.fluss.flink.source.split.SourceSplitBase;
import org.apache.fluss.flink.source.split.SourceSplitSerializer;
@@ -70,6 +71,7 @@ public class FlinkSource<OUT>
private final FlussDeserializationSchema<OUT> deserializationSchema;
@Nullable private final Predicate partitionFilters;
@Nullable private final LakeSource<LakeSplit> lakeSource;
+ private final LeaseContext leaseContext;
public FlinkSource(
Configuration flussConf,
@@ -82,7 +84,8 @@ public class FlinkSource<OUT>
long scanPartitionDiscoveryIntervalMs,
FlussDeserializationSchema<OUT> deserializationSchema,
boolean streaming,
- @Nullable Predicate partitionFilters) {
+ @Nullable Predicate partitionFilters,
+ LeaseContext leaseContext) {
this(
flussConf,
tablePath,
@@ -95,7 +98,8 @@ public class FlinkSource<OUT>
deserializationSchema,
streaming,
partitionFilters,
- null);
+ null,
+ leaseContext);
}
public FlinkSource(
@@ -110,7 +114,8 @@ public class FlinkSource<OUT>
FlussDeserializationSchema<OUT> deserializationSchema,
boolean streaming,
@Nullable Predicate partitionFilters,
- @Nullable LakeSource<LakeSplit> lakeSource) {
+ @Nullable LakeSource<LakeSplit> lakeSource,
+ LeaseContext leaseContext) {
this.flussConf = flussConf;
this.tablePath = tablePath;
this.hasPrimaryKey = hasPrimaryKey;
@@ -123,6 +128,7 @@ public class FlinkSource<OUT>
this.streaming = streaming;
this.partitionFilters = partitionFilters;
this.lakeSource = lakeSource;
+ this.leaseContext = leaseContext;
}
@Override
@@ -143,7 +149,8 @@ public class FlinkSource<OUT>
scanPartitionDiscoveryIntervalMs,
streaming,
partitionFilters,
- lakeSource);
+ lakeSource,
+ leaseContext);
}
@Override
@@ -163,7 +170,10 @@ public class FlinkSource<OUT>
scanPartitionDiscoveryIntervalMs,
streaming,
partitionFilters,
- lakeSource);
+ lakeSource,
+ new LeaseContext(
+ sourceEnumeratorState.getLeaseId(),
+ leaseContext.getKvSnapshotLeaseDurationMs()));
}
@Override
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java
index 0946db99d..4af25571f 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java
@@ -25,6 +25,7 @@ import org.apache.fluss.flink.source.deserializer.RowDataDeserializationSchema;
import org.apache.fluss.flink.source.lookup.FlinkAsyncLookupFunction;
import org.apache.fluss.flink.source.lookup.FlinkLookupFunction;
import org.apache.fluss.flink.source.lookup.LookupNormalizer;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.utils.FlinkConnectorOptionsUtils;
import org.apache.fluss.flink.utils.FlinkConversions;
import org.apache.fluss.flink.utils.PushdownUtils;
@@ -128,6 +129,8 @@ public class FlinkTableSource
private final long scanPartitionDiscoveryIntervalMs;
private final boolean isDataLakeEnabled;
+ private final LeaseContext leaseContext;
+
@Nullable private final MergeEngineType mergeEngineType;
// output type after projection pushdown
@@ -167,7 +170,8 @@ public class FlinkTableSource
long scanPartitionDiscoveryIntervalMs,
boolean isDataLakeEnabled,
@Nullable MergeEngineType mergeEngineType,
- Map<String, String> tableOptions) {
+ Map<String, String> tableOptions,
+ LeaseContext leaseContext) {
this.tablePath = tablePath;
this.flussConfig = flussConfig;
this.tableOutputType = tableOutputType;
@@ -184,6 +188,7 @@ public class FlinkTableSource
this.scanPartitionDiscoveryIntervalMs = scanPartitionDiscoveryIntervalMs;
this.isDataLakeEnabled = isDataLakeEnabled;
+ this.leaseContext = leaseContext;
this.mergeEngineType = mergeEngineType;
this.tableOptions = tableOptions;
if (isDataLakeEnabled) {
@@ -340,7 +345,8 @@ public class FlinkTableSource
new RowDataDeserializationSchema(),
streaming,
partitionFilters,
- enableLakeSource ? lakeSource : null);
+ enableLakeSource ? lakeSource : null,
+ leaseContext);
if (!streaming) {
// return a bounded source provide to make planner happy,
@@ -430,7 +436,8 @@ public class FlinkTableSource
scanPartitionDiscoveryIntervalMs,
isDataLakeEnabled,
mergeEngineType,
- tableOptions);
+ tableOptions,
+ leaseContext);
source.producedDataType = producedDataType;
source.projectedFields = projectedFields;
source.singleRowFilter = singleRowFilter;
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSource.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSource.java
index 05fe4e276..d00a74a26 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSource.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlussSource.java
@@ -21,6 +21,7 @@ import org.apache.fluss.annotation.VisibleForTesting;
import org.apache.fluss.client.initializer.OffsetsInitializer;
import org.apache.fluss.config.Configuration;
import org.apache.fluss.flink.source.deserializer.FlussDeserializationSchema;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.types.RowType;
@@ -80,7 +81,8 @@ public class FlussSource<OUT> extends FlinkSource<OUT> {
scanPartitionDiscoveryIntervalMs,
deserializationSchema,
streaming,
- null);
+ null,
+ LeaseContext.DEFAULT);
}
/**
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java
index a22c61d71..79ad2e1ab 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumerator.java
@@ -31,8 +31,10 @@ import org.apache.fluss.config.Configuration;
import org.apache.fluss.flink.lake.LakeSplitGenerator;
import org.apache.fluss.flink.lake.split.LakeSnapshotAndFlussLogSplit;
import org.apache.fluss.flink.lake.split.LakeSnapshotSplit;
+import org.apache.fluss.flink.source.event.FinishedKvSnapshotConsumeEvent;
import org.apache.fluss.flink.source.event.PartitionBucketsUnsubscribedEvent;
import org.apache.fluss.flink.source.event.PartitionsRemovedEvent;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.source.split.HybridSnapshotLogSplit;
import org.apache.fluss.flink.source.split.LogSplit;
import org.apache.fluss.flink.source.split.SourceSplitBase;
@@ -40,6 +42,7 @@ import org.apache.fluss.flink.source.state.SourceEnumeratorState;
import org.apache.fluss.lake.source.LakeSource;
import org.apache.fluss.lake.source.LakeSplit;
import org.apache.fluss.metadata.PartitionInfo;
+import org.apache.fluss.metadata.PhysicalTablePath;
import org.apache.fluss.metadata.TableBucket;
import org.apache.fluss.metadata.TableInfo;
import org.apache.fluss.metadata.TablePath;
@@ -71,9 +74,11 @@ import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
+import java.util.NavigableMap;
import java.util.Objects;
import java.util.OptionalLong;
import java.util.Set;
+import java.util.TreeMap;
import java.util.stream.Collectors;
import static org.apache.fluss.utils.Preconditions.checkNotNull;
@@ -132,6 +137,11 @@ public class FlinkSourceEnumerator
private final OffsetsInitializer startingOffsetsInitializer;
private final OffsetsInitializer stoppingOffsetsInitializer;
+ private final LeaseContext leaseContext;
+
+ /** checkpointId -> tableBuckets who finished consume kv snapshots. */
+ private final TreeMap<Long, Set<TableBucket>> consumedKvSnapshotMap = new TreeMap<>();
+
// Lazily instantiated or mutable fields.
private Connection connection;
private Admin flussAdmin;
@@ -160,7 +170,8 @@ public class FlinkSourceEnumerator
long scanPartitionDiscoveryIntervalMs,
boolean streaming,
@Nullable Predicate partitionFilters,
- @Nullable LakeSource<LakeSplit> lakeSource) {
+ @Nullable LakeSource<LakeSplit> lakeSource,
+ LeaseContext leaseContext) {
this(
tablePath,
flussConf,
@@ -174,7 +185,8 @@ public class FlinkSourceEnumerator
scanPartitionDiscoveryIntervalMs,
streaming,
partitionFilters,
- lakeSource);
+ lakeSource,
+ leaseContext);
}
public FlinkSourceEnumerator(
@@ -190,7 +202,8 @@ public class FlinkSourceEnumerator
long scanPartitionDiscoveryIntervalMs,
boolean streaming,
@Nullable Predicate partitionFilters,
- @Nullable LakeSource<LakeSplit> lakeSource) {
+ @Nullable LakeSource<LakeSplit> lakeSource,
+ LeaseContext leaseContext) {
this(
tablePath,
flussConf,
@@ -205,7 +218,8 @@ public class FlinkSourceEnumerator
streaming,
partitionFilters,
lakeSource,
- new WorkerExecutor(context));
+ new WorkerExecutor(context),
+ leaseContext);
}
FlinkSourceEnumerator(
@@ -222,7 +236,8 @@ public class FlinkSourceEnumerator
boolean streaming,
@Nullable Predicate partitionFilters,
@Nullable LakeSource<LakeSplit> lakeSource,
- WorkerExecutor workerExecutor) {
+ WorkerExecutor workerExecutor,
+ LeaseContext leaseContext) {
this.tablePath = checkNotNull(tablePath);
this.flussConf = checkNotNull(flussConf);
this.hasPrimaryKey = hasPrimaryKey;
@@ -243,6 +258,7 @@ public class FlinkSourceEnumerator
streaming ? new NoStoppingOffsetsInitializer() : OffsetsInitializer.latest();
this.lakeSource = lakeSource;
this.workerExecutor = workerExecutor;
+ this.leaseContext = leaseContext;
}
@Override
@@ -349,16 +365,7 @@ public class FlinkSourceEnumerator
private List<SourceSplitBase> initNonPartitionedSplits() {
if (hasPrimaryKey && startingOffsetsInitializer instanceof SnapshotOffsetsInitializer) {
- // get the table snapshot info
- final KvSnapshots kvSnapshots;
- try {
- kvSnapshots = flussAdmin.getLatestKvSnapshots(tablePath).get();
- } catch (Exception e) {
- throw new FlinkRuntimeException(
- String.format("Failed to get table snapshot for %s", tablePath),
- ExceptionUtils.stripCompletionException(e));
- }
- return getSnapshotAndLogSplits(kvSnapshots, null);
+ return getSnapshotAndLogSplits(getLatestKvSnapshotsAndRegister(null), null);
} else {
return getLogSplit(null, null);
}
@@ -531,22 +538,77 @@ public class FlinkSourceEnumerator
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();
+ Set<TableBucket> unavailableTableBucketSet =
+ flussAdmin
+ .createKvSnapshotLease(kvSnapshotLeaseId, kvSnapshotLeaseDurationMs)
+ .acquireSnapshots(bucketsToLease)
+ .get()
+ .getUnavailableTableBucketSet();
+ if (!unavailableTableBucketSet.isEmpty()) {
+ LOG.error(
+ "Failed to acquire kv snapshot lease for table {}: {}.",
+ tablePath,
+ unavailableTableBucketSet);
+ }
+ }
+ } catch (Exception e) {
+ throw new FlinkRuntimeException(
+ String.format("Failed to get table snapshot for %s", tablePath),
+ ExceptionUtils.stripCompletionException(e));
+ }
+
+ return new KvSnapshots(tableId, partitionId, snapshotIds, logOffsets);
+ }
+
+ private KvSnapshots getLatestKvSnapshots(@Nullable String partitionName) throws Exception {
+ if (partitionName == null) {
+ return flussAdmin.getLatestKvSnapshots(tablePath).get();
+ } else {
+ return flussAdmin.getLatestKvSnapshots(tablePath, partitionName).get();
+ }
+ }
+
private List<SourceSplitBase> getSnapshotAndLogSplits(
KvSnapshots snapshots, @Nullable String partitionName) {
long tableId = snapshots.getTableId();
@@ -889,6 +951,18 @@ public class FlinkSourceEnumerator
for (Long partitionToRemove : partitionsPendingRemove) {
assignedPartitions.remove(partitionToRemove);
}
+ } else if (sourceEvent instanceof FinishedKvSnapshotConsumeEvent) {
+ FinishedKvSnapshotConsumeEvent event = (FinishedKvSnapshotConsumeEvent) sourceEvent;
+ long checkpointId = event.getCheckpointId();
+ Set<TableBucket> tableBuckets = event.getTableBuckets();
+ if (!tableBuckets.isEmpty()) {
+ LOG.info(
+ "Received finished kv snapshot consumer event for buckets: {}, checkpoint id: {}",
+ tableBuckets,
+ checkpointId);
+ }
+
+ tableBuckets.forEach(tableBucket -> addConsumedBucket(checkpointId, tableBucket));
}
}
@@ -923,15 +997,74 @@ public class FlinkSourceEnumerator
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 (!streaming
+ && hasPrimaryKey
+ && startingOffsetsInitializer instanceof SnapshotOffsetsInitializer) {
+ // drop the kv snapshot lease for the batch mode.
+ flussAdmin
+ .createKvSnapshotLease(
+ leaseContext.getKvSnapshotLeaseId(),
+ leaseContext.getKvSnapshotLeaseDurationMs())
+ .dropLease()
+ .get();
+ }
+
if (workerExecutor != null) {
workerExecutor.close();
}
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/event/FinishedKvSnapshotConsumeEvent.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/event/FinishedKvSnapshotConsumeEvent.java
new file mode 100644
index 000000000..a5ad1f902
--- /dev/null
+++ b/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 {
+ private static final long serialVersionUID = 1L;
+
+ private final long checkpointId;
+ /** The tableBucket set who finished consume kv snapshots. */
+ private final Set<TableBucket> tableBuckets;
+
+ public FinishedKvSnapshotConsumeEvent(long checkpointId, Set<TableBucket> tableBuckets) {
+ this.checkpointId = checkpointId;
+ this.tableBuckets = tableBuckets;
+ }
+
+ public long getCheckpointId() {
+ return checkpointId;
+ }
+
+ public Set<TableBucket> getTableBuckets() {
+ return tableBuckets;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ FinishedKvSnapshotConsumeEvent that = (FinishedKvSnapshotConsumeEvent) o;
+ return checkpointId == that.checkpointId && Objects.equals(tableBuckets, that.tableBuckets);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(checkpointId, tableBuckets);
+ }
+
+ @Override
+ public String toString() {
+ return "FinishedKvSnapshotConsumeEvent{"
+ + "checkpointId="
+ + checkpointId
+ + ", tableBuckets="
+ + tableBuckets
+ + '}';
+ }
+}
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java
index 6364532bb..9a6c4e8da 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/FlinkSourceReader.java
@@ -21,10 +21,12 @@ import org.apache.fluss.config.Configuration;
import org.apache.fluss.flink.adapter.SingleThreadMultiplexSourceReaderBaseAdapter;
import org.apache.fluss.flink.lake.LakeSplitStateInitializer;
import org.apache.fluss.flink.source.emitter.FlinkRecordEmitter;
+import org.apache.fluss.flink.source.event.FinishedKvSnapshotConsumeEvent;
import org.apache.fluss.flink.source.event.PartitionBucketsUnsubscribedEvent;
import org.apache.fluss.flink.source.event.PartitionsRemovedEvent;
import org.apache.fluss.flink.source.metrics.FlinkSourceReaderMetrics;
import org.apache.fluss.flink.source.reader.fetcher.FlinkSourceFetcherManager;
+import org.apache.fluss.flink.source.split.HybridSnapshotLogSplit;
import org.apache.fluss.flink.source.split.HybridSnapshotLogSplitState;
import org.apache.fluss.flink.source.split.LogSplitState;
import org.apache.fluss.flink.source.split.SourceSplitBase;
@@ -39,9 +41,13 @@ import org.apache.flink.api.connector.source.SourceEvent;
import org.apache.flink.api.connector.source.SourceReaderContext;
import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import javax.annotation.Nullable;
+import java.util.HashSet;
+import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;
@@ -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> finishedKvSnapshotConsumeBuckets;
+
public FlinkSourceReader(
FutureCompletingBlockingQueue<RecordsWithSplitIds<RecordAndPos>> elementsQueue,
Configuration flussConfig,
@@ -77,6 +88,7 @@ public class FlinkSourceReader<OUT>
recordEmitter,
context.getConfiguration(),
context);
+ this.finishedKvSnapshotConsumeBuckets = new HashSet<>();
}
@Override
@@ -84,6 +96,47 @@ public class FlinkSourceReader<OUT>
// 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);
+
+ // TODO Reduce the external IO operation, trace by
+ // https://github.com/apache/fluss/issues/2597.
+ context.sendSourceEventToCoordinator(
+ new FinishedKvSnapshotConsumeEvent(
+ checkpointId, bucketsFinishedConsumeKvSnapshot));
+ // It won't be sent anymore in the future for this table bucket, but will be resent
+ // after failover recovery as ignoreBuckets is cleared.
+ finishedKvSnapshotConsumeBuckets.addAll(bucketsFinishedConsumeKvSnapshot);
+ }
+
+ return sourceSplitBases;
+ }
+
@Override
public void handleSourceEvents(SourceEvent sourceEvent) {
if (sourceEvent instanceof PartitionsRemovedEvent) {
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/LeaseContext.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/LeaseContext.java
new file mode 100644
index 000000000..c17fb15b7
--- /dev/null
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/LeaseContext.java
@@ -0,0 +1,91 @@
+/*
+ * 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 org.apache.fluss.flink.FlinkConnectorOptions;
+
+import org.apache.flink.configuration.ReadableConfig;
+
+import java.io.Serializable;
+import java.time.Duration;
+import java.util.Objects;
+import java.util.UUID;
+
+/** Context for lease. */
+public class LeaseContext implements Serializable {
+
+ private static final long serialVersionUID = 1L;
+
+ public static final LeaseContext DEFAULT =
+ new LeaseContext(UUID.randomUUID().toString(), Duration.ofDays(1).toMillis());
+
+ // kv snapshot lease id. null for log table.
+ private final String kvSnapshotLeaseId;
+
+ // kv snapshot lease duration. null for log table.
+ private final long kvSnapshotLeaseDurationMs;
+
+ public LeaseContext(String kvSnapshotLeaseId, long kvSnapshotLeaseDurationMs) {
+ this.kvSnapshotLeaseId = kvSnapshotLeaseId;
+ this.kvSnapshotLeaseDurationMs = kvSnapshotLeaseDurationMs;
+ }
+
+ public static LeaseContext fromConf(ReadableConfig tableOptions) {
+ return new LeaseContext(
+ tableOptions.get(FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_ID),
+ tableOptions.get(FlinkConnectorOptions.SCAN_KV_SNAPSHOT_LEASE_DURATION).toMillis());
+ }
+
+ public String getKvSnapshotLeaseId() {
+ return kvSnapshotLeaseId;
+ }
+
+ public long getKvSnapshotLeaseDurationMs() {
+ return kvSnapshotLeaseDurationMs;
+ }
+
+ @Override
+ public String toString() {
+ return "LeaseContext{"
+ + "kvSnapshotLeaseId='"
+ + kvSnapshotLeaseId
+ + '\''
+ + ", kvSnapshotLeaseDurationMs="
+ + kvSnapshotLeaseDurationMs
+ + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ LeaseContext that = (LeaseContext) o;
+
+ return Objects.equals(kvSnapshotLeaseId, that.kvSnapshotLeaseId)
+ && Objects.equals(kvSnapshotLeaseDurationMs, that.kvSnapshotLeaseDurationMs);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(kvSnapshotLeaseId, kvSnapshotLeaseDurationMs);
+ }
+}
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
index b721032d0..121041df1 100644
--- 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
@@ -18,6 +18,7 @@
package org.apache.fluss.flink.source.state;
import org.apache.fluss.annotation.VisibleForTesting;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.source.split.SourceSplitBase;
import org.apache.fluss.flink.source.split.SourceSplitSerializer;
import org.apache.fluss.lake.source.LakeSource;
@@ -68,11 +69,12 @@ public class FlussSourceEnumeratorStateSerializer
private static final int VERSION_0 = 0;
private static final int VERSION_1 = 1;
+ private static final int VERSION_2 = 2;
private static final ThreadLocal<DataOutputSerializer> SERIALIZER_CACHE =
ThreadLocal.withInitial(() -> new DataOutputSerializer(64));
- private static final int CURRENT_VERSION = VERSION_1;
+ private static final int CURRENT_VERSION = VERSION_2;
public FlussSourceEnumeratorStateSerializer(LakeSource<LakeSplit> lakeSource) {
this.lakeSource = lakeSource;
@@ -94,6 +96,9 @@ public class FlussSourceEnumeratorStateSerializer
// serialize remain hybrid lake splits
serializeRemainingHybridLakeFlussSplits(out, state);
+ // write lease context
+ serializeLeaseId(out, state);
+
final byte[] result = out.getCopyOfBuffer();
out.clear();
return result;
@@ -164,6 +169,8 @@ public class FlussSourceEnumeratorStateSerializer
@Override
public SourceEnumeratorState deserialize(int version, byte[] serialized) throws IOException {
switch (version) {
+ case VERSION_2:
+ return deserializeV2(serialized);
case VERSION_1:
return deserializeV1(serialized);
case VERSION_0:
@@ -187,10 +194,12 @@ public class FlussSourceEnumeratorStateSerializer
if (lakeSource != null) {
remainingHybridLakeFlussSplits = deserializeRemainingHybridLakeFlussSplits(in);
}
+
return new SourceEnumeratorState(
assignBucketAndPartitions.f0,
assignBucketAndPartitions.f1,
- remainingHybridLakeFlussSplits);
+ remainingHybridLakeFlussSplits,
+ LeaseContext.DEFAULT.getKvSnapshotLeaseId());
}
private SourceEnumeratorState deserializeV1(byte[] serialized) throws IOException {
@@ -203,10 +212,28 @@ public class FlussSourceEnumeratorStateSerializer
// splits. The serialized state encodes their presence via a boolean flag, so
// this logic no longer depends on the lakeSource flag. This unconditional
// deserialization is the intended behavior change compared to VERSION_0.
+
return new SourceEnumeratorState(
assignBucketAndPartitions.f0,
assignBucketAndPartitions.f1,
- remainingHybridLakeFlussSplits);
+ remainingHybridLakeFlussSplits,
+ LeaseContext.DEFAULT.getKvSnapshotLeaseId());
+ }
+
+ private SourceEnumeratorState deserializeV2(byte[] serialized) throws IOException {
+ DataInputDeserializer in = new DataInputDeserializer(serialized);
+ Tuple2<Set<TableBucket>, Map<Long, String>> assignBucketAndPartitions =
+ deserializeAssignBucketAndPartitions(in);
+ List<SourceSplitBase> remainingHybridLakeFlussSplits =
+ deserializeRemainingHybridLakeFlussSplits(in);
+
+ // deserialize lease context
+ LeaseContext leaseContext = deserializeLeaseId(in);
+ return new SourceEnumeratorState(
+ assignBucketAndPartitions.f0,
+ assignBucketAndPartitions.f1,
+ remainingHybridLakeFlussSplits,
+ leaseContext.getKvSnapshotLeaseId());
}
private Tuple2<Set<TableBucket>, Map<Long, String>> deserializeAssignBucketAndPartitions(
@@ -260,4 +287,16 @@ public class FlussSourceEnumeratorStateSerializer
return null;
}
}
+
+ private void serializeLeaseId(final DataOutputSerializer out, SourceEnumeratorState state)
+ throws IOException {
+ String leaseId = state.getLeaseId();
+ out.writeUTF(leaseId);
+ }
+
+ private LeaseContext deserializeLeaseId(final DataInputDeserializer in) throws IOException {
+ String kvSnapshotLeaseId = in.readUTF();
+ return new LeaseContext(
+ kvSnapshotLeaseId, LeaseContext.DEFAULT.getKvSnapshotLeaseDurationMs());
+ }
}
diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/SourceEnumeratorState.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/SourceEnumeratorState.java
index 6042e65f1..7ee021ff5 100644
--- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/SourceEnumeratorState.java
+++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/state/SourceEnumeratorState.java
@@ -41,13 +41,18 @@ public class SourceEnumeratorState {
// lake snapshot
@Nullable private final List<SourceSplitBase> remainingHybridLakeFlussSplits;
+ // lease context for restore.
+ private final String leaseId;
+
public SourceEnumeratorState(
Set<TableBucket> assignedBuckets,
Map<Long, String> assignedPartitions,
- @Nullable List<SourceSplitBase> remainingHybridLakeFlussSplits) {
+ @Nullable List<SourceSplitBase> remainingHybridLakeFlussSplits,
+ String leaseId) {
this.assignedBuckets = assignedBuckets;
this.assignedPartitions = assignedPartitions;
this.remainingHybridLakeFlussSplits = remainingHybridLakeFlussSplits;
+ this.leaseId = leaseId;
}
public Set<TableBucket> getAssignedBuckets() {
@@ -63,6 +68,10 @@ public class SourceEnumeratorState {
return remainingHybridLakeFlussSplits;
}
+ public String getLeaseId() {
+ return leaseId;
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) {
@@ -75,7 +84,8 @@ public class SourceEnumeratorState {
return Objects.equals(assignedBuckets, that.assignedBuckets)
&& Objects.equals(assignedPartitions, that.assignedPartitions)
&& Objects.equals(
- remainingHybridLakeFlussSplits, that.remainingHybridLakeFlussSplits);
+ remainingHybridLakeFlussSplits, that.remainingHybridLakeFlussSplits)
+ && Objects.equals(leaseId, that.leaseId);
}
@Override
@@ -92,6 +102,8 @@ public class SourceEnumeratorState {
+ assignedPartitions
+ ", remainingHybridLakeFlussSplits="
+ remainingHybridLakeFlussSplits
+ + ", leaseId="
+ + leaseId
+ '}';
}
}
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
index 5138ca869..5e22c3ff0 100644
--- 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
@@ -30,6 +30,7 @@ import org.apache.fluss.exception.NoRebalanceInProgressException;
import org.apache.fluss.exception.SecurityDisabledException;
import org.apache.fluss.metadata.DataLakeFormat;
import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.row.InternalRow;
import org.apache.fluss.server.testutils.FlussClusterExtension;
import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.server.zk.data.ServerTags;
@@ -57,7 +58,9 @@ import java.util.stream.Collectors;
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.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.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -140,7 +143,8 @@ public abstract class FlinkProcedureITCase {
"+I[sys.remove_server_tag]",
"+I[sys.rebalance]",
"+I[sys.cancel_rebalance]",
- "+I[sys.list_rebalance]");
+ "+I[sys.list_rebalance]",
+ "+I[sys.drop_kv_snapshot_lease]");
// make sure no more results is unread.
assertResultsIgnoreOrder(showProceduresIterator, expectedShowProceduresResult, true);
}
@@ -779,6 +783,43 @@ public abstract class FlinkProcedureITCase {
});
}
+ @Test
+ void testDropKvSnapshotLeaseProcedure() throws Exception {
+ tEnv.executeSql(
+ "create table testcatalog.fluss.pk_table_test_kv_snapshot_lease ("
+ + "a int not null primary key not enforced, b varchar)");
+ TablePath tablePath = TablePath.of("fluss", "pk_table_test_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);
+
+ List<String> expectedRows = Arrays.asList("+I[1, v1]", "+I[2, v2]", "+I[3, v3]");
+
+ String leaseId = "test-lease-kjhdds23";
+ org.apache.flink.util.CloseableIterator<Row> rowIter =
+ tEnv.executeSql(
+ "select * from testcatalog.fluss.pk_table_test_kv_snapshot_lease "
+ + "/*+ OPTIONS('scan.kv.snapshot.lease.id' = '"
+ + leaseId
+ + "') */")
+ .collect();
+ assertResultsIgnoreOrder(rowIter, expectedRows, false);
+
+ // Lease will not be dropped automatically as the checkpoint not trigger.
+ ZooKeeperClient zkClient = FLUSS_CLUSTER_EXTENSION.getZooKeeperClient();
+ assertThat(zkClient.getKvSnapshotLeaseMetadata(leaseId)).isPresent();
+ tEnv.executeSql(
+ String.format(
+ "Call %s.sys.drop_kv_snapshot_lease('" + leaseId + "' )",
+ CATALOG_NAME))
+ .await();
+ assertThat(zkClient.getKvSnapshotLeaseMetadata(leaseId)).isNotPresent();
+ }
+
private static Configuration initConfig() {
Configuration conf = new Configuration();
conf.setInt(ConfigOptions.DEFAULT_REPLICATION_FACTOR, 3);
diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java
index 2de475774..935190092 100644
--- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java
+++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/FlinkTableSourceITCase.java
@@ -28,6 +28,7 @@ import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.row.GenericRow;
import org.apache.fluss.row.InternalRow;
import org.apache.fluss.server.testutils.FlussClusterExtension;
+import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.utils.clock.ManualClock;
import org.apache.commons.lang3.RandomUtils;
@@ -77,6 +78,7 @@ import static org.apache.fluss.flink.utils.FlinkTestBase.writeRows;
import static org.apache.fluss.flink.utils.FlinkTestBase.writeRowsToPartition;
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.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -352,6 +354,53 @@ abstract class FlinkTableSourceITCase extends AbstractTestBase {
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-consumer-1') */")
+ .collect();
+ assertResultsIgnoreOrder(rowIter, expectedRows, false);
+
+ // now, we put rows to the table again, should read the log
+ expectedRows =
+ Arrays.asList(
+ "-U[1, v1]",
+ "+U[1, v1]",
+ "-U[2, v2]",
+ "+U[2, v2]",
+ "-U[3, v3]",
+ "+U[3, v3]");
+ writeRows(conn, tablePath, rows, false);
+ assertResultsIgnoreOrder(rowIter, expectedRows, true);
+
+ // check lease will be dropped after job finished.
+ ZooKeeperClient zkClient = FLUSS_CLUSTER_EXTENSION.getZooKeeperClient();
+ retry(
+ Duration.ofMinutes(1),
+ () -> {
+ assertThat(zkClient.getKvSnapshotLeaseMetadata("test-consumer-1"))
+ .isNotPresent();
+ assertThat(zkClient.getKvSnapshotLeasesList().contains("test-consumer-1"))
+ .isFalse();
+ });
+ }
+
// -------------------------------------------------------------------------------------
// Fluss scan start mode tests
// -------------------------------------------------------------------------------------
diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java
index 238accf0b..eb526e4e5 100644
--- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java
+++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java
@@ -27,6 +27,7 @@ import org.apache.fluss.flink.lake.split.LakeSnapshotAndFlussLogSplit;
import org.apache.fluss.flink.lake.split.LakeSnapshotSplit;
import org.apache.fluss.flink.source.event.PartitionBucketsUnsubscribedEvent;
import org.apache.fluss.flink.source.event.PartitionsRemovedEvent;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.source.split.HybridSnapshotLogSplit;
import org.apache.fluss.flink.source.split.LogSplit;
import org.apache.fluss.flink.source.split.SnapshotSplit;
@@ -113,7 +114,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null);
+ null,
+ LeaseContext.DEFAULT);
enumerator.start();
@@ -161,7 +163,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null);
+ null,
+ LeaseContext.DEFAULT);
enumerator.start();
// register all read
for (int i = 0; i < numSubtasks; i++) {
@@ -233,7 +236,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null);
+ null,
+ LeaseContext.DEFAULT);
enumerator.start();
@@ -280,7 +284,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null);
+ null,
+ LeaseContext.DEFAULT);
enumerator.start();
@@ -317,7 +322,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null);
+ null,
+ LeaseContext.DEFAULT);
enumerator.start();
@@ -378,7 +384,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null);
+ null,
+ LeaseContext.DEFAULT);
enumerator.start();
assertThat(context.getSplitsAssignmentSequence()).isEmpty();
@@ -427,7 +434,9 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
streaming,
null,
null,
- workExecutor)) {
+ workExecutor,
+ LeaseContext.DEFAULT)) {
+
Map<Long, String> partitionNameByIds =
waitUntilPartitions(zooKeeperClient, DEFAULT_TABLE_PATH);
enumerator.start();
@@ -543,7 +552,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
streaming,
null,
- null)) {
+ null,
+ LeaseContext.DEFAULT)) {
// test splits for same non-partitioned bucket, should assign to same task
TableBucket t1 = new TableBucket(tableId, 0);
@@ -656,7 +666,8 @@ class FlinkSourceEnumeratorTest extends FlinkTestBase {
streaming,
null,
lakeSource,
- workExecutor)) {
+ workExecutor,
+ LeaseContext.DEFAULT)) {
enumerator.start();
// Remove the hybrid partition to mock expire after enumerator start
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
index 2c30bf086..5273d12d1 100644
--- 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
@@ -18,6 +18,7 @@
package org.apache.fluss.flink.source.state;
import org.apache.fluss.flink.lake.split.LakeSnapshotAndFlussLogSplit;
+import org.apache.fluss.flink.source.reader.LeaseContext;
import org.apache.fluss.flink.source.split.HybridSnapshotLogSplit;
import org.apache.fluss.flink.source.split.LogSplit;
import org.apache.fluss.flink.source.split.SourceSplitBase;
@@ -81,7 +82,10 @@ class SourceEnumeratorStateSerializerTest {
SourceEnumeratorState sourceEnumeratorState =
new SourceEnumeratorState(
- assignedBuckets, assignedPartitions, remainingHybridLakeFlussSplits);
+ assignedBuckets,
+ assignedPartitions,
+ remainingHybridLakeFlussSplits,
+ "leaseId");
// serialize state with remaining hybrid lake fluss splits
byte[] serialized = serializer.serialize(sourceEnumeratorState);
@@ -107,7 +111,11 @@ class SourceEnumeratorStateSerializerTest {
assignedPartitions.put(1L, "partition1");
assignedPartitions.put(2L, "partition2");
SourceEnumeratorState sourceEnumeratorState =
- new SourceEnumeratorState(assignedBuckets, assignedPartitions, null);
+ new SourceEnumeratorState(
+ assignedBuckets,
+ assignedPartitions,
+ null,
+ LeaseContext.DEFAULT.getKvSnapshotLeaseId());
byte[] serialized = serializer.serializeV0(sourceEnumeratorState);
// then deserialize
@@ -124,7 +132,10 @@ class SourceEnumeratorStateSerializerTest {
remainingHybridLakeFlussSplits.add(logSplit);
sourceEnumeratorState =
new SourceEnumeratorState(
- assignedBuckets, assignedPartitions, remainingHybridLakeFlussSplits);
+ assignedBuckets,
+ assignedPartitions,
+ remainingHybridLakeFlussSplits,
+ LeaseContext.DEFAULT.getKvSnapshotLeaseId());
serialized = serializer.serializeV0(sourceEnumeratorState);
@@ -145,7 +156,11 @@ class SourceEnumeratorStateSerializerTest {
assignedPartitions.put(1L, "partition1");
assignedPartitions.put(2L, "partition2");
SourceEnumeratorState sourceEnumeratorState =
- new SourceEnumeratorState(assignedBuckets, assignedPartitions, null);
+ new SourceEnumeratorState(
+ assignedBuckets,
+ assignedPartitions,
+ null,
+ LeaseContext.DEFAULT.getKvSnapshotLeaseId());
byte[] serialized = serializer.serialize(sourceEnumeratorState);
// test deserialize with nonnull lake source
diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminGateway.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminGateway.java
index 3885bb387..d8ef3fde6 100644
--- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminGateway.java
+++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/AdminGateway.java
@@ -17,6 +17,8 @@
package org.apache.fluss.rpc.gateway;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.AddServerTagRequest;
import org.apache.fluss.rpc.messages.AddServerTagResponse;
import org.apache.fluss.rpc.messages.AlterClusterConfigsRequest;
@@ -39,6 +41,8 @@ import org.apache.fluss.rpc.messages.DropAclsRequest;
import org.apache.fluss.rpc.messages.DropAclsResponse;
import org.apache.fluss.rpc.messages.DropDatabaseRequest;
import org.apache.fluss.rpc.messages.DropDatabaseResponse;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.DropPartitionRequest;
import org.apache.fluss.rpc.messages.DropPartitionResponse;
import org.apache.fluss.rpc.messages.DropTableRequest;
@@ -51,6 +55,8 @@ import org.apache.fluss.rpc.messages.RebalanceRequest;
import org.apache.fluss.rpc.messages.RebalanceResponse;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsRequest;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsResponse;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.RemoveServerTagRequest;
import org.apache.fluss.rpc.messages.RemoveServerTagResponse;
import org.apache.fluss.rpc.protocol.ApiKeys;
@@ -186,6 +192,18 @@ public interface AdminGateway extends AdminReadOnlyGateway {
CompletableFuture<DeleteProducerOffsetsResponse> deleteProducerOffsets(
DeleteProducerOffsetsRequest request);
+ @RPC(api = ApiKeys.ACQUIRE_KV_SNAPSHOT_LEASE)
+ CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+ AcquireKvSnapshotLeaseRequest request);
+
+ @RPC(api = ApiKeys.RELEASE_KV_SNAPSHOT_LEASE)
+ CompletableFuture<ReleaseKvSnapshotLeaseResponse> releaseKvSnapshotLease(
+ ReleaseKvSnapshotLeaseRequest request);
+
+ @RPC(api = ApiKeys.DROP_KV_SNAPSHOT_LEASE)
+ CompletableFuture<DropKvSnapshotLeaseResponse> dropKvSnapshotLease(
+ DropKvSnapshotLeaseRequest request);
+
// todo: rename table & alter table
}
diff --git a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java
index 8aecca25f..5084d73fd 100644
--- a/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java
+++ b/fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java
@@ -97,7 +97,10 @@ public enum ApiKeys {
PREPARE_LAKE_TABLE_SNAPSHOT(1052, 0, 0, PRIVATE),
REGISTER_PRODUCER_OFFSETS(1053, 0, 0, PUBLIC),
GET_PRODUCER_OFFSETS(1054, 0, 0, PUBLIC),
- DELETE_PRODUCER_OFFSETS(1055, 0, 0, PUBLIC);
+ DELETE_PRODUCER_OFFSETS(1055, 0, 0, PUBLIC),
+ ACQUIRE_KV_SNAPSHOT_LEASE(1056, 0, 0, PUBLIC),
+ RELEASE_KV_SNAPSHOT_LEASE(1057, 0, 0, PUBLIC),
+ DROP_KV_SNAPSHOT_LEASE(1058, 0, 0, PUBLIC);
private static final Map<Integer, ApiKeys> ID_TO_TYPE =
Arrays.stream(ApiKeys.values())
diff --git a/fluss-rpc/src/main/proto/FlussApi.proto b/fluss-rpc/src/main/proto/FlussApi.proto
index 99cdf3238..19610ea59 100644
--- a/fluss-rpc/src/main/proto/FlussApi.proto
+++ b/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 = 1;
+}
+
+message ReleaseKvSnapshotLeaseRequest {
+ required string lease_id = 1;
+ repeated PbTableBucket buckets_to_release = 2;
+}
+
+message ReleaseKvSnapshotLeaseResponse {}
+
+message DropKvSnapshotLeaseRequest {
+ required string lease_id = 1;
+}
+
+message DropKvSnapshotLeaseResponse {}
+
message GetLatestLakeSnapshotRequest {
required PbTablePath table_path = 1;
}
@@ -1138,4 +1161,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 bucket_snapshots = 2;
+}
+
+message PbKvSnapshotLeaseForBucket {
+ optional int64 partition_id = 1;
+ required int32 bucket_id = 2;
+ required int64 snapshot_id = 3;
}
\ No newline at end of file
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManager.java
index b7c92289b..1b739f9a2 100644
--- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManager.java
+++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManager.java
@@ -26,6 +26,7 @@ import org.apache.fluss.server.kv.snapshot.CompletedSnapshot;
import org.apache.fluss.server.kv.snapshot.CompletedSnapshotHandle;
import org.apache.fluss.server.kv.snapshot.CompletedSnapshotHandleStore;
import org.apache.fluss.server.kv.snapshot.CompletedSnapshotStore;
+import org.apache.fluss.server.kv.snapshot.CompletedSnapshotStore.SnapshotInUseChecker;
import org.apache.fluss.server.kv.snapshot.SharedKvFileRegistry;
import org.apache.fluss.server.kv.snapshot.ZooKeeperCompletedSnapshotHandleStore;
import org.apache.fluss.server.metrics.group.CoordinatorMetricGroup;
@@ -64,19 +65,22 @@ public class CompletedSnapshotStoreManager {
private final Executor ioExecutor;
private final Function<ZooKeeperClient, CompletedSnapshotHandleStore>
makeZookeeperCompletedSnapshotHandleStore;
+ private final SnapshotInUseChecker snapshotInUseChecker;
private final CoordinatorMetricGroup coordinatorMetricGroup;
public CompletedSnapshotStoreManager(
int maxNumberOfSnapshotsToRetain,
Executor ioExecutor,
ZooKeeperClient zooKeeperClient,
- CoordinatorMetricGroup coordinatorMetricGroup) {
+ CoordinatorMetricGroup coordinatorMetricGroup,
+ SnapshotInUseChecker snapshotInUseChecker) {
this(
maxNumberOfSnapshotsToRetain,
ioExecutor,
zooKeeperClient,
ZooKeeperCompletedSnapshotHandleStore::new,
- coordinatorMetricGroup);
+ coordinatorMetricGroup,
+ snapshotInUseChecker);
}
@VisibleForTesting
@@ -86,13 +90,15 @@ public class CompletedSnapshotStoreManager {
ZooKeeperClient zooKeeperClient,
Function<ZooKeeperClient, CompletedSnapshotHandleStore>
makeZookeeperCompletedSnapshotHandleStore,
- CoordinatorMetricGroup coordinatorMetricGroup) {
+ CoordinatorMetricGroup coordinatorMetricGroup,
+ SnapshotInUseChecker snapshotInUseChecker) {
checkArgument(
maxNumberOfSnapshotsToRetain > 0, "maxNumberOfSnapshotsToRetain must be positive");
this.maxNumberOfSnapshotsToRetain = maxNumberOfSnapshotsToRetain;
this.zooKeeperClient = zooKeeperClient;
this.bucketCompletedSnapshotStores = MapUtils.newConcurrentHashMap();
this.ioExecutor = ioExecutor;
+ this.snapshotInUseChecker = snapshotInUseChecker;
this.makeZookeeperCompletedSnapshotHandleStore = makeZookeeperCompletedSnapshotHandleStore;
this.coordinatorMetricGroup = coordinatorMetricGroup;
@@ -237,7 +243,8 @@ public class CompletedSnapshotStoreManager {
sharedKvFileRegistry,
retrievedSnapshots,
completedSnapshotHandleStore,
- ioExecutor);
+ ioExecutor,
+ snapshotInUseChecker);
}
@VisibleForTesting
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
index 813c1f10e..b7ef915d1 100644
--- 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
@@ -88,6 +88,7 @@ import org.apache.fluss.server.coordinator.event.SchemaChangeEvent;
import org.apache.fluss.server.coordinator.event.TableRegistrationChangeEvent;
import org.apache.fluss.server.coordinator.event.watcher.TableChangeWatcher;
import org.apache.fluss.server.coordinator.event.watcher.TabletServerChangeWatcher;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.coordinator.rebalance.RebalanceManager;
import org.apache.fluss.server.coordinator.statemachine.ReplicaLeaderElection.ControlledShutdownLeaderElection;
import org.apache.fluss.server.coordinator.statemachine.ReplicaLeaderElection.ReassignmentLeaderElection;
@@ -178,7 +179,6 @@ public class CoordinatorEventProcessor implements EventProcessor {
private final String internalListenerName;
private final CoordinatorMetricGroup coordinatorMetricGroup;
private final RebalanceManager rebalanceManager;
-
private final CompletedSnapshotStoreManager completedSnapshotStoreManager;
private final LakeTableHelper lakeTableHelper;
@@ -192,7 +192,8 @@ public class CoordinatorEventProcessor implements EventProcessor {
CoordinatorMetricGroup coordinatorMetricGroup,
Configuration conf,
ExecutorService ioExecutor,
- MetadataManager metadataManager) {
+ MetadataManager metadataManager,
+ KvSnapshotLeaseManager kvSnapshotLeaseManager) {
this.zooKeeperClient = zooKeeperClient;
this.serverMetadataCache = serverMetadataCache;
this.coordinatorChannelManager = coordinatorChannelManager;
@@ -230,12 +231,14 @@ public class CoordinatorEventProcessor implements EventProcessor {
this.coordinatorRequestBatch =
new CoordinatorRequestBatch(
coordinatorChannelManager, coordinatorEventManager, coordinatorContext);
+
this.completedSnapshotStoreManager =
new CompletedSnapshotStoreManager(
conf.getInt(ConfigOptions.KV_MAX_RETAINED_SNAPSHOTS),
ioExecutor,
zooKeeperClient,
- coordinatorMetricGroup);
+ coordinatorMetricGroup,
+ kvSnapshotLeaseManager::snapshotLeaseExist);
this.autoPartitionManager = autoPartitionManager;
this.lakeTableTieringManager = lakeTableTieringManager;
this.coordinatorMetricGroup = coordinatorMetricGroup;
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java
index 7e06f6109..ca3708bfe 100644
--- a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java
+++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java
@@ -33,6 +33,7 @@ import org.apache.fluss.server.DynamicConfigManager;
import org.apache.fluss.server.ServerBase;
import org.apache.fluss.server.authorizer.Authorizer;
import org.apache.fluss.server.authorizer.AuthorizerLoader;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.coordinator.rebalance.RebalanceManager;
import org.apache.fluss.server.metadata.CoordinatorMetadataCache;
import org.apache.fluss.server.metadata.ServerMetadataCache;
@@ -44,6 +45,8 @@ import org.apache.fluss.server.zk.data.CoordinatorAddress;
import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.KeeperException;
import org.apache.fluss.utils.ExceptionUtils;
import org.apache.fluss.utils.ExecutorUtils;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
import org.apache.fluss.utils.concurrent.FutureUtils;
@@ -85,6 +88,7 @@ public class CoordinatorServer extends ServerBase {
private final CompletableFuture<Result> terminationFuture;
private final AtomicBoolean isShutDown = new AtomicBoolean(false);
+ private final Clock clock;
@GuardedBy("lock")
private String serverId;
@@ -141,10 +145,18 @@ public class CoordinatorServer extends ServerBase {
@GuardedBy("lock")
private LakeCatalogDynamicLoader lakeCatalogDynamicLoader;
+ @GuardedBy("lock")
+ private KvSnapshotLeaseManager kvSnapshotLeaseManager;
+
public CoordinatorServer(Configuration conf) {
+ this(conf, SystemClock.getInstance());
+ }
+
+ public CoordinatorServer(Configuration conf, Clock clock) {
super(conf);
validateConfigs(conf);
this.terminationFuture = new CompletableFuture<>();
+ this.clock = clock;
}
public static void main(String[] args) {
@@ -196,6 +208,18 @@ public class CoordinatorServer extends ServerBase {
Executors.newFixedThreadPool(
conf.get(ConfigOptions.SERVER_IO_POOL_SIZE),
new ExecutorThreadFactory("coordinator-io"));
+
+ // Initialize and start the kv snapshot lease manager
+ this.kvSnapshotLeaseManager =
+ new KvSnapshotLeaseManager(
+ conf.get(ConfigOptions.KV_SNAPSHOT_LEASE_EXPIRATION_CHECK_INTERVAL)
+ .toMillis(),
+ zkClient,
+ conf.getString(ConfigOptions.REMOTE_DATA_DIR),
+ clock,
+ serverMetricGroup);
+ kvSnapshotLeaseManager.start();
+
this.coordinatorService =
new CoordinatorService(
conf,
@@ -208,7 +232,8 @@ public class CoordinatorServer extends ServerBase {
lakeCatalogDynamicLoader,
lakeTableTieringManager,
dynamicConfigManager,
- ioExecutor);
+ ioExecutor,
+ kvSnapshotLeaseManager);
this.rpcServer =
RpcServer.create(
@@ -250,7 +275,8 @@ public class CoordinatorServer extends ServerBase {
serverMetricGroup,
conf,
ioExecutor,
- metadataManager);
+ metadataManager,
+ kvSnapshotLeaseManager);
coordinatorEventProcessor.startup();
createDefaultDatabase();
@@ -449,6 +475,11 @@ public class CoordinatorServer extends ServerBase {
if (lakeCatalogDynamicLoader != null) {
lakeCatalogDynamicLoader.close();
}
+
+ if (kvSnapshotLeaseManager != null) {
+ kvSnapshotLeaseManager.close();
+ }
+
} catch (Throwable t) {
exception = ExceptionUtils.firstOrSuppressed(t, exception);
}
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
index b9340efc3..c0f3ea499 100644
--- 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
@@ -52,6 +52,8 @@ import org.apache.fluss.metadata.TableChange;
import org.apache.fluss.metadata.TableDescriptor;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.rpc.gateway.CoordinatorGateway;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.AddServerTagRequest;
import org.apache.fluss.rpc.messages.AddServerTagResponse;
import org.apache.fluss.rpc.messages.AdjustIsrRequest;
@@ -84,6 +86,8 @@ import org.apache.fluss.rpc.messages.DropAclsRequest;
import org.apache.fluss.rpc.messages.DropAclsResponse;
import org.apache.fluss.rpc.messages.DropDatabaseRequest;
import org.apache.fluss.rpc.messages.DropDatabaseResponse;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.DropPartitionRequest;
import org.apache.fluss.rpc.messages.DropPartitionResponse;
import org.apache.fluss.rpc.messages.DropTableRequest;
@@ -99,8 +103,10 @@ import org.apache.fluss.rpc.messages.MetadataResponse;
import org.apache.fluss.rpc.messages.PbAlterConfig;
import org.apache.fluss.rpc.messages.PbHeartbeatReqForTable;
import org.apache.fluss.rpc.messages.PbHeartbeatRespForTable;
+import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForTable;
import org.apache.fluss.rpc.messages.PbPrepareLakeTableRespForTable;
import org.apache.fluss.rpc.messages.PbProducerTableOffsets;
+import org.apache.fluss.rpc.messages.PbTableBucket;
import org.apache.fluss.rpc.messages.PbTableOffsets;
import org.apache.fluss.rpc.messages.PrepareLakeTableSnapshotRequest;
import org.apache.fluss.rpc.messages.PrepareLakeTableSnapshotResponse;
@@ -108,6 +114,8 @@ import org.apache.fluss.rpc.messages.RebalanceRequest;
import org.apache.fluss.rpc.messages.RebalanceResponse;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsRequest;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsResponse;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.RemoveServerTagRequest;
import org.apache.fluss.rpc.messages.RemoveServerTagResponse;
import org.apache.fluss.rpc.netty.server.Session;
@@ -135,6 +143,8 @@ import org.apache.fluss.server.coordinator.event.EventManager;
import org.apache.fluss.server.coordinator.event.ListRebalanceProgressEvent;
import org.apache.fluss.server.coordinator.event.RebalanceEvent;
import org.apache.fluss.server.coordinator.event.RemoveServerTagEvent;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseHandler;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.coordinator.producer.ProducerOffsetsManager;
import org.apache.fluss.server.coordinator.rebalance.goal.Goal;
import org.apache.fluss.server.entity.CommitKvSnapshotData;
@@ -183,11 +193,14 @@ import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toAclBindings;
import static org.apache.fluss.server.coordinator.rebalance.goal.GoalUtils.getGoalByType;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.addTableOffsetsToResponse;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.fromTablePath;
+import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getAcquireKvSnapshotLeaseData;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getAdjustIsrData;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getCommitLakeTableSnapshotData;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getCommitRemoteLogManifestData;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getPartitionSpec;
+import static org.apache.fluss.server.utils.ServerRpcMessageUtils.getReleaseKvSnapshotLeaseData;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.groupOffsetsByTableId;
+import static org.apache.fluss.server.utils.ServerRpcMessageUtils.makeAcquireKvSnapshotLeaseResponse;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.makeCreateAclsResponse;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.makeDropAclsResponse;
import static org.apache.fluss.server.utils.ServerRpcMessageUtils.toAlterTableConfigChanges;
@@ -216,6 +229,7 @@ public final class CoordinatorService extends RpcServiceBase implements Coordina
private final ExecutorService ioExecutor;
private final LakeTableHelper lakeTableHelper;
private final ProducerOffsetsManager producerOffsetsManager;
+ private final KvSnapshotLeaseManager kvSnapshotLeaseManager;
public CoordinatorService(
Configuration conf,
@@ -228,7 +242,8 @@ public final class CoordinatorService extends RpcServiceBase implements Coordina
LakeCatalogDynamicLoader lakeCatalogDynamicLoader,
LakeTableTieringManager lakeTableTieringManager,
DynamicConfigManager dynamicConfigManager,
- ExecutorService ioExecutor) {
+ ExecutorService ioExecutor,
+ KvSnapshotLeaseManager kvSnapshotLeaseManager) {
super(
remoteFileSystem,
ServerType.COORDINATOR,
@@ -255,6 +270,8 @@ public final class CoordinatorService extends RpcServiceBase implements Coordina
// Initialize and start the producer snapshot manager
this.producerOffsetsManager = new ProducerOffsetsManager(conf, zkClient);
this.producerOffsetsManager.start();
+
+ this.kvSnapshotLeaseManager = kvSnapshotLeaseManager;
}
@Override
@@ -856,6 +873,108 @@ public final class CoordinatorService extends RpcServiceBase implements Coordina
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) {
+ String leaseId = request.getLeaseId();
+ // Capture session before entering async block since currentSession() is thread-local
+ Session session = authorizer != null ? currentSession() : null;
+ return CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ // Authorization: require READ permission on all tables acquired by this
+ // lease.
+ if (authorizer != null) {
+ Optional<KvSnapshotLeaseHandler> leaseHandlerOpt =
+ kvSnapshotLeaseManager.getLease(leaseId);
+
+ if (leaseHandlerOpt.isPresent()) {
+ KvSnapshotLeaseHandler leaseHandler = leaseHandlerOpt.get();
+ Set<Long> tableIds = leaseHandler.getTableIdToTableLease().keySet();
+ // Check WRITE permission for each table.
+ for (Long tableId : tableIds) {
+ authorizeTableWithSession(session, OperationType.READ, tableId);
+ }
+ }
+ }
+
+ DropKvSnapshotLeaseResponse response = new DropKvSnapshotLeaseResponse();
+ kvSnapshotLeaseManager.dropLease(leaseId);
+ 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 drop kv snapshot lease for" + leaseId, e);
+ }
+ },
+ ioExecutor);
+ }
+
@Override
public CompletableFuture<AlterClusterConfigsResponse> alterClusterConfigs(
AlterClusterConfigsRequest request) {
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseHandler.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseHandler.java
new file mode 100644
index 000000000..8570de048
--- /dev/null
+++ b/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 smaller than
+ * maxBucketNum, the array will be dynamically expanded to {@code maxBucketNum}.
+ *
+ * @param tableBucket table bucket
+ * @param snapshotId snapshot id
+ * @param maxBucketNum the max bucket num
+ * @return the original registered snapshotId. if -1 means the bucket is new registered
+ */
+ public long acquireBucket(TableBucket tableBucket, long snapshotId, int maxBucketNum) {
+ 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[maxBucketNum];
+ Arrays.fill(array, -1L);
+ return new KvSnapshotTableLease(tableId, array);
+ });
+ bucketSnapshot = tableLease.getBucketSnapshots();
+ // Dynamically expand the array if the maxBucketNum exceeds the current array size.
+ // This can happen when new buckets are added to an existing table.
+ if (bucketSnapshot != null && maxBucketNum > bucketSnapshot.length) {
+ bucketSnapshot = expandArray(bucketSnapshot, maxBucketNum);
+ tableLease.setBucketSnapshots(bucketSnapshot);
+ }
+ } else {
+ // For partitioned table.
+
+ // first add partition to table.
+ KvSnapshotTableLease tableLease =
+ tableIdToTableLease.computeIfAbsent(
+ tableId, k -> new KvSnapshotTableLease(tableId));
+ Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+ // then add bucket to partition.
+ bucketSnapshot =
+ partitionSnapshots.computeIfAbsent(
+ partitionId,
+ k -> {
+ Long[] array = new Long[maxBucketNum];
+ Arrays.fill(array, -1L);
+ return array;
+ });
+ // Dynamically expand the array if the maxBucketNum exceeds the current array size.
+ if (maxBucketNum > bucketSnapshot.length) {
+ bucketSnapshot = expandArray(bucketSnapshot, maxBucketNum);
+ partitionSnapshots.put(partitionId, bucketSnapshot);
+ }
+ }
+
+ if (bucketSnapshot == null) {
+ throw new IllegalArgumentException(
+ "Bucket snapshot array is null. This may indicate a conflict between "
+ + "partitioned and non-partitioned usage for the same table ID.");
+ }
+
+ long originalSnapshotId = bucketSnapshot[bucketId];
+ bucketSnapshot[bucketId] = snapshotId;
+ return originalSnapshotId;
+ }
+
+ /**
+ * Release a bucket from the lease id.
+ *
+ * @param tableBucket table bucket
+ * @return the snapshot id of the unregistered bucket, or -1 if the bucket was never registered
+ * or the bucket id exceeds the current array size
+ */
+ public long releaseBucket(TableBucket tableBucket) {
+ Long[] bucketIndex;
+ long tableId = tableBucket.getTableId();
+ Long partitionId = tableBucket.getPartitionId();
+ int bucketId = tableBucket.getBucket();
+ KvSnapshotTableLease tableLease = tableIdToTableLease.get(tableId);
+ if (partitionId == null) {
+ // For none-partitioned table.
+ bucketIndex = tableLease.getBucketSnapshots();
+ } else {
+ // For partitioned table.
+ bucketIndex = tableLease.getBucketSnapshots(partitionId);
+ }
+
+ Long snapshotId = -1L;
+ if (bucketIndex != null) {
+ // The bucket id exceeds the current array size, meaning it was never registered
+ // under this lease. Return -1 directly.
+ if (bucketId >= bucketIndex.length) {
+ return -1L;
+ }
+
+ snapshotId = bucketIndex[bucketId];
+ bucketIndex[bucketId] = -1L;
+
+ boolean needRemove = true;
+ for (Long bucket : bucketIndex) {
+ if (bucket != -1L) {
+ needRemove = false;
+ break;
+ }
+ }
+
+ if (needRemove) {
+ if (partitionId == null) {
+ tableIdToTableLease.remove(tableId);
+ } else {
+ Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+ partitionSnapshots.remove(partitionId);
+ if (partitionSnapshots.isEmpty()) {
+ tableIdToTableLease.remove(tableId);
+ }
+ }
+ }
+ }
+ return snapshotId;
+ }
+
+ public boolean isEmpty() {
+ return tableIdToTableLease.isEmpty();
+ }
+
+ /**
+ * Expand the given array to the specified new size, filling new slots with -1L.
+ *
+ * @param original the original array
+ * @param newSize the desired new size (must be greater than original.length)
+ * @return a new expanded array with original values preserved
+ */
+ private Long[] expandArray(Long[] original, int newSize) {
+ Long[] expanded = Arrays.copyOf(original, newSize);
+ Arrays.fill(expanded, original.length, newSize, -1L);
+ return expanded;
+ }
+
+ public int getLeasedSnapshotCount() {
+ int count = 0;
+ for (KvSnapshotTableLease tableLease : tableIdToTableLease.values()) {
+ count += tableLease.getLeasedSnapshotCount();
+ }
+ return count;
+ }
+
+ @Override
+ public String toString() {
+ return "KvSnapshotLease{"
+ + "expirationTime="
+ + expirationTime
+ + ", tableIdToTableLease="
+ + tableIdToTableLease
+ + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof KvSnapshotLeaseHandler)) {
+ return false;
+ }
+ KvSnapshotLeaseHandler that = (KvSnapshotLeaseHandler) o;
+ return expirationTime == that.expirationTime
+ && Objects.equals(tableIdToTableLease, that.tableIdToTableLease);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(expirationTime, tableIdToTableLease);
+ }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java
new file mode 100644
index 000000000..7ae131b0a
--- /dev/null
+++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManager.java
@@ -0,0 +1,489 @@
+/*
+ * 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.metadata.TablePartition;
+import org.apache.fluss.metrics.MetricNames;
+import org.apache.fluss.server.kv.snapshot.CompletedSnapshot;
+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 snapshotLeaseExist(CompletedSnapshot snapshot) {
+ TableBucketSnapshot tbs =
+ new TableBucketSnapshot(snapshot.getTableBucket(), snapshot.getSnapshotID());
+ return snapshotLeaseExist(tbs);
+ }
+
+ public boolean snapshotLeaseExist(TableBucketSnapshot tbs) {
+ return inReadLock(
+ managerLock,
+ () -> {
+ AtomicInteger count = refCount.get(tbs);
+ 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<TableBucketSnapshot>> tableIdToLeaseBucket)
+ throws Exception {
+ return inWriteLock(
+ managerLock,
+ () -> {
+ // 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;
+ KvSnapshotLeaseHandler kvSnapshotLeaseHandle =
+ 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 KvSnapshotLeaseHandler(newExpirationTime);
+ } else {
+ existingLease.setExpirationTime(newExpirationTime);
+ return existingLease;
+ }
+ });
+
+ Map<Long, Integer> maxBucketNumMap = new HashMap<>();
+ Map<TablePartition, Integer> maxBucketNumOfPtMap = new HashMap<>();
+ findMaxBucketNum(tableIdToLeaseBucket, maxBucketNumMap, maxBucketNumOfPtMap);
+
+ for (Map.Entry<Long, List<TableBucketSnapshot>> entry :
+ tableIdToLeaseBucket.entrySet()) {
+ List<TableBucketSnapshot> buckets = entry.getValue();
+ for (TableBucketSnapshot bucket : buckets) {
+ TableBucket tableBucket = bucket.getTableBucket();
+ long kvSnapshotId = bucket.getSnapshotId();
+ // TODO Check whether this snapshot exists, if not add it to the
+ // unavailable snapshots. Trace by:
+ // https://github.com/apache/fluss/issues/2600
+
+ int maxBucketNum;
+ if (tableBucket.getPartitionId() == null) {
+ maxBucketNum = maxBucketNumMap.get(tableBucket.getTableId());
+ } else {
+ maxBucketNum =
+ maxBucketNumOfPtMap.get(
+ new TablePartition(
+ tableBucket.getTableId(),
+ tableBucket.getPartitionId()));
+ }
+ long originalSnapshotId =
+ kvSnapshotLeaseHandle.acquireBucket(
+ tableBucket, kvSnapshotId, maxBucketNum);
+ if (originalSnapshotId == -1L) {
+ leasedBucketCount.incrementAndGet();
+ } else {
+ // clear the original ref.
+ decrementRefCount(
+ new TableBucketSnapshot(tableBucket, originalSnapshotId));
+ }
+ incrementRefCount(bucket);
+ }
+ }
+
+ if (update) {
+ metadataManager.updateLease(leaseId, kvSnapshotLeaseHandle);
+ } else {
+ metadataManager.registerLease(leaseId, kvSnapshotLeaseHandle);
+ }
+
+ return unavailableSnapshots;
+ });
+ }
+
+ /**
+ * Release kv snapshot lease.
+ *
+ * @param leaseId the lease id
+ * @param tableBucketsToRelease the table buckets to release
+ */
+ public void release(String leaseId, List<TableBucket> tableBucketsToRelease) throws Exception {
+ inWriteLock(
+ managerLock,
+ () -> {
+ KvSnapshotLeaseHandler lease = kvSnapshotLeaseMap.get(leaseId);
+ if (lease == null) {
+ return;
+ }
+
+ for (TableBucket bucket : tableBucketsToRelease) {
+ long snapshotId = lease.releaseBucket(bucket);
+ if (snapshotId != -1L) {
+ leasedBucketCount.decrementAndGet();
+ decrementRefCount(new TableBucketSnapshot(bucket, snapshotId));
+ }
+ }
+
+ if (lease.isEmpty()) {
+ dropLease(leaseId);
+ } else {
+ metadataManager.updateLease(leaseId, lease);
+ }
+ });
+ }
+
+ /**
+ * Get kv snapshot lease.
+ *
+ * @param leaseId the lease id
+ * @return the kv snapshot lease handle
+ */
+ public Optional<KvSnapshotLeaseHandler> getLease(String leaseId) throws Exception {
+ return inReadLock(managerLock, () -> metadataManager.getLease(leaseId));
+ }
+
+ /**
+ * Drop kv snapshot lease.
+ *
+ * @param leaseId the lease id
+ * @return true if clear success, false if lease not exist
+ */
+ public boolean dropLease(String leaseId) throws Exception {
+ return inWriteLock(
+ managerLock,
+ () -> {
+ KvSnapshotLeaseHandler kvSnapshotLeasehandle =
+ kvSnapshotLeaseMap.remove(leaseId);
+ if (kvSnapshotLeasehandle == null) {
+ return false;
+ }
+
+ clearRefCount(kvSnapshotLeasehandle);
+ metadataManager.deleteLease(leaseId);
+
+ LOG.info("kv snapshots of lease '{}' has been all released.", leaseId);
+ return true;
+ });
+ }
+
+ private void initializeRefCount(KvSnapshotLeaseHandler 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 TableBucketSnapshot(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 TableBucketSnapshot(
+ new TableBucket(tableId, partitionId, i), snapshots[i]));
+ }
+ }
+ }
+ }
+ }
+
+ private void clearRefCount(KvSnapshotLeaseHandler 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 TableBucketSnapshot(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 TableBucketSnapshot(
+ new TableBucket(tableId, partitionId, i), snapshots[i]));
+ leasedBucketCount.decrementAndGet();
+ }
+ }
+ }
+ }
+ }
+
+ private void incrementRefCount(TableBucketSnapshot tableBucketSnapshot) {
+ refCount.computeIfAbsent(tableBucketSnapshot, k -> new AtomicInteger(0)).incrementAndGet();
+ }
+
+ private void decrementRefCount(TableBucketSnapshot tableBucketSnapshot) {
+ refCount.computeIfPresent(
+ tableBucketSnapshot,
+ (k, v) -> {
+ int newCount = v.decrementAndGet();
+ return newCount <= 0 ? null : v;
+ });
+ }
+
+ private void expireLeases() {
+ long currentTime = clock.milliseconds();
+ // 1. First collect all expired lease IDs under read lock
+ List<String> expiredLeaseIds =
+ inReadLock(
+ managerLock,
+ () ->
+ kvSnapshotLeaseMap.entrySet().stream()
+ .filter(
+ entry ->
+ entry.getValue().getExpirationTime()
+ < currentTime)
+ .map(Map.Entry::getKey)
+ .collect(Collectors.toList()));
+
+ // 2. Then process each collected ID (dropLease acquires write lock)
+ expiredLeaseIds.forEach(
+ leaseId -> {
+ try {
+ dropLease(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
+ static void findMaxBucketNum(
+ Map<Long, List<TableBucketSnapshot>> tableIdToLeaseBucket,
+ Map<Long, Integer> maxBucketNum,
+ Map<TablePartition, Integer> maxBucketNumOfPartitionedTable) {
+ for (Map.Entry<Long, List<TableBucketSnapshot>> entry : tableIdToLeaseBucket.entrySet()) {
+ List<TableBucketSnapshot> buckets = entry.getValue();
+ for (TableBucketSnapshot bucket : buckets) {
+ TableBucket tableBucket = bucket.getTableBucket();
+ Long partitionId = tableBucket.getPartitionId();
+ if (partitionId == null) {
+ maxBucketNum.merge(
+ tableBucket.getTableId(), tableBucket.getBucket() + 1, Math::max);
+ } else {
+ maxBucketNumOfPartitionedTable.merge(
+ new TablePartition(tableBucket.getTableId(), partitionId),
+ tableBucket.getBucket() + 1,
+ Math::max);
+ }
+ }
+ }
+ }
+
+ public void close() {
+ if (scheduledExecutor != null) {
+ scheduledExecutor.shutdown();
+ }
+ }
+
+ @VisibleForTesting
+ int getLeaseCount() {
+ return inReadLock(managerLock, kvSnapshotLeaseMap::size);
+ }
+
+ @VisibleForTesting
+ int getLeasedBucketCount() {
+ return leasedBucketCount.get();
+ }
+
+ @VisibleForTesting
+ int getRefCount(TableBucketSnapshot tableBucketSnapshot) {
+ return inReadLock(
+ managerLock,
+ () -> {
+ AtomicInteger count = refCount.get(tableBucketSnapshot);
+ return count == null ? 0 : count.get();
+ });
+ }
+
+ @VisibleForTesting
+ KvSnapshotLeaseHandler getKvSnapshotLeaseData(String leaseId) {
+ return inReadLock(managerLock, () -> kvSnapshotLeaseMap.get(leaseId));
+ }
+
+ @VisibleForTesting
+ KvSnapshotLeaseMetadataManager getMetadataManager() {
+ return metadataManager;
+ }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseMetadataManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseMetadataManager.java
new file mode 100644
index 000000000..d5333885f
--- /dev/null
+++ b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseMetadataManager.java
@@ -0,0 +1,281 @@
+/*
+ * 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.fs.FSDataInputStream;
+import org.apache.fluss.fs.FSDataOutputStream;
+import org.apache.fluss.fs.FileSystem;
+import org.apache.fluss.fs.FsPath;
+import org.apache.fluss.server.zk.ZooKeeperClient;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLeaseJsonSerde;
+import org.apache.fluss.utils.FlussPaths;
+import org.apache.fluss.utils.IOUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.util.ArrayList;
+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 KvSnapshotLeaseHandler} 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.
+ *
+ * <p>Follows the createdFiles pattern (similar to {@code ProducerOffsetsStore.tryStoreOffsets})
+ * to track files created in this operation. On failure, only the newly created files are
+ * cleaned up, avoiding accidental removal of previously successful metadata files.
+ *
+ * @param leaseId the lease id.
+ * @param lease the kv snapshot lease.
+ */
+ public void registerLease(String leaseId, KvSnapshotLeaseHandler lease) throws Exception {
+ List<FsPath> createdFiles = new ArrayList<>();
+ try {
+ Map<Long, FsPath> tableIdToRemoteMetadataFsPath =
+ generateMetadataFile(leaseId, lease, createdFiles);
+
+ // generate remote fsPath of metadata.
+ KvSnapshotLeaseMetadata leaseMetadata =
+ new KvSnapshotLeaseMetadata(
+ lease.getExpirationTime(), tableIdToRemoteMetadataFsPath);
+
+ // register kv snapshot metadata to zk.
+ zkClient.registerKvSnapshotLeaseMetadata(leaseId, leaseMetadata);
+ } catch (Exception e) {
+ LOG.warn(
+ "Failed to register kv snapshot lease metadata for lease {}, "
+ + "cleaning up {} created files.",
+ leaseId,
+ createdFiles.size(),
+ e);
+ cleanupFilesSafely(createdFiles);
+ throw e;
+ }
+ }
+
+ /**
+ * Update a kv snapshot lease to zk and remote fs.
+ *
+ * <p>Follows the createdFiles pattern (similar to {@code ProducerOffsetsStore.tryStoreOffsets})
+ * to track files created in this operation. On failure, only the newly created files are
+ * cleaned up, avoiding accidental removal of previously successful metadata files.
+ *
+ * @param leaseId the lease id.
+ * @param kvSnapshotLeaseHandle the kv snapshot lease.
+ */
+ public void updateLease(String leaseId, KvSnapshotLeaseHandler kvSnapshotLeaseHandle)
+ throws Exception {
+ // TODO change this to incremental update to avoid create too many remote metadata files.
+ Optional<KvSnapshotLeaseMetadata> originalLeaseMetadata;
+ List<FsPath> createdFiles = new ArrayList<>();
+ try {
+ originalLeaseMetadata = zkClient.getKvSnapshotLeaseMetadata(leaseId);
+
+ Map<Long, FsPath> tableIdToNewRemoteMetadataFsPath =
+ generateMetadataFile(leaseId, kvSnapshotLeaseHandle, createdFiles);
+
+ // generate new kv snapshot lease metadata.
+ KvSnapshotLeaseMetadata newLeaseMetadata =
+ new KvSnapshotLeaseMetadata(
+ kvSnapshotLeaseHandle.getExpirationTime(),
+ tableIdToNewRemoteMetadataFsPath);
+
+ zkClient.updateKvSnapshotLeaseMetadata(leaseId, newLeaseMetadata);
+ } catch (Exception e) {
+ LOG.warn(
+ "Failed to update kv snapshot lease metadata for lease {}, "
+ + "cleaning up {} created files.",
+ leaseId,
+ createdFiles.size(),
+ e);
+ cleanupFilesSafely(createdFiles);
+ throw e;
+ }
+
+ // Best-effort cleanup of old metadata files. Failures are logged but not propagated,
+ // as the ZK metadata has already been updated to point to the new files.
+ try {
+ originalLeaseMetadata.ifPresent(this::discardLeaseMetadata);
+ } catch (Exception e) {
+ LOG.warn(
+ "Failed to discard original lease metadata for lease {}, "
+ + "orphaned files may exist.",
+ leaseId,
+ e);
+ }
+ }
+
+ /**
+ * Get a kv snapshot lease from zk and remote fs.
+ *
+ * @param leaseId the lease id.
+ * @return the kv snapshot lease.
+ */
+ public Optional<KvSnapshotLeaseHandler> getLease(String leaseId) throws Exception {
+ Optional<KvSnapshotLeaseMetadata> kvSnapshotLeaseMetadataOpt =
+ zkClient.getKvSnapshotLeaseMetadata(leaseId);
+ if (!kvSnapshotLeaseMetadataOpt.isPresent()) {
+ return Optional.empty();
+ }
+
+ KvSnapshotLeaseMetadata kvSnapshotLeaseMetadata = kvSnapshotLeaseMetadataOpt.get();
+ KvSnapshotLeaseHandler kvSnapshotLeasehandle =
+ buildKvSnapshotLease(kvSnapshotLeaseMetadata);
+ return Optional.of(kvSnapshotLeasehandle);
+ }
+
+ /**
+ * Delete a kv snapshot lease from zk and remote fs.
+ *
+ * @param leaseId the lease id.
+ */
+ public void deleteLease(String leaseId) throws Exception {
+ Optional<KvSnapshotLeaseMetadata> leaseMetadataOpt =
+ zkClient.getKvSnapshotLeaseMetadata(leaseId);
+
+ // delete zk metadata.
+ zkClient.deleteKvSnapshotLease(leaseId);
+
+ // Best-effort cleanup of remote metadata files. Failures are logged but not
+ // propagated, as the ZK metadata has already been deleted.
+ try {
+ leaseMetadataOpt.ifPresent(this::discardLeaseMetadata);
+ } catch (Exception e) {
+ LOG.warn(
+ "Failed to discard lease metadata for lease {}, " + "orphaned files may exist.",
+ leaseId,
+ e);
+ }
+ }
+
+ /**
+ * Generate metadata files for all tables in the lease. Each created file is tracked in the
+ * provided {@code createdFiles} list for safe cleanup on failure.
+ */
+ private Map<Long, FsPath> generateMetadataFile(
+ String leaseId, KvSnapshotLeaseHandler lease, List<FsPath> createdFiles)
+ throws Exception {
+ Map<Long, FsPath> tableIdToMetadataFile = new HashMap<>();
+ for (Map.Entry<Long, KvSnapshotTableLease> entry :
+ lease.getTableIdToTableLease().entrySet()) {
+ long tableId = entry.getKey();
+ FsPath path = generateMetadataFile(tableId, leaseId, entry.getValue());
+ createdFiles.add(path);
+ tableIdToMetadataFile.put(tableId, path);
+ }
+ return tableIdToMetadataFile;
+ }
+
+ private FsPath generateMetadataFile(
+ long tableId, String leaseId, KvSnapshotTableLease tableLease) throws Exception {
+ // get the remote file path to store the kv snapshot lease metadata of a table
+ FsPath remoteKvSnapshotLeaseFile =
+ FlussPaths.remoteKvSnapshotLeaseFile(remoteDataDir, leaseId, tableId);
+ // check whether the parent directory exists, if not, create the directory
+ FileSystem fileSystem = remoteKvSnapshotLeaseFile.getFileSystem();
+ if (!fileSystem.exists(remoteKvSnapshotLeaseFile.getParent())) {
+ fileSystem.mkdirs(remoteKvSnapshotLeaseFile.getParent());
+ }
+
+ // serialize table lease to json bytes, and write to file.
+ byte[] jsonBytes = KvSnapshotTableLeaseJsonSerde.toJson(tableLease);
+ try (FSDataOutputStream outputStream =
+ fileSystem.create(remoteKvSnapshotLeaseFile, FileSystem.WriteMode.OVERWRITE)) {
+ outputStream.write(jsonBytes);
+ }
+ return remoteKvSnapshotLeaseFile;
+ }
+
+ private KvSnapshotLeaseHandler buildKvSnapshotLease(
+ KvSnapshotLeaseMetadata kvSnapshotLeaseMetadata) throws Exception {
+ Map<Long, FsPath> tableIdToRemoteMetadataFilePath =
+ kvSnapshotLeaseMetadata.getTableIdToRemoteMetadataFilePath();
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+ for (Map.Entry<Long, FsPath> entry : tableIdToRemoteMetadataFilePath.entrySet()) {
+ long tableId = entry.getKey();
+ FsPath remoteMetadataFilePath = entry.getValue();
+ tableIdToTableLease.put(tableId, buildKvSnapshotTableLease(remoteMetadataFilePath));
+ }
+ return new KvSnapshotLeaseHandler(
+ kvSnapshotLeaseMetadata.getExpirationTime(), tableIdToTableLease);
+ }
+
+ private KvSnapshotTableLease buildKvSnapshotTableLease(FsPath remoteMetadataFilePath)
+ throws Exception {
+ checkNotNull(remoteMetadataFilePath);
+ FSDataInputStream inputStream =
+ remoteMetadataFilePath.getFileSystem().open(remoteMetadataFilePath);
+ try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
+ IOUtils.copyBytes(inputStream, outputStream, true);
+ return KvSnapshotTableLeaseJsonSerde.fromJson(outputStream.toByteArray());
+ }
+ }
+
+ private void discardLeaseMetadata(KvSnapshotLeaseMetadata kvSnapshotLeaseMetadata) {
+ cleanupFilesSafely(
+ new ArrayList<>(
+ kvSnapshotLeaseMetadata.getTableIdToRemoteMetadataFilePath().values()));
+ }
+
+ /**
+ * Safely clean up a list of remote files. Each file deletion is independent - failures for
+ * individual files are logged but do not prevent cleanup of remaining files.
+ */
+ private void cleanupFilesSafely(List<FsPath> files) {
+ for (FsPath file : files) {
+ try {
+ FileSystem fileSystem = file.getFileSystem();
+ if (fileSystem.exists(file)) {
+ fileSystem.delete(file, false);
+ }
+ } catch (Exception e) {
+ LOG.warn(
+ "Error deleting remote file path of kv snapshot lease metadata at {}",
+ file,
+ e);
+ }
+ }
+ }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java
index b67f98466..cf398636b 100644
--- a/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java
+++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStore.java
@@ -31,11 +31,14 @@ import javax.annotation.concurrent.ThreadSafe;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
-import java.util.Collections;
import java.util.Deque;
+import java.util.HashMap;
+import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
+import java.util.Map;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.Executor;
import java.util.concurrent.locks.ReentrantLock;
@@ -66,6 +69,7 @@ public class CompletedSnapshotStore {
private final Executor ioExecutor;
private final SnapshotsCleaner snapshotsCleaner;
+ private final SnapshotInUseChecker snapshotInUseChecker;
private final ReentrantLock lock = new ReentrantLock();
@@ -75,17 +79,31 @@ public class CompletedSnapshotStore {
*/
private final ArrayDeque<CompletedSnapshot> completedSnapshots;
+ /**
+ * Snapshots that are still in use by a lease but have been moved out of the standard retention
+ * window. These snapshots are protected from deletion (both metadata and SST files) until their
+ * lease expires.
+ *
+ * <p>When a snapshot should be subsumed (beyond the retention window) but cannot be (because
+ * it's leased), it is moved here instead of staying in {@link #completedSnapshots}. This
+ * ensures the effective {@code lowestSnapshotID} is computed from retained snapshots only,
+ * allowing SST files from non-leased subsumed snapshots to be cleaned up properly.
+ */
+ @VisibleForTesting final Map<Long, CompletedSnapshot> stillInUseSnapshots = new HashMap<>();
+
public CompletedSnapshotStore(
int maxNumberOfSnapshotsToRetain,
SharedKvFileRegistry sharedKvFileRegistry,
Collection<CompletedSnapshot> completedSnapshots,
CompletedSnapshotHandleStore completedSnapshotHandleStore,
- Executor executor) {
+ Executor executor,
+ SnapshotInUseChecker snapshotInUseChecker) {
this.maxNumberOfSnapshotsToRetain = maxNumberOfSnapshotsToRetain;
this.sharedKvFileRegistry = sharedKvFileRegistry;
this.completedSnapshots = new ArrayDeque<>();
this.completedSnapshots.addAll(completedSnapshots);
this.completedSnapshotHandleStore = completedSnapshotHandleStore;
+ this.snapshotInUseChecker = snapshotInUseChecker;
this.ioExecutor = executor;
this.snapshotsCleaner = new SnapshotsCleaner();
}
@@ -103,7 +121,7 @@ public class CompletedSnapshotStore {
}
public long getNumSnapshots() {
- return inLock(lock, () -> completedSnapshots.size());
+ return inLock(lock, completedSnapshots::size);
}
/**
@@ -113,7 +131,7 @@ public class CompletedSnapshotStore {
* @param snapshot Completed snapshot to add.
*/
@VisibleForTesting
- CompletedSnapshot addSnapshotAndSubsumeOldestOne(
+ void addSnapshotAndSubsumeOldestOne(
final CompletedSnapshot snapshot,
SnapshotsCleaner snapshotsCleaner,
Runnable postCleanup)
@@ -128,62 +146,93 @@ public class CompletedSnapshotStore {
snapshot.getTableBucket(), snapshot.getSnapshotID(), completedSnapshotHandle);
// Now add the new one. If it fails, we don't want to lose existing data.
- return inLock(
+ inLock(
lock,
() -> {
completedSnapshots.addLast(snapshot);
// 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 -> {
+ if (snapshotInUseChecker.isInUse(completedSnapshot)) {
+ LOG.debug(
+ "Snapshot {} is still in use, move it to stillInUseSnapshots",
+ completedSnapshot.getSnapshotID());
+ stillInUseSnapshots.put(
+ completedSnapshot.getSnapshotID(), completedSnapshot);
+ } else {
+ remove(
+ completedSnapshot.getTableBucket(),
+ completedSnapshot.getSnapshotID());
+ snapshotsCleaner.addSubsumedSnapshot(completedSnapshot);
+ }
+ });
+
+ // Check if any previously still-in-use snapshots can now be released
+ // (lease expired).
+ removeUnusedSnapshots(snapshotsCleaner);
+
+ // 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);
});
}
+ private void removeUnusedSnapshots(SnapshotsCleaner snapshotsCleaner) throws Exception {
+ Iterator<Map.Entry<Long, CompletedSnapshot>> stillInUseIter =
+ stillInUseSnapshots.entrySet().iterator();
+ while (stillInUseIter.hasNext()) {
+ Map.Entry<Long, CompletedSnapshot> entry = stillInUseIter.next();
+ CompletedSnapshot snapshot = entry.getValue();
+ if (!snapshotInUseChecker.isInUse(snapshot)) {
+ stillInUseIter.remove();
+
+ remove(snapshot.getTableBucket(), snapshot.getSnapshotID());
+ snapshotsCleaner.addSubsumedSnapshot(snapshot);
+ LOG.debug(
+ "Released snapshot {} from stillInUseSnapshots (lease expired)",
+ snapshot.getSnapshotID());
+ }
+ }
+ }
+
+ /**
+ * Returns a list of all the completed snapshots in this store. Note this doesn't include the
+ * snapshots that are still in use but have been moved out of the standard retention window,
+ * which can be obtained from {@link #stillInUseSnapshots}.
+ */
+ @VisibleForTesting
public List<CompletedSnapshot> getAllSnapshots() {
return inLock(lock, () -> new ArrayList<>(completedSnapshots));
}
- private static Optional<CompletedSnapshot> subsume(
+ private static void subsume(
Deque<CompletedSnapshot> snapshots, int numRetain, SubsumeAction subsumeAction) {
if (snapshots.isEmpty()) {
- return Optional.empty();
+ return;
}
CompletedSnapshot latest = snapshots.peekLast();
- Optional<CompletedSnapshot> lastSubsumedSnapshot = Optional.empty();
Iterator<CompletedSnapshot> iterator = snapshots.iterator();
while (snapshots.size() > numRetain && iterator.hasNext()) {
CompletedSnapshot next = iterator.next();
if (canSubsume(next, latest)) {
- // always return the subsumed snapshot with larger snapshot id.
- if (!lastSubsumedSnapshot.isPresent()
- || next.getSnapshotID() > lastSubsumedSnapshot.get().getSnapshotID()) {
- lastSubsumedSnapshot = Optional.of(next);
- }
iterator.remove();
try {
subsumeAction.subsume(next);
@@ -192,7 +241,6 @@ public class CompletedSnapshotStore {
}
}
}
- return lastSubsumedSnapshot;
}
@FunctionalInterface
@@ -200,6 +248,12 @@ public class CompletedSnapshotStore {
void subsume(CompletedSnapshot snapshot) throws Exception;
}
+ /** A function to check whether a snapshot still in use. */
+ @FunctionalInterface
+ public interface SnapshotInUseChecker {
+ boolean isInUse(CompletedSnapshot snapshot);
+ }
+
private static boolean canSubsume(CompletedSnapshot next, CompletedSnapshot latest) {
// if the snapshot is equal to the latest snapshot, it means it can't be subsumed
if (next == latest) {
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistry.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistry.java
index 9d13d3e0e..558e847d9 100644
--- a/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistry.java
+++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistry.java
@@ -23,11 +23,13 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
+import java.util.Set;
import java.util.concurrent.Executor;
import static org.apache.fluss.utils.Preconditions.checkNotNull;
@@ -127,10 +129,31 @@ public class SharedKvFileRegistry implements AutoCloseable {
}
public void unregisterUnusedKvFile(long lowestSnapshotID) {
+ unregisterUnusedKvFile(lowestSnapshotID, Collections.emptySet());
+ }
+
+ /**
+ * Unregister and delete KV files that are no longer used by any retained or still-in-use
+ * (leased) snapshot.
+ *
+ * <p>A KV file is safe to delete only if:
+ *
+ * <ol>
+ * <li>Its {@code lastUsedSnapshotID < lowestSnapshotID} (not used by any retained snapshot)
+ * <li>No snapshot in {@code stillInUseSnapshotIds} falls within the file's referencing range
+ * [{@code createdBySnapshotID}, {@code lastUsedSnapshotID}]
+ * </ol>
+ *
+ * @param lowestSnapshotID the effective lowest snapshot ID from retained (non-leased) snapshots
+ * @param stillInUseSnapshotIds snapshot IDs that are protected by leases
+ */
+ public void unregisterUnusedKvFile(long lowestSnapshotID, Set<Long> stillInUseSnapshotIds) {
// delete kv files that aren't used
LOG.debug(
- "Discard kv files created before snapshot {} and not used afterwards",
- lowestSnapshotID);
+ "Discard kv files created before snapshot {} and not used afterwards, "
+ + "stillInUse snapshots: {}",
+ lowestSnapshotID,
+ stillInUseSnapshotIds);
List<KvFileHandle> subsumed = new ArrayList<>();
// Iterate over all the registered kv file handles.
// Using a simple loop and NOT index by snapshotID because:
@@ -140,7 +163,8 @@ public class SharedKvFileRegistry implements AutoCloseable {
Iterator<SharedKvEntry> it = registeredKvEntries.values().iterator();
while (it.hasNext()) {
SharedKvEntry entry = it.next();
- if (entry.lastUsedSnapshotID < lowestSnapshotID) {
+ if (entry.lastUsedSnapshotID < lowestSnapshotID
+ && !isReferencedByStillInUse(entry, stillInUseSnapshotIds)) {
subsumed.add(entry.kvFileHandle);
it.remove();
fileSize -= entry.kvFileHandle.getSize();
@@ -153,6 +177,23 @@ public class SharedKvFileRegistry implements AutoCloseable {
}
}
+ /**
+ * Check if a KV file entry is referenced by any still-in-use (leased) snapshot. A file is
+ * referenced by a snapshot if the snapshot ID falls within the file's usage range [{@code
+ * createdBySnapshotID}, {@code lastUsedSnapshotID}].
+ *
+ * <p>Note: {@code createdBySnapshotID} serves as the first snapshot ID that uses this file
+ * (i.e., firstUsedSnapshotID), while {@code lastUsedSnapshotID} is the last.
+ */
+ private boolean isReferencedByStillInUse(SharedKvEntry entry, Set<Long> stillInUseSnapshotIds) {
+ for (Long snapshotId : stillInUseSnapshotIds) {
+ if (snapshotId >= entry.createdBySnapshotID && snapshotId <= entry.lastUsedSnapshotID) {
+ return true;
+ }
+ }
+ return false;
+ }
+
public void registerAll(KvSnapshotHandle kvSnapshotHandle, long snapshotID) {
if (kvSnapshotHandle == null) {
return;
@@ -211,7 +252,17 @@ public class SharedKvFileRegistry implements AutoCloseable {
private static final class SharedKvEntry {
+ /**
+ * The snapshot ID that first created/registered this KV file. This effectively serves as
+ * {@code firstUsedSnapshotID} — the lower bound of the snapshot range referencing this
+ * file.
+ */
private final long createdBySnapshotID;
+
+ /**
+ * The snapshot ID that last referenced this KV file — the upper bound of the snapshot range
+ * referencing this file.
+ */
private long lastUsedSnapshotID;
/** The shared kv file handle. */
KvFileHandle kvFileHandle;
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
index 3a8dd5c24..0ecc719f1 100644
--- 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
@@ -34,6 +34,7 @@ import org.apache.fluss.metadata.PartitionSpec;
import org.apache.fluss.metadata.PhysicalTablePath;
import org.apache.fluss.metadata.ResolvedPartitionSpec;
import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableBucketSnapshot;
import org.apache.fluss.metadata.TableChange;
import org.apache.fluss.metadata.TableDescriptor;
import org.apache.fluss.metadata.TableInfo;
@@ -55,6 +56,8 @@ import org.apache.fluss.rpc.entity.LookupResultForBucket;
import org.apache.fluss.rpc.entity.PrefixLookupResultForBucket;
import org.apache.fluss.rpc.entity.ProduceLogResultForBucket;
import org.apache.fluss.rpc.entity.PutKvResultForBucket;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.AdjustIsrRequest;
import org.apache.fluss.rpc.messages.AdjustIsrResponse;
import org.apache.fluss.rpc.messages.AlterTableRequest;
@@ -107,6 +110,8 @@ import org.apache.fluss.rpc.messages.PbFetchLogRespForBucket;
import org.apache.fluss.rpc.messages.PbFetchLogRespForTable;
import org.apache.fluss.rpc.messages.PbKeyValue;
import org.apache.fluss.rpc.messages.PbKvSnapshot;
+import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForBucket;
+import org.apache.fluss.rpc.messages.PbKvSnapshotLeaseForTable;
import org.apache.fluss.rpc.messages.PbLakeSnapshotForBucket;
import org.apache.fluss.rpc.messages.PbLakeTableOffsetForBucket;
import org.apache.fluss.rpc.messages.PbLakeTableSnapshotInfo;
@@ -149,6 +154,7 @@ 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.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;
@@ -1968,6 +1974,68 @@ public class ServerRpcMessageUtils {
return pbRebalancePlanForBucket;
}
+ public static Map<Long, List<TableBucketSnapshot>> getAcquireKvSnapshotLeaseData(
+ AcquireKvSnapshotLeaseRequest request) {
+ Map<Long, List<TableBucketSnapshot>> tableIdToLeasedBucket = new HashMap<>();
+ for (PbKvSnapshotLeaseForTable snapshotToLease : request.getSnapshotsToLeasesList()) {
+ long tableId = snapshotToLease.getTableId();
+ List<TableBucketSnapshot> bucketList = new ArrayList<>();
+ for (PbKvSnapshotLeaseForBucket leaseForBucket :
+ snapshotToLease.getBucketSnapshotsList()) {
+ bucketList.add(getKvSnapshotLeaseForBucket(tableId, leaseForBucket));
+ }
+ tableIdToLeasedBucket.put(tableId, bucketList);
+ }
+ return tableIdToLeasedBucket;
+ }
+
+ public static List<TableBucket> getReleaseKvSnapshotLeaseData(
+ ReleaseKvSnapshotLeaseRequest request) {
+ List<TableBucket> bucketList = new ArrayList<>();
+ for (PbTableBucket pbTableBucket : request.getBucketsToReleasesList()) {
+ bucketList.add(
+ new TableBucket(
+ pbTableBucket.getTableId(),
+ pbTableBucket.hasPartitionId() ? pbTableBucket.getPartitionId() : null,
+ pbTableBucket.getBucketId()));
+ }
+ return bucketList;
+ }
+
+ public static AcquireKvSnapshotLeaseResponse makeAcquireKvSnapshotLeaseResponse(
+ Map<TableBucket, Long> unavailableSnapshots) {
+ AcquireKvSnapshotLeaseResponse response = new AcquireKvSnapshotLeaseResponse();
+ 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.addUnavailableSnapshot()
+ .setTableId(entry.getKey())
+ .addAllBucketSnapshots(entry.getValue());
+ }
+ return response;
+ }
+
+ private static TableBucketSnapshot getKvSnapshotLeaseForBucket(
+ long tableId, PbKvSnapshotLeaseForBucket leaseForBucket) {
+ return new TableBucketSnapshot(
+ new TableBucket(
+ tableId,
+ leaseForBucket.hasPartitionId() ? leaseForBucket.getPartitionId() : null,
+ leaseForBucket.getBucketId()),
+ leaseForBucket.getSnapshotId());
+ }
+
private static <T> Map<TableBucket, T> mergeResponse(
Map<TableBucket, T> response, Map<TableBucket, T> errors) {
if (errors.isEmpty()) {
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java
index 86f3b952f..ff0e4fde1 100644
--- a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java
+++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java
@@ -63,6 +63,8 @@ import org.apache.fluss.server.zk.data.ZkData.ConfigEntityZNode;
import org.apache.fluss.server.zk.data.ZkData.CoordinatorZNode;
import org.apache.fluss.server.zk.data.ZkData.DatabaseZNode;
import org.apache.fluss.server.zk.data.ZkData.DatabasesZNode;
+import org.apache.fluss.server.zk.data.ZkData.KvSnapshotLeaseZNode;
+import org.apache.fluss.server.zk.data.ZkData.KvSnapshotLeasesZNode;
import org.apache.fluss.server.zk.data.ZkData.LakeTableZNode;
import org.apache.fluss.server.zk.data.ZkData.LeaderAndIsrZNode;
import org.apache.fluss.server.zk.data.ZkData.PartitionIdZNode;
@@ -85,6 +87,7 @@ import org.apache.fluss.server.zk.data.ZkData.TablesZNode;
import org.apache.fluss.server.zk.data.ZkData.WriterIdZNode;
import org.apache.fluss.server.zk.data.lake.LakeTable;
import org.apache.fluss.server.zk.data.lake.LakeTableSnapshot;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata;
import org.apache.fluss.server.zk.data.producer.ProducerOffsets;
import org.apache.fluss.shaded.curator5.org.apache.curator.framework.CuratorFramework;
import org.apache.fluss.shaded.curator5.org.apache.curator.framework.api.BackgroundCallback;
@@ -1039,6 +1042,36 @@ public class ZooKeeperClient implements AutoCloseable {
return snapshots;
}
+ public List<String> getKvSnapshotLeasesList() throws Exception {
+ return getChildren(KvSnapshotLeasesZNode.path());
+ }
+
+ public void registerKvSnapshotLeaseMetadata(
+ String leaseId, KvSnapshotLeaseMetadata leaseMetadata) throws Exception {
+ String path = KvSnapshotLeaseZNode.path(leaseId);
+ zkClient.create()
+ .creatingParentsIfNeeded()
+ .withMode(CreateMode.PERSISTENT)
+ .forPath(path, KvSnapshotLeaseZNode.encode(leaseMetadata));
+ }
+
+ public void updateKvSnapshotLeaseMetadata(String leaseId, KvSnapshotLeaseMetadata leaseMetadata)
+ throws Exception {
+ String path = KvSnapshotLeaseZNode.path(leaseId);
+ zkClient.setData().forPath(path, KvSnapshotLeaseZNode.encode(leaseMetadata));
+ }
+
+ public Optional<KvSnapshotLeaseMetadata> getKvSnapshotLeaseMetadata(String leaseId)
+ throws Exception {
+ String path = KvSnapshotLeaseZNode.path(leaseId);
+ return getOrEmpty(path).map(KvSnapshotLeaseZNode::decode);
+ }
+
+ public void deleteKvSnapshotLease(String leaseId) throws Exception {
+ String path = KvSnapshotLeaseZNode.path(leaseId);
+ zkClient.delete().forPath(path);
+ }
+
// --------------------------------------------------------------------------------------------
// Writer
// --------------------------------------------------------------------------------------------
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/ZkData.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/ZkData.java
index 43c1acbf4..b1e0e9493 100644
--- a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/ZkData.java
+++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/ZkData.java
@@ -26,6 +26,8 @@ import org.apache.fluss.security.acl.Resource;
import org.apache.fluss.security.acl.ResourceType;
import org.apache.fluss.server.zk.data.lake.LakeTable;
import org.apache.fluss.server.zk.data.lake.LakeTableJsonSerde;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadataJsonSerde;
import org.apache.fluss.server.zk.data.producer.ProducerOffsets;
import org.apache.fluss.server.zk.data.producer.ProducerOffsetsJsonSerde;
import org.apache.fluss.utils.json.JsonSerdeUtils;
@@ -906,4 +908,38 @@ public final class ZkData {
return JsonSerdeUtils.readValue(json, ProducerOffsetsJsonSerde.INSTANCE);
}
}
+
+ // ------------------------------------------------------------------------------------------
+ // ZNodes for Consumers.
+ // ------------------------------------------------------------------------------------------
+
+ /** The root znode for leases. It will record all the info of fluss leases. */
+ public static final class LeasesNode {
+ public static String path() {
+ return "/leases";
+ }
+ }
+
+ /** The root znode for kv snapshot leases. */
+ public static final class KvSnapshotLeasesZNode {
+ public static String path() {
+ return LeasesNode.path() + "/kv_snapshot";
+ }
+ }
+
+ /** The znode for kv snapshot lease zk data. */
+ public static final class KvSnapshotLeaseZNode {
+ public static String path(String leaseId) {
+ return KvSnapshotLeasesZNode.path() + "/" + leaseId;
+ }
+
+ public static byte[] encode(KvSnapshotLeaseMetadata kvSnapshotLeaseMetadata) {
+ return JsonSerdeUtils.writeValueAsBytes(
+ kvSnapshotLeaseMetadata, KvSnapshotLeaseMetadataJsonSerde.INSTANCE);
+ }
+
+ public static KvSnapshotLeaseMetadata decode(byte[] json) {
+ return JsonSerdeUtils.readValue(json, KvSnapshotLeaseMetadataJsonSerde.INSTANCE);
+ }
+ }
}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadata.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadata.java
new file mode 100644
index 000000000..b6adff1cb
--- /dev/null
+++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadata.java
@@ -0,0 +1,71 @@
+/*
+ * 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 java.util.Map;
+import java.util.Objects;
+
+/** 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;
+ }
+
+ @Override
+ public String toString() {
+ return "KvSnapshotLeaseMetadata{"
+ + "expirationTime="
+ + expirationTime
+ + ", tableIdToRemoteMetadataFilePath="
+ + tableIdToRemoteMetadataFilePath
+ + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ KvSnapshotLeaseMetadata that = (KvSnapshotLeaseMetadata) o;
+ return expirationTime == that.expirationTime
+ && tableIdToRemoteMetadataFilePath.equals(that.tableIdToRemoteMetadataFilePath);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(expirationTime, tableIdToRemoteMetadataFilePath);
+ }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerde.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerde.java
new file mode 100644
index 000000000..30c9367da
--- /dev/null
+++ b/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)
+ throws IOException {
+ generator.writeStartObject();
+ generator.writeNumberField(VERSION_KEY, VERSION);
+ generator.writeNumberField(EXPIRATION_TIME, kvSnapshotLeaseMetadata.getExpirationTime());
+
+ generator.writeFieldName(TABLES);
+ generator.writeStartArray();
+ for (Map.Entry<Long, FsPath> entry :
+ kvSnapshotLeaseMetadata.getTableIdToRemoteMetadataFilePath().entrySet()) {
+ generator.writeStartObject();
+ generator.writeNumberField(TABLE_ID, entry.getKey());
+ generator.writeStringField(KV_SNAPSHOT_PATH, entry.getValue().getPath());
+ generator.writeEndObject();
+ }
+ // end tables
+ generator.writeEndArray();
+
+ // end root
+ generator.writeEndObject();
+ }
+
+ @Override
+ public KvSnapshotLeaseMetadata deserialize(JsonNode node) {
+ long expirationTime = node.get(EXPIRATION_TIME).asLong();
+
+ Map<Long, FsPath> tableIdToRemoteMetadataFilePath = new HashMap<>();
+ JsonNode tablesNode = node.get(TABLES);
+
+ for (JsonNode tableNode : tablesNode) {
+ long tableId = tableNode.get(TABLE_ID).asLong();
+ String kvSnapshotPath = tableNode.get(KV_SNAPSHOT_PATH).asText();
+ tableIdToRemoteMetadataFilePath.put(tableId, new FsPath(kvSnapshotPath));
+ }
+
+ return new KvSnapshotLeaseMetadata(expirationTime, tableIdToRemoteMetadataFilePath);
+ }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLease.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLease.java
new file mode 100644
index 000000000..8ff2e8c9c
--- /dev/null
+++ b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLease.java
@@ -0,0 +1,190 @@
+/*
+ * 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 javax.annotation.Nullable;
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.util.Arrays;
+import java.util.HashMap;
+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 @Nullable Long[] bucketSnapshots;
+ private final Map<Long, Long[]> partitionSnapshots;
+
+ public KvSnapshotTableLease(long tableId) {
+ this(tableId, null, new HashMap<>());
+ }
+
+ public KvSnapshotTableLease(long tableId, Long[] bucketSnapshots) {
+ this(tableId, bucketSnapshots, new HashMap<>());
+ }
+
+ 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 void setBucketSnapshots(@Nullable Long[] bucketSnapshots) {
+ this.bucketSnapshots = 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 = 31 * hash + (Objects.hashCode(key) ^ Arrays.hashCode(value));
+ }
+ return hash;
+ }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLeaseJsonSerde.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLeaseJsonSerde.java
new file mode 100644
index 000000000..09dfbb3f3
--- /dev/null
+++ b/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)
+ throws IOException {
+ generator.writeStartObject();
+ generator.writeNumberField(VERSION_KEY, VERSION);
+ generator.writeNumberField(TABLE_ID, kvSnapshotTableLease.getTableId());
+
+ if (kvSnapshotTableLease.getBucketSnapshots() != null) {
+ // for none-partition table.
+ generator.writeArrayFieldStart(BUCKET_SNAPSHOTS);
+ for (Long snapshot : kvSnapshotTableLease.getBucketSnapshots()) {
+ generator.writeNumber(snapshot);
+ }
+ generator.writeEndArray();
+ } else {
+ // for partition table.
+ Map<Long, Long[]> partitionSnapshots = kvSnapshotTableLease.getPartitionSnapshots();
+ if (partitionSnapshots != null && !partitionSnapshots.isEmpty()) {
+ generator.writeArrayFieldStart(PARTITION_SNAPSHOTS);
+ for (Map.Entry<Long, Long[]> entry : partitionSnapshots.entrySet()) {
+ generator.writeStartObject();
+ generator.writeNumberField(PARTITION_ID, entry.getKey());
+ generator.writeArrayFieldStart(BUCKET_SNAPSHOTS);
+ for (Long snapshot : entry.getValue()) {
+ generator.writeNumber(snapshot);
+ }
+ generator.writeEndArray();
+ generator.writeEndObject();
+ }
+ generator.writeEndArray();
+ }
+ }
+ }
+
+ @Override
+ public KvSnapshotTableLease deserialize(JsonNode node) {
+ long tableId = node.get(TABLE_ID).asLong();
+ if (node.has(BUCKET_SNAPSHOTS)) {
+ // for none-partition table.
+ Long[] bucketSnapshots = new Long[node.get(BUCKET_SNAPSHOTS).size()];
+ for (int i = 0; i < bucketSnapshots.length; i++) {
+ bucketSnapshots[i] = node.get(BUCKET_SNAPSHOTS).get(i).asLong();
+ }
+ return new KvSnapshotTableLease(tableId, bucketSnapshots);
+ } else {
+ // for partition table.
+ Map<Long, Long[]> partitionSnapshots = new HashMap<>();
+ JsonNode partitionSnapshotsNode = node.get(PARTITION_SNAPSHOTS);
+ for (JsonNode partitionSnapshotNode : partitionSnapshotsNode) {
+ long partitionId = partitionSnapshotNode.get(PARTITION_ID).asLong();
+ Long[] bucketSnapshots =
+ new Long[partitionSnapshotNode.get(BUCKET_SNAPSHOTS).size()];
+ for (int i = 0; i < bucketSnapshots.length; i++) {
+ bucketSnapshots[i] =
+ partitionSnapshotNode.get(BUCKET_SNAPSHOTS).get(i).asLong();
+ }
+ partitionSnapshots.put(partitionId, bucketSnapshots);
+ }
+ return new KvSnapshotTableLease(tableId, partitionSnapshots);
+ }
+ }
+
+ /** Serialize the {@link KvSnapshotTableLease} to json bytes using current version. */
+ public static byte[] toJson(KvSnapshotTableLease kvSnapshotTableLease) {
+ return JsonSerdeUtils.writeValueAsBytes(kvSnapshotTableLease, INSTANCE);
+ }
+
+ /** Deserialize the json bytes to {@link KvSnapshotTableLease}. */
+ public static KvSnapshotTableLease fromJson(byte[] json) {
+ return JsonSerdeUtils.readValue(json, INSTANCE);
+ }
+}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManagerTest.java
index b8fec59ca..2558be96f 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManagerTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CompletedSnapshotStoreManagerTest.java
@@ -211,7 +211,8 @@ class CompletedSnapshotStoreManagerTest {
ioExecutor,
zookeeperClient,
zooKeeperClient -> completedSnapshotHandleStore,
- TestingMetricGroups.COORDINATOR_METRICS);
+ TestingMetricGroups.COORDINATOR_METRICS,
+ bucket -> true);
// Verify that only the valid snapshot remains
CompletedSnapshotStore completedSnapshotStore =
@@ -227,7 +228,8 @@ class CompletedSnapshotStoreManagerTest {
maxNumberOfSnapshotsToRetain,
ioExecutor,
zookeeperClient,
- TestingMetricGroups.COORDINATOR_METRICS);
+ TestingMetricGroups.COORDINATOR_METRICS,
+ bucket -> true);
}
private CompletedSnapshot getLatestCompletedSnapshot(
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java
index d4ec8f8db..c9e8bf5a0 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessorTest.java
@@ -48,6 +48,7 @@ import org.apache.fluss.server.coordinator.event.AdjustIsrReceivedEvent;
import org.apache.fluss.server.coordinator.event.CommitKvSnapshotEvent;
import org.apache.fluss.server.coordinator.event.CommitRemoteLogManifestEvent;
import org.apache.fluss.server.coordinator.event.CoordinatorEventManager;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.coordinator.statemachine.BucketState;
import org.apache.fluss.server.coordinator.statemachine.ReplicaState;
import org.apache.fluss.server.entity.AdjustIsrResultForBucket;
@@ -78,6 +79,7 @@ import org.apache.fluss.server.zk.data.ZkData.TableIdsZNode;
import org.apache.fluss.testutils.common.AllCallbackWrapper;
import org.apache.fluss.types.DataTypes;
import org.apache.fluss.utils.ExceptionUtils;
+import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
import org.apache.fluss.utils.types.Tuple2;
@@ -154,6 +156,7 @@ class CoordinatorEventProcessorTest {
private LakeTableTieringManager lakeTableTieringManager;
private CompletedSnapshotStoreManager completedSnapshotStoreManager;
private CoordinatorMetadataCache serverMetadataCache;
+ private KvSnapshotLeaseManager kvSnapshotLeaseManager;
@BeforeAll
static void baseBeforeAll() throws Exception {
@@ -193,7 +196,17 @@ class CoordinatorEventProcessorTest {
new AutoPartitionManager(serverMetadataCache, metadataManager, new Configuration());
lakeTableTieringManager = new LakeTableTieringManager();
Configuration conf = new Configuration();
- conf.setString(ConfigOptions.REMOTE_DATA_DIR, "/tmp/fluss/remote-data");
+ String remoteDataDir = "/tmp/fluss/remote-data";
+ conf.setString(ConfigOptions.REMOTE_DATA_DIR, remoteDataDir);
+ kvSnapshotLeaseManager =
+ new KvSnapshotLeaseManager(
+ Duration.ofMinutes(10).toMillis(),
+ zookeeperClient,
+ remoteDataDir,
+ SystemClock.getInstance(),
+ TestingMetricGroups.COORDINATOR_METRICS);
+ kvSnapshotLeaseManager.start();
+
eventProcessor = buildCoordinatorEventProcessor();
eventProcessor.startup();
metadataManager.createDatabase(
@@ -1051,7 +1064,8 @@ class CoordinatorEventProcessorTest {
TestingMetricGroups.COORDINATOR_METRICS,
new Configuration(),
Executors.newFixedThreadPool(1, new ExecutorThreadFactory("test-coordinator-io")),
- metadataManager);
+ metadataManager,
+ kvSnapshotLeaseManager);
}
private void initCoordinatorChannel() throws Exception {
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java
index a0a49e027..38e9730ca 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TestCoordinatorGateway.java
@@ -22,6 +22,8 @@ import org.apache.fluss.exception.IneligibleReplicaException;
import org.apache.fluss.exception.NetworkException;
import org.apache.fluss.metadata.TableBucket;
import org.apache.fluss.rpc.gateway.CoordinatorGateway;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.AcquireKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.AddServerTagRequest;
import org.apache.fluss.rpc.messages.AddServerTagResponse;
import org.apache.fluss.rpc.messages.AdjustIsrRequest;
@@ -60,6 +62,8 @@ import org.apache.fluss.rpc.messages.DropAclsRequest;
import org.apache.fluss.rpc.messages.DropAclsResponse;
import org.apache.fluss.rpc.messages.DropDatabaseRequest;
import org.apache.fluss.rpc.messages.DropDatabaseResponse;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.DropKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.DropPartitionRequest;
import org.apache.fluss.rpc.messages.DropPartitionResponse;
import org.apache.fluss.rpc.messages.DropTableRequest;
@@ -100,6 +104,8 @@ import org.apache.fluss.rpc.messages.RebalanceRequest;
import org.apache.fluss.rpc.messages.RebalanceResponse;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsRequest;
import org.apache.fluss.rpc.messages.RegisterProducerOffsetsResponse;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseRequest;
+import org.apache.fluss.rpc.messages.ReleaseKvSnapshotLeaseResponse;
import org.apache.fluss.rpc.messages.RemoveServerTagRequest;
import org.apache.fluss.rpc.messages.RemoveServerTagResponse;
import org.apache.fluss.rpc.messages.TableExistsRequest;
@@ -397,6 +403,24 @@ public class TestCoordinatorGateway implements CoordinatorGateway {
throw new UnsupportedOperationException();
}
+ @Override
+ public CompletableFuture<AcquireKvSnapshotLeaseResponse> acquireKvSnapshotLease(
+ AcquireKvSnapshotLeaseRequest request) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public CompletableFuture<ReleaseKvSnapshotLeaseResponse> releaseKvSnapshotLease(
+ ReleaseKvSnapshotLeaseRequest request) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public CompletableFuture<DropKvSnapshotLeaseResponse> dropKvSnapshotLease(
+ DropKvSnapshotLeaseRequest request) {
+ return null;
+ }
+
@Override
public CompletableFuture<ListAclsResponse> listAcls(ListAclsRequest request) {
throw new UnsupportedOperationException();
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseHandlerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseHandlerTest.java
new file mode 100644
index 000000000..b1905a614
--- /dev/null
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseHandlerTest.java
@@ -0,0 +1,287 @@
+/*
+ * 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 org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Test for {@link KvSnapshotLeaseHandler}. */
+public class KvSnapshotLeaseHandlerTest {
+
+ @Test
+ void testConstructorAndGetters() {
+ long expirationTime = 1000L;
+ KvSnapshotLeaseHandler kvSnapshotLeaseHandle = new KvSnapshotLeaseHandler(expirationTime);
+
+ assertThat(kvSnapshotLeaseHandle.getExpirationTime()).isEqualTo(expirationTime);
+ assertThat(kvSnapshotLeaseHandle.getTableIdToTableLease()).isEmpty();
+ assertThat(kvSnapshotLeaseHandle.getLeasedSnapshotCount()).isEqualTo(0);
+ }
+
+ @Test
+ void testRegisterBucketForNonPartitionedTable() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+ int bucketId = 0;
+
+ long originalSnapshot = acquireBucket(lease, new TableBucket(tableId, bucketId), 123L);
+
+ assertThat(originalSnapshot).isEqualTo(-1L);
+ assertThat(lease.getTableIdToTableLease()).containsKey(tableId);
+ KvSnapshotTableLease tableLease = lease.getTableIdToTableLease().get(tableId);
+ Long[] bucketSnapshots = tableLease.getBucketSnapshots();
+ assertThat(bucketSnapshots).isNotNull();
+ assertThat(bucketSnapshots).hasSize(1);
+ assertThat(bucketSnapshots[bucketId]).isEqualTo(123L);
+
+ // Register again same bucket → should be update
+ originalSnapshot = acquireBucket(lease, new TableBucket(tableId, bucketId), 456L);
+ assertThat(originalSnapshot).isEqualTo(123L);
+ tableLease = lease.getTableIdToTableLease().get(tableId);
+ bucketSnapshots = tableLease.getBucketSnapshots();
+ assertThat(bucketSnapshots).isNotNull();
+ assertThat(bucketSnapshots[bucketId]).isEqualTo(456L);
+ }
+
+ @Test
+ void testAcquireBucketDynamicExpansionForNonPartitionedTable() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+
+ // Acquire bucket 0 → creates array of size 1
+ acquireBucket(lease, new TableBucket(tableId, 0), 100L);
+ Long[] bucketSnapshots = lease.getTableIdToTableLease().get(tableId).getBucketSnapshots();
+ assertThat(bucketSnapshots).hasSize(1);
+ assertThat(bucketSnapshots[0]).isEqualTo(100L);
+
+ // Acquire bucket 2 → array should expand from size 1 to size 3
+ long originalSnapshot = acquireBucket(lease, new TableBucket(tableId, 2), 200L);
+ assertThat(originalSnapshot).isEqualTo(-1L);
+ bucketSnapshots = lease.getTableIdToTableLease().get(tableId).getBucketSnapshots();
+ assertThat(bucketSnapshots).hasSize(3);
+ assertThat(bucketSnapshots[0]).isEqualTo(100L);
+ assertThat(bucketSnapshots[1]).isEqualTo(-1L);
+ assertThat(bucketSnapshots[2]).isEqualTo(200L);
+
+ // Acquire bucket 1 → no expansion needed, just fill in
+ originalSnapshot = acquireBucket(lease, new TableBucket(tableId, 1), 150L);
+ assertThat(originalSnapshot).isEqualTo(-1L);
+ bucketSnapshots = lease.getTableIdToTableLease().get(tableId).getBucketSnapshots();
+ assertThat(bucketSnapshots).hasSize(3);
+ assertThat(bucketSnapshots[1]).isEqualTo(150L);
+ }
+
+ @Test
+ void testAcquireBucketDynamicExpansionForPartitionedTable() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+ long partitionId = 1000L;
+
+ // Acquire partition 1000, bucket 0 → creates array of size 1
+ acquireBucket(lease, new TableBucket(tableId, partitionId, 0), 100L);
+ Long[] partitionBuckets =
+ lease.getTableIdToTableLease()
+ .get(tableId)
+ .getPartitionSnapshots()
+ .get(partitionId);
+ assertThat(partitionBuckets).hasSize(1);
+ assertThat(partitionBuckets[0]).isEqualTo(100L);
+
+ // Acquire partition 1000, bucket 2 → array should expand from size 1 to size 3
+ long originalSnapshot =
+ acquireBucket(lease, new TableBucket(tableId, partitionId, 2), 200L);
+ assertThat(originalSnapshot).isEqualTo(-1L);
+ partitionBuckets =
+ lease.getTableIdToTableLease()
+ .get(tableId)
+ .getPartitionSnapshots()
+ .get(partitionId);
+ assertThat(partitionBuckets).hasSize(3);
+ assertThat(partitionBuckets[0]).isEqualTo(100L);
+ assertThat(partitionBuckets[1]).isEqualTo(-1L);
+ assertThat(partitionBuckets[2]).isEqualTo(200L);
+ }
+
+ @Test
+ void testRegisterBucketForPartitionedTable() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+
+ long originalSnapshot = acquireBucket(lease, new TableBucket(tableId, 1000L, 0), 111L);
+ assertThat(originalSnapshot).isEqualTo(-1L);
+ // Acquire bucket 1 for partition 1000 → array expands from size 1 to size 2
+ originalSnapshot = acquireBucket(lease, new TableBucket(tableId, 1000L, 1), 122L);
+ assertThat(originalSnapshot).isEqualTo(-1L);
+ originalSnapshot = acquireBucket(lease, new TableBucket(tableId, 1001L, 0), 122L);
+ assertThat(originalSnapshot).isEqualTo(-1L);
+
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = lease.getTableIdToTableLease();
+ assertThat(tableIdToTableLease).containsKey(tableId);
+ KvSnapshotTableLease tableLease = tableIdToTableLease.get(tableId);
+ Map<Long, Long[]> partitionSnapshots = tableLease.getPartitionSnapshots();
+ assertThat(partitionSnapshots).containsKeys(1000L, 1001L);
+ assertThat(partitionSnapshots.get(1000L)[0]).isEqualTo(111L);
+ assertThat(partitionSnapshots.get(1000L)[1]).isEqualTo(122L);
+ // Partition 1001 only has bucket 0, so its array size is 1
+ assertThat(partitionSnapshots.get(1001L)).hasSize(1);
+ assertThat(partitionSnapshots.get(1001L)[0]).isEqualTo(122L);
+
+ // test update.
+ originalSnapshot = acquireBucket(lease, new TableBucket(tableId, 1000L, 0), 222L);
+ assertThat(originalSnapshot).isEqualTo(111L);
+ assertThat(partitionSnapshots.get(1000L)[0]).isEqualTo(222L);
+ }
+
+ @Test
+ void testReleaseBucket() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+
+ // Register
+ acquireBucket(lease, new TableBucket(tableId, 0), 123L);
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(1);
+
+ // Unregister
+ long snapshotId = releaseBucket(lease, new TableBucket(tableId, 0));
+ assertThat(snapshotId).isEqualTo(123L);
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(0);
+ assertThat(lease.isEmpty()).isTrue();
+ }
+
+ @Test
+ void testReleaseBucketExceedingArraySize() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+
+ // Register bucket 0 only → array size is 1
+ acquireBucket(lease, new TableBucket(tableId, 0), 100L);
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(1);
+
+ // Release bucket 5 which exceeds the array size → should return -1
+ long snapshotId = releaseBucket(lease, new TableBucket(tableId, 5));
+ assertThat(snapshotId).isEqualTo(-1L);
+ // The original lease should remain unchanged
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(1);
+ assertThat(lease.isEmpty()).isFalse();
+ }
+
+ @Test
+ void testReleaseBucketExceedingArraySizeForPartitionedTable() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ long tableId = 1L;
+ long partitionId = 1000L;
+
+ // Register partition 1000, bucket 0 only → array size is 1
+ acquireBucket(lease, new TableBucket(tableId, partitionId, 0), 100L);
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(1);
+
+ // Release partition 1000, bucket 3 which exceeds the array size → should return -1
+ long snapshotId = releaseBucket(lease, new TableBucket(tableId, partitionId, 3));
+ assertThat(snapshotId).isEqualTo(-1L);
+ // The original lease should remain unchanged
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(1);
+ assertThat(lease.isEmpty()).isFalse();
+ }
+
+ @Test
+ void testGetLeasedSnapshotCount() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+
+ // Non-partitioned
+ acquireBucket(lease, new TableBucket(1L, 0), 100L);
+ acquireBucket(lease, new TableBucket(1L, 1), 101L);
+
+ // Partitioned
+ acquireBucket(lease, new TableBucket(2L, 20L, 0), 200L);
+ acquireBucket(lease, new TableBucket(2L, 21L, 1), 201L);
+
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(4);
+
+ // Unregister one
+ releaseBucket(lease, new TableBucket(1L, 0));
+ assertThat(lease.getLeasedSnapshotCount()).isEqualTo(3);
+ }
+
+ @Test
+ void testEqualsAndHashCode() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ assertThat(lease).isEqualTo(lease);
+ assertThat(lease.hashCode()).isEqualTo(lease.hashCode());
+
+ KvSnapshotLeaseHandler c1 = new KvSnapshotLeaseHandler(1000L);
+ KvSnapshotLeaseHandler c2 = new KvSnapshotLeaseHandler(2000L);
+ assertThat(c1).isNotEqualTo(c2);
+
+ // Create two leases with same logical content but different array objects
+ Map<Long, KvSnapshotTableLease> map1 = new HashMap<>();
+ ConcurrentHashMap<Long, Long[]> partitionSnapshots1 = MapUtils.newConcurrentHashMap();
+ partitionSnapshots1.put(2001L, new Long[] {100L, -1L});
+ partitionSnapshots1.put(2002L, new Long[] {-1L, 101L});
+ map1.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L, -1L}, partitionSnapshots1));
+ Map<Long, KvSnapshotTableLease> map2 = new HashMap<>();
+ ConcurrentHashMap<Long, Long[]> partitionSnapshots2 = MapUtils.newConcurrentHashMap();
+ partitionSnapshots2.put(2001L, new Long[] {100L, -1L});
+ partitionSnapshots2.put(2002L, new Long[] {-1L, 101L});
+ map2.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L, -1L}, partitionSnapshots2));
+ c1 = new KvSnapshotLeaseHandler(1000L, map1);
+ c2 = new KvSnapshotLeaseHandler(1000L, map2);
+ assertThat(c1).isEqualTo(c2);
+ assertThat(c1.hashCode()).isEqualTo(c2.hashCode());
+
+ // different array content.
+ map1 = new HashMap<>();
+ map1.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L, -1L}));
+ map2 = new HashMap<>();
+ map2.put(1L, new KvSnapshotTableLease(1L, new Long[] {200L, -1L}));
+ c1 = new KvSnapshotLeaseHandler(1000L, map1);
+ c2 = new KvSnapshotLeaseHandler(1000L, map2);
+ assertThat(c1).isNotEqualTo(c2);
+ }
+
+ @Test
+ void testToString() {
+ KvSnapshotLeaseHandler lease = new KvSnapshotLeaseHandler(1000L);
+ acquireBucket(lease, new TableBucket(1L, 0), 100L);
+ acquireBucket(lease, new TableBucket(1L, 1), 101L);
+ acquireBucket(lease, new TableBucket(2L, 0L, 0), 200L);
+ acquireBucket(lease, new TableBucket(2L, 1L, 1), 201L);
+ assertThat(lease.toString())
+ .isEqualTo(
+ "KvSnapshotLease{expirationTime=1000, tableIdToTableLease={"
+ + "1=KvSnapshotTableLease{tableId=1, bucketSnapshots=[100, 101], partitionSnapshots={}}, "
+ + "2=KvSnapshotTableLease{tableId=2, bucketSnapshots=null, partitionSnapshots={"
+ + "0=[200], 1=[-1, 201]}}}}");
+ }
+
+ private long acquireBucket(KvSnapshotLeaseHandler lease, TableBucket tb, long kvSnapshotId) {
+ return lease.acquireBucket(tb, kvSnapshotId, tb.getBucket() + 1);
+ }
+
+ private long releaseBucket(KvSnapshotLeaseHandler lease, TableBucket tb) {
+ return lease.releaseBucket(tb);
+ }
+}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManagerTest.java
new file mode 100644
index 000000000..161a95510
--- /dev/null
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseManagerTest.java
@@ -0,0 +1,612 @@
+/*
+ * 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.metadata.TableBucketSnapshot;
+import org.apache.fluss.metadata.TablePartition;
+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.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.PARTITION_TABLE_ID;
+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 long PARTITION_ID_2 = 19002L;
+
+ 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 ManualClock manualClock;
+ private ManuallyTriggeredScheduledExecutorService clearLeaseScheduler;
+ private KvSnapshotLeaseManager kvSnapshotLeaseManager;
+
+ private @TempDir Path tempDir;
+
+ @BeforeAll
+ static void beforeAll() {
+ zookeeperClient =
+ ZOO_KEEPER_EXTENSION_WRAPPER
+ .getCustomExtension()
+ .getZooKeeperClient(NOPErrorHandler.INSTANCE);
+ }
+
+ @BeforeEach
+ void beforeEach() throws Exception {
+ manualClock = new ManualClock(System.currentTimeMillis());
+ clearLeaseScheduler = new ManuallyTriggeredScheduledExecutorService();
+ kvSnapshotLeaseManager =
+ new KvSnapshotLeaseManager(
+ Duration.ofDays(7).toMillis(),
+ zookeeperClient,
+ tempDir.toString(),
+ 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 TableBucketSnapshot(t0b0, 0L),
+ new TableBucketSnapshot(t0b1, 0L),
+ new TableBucketSnapshot(t1p0b0, 0L),
+ new TableBucketSnapshot(t1p0b1, 0L),
+ new TableBucketSnapshot(t1p1b0, 0L))))
+ .isTrue();
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+
+ // test initialize from zookeeper when coordinator is started.
+ KvSnapshotLeaseHandler kvSnapshotLeasehandle = new KvSnapshotLeaseHandler(1000L);
+ acquire(kvSnapshotLeasehandle, new TableBucketSnapshot(t0b0, 0L));
+ acquire(kvSnapshotLeasehandle, new TableBucketSnapshot(t0b1, 0L));
+ acquire(kvSnapshotLeasehandle, new TableBucketSnapshot(t1p0b0, 0L));
+ acquire(kvSnapshotLeasehandle, new TableBucketSnapshot(t1p0b1, 0L));
+ acquire(kvSnapshotLeasehandle, new TableBucketSnapshot(t1p1b0, 0L));
+ kvSnapshotLeaseManager.getMetadataManager().registerLease("lease1", kvSnapshotLeasehandle);
+
+ kvSnapshotLeaseManager.start();
+
+ assertThat(
+ snapshotLeaseExists(
+ Arrays.asList(
+ new TableBucketSnapshot(t0b0, 0L),
+ new TableBucketSnapshot(t0b1, 0L),
+ new TableBucketSnapshot(t1p0b0, 0L),
+ new TableBucketSnapshot(t1p0b1, 0L),
+ new TableBucketSnapshot(t1p1b0, 0L))))
+ .isTrue();
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(5);
+
+ // check detail content.
+ // Array sizes are dynamically determined by the maximum bucket id:
+ // - DATA1_TABLE_ID_PK: bucket 0 and 1 → array size 2
+ // - PARTITION_ID_1: bucket 0 and 1 → array size 2
+ // - PARTITION_ID_2: bucket 0 only → array size 1
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+ tableIdToTableLease.put(
+ DATA1_TABLE_ID_PK,
+ new KvSnapshotTableLease(DATA1_TABLE_ID_PK, new Long[] {0L, 0L}));
+ KvSnapshotTableLease leaseForPartitionTable = new KvSnapshotTableLease(PARTITION_TABLE_ID);
+ leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_1, new Long[] {0L, 0L});
+ leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_2, new Long[] {0L});
+ tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+ KvSnapshotLeaseHandler expectedLease =
+ new KvSnapshotLeaseHandler(1000L, tableIdToTableLease);
+ assertThat(kvSnapshotLeaseManager.getKvSnapshotLeaseData("lease1"))
+ .isEqualTo(expectedLease);
+ assertThat(kvSnapshotLeaseManager.getMetadataManager().getLease("lease1"))
+ .hasValue(expectedLease);
+ }
+
+ @Test
+ void testAcquireAndRelease() throws Exception {
+ Map<Long, List<TableBucketSnapshot>> 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 TableBucketSnapshot(t1p0b0, 1L), new TableBucketSnapshot(t1p0b1, 1L)));
+ acquire("lease2", tableIdToRegisterBucket);
+
+ assertThat(
+ snapshotLeaseExists(
+ Arrays.asList(
+ new TableBucketSnapshot(t0b0, 0L),
+ new TableBucketSnapshot(t0b1, 0L),
+ new TableBucketSnapshot(t1p0b0, 0),
+ new TableBucketSnapshot(t1p0b1, 0L),
+ new TableBucketSnapshot(t1p1b0, 0L),
+ new TableBucketSnapshot(t1p0b0, 1L),
+ new TableBucketSnapshot(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 TableBucketSnapshot(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 TableBucketSnapshot(newTableBucket, 1L)));
+ acquire("lease1", tableIdToRegisterBucket);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(8);
+
+ // release
+ release("lease1", Collections.singletonList(newTableBucket));
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(7);
+
+ // release a non-exist bucket.
+ release(
+ "lease1",
+ Collections.singletonList(new TableBucket(DATA1_TABLE_ID_PK, PARTITION_ID_1, 2)));
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(7);
+
+ // check detail content for lease1.
+ // - DATA1_TABLE_ID_PK: bucket 0,1 initially (size 2), then bucket 2 acquired and
+ // released → expanded to size 3: {1L, 0L, -1L}
+ // - PARTITION_ID_1: bucket 0,1 → array size 2
+ // - PARTITION_ID_2: bucket 0 only → array size 1
+ 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});
+ leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_2, new Long[] {0L});
+ tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+ KvSnapshotLeaseHandler expectedLease =
+ new KvSnapshotLeaseHandler(manualClock.milliseconds() + 1000L, tableIdToTableLease);
+ assertThat(kvSnapshotLeaseManager.getKvSnapshotLeaseData("lease1"))
+ .isEqualTo(expectedLease);
+ assertThat(kvSnapshotLeaseManager.getMetadataManager().getLease("lease1"))
+ .hasValue(expectedLease);
+
+ // check detail content for lease2.
+ // - PARTITION_ID_1: bucket 0,1 → array size 2
+ tableIdToTableLease = new HashMap<>();
+ leaseForPartitionTable = new KvSnapshotTableLease(PARTITION_TABLE_ID);
+ leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_1, new Long[] {1L, 1L});
+ tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+ KvSnapshotLeaseHandler expectedLease2 =
+ new KvSnapshotLeaseHandler(manualClock.milliseconds() + 1000L, tableIdToTableLease);
+ assertThat(kvSnapshotLeaseManager.getKvSnapshotLeaseData("lease2"))
+ .isEqualTo(expectedLease2);
+ assertThat(kvSnapshotLeaseManager.getMetadataManager().getLease("lease2"))
+ .hasValue(expectedLease2);
+ }
+
+ @Test
+ void testDropLease() throws Exception {
+ Map<Long, List<TableBucketSnapshot>> tableIdToRegisterBucket = initRegisterBuckets();
+ acquire("lease1", tableIdToRegisterBucket);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(5);
+ assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(1);
+ assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isPresent();
+
+ // release all will clear this lease.
+ release("lease1", Arrays.asList(t0b0, t0b1, t1p0b0, t1p0b1, t1p1b0));
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+ assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(0);
+ assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isNotPresent();
+ }
+
+ @Test
+ void testClear() throws Exception {
+ Map<Long, List<TableBucketSnapshot>> 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 TableBucketSnapshot(t0b0, 0L), // same ref.
+ new TableBucketSnapshot(t1p0b0, 1L),
+ new TableBucketSnapshot(t1p0b1, 1L)));
+ acquire("lease2", tableIdToRegisterBucket);
+
+ assertThat(
+ snapshotLeaseExists(
+ Arrays.asList(
+ new TableBucketSnapshot(t0b0, 0L),
+ new TableBucketSnapshot(t0b1, 0L),
+ new TableBucketSnapshot(t1p0b0, 0L),
+ new TableBucketSnapshot(t1p0b1, 0L),
+ new TableBucketSnapshot(t1p1b0, 0L),
+ new TableBucketSnapshot(t1p0b0, 1L),
+ new TableBucketSnapshot(t1p0b1, 1L))))
+ .isTrue();
+ assertThat(kvSnapshotLeaseManager.getRefCount(new TableBucketSnapshot(t0b0, 0L)))
+ .isEqualTo(2);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(8);
+ assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(2);
+
+ kvSnapshotLeaseManager.dropLease("lease1");
+ assertThat(kvSnapshotLeaseManager.getRefCount(new TableBucketSnapshot(t0b0, 0L)))
+ .isEqualTo(1);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(3);
+ assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(1);
+ assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease1")).isEmpty();
+
+ kvSnapshotLeaseManager.dropLease("lease2");
+ assertThat(kvSnapshotLeaseManager.getRefCount(new TableBucketSnapshot(t0b0, 0L)))
+ .isEqualTo(0);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+ assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(0);
+ assertThat(zookeeperClient.getKvSnapshotLeaseMetadata("lease2")).isEmpty();
+
+ assertThat(kvSnapshotLeaseManager.dropLease("non-exist")).isFalse();
+ }
+
+ @Test
+ void testExpireLeases() throws Exception {
+ // test lease expire by expire thread.
+ Map<Long, List<TableBucketSnapshot>> 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 TableBucketSnapshot(t0b0, 0L), // same ref.
+ new TableBucketSnapshot(t1p0b0, 1L),
+ new TableBucketSnapshot(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<TableBucketSnapshot>> tableIdToRegisterBucket = new HashMap<>();
+ tableIdToRegisterBucket.put(
+ DATA1_TABLE_ID_PK,
+ Arrays.asList(
+ new TableBucketSnapshot(t0b0, 1000L),
+ new TableBucketSnapshot(t0b1, 1000L)));
+
+ // TODO this case will return t0b0, t0b1, after we add the checker the check whether
+ // snapshot exists.
+ assertThat(
+ kvSnapshotLeaseManager
+ .acquireLease("lease1", 1000L, tableIdToRegisterBucket)
+ .keySet())
+ .isEmpty();
+ }
+
+ // ------------------------------------------------------------------------
+ // findMaxBucketNum Tests
+ // ------------------------------------------------------------------------
+
+ @Test
+ void testFindMaxBucketNumForNonPartitionedTable() {
+ Map<Long, List<TableBucketSnapshot>> tableIdToLeaseBucket = new HashMap<>();
+ // Table 1: buckets 0, 3, 1 → max bucket id = 3, max bucket num = 4
+ tableIdToLeaseBucket.put(
+ 1L,
+ Arrays.asList(
+ new TableBucketSnapshot(new TableBucket(1L, 0), 0L),
+ new TableBucketSnapshot(new TableBucket(1L, 3), 0L),
+ new TableBucketSnapshot(new TableBucket(1L, 1), 0L)));
+ // Table 2: bucket 5 only → max bucket num = 6
+ tableIdToLeaseBucket.put(
+ 2L, Collections.singletonList(new TableBucketSnapshot(new TableBucket(2L, 5), 0L)));
+
+ Map<Long, Integer> maxBucketNum = new HashMap<>();
+ Map<TablePartition, Integer> maxBucketNumOfPt = new HashMap<>();
+ KvSnapshotLeaseManager.findMaxBucketNum(
+ tableIdToLeaseBucket, maxBucketNum, maxBucketNumOfPt);
+
+ assertThat(maxBucketNum).containsEntry(1L, 4);
+ assertThat(maxBucketNum).containsEntry(2L, 6);
+ assertThat(maxBucketNumOfPt).isEmpty();
+ }
+
+ @Test
+ void testFindMaxBucketNumForPartitionedTable() {
+ Map<Long, List<TableBucketSnapshot>> tableIdToLeaseBucket = new HashMap<>();
+ long tableId = 10L;
+ long partitionId1 = 100L;
+ long partitionId2 = 200L;
+ // Partition 100: buckets 2, 0 → max bucket id = 2, max bucket num = 3
+ // Partition 200: bucket 4 only → max bucket num = 5
+ tableIdToLeaseBucket.put(
+ tableId,
+ Arrays.asList(
+ new TableBucketSnapshot(new TableBucket(tableId, partitionId1, 2), 0L),
+ new TableBucketSnapshot(new TableBucket(tableId, partitionId1, 0), 0L),
+ new TableBucketSnapshot(new TableBucket(tableId, partitionId2, 4), 0L)));
+
+ Map<Long, Integer> maxBucketNum = new HashMap<>();
+ Map<TablePartition, Integer> maxBucketNumOfPt = new HashMap<>();
+ KvSnapshotLeaseManager.findMaxBucketNum(
+ tableIdToLeaseBucket, maxBucketNum, maxBucketNumOfPt);
+
+ assertThat(maxBucketNum).isEmpty();
+ assertThat(maxBucketNumOfPt).containsEntry(new TablePartition(tableId, partitionId1), 3);
+ assertThat(maxBucketNumOfPt).containsEntry(new TablePartition(tableId, partitionId2), 5);
+ }
+
+ @Test
+ void testFindMaxBucketNumMixed() {
+ Map<Long, List<TableBucketSnapshot>> tableIdToLeaseBucket = new HashMap<>();
+ // Non-partitioned table 1: buckets 5, 1, 3 → max bucket num = 6
+ tableIdToLeaseBucket.put(
+ 1L,
+ Arrays.asList(
+ new TableBucketSnapshot(new TableBucket(1L, 5), 0L),
+ new TableBucketSnapshot(new TableBucket(1L, 1), 0L),
+ new TableBucketSnapshot(new TableBucket(1L, 3), 0L)));
+ // Partitioned table 2: partition 100 buckets 0, 2 → max bucket num = 3
+ long partitionId = 100L;
+ tableIdToLeaseBucket.put(
+ 2L,
+ Arrays.asList(
+ new TableBucketSnapshot(new TableBucket(2L, partitionId, 0), 0L),
+ new TableBucketSnapshot(new TableBucket(2L, partitionId, 2), 0L)));
+
+ Map<Long, Integer> maxBucketNum = new HashMap<>();
+ Map<TablePartition, Integer> maxBucketNumOfPt = new HashMap<>();
+ KvSnapshotLeaseManager.findMaxBucketNum(
+ tableIdToLeaseBucket, maxBucketNum, maxBucketNumOfPt);
+
+ assertThat(maxBucketNum).hasSize(1).containsEntry(1L, 6);
+ assertThat(maxBucketNumOfPt)
+ .hasSize(1)
+ .containsEntry(new TablePartition(2L, partitionId), 3);
+ }
+
+ private Map<Long, List<TableBucketSnapshot>> initRegisterBuckets() {
+ Map<Long, List<TableBucketSnapshot>> tableIdToRegisterBucket = new HashMap<>();
+ tableIdToRegisterBucket.put(
+ DATA1_TABLE_ID_PK,
+ Arrays.asList(
+ new TableBucketSnapshot(t0b0, 0L), new TableBucketSnapshot(t0b1, 0L)));
+ tableIdToRegisterBucket.put(
+ PARTITION_TABLE_ID,
+ Arrays.asList(
+ new TableBucketSnapshot(t1p0b0, 0L),
+ new TableBucketSnapshot(t1p0b1, 0L),
+ new TableBucketSnapshot(t1p1b0, 0L)));
+ return tableIdToRegisterBucket;
+ }
+
+ private boolean snapshotLeaseNotExists(List<TableBucketSnapshot> bucketList) {
+ return bucketList.stream()
+ .noneMatch(bucket -> kvSnapshotLeaseManager.snapshotLeaseExist(bucket));
+ }
+
+ private boolean snapshotLeaseExists(List<TableBucketSnapshot> bucketList) {
+ return bucketList.stream()
+ .allMatch(bucket -> kvSnapshotLeaseManager.snapshotLeaseExist(bucket));
+ }
+
+ private void acquire(String leaseId, Map<Long, List<TableBucketSnapshot>> tableIdToLeaseBucket)
+ throws Exception {
+ kvSnapshotLeaseManager.acquireLease(leaseId, 1000L, tableIdToLeaseBucket);
+ }
+
+ private void release(String leaseId, List<TableBucket> tableBucketsToRelease) throws Exception {
+ kvSnapshotLeaseManager.release(leaseId, tableBucketsToRelease);
+ }
+
+ private long acquire(
+ KvSnapshotLeaseHandler kvSnapshotLeaseHandle, TableBucketSnapshot leaseForBucket) {
+ TableBucket tableBucket = leaseForBucket.getTableBucket();
+ return kvSnapshotLeaseHandle.acquireBucket(
+ tableBucket, leaseForBucket.getSnapshotId(), tableBucket.getBucket() + 1);
+ }
+
+ // ------------------------------------------------------------------------
+ // Bucket Expansion Tests
+ // ------------------------------------------------------------------------
+
+ @Test
+ void testAcquireBucketExpansionForNonPartitionedTable() throws Exception {
+ // Acquire bucket 0 for non-partitioned table → creates array of size 1
+ Map<Long, List<TableBucketSnapshot>> buckets = new HashMap<>();
+ buckets.put(
+ DATA1_TABLE_ID_PK, Collections.singletonList(new TableBucketSnapshot(t0b0, 0L)));
+ acquire("lease1", buckets);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(1);
+
+ // Acquire bucket 1 → should expand array from size 1 to size 2
+ buckets = new HashMap<>();
+ buckets.put(
+ DATA1_TABLE_ID_PK, Collections.singletonList(new TableBucketSnapshot(t0b1, 0L)));
+ acquire("lease1", buckets);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(2);
+
+ // Acquire bucket 5 (skipping buckets 2,3,4) → should expand array from size 2 to size 6
+ TableBucket t0b5 = new TableBucket(DATA1_TABLE_ID_PK, 5);
+ zookeeperClient.registerTableBucketSnapshot(t0b5, new BucketSnapshot(0L, 0L, "test-path"));
+ buckets = new HashMap<>();
+ buckets.put(
+ DATA1_TABLE_ID_PK, Collections.singletonList(new TableBucketSnapshot(t0b5, 0L)));
+ acquire("lease1", buckets);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(3);
+
+ // Verify the array content: bucket 0,1 = 0L, bucket 2-4 = -1L, bucket 5 = 0L
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+ tableIdToTableLease.put(
+ DATA1_TABLE_ID_PK,
+ new KvSnapshotTableLease(
+ DATA1_TABLE_ID_PK, new Long[] {0L, 0L, -1L, -1L, -1L, 0L}));
+ KvSnapshotLeaseHandler expectedLease =
+ new KvSnapshotLeaseHandler(manualClock.milliseconds() + 1000L, tableIdToTableLease);
+ assertThat(kvSnapshotLeaseManager.getKvSnapshotLeaseData("lease1"))
+ .isEqualTo(expectedLease);
+ }
+
+ @Test
+ void testAcquireBucketExpansionForPartitionedTable() throws Exception {
+ // Acquire partition bucket 0 → creates array of size 1
+ Map<Long, List<TableBucketSnapshot>> buckets = new HashMap<>();
+ buckets.put(
+ PARTITION_TABLE_ID, Collections.singletonList(new TableBucketSnapshot(t1p0b0, 0L)));
+ acquire("lease1", buckets);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(1);
+
+ // Acquire partition bucket 3 (skipping bucket 1,2) → should expand array to size 4
+ TableBucket t1p0b3 = new TableBucket(PARTITION_TABLE_ID, PARTITION_ID_1, 3);
+ zookeeperClient.registerTableBucketSnapshot(
+ t1p0b3, new BucketSnapshot(0L, 0L, "test-path"));
+ buckets = new HashMap<>();
+ buckets.put(
+ PARTITION_TABLE_ID, Collections.singletonList(new TableBucketSnapshot(t1p0b3, 0L)));
+ acquire("lease1", buckets);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(2);
+
+ // Verify the array content: bucket 0 = 0L, bucket 1,2 = -1L, bucket 3 = 0L
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+ KvSnapshotTableLease leaseForPartitionTable = new KvSnapshotTableLease(PARTITION_TABLE_ID);
+ leaseForPartitionTable.addPartitionSnapshots(PARTITION_ID_1, new Long[] {0L, -1L, -1L, 0L});
+ tableIdToTableLease.put(PARTITION_TABLE_ID, leaseForPartitionTable);
+ KvSnapshotLeaseHandler expectedLease =
+ new KvSnapshotLeaseHandler(manualClock.milliseconds() + 1000L, tableIdToTableLease);
+ assertThat(kvSnapshotLeaseManager.getKvSnapshotLeaseData("lease1"))
+ .isEqualTo(expectedLease);
+ }
+
+ @Test
+ void testReleaseBucketExceedingArraySize() throws Exception {
+ // Acquire bucket 0 only → array size 1
+ Map<Long, List<TableBucketSnapshot>> buckets = new HashMap<>();
+ buckets.put(
+ DATA1_TABLE_ID_PK, Collections.singletonList(new TableBucketSnapshot(t0b0, 0L)));
+ acquire("lease1", buckets);
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(1);
+
+ // Release bucket 1 which exceeds array size → should not affect existing leases
+ release("lease1", Collections.singletonList(t0b1));
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(1);
+
+ // Release bucket 0 → should remove the lease entry
+ release("lease1", Collections.singletonList(t0b0));
+ assertThat(kvSnapshotLeaseManager.getLeasedBucketCount()).isEqualTo(0);
+ assertThat(kvSnapshotLeaseManager.getLeaseCount()).isEqualTo(0);
+ }
+}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseMetadataManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseMetadataManagerTest.java
new file mode 100644
index 000000000..6f10400f6
--- /dev/null
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/lease/KvSnapshotLeaseMetadataManagerTest.java
@@ -0,0 +1,281 @@
+/*
+ * 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.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.server.zk.data.lease.KvSnapshotLeaseMetadata;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLease;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotTableLeaseJsonSerde;
+import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.utils.IOUtils;
+import org.apache.fluss.utils.MapUtils;
+
+import org.assertj.core.api.Assertions;
+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.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Stream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** 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 metadataManager;
+
+ @BeforeAll
+ static void beforeAll() {
+ zookeeperClient =
+ ZOO_KEEPER_EXTENSION_WRAPPER
+ .getCustomExtension()
+ .getZooKeeperClient(NOPErrorHandler.INSTANCE);
+ }
+
+ @BeforeEach
+ void beforeEach() {
+ metadataManager = new KvSnapshotLeaseMetadataManager(zookeeperClient, tempDir.toString());
+ }
+
+ @AfterEach
+ void afterEach() {
+ ZOO_KEEPER_EXTENSION_WRAPPER.getCustomExtension().cleanupRoot();
+ }
+
+ @Test
+ void testGetLeasesList() throws Exception {
+ List<String> leasesList = metadataManager.getLeasesList();
+ assertThat(leasesList).isEmpty();
+
+ metadataManager.registerLease("leaseId1", new KvSnapshotLeaseHandler(1000L));
+
+ Map<Long, KvSnapshotTableLease> tableIdToTableLease = new HashMap<>();
+ tableIdToTableLease.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L, -1L}));
+ metadataManager.registerLease(
+ "leaseId2", new KvSnapshotLeaseHandler(2000L, tableIdToTableLease));
+ leasesList = metadataManager.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}));
+
+ ConcurrentHashMap<Long, Long[]> partitionSnapshots = MapUtils.newConcurrentHashMap();
+ partitionSnapshots.put(1000L, new Long[] {111L, 122L});
+ partitionSnapshots.put(1001L, new Long[] {122L, -1L});
+ tableIdToTableLease.put(2L, new KvSnapshotTableLease(2L, partitionSnapshots));
+
+ KvSnapshotLeaseHandler expectedLease =
+ new KvSnapshotLeaseHandler(1000L, tableIdToTableLease);
+ metadataManager.registerLease("leaseId1", expectedLease);
+
+ Optional<KvSnapshotLeaseHandler> lease = metadataManager.getLease("leaseId1");
+ assertThat(lease).hasValue(expectedLease);
+ // assert zk and remote fs.
+ assertRemoteFsAndZkEquals("leaseId1", expectedLease);
+
+ // test update lease.
+ tableIdToTableLease.remove(1L);
+ expectedLease = new KvSnapshotLeaseHandler(2000L, tableIdToTableLease);
+ metadataManager.updateLease("leaseId1", expectedLease);
+ lease = metadataManager.getLease("leaseId1");
+ assertThat(lease).hasValue(expectedLease);
+ // assert zk and remote fs.
+ assertRemoteFsAndZkEquals("leaseId1", expectedLease);
+
+ // test delete lease.
+ metadataManager.deleteLease("leaseId1");
+ lease = metadataManager.getLease("leaseId1");
+ assertThat(lease).isEmpty();
+ }
+
+ @Test
+ void testRegisterLeaseFailureOnlyCleanupCreatedFiles() throws Exception {
+ // Step 1: Register a lease successfully with table 1.
+ Map<Long, KvSnapshotTableLease> tableLeases1 = new HashMap<>();
+ tableLeases1.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L}));
+ KvSnapshotLeaseHandler lease1 = new KvSnapshotLeaseHandler(1000L, tableLeases1);
+ metadataManager.registerLease("leaseId1", lease1);
+
+ // Record the original file path from ZK metadata.
+ Optional<KvSnapshotLeaseMetadata> originalMeta =
+ zookeeperClient.getKvSnapshotLeaseMetadata("leaseId1");
+ assertThat(originalMeta).isPresent();
+ FsPath originalFile = originalMeta.get().getTableIdToRemoteMetadataFilePath().get(1L);
+ assertThat(originalFile.getFileSystem().exists(originalFile)).isTrue();
+
+ // Step 2: Try to register again with the same leaseId but a different table.
+ // This should fail because the ZK node already exists.
+ Map<Long, KvSnapshotTableLease> tableLeases2 = new HashMap<>();
+ tableLeases2.put(2L, new KvSnapshotTableLease(2L, new Long[] {200L}));
+ KvSnapshotLeaseHandler lease2 = new KvSnapshotLeaseHandler(2000L, tableLeases2);
+
+ assertThatThrownBy(() -> metadataManager.registerLease("leaseId1", lease2))
+ .isInstanceOf(Exception.class);
+
+ // Step 3: Verify the original file (from the first successful register) still exists.
+ assertThat(originalFile.getFileSystem().exists(originalFile)).isTrue();
+
+ // Step 4: Verify the file created for table 2 was cleaned up.
+ Path table2Dir = tempDir.resolve("lease/kv-snapshot/leaseId1/2");
+ if (Files.exists(table2Dir)) {
+ try (Stream<Path> files = Files.list(table2Dir)) {
+ assertThat(files.filter(p -> p.toString().endsWith(".metadata")).count())
+ .isEqualTo(0);
+ }
+ }
+ }
+
+ @Test
+ void testUpdateLeaseFailureOnlyCleanupCreatedFiles() throws Exception {
+ // Step 1: Register a lease successfully with table 1.
+ Map<Long, KvSnapshotTableLease> tableLeases = new HashMap<>();
+ tableLeases.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L}));
+ KvSnapshotLeaseHandler originalLease = new KvSnapshotLeaseHandler(1000L, tableLeases);
+ metadataManager.registerLease("leaseId1", originalLease);
+
+ // Record the original file path.
+ Optional<KvSnapshotLeaseMetadata> originalMeta =
+ zookeeperClient.getKvSnapshotLeaseMetadata("leaseId1");
+ assertThat(originalMeta).isPresent();
+ FsPath originalFile = originalMeta.get().getTableIdToRemoteMetadataFilePath().get(1L);
+ assertThat(originalFile.getFileSystem().exists(originalFile)).isTrue();
+
+ // Step 2: Delete the ZK node to simulate a failure scenario.
+ // updateKvSnapshotLeaseMetadata will throw NoNodeException.
+ zookeeperClient.deleteKvSnapshotLease("leaseId1");
+
+ // Step 3: Try to update the lease (will fail at ZK update).
+ // generateMetadataFile creates new files, but they should be cleaned up on failure.
+ Map<Long, KvSnapshotTableLease> newTableLeases = new HashMap<>();
+ newTableLeases.put(1L, new KvSnapshotTableLease(1L, new Long[] {200L}));
+ newTableLeases.put(2L, new KvSnapshotTableLease(2L, new Long[] {300L}));
+ KvSnapshotLeaseHandler updatedLease = new KvSnapshotLeaseHandler(2000L, newTableLeases);
+
+ assertThatThrownBy(() -> metadataManager.updateLease("leaseId1", updatedLease))
+ .isInstanceOf(Exception.class);
+
+ // Step 4: Verify the original file (from the first register) is preserved on disk.
+ // It was NOT in createdFiles, so it should NOT be cleaned up.
+ assertThat(originalFile.getFileSystem().exists(originalFile)).isTrue();
+
+ // Step 5: Verify the newly created files for table 2 were cleaned up.
+ Path table2Dir = tempDir.resolve("lease/kv-snapshot/leaseId1/2");
+ if (Files.exists(table2Dir)) {
+ try (Stream<Path> files = Files.list(table2Dir)) {
+ assertThat(files.filter(p -> p.toString().endsWith(".metadata")).count())
+ .isEqualTo(0);
+ }
+ }
+
+ // Verify the newly created file for table 1 (with new UUID) was also cleaned up.
+ // Only the original file should remain under table 1 directory.
+ Path table1Dir = tempDir.resolve("lease/kv-snapshot/leaseId1/1");
+ try (Stream<Path> files = Files.list(table1Dir)) {
+ long metadataFileCount = files.filter(p -> p.toString().endsWith(".metadata")).count();
+ // Only the original file should remain.
+ assertThat(metadataFileCount).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void testUpdateLeaseSuccessCleansUpOldFiles() throws Exception {
+ // Step 1: Register a lease with table 1.
+ Map<Long, KvSnapshotTableLease> tableLeases = new HashMap<>();
+ tableLeases.put(1L, new KvSnapshotTableLease(1L, new Long[] {100L}));
+ KvSnapshotLeaseHandler originalLease = new KvSnapshotLeaseHandler(1000L, tableLeases);
+ metadataManager.registerLease("leaseId1", originalLease);
+
+ // Record the original file path.
+ Optional<KvSnapshotLeaseMetadata> originalMeta =
+ zookeeperClient.getKvSnapshotLeaseMetadata("leaseId1");
+ assertThat(originalMeta).isPresent();
+ FsPath originalFile = originalMeta.get().getTableIdToRemoteMetadataFilePath().get(1L);
+ assertThat(originalFile.getFileSystem().exists(originalFile)).isTrue();
+
+ // Step 2: Update the lease successfully with different data.
+ Map<Long, KvSnapshotTableLease> newTableLeases = new HashMap<>();
+ newTableLeases.put(1L, new KvSnapshotTableLease(1L, new Long[] {200L}));
+ KvSnapshotLeaseHandler updatedLease = new KvSnapshotLeaseHandler(2000L, newTableLeases);
+ metadataManager.updateLease("leaseId1", updatedLease);
+
+ // Step 3: Verify the old file was cleaned up after successful update.
+ assertThat(originalFile.getFileSystem().exists(originalFile)).isFalse();
+
+ // Step 4: Verify the new file exists and ZK points to it.
+ Optional<KvSnapshotLeaseMetadata> newMeta =
+ zookeeperClient.getKvSnapshotLeaseMetadata("leaseId1");
+ assertThat(newMeta).isPresent();
+ FsPath newFile = newMeta.get().getTableIdToRemoteMetadataFilePath().get(1L);
+ assertThat(newFile.getFileSystem().exists(newFile)).isTrue();
+ // The new file should be different from the old one (different UUID).
+ assertThat(newFile).isNotEqualTo(originalFile);
+ }
+
+ private void assertRemoteFsAndZkEquals(String leaseId, KvSnapshotLeaseHandler expectedLease)
+ throws Exception {
+ Optional<KvSnapshotLeaseMetadata> leaseMetadataOpt =
+ zookeeperClient.getKvSnapshotLeaseMetadata(leaseId);
+ assertThat(leaseMetadataOpt).isPresent();
+ KvSnapshotLeaseMetadata leaseMetadata = leaseMetadataOpt.get();
+ assertThat(leaseMetadata.getExpirationTime()).isEqualTo(expectedLease.getExpirationTime());
+ Map<Long, FsPath> actualFsPathSet = leaseMetadata.getTableIdToRemoteMetadataFilePath();
+ Map<Long, KvSnapshotTableLease> expectedTableLeases =
+ expectedLease.getTableIdToTableLease();
+ assertThat(actualFsPathSet).hasSize(expectedTableLeases.size());
+ for (Map.Entry<Long, FsPath> actualEntry : actualFsPathSet.entrySet()) {
+ long tableId = actualEntry.getKey();
+ FsPath actualMetadataPath = actualEntry.getValue();
+ assertThat(actualMetadataPath).isNotNull();
+ KvSnapshotTableLease expectedTableLease = expectedTableLeases.get(tableId);
+ assertThat(expectedTableLease).isNotNull();
+ FSDataInputStream inputStream =
+ actualMetadataPath.getFileSystem().open(actualMetadataPath);
+ try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
+ IOUtils.copyBytes(inputStream, outputStream, true);
+ Assertions.assertThat(
+ KvSnapshotTableLeaseJsonSerde.fromJson(outputStream.toByteArray()))
+ .isEqualTo(expectedTableLease);
+ }
+ }
+ }
+}
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
index fc36e3034..43ed3ae78 100644
--- 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
@@ -26,6 +26,7 @@ import org.apache.fluss.server.coordinator.LakeCatalogDynamicLoader;
import org.apache.fluss.server.coordinator.LakeTableTieringManager;
import org.apache.fluss.server.coordinator.MetadataManager;
import org.apache.fluss.server.coordinator.TestCoordinatorChannelManager;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.metadata.CoordinatorMetadataCache;
import org.apache.fluss.server.metrics.group.TestingMetricGroups;
import org.apache.fluss.server.zk.NOPErrorHandler;
@@ -33,6 +34,7 @@ import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.server.zk.ZooKeeperExtension;
import org.apache.fluss.server.zk.data.RebalanceTask;
import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
import org.junit.jupiter.api.AfterEach;
@@ -41,6 +43,7 @@ import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
+import java.time.Duration;
import java.util.HashMap;
import java.util.concurrent.Executors;
@@ -63,6 +66,7 @@ public class RebalanceManagerTest {
private AutoPartitionManager autoPartitionManager;
private LakeTableTieringManager lakeTableTieringManager;
private RebalanceManager rebalanceManager;
+ private KvSnapshotLeaseManager kvSnapshotLeaseManager;
@BeforeAll
static void baseBeforeAll() {
@@ -76,6 +80,17 @@ public class RebalanceManagerTest {
void beforeEach() {
serverMetadataCache = new CoordinatorMetadataCache();
testCoordinatorChannelManager = new TestCoordinatorChannelManager();
+
+ String remoteDataDir = "/tmp/fluss/remote-data";
+ kvSnapshotLeaseManager =
+ new KvSnapshotLeaseManager(
+ Duration.ofMinutes(10).toMillis(),
+ zookeeperClient,
+ remoteDataDir,
+ SystemClock.getInstance(),
+ TestingMetricGroups.COORDINATOR_METRICS);
+ kvSnapshotLeaseManager.start();
+
autoPartitionManager =
new AutoPartitionManager(serverMetadataCache, metadataManager, new Configuration());
lakeTableTieringManager = new LakeTableTieringManager();
@@ -126,6 +141,7 @@ public class RebalanceManagerTest {
TestingMetricGroups.COORDINATOR_METRICS,
new Configuration(),
Executors.newFixedThreadPool(1, new ExecutorThreadFactory("test-coordinator-io")),
- metadataManager);
+ metadataManager,
+ kvSnapshotLeaseManager);
}
}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachineTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachineTest.java
index bafb477c5..26dfeb8dc 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachineTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachineTest.java
@@ -35,6 +35,7 @@ import org.apache.fluss.server.coordinator.LakeTableTieringManager;
import org.apache.fluss.server.coordinator.MetadataManager;
import org.apache.fluss.server.coordinator.TestCoordinatorChannelManager;
import org.apache.fluss.server.coordinator.event.CoordinatorEventManager;
+import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.coordinator.statemachine.ReplicaLeaderElection.ControlledShutdownLeaderElection;
import org.apache.fluss.server.coordinator.statemachine.TableBucketStateMachine.ElectionResult;
import org.apache.fluss.server.metadata.CoordinatorMetadataCache;
@@ -45,6 +46,7 @@ import org.apache.fluss.server.zk.ZooKeeperExtension;
import org.apache.fluss.server.zk.data.LeaderAndIsr;
import org.apache.fluss.shaded.guava32.com.google.common.collect.Sets;
import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
import org.junit.jupiter.api.BeforeAll;
@@ -86,6 +88,7 @@ class TableBucketStateMachineTest {
private AutoPartitionManager autoPartitionManager;
private LakeTableTieringManager lakeTableTieringManager;
private CoordinatorMetadataCache serverMetadataCache;
+ private KvSnapshotLeaseManager kvSnapshotLeaseManager;
@BeforeAll
static void baseBeforeAll() {
@@ -99,7 +102,8 @@ class TableBucketStateMachineTest {
void beforeEach() throws IOException {
Configuration conf = new Configuration();
conf.setString(ConfigOptions.COORDINATOR_HOST, "localhost");
- conf.setString(ConfigOptions.REMOTE_DATA_DIR, "/tmp/fluss/remote-data");
+ String remoteDir = "/tmp/fluss/remote-data";
+ conf.setString(ConfigOptions.REMOTE_DATA_DIR, remoteDir);
coordinatorContext = new CoordinatorContext();
testCoordinatorChannelManager = new TestCoordinatorChannelManager();
coordinatorRequestBatch =
@@ -119,6 +123,15 @@ class TableBucketStateMachineTest {
new LakeCatalogDynamicLoader(new Configuration(), null, true)),
new Configuration());
lakeTableTieringManager = new LakeTableTieringManager();
+
+ kvSnapshotLeaseManager =
+ new KvSnapshotLeaseManager(
+ Duration.ofMinutes(10).toMillis(),
+ zookeeperClient,
+ remoteDir,
+ SystemClock.getInstance(),
+ TestingMetricGroups.COORDINATOR_METRICS);
+ kvSnapshotLeaseManager.start();
}
@Test
@@ -268,7 +281,8 @@ class TableBucketStateMachineTest {
new MetadataManager(
zookeeperClient,
new Configuration(),
- new LakeCatalogDynamicLoader(new Configuration(), null, true)));
+ new LakeCatalogDynamicLoader(new Configuration(), null, true)),
+ kvSnapshotLeaseManager);
CoordinatorEventManager eventManager =
new CoordinatorEventManager(
coordinatorEventProcessor, TestingMetricGroups.COORDINATOR_METRICS);
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStoreTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStoreTest.java
index 1ab73cfee..659ab3fbe 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStoreTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/CompletedSnapshotStoreTest.java
@@ -505,6 +505,226 @@ class CompletedSnapshotStoreTest {
}
}
+ @Test
+ void testLeasedSnapshotDoesNotBlockSubsumedSnapshotCleanup() throws Exception {
+ // Simulate that snapshot 1 has a lease and cannot be subsumed.
+ // The snapshotInUse checker returns true for snapshot 1, meaning it is in-use.
+ CompletedSnapshotStore store =
+ createCompletedSnapshotStore(
+ 1,
+ defaultHandleStore,
+ Collections.emptyList(),
+ bucket -> bucket.getSnapshotID() == 1L);
+
+ CompletedSnapshot cp1 = getSnapshot(1L);
+ CompletedSnapshot cp2 = getSnapshot(2L);
+ CompletedSnapshot cp3 = getSnapshot(3L);
+
+ store.add(cp1);
+ assertThat(store.getAllSnapshots()).hasSize(1);
+ assertThat(store.getAllSnapshots().get(0).getSnapshotID()).isEqualTo(1L);
+ assertThat(store.stillInUseSnapshots).isEmpty();
+
+ store.add(cp2);
+
+ // After adding cp1 and cp2 with maxRetain=1:
+ // cp1 can't be subsumed (leased) → moved to stillInUseSnapshots
+ // cp2 is latest → retained in completedSnapshots
+ assertThat(store.getAllSnapshots()).hasSize(1);
+ assertThat(store.getAllSnapshots().get(0).getSnapshotID()).isEqualTo(2L);
+ assertThat(store.stillInUseSnapshots).containsKey(1L);
+
+ store.add(cp3);
+
+ // After adding cp3 with maxRetain=1:
+ // cp1 is still in stillInUseSnapshots (lease not expired)
+ // cp2 CAN be subsumed → subsumed and cleaned.
+ // With the fix, cleanSubsumedSnapshots uses upTo = snapshot.getSnapshotID() + 1 = 4,
+ // and stillInUseIds = {1L}, so cp2 (ID=2, not in stillInUse) will be properly cleaned.
+ // cp3 is latest → retained in completedSnapshots
+ List<CompletedSnapshot> remaining = store.getAllSnapshots();
+ assertThat(remaining).hasSize(1);
+ assertThat(remaining.get(0).getSnapshotID()).isEqualTo(3L);
+ assertThat(store.stillInUseSnapshots).containsKey(1L);
+ }
+
+ @Test
+ void testLeasedSnapshotSSTFileCleanupWithStillInUse() throws Exception {
+ // Lease protects snapshot 1.
+ SharedKvFileRegistry registry = new SharedKvFileRegistry();
+
+ TestKvHandle fileA = new TestKvHandle("fileA");
+ TestKvHandle fileB = new TestKvHandle("fileB");
+ TestKvHandle fileC = new TestKvHandle("fileC");
+ TestKvHandle fileD = new TestKvHandle("fileD");
+ TestKvHandle fileE = new TestKvHandle("fileE");
+ TestKvHandle fileF = new TestKvHandle("fileF");
+
+ // Snapshot 1: fileA (new), fileB (new)
+ CompletedSnapshot cp1 =
+ getSnapshotWithSharedFiles(
+ 1L,
+ Arrays.asList(
+ KvFileHandleAndLocalPath.of(fileA, "a.sst"),
+ KvFileHandleAndLocalPath.of(fileB, "b.sst")));
+ // Snapshot 2: fileA (placeholder), fileC (new)
+ CompletedSnapshot cp2 =
+ getSnapshotWithSharedFiles(
+ 2L,
+ Arrays.asList(
+ KvFileHandleAndLocalPath.of(
+ new PlaceholderKvFileHandler(fileA), "a.sst"),
+ KvFileHandleAndLocalPath.of(fileC, "c.sst")));
+ // Snapshot 3: fileA (placeholder), fileD (new)
+ CompletedSnapshot cp3 =
+ getSnapshotWithSharedFiles(
+ 3L,
+ Arrays.asList(
+ KvFileHandleAndLocalPath.of(
+ new PlaceholderKvFileHandler(fileA), "a.sst"),
+ KvFileHandleAndLocalPath.of(fileD, "d.sst")));
+ // Snapshot 4: fileA (placeholder), fileE (new)
+ CompletedSnapshot cp4 =
+ getSnapshotWithSharedFiles(
+ 4L,
+ Arrays.asList(
+ KvFileHandleAndLocalPath.of(
+ new PlaceholderKvFileHandler(fileA), "a.sst"),
+ KvFileHandleAndLocalPath.of(fileE, "e.sst")));
+ // Snapshot 5: fileA (placeholder), fileF (new)
+ CompletedSnapshot cp5 =
+ getSnapshotWithSharedFiles(
+ 5L,
+ Arrays.asList(
+ KvFileHandleAndLocalPath.of(
+ new PlaceholderKvFileHandler(fileA), "a.sst"),
+ KvFileHandleAndLocalPath.of(fileF, "f.sst")));
+
+ CompletedSnapshotStore store =
+ new CompletedSnapshotStore(
+ 2,
+ registry,
+ Collections.emptyList(),
+ defaultHandleStore,
+ executorService,
+ bucket -> bucket.getSnapshotID() == 1L);
+
+ store.add(cp1);
+ store.add(cp2);
+ store.add(cp3);
+ store.add(cp4);
+ store.add(cp5);
+
+ // Wait for async cleanup
+ Thread.sleep(200);
+
+ // Verify snapshot state: completedSnapshots = [4, 5], stillInUse = {1}
+ assertThat(store.getAllSnapshots()).hasSize(2);
+ assertThat(
+ store.getAllSnapshots().stream()
+ .map(CompletedSnapshot::getSnapshotID)
+ .collect(Collectors.toList()))
+ .containsExactly(4L, 5L);
+ assertThat(store.stillInUseSnapshots).hasSize(1);
+ assertThat(store.stillInUseSnapshots).containsKey(1L);
+
+ // Verify SST file cleanup:
+ // fileA: shared across all, lastUsed=5, NOT deleted
+ assertThat(fileA.discarded).isFalse();
+ // fileB: exclusive to snapshot 1, protected by lease
+ assertThat(fileB.discarded).isFalse();
+ // fileC: exclusive to snapshot 2, not protected, DELETED
+ assertThat(fileC.discarded).isTrue();
+ // fileD: exclusive to snapshot 3, not protected, DELETED
+ assertThat(fileD.discarded).isTrue();
+ // fileE: exclusive to snapshot 4, still in retention window, NOT deleted
+ assertThat(fileE.discarded).isFalse();
+ // fileF: exclusive to snapshot 5, still in retention window, NOT deleted
+ assertThat(fileF.discarded).isFalse();
+ }
+
+ @Test
+ void testLeaseExpirySSTCleanup() throws Exception {
+ // Use a mutable set to control which snapshots are leased.
+ Set<Long> leasedSnapshotIds = new HashSet<>();
+ leasedSnapshotIds.add(1L);
+
+ SharedKvFileRegistry registry = new SharedKvFileRegistry();
+
+ TestKvHandle fileA = new TestKvHandle("fileA");
+ TestKvHandle fileB = new TestKvHandle("fileB"); // exclusive to snapshot 1
+
+ CompletedSnapshot cp1 =
+ getSnapshotWithSharedFiles(
+ 1L,
+ Arrays.asList(
+ KvFileHandleAndLocalPath.of(fileA, "a.sst"),
+ KvFileHandleAndLocalPath.of(fileB, "b.sst")));
+ CompletedSnapshot cp2 =
+ getSnapshotWithSharedFiles(
+ 2L,
+ Collections.singletonList(
+ KvFileHandleAndLocalPath.of(
+ new PlaceholderKvFileHandler(fileA), "a.sst")));
+
+ CompletedSnapshotStore store =
+ new CompletedSnapshotStore(
+ 1,
+ registry,
+ Collections.emptyList(),
+ defaultHandleStore,
+ executorService,
+ bucket -> leasedSnapshotIds.contains(bucket.getSnapshotID()));
+
+ store.add(cp1);
+ store.add(cp2);
+
+ // cp1 leased → moved to stillInUseSnapshots, completedSnapshots = [2]
+ assertThat(store.stillInUseSnapshots).containsKey(1L);
+ assertThat(store.getAllSnapshots()).hasSize(1);
+ // fileB is protected by lease
+ assertThat(fileB.discarded).isFalse();
+
+ // Simulate lease expiry.
+ leasedSnapshotIds.remove(1L);
+
+ // Add a new snapshot to trigger the release check.
+ CompletedSnapshot cp3 =
+ getSnapshotWithSharedFiles(
+ 3L,
+ Collections.singletonList(
+ KvFileHandleAndLocalPath.of(
+ new PlaceholderKvFileHandler(fileA), "a.sst")));
+ store.add(cp3);
+ Thread.sleep(200);
+
+ // After lease expiry: stillInUseSnapshots should be empty
+ assertThat(store.stillInUseSnapshots).isEmpty();
+ assertThat(store.getAllSnapshots()).hasSize(1);
+ assertThat(store.getAllSnapshots().get(0).getSnapshotID()).isEqualTo(3L);
+
+ // fileB should now be cleaned up (no longer protected)
+ assertThat(fileB.discarded).isTrue();
+ // fileA is still used by the latest snapshot
+ assertThat(fileA.discarded).isFalse();
+ }
+
+ private CompletedSnapshotStore createCompletedSnapshotStore(
+ int numToRetain,
+ CompletedSnapshotHandleStore snapshotHandleStore,
+ Collection<CompletedSnapshot> completedSnapshots,
+ CompletedSnapshotStore.SnapshotInUseChecker snapshotInUseChecker) {
+
+ SharedKvFileRegistry sharedKvFileRegistry = new SharedKvFileRegistry();
+ return new CompletedSnapshotStore(
+ numToRetain,
+ sharedKvFileRegistry,
+ completedSnapshots,
+ snapshotHandleStore,
+ executorService,
+ snapshotInUseChecker);
+ }
+
private List<CompletedSnapshot> mapToCompletedSnapshot(
List<Tuple2<CompletedSnapshotHandle, String>> snapshotHandles) {
return snapshotHandles.stream()
@@ -529,6 +749,16 @@ class CompletedSnapshotStoreTest {
new KvSnapshotHandle(Collections.emptyList(), Collections.emptyList(), 0));
}
+ private CompletedSnapshot getSnapshotWithSharedFiles(
+ long id, List<KvFileHandleAndLocalPath> sharedFiles) {
+ TableBucket tableBucket = new TableBucket(1, 1);
+ return new CompletedSnapshot(
+ tableBucket,
+ id,
+ new FsPath(tempDir.toString(), "snapshot_" + id),
+ new KvSnapshotHandle(sharedFiles, Collections.emptyList(), 0));
+ }
+
private void testSnapshotRetention(
int numToRetain,
List<CompletedSnapshot> completed,
@@ -557,7 +787,8 @@ class CompletedSnapshotStoreTest {
sharedKvFileRegistry,
completedSnapshots,
snapshotHandleStore,
- executorService);
+ executorService,
+ bucket -> false);
}
private List<Tuple2<CompletedSnapshotHandle, String>> createSnapshotHandles(int num) {
@@ -582,4 +813,21 @@ class CompletedSnapshotStoreTest {
}
return stateHandles;
}
+
+ /** A test KvFileHandle that tracks whether it has been discarded. */
+ private static class TestKvHandle extends KvFileHandle {
+
+ private static final long serialVersionUID = 1L;
+
+ boolean discarded;
+
+ TestKvHandle(String path) {
+ super(path, 10);
+ }
+
+ @Override
+ public void discard() throws Exception {
+ this.discarded = true;
+ }
+ }
}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvTabletSnapshotTargetTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvTabletSnapshotTargetTest.java
index 4077e0e65..98ca31605 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvTabletSnapshotTargetTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/KvTabletSnapshotTargetTest.java
@@ -480,7 +480,8 @@ class KvTabletSnapshotTargetTest {
sharedKvFileRegistry,
Collections.emptyList(),
snapshotHandleStore,
- executor);
+ executor,
+ (consumeKvSnapshotForBucket) -> false); // only retain the latest snapshot.
RocksIncrementalSnapshot rocksIncrementalSnapshot =
createIncrementalSnapshot(snapshotFailType);
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistryTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistryTest.java
index 6a178b40c..a8082cb92 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistryTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/SharedKvFileRegistryTest.java
@@ -19,6 +19,10 @@ package org.apache.fluss.server.kv.snapshot;
import org.junit.jupiter.api.Test;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+
import static org.assertj.core.api.Assertions.assertThat;
/** Test for {@link org.apache.fluss.server.kv.snapshot.SharedKvFileRegistry}. */
@@ -77,6 +81,92 @@ class SharedKvFileRegistryTest {
sharedStateRegistry.unregisterUnusedKvFile(Long.MAX_VALUE);
}
+ /**
+ * Test that stillInUse snapshot IDs protect KV files from being deleted during unregister.
+ *
+ * <p>Scenario: file created by snapshot 1, used by snapshot 1 through 3. If snapshot 1 is
+ * leased (in stillInUse), the file should NOT be deleted even if lowestSnapshotID > 3.
+ */
+ @Test
+ void testUnregisterWithStillInUseProtection() {
+ SharedKvFileRegistry registry = new SharedKvFileRegistry();
+
+ // Register fileA: created by snapshot 1
+ TestKvHandle fileA = new TestKvHandle("fileA");
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(fileA), fileA, 1L);
+
+ // fileA used by snapshot 2 and 3 (placeholder references)
+ PlaceholderKvFileHandler placeholder2 = new PlaceholderKvFileHandler(fileA);
+ registry.registerReference(
+ SharedKvFileRegistryKey.fromKvFileHandle(placeholder2), placeholder2, 2L);
+ PlaceholderKvFileHandler placeholder3 = new PlaceholderKvFileHandler(fileA);
+ registry.registerReference(
+ SharedKvFileRegistryKey.fromKvFileHandle(placeholder3), placeholder3, 3L);
+
+ // Register fileB: created by snapshot 2, only used by snapshot 2
+ TestKvHandle fileB = new TestKvHandle("fileB");
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(fileB), fileB, 2L);
+
+ // Register fileC: created by snapshot 4, used by snapshot 4
+ TestKvHandle fileC = new TestKvHandle("fileC");
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(fileC), fileC, 4L);
+
+ // lowestSnapshotID = 4, stillInUse = {1}
+ // fileA: createdBy=1, lastUsed=3 → lastUsed(3) < 4, but snapshot 1 ∈ [1,3] → NOT deleted
+ // fileB: createdBy=2, lastUsed=2 → lastUsed(2) < 4, snapshot 1 NOT ∈ [2,2] → DELETED
+ // fileC: createdBy=4, lastUsed=4 → lastUsed(4) >= 4 → NOT deleted
+ Set<Long> stillInUse = new HashSet<>();
+ stillInUse.add(1L);
+ registry.unregisterUnusedKvFile(4L, stillInUse);
+
+ assertThat(fileA.discarded).isFalse();
+ assertThat(fileB.discarded).isTrue();
+ assertThat(fileC.discarded).isFalse();
+
+ // Now remove the lease: stillInUse = empty
+ // fileA: lastUsed(3) < 4, no stillInUse protection → DELETED
+ registry.unregisterUnusedKvFile(4L, Collections.emptySet());
+ assertThat(fileA.discarded).isTrue();
+ // fileC still NOT deleted
+ assertThat(fileC.discarded).isFalse();
+ }
+
+ /** Test with multiple leased snapshots protecting different file ranges. */
+ @Test
+ void testUnregisterWithMultipleStillInUseSnapshots() {
+ SharedKvFileRegistry registry = new SharedKvFileRegistry();
+
+ // fileA: createdBy=1, lastUsed=3
+ TestKvHandle fileA = new TestKvHandle("fileA");
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(fileA), fileA, 1L);
+ PlaceholderKvFileHandler pA3 = new PlaceholderKvFileHandler(fileA);
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(pA3), pA3, 3L);
+
+ // fileB: createdBy=2, lastUsed=4
+ TestKvHandle fileB = new TestKvHandle("fileB");
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(fileB), fileB, 2L);
+ PlaceholderKvFileHandler pB4 = new PlaceholderKvFileHandler(fileB);
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(pB4), pB4, 4L);
+
+ // fileC: createdBy=5, lastUsed=5
+ TestKvHandle fileC = new TestKvHandle("fileC");
+ registry.registerReference(SharedKvFileRegistryKey.fromKvFileHandle(fileC), fileC, 5L);
+
+ // lowestSnapshotID = 6, stillInUse = {1, 3}
+ // fileA: lastUsed(3) < 6, snapshot 1 ∈ [1,3] → protected
+ // fileA: lastUsed(3) < 6, snapshot 3 ∈ [1,3] → protected
+ // fileB: lastUsed(4) < 6, snapshot 1 NOT ∈ [2,4], snapshot 3 ∈ [2,4] → protected
+ // fileC: lastUsed(5) < 6, snapshot 1 NOT ∈ [5,5], snapshot 3 NOT ∈ [5,5] → DELETED
+ Set<Long> stillInUse = new HashSet<>();
+ stillInUse.add(1L);
+ stillInUse.add(3L);
+ registry.unregisterUnusedKvFile(6L, stillInUse);
+
+ assertThat(fileA.discarded).isFalse();
+ assertThat(fileB.discarded).isFalse();
+ assertThat(fileC.discarded).isTrue();
+ }
+
private static class TestKvHandle extends KvFileHandle {
private static final long serialVersionUID = 1;
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/ZooKeeperCompletedSnapshotStoreTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/ZooKeeperCompletedSnapshotStoreTest.java
index c6e90ff63..c207726a9 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/ZooKeeperCompletedSnapshotStoreTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/snapshot/ZooKeeperCompletedSnapshotStoreTest.java
@@ -139,6 +139,7 @@ class ZooKeeperCompletedSnapshotStoreTest {
sharedKvFileRegistry,
Collections.emptyList(),
snapshotsInZooKeeper,
- Executors.directExecutor());
+ Executors.directExecutor(),
+ (consumeKvSnapshotForBucket) -> false); // only retain the latest snapshot.
}
}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java b/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java
index 8eaf29a0b..90f28dd95 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/testutils/FlussClusterExtension.java
@@ -268,7 +268,7 @@ public final class FlussClusterExtension
conf.setString(ConfigOptions.ZOOKEEPER_ADDRESS, zooKeeperServer.getConnectString());
conf.setString(ConfigOptions.BIND_LISTENERS, coordinatorServerListeners);
setRemoteDataDir(conf);
- coordinatorServer = new CoordinatorServer(conf);
+ coordinatorServer = new CoordinatorServer(conf, clock);
coordinatorServer.start();
coordinatorServerInfo =
// TODO, Currently, we use 0 as coordinator server id.
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java
index 94cb34b23..4cf0a511e 100644
--- a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java
+++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java
@@ -23,6 +23,7 @@ import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket;
import org.apache.fluss.cluster.rebalance.ServerTag;
import org.apache.fluss.config.ConfigOptions;
import org.apache.fluss.config.Configuration;
+import org.apache.fluss.fs.FsPath;
import org.apache.fluss.metadata.DatabaseSummary;
import org.apache.fluss.metadata.Schema;
import org.apache.fluss.metadata.SchemaInfo;
@@ -41,6 +42,7 @@ import org.apache.fluss.server.zk.data.ServerTags;
import org.apache.fluss.server.zk.data.TableAssignment;
import org.apache.fluss.server.zk.data.TableRegistration;
import org.apache.fluss.server.zk.data.TabletServerRegistration;
+import org.apache.fluss.server.zk.data.lease.KvSnapshotLeaseMetadata;
import org.apache.fluss.shaded.curator5.org.apache.curator.CuratorZookeeperClient;
import org.apache.fluss.shaded.curator5.org.apache.curator.framework.CuratorFramework;
import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.KeeperException;
@@ -489,6 +491,36 @@ class ZooKeeperClientTest {
assertThat(zookeeperClient.getTableBucketSnapshot(table1Bucket2, 1)).isEmpty();
}
+ @Test
+ void testKvSnapshotLease() throws Exception {
+ Map<Long, FsPath> tableIdToRemotePath = new HashMap<>();
+ tableIdToRemotePath.put(150002L, new FsPath("/test/cp1"));
+ KvSnapshotLeaseMetadata leaseMetadata =
+ new KvSnapshotLeaseMetadata(1000L, tableIdToRemotePath);
+
+ assertThat(zookeeperClient.getKvSnapshotLeasesList()).isEmpty();
+ zookeeperClient.registerKvSnapshotLeaseMetadata("lease1", leaseMetadata);
+ assertThat(zookeeperClient.getKvSnapshotLeasesList()).containsExactly("lease1");
+
+ Optional<KvSnapshotLeaseMetadata> metadataOpt =
+ zookeeperClient.getKvSnapshotLeaseMetadata("lease1");
+ assertThat(metadataOpt.isPresent()).isTrue();
+ assertThat(metadataOpt.get()).isEqualTo(leaseMetadata);
+
+ tableIdToRemotePath = new HashMap<>();
+ tableIdToRemotePath.put(150002L, new FsPath("/test/cp2"));
+ leaseMetadata = new KvSnapshotLeaseMetadata(1000L, tableIdToRemotePath);
+ zookeeperClient.updateKvSnapshotLeaseMetadata("lease1", leaseMetadata);
+
+ metadataOpt = zookeeperClient.getKvSnapshotLeaseMetadata("lease1");
+ assertThat(metadataOpt.isPresent()).isTrue();
+ assertThat(metadataOpt.get()).isEqualTo(leaseMetadata);
+
+ zookeeperClient.deleteKvSnapshotLease("lease1");
+ metadataOpt = zookeeperClient.getKvSnapshotLeaseMetadata("lease1");
+ assertThat(metadataOpt).isNotPresent();
+ }
+
@Test
void testGetWriterIdAndIncrement() throws Exception {
// init
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerdeTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerdeTest.java
new file mode 100644
index 000000000..b4ed0654d
--- /dev/null
+++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotLeaseMetadataJsonSerdeTest.java
@@ -0,0 +1,57 @@
+/*
+ * 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.utils.json.JsonSerdeTestBase;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Test for {@link KvSnapshotLeaseMetadataJsonSerde}. */
+public class KvSnapshotLeaseMetadataJsonSerdeTest
+ extends JsonSerdeTestBase<KvSnapshotLeaseMetadata> {
+
+ KvSnapshotLeaseMetadataJsonSerdeTest() {
+ super(KvSnapshotLeaseMetadataJsonSerde.INSTANCE);
+ }
+
+ @Override
+ protected KvSnapshotLeaseMetadata[] createObjects() {
+ KvSnapshotLeaseMetadata[] kvSnapshotLeaseMetadata = new KvSnapshotLeaseMetadata[2];
+ Map<Long, FsPath> tableIdToRemoteMetadataFilePath = new HashMap<>();
+ tableIdToRemoteMetadataFilePath.put(1L, new FsPath("/path/to/metadata1"));
+ tableIdToRemoteMetadataFilePath.put(2L, new FsPath("/path/to/metadata2"));
+ kvSnapshotLeaseMetadata[0] =
+ new KvSnapshotLeaseMetadata(1735538268L, tableIdToRemoteMetadataFilePath);
+ kvSnapshotLeaseMetadata[1] =
+ new KvSnapshotLeaseMetadata(1735538268L, Collections.emptyMap());
+ return kvSnapshotLeaseMetadata;
+ }
+
+ @Override
+ protected String[] expectedJsons() {
+ return new String[] {
+ "{\"version\":1,\"expiration_time\":1735538268,\"tables\":"
+ + "[{\"table_id\":1,\"lease_metadata_path\":\"/path/to/metadata1\"},"
+ + "{\"table_id\":2,\"lease_metadata_path\":\"/path/to/metadata2\"}]}",
+ "{\"version\":1,\"expiration_time\":1735538268,\"tables\":[]}"
+ };
+ }
+}
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLeaseJsonSerdeTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLeaseJsonSerdeTest.java
new file mode 100644
index 000000000..de77bb777
--- /dev/null
+++ b/fluss-server/src/test/java/org/apache/fluss/server/zk/data/lease/KvSnapshotTableLeaseJsonSerdeTest.java
@@ -0,0 +1,54 @@
+/*
+ * 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.json.JsonSerdeTestBase;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/** Test for {@link KvSnapshotTableLeaseJsonSerde}. */
+public class KvSnapshotTableLeaseJsonSerdeTest extends JsonSerdeTestBase<KvSnapshotTableLease> {
+
+ KvSnapshotTableLeaseJsonSerdeTest() {
+ super(KvSnapshotTableLeaseJsonSerde.INSTANCE);
+ }
+
+ @Override
+ protected KvSnapshotTableLease[] createObjects() {
+ KvSnapshotTableLease[] kvSnapshotTableLeases = new KvSnapshotTableLease[2];
+ kvSnapshotTableLeases[0] = new KvSnapshotTableLease(1L, new Long[] {1L, -1L, 1L, 2L});
+
+ Map<Long, Long[]> partitionSnapshots = new HashMap<>();
+ partitionSnapshots.put(2001L, new Long[] {10L, -1L, 20L, 30L});
+ partitionSnapshots.put(2002L, new Long[] {15L, -1L, 25L, 35L});
+ kvSnapshotTableLeases[1] = new KvSnapshotTableLease(2L, partitionSnapshots);
+
+ return kvSnapshotTableLeases;
+ }
+
+ @Override
+ protected String[] expectedJsons() {
+ return new String[] {
+ "{\"version\":1,\"table_id\":1,\"bucket_snapshots\":[1,-1,1,2]}",
+ "{\"version\":1,\"table_id\":2,\"partition_snapshots\":["
+ + "{\"partition_id\":2001,\"bucket_snapshots\":[10,-1,20,30]},"
+ + "{\"partition_id\":2002,\"bucket_snapshots\":[15,-1,25,35]}]}"
+ };
+ }
+}
diff --git a/fluss-test-coverage/pom.xml b/fluss-test-coverage/pom.xml
index d25437225..91e355410 100644
--- a/fluss-test-coverage/pom.xml
+++ b/fluss-test-coverage/pom.xml
@@ -407,6 +407,9 @@
<exclude>
org.apache.fluss.flink.metrics.*
</exclude>
+ <exclude>
+ org.apache.fluss.flink.procedure.*
+ </exclude>
<!-- end exclude for flink-connector -->
<!-- exclude flink 1.18 compatibility class for sink -->
<exclude>org.apache.flink.streaming.api.functions.sink.v2.*</exclude>
diff --git a/website/docs/_configs/_partial_config.mdx b/website/docs/_configs/_partial_config.mdx
index d410834f6..74bce1ad6 100644
--- a/website/docs/_configs/_partial_config.mdx
+++ b/website/docs/_configs/_partial_config.mdx
@@ -78,7 +78,7 @@
| `client.lookup.batch-timeout` | `0 s` | Duration | The maximum time to wait for the lookup batch to full, if this timeout is reached, the lookup batch will be closed to send. |
| `client.lookup.max-retries` | `2147483647` | Integer | Setting a value greater than zero will cause the client to resend any lookup request that fails with a potentially transient error. |
| `client.scanner.remote-log.prefetch-num` | `4` | Integer | 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.scanner.io.tmpdir` | `/tmp/fluss` | String | Local directory that is used by client for storing the data files (like kv snapshot, log segment files) to read temporarily |
+| `client.scanner.io.tmpdir` | `/var/folders/ck/q6rhkhq9759b3nwc7501p6780000gp/T//fluss` | String | Local directory that is used by client for storing the data files (like kv snapshot, log segment files) to read temporarily |
| `client.remote-file.download-thread-num` | `3` | Integer | The number of threads the client uses to download remote files. |
| `client.filesystem.security.token.renewal.backoff` | `1 hours` | Duration | The time period how long to wait before retrying to obtain new security tokens for filesystem after a failure. |
| `client.filesystem.security.token.renewal.time-ratio` | `0.75` | Double | Ratio of the token's expiration time when new credentials for access filesystem should be re-obtained. |
@@ -95,6 +95,8 @@
| `coordinator.host` | `none` | String | The config parameter defining the network address to connect to for communication with the coordinator server. If the coordinator server is used as a bootstrap server (discover all the servers in the cluster), the value of this config option should be a static hostname or address.This option is deprecated. Please use bind.listeners instead, which provides a more flexible configuration for multiple ports |
| `coordinator.port` | `9123` | String | The config parameter defining the network port to connect to for communication with the coordinator server. Like coordinator.host, if the coordinator server is used as a bootstrap server (discover all the servers in the cluster), the value of this config option should be a static port. Otherwise, the value can be set to "0" for a dynamic service name resolution. The value accepts a list of ports (“50100,50101”), ranges (“50100-50200”) or a combina [...]
| `coordinator.io-pool.size` | `10` | Integer | The size of the IO thread pool to run blocking operations for coordinator server. This includes discard unnecessary snapshot files. Increase this value if you experience slow unnecessary snapshot files clean. The default value is 10. This option is deprecated. Please use server.io-pool.size instead. |
+| `coordinator.producer-offsets.ttl` | `24 hours` | Duration | The TTL (time-to-live) for producer offsets. Producer offsets older than this TTL will be automatically cleaned up by the coordinator server. Default is 24 hours. |
+| `coordinator.producer-offsets.cleanup-interval` | `1 hours` | Duration | The interval for cleaning up expired producer offsets and orphan files in remote storage. Default is 1 hour. |
## Data Configurations
@@ -137,7 +139,8 @@
| `kv.snapshot.interval` | `10 min` | Duration | The interval to perform periodic snapshot for kv data. The default setting is 10 minutes. |
| `kv.snapshot.scheduler-thread-num` | `1` | Integer | The number of threads that the server uses to schedule snapshot kv data for all the replicas in the server. |
| `kv.snapshot.transfer-thread-num` | `4` | Integer | The number of threads the server uses to transfer (download and upload) kv snapshot files. This option is deprecated. Please use server.io-pool.size instead. |
-| `kv.snapshot.num-retained` | `1` | Integer | The maximum number of completed snapshots to retain. |
+| `kv.snapshot.num-retained` | `2` | Integer | The maximum number of completed snapshots to retain. |
+| `kv.snapshot.lease.expiration-check-interval` | `10 min` | Duration | The interval to check the expiration of kv snapshot lease. The default setting is 10 minutes. |
| `kv.rocksdb.thread.num` | `2` | Integer | The maximum number of concurrent background flush and compaction jobs (per bucket of table). The default value is `2`. |
| `kv.rocksdb.files.open` | `-1` | Integer | The maximum number of open files (per bucket of table) that can be used by the DB, `-1` means no limit. The default value is `-1`. |
| `kv.rocksdb.log.max-file-size` | `25 mb` | MemorySize | The maximum size of RocksDB's file used for information logging. If the log files becomes larger than this, a new file will be created. If 0, all logs will be written to one log file. The default maximum file size is `25MB`. |
@@ -206,6 +209,12 @@
| :--- | :--- | :--- | :--- |
| `metrics.reporters` | `none` | String | An optional list of reporter names. If configured, only reporters whose name matches in the list will be started |
| `metrics.reporter.prometheus.port` | `9249` | String | The port the Prometheus reporter listens on.In order to be able to run several instances of the reporter on one host (e.g. when one TabletServer is colocated with the CoordinatorServer) it is advisable to use a port range like 9250-9260. |
+| `metrics.reporter.prometheus-push.host-url` | `none` | String | The PushGateway server host URL including scheme, host name, and port. |
+| `metrics.reporter.prometheus-push.job-name` | `none` | String | The job name under which metrics will be pushed |
+| `metrics.reporter.prometheus-push.random-job-name-suffix` | `true` | Boolean | Specifies whether a random suffix should be appended to the job name. This is useful when multiple instances of the reporter are running on the same host. |
+| `metrics.reporter.prometheus-push.delete-on-shutdown` | `true` | Boolean | Specifies whether to delete metrics from the PushGateway on shutdown. Fluss will try its best to delete the metrics but this is not guaranteed. |
+| `metrics.reporter.prometheus-push.grouping-key` | `none` | String | Specifies the grouping key which is the group and global labels of all metrics. The label name and value are separated by '=', and labels are separated by ';', e.g., k1=v1;k2=v2. |
+| `metrics.reporter.prometheus-push.push-interval` | `10 s` | Duration | The interval of pushing metrics to Prometheus PushGateway. |
| `metrics.reporter.jmx.port` | `none` | String | The port for the JMXServer that JMX clients can connect to. If not set, the JMXServer won't start. In order to be able to run several instances of the reporter on one host (e.g. when one TabletServer is colocated with the CoordinatorServer) it is advisable to use a port range like 9990-9999. |
## Netty Configurations
@@ -271,10 +280,11 @@
| `table.log.arrow.compression.type` | `ZSTD` | ArrowCompressionType | The compression type of the log records if the log format is set to `ARROW`. The candidate compression type is [NONE, LZ4_FRAME, ZSTD] |
| `table.log.arrow.compression.zstd.level` | `3` | Integer | The compression level of ZSTD for the log records if the log format is set to `ARROW` and the compression type is set to `ZSTD`. The valid range is 1 to 22. |
| `table.kv.format` | `COMPACTED` | KvFormat | The format of the kv records in kv store. The default value is `compacted`. The supported formats are `compacted` and `indexed`. |
+| `table.kv.format-version` | `none` | String | The version of the kv format. Automatically set by the coordinator during table creation if not configured by users. Note: The datalake encoding and bucketing strategy mentioned below only takes effect when 'datalake.format' is configured at cluster level. Version Behaviors: (1) Version 1: Tables created before 'table.kv.format-version' was introduced are treated as version 1. Uses datalake's encoder (e.g., Paimon/Iceberg) for both primary [...]
| `table.auto-partition.enabled` | `false` | Boolean | Whether enable auto partition for the table. Disable by default. When auto partition is enabled, the partitions of the table will be created automatically. |
| `table.auto-partition.key` | `none` | String | This configuration defines the time-based partition key to be used for auto-partitioning when a table is partitioned with multiple keys. Auto-partitioning utilizes a time-based partition key to handle partitions automatically, including creating new ones and removing outdated ones, by comparing the time value of the partition with the current system time. In the case of a table using multiple partition keys (such as a composite partitionin [...]
| `table.auto-partition.time-unit` | `DAY` | AutoPartitionTimeUnit | The time granularity for auto created partitions. The default value is `DAY`. Valid values are `HOUR`, `DAY`, `MONTH`, `QUARTER`, `YEAR`. If the value is `HOUR`, the partition format for auto created is yyyyMMddHH. If the value is `DAY`, the partition format for auto created is yyyyMMdd. If the value is `MONTH`, the partition format for auto created is yyyyMM. If the value is `QUARTER`, the partition format for auto cre [...]
-| `table.auto-partition.time-zone` | `Etc/UTC` | String | The time zone for auto partitions, which is by default the same as the system time zone. |
+| `table.auto-partition.time-zone` | `Asia/Shanghai` | String | The time zone for auto partitions, which is by default the same as the system time zone. |
| `table.auto-partition.num-precreate` | `2` | Integer | The number of partitions to pre-create for auto created partitions in each check for auto partition. For example, if the current check time is 2024-11-11 and the value is configured as 3, then partitions 20241111, 20241112, 20241113 will be pre-created. If any one partition exists, it'll skip creating the partition. The default value is 2, which means 2 partitions will be pre-created. If the `table.auto-partition.time-unit` is `DAY [...]
| `table.auto-partition.num-retention` | `7` | Integer | The number of history partitions to retain for auto created partitions in each check for auto partition. For example, if the current check time is 2024-11-11, time-unit is DAY, and the value is configured as 3, then the history partitions 20241108, 20241109, 20241110 will be retained. The partitions earlier than 20241108 will be deleted. The default value is 7. |
| `table.log.ttl` | `168 hours` | Duration | The time to live for log segments. The configuration controls the maximum time we will retain a log before we will delete old segments to free up space. If set to -1, the log will not be deleted. |
diff --git a/website/docs/engine-flink/options.md b/website/docs/engine-flink/options.md
index 7f9463a73..2ed7f5a44 100644
--- a/website/docs/engine-flink/options.md
+++ b/website/docs/engine-flink/options.md
@@ -93,20 +93,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 [...]
-| 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 [...]
+| scan.kv.snapshot.lease.id | String | UUID | 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. [...]
+| 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. [...]
+| client.remote-file.download-thread-num | Integer | 3 | The number of threads the client uses to download remote files. [...]
## Lookup Options
diff --git a/website/docs/engine-flink/procedures.md b/website/docs/engine-flink/procedures.md
index 182ef033e..9cd9cb1d6 100644
--- a/website/docs/engine-flink/procedures.md
+++ b/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 drop leased kv snapshots.
+
+### drop_kv_snapshot_lease
+
+Drop KV snapshots leased under a specified leaseId. This is typically used for handle the scenario of lease
+remnants. After an abnormal job termination (e.g., crash or forced cancellation), the registered lease may not
+be released automatically and could require manual cleanup.
+
+**Syntax:**
+
+```sql
+CALL [catalog_name.]sys.drop_kv_snapshot_lease(
+ leaseId => 'STRING'
+)
+```
+
+**Parameters:**
+
+- `leaseId` (required): The lease identifier of the KV snapshots to release. This should match the lease ID used when acquiring the KV snapshots.
+
+**Returns:** An array with a single element `'success'` if the operation completes successfully.
+
+**Example:**
+
+```sql title="Flink SQL"
+-- Use the Fluss catalog (replace 'fluss_catalog' with your catalog name if different)
+USE fluss_catalog;
+
+-- Drop KV snapshots leased under the given leaseId
+CALL sys.drop_kv_snapshot_lease('test-lease-id');
```
\ No newline at end of file
diff --git a/website/docs/maintenance/configuration.md b/website/docs/maintenance/configuration.md
index 57f7017c8..80d40fddb 100644
--- a/website/docs/maintenance/configuration.md
+++ b/website/docs/maintenance/configuration.md
@@ -139,6 +139,7 @@ during the Fluss cluster working.
| kv.snapshot.scheduler-thread-num | Integer | 1 | The number of threads that the server uses to schedule snapshot kv data for all the replicas in the server. [...]
| kv.snapshot.transfer-thread-num | Integer | 4 | **Deprecated**: This option is deprecated. Please use `server.io-pool.size` instead. The number of threads the server uses to transfer (download and upload) kv snapshot files. [...]
| kv.snapshot.num-retained | Integer | 1 | The maximum number of completed snapshots to retain. [...]
+| kv.snapshot.lease-expiration-check-interval | Duration | 10min | The interval to check the expiration of kv snapshot leases. The default setting is 10 minutes. [...]
| kv.rocksdb.thread.num | Integer | 2 | The maximum number of concurrent background flush and compaction jobs (per bucket of table). The default value is `2`. [...]
| kv.rocksdb.files.open | Integer | -1 | The maximum number of open files (per bucket of table) that can be used by the DB, `-1` means no limit. The default value is `-1`. [...]
| kv.rocksdb.log.max-file-size | MemorySize | 25mb | The maximum size of RocksDB's file used for information logging. If the log files becomes larger than this, a new file will be created. If 0, all logs will be written to one log file. The default maximum file size is `25MB`. [...]
@@ -163,7 +164,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](operatio [...]
+| 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/updat [...]
| kv.recover.log-record-batch.max-size | MemorySize | 16mb | The max fetch size for fetching log to apply to kv during recovering kv. [...]
## Metrics
diff --git a/website/docs/maintenance/observability/monitor-metrics.md b/website/docs/maintenance/observability/monitor-metrics.md
index 4b082b2f2..bf62746b4 100644
--- a/website/docs/maintenance/observability/monitor-metrics.md
+++ b/website/docs/maintenance/observability/monitor-metrics.md
@@ -294,7 +294,7 @@ Some metrics might not be exposed when using other JVM implementations (e.g. IBM
</thead>
<tbody>
<tr>
- <th rowspan="13"><strong>coordinator</strong></th>
+ <th rowspan="15"><strong>coordinator</strong></th>
<td style={{textAlign: 'center', verticalAlign: 'middle' }} rowspan="8">-</td>
<td>activeCoordinatorCount</td>
<td>The number of active CoordinatorServer in this cluster.</td>
@@ -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 snapshot leases in this cluster.</td>
+ <td>Gauge</td>
+ </tr>
+ <tr>
+ <td>leasedKvSnapshotCount</td>
+ <td>The total number of leased kv snapshots in this cluster.</td>
+ <td>Gauge</td>
+ </tr>
<tr>
<td>replicasToDeleteCount</td>
<td>The total number of replicas in the progress to be deleted in this cluster.</td>