From 87b69a318bfacb8a08046c3b100fa8f3ce351730 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 30 Mar 2026 15:21:30 -0700 Subject: [PATCH] feat: implement eager migration in rebalancer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When rebalancer.eager_migrate is enabled, the rebalancer POSTs /_internal/migrate to source nodes after reassigning buckets, triggering immediate workstream eviction instead of waiting for lazy resume on the next request. Only idle workstreams are eagerly migrated — active ones (running, thinking, attention) are left alone to avoid disrupting in-flight work. Failed migrations are logged and skipped (the lazy path handles them eventually). --- tests/test_rebalancer.py | 133 ++++++++++++++++++++++++++++++++ turnstone/console/rebalancer.py | 85 ++++++++++++++++++++ turnstone/console/server.py | 2 + 3 files changed, 220 insertions(+) diff --git a/tests/test_rebalancer.py b/tests/test_rebalancer.py index fa37c208..99bfd5f3 100644 --- a/tests/test_rebalancer.py +++ b/tests/test_rebalancer.py @@ -339,3 +339,136 @@ class TestGetStatus: assert status["version"] == 1 assert status["last_result"] is not None assert status["last_result"]["seeded"] is True + + +class TestEagerMigration: + def test_eager_migrate_posts_to_source_nodes(self, storage): + """When eager_migrate=True, rebalancer POSTs /_internal/migrate for idle workstreams.""" + import httpx + + # Seed ring with 2 nodes + _register_nodes(storage, 2) + rb = Rebalancer(storage=storage, eager_migrate=True) + rb.rebalance_once() # seeds + + # Create a workstream on node-0's bucket range + # Find a bucket assigned to node-0 + buckets = storage.list_ring_buckets() + node0_bucket = None + for b in buckets: + if b["node_id"] == "node-0": + node0_bucket = b["bucket"] + break + assert node0_bucket is not None + + ws_id = f"{node0_bucket:04x}" + "a" * 28 + storage.register_workstream(ws_id, node_id="node-0", name="test") + storage.increment_bucket_count(node0_bucket) + + # Add a 3rd node — this will trigger rebalance + meta = json.dumps({"weight": 1, "started": "2026-01-01T00:00:00Z"}) + storage.register_service("server", "node-2", "http://node-2:8080", metadata=meta) + + # Track migrate calls + migrate_calls: list[tuple[str, str]] = [] # (url, ws_id) + + class FakeTransport(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + migrate_calls.append((str(request.url), body.get("ws_id", ""))) + return httpx.Response(200, json={"status": "ok", "ws_id": body["ws_id"]}) + + # Monkey-patch httpx.Client to use our fake transport + original_init = httpx.Client.__init__ + + def patched_init(self_client, **kwargs): + kwargs["transport"] = FakeTransport() + original_init(self_client, **kwargs) + + import unittest.mock + + with unittest.mock.patch.object(httpx.Client, "__init__", patched_init): + result = rb.rebalance_once(trigger="test") + + # If the bucket moved to a different node, the workstream should be migrated + new_buckets = storage.list_ring_buckets() + new_owner = None + for b in new_buckets: + if b["bucket"] == node0_bucket: + new_owner = b["node_id"] + break + + if new_owner != "node-0": + # Bucket moved — migration should have happened + assert result.migrations > 0 + assert any(ws_id in call[1] for call in migrate_calls) + else: + # Bucket stayed — no migration needed for this ws + assert result.migrations >= 0 # other workstreams might have been migrated + + def test_eager_migrate_skips_active_workstreams(self, storage): + """Active workstreams are not eagerly migrated (would disrupt in-flight work).""" + import httpx + + _register_nodes(storage, 2) + rb = Rebalancer(storage=storage, eager_migrate=True) + rb.rebalance_once() # seeds + + # Find a bucket on node-0 + buckets = storage.list_ring_buckets() + node0_bucket = None + for b in buckets: + if b["node_id"] == "node-0": + node0_bucket = b["bucket"] + break + assert node0_bucket is not None + + # Create an ACTIVE workstream (state="running") + ws_id = f"{node0_bucket:04x}" + "b" * 28 + storage.register_workstream(ws_id, node_id="node-0", name="active-ws") + storage.update_workstream_state(ws_id, "running") + storage.increment_bucket_count(node0_bucket, active=True) + + # Add 3rd node to trigger rebalance + meta = json.dumps({"weight": 1, "started": "2026-01-01T00:00:00Z"}) + storage.register_service("server", "node-2", "http://node-2:8080", metadata=meta) + + migrate_calls: list[str] = [] + + class FakeTransport(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + migrate_calls.append(body.get("ws_id", "")) + return httpx.Response(200, json={"status": "ok"}) + + original_init = httpx.Client.__init__ + + def patched_init(self_client, **kwargs): + kwargs["transport"] = FakeTransport() + original_init(self_client, **kwargs) + + import unittest.mock + + with unittest.mock.patch.object(httpx.Client, "__init__", patched_init): + rb.rebalance_once(trigger="test") + + # The active workstream should NOT have been migrated + assert ws_id not in migrate_calls + + def test_eager_migrate_disabled_by_default(self, storage): + """When eager_migrate=False (default), no migrate calls happen.""" + _register_nodes(storage, 2) + rb = Rebalancer(storage=storage) # eager_migrate defaults to False + rb.rebalance_once() # seeds + + # Create workstream and trigger rebalance + buckets = storage.list_ring_buckets() + node0_bucket = next(b["bucket"] for b in buckets if b["node_id"] == "node-0") + ws_id = f"{node0_bucket:04x}" + "c" * 28 + storage.register_workstream(ws_id, node_id="node-0", name="test") + + 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(trigger="test") + assert result.migrations == 0 # no eager migration when disabled diff --git a/turnstone/console/rebalancer.py b/turnstone/console/rebalancer.py index 08edb0e8..a353d9c0 100644 --- a/turnstone/console/rebalancer.py +++ b/turnstone/console/rebalancer.py @@ -36,6 +36,7 @@ class RebalanceResult: """Summary of a single rebalance pass.""" moves: int = 0 + migrations: int = 0 trigger: str = "periodic" duration_ms: float = 0.0 nodes: int = 0 @@ -59,6 +60,8 @@ class Rebalancer: threshold: float = 0.10, vnodes_per_unit: int = 150, lock_ttl: int = 120, + eager_migrate: bool = False, + api_token: str = "", ) -> None: self._storage = storage self._router = router @@ -67,6 +70,8 @@ class Rebalancer: self._threshold = threshold self._vnodes_per_unit = vnodes_per_unit self._lock_ttl = lock_ttl + self._eager_migrate = eager_migrate + self._api_token = api_token self._stop_event = threading.Event() self._trigger_event = threading.Event() self._thread: threading.Thread | None = None @@ -332,13 +337,23 @@ class Rebalancer: if self._router is not None: self._router.refresh_cache() + # 13. Eager migration: evict workstreams on moved buckets from source nodes + migrations = 0 + if self._eager_migrate and filtered_moves: + migrations = self._eager_migrate_workstreams( + filtered_moves, + nodes_raw, + ) + result.moves = total_moved + result.migrations = migrations result.noop = False result.duration_ms = (time.monotonic() - t0) * 1000 log.info( "rebalancer.rebalanced", moves=total_moved, + migrations=migrations, nodes=len(ring_nodes), trigger=trigger, duration_ms=round(result.duration_ms, 1), @@ -441,6 +456,76 @@ class Rebalancer: if active_delta != 0: self._storage.adjust_bucket_active(bucket, active_delta) + def _eager_migrate_workstreams( + self, + moves: list[tuple[int, str, str]], + nodes_raw: list[dict[str, str]], + ) -> int: + """POST /_internal/migrate to source nodes for workstreams on moved buckets. + + Only migrates idle workstreams — active ones would be disrupted. + Returns the number of successful migrations. + """ + import httpx + + # Build node URL map from the services data already loaded + node_urls: dict[str, str] = {s["service_id"]: s["url"] for s in nodes_raw} + + # Moved buckets grouped by source node + moved_buckets: dict[str, set[int]] = defaultdict(set) + for bucket, from_node, _to_node in moves: + moved_buckets[from_node].add(bucket) + + # Find workstreams on moved buckets (idle only — don't disrupt active work) + ws_data = self._storage.list_workstream_routing_data() + to_migrate: list[tuple[str, str]] = [] # (ws_id, source_node_url) + for ws_id, state in ws_data: + if len(ws_id) < 4 or state in _ACTIVE_STATES: + continue + bucket = bucket_of(ws_id) + for node_id, buckets in moved_buckets.items(): + if bucket in buckets: + url = node_urls.get(node_id) + if url: + to_migrate.append((ws_id, url)) + break + + if not to_migrate: + return 0 + + headers: dict[str, str] = {} + if self._api_token: + headers["Authorization"] = f"Bearer {self._api_token}" + + migrated = 0 + with httpx.Client(timeout=10, headers=headers) as client: + for ws_id, source_url in to_migrate: + try: + resp = client.post( + f"{source_url}/v1/api/_internal/migrate", + json={"ws_id": ws_id}, + ) + if resp.status_code == 200: + migrated += 1 + elif resp.status_code == 409: + log.debug( + "rebalancer.migrate.refused", + ws_id=ws_id[:8], + reason="last_workstream", + ) + # 404 = already gone, that's fine + except httpx.HTTPError: + log.warning( + "rebalancer.migrate.failed", + ws_id=ws_id[:8], + source=source_url, + exc_info=True, + ) + + if migrated: + log.info("rebalancer.migrations", count=migrated, total=len(to_migrate)) + return migrated + def _build_ring_nodes(services: list[dict[str, str]]) -> list[RingNode]: """Convert service registry rows into RingNode instances.""" diff --git a/turnstone/console/server.py b/turnstone/console/server.py index d12379d4..30206627 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -6107,6 +6107,8 @@ def main() -> None: interval=_rcs.get("rebalancer.interval", 60), threshold=_rcs.get("rebalancer.threshold", 0.10), vnodes_per_unit=_rcs.get("ring.vnodes_per_unit", 150), + eager_migrate=_rcs.get("rebalancer.eager_migrate", False), + api_token=proxy_auth_token, ) log.info("rebalancer.configured") except Exception: