ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

feat: minimal-transfer rebalancer algorithm

Browse Source 
Replace the full-rehash algorithm (diff ideal vs current across all
65536 buckets) with a donor/recipient algorithm that only moves
buckets from overloaded nodes to underloaded nodes.

Key improvements:
- Adding node C to {A, B} only moves buckets TO C, never between
  A and B. Previously the HashRing rehash could shuffle between
  existing nodes.
- Seeding uses weight-proportional distribution instead of HashRing
  virtual nodes, producing an exact split that doesn't trigger
  immediate correction on the next cycle.
- Dead-node buckets are redistributed to the most underloaded
  survivors, not rehashed across the whole ring.
- HashRing class is no longer used by the rebalancer (still
  available for other uses like the Go rewrite reference).

The threshold check still gates live-to-live moves. Dead-node
recovery remains unconditional.
This commit is contained in:
Patrick Buckley
2026-03-30 15:55:29 -07:00
committed by Patrick Buckley
parent bd782f804e
commit c2750de7a4
2 changed files with 190 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_rebalancer.py
+67
View File
@@ -472,3 +472,70 @@ class TestEagerMigration:
result = rb.rebalance_once(trigger="test")
assert result.migrations == 0 # no eager migration when disabled
class TestMinimalTransfer:
def test_new_node_only_receives_never_shuffles(self, storage):
"""Adding a 3rd node moves buckets TO it, never between existing nodes.
This is the key property of the minimal-transfer algorithm: nodes A
and B should not exchange buckets with each other — only donate to C.
"""
_register_nodes(storage, 2)
rb = Rebalancer(storage=storage, threshold=0.05)
rb.rebalance_once() # seeds: node-0 gets 32768, node-1 gets 32768
# Record which node owns each bucket before adding node-2
before = {r["bucket"]: r["node_id"] for r in storage.list_ring_buckets()}
# Add a third node
meta = json.dumps({"weight": 1, "started": "2026-01-01T00:00:00Z"})
storage.register_service("server", "node-2", "http://node-2:8080", metadata=meta)
result = rb.rebalance_once()
after = {r["bucket"]: r["node_id"] for r in storage.list_ring_buckets()}
# Verify: every bucket that moved went TO node-2
for bucket in range(RING_SIZE):
old = before[bucket]
new = after[bucket]
if old != new:
assert new == "node-2", (
f"bucket {bucket} moved {old} -> {new}, "
"expected all moves to target node-2"
)
# Verify: node-2 got roughly 1/3 of all buckets
node2_count = sum(1 for nid in after.values() if nid == "node-2")
assert 19000 < node2_count < 24000, f"node-2 got {node2_count} buckets"
assert result.moves > 0
def test_remove_node_distributes_proportionally(self, storage):
"""Removing a node distributes its buckets to remaining nodes
proportionally — doesn't shuffle between survivors."""
_register_nodes(storage, 3)
rb = Rebalancer(storage=storage, threshold=0.05)
rb.rebalance_once() # seeds
before = {r["bucket"]: r["node_id"] for r in storage.list_ring_buckets()}
# Remove node-2
storage.deregister_service("server", "node-2")
result = rb.rebalance_once()
after = {r["bucket"]: r["node_id"] for r in storage.list_ring_buckets()}
# Every moved bucket should have been owned by node-2 (the dead node)
for bucket in range(RING_SIZE):
old = before[bucket]
new = after[bucket]
if old != new:
assert old == "node-2", (
f"bucket {bucket} moved {old} -> {new}, "
"but only node-2's buckets should move"
)
# node-2 should have zero buckets now
node2_count = sum(1 for nid in after.values() if nid == "node-2")
assert node2_count == 0
assert result.moves > 0
turnstone/console/rebalancer.py
+123 -57
View File
@@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any
import structlog
from turnstone.core.hash_ring import RING_SIZE, HashRing, RingNode, bucket_of
from turnstone.core.hash_ring import RING_SIZE, RingNode, bucket_of
if TYPE_CHECKING:
from turnstone.console.collector import ClusterCollector
@@ -251,8 +251,7 @@ class Rebalancer:
# 3. If table is empty — first run, seed all 65536 buckets
if not current_rows:
ring = HashRing(ring_nodes, vnodes_per_unit=self._vnodes_per_unit)
assignments = ring.assignments()
assignments = _weight_based_assignments(ring_nodes)
self._storage.seed_ring_buckets(assignments)
self._bump_version()
if self._router is not None:
@@ -267,74 +266,114 @@ class Rebalancer:
)
return result
# 4. Single node with all buckets assigned — noop
# 4. Build current assignment map and per-node bucket lists
current_map: dict[int, str] = {r["bucket"]: r["node_id"] for r in current_rows}
live_ids = {n.node_id for n in ring_nodes}
# Single node with all buckets assigned — noop
if len(live_ids) == 1 and all(nid in live_ids for nid in current_map.values()):
result.duration_ms = (time.monotonic() - t0) * 1000
return result
# 5. Compute ideal assignments
ring = HashRing(ring_nodes, vnodes_per_unit=self._vnodes_per_unit)
ideal = ring.assignments()
ideal_map: dict[int, str] = {b: nid for b, nid in ideal}
# 6. Find buckets that need to move
moves_needed: list[tuple[int, str, str]] = [] # (bucket, from_node, to_node)
for bucket in range(RING_SIZE):
cur = current_map.get(bucket)
want = ideal_map.get(bucket)
if cur is None or want is None:
continue
# Only move if current owner is dead OR ideal owner differs
if cur != want:
# If current owner is still alive and threshold applies,
# we check whether deviation is big enough to warrant moves
moves_needed.append((bucket, cur, want))
if not moves_needed:
result.duration_ms = (time.monotonic() - t0) * 1000
return result
# 7. Reconcile bucket_stats before computing transfer priority
# 5. Reconcile bucket_stats before computing transfer costs
self._reconcile_bucket_stats()
# 8. Load stats for transfer priority
stats_rows = self._storage.list_bucket_stats()
stats_map: dict[int, tuple[int, int]] = {} # bucket -> (ws_count, active_count)
stats_map: dict[int, tuple[int, int]] = {}
for s in stats_rows:
stats_map[s["bucket"]] = (s["ws_count"], s["active_count"])
# 9. Sort moves by cost: empty first, then idle, then active
moves_needed.sort(
key=lambda m: stats_map.get(m[0], (0, 0)),
)
# 6. Group buckets by current owner
buckets_by_node: dict[str, list[int]] = defaultdict(list)
for bucket, nid in current_map.items():
buckets_by_node[nid].append(bucket)
# 10. Check if deviation is significant enough to warrant rebalancing.
# Count buckets per node (current vs ideal) and check max deviation.
cur_counts: dict[str, int] = defaultdict(int)
ideal_counts: dict[str, int] = defaultdict(int)
for nid in current_map.values():
cur_counts[nid] += 1
for nid in ideal_map.values():
ideal_counts[nid] += 1
# 7. Compute ideal bucket count per node from weights
total_weight = sum(n.weight for n in ring_nodes)
ideal_counts: dict[str, int] = {}
remainder_pool: list[str] = []
assigned_ideal = 0
for n in ring_nodes:
ideal_n = int((n.weight / total_weight) * RING_SIZE)
ideal_counts[n.node_id] = ideal_n
assigned_ideal += ideal_n
remainder_pool.append(n.node_id)
# Distribute remainder buckets (rounding error) to heaviest nodes
leftover = RING_SIZE - assigned_ideal
remainder_pool.sort(key=lambda nid: ideal_counts[nid], reverse=True)
for i in range(leftover):
ideal_counts[remainder_pool[i % len(remainder_pool)]] += 1
# Always reassign buckets owned by dead nodes
dead_moves = [m for m in moves_needed if m[1] not in live_ids]
live_moves = [m for m in moves_needed if m[1] in live_ids]
# 8. Always reassign dead-node buckets first (unconditional)
dead_node_ids = {nid for nid in buckets_by_node if nid not in live_ids}
filtered_moves: list[tuple[int, str, str]] = [] # (bucket, from, to)
# For live-to-live moves, check threshold
filtered_moves: list[tuple[int, str, str]] = list(dead_moves)
if live_moves:
max_deviation = 0.0
for nid in live_ids:
ideal_n = ideal_counts.get(nid, 0)
actual_n = cur_counts.get(nid, 0)
if ideal_n > 0:
dev = abs(actual_n - ideal_n) / ideal_n
max_deviation = max(max_deviation, dev)
if max_deviation >= self._threshold:
filtered_moves.extend(live_moves)
if dead_node_ids:
# Dead nodes are implicit donors — all their buckets must move.
# Distribute to the most underloaded live nodes.
dead_buckets: list[int] = []
for nid in dead_node_ids:
dead_buckets.extend(buckets_by_node.pop(nid))
# Sort by cost (cheapest first)
dead_buckets.sort(key=lambda b: stats_map.get(b, (0, 0)))
# Assign to live nodes that are most below their ideal
for bucket in dead_buckets:
# Pick the node with the largest deficit
best = min(
live_ids,
key=lambda nid: len(buckets_by_node.get(nid, [])) - ideal_counts.get(nid, 0),
)
filtered_moves.append((bucket, "", best))
buckets_by_node[best].append(bucket)
# 9. Identify donors and recipients among live nodes
actual_counts = {nid: len(bkts) for nid, bkts in buckets_by_node.items()}
donors: list[str] = []
recipients: list[str] = []
for nid in live_ids:
actual = actual_counts.get(nid, 0)
ideal = ideal_counts.get(nid, 0)
if ideal > 0 and actual > ideal * (1 + self._threshold):
donors.append(nid)
elif ideal > 0 and actual < ideal * (1 - self._threshold):
recipients.append(nid)
# 10. Transfer from donors to recipients — minimal moves only
if donors and recipients:
# Sort donors by excess descending, recipients by deficit descending
donors.sort(key=lambda nid: actual_counts[nid] - ideal_counts[nid], reverse=True)
recipients.sort(
key=lambda nid: ideal_counts[nid] - actual_counts.get(nid, 0),
reverse=True,
)
for donor_id in donors:
donor_excess = len(buckets_by_node[donor_id]) - ideal_counts[donor_id]
if donor_excess <= 0:
continue
# Sort this donor's buckets by cost (cheapest to move first)
donor_buckets = sorted(
buckets_by_node[donor_id],
key=lambda b: stats_map.get(b, (0, 0)),
)
moved_from_donor = 0
for recipient_id in recipients:
recipient_deficit = ideal_counts[recipient_id] - len(
buckets_by_node.get(recipient_id, [])
)
if recipient_deficit <= 0:
continue
# Transfer min(donor_excess - moved, recipient_deficit) buckets
to_move = min(donor_excess - moved_from_donor, recipient_deficit)
for _ in range(to_move):
if not donor_buckets:
break
bucket = donor_buckets.pop(0)
filtered_moves.append((bucket, donor_id, recipient_id))
buckets_by_node[donor_id].remove(bucket)
buckets_by_node.setdefault(recipient_id, []).append(bucket)
moved_from_donor += 1
if moved_from_donor >= donor_excess:
break
if not filtered_moves:
result.duration_ms = (time.monotonic() - t0) * 1000
@@ -524,6 +563,33 @@ class Rebalancer:
return migrated
def _weight_based_assignments(nodes: list[RingNode]) -> list[tuple[int, str]]:
"""Compute bucket assignments proportional to node weights.
Distributes all 65536 buckets across nodes proportionally to their
weights, with deterministic rounding. Used for seeding — produces
an exact weight-proportional split that the donor/recipient
algorithm won't try to "correct" on the next run.
"""
total_weight = sum(n.weight for n in nodes)
assignments: list[tuple[int, str]] = []
# Sort nodes for determinism
sorted_nodes = sorted(nodes, key=lambda n: n.node_id)
bucket = 0
for i, node in enumerate(sorted_nodes):
if i == len(sorted_nodes) - 1:
# Last node gets the remainder (avoids rounding gaps)
count = RING_SIZE - bucket
else:
count = round((node.weight / total_weight) * RING_SIZE)
for _ in range(count):
if bucket >= RING_SIZE:
break
assignments.append((bucket, node.node_id))
bucket += 1
return assignments
def _build_ring_nodes(services: list[dict[str, str]]) -> list[RingNode]:
"""Convert service registry rows into RingNode instances."""
nodes: list[RingNode] = []