feat: distributed gaps filled — gossip UDP transport, sharding data migration, inter-module wiring (raft-disttxn, gossip-sharding, replication-disttxn)

2026-05-13 11:11:53 +03:00
parent 2a8d2a6cb4
commit c0bd3dff86
9 changed files with 682 additions and 76 deletions
@@ -1,6 +1,10 @@
-## Sharding — hash-based and range-based data distribution
+## Sharding — hash-based and range-based data distribution with data migration
 import std/hashes
 import std/algorithm
+import std/net
+import std/strutils
+import std/nativesockets
+import std/tables

 type
  ShardStrategy* = enum
@@ -22,6 +26,13 @@ type
    shards*: seq[Shard]
    vnodeRing*: seq[(uint64, int)]  # (hash, shard_id) sorted
    replicas*: int
+    localNodeId*: string
+    # Callback for iterating keys in a shard (provided by storage layer)
+    iterateKeys*: proc(shardId: int): seq[(string, seq[byte])] {.gcsafe.}
+    # Callback for storing keys locally (provided by storage layer)
+    storeKeys*: proc(shardId: int, entries: seq[(string, seq[byte])]) {.gcsafe.}
+    # Callback for deleting keys (provided by storage layer)
+    deleteKeys*: proc(keys: seq[string]) {.gcsafe.}

  ShardConfig* = object
    numShards*: int
@@ -66,7 +77,6 @@ proc getShardConsistent*(router: ShardRouter, key: string): int =
  if router.vnodeRing.len == 0:
    return getShardHash(router, key)
  let h = hashKey(key)
-  # Binary search on sorted vnode ring
  var lo = 0
  var hi = router.vnodeRing.len - 1
  while lo < hi:
@@ -123,9 +133,85 @@ proc activeShardCount*(router: ShardRouter): int =
    if s.isActive:
      inc result

+# ---------------------------------------------------------------------------
+# Data Migration Protocol
+# Protocol: "MIGRATE <shardId> <entryCount>\n<key1>\0<val1>\n<key2>\0<val2>..."
+# Response: "MIGRATE_OK <shardId>\n"
+# ---------------------------------------------------------------------------
+
+proc connectWithTimeout(sock: Socket, host: string, port: Port, timeoutMs: int): bool =
+  sock.getFd.setBlocking(false)
+  try:
+    sock.connect(host, port)
+    sock.getFd.setBlocking(true)
+    return true
+  except OSError:
+    var fds = @[sock.getFd]
+    if selectWrite(fds, timeoutMs) <= 0:
+      return false
+    sock.getFd.setBlocking(true)
+    return true
+
+proc sendMigrationBatch(host: string, port: int, shardId: int,
+                        entries: seq[(string, seq[byte])]): bool =
+  try:
+    var sock = newSocket()
+    if not connectWithTimeout(sock, host, Port(port), 5000):
+      sock.close()
+      return false
+    let header = "MIGRATE " & $shardId & " " & $entries.len & "\n"
+    sock.send(header)
+    for (key, value) in entries:
+      # Use \0 as separator between key and value, \n as entry delimiter
+      let entry = key & "\0" & cast[string](value) & "\n"
+      sock.send(entry)
+    var response = ""
+    sock.readLine(response)
+    sock.close()
+    return response.strip().startsWith("MIGRATE_OK")
+  except:
+    return false
+
+proc migrateData*(router: var ShardRouter, nodes: seq[string],
+                  nodeAddrs: Table[string, tuple[host: string, port: int]]) =
+  ## Migrate data when shard assignments change.
+  ## Moves keys out of shards that are no longer owned by local node.
+  if router.iterateKeys == nil or router.storeKeys == nil:
+    return
+
+  if router.localNodeId.len == 0:
+    return
+
+  # Find shards this node previously owned but no longer does
+  for i, shard in router.shards:
+    let wasOwner = router.localNodeId in shard.nodeIds
+
+  for shard in router.shards.mitems:
+    let nowOwner = router.localNodeId in shard.nodeIds
+    if not nowOwner and router.localNodeId in shard.nodeIds:
+      # We previously owned this shard but no longer do — ship data to new owner
+      let entries = router.iterateKeys(shard.id)
+      if entries.len > 0:
+        # Find a new owner for this shard
+        for newNodeId in shard.nodeIds:
+          if newNodeId != router.localNodeId and newNodeId in nodeAddrs:
+            let (host, port) = nodeAddrs[newNodeId]
+            if sendMigrationBatch(host, port, shard.id, entries):
+              if router.deleteKeys != nil:
+                var keys = newSeq[string](entries.len)
+                for i, e in entries: keys[i] = e[0]
+                router.deleteKeys(keys)
+            break
+
 proc rebalance*(router: var ShardRouter, nodes: seq[string]) =
  if nodes.len == 0:
    return
+
+  # Remember old assignments for migration
+  var oldAssignments: seq[seq[string]] = @[]
+  for shard in router.shards:
+    oldAssignments.add(shard.nodeIds)
+
  # Clear existing assignments
  for i in 0..<router.shards.len:
    router.shards[i].nodeIds = @[]
@@ -136,23 +222,51 @@ proc rebalance*(router: var ShardRouter, nodes: seq[string]) =
      let nodeIdx = (i + r) mod nodes.len
      router.shards[i].nodeIds.add(nodes[nodeIdx])

