[server] Support primary-key writes to historical partitions

fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/server/NettyServerHandler.java

@@ -27,6 +27,7 @@
 import org.apache.fluss.rpc.messages.AuthenticateResponse;
 import org.apache.fluss.rpc.messages.FetchLogRequest;
 import org.apache.fluss.rpc.messages.LookupRequest;
+import org.apache.fluss.rpc.messages.PutKvRequest;
 import org.apache.fluss.rpc.protocol.ApiError;
 import org.apache.fluss.rpc.protocol.ApiKeys;
 import org.apache.fluss.rpc.protocol.ApiManager;
@@ -57,6 +58,7 @@
 import static org.apache.fluss.rpc.protocol.MessageCodec.encodeServerFailure;
 import static org.apache.fluss.rpc.protocol.MessageCodec.encodeSuccessResponse;
 import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.hasHistoricalLookup;
+import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.hasHistoricalPut;
 
 /** Implementation of the channel handler to process inbound requests for RPC server. */
 public final class NettyServerHandler extends ChannelInboundHandlerAdapter {
@@ -308,6 +310,8 @@ private Optional<RequestsMetrics.Metrics> getMetrics(FlussRequest request) {
             isFromFollower = fetchLogRequest.getFollowerServerId() >= 0;
         } else if (request.getApiKey() == ApiKeys.LOOKUP.id) {
             isHistorical = hasHistoricalLookup((LookupRequest) requestMessage);
+        } else if (request.getApiKey() == ApiKeys.PUT_KV.id) {
+            isHistorical = hasHistoricalPut((PutKvRequest) requestMessage);
         }
         return requestsMetrics.getMetrics(request.getApiKey(), isFromFollower, isHistorical);
     }

fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/server/RequestsMetrics.java

@@ -59,6 +59,9 @@ private RequestsMetrics(MetricGroup serverMetricsGroup, Collection<ApiKeys> apiK
             if (apiKey == ApiKeys.LOOKUP) {
                 addMetrics(serverMetricsGroup, toRequestName(apiKey, false, true));
             }
+            if (apiKey == ApiKeys.PUT_KV) {
+                addMetrics(serverMetricsGroup, toRequestName(apiKey, false, true));
+            }
         }
         this.requestMetricGroup = serverMetricsGroup.addGroup("request");
     }
@@ -103,7 +106,7 @@ private static String toRequestName(
             case PRODUCE_LOG:
                 return "produceLog";
             case PUT_KV:
-                return "putKv";
+                return isHistorical ? "historicalPutKv" : "putKv";
             case LOOKUP:
                 return isHistorical ? "historicalLookup" : "lookup";
             case PREFIX_LOOKUP:

fluss-rpc/src/main/java/org/apache/fluss/rpc/util/CommonRpcMessageUtils.java

@@ -34,6 +34,7 @@
 import org.apache.fluss.rpc.messages.PbPartitionSpec;
 import org.apache.fluss.rpc.messages.PbRemoteLogFetchInfo;
 import org.apache.fluss.rpc.messages.PbRemoteLogSegment;
+import org.apache.fluss.rpc.messages.PutKvRequest;
 import org.apache.fluss.rpc.protocol.ApiError;
 import org.apache.fluss.security.acl.AccessControlEntry;
 import org.apache.fluss.security.acl.AccessControlEntryFilter;
@@ -71,6 +72,17 @@ public static boolean hasHistoricalLookup(LookupRequest lookupRequest) {
                 && lookupRequest.getBucketsReqAt(0).hasOriginalPartitionName();
     }
 
+    /**
+     * Returns whether the put-KV request is for historical partition writes.
+     *
+     * <p>Normal and historical write buckets cannot be mixed in the same request, so the first
+     * bucket determines the request type.
+     */
+    public static boolean hasHistoricalPut(PutKvRequest putKvRequest) {
+        return putKvRequest.getBucketsReqsCount() > 0
+                && putKvRequest.getBucketsReqAt(0).hasOriginalPartitionName();
+    }
+
     public static List<PbAclInfo> toPbAclInfos(Collection<AclBinding> aclBindings) {
         return aclBindings.stream()
                 .map(CommonRpcMessageUtils::toPbAclInfo)

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

@@ -933,6 +933,8 @@ message PbPutKvReqForBucket {
   optional int64 partition_id = 1;
   required int32 bucket_id = 2;
   required bytes records = 3;
+  // The original partition name for historical PK writes. It is unset for normal writes.
+  optional string original_partition_name = 4;
 }
 
 message PbPutKvRespForBucket {

fluss-rust/crates/fluss/src/rpc/message/put_kv.rs

@@ -51,6 +51,7 @@ impl PutKvRequest {
                 partition_id: ready_batch.table_bucket.partition_id(),
                 bucket_id: ready_batch.table_bucket.bucket_id(),
                 records: ready_batch.write_batch.build()?,
+                original_partition_name: None,
             })
         }
 

fluss-server/src/main/java/org/apache/fluss/server/entity/PutKvDataForBucket.java

@@ -0,0 +1,59 @@
+/*
+ * 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.entity;
+
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.record.KvRecordBatch;
+
+import javax.annotation.Nullable;
+
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/** Put KV request data and historical partition context for one table bucket. */
+public final class PutKvDataForBucket {
+
+    private final TableBucket tableBucket;
+    private final KvRecordBatch records;
+    private final @Nullable String originalPartitionName;
+
+    /** Creates decoded put-KV data for one table bucket. */
+    public PutKvDataForBucket(
+            TableBucket tableBucket,
+            KvRecordBatch records,
+            @Nullable String originalPartitionName) {
+        this.tableBucket = checkNotNull(tableBucket, "tableBucket must not be null.");
+        this.records = checkNotNull(records, "records must not be null.");
+        this.originalPartitionName = originalPartitionName;
+    }
+
+    /** Returns the physical table bucket targeted by this request data. */
+    public TableBucket tableBucket() {
+        return tableBucket;
+    }
+
+    /** Returns the encoded KV records for this table bucket. */
+    public KvRecordBatch records() {
+        return records;
+    }
+
+    /** Returns the original partition name for a historical write, or null. */
+    public @Nullable String originalPartitionName() {
+        return originalPartitionName;
+    }
+}

fluss-server/src/main/java/org/apache/fluss/server/kv/KvManager.java

@@ -386,8 +386,6 @@ public void dropKv(TableBucket tableBucket) {
                                 dropKvTablet.getKvTabletDir().getAbsolutePath()),
                         e);
             }
