[kv] PrimaryKey table support standby replica to reduce recovery time

Author: swuferhong

fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java

@@ -73,6 +73,7 @@ public class MetricNames {
     public static final String SERVER_LOGICAL_STORAGE_KV_SIZE = "kvSize";
     public static final String SERVER_PHYSICAL_STORAGE_LOCAL_SIZE = "localSize";
     public static final String SERVER_PHYSICAL_STORAGE_REMOTE_LOG_SIZE = "remoteLogSize";
+    public static final String SERVER_PHYSICAL_STORAGE_STANDBY_SIZE = "standbySize";
 
     // --------------------------------------------------------------------------------------------
     // metrics for table

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

@@ -332,6 +332,7 @@ message NotifyKvSnapshotOffsetRequest {
   required int32 bucket_id = 3;
   required int32 coordinator_epoch = 4;
   required int64 min_retain_offset = 5;
+  optional int64 snapshot_id = 6;
 }
 
 message NotifyKvSnapshotOffsetResponse {
@@ -745,6 +746,7 @@ message PbAdjustIsrReqForBucket {
   repeated int32 new_isr = 4 [packed = true];
   required int32 coordinator_epoch = 5;
   required int32 bucket_epoch = 6;
+  repeated int32 standby_replicas = 7 [packed = true];
 }
 
 message PbAdjustIsrRespForTable {
@@ -762,6 +764,7 @@ message PbAdjustIsrRespForBucket {
   repeated int32 isr = 7 [packed = true];
   optional int32 bucket_epoch = 8;
   optional int32 coordinator_epoch = 9;
+  repeated int32 standby_replicas = 10 [packed = true];
 }
 
 message PbListOffsetsRespForBucket {
@@ -782,6 +785,7 @@ message PbNotifyLeaderAndIsrReqForBucket {
   repeated int32 replicas = 5 [packed = true];
   repeated int32 isr = 6 [packed = true];
   required int32 bucket_epoch = 7;
+  repeated int32 standby_replicas = 8 [packed = true];
 }
 
 message PbNotifyLeaderAndIsrRespForBucket {

fluss-server/src/main/java/org/apache/fluss/server/TabletManagerBase.java

@@ -173,7 +173,7 @@ protected int runInThreadPool(Runnable[] runnableJobs, String poolName) throws T
      * @param tableBucket the table bucket
      * @return the tablet directory
      */
-    protected File getOrCreateTabletDir(PhysicalTablePath tablePath, TableBucket tableBucket) {
+    public File getOrCreateTabletDir(PhysicalTablePath tablePath, TableBucket tableBucket) {
         File tabletDir = getTabletDir(tablePath, tableBucket);
         if (tabletDir.exists()) {
             return tabletDir;

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

@@ -959,6 +959,7 @@ private List<AdjustIsrResultForBucket> tryProcessAdjustIsr(
                             // TODO: reject the request if there is a replica in ISR is not online,
                             //  see KIP-841.
                             tryAdjustLeaderAndIsr.isr(),
+                            tryAdjustLeaderAndIsr.standbyReplicas(),
                             coordinatorContext.getCoordinatorEpoch(),
                             currentLeaderAndIsr.bucketEpoch() + 1);
             newLeaderAndIsrList.put(tableBucket, newLeaderAndIsr);
@@ -1071,7 +1072,9 @@ private void tryProcessCommitKvSnapshot(
                         completedSnapshotStore.add(completedSnapshot);
                         coordinatorEventManager.put(
                                 new NotifyKvSnapshotOffsetEvent(
-                                        tb, completedSnapshot.getLogOffset()));
+                                        tb,
+                                        completedSnapshot.getLogOffset(),
+                                        completedSnapshot.getSnapshotID()));
                         callback.complete(new CommitKvSnapshotResponse());
                     } catch (Exception e) {
                         callback.completeExceptionally(e);
@@ -1081,7 +1084,6 @@ private void tryProcessCommitKvSnapshot(
 
     private void processNotifyKvSnapshotOffsetEvent(NotifyKvSnapshotOffsetEvent event) {
         TableBucket tb = event.getTableBucket();
-        long logOffset = event.getLogOffset();
         coordinatorRequestBatch.newBatch();
         coordinatorContext
                 .getBucketLeaderAndIsr(tb)
@@ -1092,7 +1094,8 @@ private void processNotifyKvSnapshotOffsetEvent(NotifyKvSnapshotOffsetEvent even
                                                 coordinatorContext.getFollowers(
                                                         tb, leaderAndIsr.leader()),
                                                 tb,
-                                                logOffset));
+                                                event.getLogOffset(),
+                                                event.getSnapshotId()));
         coordinatorRequestBatch.sendNotifyKvSnapshotOffsetRequest(
                 coordinatorContext.getCoordinatorEpoch());
     }

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

@@ -361,15 +361,18 @@ public void addNotifyRemoteLogOffsetsRequestForTabletServers(
     }
 
     public void addNotifyKvSnapshotOffsetRequestForTabletServers(
-            List<Integer> tabletServers, TableBucket tableBucket, long minRetainOffset) {
+            List<Integer> tabletServers,
+            TableBucket tableBucket,
+            long minRetainOffset,
+            long snapshotId) {
         tabletServers.stream()
                 .filter(s -> s >= 0)
                 .forEach(
                         id ->
                                 notifyKvSnapshotOffsetRequestMap.put(
                                         id,
                                         makeNotifyKvSnapshotOffsetRequest(
-                                                tableBucket, minRetainOffset)));
+                                                tableBucket, minRetainOffset, snapshotId)));
     }
 
     public void addNotifyLakeTableOffsetRequestForTableServers(

fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyKvSnapshotOffsetEvent.java

@@ -25,10 +25,12 @@ public class NotifyKvSnapshotOffsetEvent implements CoordinatorEvent {
 
     private final TableBucket tableBucket;
     private final long logOffset;
+    private final long snapshotId;
 
-    public NotifyKvSnapshotOffsetEvent(TableBucket tableBucket, long logOffset) {
+    public NotifyKvSnapshotOffsetEvent(TableBucket tableBucket, long logOffset, long snapshotId) {
         this.tableBucket = tableBucket;
         this.logOffset = logOffset;
+        this.snapshotId = snapshotId;
     }
 
     public TableBucket getTableBucket() {
@@ -38,4 +40,8 @@ public TableBucket getTableBucket() {
     public long getLogOffset() {
         return logOffset;
     }
+
+    public long getSnapshotId() {
+        return snapshotId;
+    }
 }

fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/ReplicaLeaderElectionAlgorithms.java

@@ -21,6 +21,7 @@
 import org.apache.fluss.server.zk.data.LeaderAndIsr;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
@@ -36,23 +37,39 @@ public class ReplicaLeaderElectionAlgorithms {
      * @param assignments the assignments
      * @param aliveReplicas the alive replicas
      * @param coordinatorEpoch the coordinator epoch
+     * @param isPrimaryKeyTable whether this table bucket is primary key table
      * @return the election result
      */
     public static Optional<ElectionResult> initReplicaLeaderElection(
-            List<Integer> assignments, List<Integer> aliveReplicas, int coordinatorEpoch) {
-        // currently, we always use the first replica in assignment, which also in aliveReplicas and
-        // isr as the leader replica.
-        for (int assignment : assignments) {
-            if (aliveReplicas.contains(assignment)) {
-                return Optional.of(
-                        new ElectionResult(
-                                aliveReplicas,
-                                new LeaderAndIsr(
-                                        assignment, 0, aliveReplicas, coordinatorEpoch, 0)));
-            }
+            List<Integer> assignments,
+            List<Integer> aliveReplicas,
+            int coordinatorEpoch,
+            boolean isPrimaryKeyTable) {
+        // First we will filter out the assignment list to only contain the alive replicas.
+        List<Integer> availableReplicas =
+                assignments.stream().filter(aliveReplicas::contains).collect(Collectors.toList());
+
+        // If the assignment list is empty, we return empty.
+        if (availableReplicas.isEmpty()) {
+            return Optional.empty();
         }
 
-        return Optional.empty();
+        //  Then we will use the first replica in assignment as the leader replica.
+        int leader = availableReplicas.get(0);
+
+        // If this table is primaryKey table, we will use the second replica in assignment as the
+        // standby if exists.
+        List<Integer> standbyReplica = new ArrayList<>();
+        if (isPrimaryKeyTable) {
+            if (availableReplicas.size() > 1) {
+                standbyReplica.add(availableReplicas.get(1));
+            }
+        }
+        return Optional.of(
+                new ElectionResult(
+                        aliveReplicas,
+                        new LeaderAndIsr(
+                                leader, 0, aliveReplicas, standbyReplica, coordinatorEpoch, 0)));
     }
 
     /**
@@ -61,22 +78,77 @@ public static Optional<ElectionResult> initReplicaLeaderElection(
      * @param assignments the assignments
      * @param aliveReplicas the alive replicas
      * @param leaderAndIsr the original leaderAndIsr
+     * @param isPrimaryKeyTable whether this table bucket is primary key table
      * @return the election result
      */
     public static Optional<ElectionResult> defaultReplicaLeaderElection(
-            List<Integer> assignments, List<Integer> aliveReplicas, LeaderAndIsr leaderAndIsr) {
-        // currently, we always use the first replica in assignment, which also in aliveReplicas and
-        // isr as the leader replica.
+            List<Integer> assignments,
+            List<Integer> aliveReplicas,
+            LeaderAndIsr leaderAndIsr,
+            boolean isPrimaryKeyTable) {
         List<Integer> isr = leaderAndIsr.isr();
-        for (int assignment : assignments) {
-            if (aliveReplicas.contains(assignment) && isr.contains(assignment)) {
-                return Optional.of(
-                        new ElectionResult(
-                                aliveReplicas, leaderAndIsr.newLeaderAndIsr(assignment, isr)));
+        // First we will filter out the assignment list to only contain the alive replicas and isr.
+        List<Integer> availableReplicas =
+                assignments.stream()
+                        .filter(replica -> aliveReplicas.contains(replica) && isr.contains(replica))
+                        .collect(Collectors.toList());
+
+        // If the assignment list is empty, we return empty.
+        if (availableReplicas.isEmpty()) {
+            return Optional.empty();
+        }
+
+        // For log table, we will use the first replica in assignment as the leader replica.
+        if (!isPrimaryKeyTable) {
+            int leader = availableReplicas.get(0);
+            return Optional.of(
+                    new ElectionResult(
+                            aliveReplicas,
+                            leaderAndIsr.newLeaderAndIsr(leader, isr, Collections.emptyList())));
+        }
+
+        int currentStandby =
+                leaderAndIsr.standbyReplicas().isEmpty()
+                        ? -1
+                        : leaderAndIsr.standbyReplicas().get(0);
+        int newLeader;
+        int newStandby = -1;
+        if (currentStandby != -1 && availableReplicas.contains(currentStandby)) {
+            // Standby exists and is available: promote it to leader.
+            newLeader = currentStandby;
+            // Find a new standby for the candidate replica (not leader).
+            List<Integer> candidatesForStandby =
+                    availableReplicas.stream()
+                            .filter(replica -> replica != newLeader)
+                            .collect(Collectors.toList());
+            if (!candidatesForStandby.isEmpty()) {
+                newStandby = candidatesForStandby.get(0);
+            }
+        } else {
+            // Standby does not exist or is not available: promote the first replica in assignment
+            // to leader.
+            newLeader = availableReplicas.get(0);
+            // Find a new standby for the candidate replica (not leader).
+            List<Integer> candidatesForStandby =
+                    availableReplicas.stream()
+                            .filter(replica -> replica != newLeader)
+                            .collect(Collectors.toList());
+            if (!candidatesForStandby.isEmpty()) {
+                newStandby = candidatesForStandby.get(0);
             }
         }
 
-        return Optional.empty();
+        if (newStandby == -1) {
+            newStandby =
+                    availableReplicas.stream()
+                            .filter(replica -> replica != newLeader)
+                            .collect(Collectors.toList())
+                            .get(0);
+        }
+
+        LeaderAndIsr newLeaderAndIsr =
+                leaderAndIsr.newLeaderAndIsr(newLeader, isr, Collections.singletonList(newStandby));
+        return Optional.of(new ElectionResult(aliveReplicas, newLeaderAndIsr));
     }
 
     /**
@@ -92,25 +164,79 @@ public static Optional<ElectionResult> controlledShutdownReplicaLeaderElection(
             List<Integer> assignments,
             List<Integer> aliveReplicas,
             LeaderAndIsr leaderAndIsr,
-            Set<Integer> shutdownTabletServers) {
+            Set<Integer> shutdownTabletServers,
+            boolean isPrimaryKeyTable) {
         List<Integer> originIsr = leaderAndIsr.isr();
         Set<Integer> isrSet = new HashSet<>(originIsr);
-        for (Integer id : assignments) {
-            if (aliveReplicas.contains(id)
-                    && isrSet.contains(id)
-                    && !shutdownTabletServers.contains(id)) {
-                Set<Integer> newAliveReplicas = new HashSet<>(aliveReplicas);
-                newAliveReplicas.removeAll(shutdownTabletServers);
-                List<Integer> newIsr =
-                        originIsr.stream()
-                                .filter(replica -> !shutdownTabletServers.contains(replica))
-                                .collect(Collectors.toList());
-                return Optional.of(
-                        new ElectionResult(
-                                new ArrayList<>(newAliveReplicas),
-                                leaderAndIsr.newLeaderAndIsr(id, newIsr)));
+        // First we will filter out the assignment list to only contain the alive replicas, isr, and
+        // not is shutdownTabletServers set.
+        List<Integer> availableReplicas =
+                assignments.stream()
+                        .filter(
+                                replica ->
+                                        aliveReplicas.contains(replica)
+                                                && isrSet.contains(replica)
+                                                && !shutdownTabletServers.contains(replica))
+                        .collect(Collectors.toList());
+        // If the assignment list is empty, we return empty.
+        if (availableReplicas.isEmpty()) {
+            return Optional.empty();
+        }
+
+        int currentStandby =
+                leaderAndIsr.standbyReplicas().isEmpty()
+                        ? -1
+                        : leaderAndIsr.standbyReplicas().get(0);
+        int newLeader;
+        int newStandby = -1;
+        if (!isPrimaryKeyTable) {
+            // For log table, we will use the first replica in availableReplicas as the leader
+            // replica.
+            newLeader = availableReplicas.get(0);
+        } else if (currentStandby != -1 && availableReplicas.contains(currentStandby)) {
+            // For pk table, Standby exists and is available: promote it to leader.
+            newLeader = currentStandby;
+            // Find a new standby for the candidate replica (not leader).
+            List<Integer> candidatesForStandby =
+                    availableReplicas.stream()
+                            .filter(replica -> replica != newLeader)
+                            .collect(Collectors.toList());
+            if (!candidatesForStandby.isEmpty()) {
+                newStandby = candidatesForStandby.get(0);
+            }
+        } else {
+            // Standby does not exist or is not available: promote the first replica in assignment
+            // to leader.
+            newLeader = availableReplicas.get(0);
+            // Find a new standby for the candidate replica (not leader).
+            List<Integer> candidatesForStandby =
+                    availableReplicas.stream()
+                            .filter(replica -> replica != newLeader)
+                            .collect(Collectors.toList());
+            if (!candidatesForStandby.isEmpty()) {
+                newStandby = candidatesForStandby.get(0);
             }
         }
-        return Optional.empty();
+
+        Set<Integer> newAliveReplicas = new HashSet<>(aliveReplicas);
+        newAliveReplicas.removeAll(shutdownTabletServers);
+        List<Integer> newIsr =
+                originIsr.stream()
+                        .filter(replica -> !shutdownTabletServers.contains(replica))
+                        .collect(Collectors.toList());
+
+        if (newStandby == -1) {
+            newStandby =
+                    availableReplicas.stream()
+                            .filter(replica -> replica != newLeader)
+                            .collect(Collectors.toList())
+                            .get(0);
+        }
+
+        return Optional.of(
+                new ElectionResult(
+                        new ArrayList<>(newAliveReplicas),
+                        leaderAndIsr.newLeaderAndIsr(
+                                newLeader, newIsr, Collections.singletonList(newStandby))));
     }
 }

fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/ReplicaStateMachine.java

@@ -453,7 +453,8 @@ private Map<TableBucketReplica, LeaderAndIsr> doRemoveReplicaFromIsr(
                                     .collect(Collectors.toList());
             LeaderAndIsr adjustLeaderAndIsr =
                     newLeader == LeaderAndIsr.NO_LEADER
-                            ? leaderAndIsr.newLeaderAndIsr(newLeader, newIsr)
+                            ? leaderAndIsr.newLeaderAndIsr(
+                                    newLeader, newIsr, leaderAndIsr.standbyReplicas())
                             : leaderAndIsr.newLeaderAndIsr(newIsr);
             adjustedLeaderAndIsr.put(tableBucketReplica, adjustLeaderAndIsr);
             toUpdateLeaderAndIsrList.put(tableBucket, adjustLeaderAndIsr);