fix: resolve all 55 identified bugs across the codebase

Comprehensive bug fix pass across all subsystems: - Critical: lexer infinite loop, WAL lock discipline, checkpoint safety, compaction verification, HMAC key truncation, healthCheck leak - High: MVCC lock safety, B-Tree delete rebalancing, Raft stale term handling, replication socket leaks and ack cleanup, sharding migration with old assignments, 2PC recovery, JWT/SCRAM auth fixes, wire protocol bounds checks - Medium: config error handling, MERGE parser completeness, IR/codegen correctness, UDF bounds checks, LIMIT cost accuracy, mmap safety - Low: timeout handling, float parsing edge cases, uint32 truncation, cache entry cleanup Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-21 10:59:51 +03:00
parent 7fc6adbe47
commit 2d310a33a1
23 changed files with 765 additions and 207 deletions
@@ -118,10 +118,10 @@ proc loadConfigFromJson*(path: string, cfg: var BaraConfig) =
      if s.hasKey("query_timeout_ms"): cfg.queryTimeoutMs = s["query_timeout_ms"].getInt()
      if s.hasKey("slow_query_threshold_ms"): cfg.slowQueryThresholdMs = s["slow_query_threshold_ms"].getInt()
      if s.hasKey("slow_query_log_path"): cfg.slowQueryLogPath = s["slow_query_log_path"].getStr()
-  except JsonParsingError:
-    discard
-  except KeyError:
-    discard
+  except JsonParsingError as e:
+    echo "[WARN] Failed to parse config file ", path, ": ", e.msg
+  except KeyError as e:
+    echo "[WARN] Missing key in config file ", path, ": ", e.msg

 # ----------------------------------------------------------------------
 # Environment Variables
@@ -182,4 +182,4 @@ proc loadConfig*(): BaraConfig =
 proc getEffectiveJwtSecret*(cfg: BaraConfig): string =
  if cfg.jwtSecret.len > 0:
    return cfg.jwtSecret
-  return "baradb-default-secret-change-in-production!"
+  return ""
@@ -23,6 +23,8 @@ type
    prepared*: bool
    committed*: bool
    aborted*: bool
+    commitPending*: bool
+    rollbackPending*: bool
    errorMsg*: string

  DistributedTransaction* = ref object