-        } else {
-            LOG.warn("Fail to delete kv bucket {}.", tableBucket.getBucket());
         }
     }
 

fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateAccessor.java

@@ -0,0 +1,53 @@
+/*
+ * 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.kv;
+
+import org.apache.fluss.annotation.Internal;
+import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.Key;
+import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.TruncateReason;
+
+import javax.annotation.Nullable;
+
+/** Accessor for the local state used while processing KV records. */
+@Internal
+public interface KvStateAccessor {
+
+    /**
+     * Encodes the logical primary key into the physical key used by this state.
+     *
+     * <p>Normal KV state keeps the primary key unchanged. Historical KV state also encodes the
+     * original partition context because multiple original partitions share one historical KV
+     * tablet.
+     */
+    Key encodeKey(byte[] primaryKey);
+
+    /** Looks up an encoded key from local state and an optional external fallback. */
+    KvStateLookupResult lookup(Key key) throws Exception;
+
+    /** Adds an insert mutation to the prewrite buffer. */
+    void insert(Key key, byte[] value, long logOffset);
+
+    /** Adds an update mutation to the prewrite buffer. */
+    void update(Key key, @Nullable byte[] value, long logOffset);
+
+    /** Adds a delete mutation to the prewrite buffer. */
+    void delete(Key key, long logOffset);
+
+    /** Truncates pending mutations to the given log offset. */
+    void truncateTo(long logOffset, TruncateReason reason);
+}

fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateLookupResult.java

@@ -0,0 +1,131 @@
+/*
+ * 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.kv;
+
+import org.apache.fluss.annotation.Internal;
+
+import javax.annotation.Nullable;
+
+import java.util.Arrays;
+
+import static org.apache.fluss.utils.Preconditions.checkArgument;
+import static org.apache.fluss.utils.Preconditions.checkNotNull;
+
+/**
+ * Result of looking up a key from local KV state.
+ *
+ * <p>{@link Status#NOT_FOUND} only means that the key was not found in local KV state, so a
+ * historical lookup may still query lake storage. {@link Status#DELETED} confirms that the key has
+ * been deleted and must stop the fallback from exposing an older value in lake storage.
+ */
+@Internal
+public final class KvStateLookupResult {
+
+    /** Status of a local KV state lookup. */
+    public enum Status {
+        /** The key was not found in local KV state. */
+        NOT_FOUND,
+
+        /** Local state contains a non-empty encoded value. */
+        PRESENT,
+
+        /** The key is known to have been deleted. */
+        DELETED
+    }
+
+    private static final KvStateLookupResult NOT_FOUND =
+            new KvStateLookupResult(Status.NOT_FOUND, null);
+    private static final KvStateLookupResult DELETED =
+            new KvStateLookupResult(Status.DELETED, null);
+
+    private final Status status;
+    private final @Nullable byte[] value;
+
+    private KvStateLookupResult(Status status, @Nullable byte[] value) {
+        this.status = status;
+        this.value = value;
+    }
+
+    /** Returns a result indicating that the key was not found in local KV state. */
+    public static KvStateLookupResult notFound() {
+        return NOT_FOUND;
+    }
+
+    /** Returns a result containing a non-empty encoded value. */
+    public static KvStateLookupResult present(byte[] value) {
+        checkNotNull(value, "value must not be null");
+        checkArgument(value.length > 0, "value must not be empty");
+        return new KvStateLookupResult(Status.PRESENT, value);
+    }
+
+    /** Returns a result indicating that the key is known to have been deleted. */
+    public static KvStateLookupResult deleted() {
+        return DELETED;
+    }
+
+    /** Returns the lookup status. */
+    public Status status() {
+        return status;
+    }
+
+    /** Returns whether this result contains an encoded value. */
+    public boolean isPresent() {
+        return status == Status.PRESENT;
+    }
+
+    /** Returns whether the key is known to have been deleted. */
+    public boolean isDeleted() {
+        return status == Status.DELETED;
+    }
+
+    /**
+     * Returns the encoded value, or null when the key is absent or deleted.
+     *
+     * <p>The returned byte array is owned by the underlying state and must not be modified.
+     */
+    public @Nullable byte[] value() {
+        return value;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        KvStateLookupResult that = (KvStateLookupResult) o;
+        return status == that.status && Arrays.equals(value, that.value);
+    }
+
+    @Override
+    public int hashCode() {
+        return 31 * status.hashCode() + Arrays.hashCode(value);
+    }
+
+    @Override
+    public String toString() {
+        return "KvStateLookupResult{"
+                + "status="
+                + status
+                + ", value="
+                + Arrays.toString(value)
+                + '}';
+    }
+}