[kv] Support kv snapshot consumer

Author: swuferhong

fluss-client/src/main/java/org/apache/fluss/client/FlussConnection.java

@@ -142,7 +142,7 @@ public Admin getOrCreateAdmin() {
         if (admin == null) {
             synchronized (this) {
                 if (admin == null) {
-                    admin = new FlussAdmin(rpcClient, metadataUpdater);
+                    admin = new FlussAdmin(rpcClient, metadataUpdater, conf);
                 }
             }
         }

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

@@ -21,6 +21,7 @@
 import org.apache.fluss.client.metadata.KvSnapshotMetadata;
 import org.apache.fluss.client.metadata.KvSnapshots;
 import org.apache.fluss.client.metadata.LakeSnapshot;
+import org.apache.fluss.client.metadata.RegisterKvSnapshotResult;
 import org.apache.fluss.cluster.ServerNode;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.cluster.AlterConfig;
@@ -33,6 +34,7 @@
 import org.apache.fluss.exception.InvalidPartitionException;
 import org.apache.fluss.exception.InvalidReplicationFactorException;
 import org.apache.fluss.exception.InvalidTableException;
+import org.apache.fluss.exception.KvSnapshotConsumerNotExistException;
 import org.apache.fluss.exception.KvSnapshotNotExistException;
 import org.apache.fluss.exception.LakeTableSnapshotNotExistException;
 import org.apache.fluss.exception.NonPrimaryKeyTableException;
@@ -60,6 +62,8 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -400,6 +404,62 @@ CompletableFuture<Void> dropPartition(
     CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
             TableBucket bucket, long snapshotId);
 
+    /**
+     * Register the kv snapshot consumer of the given tableBucket set asynchronously. After
+     * registered, the kv snapshot will not be deleted until unregistered or the consumer expired.
+     *
+     * <p>For the details bucket to consume, you can call {@link #getLatestKvSnapshots(TablePath)}
+     * or {@link #getLatestKvSnapshots(TablePath, String)} first.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on returned future.
+     *
+     * <ul>
+     *   <li>{@link TableNotExistException} if the table does not exist.
+     *   <li>{@link PartitionNotExistException} if the partition does not exist
+     * </ul>
+     *
+     * @param consumerId the consumer id.
+     * @param consumeBuckets the tableBuckets to consume, a map from TableBucket to kvSnapshotId.
+     * @return the result of registering kv snapshot consumer.
+     */
+    CompletableFuture<RegisterKvSnapshotResult> registerKvSnapshotConsumer(
+            String consumerId, Map<TableBucket, Long> consumeBuckets);
+
+    /**
+     * Unregister the kv snapshot consumer of the given tableBucket set asynchronously.
+     *
+     * <p>As the following situation happens, unregister operation will be partially ignored or
+     * completely ignored
+     *
+     * <ul>
+     *   <li>If the table/partition not exists, the unregister operation will be ignored for this
+     *       table/partition, others will success.
+     *   <li>If the partition bucket/kvSnapshotId not exists, the unregister operation will be
+     *       ignored for this bucket/kvSnapshotId, others will success.
+     *   <li>If the consumerId not exists, the unregister operation will be ignored.
+     * </ul>
+     *
+     * @param consumerId the consumer id.
+     * @param bucketsToUnregister the tableBucket to unregister.
+     */
+    CompletableFuture<Void> unregisterKvSnapshotConsumer(
+            String consumerId, Set<TableBucket> bucketsToUnregister);
+
+    /**
+     * Clear the kv snapshot consumer asynchronously. After cleared, all the registered kv snapshot
+     * consume info for this consumer will be cleared.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on returned future.
+     *
+     * <ul>
+     *   <li>{@link TableNotExistException} if the table does not exist.
+     *   <li>{@link KvSnapshotConsumerNotExistException} if the kv snapshot consumer does not exist.
+     * </ul>
+     *
+     * @param consumerId the consumer id.
+     */
+    CompletableFuture<Void> clearKvSnapshotConsumer(String consumerId);
+
     /**
      * Get table lake snapshot info of the given table asynchronously.
      *

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

@@ -21,9 +21,12 @@
 import org.apache.fluss.client.metadata.KvSnapshots;
 import org.apache.fluss.client.metadata.LakeSnapshot;
 import org.apache.fluss.client.metadata.MetadataUpdater;
+import org.apache.fluss.client.metadata.RegisterKvSnapshotResult;
 import org.apache.fluss.client.utils.ClientRpcMessageUtils;
 import org.apache.fluss.cluster.Cluster;
 import org.apache.fluss.cluster.ServerNode;
+import org.apache.fluss.config.ConfigOptions;
+import org.apache.fluss.config.Configuration;
 import org.apache.fluss.config.cluster.AlterConfig;
 import org.apache.fluss.config.cluster.ConfigEntry;
 import org.apache.fluss.exception.LeaderNotAvailableException;
@@ -46,6 +49,7 @@
 import org.apache.fluss.rpc.gateway.TabletServerGateway;
 import org.apache.fluss.rpc.messages.AlterClusterConfigsRequest;
 import org.apache.fluss.rpc.messages.AlterTableRequest;
+import org.apache.fluss.rpc.messages.ClearKvSnapshotConsumerRequest;
 import org.apache.fluss.rpc.messages.CreateAclsRequest;
 import org.apache.fluss.rpc.messages.CreateDatabaseRequest;
 import org.apache.fluss.rpc.messages.CreateTableRequest;
@@ -87,13 +91,16 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAlterTableRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeCreatePartitionRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeDropPartitionRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeListOffsetsRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makePbPartitionSpec;
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeRegisterKvSnapshotConsumerRequest;
+import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeUnregisterKvSnapshotConsumerRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.toConfigEntries;
 import static org.apache.fluss.client.utils.MetadataUtils.sendMetadataRequestAndRebuildCluster;
 import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toAclBindings;
@@ -112,15 +119,18 @@ public class FlussAdmin implements Admin {
     private final AdminGateway gateway;
     private final AdminReadOnlyGateway readOnlyGateway;
     private final MetadataUpdater metadataUpdater;
+    private final Configuration clientConfig;
 
-    public FlussAdmin(RpcClient client, MetadataUpdater metadataUpdater) {
+    public FlussAdmin(
+            RpcClient client, MetadataUpdater metadataUpdater, Configuration clientConfig) {
         this.gateway =
                 GatewayClientProxy.createGatewayProxy(
                         metadataUpdater::getCoordinatorServer, client, AdminGateway.class);
         this.readOnlyGateway =
                 GatewayClientProxy.createGatewayProxy(
                         metadataUpdater::getRandomTabletServer, client, AdminGateway.class);
         this.metadataUpdater = metadataUpdater;
+        this.clientConfig = clientConfig;
     }
 
     @Override
@@ -377,6 +387,38 @@ public CompletableFuture<KvSnapshotMetadata> getKvSnapshotMetadata(
                 .thenApply(ClientRpcMessageUtils::toKvSnapshotMetadata);
     }
 
+    @Override
+    public CompletableFuture<RegisterKvSnapshotResult> registerKvSnapshotConsumer(
+            String consumerId, Map<TableBucket, Long> consumeBuckets) {
+        if (consumeBuckets.isEmpty()) {
+            throw new IllegalArgumentException("consumeBuckets is empty");
+        }
+
+        long expirationTime =
+                clientConfig
+                        .get(ConfigOptions.CLIENT_SCANNER_KV_SNAPSHOT_CONSUMER_EXPIRATION_TIME)
+                        .toMillis();
+        return gateway.registerKvSnapshotConsumer(
+                        makeRegisterKvSnapshotConsumerRequest(
+                                consumerId, consumeBuckets, expirationTime))
+                .thenApply(ClientRpcMessageUtils::toRegisterKvSnapshotResult);
+    }
+
+    @Override
+    public CompletableFuture<Void> unregisterKvSnapshotConsumer(
+            String consumerId, Set<TableBucket> bucketsToUnregister) {
+        return gateway.unregisterKvSnapshotConsumer(
+                        makeUnregisterKvSnapshotConsumerRequest(consumerId, bucketsToUnregister))
+                .thenApply(r -> null);
+    }
+
+    @Override
+    public CompletableFuture<Void> clearKvSnapshotConsumer(String consumerId) {
+        ClearKvSnapshotConsumerRequest request =
+                new ClearKvSnapshotConsumerRequest().setConsumerId(consumerId);
+        return gateway.clearKvSnapshotConsumer(request).thenApply(r -> null);
+    }
+
     @Override
     public CompletableFuture<LakeSnapshot> getLatestLakeSnapshot(TablePath tablePath) {
         GetLatestLakeSnapshotRequest request = new GetLatestLakeSnapshotRequest();

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 Set<Integer> getBucketIds() {
         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.
      */

fluss-client/src/main/java/org/apache/fluss/client/metadata/RegisterKvSnapshotResult.java

@@ -0,0 +1,46 @@
+/*
+ * 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.Set;
+
+/**
+ * A class to represent the result of registering kv snapshot. It contains:
+ *
+ * <ul>
+ *   <li>An set of failed tableBuckets. Such as the specify snapshotId is not exist for this table
+ *       bucket.
+ * </ul>
+ *
+ * @since 0.9
+ */
+@PublicEvolving
+public class RegisterKvSnapshotResult {
+    private final Set<TableBucket> failedTableBucketSet;
+
+    public RegisterKvSnapshotResult(Set<TableBucket> failedTableBucketSet) {
+        this.failedTableBucketSet = failedTableBucketSet;
+    }
+
+    public Set<TableBucket> getFailedTableBucketSet() {
+        return failedTableBucketSet;
+    }
+}

fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java

@@ -23,6 +23,7 @@
 import org.apache.fluss.client.metadata.KvSnapshotMetadata;
 import org.apache.fluss.client.metadata.KvSnapshots;
 import org.apache.fluss.client.metadata.LakeSnapshot;
+import org.apache.fluss.client.metadata.RegisterKvSnapshotResult;
 import org.apache.fluss.client.write.KvWriteBatch;
 import org.apache.fluss.client.write.ReadyWriteBatch;
 import org.apache.fluss.config.cluster.AlterConfigOpType;
@@ -50,6 +51,8 @@
 import org.apache.fluss.rpc.messages.MetadataRequest;
 import org.apache.fluss.rpc.messages.PbAddColumn;
 import org.apache.fluss.rpc.messages.PbAlterConfig;
+import org.apache.fluss.rpc.messages.PbBucket;
+import org.apache.fluss.rpc.messages.PbConsumeKvSnapshotForBucket;
 import org.apache.fluss.rpc.messages.PbDescribeConfig;
 import org.apache.fluss.rpc.messages.PbDropColumn;
 import org.apache.fluss.rpc.messages.PbKeyValue;
@@ -63,9 +66,13 @@
 import org.apache.fluss.rpc.messages.PbPutKvReqForBucket;
 import org.apache.fluss.rpc.messages.PbRemotePathAndLocalFile;
 import org.apache.fluss.rpc.messages.PbRenameColumn;
+import org.apache.fluss.rpc.messages.PbTable;
 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.RegisterKvSnapshotConsumerRequest;
+import org.apache.fluss.rpc.messages.RegisterKvSnapshotConsumerResponse;
+import org.apache.fluss.rpc.messages.UnregisterKvSnapshotConsumerRequest;
 import org.apache.fluss.utils.json.DataTypeJsonSerde;
 import org.apache.fluss.utils.json.JsonSerdeUtils;
 
@@ -75,6 +82,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -370,6 +378,73 @@ public static AlterTableRequest makeAlterTableRequest(
         return request;
     }
 
+    public static RegisterKvSnapshotConsumerRequest makeRegisterKvSnapshotConsumerRequest(
+            String consumerId, Map<TableBucket, Long> consumeBuckets, long expirationTime) {
+        RegisterKvSnapshotConsumerRequest request = new RegisterKvSnapshotConsumerRequest();
+        request.setConsumerId(consumerId).setExpirationTime(expirationTime);
+
+        Map<Long, List<PbConsumeKvSnapshotForBucket>> pbConsumeTables = new HashMap<>();
+        for (Map.Entry<TableBucket, Long> entry : consumeBuckets.entrySet()) {
+            TableBucket tableBucket = entry.getKey();
+            Long snapshotId = entry.getValue();
+            PbConsumeKvSnapshotForBucket pbConsumeKvSnapshotForBucket =
+                    new PbConsumeKvSnapshotForBucket()
+                            .setBucketId(tableBucket.getBucket())
+                            .setSnapshotId(snapshotId);
+            if (tableBucket.getPartitionId() != null) {
+                pbConsumeKvSnapshotForBucket.setPartitionId(tableBucket.getPartitionId());
+            }
+            pbConsumeTables
+                    .computeIfAbsent(tableBucket.getTableId(), k -> new ArrayList<>())
+                    .add(pbConsumeKvSnapshotForBucket);
+        }
+
+        for (Map.Entry<Long, List<PbConsumeKvSnapshotForBucket>> entry :
+                pbConsumeTables.entrySet()) {
+            request.addRegisterTable()
+                    .setTableId(entry.getKey())
+                    .addAllBucketsReqs(entry.getValue());
+        }
+        return request;
+    }
+
+    public static RegisterKvSnapshotResult toRegisterKvSnapshotResult(
+            RegisterKvSnapshotConsumerResponse response) {
+        Set<TableBucket> failedTableBucketSet = new HashSet<>();
+        for (PbTable failedTable : response.getFailedTablesList()) {
+            long tableId = failedTable.getTableId();
+            for (PbBucket pbBucket : failedTable.getBucketsList()) {
+                TableBucket tableBucket =
+                        new TableBucket(
+                                tableId,
+                                pbBucket.hasPartitionId() ? pbBucket.getPartitionId() : null,
+                                pbBucket.getBucketId());
+                failedTableBucketSet.add(tableBucket);
+            }
+        }
+        return new RegisterKvSnapshotResult(failedTableBucketSet);
+    }
+
+    public static UnregisterKvSnapshotConsumerRequest makeUnregisterKvSnapshotConsumerRequest(
+            String consumerId, Set<TableBucket> bucketsToUnregister) {
+        UnregisterKvSnapshotConsumerRequest request = new UnregisterKvSnapshotConsumerRequest();
+        request.setConsumerId(consumerId);
+
+        Map<Long, List<PbBucket>> pbConsumeTables = new HashMap<>();
+        for (TableBucket tb : bucketsToUnregister) {
+            PbBucket pbBucket = new PbBucket().setBucketId(tb.getBucket());
+            if (tb.getPartitionId() != null) {
+                pbBucket.setPartitionId(tb.getPartitionId());
+            }
+            pbConsumeTables.computeIfAbsent(tb.getTableId(), k -> new ArrayList<>()).add(pbBucket);
+        }
+
+        for (Map.Entry<Long, List<PbBucket>> entry : pbConsumeTables.entrySet()) {
+            request.addUnregisterTable().setTableId(entry.getKey()).addAllBuckets(entry.getValue());
+        }
+        return request;
+    }
+
     public static List<PartitionInfo> toPartitionInfos(ListPartitionInfosResponse response) {
         return response.getPartitionsInfosList().stream()
                 .map(