dimgigov
42c675224c
- disttxn: rollback now correctly updates participant state (was modifying a copy) - disttxn: bare except replaced with CatchableError - sharding: migrateData condition was always false, now correctly migrates data - sharding: unsafe cast[string] replaced with manual byte-to-char conversion - sharding: bare except replaced with CatchableError - server: parseUInt replaced with parseBiggestUint for uint64 portability
352 lines
11 KiB
Nim
352 lines
11 KiB
Nim
## 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
|
|
ssHash = "hash"
|
|
ssRange = "range"
|
|
ssConsistent = "consistent"
|
|
|
|
Shard* = object
|
|
id*: int
|
|
name*: string
|
|
minKey*: string
|
|
maxKey*: string
|
|
nodeIds*: seq[string] # replica node ids
|
|
isActive*: bool
|
|
entryCount*: int
|
|
|
|
ShardRouter* = ref object
|
|
strategy*: ShardStrategy
|
|
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
|
|
replicas*: int
|
|
strategy*: ShardStrategy
|
|
|
|
proc defaultShardConfig*(): ShardConfig =
|
|
ShardConfig(numShards: 4, replicas: 1, strategy: ssHash)
|
|
|
|
proc newShardRouter*(config: ShardConfig = defaultShardConfig()): ShardRouter =
|
|
result = ShardRouter(
|
|
strategy: config.strategy,
|
|
shards: @[],
|
|
vnodeRing: @[],
|
|
replicas: config.replicas,
|
|
)
|
|
for i in 0..<config.numShards:
|
|
result.shards.add(Shard(
|
|
id: i,
|
|
name: "shard_" & $i,
|
|
minKey: "",
|
|
maxKey: "",
|
|
nodeIds: @[],
|
|
isActive: true,
|
|
entryCount: 0,
|
|
))
|
|
|
|
proc hashKey*(key: string): uint64 =
|
|
return uint64(hash(key))
|
|
|
|
proc getShardHash*(router: ShardRouter, key: string): int =
|
|
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:
|
|
return i
|
|
return -1 # key outside all defined ranges
|
|
|
|
proc getShardConsistent*(router: ShardRouter, key: string): int =
|
|
if router.vnodeRing.len == 0:
|
|
return getShardHash(router, key)
|
|
let h = hashKey(key)
|
|
var lo = 0
|
|
var hi = router.vnodeRing.len - 1
|
|
while lo < hi:
|
|
let mid = (lo + hi) div 2
|
|
if router.vnodeRing[mid][0] < h:
|
|
lo = mid + 1
|
|
else:
|
|
hi = mid
|
|
if lo < router.vnodeRing.len and h <= router.vnodeRing[lo][0]:
|
|
return router.vnodeRing[lo][1]
|
|
return router.vnodeRing[0][1]
|
|
|
|
proc getShard*(router: ShardRouter, key: string): int =
|
|
case router.strategy
|
|
of ssHash: router.getShardHash(key)
|
|
of ssRange: router.getShardRange(key)
|
|
of ssConsistent: router.getShardConsistent(key)
|
|
|
|
proc addVirtualNodes*(router: var ShardRouter, numVnodes: int = 100) =
|
|
for shard in router.shards:
|
|
for i in 0..<numVnodes:
|
|
let vnodeKey = shard.name & "_vnode_" & $i
|
|
router.vnodeRing.add((hashKey(vnodeKey), shard.id))
|
|
router.vnodeRing.sort(proc(a, b: (uint64, int)): int = cmp(a[0], b[0]))
|
|
|
|
proc setRangeBounds*(router: var ShardRouter, bounds: seq[(string, string)]) =
|
|
for i, bound in bounds:
|
|
if i < router.shards.len:
|
|
router.shards[i].minKey = bound[0]
|
|
router.shards[i].maxKey = bound[1]
|
|
|
|
proc assignNode*(router: var ShardRouter, shardId: int, nodeId: string) =
|
|
if shardId < router.shards.len:
|
|
router.shards[shardId].nodeIds.add(nodeId)
|
|
|
|
proc getShardForNode*(router: ShardRouter, nodeId: string): seq[int] =
|
|
result = @[]
|
|
for i, shard in router.shards:
|
|
if nodeId in shard.nodeIds:
|
|
result.add(i)
|
|
|
|
proc replicasOf*(router: ShardRouter, key: string): seq[string] =
|
|
let shardId = router.getShard(key)
|
|
if shardId >= 0 and shardId < router.shards.len:
|
|
return router.shards[shardId].nodeIds
|
|
return @[]
|
|
|
|
proc allShards*(router: ShardRouter): seq[Shard] =
|
|
return router.shards
|
|
|
|
proc activeShardCount*(router: ShardRouter): int =
|
|
result = 0
|
|
for s in router.shards:
|
|
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
|
|
var valStr = newString(value.len)
|
|
for i, b in value: valStr[i] = char(b)
|
|
let entry = key & "\0" & valStr & "\n"
|
|
sock.send(entry)
|
|
var response = ""
|
|
sock.readLine(response)
|
|
sock.close()
|
|
return response.strip().startsWith("MIGRATE_OK")
|
|
except CatchableError:
|
|
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
|
|
|
|
for shard in router.shards.mitems:
|
|
let isOwner = router.localNodeId in shard.nodeIds
|
|
if 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:
|
|
# 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 = @[]
|
|
|
|
# Round-robin assign replicas
|
|
for i in 0..<router.shards.len:
|
|
for r in 0..<router.replicas:
|
|
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
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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, 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,
|
|
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 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] = @[]
|
|
for n in cm.nodes:
|
|
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,
|
|
host: string = "", port: int = 0) =
|
|
echo "[cluster] node joined: ", nodeId
|
|
cm.addNode(nodeId, host, port)
|
|
|
|
proc onNodeLeave*(cm: ClusterMembership, nodeId: string) =
|
|
echo "[cluster] node left: ", nodeId
|
|
# 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"
|