feat(raft): follower InstallSnapshot receive and restore

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)

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:
+  of rmkInstallSnapshot:
-    # Wire protocol only (v1.3); snapshot transfer behavior lands in a
+    # Same election-timer rule as AppendEntries: only a plausible current
-    # follow-up task. Ignore until then.
+    # 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.} =

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)

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)
 

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"