@@ -94,19 +96,15 @@ proc sendDistTxnRpc(host: string, port: int, txnId: uint64, action: string, time
  ## Send 2PC RPC to participant node via TCP text protocol.
  ## Protocol: "DISTTXN <txnId> <action>\n" where action = PREPARE|COMMIT|ROLLBACK
  ## Response: "OK\n" or "ERR <msg>\n"
-  try:
-    var sock = newSocket()
-    if not connectWithTimeout(sock, host, Port(port), timeoutMs):
-      sock.close()
-      return false
-    let msg = "DISTTXN " & $txnId & " " & action & "\n"
-    sock.send(msg)
-    var response = ""
-    sock.readLine(response)
-    sock.close()
-    return response.strip() == "OK"
-  except CatchableError:
+  var sock = newSocket()
+  defer: sock.close()
+  if not connectWithTimeout(sock, host, Port(port), timeoutMs):
    return false
+  let msg = "DISTTXN " & $txnId & " " & action & "\n"
+  sock.send(msg)
+  var response = ""
+  sock.readLine(response)
+  return response.strip() == "OK"

 type
  ParticipantInfo = object
@@ -147,13 +145,20 @@ proc prepare*(txn: DistributedTransaction): bool =
    for nodeId, _ in txn.participants.mpairs:
      txn.participants[nodeId].prepared = true
  else:
-    # Rollback already-prepared participants
+    # Rollback already-prepared participants; track failures for recovery
+    var rollbackFailed = false
    for nodeId in preparedNodes:
      if txn.participants.hasKey(nodeId):
-        discard sendDistTxnRpc(txn.participants[nodeId].host, txn.participants[nodeId].port, txn.id, "ROLLBACK")
-        txn.participants[nodeId].prepared = false
-        txn.participants[nodeId].aborted = true
+        let ok = sendDistTxnRpc(txn.participants[nodeId].host, txn.participants[nodeId].port, txn.id, "ROLLBACK")
+        if ok:
+          txn.participants[nodeId].prepared = false
+          txn.participants[nodeId].aborted = true
+        else:
+          txn.participants[nodeId].rollbackPending = true
+          rollbackFailed = true
    txn.state = dtsAborted
+    if rollbackFailed:
+      echo "[WARN] 2PC rollback failed for some participants of txn ", txn.id, " — recovery needed"
  release(txn.lock)
  return allOk

@@ -190,10 +195,15 @@ proc commit*(txn: DistributedTransaction): bool =
    for nodeId, _ in txn.participants.mpairs:
      txn.participants[nodeId].committed = true
  elif committedNodes.len > 0:
+    # Partial commit — mark committed, flag uncommitted for recovery
    txn.state = dtsCommitted
    for nodeId in committedNodes:
      if txn.participants.hasKey(nodeId):
        txn.participants[nodeId].committed = true
+    for nodeId, p in txn.participants.mpairs:
+      if not p.committed:
+        p.commitPending = true
+    echo "[WARN] 2PC partial commit for txn ", txn.id, " (", committedNodes.len, "/", txn.participants.len, ") — recovery needed"
  else:
    txn.state = dtsAborted
  release(txn.lock)
@@ -288,7 +298,10 @@ proc execute*(saga: Saga): bool =
      # Rollback: compensate completed steps in reverse order
      for j in countdown(saga.completedSteps.len - 1, 0):
        let idx = saga.completedSteps[j]
-        saga.steps[idx].compensate()
+        try:
+          saga.steps[idx].compensate()
+        except CatchableError as e:
+          echo "[ERROR] Saga compensation failed for step ", idx, ": ", e.msg
      return false
  return true

@@ -293,10 +293,14 @@ proc handleIncomingGossip(gp: GossipProtocol, data: string, senderAddr: string)
      gp.members[msg.senderId].lastSeen = getMonoTime().ticks()
    elif msg.senderId != gp.self.id:
      var host = senderAddr
+      var port = gp.gossipPort
      if ':' in host:
-        host = host.split(":")[0]
+        let parts = host.split(":")
+        host = parts[0]
+        if parts[1].len > 0:
+          port = try: parseInt(parts[1]) except: gp.gossipPort
      let newNode = GossipNode(
-        id: msg.senderId, host: host, port: gp.gossipPort,
+        id: msg.senderId, host: host, port: port,
        state: nsAlive, incarnation: msg.senderIncarnation,
        lastSeen: getMonoTime().ticks(),
      )
@@ -137,8 +137,9 @@ proc isVisible(tm: TxnManager, txn: Transaction, version: VersionedRecord): bool
      if deleter in tm.activeTxns:
        if tm.activeTxns[deleter].state == tsCommitted:
          return false
-      else:
-        return false
+      elif deleter in tm.committedTxnsSet:
+        return false  # deleter committed after snapshot, conservatively show record
+      # Deleter not in active, committed, or aborted — unknown state, show record

  return true

@@ -195,6 +196,7 @@ proc write*(tm: TxnManager, txn: Transaction, key: string, value: seq[byte]): bo
          if victimId in tm.activeTxns:
            tm.activeTxns[victimId].state = tsAborted
            tm.activeTxns.del(victimId)
+        tm.deadlockDetector.removeWait(uint64(txn.id), uint64(otherId))
        release(tm.lock)
        return false  # write-write conflict with uncommitted txn

@@ -255,6 +257,7 @@ proc delete*(tm: TxnManager, txn: Transaction, key: string): bool =
          if victimId in tm.activeTxns:
            tm.activeTxns[victimId].state = tsAborted
            tm.activeTxns.del(victimId)
+        tm.deadlockDetector.removeWait(uint64(txn.id), uint64(otherId))
        release(tm.lock)
        return false

@@ -326,7 +329,7 @@ proc commit*(tm: TxnManager, txn: Transaction): bool =
  release(tm.lock)
  return true

-proc compactVersions*(tm: TxnManager) =
+proc compactVersions(tm: TxnManager) =
  ## Remove old overwritten versions that are no longer visible to any active transaction.
  for key, versions in tm.globalVersions.mpairs:
    if versions.len <= 3:
@@ -184,6 +184,14 @@ proc lastLogTerm*(node: RaftNode): uint64 =
    return 0
  return node.log[^1].term

+proc findLogEntryByIndex(node: RaftNode, index: uint64): int =
+  ## Find array position for a logical log index.
+  ## Returns -1 if not found. Does NOT assume index - 1 == array position.
+  for i, entry in node.log:
+    if entry.index == index:
+      return i
+  return -1
+
 proc applyCommitted(node: RaftNode) =
  while node.lastApplied < node.commitIndex:
    let idx = int(node.lastApplied)
@@ -211,6 +219,8 @@ proc becomeFollower*(node: RaftNode, term: uint64) =
  node.currentTerm = term
  node.votedFor = ""
  node.votesReceived.clear()
+  node.nextIndex.clear()
+  node.matchIndex.clear()
  node.saveState()

 proc becomeCandidate*(node: RaftNode) =
@@ -266,26 +276,27 @@ proc handleAppendEntries*(node: RaftNode, msg: RaftMessage): RaftMessage =
  if msg.term < node.currentTerm:
    return reply

-  if msg.term >= node.currentTerm:
+  if msg.term > node.currentTerm:
    node.becomeFollower(msg.term)
-    node.leaderId = msg.senderId
+  node.leaderId = msg.senderId

  # Check if log contains entry at prevLogIndex with prevLogTerm
  if msg.prevLogIndex > 0:
-    if msg.prevLogIndex > uint64(node.log.len):
+    let prevPos = node.findLogEntryByIndex(msg.prevLogIndex)
+    if prevPos < 0:
      return reply
-    if node.log[msg.prevLogIndex - 1].term != msg.prevLogTerm:
+    if node.log[prevPos].term != msg.prevLogTerm:
      # Delete conflicting entries
-      node.log.setLen(int(msg.prevLogIndex - 1))
+      node.log.setLen(prevPos)
      return reply

  # Append new entries
  var logChanged = false
  for entry in msg.entries:
-    let idx = int(entry.index - 1)
-    if idx < node.log.len:
-      if node.log[idx].term != entry.term:
-        node.log.setLen(idx)
+    let pos = node.findLogEntryByIndex(entry.index)
+    if pos >= 0:
+      if node.log[pos].term != entry.term:
+        node.log.setLen(pos)
        node.log.add(entry)
        logChanged = true
    else:
@@ -354,6 +365,9 @@ proc handleVoteReply*(node: RaftNode, reply: RaftMessage) =
    node.becomeFollower(reply.term)
    return

+  if reply.term < node.currentTerm:
+    return
+
  if node.state != rsCandidate:
    return

@@ -367,6 +381,9 @@ proc handleAppendReply*(node: RaftNode, peerId: string, reply: RaftMessage) =
    node.becomeFollower(reply.term)
    return

+  if reply.term < node.currentTerm:
+    return
+
  if node.state != rsLeader:
    return

@@ -65,9 +65,11 @@ proc loadExistingDatabases*(reg: DatabaseRegistry) =
        info("Loading database '" & dbName & "' from " & dbDir)
        let db = newLSMTree(dbDir)
        let ctx = reg.ctxFactory(db, reg)
+        acquire(reg.lock)
        reg.databases[dbName] = DatabaseInfo(
          name: dbName, db: db, ctx: ctx, activeConnections: 0
        )
+        release(reg.lock)

 proc setDatabase*(reg: DatabaseRegistry, name: string, db: LSMTree, ctx: ContextRef) =
  acquire(reg.lock)
@@ -80,14 +82,20 @@ proc ensureDefaultDatabase*(reg: DatabaseRegistry) =
    raise newException(ValueError, "Context factory not set. Call setContextFactory first.")

  let defaultDbName = reg.defaultDbName
-  if defaultDbName notin reg.databases:
+  acquire(reg.lock)
+  let exists = defaultDbName in reg.databases
+  release(reg.lock)
+
+  if not exists:
    let dbDir = reg.dataRoot / defaultDbName
    info("Creating default database at " & dbDir)
    let db = newLSMTree(dbDir)
    let ctx = reg.ctxFactory(db, reg)
+    acquire(reg.lock)
    reg.databases[defaultDbName] = DatabaseInfo(
      name: defaultDbName, db: db, ctx: ctx, activeConnections: 0
    )
+    release(reg.lock)

 proc getOrCreateDatabase*(reg: DatabaseRegistry, name: string): DatabaseInfo =
  if not isValidDbName(name):
@@ -135,29 +143,30 @@ proc dropDatabase*(reg: DatabaseRegistry, name: string): bool =
    return false

  acquire(reg.lock)
-  defer: release(reg.lock)
-
  if name notin reg.databases:
+    release(reg.lock)
    return false

  if name == reg.defaultDbName:
+    release(reg.lock)
    raise newException(ValueError, "Cannot drop the default database")

  let info = reg.databases[name]
  if info.activeConnections > 0:
+    release(reg.lock)
    raise newException(ValueError,
      "Cannot drop database '" & name & "': " &
      $info.activeConnections & " active connections")

-  # Close LSMTree
-  info.db.close()
+  # Remove from registry first so no new references can be obtained
+  reg.databases.del(name)
+  release(reg.lock)

-  # Remove data directory
+  # Close LSMTree and remove directory outside the lock
+  info.db.close()
  let dbDir = reg.dataRoot / name
  if dirExists(dbDir):
    removeDir(dbDir)
-
-  reg.databases.del(name)
  true

 proc listDatabases*(reg: DatabaseRegistry): seq[string] =
@@ -3,6 +3,7 @@ import std/tables
 import std/sets
 import std/locks
 import std/net
+import std/posix
 import std/strutils
 import std/nativesockets
 import std/monotimes
@@ -88,29 +89,29 @@ proc connectWithTimeout(sock: Socket, host: string, port: Port, timeoutMs: int):
    var fds = @[sock.getFd]
    if selectWrite(fds, timeoutMs) <= 0:
      return false
+    # Verify connection actually succeeded via SO_ERROR
+    var err: cint = 0
+    var errLen = cint(sizeof(err)).SockLen
+    discard posix.getsockopt(sock.getFd, 1'i32, 4'i32, addr err, addr errLen)
    sock.getFd.setBlocking(true)
-    return true
+    return err == 0

 proc shipToReplica(replica: Replica, lsn: uint64, data: seq[byte]): bool =
  ## Send replication data to a replica via TCP.
  ## Protocol: "REP <lsn> <dataLen>\n<data>"
  ## Response: "ACK <lsn>\n" on success
-  try:
-    var sock = newSocket()
-    if not connectWithTimeout(sock, replica.host, Port(replica.port), 500):
-      sock.close()
-      return false
-    let header = "REP " & $lsn & " " & $data.len & "\n"
-    sock.send(header)
-    if data.len > 0:
-      sock.send(cast[string](data))
-    var response = ""
-    sock.readLine(response)
-    sock.close()
-    let parts = response.strip().split(" ")
-    return parts.len >= 2 and parts[0] == "ACK"
-  except:
+  var sock = newSocket()
+  defer: sock.close()
+  if not connectWithTimeout(sock, replica.host, Port(replica.port), 500):
    return false
+  let header = "REP " & $lsn & " " & $data.len & "\n"
+  sock.send(header)
+  if data.len > 0:
+    sock.send(cast[string](data))
+  var response = ""
+  sock.readLine(response)
+  let parts = response.strip().split(" ")
+  return parts.len >= 2 and parts[0] == "ACK"

 proc writeLsn*(rm: ReplicationManager, data: seq[byte]): uint64 =
  acquire(rm.lock)
@@ -135,11 +136,20 @@ proc writeLsn*(rm: ReplicationManager, data: seq[byte]): uint64 =
        rm.pendingAcks[lsn].incl(replica.id)
    release(rm.lock)
    var ackCount = 0
+    var ackedIds: seq[string] = @[]
    for replica in replicasToShip:
      if shipToReplica(replica, lsn, data):
        inc ackCount
+        ackedIds.add(replica.id)
+    # Clean up pendingAcks for successfully acked replicas
+    acquire(rm.lock)
+    if lsn in rm.pendingAcks:
+      for id in ackedIds:
+        rm.pendingAcks[lsn].excl(id)
+      if rm.pendingAcks[lsn].len == 0:
+        rm.pendingAcks.del(lsn)
+    release(rm.lock)
    if replicasToShip.len > 0 and ackCount < replicasToShip.len:
-      # Not all replicas acked — log but still return LSN (caller decides)
      when defined(debug):
        echo "Replication sync: only ", ackCount, "/", replicasToShip.len, " replicas acked for LSN ", lsn
    return lsn
@@ -153,11 +163,21 @@ proc writeLsn*(rm: ReplicationManager, data: seq[byte]): uint64 =
          inc count
    release(rm.lock)
    var ackCount = 0
+    var ackedIds: seq[string] = @[]
    for replica in replicasToShip:
      if shipToReplica(replica, lsn, data):
        inc ackCount
+        ackedIds.add(replica.id)
        if ackCount >= rm.syncReplicaCount:
          break
+    # Clean up pendingAcks for successfully acked replicas
+    acquire(rm.lock)
+    if lsn in rm.pendingAcks:
+      for id in ackedIds:
+        rm.pendingAcks[lsn].excl(id)
+      if rm.pendingAcks[lsn].len == 0:
+        rm.pendingAcks.del(lsn)
+    release(rm.lock)
    if replicasToShip.len > 0 and ackCount == 0 and rm.syncReplicaCount > 0:
      when defined(debug):
        echo "Replication semi-sync: no replicas acked for LSN ", lsn
@@ -226,56 +246,84 @@ proc switchMode*(rm: ReplicationManager, mode: ReplicationMode) =

 proc healthCheck*(rm: ReplicationManager) =
  acquire(rm.lock)
+  var replicas: seq[(string, Replica)] = @[]
  for id, replica in rm.replicas:
    if replica.connected:
-      # Probe connection by sending a heartbeat
-      var sock = newSocket()
+      replicas.add((id, replica))
+  release(rm.lock)
+
+  for (id, replica) in replicas:
+    var connected = true
+    var sock = newSocket()
+    try:
      if not connectWithTimeout(sock, replica.host, Port(replica.port), 1000):
-        replica.connected = false
-        replica.state = rsDisconnected
+        connected = false
      else:
+        defer: sock.close()
        sock.send("PING\n")
        var response = ""
        try:
          sock.readLine(response)
          if response.strip() != "PONG":
-            replica.connected = false
-            replica.state = rsDisconnected
+            connected = false
        except:
-          replica.connected = false
-          replica.state = rsDisconnected
-        sock.close()
-  release(rm.lock)
+          connected = false
+    except:
+      connected = false
+    finally:
+      sock.close()
+
+    if not connected:
+      acquire(rm.lock)
+      if id in rm.replicas:
+        rm.replicas[id].connected = false
+        rm.replicas[id].state = rsDisconnected
+      release(rm.lock)

 proc reconnectReplica*(rm: ReplicationManager, id: string): bool =
-  acquire(rm.lock)
  result = false
+  var replica: Replica
+  var found = false
+  acquire(rm.lock)
  if id in rm.replicas:
-    let replica = rm.replicas[id]
-    if not replica.connected and replica.host.len > 0 and replica.port > 0:
-      var sock = newSocket()
-      if connectWithTimeout(sock, replica.host, Port(replica.port), 2000):
-        replica.connected = true
-        replica.state = rsStreaming
-        replica.lastSeen = getMonoTime().ticks()
-        result = true
-      sock.close()
+    replica = rm.replicas[id]
+    found = true
  release(rm.lock)
+  if not found: return false
+  if replica.connected or replica.host.len == 0 or replica.port == 0: return false
+
+  var sock = newSocket()
+  defer: sock.close()
+  if connectWithTimeout(sock, replica.host, Port(replica.port), 2000):
+    acquire(rm.lock)
+    if id in rm.replicas:
+      rm.replicas[id].connected = true
+      rm.replicas[id].state = rsStreaming
+      rm.replicas[id].lastSeen = getMonoTime().ticks()
+      result = true
+    release(rm.lock)

 proc reconnectAll*(rm: ReplicationManager): int =
  acquire(rm.lock)
-  result = 0
+  var candidates: seq[(string, Replica)] = @[]
  for id, replica in rm.replicas:
    if not replica.connected and replica.host.len > 0 and replica.port > 0:
-      var sock = newSocket()
-      if connectWithTimeout(sock, replica.host, Port(replica.port), 2000):
-        replica.connected = true
-        replica.state = rsStreaming
-        replica.lastSeen = getMonoTime().ticks()
-        inc result
-      sock.close()
+      candidates.add((id, replica))
  release(rm.lock)

+  result = 0
+  for (id, replica) in candidates:
+    var sock = newSocket()
+    defer: sock.close()
+    if connectWithTimeout(sock, replica.host, Port(replica.port), 2000):
+      acquire(rm.lock)
+      if id in rm.replicas:
+        rm.replicas[id].connected = true
+        rm.replicas[id].state = rsStreaming
+        rm.replicas[id].lastSeen = getMonoTime().ticks()
+      release(rm.lock)
+      inc result
+
 proc startHealthCheck*(rm: ReplicationManager, intervalMs: int = 5000) {.async.} =
  while true:
    await sleepAsync(intervalMs)
@@ -317,6 +317,8 @@ proc recvExactWithTimeout(client: AsyncSocket, size: int, timeoutMs: int): Futur
  let ok = await withTimeout(fut, timeoutMs)
  if ok:
    return fut.read()
+  # Timeout: caller will close the socket, which cancels the pending recv
+  return ""

 proc slowQueryLog(logPath: string, query: string, durationMs: int, clientId: int) =
  if logPath.len == 0:
@@ -505,6 +507,7 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
            let info = getDatabaseInfo(server.registry, jwtDatabase)
            if info != nil:
              let targetCtx = cast[ExecutionContext](cast[pointer](info.ctx))
+              let oldDb = connCtx.currentDatabase
              connCtx.db = info.db
              connCtx.tables = targetCtx.tables
              connCtx.btrees = targetCtx.btrees
@@ -517,6 +520,8 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
              connCtx.autoIncCounters = targetCtx.autoIncCounters
              connCtx.sequences = targetCtx.sequences
              connCtx.currentDatabase = jwtDatabase
+              if oldDb.len > 0 and oldDb != jwtDatabase:
+                decrementConnections(server.registry, oldDb)
              incrementConnections(server.registry, jwtDatabase)
        else:
          let err = makeErrorMessage(header.requestId, 403, "Invalid token")
@@ -64,12 +64,16 @@ proc hashKey*(key: string): uint64 =
  return uint64(hash(key))

 proc getShardHash*(router: ShardRouter, key: string): int =
+  if router.shards.len == 0: return -1
  let h = hashKey(key)
  return int(h mod uint64(router.shards.len))

 proc getShardRange*(router: ShardRouter, key: string): int =
  for i, shard in router.shards:
-    if key >= shard.minKey and key <= shard.maxKey:
+    let lastShard = (i == router.shards.len - 1)
+    if key >= shard.minKey and (lastShard or key < shard.maxKey):
+      return i
+    if lastShard and key == shard.maxKey:
      return i
  return -1  # key outside all defined ranges

@@ -175,18 +179,22 @@ proc sendMigrationBatch(host: string, port: int, shardId: int,
    return false

 proc migrateData*(router: var ShardRouter, nodes: seq[string],
-                  nodeAddrs: Table[string, tuple[host: string, port: int]]) =
+                  nodeAddrs: Table[string, tuple[host: string, port: int]],
+                  oldAssignments: seq[seq[string]] = @[]) =
  ## Migrate data when shard assignments change.
  ## Moves keys out of shards that are no longer owned by local node.
+  ## oldAssignments[i] contains the previous nodeIds for shard i.
  if router.iterateKeys == nil or router.storeKeys == nil:
    return

  if router.localNodeId.len == 0:
    return

-  for shard in router.shards.mitems:
+  for i in 0..<router.shards.len:
+    let shard = router.shards[i]
+    let wasOwner = oldAssignments.len > i and router.localNodeId in oldAssignments[i]
    let isOwner = router.localNodeId in shard.nodeIds
-    if not isOwner:
+    if wasOwner and not isOwner:
      # We previously owned this shard but no longer do — ship data to new owner
      let entries = router.iterateKeys(shard.id)
      if entries.len > 0:
@@ -201,9 +209,10 @@ proc migrateData*(router: var ShardRouter, nodes: seq[string],
                router.deleteKeys(keys)
            break

-proc rebalance*(router: var ShardRouter, nodes: seq[string]) =
+proc rebalance*(router: var ShardRouter, nodes: seq[string]): seq[seq[string]] =
+  ## Rebalance shard assignments across nodes. Returns old assignments for migration.
  if nodes.len == 0:
-    return
+    return @[]

  # Remember old assignments for migration
  var oldAssignments: seq[seq[string]] = @[]
@@ -220,6 +229,8 @@ proc rebalance*(router: var ShardRouter, nodes: seq[string]) =
      let nodeIdx = (i + r) mod nodes.len
      router.shards[i].nodeIds.add(nodes[nodeIdx])

+  return oldAssignments
+
 proc applyMigrationBatch*(router: var ShardRouter, shardId: int,
                          entries: seq[(string, seq[byte])]) =
  if router.storeKeys != nil:
@@ -261,10 +272,10 @@ proc addNode*(cm: ClusterMembership, nodeId: string,
  if host.len > 0:
    cm.nodeAddrs[nodeId] = (host, port)
  if cm.nodes.len >= 2:
-    cm.router.rebalance(cm.nodes)
+    let oldAssignments = cm.router.rebalance(cm.nodes)
    # Migrate data if we have migration callbacks and node addresses
    if cm.router.iterateKeys != nil:
-      cm.router.migrateData(cm.nodes, cm.nodeAddrs)
+      cm.router.migrateData(cm.nodes, cm.nodeAddrs, oldAssignments)

 proc removeNode*(cm: ClusterMembership, nodeId: string) =
  var newNodes: seq[string] = @[]
@@ -274,7 +285,9 @@ proc removeNode*(cm: ClusterMembership, nodeId: string) =
  cm.nodes = newNodes
  cm.nodeAddrs.del(nodeId)
  if cm.nodes.len >= 1:
-    cm.router.rebalance(cm.nodes)
+    let oldAssignments = cm.router.rebalance(cm.nodes)
+    if cm.router.iterateKeys != nil:
+      cm.router.migrateData(cm.nodes, cm.nodeAddrs, oldAssignments)

 proc onNodeJoin*(cm: ClusterMembership, nodeId: string,
                 host: string = "", port: int = 0) =