feat(raft): leader InstallSnapshot send on unrecoverable lag
This commit is contained in:
+121
-2
@@ -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
|
||||
|
||||
@@ -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_<snapId>.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)
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user