diff --git a/src/barabadb/core/raft.nim b/src/barabadb/core/raft.nim index 2e144ed..3c8ebe8 100644 --- a/src/barabadb/core/raft.nim +++ b/src/barabadb/core/raft.nim @@ -83,6 +83,8 @@ type rmkRequestVoteReply rmkAppendEntries rmkAppendEntriesReply + rmkInstallSnapshot + rmkInstallSnapshotReply RaftMessage* = object kind*: RaftMessageKind @@ -99,6 +101,12 @@ type # Reply success*: bool matchIdx*: uint64 + # InstallSnapshot (prevLogIndex/prevLogTerm reuse: snapshot base index/term; + # reply uses success/matchIdx as usual) + snapId*: uint64 # snapshot generation, matches leader's base at build time + snapOffset*: uint64 # byte offset of this chunk within the archive + snapData*: seq[byte] # chunk payload (<= snapChunkBytes) + snapDone*: bool # last chunk RaftCluster* = ref object nodes*: Table[string, RaftNode] @@ -694,6 +702,14 @@ proc serialize*(msg: RaftMessage): seq[byte] = stream.write(msg.leaderCommit) stream.write(char(if msg.success: 1 else: 0)) stream.write(msg.matchIdx) + # InstallSnapshot trailing fields (appended for wire backward compatibility; + # pre-v1.3 peers stop reading at matchIdx and ignore these bytes) + stream.write(msg.snapId) + stream.write(msg.snapOffset) + stream.write(uint32(msg.snapData.len)) + if msg.snapData.len > 0: + stream.writeData(addr msg.snapData[0], msg.snapData.len) + stream.write(char(if msg.snapDone: 1 else: 0)) let strData = stream.data result = newSeq[byte](strData.len) for i in 0 ..< strData.len: @@ -722,6 +738,19 @@ proc deserializeRaftMessage*(data: seq[byte]): RaftMessage = result.leaderCommit = stream.readUint64() result.success = stream.readChar() != '\0' result.matchIdx = stream.readUint64() + # Optional trailing InstallSnapshot fields (absent in pre-v1.3 buffers) + if not stream.atEnd: + result.snapId = stream.readUint64() + if not stream.atEnd: + result.snapOffset = stream.readUint64() + if not stream.atEnd: + let dataLen = int(stream.readUint32()) + result.snapData = newSeq[byte](dataLen) + if dataLen > 0: + if stream.readData(addr result.snapData[0], dataLen) != dataLen: + raise newException(IOError, "Incomplete snapshot data read from stream") + if not stream.atEnd: + result.snapDone = stream.readChar() != '\0' stream.close() # --------------------------------------------------------------------------- @@ -810,6 +839,10 @@ proc processMessage*(net: RaftNetwork, msg: RaftMessage) {.async.} = await net.send(msg.senderId, reply) of rmkAppendEntriesReply: net.node.handleAppendReply(msg.senderId, msg) + of rmkInstallSnapshot, rmkInstallSnapshotReply: + # Wire protocol only (v1.3); snapshot transfer behavior lands in a + # follow-up task. Ignore until then. + discard proc recvExact*(client: AsyncSocket, size: int): Future[string] {.async.} = ## Reads exactly `size` bytes from `client`. A short return means the peer diff --git a/tests/test_all.nim b/tests/test_all.nim index 9df363d..39e9053 100644 --- a/tests/test_all.nim +++ b/tests/test_all.nim @@ -8,6 +8,7 @@ import std/asyncnet import std/monotimes import std/base64 import std/json +import std/streams import barabadb/core/types import barabadb/core/mvcc @@ -2607,6 +2608,101 @@ suite "Raft Network Transport": check replyMsg.kind == rmkRequestVoteReply check replyMsg.success +suite "Raft InstallSnapshot Protocol": + test "InstallSnapshot fields survive serialize/deserialize round-trip": + let msg = RaftMessage( + kind: rmkInstallSnapshot, + term: 9, + senderId: "leader-1", + prevLogIndex: 42, # snapshot base index + prevLogTerm: 7, # snapshot base term + snapId: 3, + snapOffset: 4096, + snapData: @[byte 1, 2, 3, 250, 0, 17], + snapDone: true) + let decoded = deserializeRaftMessage(serialize(msg)) + check decoded.kind == rmkInstallSnapshot + check decoded.term == 9 + check decoded.senderId == "leader-1" + check decoded.prevLogIndex == 42 + check decoded.prevLogTerm == 7 + check decoded.snapId == 3 + check decoded.snapOffset == 4096 + check decoded.snapData == @[byte 1, 2, 3, 250, 0, 17] + check decoded.snapDone + + test "InstallSnapshotReply fields survive serialize/deserialize round-trip": + let msg = RaftMessage( + kind: rmkInstallSnapshotReply, + term: 9, + senderId: "follower-2", + success: true, + matchIdx: 42, + snapId: 3, + snapOffset: 8192, + snapData: @[], + snapDone: false) + let decoded = deserializeRaftMessage(serialize(msg)) + check decoded.kind == rmkInstallSnapshotReply + check decoded.term == 9 + check decoded.senderId == "follower-2" + check decoded.success + check decoded.matchIdx == 42 + check decoded.snapId == 3 + check decoded.snapOffset == 8192 + check decoded.snapData.len == 0 + check not decoded.snapDone + + test "old wire layout (no snapshot fields) deserializes with zero defaults": + # Manually serialize a message in the pre-InstallSnapshot layout: + # magic, version, kind, term, senderId, lastLogIndex, lastLogTerm, + # prevLogIndex, prevLogTerm, entries, leaderCommit, success, matchIdx. + let s = newStringStream() + s.write("RAFT") + s.write(1'u32) # RaftProtoVersion + s.write(uint32(ord(rmkAppendEntries))) + s.write(5'u64) # term + let sender = "old-leader" + s.write(uint32(sender.len)) + s.writeData(sender[0].unsafeAddr, sender.len) + s.write(11'u64) # lastLogIndex + s.write(4'u64) # lastLogTerm + s.write(10'u64) # prevLogIndex + s.write(4'u64) # prevLogTerm + s.write(0'u32) # entries count + s.write(10'u64) # leaderCommit + s.write(char(1)) # success + s.write(10'u64) # matchIdx + let strData = s.data + var buf = newSeq[byte](strData.len) + for i in 0 ..< strData.len: + buf[i] = byte(strData[i]) + s.close() + + let decoded = deserializeRaftMessage(buf) + check decoded.kind == rmkAppendEntries + check decoded.term == 5 + check decoded.senderId == "old-leader" + check decoded.matchIdx == 10 + check decoded.snapId == 0 + check decoded.snapOffset == 0 + check decoded.snapData.len == 0 + check not decoded.snapDone + + test "old message kinds still round-trip unchanged": + let msg = RaftMessage(kind: rmkRequestVote, term: 2, senderId: "cand", + lastLogIndex: 5, lastLogTerm: 1) + let decoded = deserializeRaftMessage(serialize(msg)) + check decoded.kind == rmkRequestVote + check decoded.term == 2 + check decoded.senderId == "cand" + check decoded.lastLogIndex == 5 + check decoded.lastLogTerm == 1 + check decoded.snapId == 0 + check decoded.snapOffset == 0 + check decoded.snapData.len == 0 + check not decoded.snapDone + suite "Raft TLS Transport": test "2-node election over TLS": let certDir = getTempDir() / "baradb_test_raft_tls"