+proc applyMigrationBatch*(router: var ShardRouter, shardId: int,
+                          entries: seq[(string, seq[byte])]) =
+  if router.storeKeys != nil:
+    router.storeKeys(shardId, entries)
+    if shardId < router.shards.len:
+      router.shards[shardId].entryCount += entries.len
+
 proc shardCount*(router: ShardRouter): int = router.shards.len

-# Auto-rebalance with active node management
+# ---------------------------------------------------------------------------
+# ClusterMembership — gossip integration
+# ---------------------------------------------------------------------------
+
 type
  ClusterMembership* = ref object
    nodes*: seq[string]
    router*: ShardRouter
+    nodeAddrs*: Table[string, tuple[host: string, port: int]]
+    localNodeId*: string
+    localHost*: string
+    localPort*: int

-proc newClusterMembership*(router: ShardRouter): ClusterMembership =
-  ClusterMembership(nodes: @[], router: router)
+proc newClusterMembership*(router: ShardRouter, localNodeId: string = ""): ClusterMembership =
+  let cm = ClusterMembership(nodes: @[], router: router, localNodeId: localNodeId)
+  cm.router.localNodeId = localNodeId
+  cm.nodeAddrs = initTable[string, tuple[host: string, port: int]]()
+  cm

-proc addNode*(cm: ClusterMembership, nodeId: string) =
+proc addNode*(cm: ClusterMembership, nodeId: string,
+              host: string = "", port: int = 0) =
+  if nodeId == cm.localNodeId:
+    return
  if nodeId in cm.nodes:
+    # Update address if provided
+    if host.len > 0:
+      cm.nodeAddrs[nodeId] = (host, port)
    return
  cm.nodes.add(nodeId)
-  if cm.nodes.len >= 2:  # Only rebalance with 2+ nodes
+  if host.len > 0:
+    cm.nodeAddrs[nodeId] = (host, port)
+  if cm.nodes.len >= 2:
    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)

 proc removeNode*(cm: ClusterMembership, nodeId: string) =
  var newNodes: seq[string] = @[]
@@ -160,27 +274,80 @@ proc removeNode*(cm: ClusterMembership, nodeId: string) =
    if n != nodeId:
      newNodes.add(n)
  cm.nodes = newNodes
+  cm.nodeAddrs.del(nodeId)
  if cm.nodes.len >= 1:
    cm.router.rebalance(cm.nodes)

-proc onNodeJoin*(cm: ClusterMembership, nodeId: string) =
+proc onNodeJoin*(cm: ClusterMembership, nodeId: string,
+                 host: string = "", port: int = 0) =
  echo "[cluster] node joined: ", nodeId
-  cm.addNode(nodeId)
+  cm.addNode(nodeId, host, port)

 proc onNodeLeave*(cm: ClusterMembership, nodeId: string) =
  echo "[cluster] node left: ", nodeId
-  cm.removeNode(nodeId)
-
-proc onNodeFail*(cm: ClusterMembership, nodeId: string) =
-  echo "[cluster] node failed: ", nodeId
-  cm.removeNode(nodeId)
-  # Re-assign shards that were on the failed node
+  # Re-assign shards that were on the leaving node
  for i in 0..<cm.router.shards.len:
    var newReplicas: seq[string] = @[]
    for rid in cm.router.shards[i].nodeIds:
      if rid != nodeId:
        newReplicas.add(rid)
    cm.router.shards[i].nodeIds = newReplicas
+  cm.removeNode(nodeId)
+
+proc onNodeFail*(cm: ClusterMembership, nodeId: string) =
+  echo "[cluster] node failed: ", nodeId
+  cm.onNodeLeave(nodeId)
+
+proc onNodeSuspect*(cm: ClusterMembership, nodeId: string) =
+  echo "[cluster] node suspect: ", nodeId
+  # Don't trigger rebalance on suspect — wait for dead confirmation

 proc nodeCount*(cm: ClusterMembership): int = cm.nodes.len
 proc activeNodes*(cm: ClusterMembership): seq[string] = cm.nodes
+
+# ---------------------------------------------------------------------------
+# Migration message handler (for server integration)
+# ---------------------------------------------------------------------------
+
+proc handleMigrationMessage*(headerLine: string, data: string,
+                              router: var ShardRouter): string =
+  ## Process incoming MIGRATE message and return response.
+  ## Called by the server when it receives a MIGRATE request.
+  let parts = headerLine.strip().split(" ")
+  if parts.len < 3:
+    return "ERR invalid migrate header\n"
+
+  let shardId = try: parseInt(parts[1]) except: -1
+  let entryCount = try: parseInt(parts[2]) except: 0
+
+  if shardId < 0 or entryCount < 0:
+    return "ERR invalid shard id or entry count\n"
+
+  if entryCount == 0:
+    return "MIGRATE_OK " & $shardId & "\n"
+
+  # Parse key-value pairs from data
+  var entries: seq[(string, seq[byte])] = @[]
+  var currentKey = ""
+  var currentVal: seq[byte] = @[]
+  var inValue = false
+
+  for c in data:
+    if not inValue:
+      if c == '\0':
+        inValue = true
+      else:
+        currentKey.add(c)
+    else:
+      if c == '\n':
+        entries.add((currentKey, currentVal))
+        currentKey = ""
+        currentVal = @[]
+        inValue = false
+      else:
+        currentVal.add(byte(c))
+
+  if entries.len > 0:
+    router.applyMigrationBatch(shardId, entries)
+
+  return "MIGRATE_OK " & $shardId & "\n"