[server] add lake-aware auto-partition retention

fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java

@@ -1775,6 +1775,15 @@ public class ConfigOptions {
                                     + "The partitions earlier than 20241108 will be deleted. "
                                     + "The default value is 7.");
 
+    public static final ConfigOption<Boolean> TABLE_AUTO_PARTITION_RETENTION_ENSURE_TIERED =
+            key("table.auto-partition.retention.ensure-tiered")
+                    .booleanType()
+                    .defaultValue(false)
+                    .withDescription(
+                            "Whether an expired auto partition must be fully committed to the data lake before deletion. "
+                                    + "When enabled, writes to the partition are frozen first, and the partition is deleted only after the lake contains all frozen log offsets. "
+                                    + "The default value is false.");
+
     public static final ConfigOption<Duration> TABLE_LOG_TTL =
             key("table.log.ttl")
                     .durationType()

fluss-common/src/main/java/org/apache/fluss/utils/AutoPartitionStrategy.java

@@ -34,6 +34,7 @@ public class AutoPartitionStrategy {
     private final String timeFormat;
     private final int numPreCreate;
     private final int numToRetain;
+    private final boolean ensureLakeTieredBeforeDrop;
     private final TimeZone timeZone;
 
     private AutoPartitionStrategy(
@@ -43,13 +44,15 @@ private AutoPartitionStrategy(
             String timeFormat,
             int numPreCreate,
             int numToRetain,
+            boolean ensureLakeTieredBeforeDrop,
             TimeZone timeZone) {
         this.autoPartitionEnabled = autoPartitionEnabled;
         this.key = key;
         this.timeUnit = autoPartitionTimeUnit;
         this.timeFormat = timeFormat;
         this.numPreCreate = numPreCreate;
         this.numToRetain = numToRetain;
+        this.ensureLakeTieredBeforeDrop = ensureLakeTieredBeforeDrop;
         this.timeZone = timeZone;
     }
 
@@ -65,6 +68,7 @@ public static AutoPartitionStrategy from(Configuration conf) {
                 conf.getOptional(ConfigOptions.TABLE_AUTO_PARTITION_TIME_FORMAT).orElse(null),
                 conf.getInt(ConfigOptions.TABLE_AUTO_PARTITION_NUM_PRECREATE),
                 conf.getInt(ConfigOptions.TABLE_AUTO_PARTITION_NUM_RETENTION),
+                conf.getBoolean(ConfigOptions.TABLE_AUTO_PARTITION_RETENTION_ENSURE_TIERED),
                 TimeZone.getTimeZone(conf.getString(ConfigOptions.TABLE_AUTO_PARTITION_TIMEZONE)));
     }
 
@@ -92,6 +96,10 @@ public int numToRetain() {
         return numToRetain;
     }
 
+    public boolean ensureLakeTieredBeforeDrop() {
+        return ensureLakeTieredBeforeDrop;
+    }
+
     public TimeZone timeZone() {
         return timeZone;
     }

fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/TabletServerGateway.java

@@ -20,6 +20,8 @@
 import org.apache.fluss.rpc.RpcGateway;
 import org.apache.fluss.rpc.messages.FetchLogRequest;
 import org.apache.fluss.rpc.messages.FetchLogResponse;
+import org.apache.fluss.rpc.messages.FreezePartitionRequest;
+import org.apache.fluss.rpc.messages.FreezePartitionResponse;
 import org.apache.fluss.rpc.messages.GetTableStatsRequest;
 import org.apache.fluss.rpc.messages.GetTableStatsResponse;
 import org.apache.fluss.rpc.messages.InitWriterRequest;
@@ -84,6 +86,10 @@ CompletableFuture<NotifyLeaderAndIsrResponse> notifyLeaderAndIsr(
     @RPC(api = ApiKeys.STOP_REPLICA)
     CompletableFuture<StopReplicaResponse> stopReplica(StopReplicaRequest stopBucketReplicaRequest);
 
+    /** Freeze writes to partition bucket leaders for lake-aware retention. */
+    @RPC(api = ApiKeys.FREEZE_PARTITION)
+    CompletableFuture<FreezePartitionResponse> freezePartition(FreezePartitionRequest request);
+
     /**
      * Produce log data to the specified table bucket.
      *

fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java

@@ -106,7 +106,8 @@ public enum ApiKeys {
     SCAN_KV(1061, 0, 0, PUBLIC),
     GET_CLUSTER_HEALTH(1062, 0, 0, PUBLIC),
     LIST_REMOTE_LOG_MANIFESTS(1063, 0, 0, PUBLIC),
-    LIST_KV_SNAPSHOTS(1064, 0, 0, PUBLIC);
+    LIST_KV_SNAPSHOTS(1064, 0, 0, PUBLIC),
+    FREEZE_PARTITION(1065, 0, 0, PRIVATE);
 
     private static final Map<Integer, ApiKeys> ID_TO_TYPE =
             Arrays.stream(ApiKeys.values())

fluss-rpc/src/main/proto/FlussApi.proto

@@ -375,6 +375,16 @@ message StopReplicaResponse {
   repeated PbStopReplicaRespForBucket stop_replicas_resp = 1;
 }
 
+// Freeze writes to partition bucket leaders before lake-aware retention.
+message FreezePartitionRequest {
+  required int32 coordinator_epoch = 1;
+  repeated PbFreezePartitionReqForBucket buckets_req = 2;
+}
+
+message FreezePartitionResponse {
+  repeated PbFreezePartitionRespForBucket buckets_resp = 1;
+}
+
 // adjust isr request and response
 message AdjustIsrRequest {
   required int32 serverId = 1;
@@ -1065,6 +1075,19 @@ message PbStopReplicaRespForBucket {
   optional string error_message = 3;
 }
 
+message PbFreezePartitionReqForBucket {
+  required PbTableBucket table_bucket = 1;
+  required int32 leader_epoch = 2;
+}
+
+message PbFreezePartitionRespForBucket {
+  required PbTableBucket table_bucket = 1;
+  optional int32 error_code = 2;
+  optional string error_message = 3;
+  optional int64 high_watermark = 4;
+  optional int64 log_end_offset = 5;
+}
+
 message PbKvSnapshot {
   required int32 bucket_id = 1;
   // null if there is no snapshot for this bucket

fluss-rpc/src/test/java/org/apache/fluss/rpc/TestingTabletGatewayService.java

@@ -25,6 +25,8 @@
 import org.apache.fluss.rpc.messages.DescribeClusterConfigsResponse;
 import org.apache.fluss.rpc.messages.FetchLogRequest;
 import org.apache.fluss.rpc.messages.FetchLogResponse;
+import org.apache.fluss.rpc.messages.FreezePartitionRequest;
+import org.apache.fluss.rpc.messages.FreezePartitionResponse;
 import org.apache.fluss.rpc.messages.GetClusterHealthRequest;
 import org.apache.fluss.rpc.messages.GetClusterHealthResponse;
 import org.apache.fluss.rpc.messages.GetDatabaseInfoRequest;
@@ -112,6 +114,12 @@ public CompletableFuture<StopReplicaResponse> stopReplica(
         return null;
     }
 
+    @Override
+    public CompletableFuture<FreezePartitionResponse> freezePartition(
+            FreezePartitionRequest request) {
+        return null;
+    }
+
     @Override
     public CompletableFuture<ProduceLogResponse> produceLog(ProduceLogRequest request) {
         return null;

fluss-server/src/main/java/org/apache/fluss/server/coordinator/AutoPartitionManager.java

@@ -93,6 +93,7 @@ public class AutoPartitionManager implements AutoCloseable {
     private final MetadataManager metadataManager;
     private final RemoteDirDynamicLoader remoteDirDynamicLoader;
     private final ReplicaCapacityController replicaCapacityController;
+    private final @Nullable LakeAwarePartitionRetentionManager lakeAwareRetentionManager;
     private final Clock clock;
 
     private final long periodicInterval;
@@ -113,6 +114,7 @@ public class AutoPartitionManager implements AutoCloseable {
 
     private final Lock lock = new ReentrantLock();
 
+    @VisibleForTesting
     public AutoPartitionManager(
             ServerMetadataCache metadataCache,
             MetadataManager metadataManager,
@@ -129,7 +131,29 @@ public AutoPartitionManager(
                 // TODO: Reuse the CoordinatorServer shared scheduler for this lightweight
                 // coordinator periodic task instead of creating a component-owned scheduler.
                 Executors.newScheduledThreadPool(
-                        1, new ExecutorThreadFactory("periodic-auto-partition-manager")));
+                        1, new ExecutorThreadFactory("periodic-auto-partition-manager")),
+                null);
+    }
+
+    public AutoPartitionManager(
+            ServerMetadataCache metadataCache,
+            MetadataManager metadataManager,
+            RemoteDirDynamicLoader remoteDirDynamicLoader,
+            Configuration conf,
+            ReplicaCapacityController replicaCapacityController,
+            LakeAwarePartitionRetentionManager lakeAwareRetentionManager) {
+        this(
+                metadataCache,
+                metadataManager,
+                remoteDirDynamicLoader,
+                conf,
+                replicaCapacityController,
+                SystemClock.getInstance(),
+                // TODO: Reuse the CoordinatorServer shared scheduler for this lightweight
+                // coordinator periodic task instead of creating a component-owned scheduler.
+                Executors.newScheduledThreadPool(
+                        1, new ExecutorThreadFactory("periodic-auto-partition-manager")),
+                checkNotNull(lakeAwareRetentionManager));
     }
 
     @VisibleForTesting
@@ -141,10 +165,32 @@ public AutoPartitionManager(
             ReplicaCapacityController replicaCapacityController,
             Clock clock,
             ScheduledExecutorService periodicExecutor) {
+        this(
+                metadataCache,
+                metadataManager,
+                remoteDirDynamicLoader,
+                conf,
+                replicaCapacityController,
+                clock,
+                periodicExecutor,
+                null);
+    }
+
+    @VisibleForTesting
+    AutoPartitionManager(
+            ServerMetadataCache metadataCache,
+            MetadataManager metadataManager,
+            RemoteDirDynamicLoader remoteDirDynamicLoader,
+            Configuration conf,
+            ReplicaCapacityController replicaCapacityController,
+            Clock clock,
+            ScheduledExecutorService periodicExecutor,
+            @Nullable LakeAwarePartitionRetentionManager lakeAwareRetentionManager) {
         this.metadataCache = metadataCache;
         this.metadataManager = metadataManager;
         this.remoteDirDynamicLoader = remoteDirDynamicLoader;
         this.replicaCapacityController = replicaCapacityController;
+        this.lakeAwareRetentionManager = lakeAwareRetentionManager;
         this.clock = clock;
         this.periodicExecutor = periodicExecutor;
         this.periodicInterval = conf.get(ConfigOptions.AUTO_PARTITION_CHECK_INTERVAL).toMillis();
@@ -388,8 +434,7 @@ private void doAutoPartition(Instant now, Set<Long> tableIds, boolean forceDoAut
             }
 
             dropPartitions(
-                    tablePath,
-                    tableInfo.getPartitionKeys(),
+                    tableInfo,
                     now,
                     tableInfo.getTableConfig().getAutoPartitionStrategy(),
                     currentPartitions);
@@ -502,11 +547,12 @@ private List<ResolvedPartitionSpec> partitionNamesToPreCreate(
     }
 
     private void dropPartitions(
-            TablePath tablePath,
-            List<String> partitionKeys,
+            TableInfo tableInfo,
             Instant currentInstant,
             AutoPartitionStrategy autoPartitionStrategy,
             NavigableMap<String, Set<String>> currentPartitions) {
+        boolean ensureTiered = autoPartitionStrategy.ensureLakeTieredBeforeDrop();
+
         int numToRetain = autoPartitionStrategy.numToRetain();
         // negative value means not to drop partitions
         if (numToRetain < 0) {
@@ -541,15 +587,17 @@ private void dropPartitions(
                 currentPartitions.headMap(lastRetainPartitionTime).entrySet().iterator();
         while (iterator.hasNext()) {
             Map.Entry<String, Set<String>> entry = iterator.next();
-            dropPartitions(tablePath, partitionKeys, iterator, entry);
+            dropPartitions(tableInfo, ensureTiered, iterator, entry);
         }
     }
 
     private void dropPartitions(
-            TablePath tablePath,
-            List<String> partitionKeys,
+            TableInfo tableInfo,
+            boolean ensureTiered,
             Iterator<Map.Entry<String, Set<String>>> iterator,
             Map.Entry<String, Set<String>> entry) {
+        TablePath tablePath = tableInfo.getTablePath();
+        List<String> partitionKeys = tableInfo.getPartitionKeys();
         Iterator<String> dropIterator;
         if (entry.getValue() == null) {
             dropIterator = new HashSet<>(Collections.singleton(entry.getKey())).iterator();
@@ -559,6 +607,11 @@ private void dropPartitions(
 
         while (dropIterator.hasNext()) {
             String partitionName = dropIterator.next();
+            if (ensureTiered) {
+                checkNotNull(lakeAwareRetentionManager).startRetention(tableInfo, partitionName);
+                continue;
+            }
+
             try {
                 metadataManager.dropPartition(
                         tablePath,
@@ -577,7 +630,9 @@ private void dropPartitions(
                     partitionName,
                     tablePath);
         }
-        iterator.remove();
+        if (!ensureTiered) {
+            iterator.remove();
+        }
     }
 
     @VisibleForTesting
@@ -590,6 +645,9 @@ protected Integer getAutoCreateDayDelayMinutes(long tableId) {
     public void close() throws Exception {
         if (isClosed.compareAndSet(false, true)) {
             periodicExecutor.shutdownNow();
+            if (lakeAwareRetentionManager != null) {
+                lakeAwareRetentionManager.close();
+            }
         }
     }
 }

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorChannelManager.java

@@ -20,9 +20,12 @@
 import org.apache.fluss.annotation.VisibleForTesting;
 import org.apache.fluss.cluster.ServerNode;
 import org.apache.fluss.cluster.ServerType;
+import org.apache.fluss.exception.TabletServerNotAvailableException;
 import org.apache.fluss.rpc.RpcClient;
 import org.apache.fluss.rpc.gateway.TabletServerGateway;
 import org.apache.fluss.rpc.messages.ApiMessage;
+import org.apache.fluss.rpc.messages.FreezePartitionRequest;
+import org.apache.fluss.rpc.messages.FreezePartitionResponse;
 import org.apache.fluss.rpc.messages.NotifyKvSnapshotOffsetRequest;
 import org.apache.fluss.rpc.messages.NotifyKvSnapshotOffsetResponse;
 import org.apache.fluss.rpc.messages.NotifyLakeTableOffsetRequest;
@@ -40,7 +43,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.annotation.concurrent.NotThreadSafe;
+import javax.annotation.concurrent.ThreadSafe;
 
 import java.util.Collection;
 import java.util.Optional;
@@ -53,7 +56,7 @@
  * Using by coordinator server. It's a manager to manage the rpc channels to tablet servers and send
  * request to the servers.
  */
-@NotThreadSafe
+@ThreadSafe
 public class CoordinatorChannelManager {
 
     private static final Logger LOG = LoggerFactory.getLogger(CoordinatorChannelManager.class);
@@ -121,6 +124,20 @@ public void sendStopBucketReplicaRequest(
                 responseConsumer);
     }
 
+    /** Send a request to freeze partition writes on a tablet server. */
+    public CompletableFuture<FreezePartitionResponse> sendFreezePartitionRequest(
+            int receiveServerId, FreezePartitionRequest request) {
+        Optional<TabletServerGateway> gateway = getTabletServerGateway(receiveServerId);
+        if (gateway.isPresent()) {
+            return gateway.get().freezePartition(request);
+        }
+        CompletableFuture<FreezePartitionResponse> response = new CompletableFuture<>();
+        response.completeExceptionally(
+                new TabletServerNotAvailableException(
+                        "Tablet server " + receiveServerId + " is not available."));
+        return response;
+    }
+
     /** Send UpdateMetadataRequest to the server and handle the response. */
     public void sendUpdateMetadataRequest(
             int receiveServerId,

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorServer.java

@@ -329,14 +329,21 @@ protected void initCoordinatorLeader() throws Exception {
 
             this.coordinatorChannelManager = new CoordinatorChannelManager(rpcClient);
 
+            LakeAwarePartitionRetentionManager lakeAwarePartitionRetentionManager =
+                    new LakeAwarePartitionRetentionManager(
+                            metadataManager,
+                            zkClient,
+                            coordinatorChannelManager,
+                            zkEpoch.getCoordinatorEpoch(),
+                            ioExecutor);
             this.autoPartitionManager =
                     new AutoPartitionManager(
                             metadataCache,
                             metadataManager,
                             remoteDirDynamicLoader,
                             conf,
-                            replicaCapacityController);
-            autoPartitionManager.start();
+                            replicaCapacityController,
+                            lakeAwarePartitionRetentionManager);
 
             // start coordinator event processor after we register coordinator leader to zk
             // so that the event processor can get the coordinator leader node from zk during
@@ -360,6 +367,7 @@ protected void initCoordinatorLeader() throws Exception {
                             scheduler,
                             clock);
             coordinatorEventProcessor.startup();
+            autoPartitionManager.start();
 
             // As the active leader, this server is the sole writer of dynamic configs and holds the
             // latest values, so stop consuming change notifications to avoid rolling a value back.