From efa04e4b363694cf7828afec2e9befeb6f1a25e1 Mon Sep 17 00:00:00 2001 From: dimgigov Date: Fri, 31 Jul 2026 02:33:32 +0300 Subject: [PATCH] feat(raft): follower InstallSnapshot receive and restore --- src/barabadb/core/config.nim | 3 + src/barabadb/core/raft.nim | 105 ++++++++++++++++++++++++- src/barabadb/core/registry.nim | 26 +++++++ src/baradadb.nim | 31 ++++++++ tests/test_all.nim | 136 +++++++++++++++++++++++++++++++++ 5 files changed, 298 insertions(+), 3 deletions(-) diff --git a/src/barabadb/core/config.nim b/src/barabadb/core/config.nim index 2bd5185..744ac01 100644 --- a/src/barabadb/core/config.nim +++ b/src/barabadb/core/config.nim @@ -44,6 +44,7 @@ type raftPeerClientAddrs*: Table[string, tuple[host: string, port: int]] raftWriteTimeoutMs*: int raftLogMaxEntries*: int + raftSnapChunkKb*: int raftTlsEnabled*: bool raftTlsCertFile*: string raftTlsKeyFile*: string @@ -91,6 +92,7 @@ proc defaultConfig*(): BaraConfig = raftPeerClientAddrs: initTable[string, tuple[host: string, port: int]](), raftWriteTimeoutMs: 5_000, raftLogMaxEntries: 256, + raftSnapChunkKb: 256, raftTlsEnabled: false, raftTlsCertFile: "", raftTlsKeyFile: "", @@ -223,6 +225,7 @@ proc loadConfigFromEnv*(cfg: var BaraConfig) = cfg.raftNodeId = getEnv("BARADB_RAFT_NODE_ID", cfg.raftNodeId) cfg.raftWriteTimeoutMs = parseEnvInt(getEnv("BARADB_RAFT_WRITE_TIMEOUT_MS", ""), cfg.raftWriteTimeoutMs) cfg.raftLogMaxEntries = parseEnvInt(getEnv("BARADB_RAFT_LOG_MAX_ENTRIES", ""), cfg.raftLogMaxEntries) + cfg.raftSnapChunkKb = parseEnvInt(getEnv("BARADB_RAFT_SNAP_CHUNK_KB", ""), cfg.raftSnapChunkKb) cfg.raftTlsEnabled = parseEnvBool(getEnv("BARADB_RAFT_TLS_ENABLED", ""), cfg.raftTlsEnabled) cfg.raftTlsCertFile = getEnv("BARADB_RAFT_TLS_CERT_FILE", cfg.raftTlsCertFile) cfg.raftTlsKeyFile = getEnv("BARADB_RAFT_TLS_KEY_FILE", cfg.raftTlsKeyFile) diff --git a/src/barabadb/core/raft.nim b/src/barabadb/core/raft.nim index 3c8ebe8..71e1056 100644 --- a/src/barabadb/core/raft.nim +++ b/src/barabadb/core/raft.nim @@ -77,6 +77,15 @@ type peerAddrs*: Table[string, tuple[host: string, port: int]] raftPort*: int dataDir*: string + ## InstallSnapshot follower receive. snapChunkBytes caps a single chunk + ## (from BARADB_RAFT_SNAP_CHUNK_KB, default 262144); snapIncomingId / + ## snapIncomingFile track the archive currently being assembled under + ## dataDir/snap_incoming/. + snapChunkBytes*: int + restoreSnapshot*: proc(archivePath: string, baseIndex: uint64, + baseTerm: uint64): bool {.gcsafe.} + snapIncomingId*: uint64 + snapIncomingFile*: string RaftMessageKind* = enum rmkRequestVote @@ -208,6 +217,9 @@ proc newRaftNode*(id: string, peers: seq[string], raftPort: int = 0, peerAddrs: initTable[string, tuple[host: string, port: int]](), raftPort: raftPort, dataDir: dataDir, + snapChunkBytes: 262144, + snapIncomingId: 0, + snapIncomingFile: "", ) result.loadState() @@ -427,6 +439,87 @@ proc handleAppendEntries*(node: RaftNode, msg: RaftMessage): RaftMessage = reply.matchIdx = node.lastLogIndex return reply +proc handleInstallSnapshot*(node: RaftNode, msg: RaftMessage): RaftMessage = + ## Follower side of InstallSnapshot: assemble the chunk stream into a temp + ## archive under `dataDir/snap_incoming/`, then hand the completed archive + ## to the restoreSnapshot callback. Chunks arrive in order from a single + ## leader over one socket, so we append sequentially and only sanity-check + ## that snapOffset equals the number of bytes assembled so far. + ## + ## NOTE: this runs on the async event loop and restoreSnapshot performs + ## blocking disk I/O (archive extract + DB reopen). Implementations must be + ## fast, or defer the heavy work; the baradadb.nim wiring decides. + var reply = RaftMessage( + kind: rmkInstallSnapshotReply, + term: node.currentTerm, + senderId: node.id, + success: false, + matchIdx: node.lastSnapshotIndex, + ) + if msg.term < node.currentTerm: + return reply + if msg.term > node.currentTerm: + node.becomeFollower(msg.term) + node.leaderId = msg.senderId + + # Chunk size cap (deferred from the wire-protocol task). + if msg.snapData.len > node.snapChunkBytes or node.dataDir.len == 0: + return reply + + let snapDir = node.dataDir / "snap_incoming" + if msg.snapId != node.snapIncomingId: + # New snapshot generation: discard any partial assembly and restart. + if msg.snapOffset != 0: + return reply + createDir(snapDir) + node.snapIncomingId = msg.snapId + node.snapIncomingFile = snapDir / "snap_" & $msg.snapId & ".tar.gz" + let f = open(node.snapIncomingFile, fmWrite) # truncate any leftover + f.close() + + if node.snapIncomingFile.len == 0: + return reply + + let assembled = getFileSize(node.snapIncomingFile) + if msg.snapOffset != uint64(assembled): + # Gap or overlap: reset so the leader restarts the transfer. + removeFile(node.snapIncomingFile) + node.snapIncomingId = 0 + node.snapIncomingFile = "" + return reply + + if msg.snapData.len > 0: + let f = open(node.snapIncomingFile, fmAppend) + try: + discard f.writeBuffer(addr msg.snapData[0], msg.snapData.len) + finally: + f.close() + + if not msg.snapDone: + reply.success = true + return reply + + # Transfer complete: restore the data dir and adopt the snapshot base. + if node.restoreSnapshot == nil or + not node.restoreSnapshot(node.snapIncomingFile, + msg.prevLogIndex, msg.prevLogTerm): + removeFile(node.snapIncomingFile) + node.snapIncomingId = 0 + node.snapIncomingFile = "" + return reply + + node.lastSnapshotIndex = msg.prevLogIndex + node.lastSnapshotTerm = msg.prevLogTerm + node.commitIndex = node.lastSnapshotIndex + node.lastApplied = node.lastSnapshotIndex + node.log = @[] + node.snapIncomingId = 0 + node.snapIncomingFile = "" + node.saveState() + reply.success = true + reply.matchIdx = node.lastSnapshotIndex + return reply + proc requestVote*(node: RaftNode): seq[RaftMessage] = result = @[] for peer in node.peers: @@ -839,9 +932,15 @@ 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. + of rmkInstallSnapshot: + # Same election-timer rule as AppendEntries: only a plausible current + # leader resets it. + if msg.term >= net.node.currentTerm: + net.timer.resetTimeout() + 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 proc recvExact*(client: AsyncSocket, size: int): Future[string] {.async.} = diff --git a/src/barabadb/core/registry.nim b/src/barabadb/core/registry.nim index 154c05a..8da13df 100644 --- a/src/barabadb/core/registry.nim +++ b/src/barabadb/core/registry.nim @@ -203,6 +203,32 @@ proc getDatabaseInfo*(reg: DatabaseRegistry, name: string): DatabaseInfo = return reg.databases[name] return nil +proc reopenDatabase*(reg: DatabaseRegistry, name: string): bool = + ## Reopen a database from its on-disk directory, swapping the new LSMTree + ## and ctx into the EXISTING DatabaseInfo slot so captured references (e.g. + ## the raft applyCommand closure) see the new state. + ## Minimal API added for raft InstallSnapshot restore: the caller must have + ## closed info.db first (snapshot restore closes it before swapping the data + ## directory); this proc does not close. + ## Returns false if the database is unknown or the reopen fails. + acquire(reg.lock) + let info = if name in reg.databases: reg.databases[name] else: nil + release(reg.lock) + if info == nil: + return false + try: + let dbDir = reg.dataRoot / name + let db = openLsmForRegistry(reg, dbDir) + let ctx = reg.ctxFactory(db, reg) + info.db = db + info.ctx = ctx + return true + except CatchableError as e: + # echo instead of logging: callers include gcsafe raft callbacks, and + # core/logging's info/warn are not gcsafe. + echo "[registry] Error reopening database '", name, "': ", e.msg + return false + proc closeAll*(reg: DatabaseRegistry) = acquire(reg.lock) defer: release(reg.lock) diff --git a/src/baradadb.nim b/src/baradadb.nim index ba3fd4d..5d954e3 100644 --- a/src/baradadb.nim +++ b/src/baradadb.nim @@ -23,6 +23,7 @@ import barabadb/core/gossip import barabadb/core/replication import barabadb/core/disttxn import barabadb/core/registry +import barabadb/core/backup import barabadb/tools/repair import barabadb/tools/migrate @@ -365,6 +366,8 @@ proc main() = raftNode.peerAddrs = config.raftPeerAddrs if config.raftLogMaxEntries > 0: raftNode.logMaxEntries = config.raftLogMaxEntries + if config.raftSnapChunkKb > 0: + raftNode.snapChunkBytes = config.raftSnapChunkKb * 1024 tcpServer.raftNode = raftNode # C3b: executeQuery rejects writes on followers httpServer.raftNode = raftNode # /metrics + /health raft gauges # Wire state machine: committed entries update LSM + secondary indexes @@ -382,6 +385,34 @@ proc main() = elif cmd == "ddl": applyReplicatedDdl(ctx, cast[string](data)) + # Follower InstallSnapshot restore: swap the default DB's data directory + # with the received archive, then reopen it into the same DatabaseInfo + # slot (the applyCommand closure above keeps working through the swap). + # Runs on the raft async event loop and performs blocking disk I/O + # (tar extract + LSM close/reopen); snapshot installs are rare, so we + # accept the stall rather than adding a worker round-trip. + let defaultDbDir = config.dataDir / "databases" / "default" + raftNode.restoreSnapshot = proc(archivePath: string, baseIndex: uint64, + baseTerm: uint64): bool {.gcsafe.} = + # NOTE: core/logging's info/warn are not gcsafe (global logger), so + # this callback stays silent; restoreDataDir echoes progress itself. + echo "[raft] Installing snapshot (base index ", baseIndex, + ", base term ", baseTerm, ")" + # gcsafe cast: this runs on the raft event-loop thread (same thread as + # the rest of the server); the registry ctxFactory type is not marked + # gcsafe, which would otherwise reject the call. + {.cast(gcsafe).}: + try: + defaultDbInfo.db.close() + # restoreDataDir moves the old dir aside and extracts the archive; on + # extraction failure it rolls back automatically. Reopen whatever is + # on disk either way so the node is not left with a closed DB. + let restored = restoreDataDir(archivePath, defaultDbDir) + result = registry.reopenDatabase("default") and restored + except CatchableError as e: + echo "[raft] Snapshot restore failed: ", e.msg + result = false + # Wire RAFT ↔ DistTxn wireRaftDistTxn(raftNode, tcpServer) diff --git a/tests/test_all.nim b/tests/test_all.nim index 39e9053..987a7cd 100644 --- a/tests/test_all.nim +++ b/tests/test_all.nim @@ -14,6 +14,7 @@ import barabadb/core/types import barabadb/core/mvcc import barabadb/core/deadlock import barabadb/core/config +import barabadb/core/backup import barabadb/core/server import barabadb/core/columnar import barabadb/core/raft @@ -2703,6 +2704,141 @@ suite "Raft InstallSnapshot Protocol": check decoded.snapData.len == 0 check not decoded.snapDone +suite "Raft InstallSnapshot Receive": + test "follower assembles chunks, restores snapshot, resets state": + proc scenario() = + let tmp = getTempDir() / "baradb_snaprx_ok_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + # Real tar.gz fixture with a marker file + let srcDb = tmp / "srcdb" + createDir(srcDb) + writeFile(srcDb / "marker.txt", "snapshot-payload") + let archivePath = tmp / "snap.tar.gz" + check backupDataDir(srcDb, archivePath) + let archiveBytes = readFile(archivePath) + check archiveBytes.len > 0 + + let raftDir = tmp / "raft" + var node = newRaftNode("follower-1", @[], dataDir = raftDir) + node.currentTerm = 5 + node.log.add(LogEntry(term: 3, index: 10, command: "put", data: @[byte 1])) + node.commitIndex = 10 + + var gotPath = "" + var gotBaseIndex = 0'u64 + var gotBaseTerm = 0'u64 + node.restoreSnapshot = proc(p: string, bi: uint64, bt: uint64): bool {.gcsafe.} = + gotPath = p + gotBaseIndex = bi + gotBaseTerm = bt + # Assembled archive must match the original byte-for-byte + result = readFile(p) == archiveBytes + + let half = archiveBytes.len div 2 + let reply1 = node.handleInstallSnapshot(RaftMessage( + kind: rmkInstallSnapshot, term: 5, senderId: "leader-1", + prevLogIndex: 40, prevLogTerm: 4, + snapId: 7, snapOffset: 0, + snapData: cast[seq[byte]](archiveBytes[0 ..< half]), snapDone: false)) + check reply1.kind == rmkInstallSnapshotReply + check reply1.success + + let reply2 = node.handleInstallSnapshot(RaftMessage( + kind: rmkInstallSnapshot, term: 5, senderId: "leader-1", + prevLogIndex: 40, prevLogTerm: 4, + snapId: 7, snapOffset: uint64(half), + snapData: cast[seq[byte]](archiveBytes[half .. ^1]), snapDone: true)) + check reply2.success + check reply2.matchIdx == 40 + + check gotPath.len > 0 + check "snap_incoming" in gotPath + check gotBaseIndex == 40 + check gotBaseTerm == 4 + check node.lastSnapshotIndex == 40 + check node.lastSnapshotTerm == 4 + check node.commitIndex == 40 + check node.lastApplied == 40 + check node.log.len == 0 + + # State was persisted: a fresh node on the same dir sees the snapshot base + let reloaded = newRaftNode("follower-1", @[], dataDir = raftDir) + check reloaded.lastSnapshotIndex == 40 + check reloaded.lastSnapshotTerm == 4 + check reloaded.log.len == 0 + scenario() + + test "failed restore leaves state untouched and removes temp file": + proc scenario() = + let tmp = getTempDir() / "baradb_snaprx_fail_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + let raftDir = tmp / "raft" + var node = newRaftNode("follower-1", @[], dataDir = raftDir) + node.currentTerm = 5 + node.log.add(LogEntry(term: 2, index: 3, command: "put", data: @[byte 9])) + node.restoreSnapshot = proc(p: string, bi: uint64, bt: uint64): bool {.gcsafe.} = + false + + let reply = node.handleInstallSnapshot(RaftMessage( + kind: rmkInstallSnapshot, term: 5, senderId: "leader-1", + prevLogIndex: 8, prevLogTerm: 2, + snapId: 1, snapOffset: 0, + snapData: @[byte 1, 2, 3], snapDone: true)) + check not reply.success + check node.lastSnapshotIndex == 0 + check node.lastSnapshotTerm == 0 + check node.commitIndex == 0 + check node.log.len == 1 + check not fileExists(raftDir / "snap_incoming" / "snap_1.tar.gz") + scenario() + + test "oversized chunk is rejected": + let tmp = getTempDir() / "baradb_snaprx_cap_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + var node = newRaftNode("follower-1", @[], dataDir = tmp / "raft") + node.currentTerm = 1 + node.snapChunkBytes = 4 + let reply = node.handleInstallSnapshot(RaftMessage( + kind: rmkInstallSnapshot, term: 1, senderId: "leader-1", + prevLogIndex: 1, prevLogTerm: 1, + snapId: 1, snapOffset: 0, + snapData: @[byte 1, 2, 3, 4, 5], snapDone: false)) + check not reply.success + check node.snapIncomingId == 0 + + test "out-of-order offset is rejected and assembly restarts on new snapId": + let tmp = getTempDir() / "baradb_snaprx_off_" & $getCurrentProcessId() + removeDir(tmp) + createDir(tmp) + defer: removeDir(tmp) + + let raftDir = tmp / "raft" + var node = newRaftNode("follower-1", @[], dataDir = raftDir) + node.currentTerm = 1 + let ok1 = node.handleInstallSnapshot(RaftMessage( + kind: rmkInstallSnapshot, term: 1, senderId: "leader-1", + prevLogIndex: 2, prevLogTerm: 1, + snapId: 3, snapOffset: 0, + snapData: @[byte 65, 66], snapDone: false)) + check ok1.success + # Gap: offset 5 while only 2 bytes assembled + let bad = node.handleInstallSnapshot(RaftMessage( + kind: rmkInstallSnapshot, term: 1, senderId: "leader-1", + prevLogIndex: 2, prevLogTerm: 1, + snapId: 3, snapOffset: 5, + snapData: @[byte 67], snapDone: false)) + check not bad.success + check node.snapIncomingId == 0 + suite "Raft TLS Transport": test "2-node election over TLS": let certDir = getTempDir() / "baradb_test_raft_tls"