From 862d62590e186a2f47c0fa6aec24138f0c09dfc4 Mon Sep 17 00:00:00 2001 From: dimgigov Date: Fri, 31 Jul 2026 02:53:30 +0300 Subject: [PATCH] feat(raft): leader InstallSnapshot send on unrecoverable lag --- src/barabadb/core/raft.nim | 123 +++++++++++++++++++- src/baradadb.nim | 13 +++ tests/test_all.nim | 226 +++++++++++++++++++++++++++++++++++++ 3 files changed, 360 insertions(+), 2 deletions(-) diff --git a/src/barabadb/core/raft.nim b/src/barabadb/core/raft.nim index 71e1056..b96674d 100644 --- a/src/barabadb/core/raft.nim +++ b/src/barabadb/core/raft.nim @@ -86,6 +86,16 @@ type baseTerm: uint64): bool {.gcsafe.} snapIncomingId*: uint64 snapIncomingFile*: string + ## Leader InstallSnapshot send. buildSnapshot archives the current data + ## dir into destPath (wired in baradadb.nim via backupDataDir). + ## snapRejectStreak counts consecutive floor-level AppendEntries rejects + ## per peer; at 2 the peer is queued in snapPending and the network layer + ## (processMessage) kicks off sendSnapshot. snapSending is the + ## single-flight guard: at most one snapshot transfer per peer. + buildSnapshot*: proc(destPath: string): bool {.gcsafe.} + snapRejectStreak*: Table[string, int] + snapPending*: HashSet[string] + snapSending*: HashSet[string] RaftMessageKind* = enum rmkRequestVote @@ -220,6 +230,9 @@ proc newRaftNode*(id: string, peers: seq[string], raftPort: int = 0, snapChunkBytes: 262144, snapIncomingId: 0, snapIncomingFile: "", + snapRejectStreak: initTable[string, int](), + snapPending: initHashSet[string](), + snapSending: initHashSet[string](), ) result.loadState() @@ -333,6 +346,9 @@ proc becomeFollower*(node: RaftNode, term: uint64) = node.votesReceived.clear() node.nextIndex.clear() node.matchIndex.clear() + # Leader-only snapshot-send state is meaningless once we step down + node.snapRejectStreak.clear() + node.snapPending.clear() node.saveState() proc becomeCandidate*(node: RaftNode) = @@ -354,6 +370,8 @@ proc becomeLeader*(node: RaftNode) = for peer in node.peers: node.nextIndex[peer] = node.lastLogIndex + 1 node.matchIndex[peer] = 0 + node.snapRejectStreak.clear() + node.snapPending.clear() proc handleRequestVote*(node: RaftNode, msg: RaftMessage): RaftMessage = var reply = RaftMessage( @@ -600,6 +618,8 @@ proc handleAppendReply*(node: RaftNode, peerId: string, reply: RaftMessage) = if reply.success: node.matchIndex[peerId] = reply.matchIdx node.nextIndex[peerId] = reply.matchIdx + 1 + node.snapRejectStreak.del(peerId) + node.snapPending.excl(peerId) # Update commit index using true majority calculation let majority = (node.peers.len + 1 + 1) div 2 # majority of cluster (peers + leader) @@ -629,8 +649,41 @@ proc handleAppendReply*(node: RaftNode, peerId: string, reply: RaftMessage) = let floor = node.lastSnapshotIndex + 1 if node.nextIndex.getOrDefault(peerId, 1) > floor: dec node.nextIndex[peerId] + # Not a floor-level reject, so it breaks any floor-reject streak. + node.snapRejectStreak.del(peerId) else: node.nextIndex[peerId] = floor + # Stuck at the compaction floor: the entries the follower needs have + # been compacted away, so AppendEntries can never catch it up. Count + # consecutive floor rejects; at 2, queue an InstallSnapshot transfer + # (the network layer picks this up after handleAppendReply returns). + node.snapRejectStreak[peerId] = + node.snapRejectStreak.getOrDefault(peerId, 0) + 1 + if node.snapRejectStreak[peerId] >= 2: + node.snapPending.incl(peerId) + +proc handleInstallSnapshotReply*(node: RaftNode, peerId: string, + reply: RaftMessage) = + ## Leader side: follower's answer to a completed InstallSnapshot transfer. + ## success=true adopts the snapshot base (reply.matchIdx) as the peer's + ## match point; success=false leaves all state alone — the normal + ## AppendEntries reject path re-triggers another snapshot if the peer is + ## still stuck at the floor. + if reply.term > node.currentTerm: + node.becomeFollower(reply.term) + return + + if reply.term < node.currentTerm: + return + + if node.state != rsLeader: + return + + if reply.success: + node.matchIndex[peerId] = reply.matchIdx + node.nextIndex[peerId] = reply.matchIdx + 1 + node.snapRejectStreak.del(peerId) + node.snapPending.excl(peerId) proc state*(node: RaftNode): RaftState = node.state proc isLeader*(node: RaftNode): bool = node.state == rsLeader @@ -915,6 +968,69 @@ proc broadcast*(net: RaftNetwork, msgs: seq[RaftMessage]) {.async.} = if i < msgs.len: await net.send(peer, msgs[i]) +proc sendSnapshot*(net: RaftNetwork, peerId: string) {.async.} = + ## Leader side of InstallSnapshot: build an archive of the current data dir + ## via the buildSnapshot callback and stream it to a lagging peer in + ## snapChunkBytes chunks. Triggered (via asyncCheck from processMessage) + ## when handleAppendReply queues the peer in snapPending after consecutive + ## floor-level rejects. Single-flight per peer via node.snapSending. + ## + ## Runs on the raft event loop; buildSnapshot performs blocking disk I/O + ## (tar+gzip). Snapshot sends are rare, so we accept the stall rather than + ## adding a worker round-trip (same trade-off as restoreSnapshot). + let node = net.node + if peerId in node.snapSending: + return + if node.state != rsLeader or node.buildSnapshot == nil or + node.dataDir.len == 0: + return + let snapId = node.lastSnapshotIndex + if snapId == 0: + # snapId 0 can never be accepted (a follower's initial snapIncomingId is + # 0), and sends only trigger after compaction anyway — guard regardless. + warn("sendSnapshot: lastSnapshotIndex is 0; skipping snapshot send to " & peerId) + return + node.snapSending.incl(peerId) + defer: node.snapSending.excl(peerId) + + let baseIndex = node.lastSnapshotIndex + let baseTerm = node.lastSnapshotTerm + let destPath = node.dataDir / ("snap_out_" & $snapId & ".tar.gz") + defer: + if fileExists(destPath): + removeFile(destPath) + + if not node.buildSnapshot(destPath): + warn("sendSnapshot: buildSnapshot failed; aborting snapshot send to " & peerId) + return + + var f: File + if not open(f, destPath, fmRead): + warn("sendSnapshot: cannot open built archive " & destPath) + return + defer: f.close() + + let total = uint64(getFileSize(destPath)) + var offset = 0'u64 + while true: + var chunk = newSeq[byte](node.snapChunkBytes) + let n = f.readBytes(chunk, 0, chunk.len) + let done = offset + uint64(n) >= total + await net.send(peerId, RaftMessage( + kind: rmkInstallSnapshot, + term: node.currentTerm, + senderId: node.id, + prevLogIndex: baseIndex, # snapshot base index/term (T7 wire layout) + prevLogTerm: baseTerm, + snapId: snapId, + snapOffset: offset, + snapData: chunk[0 ..< n], + snapDone: done, + )) + if done: + break + offset += uint64(n) + proc processMessage*(net: RaftNetwork, msg: RaftMessage) {.async.} = case msg.kind of rmkRequestVote: @@ -932,6 +1048,10 @@ proc processMessage*(net: RaftNetwork, msg: RaftMessage) {.async.} = await net.send(msg.senderId, reply) of rmkAppendEntriesReply: net.node.handleAppendReply(msg.senderId, msg) + # Floor-reject streak reached the threshold: this peer needs a snapshot. + if msg.senderId in net.node.snapPending: + net.node.snapPending.excl(msg.senderId) + asyncCheck net.sendSnapshot(msg.senderId) of rmkInstallSnapshot: # Same election-timer rule as AppendEntries: only a plausible current # leader resets it. @@ -940,8 +1060,7 @@ proc processMessage*(net: RaftNetwork, msg: RaftMessage) {.async.} = let reply = net.node.handleInstallSnapshot(msg) await net.send(msg.senderId, reply) of rmkInstallSnapshotReply: - # Leader side of snapshot transfer lands in a follow-up task. - discard + net.node.handleInstallSnapshotReply(msg.senderId, msg) proc recvExact*(client: AsyncSocket, size: int): Future[string] {.async.} = ## Reads exactly `size` bytes from `client`. A short return means the peer diff --git a/src/baradadb.nim b/src/baradadb.nim index 5d954e3..5bdb894 100644 --- a/src/baradadb.nim +++ b/src/baradadb.nim @@ -413,6 +413,19 @@ proc main() = echo "[raft] Snapshot restore failed: ", e.msg result = false + # Leader InstallSnapshot send: archive the default DB's data directory + # into the path raft picks (dataDir/raft/snap_out_.tar.gz). Like + # restoreSnapshot this runs on the raft event loop and blocks on disk I/O + # (tar+gzip); snapshot sends are rare, so we accept the stall. + raftNode.buildSnapshot = proc(destPath: string): bool {.gcsafe.} = + echo "[raft] Building snapshot archive ", destPath + {.cast(gcsafe).}: + try: + result = backupDataDir(defaultDbDir, destPath) + except CatchableError as e: + echo "[raft] Snapshot build failed: ", e.msg + result = false + # Wire RAFT ↔ DistTxn wireRaftDistTxn(raftNode, tcpServer) diff --git a/tests/test_all.nim b/tests/test_all.nim index 987a7cd..29f04a9 100644 --- a/tests/test_all.nim +++ b/tests/test_all.nim @@ -1,6 +1,7 @@ ## BaraDB — Test Suite import std/unittest import std/tables +import std/sets import std/strutils import std/os import std/asyncdispatch @@ -2839,6 +2840,231 @@ suite "Raft InstallSnapshot Receive": check not bad.success check node.snapIncomingId == 0 +suite "Raft InstallSnapshot Send": + test "two consecutive floor rejects queue a snapshot send": + var node = newRaftNode("leader", @["peer-1"]) + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.lastSnapshotTerm = 4 + node.nextIndex["peer-1"] = 101 + node.matchIndex["peer-1"] = 0 + + let reject = RaftMessage(kind: rmkAppendEntriesReply, term: 5, + senderId: "peer-1", success: false) + node.handleAppendReply("peer-1", reject) + check node.snapRejectStreak["peer-1"] == 1 + check "peer-1" notin node.snapPending + check node.nextIndex["peer-1"] == 101 # pinned at the compaction floor + + node.handleAppendReply("peer-1", reject) + check node.snapRejectStreak["peer-1"] == 2 + check "peer-1" in node.snapPending + check node.nextIndex["peer-1"] == 101 + + test "non-floor reject decrements nextIndex without touching the streak": + var node = newRaftNode("leader", @["peer-1"]) + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.lastSnapshotTerm = 4 + node.nextIndex["peer-1"] = 105 + + node.handleAppendReply("peer-1", RaftMessage( + kind: rmkAppendEntriesReply, term: 5, senderId: "peer-1", success: false)) + check node.nextIndex["peer-1"] == 104 + check "peer-1" notin node.snapRejectStreak + check "peer-1" notin node.snapPending + + test "successful AppendEntries reply resets the streak and cancels a pending snapshot": + var node = newRaftNode("leader", @["peer-1"]) + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.lastSnapshotTerm = 4 + node.nextIndex["peer-1"] = 101 + node.matchIndex["peer-1"] = 0 + node.snapRejectStreak["peer-1"] = 1 + node.snapPending.incl("peer-1") + + node.handleAppendReply("peer-1", RaftMessage( + kind: rmkAppendEntriesReply, term: 5, senderId: "peer-1", + success: true, matchIdx: 101)) + check "peer-1" notin node.snapRejectStreak + check "peer-1" notin node.snapPending + check node.matchIndex["peer-1"] == 101 + check node.nextIndex["peer-1"] == 102 + + test "InstallSnapshotReply success advances match/next index and clears streak": + var node = newRaftNode("leader", @["peer-1"]) + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.lastSnapshotTerm = 4 + node.nextIndex["peer-1"] = 101 + node.matchIndex["peer-1"] = 0 + node.snapRejectStreak["peer-1"] = 2 + node.snapPending.incl("peer-1") + + node.handleInstallSnapshotReply("peer-1", RaftMessage( + kind: rmkInstallSnapshotReply, term: 5, senderId: "peer-1", + success: true, matchIdx: 100)) + check node.matchIndex["peer-1"] == 100 + check node.nextIndex["peer-1"] == 101 + check "peer-1" notin node.snapRejectStreak + check "peer-1" notin node.snapPending + + test "InstallSnapshotReply failure leaves leader state untouched": + var node = newRaftNode("leader", @["peer-1"]) + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.lastSnapshotTerm = 4 + node.nextIndex["peer-1"] = 101 + node.matchIndex["peer-1"] = 0 + node.snapRejectStreak["peer-1"] = 2 + + node.handleInstallSnapshotReply("peer-1", RaftMessage( + kind: rmkInstallSnapshotReply, term: 5, senderId: "peer-1", + success: false, matchIdx: 0)) + check node.matchIndex["peer-1"] == 0 + check node.nextIndex["peer-1"] == 101 + check node.snapRejectStreak["peer-1"] == 2 + + test "InstallSnapshotReply term handling matches AppendEntriesReply": + var node = newRaftNode("leader", @["peer-1"]) + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.nextIndex["peer-1"] = 101 + node.matchIndex["peer-1"] = 0 + + # Stale term: ignored entirely + node.handleInstallSnapshotReply("peer-1", RaftMessage( + kind: rmkInstallSnapshotReply, term: 4, senderId: "peer-1", + success: true, matchIdx: 100)) + check node.matchIndex["peer-1"] == 0 + check node.state == rsLeader + + # Higher term: step down + node.handleInstallSnapshotReply("peer-1", RaftMessage( + kind: rmkInstallSnapshotReply, term: 7, senderId: "peer-1", + success: true, matchIdx: 100)) + check node.state == rsFollower + check node.currentTerm == 7 + + test "floor rejects trigger sendSnapshot end-to-end via processMessage": + proc scenario() = + let tmp = getTempDir() / "baradb_snaptx_e2e_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + var payload = "" + for i in 0 ..< 200: + payload.add(char(32 + (i mod 90))) + + var leader = newRaftNode("leader", @["peer-1"], raftPort = 29331, + dataDir = tmp / "raft-l") + createDir(tmp / "raft-l") # newRaftNode only reads; sendSnapshot writes here + leader.currentTerm = 5 + leader.state = rsLeader + leader.lastSnapshotIndex = 100 + leader.lastSnapshotTerm = 4 + leader.nextIndex["peer-1"] = 101 + leader.matchIndex["peer-1"] = 0 + leader.snapChunkBytes = 64 # 200 bytes -> 4 chunks + leader.peerAddrs["peer-1"] = ("127.0.0.1", 29332) + var buildCalls = 0 + leader.buildSnapshot = proc(destPath: string): bool {.gcsafe.} = + inc buildCalls + check "snap_out_100" in destPath + writeFile(destPath, payload) + true + + var follower = newRaftNode("peer-1", @["leader"], raftPort = 29332, + dataDir = tmp / "raft-f") + follower.currentTerm = 1 + var gotBaseIndex = 0'u64 + var gotBaseTerm = 0'u64 + follower.restoreSnapshot = proc(p: string, bi: uint64, + bt: uint64): bool {.gcsafe.} = + gotBaseIndex = bi + gotBaseTerm = bt + result = readFile(p) == payload + + let netL = newRaftNetwork(leader) + let netF = newRaftNetwork(follower) + asyncCheck netF.run() + waitFor sleepAsync(50) + + # Two floor-level rejects through the real message path; the second one + # must trigger an async snapshot send (leader itself never listens). + let reject = RaftMessage(kind: rmkAppendEntriesReply, term: 5, + senderId: "peer-1", success: false) + waitFor netL.processMessage(reject) + check "peer-1" notin leader.snapPending + waitFor netL.processMessage(reject) + + var waited = 0 + while follower.lastSnapshotIndex != 100 and waited < 3000: + waitFor sleepAsync(50) + waited += 50 + + netF.stop() + waitFor sleepAsync(50) + + check buildCalls == 1 + check follower.lastSnapshotIndex == 100 + check follower.lastSnapshotTerm == 4 + check gotBaseIndex == 100 + check gotBaseTerm == 4 + # Temp archive cleaned up after the transfer + check not fileExists(tmp / "raft-l" / "snap_out_100.tar.gz") + scenario() + + test "sendSnapshot single-flight guard skips a concurrent send": + let tmp = getTempDir() / "baradb_snaptx_guard_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + var node = newRaftNode("leader", @["peer-1"], dataDir = tmp / "raft") + node.currentTerm = 5 + node.state = rsLeader + node.lastSnapshotIndex = 100 + node.lastSnapshotTerm = 4 + var buildCalls = 0 + node.buildSnapshot = proc(destPath: string): bool {.gcsafe.} = + inc buildCalls + writeFile(destPath, "x") + true + + let net = newRaftNetwork(node) + node.snapSending.incl("peer-1") # a send is already in flight + waitFor net.sendSnapshot("peer-1") + check buildCalls == 0 + + test "sendSnapshot skips when there is no compacted snapshot": + let tmp = getTempDir() / "baradb_snaptx_zero_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + var node = newRaftNode("leader", @["peer-1"], dataDir = tmp / "raft") + node.currentTerm = 5 + node.state = rsLeader + # lastSnapshotIndex == 0: snapId 0 can never be received by a follower + var buildCalls = 0 + node.buildSnapshot = proc(destPath: string): bool {.gcsafe.} = + inc buildCalls + true + + let net = newRaftNetwork(node) + waitFor net.sendSnapshot("peer-1") + check buildCalls == 0 + check "peer-1" notin node.snapSending + suite "Raft TLS Transport": test "2-node election over TLS": let certDir = getTempDir() / "baradb_test_raft_tls"