mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
refactor(routing): replace hash-ring rebalancer with rendezvous (HRW)… (#384)
* refactor(routing): replace hash-ring rebalancer with rendezvous (HRW) hashing Routing was a stored bucket table maintained by a central rebalancer daemon, which shared its liveness primitive (services.last_heartbeat) with the collector — when a heartbeat-fresh node went into a zombie HTTP-handler-broken state, neither the collector nor the rebalancer could self-correct, and the router kept directing traffic at it. Rendezvous hashing makes the route a pure function of (ws_id, live_services) so the heartbeat is the single source of truth and any liveness-eviction propagates to the next route call without a separate state-publication step. The rebalancer's central state has no analogue: the new router computes the per-key node winner on every call, the collector pushes membership updates into the router cache from its discovery thread, and per-route overrides survive on workstream_overrides. Eager workstream migration goes away; in-flight workstreams lazily rehydrate from storage on the new owner — already the dead-node behaviour. * fix(tools): describe rendezvous re-routing on spawn/inspect node_id The first pass overclaimed `node_id` "stays canonical for this workstream's lifetime" — under rendezvous routing the active owner re-derives per-call from live membership, so a node join/drop after spawn can shift it. Tool descriptions now say `node_id` is the spawn-time binding; subsequent ops re-route via rendezvous over the current live-node set; the new owner lazily rehydrates from shared storage; coordinators should re-read with inspect_workstream rather than caching the value.
This commit is contained in:
Patrick Buckley
@@ -90,7 +90,7 @@ Built-in tools for shell, files, search, web, memory, notifications, and autonom
|
||||
|
||||
**Single-node**: Client → Server (direct HTTP + SSE). No external dependencies beyond the database.
|
||||
|
||||
**Multi-node**: Client → Console (hash ring routing proxy) → Server nodes. The console maintains a 65536-entry bucket cache for O(1) workstream routing. A rebalancer daemon redistributes buckets when nodes join or leave.
|
||||
**Multi-node**: Client → Console (rendezvous routing proxy) → Server nodes. The console picks the target node for each workstream via rendezvous (HRW) hashing over the live service registry — pure function of `(ws_id, live_nodes)`, no stored bucket state, deterministic across readers. A node join or drop only re-routes the keys that score highest on the affected node.
|
||||
|
||||
| Component | Purpose |
|
||||
|-----------|---------|
|
||||
|
||||
@@ -2121,15 +2121,15 @@ turnstone_tool_calls_total{tool="read_file"} 3
|
||||
## Console Routing Proxy Endpoints
|
||||
|
||||
These endpoints are served by the console (`turnstone-console`) and proxy
|
||||
requests to the correct server node via the hash ring bucket cache. In
|
||||
multi-node deployments, clients (SDK, channel gateway) talk to the console
|
||||
instead of individual server nodes.
|
||||
requests to the correct server node via rendezvous (HRW) hashing over the
|
||||
live service registry. In multi-node deployments, clients (SDK, channel
|
||||
gateway) talk to the console instead of individual server nodes.
|
||||
|
||||
### `POST /v1/api/route/workstreams/new`
|
||||
|
||||
Create a workstream via hash-ring routing. The console generates the `ws_id`,
|
||||
routes to the assigned node, and includes `node_url` in the response for
|
||||
direct SSE connections.
|
||||
Create a workstream via rendezvous routing. The console generates the `ws_id`,
|
||||
routes to the rendezvous-selected node, and includes `node_url` in the
|
||||
response for direct SSE connections.
|
||||
|
||||
### `POST /v1/api/route/send`
|
||||
|
||||
@@ -2164,5 +2164,4 @@ Used by channel adapters to open direct SSE connections to the correct server no
|
||||
|
||||
Prometheus metrics for the console routing layer. Includes:
|
||||
`turnstone_router_requests_total`, `turnstone_router_request_duration_seconds`,
|
||||
`turnstone_ring_membership_size`, `turnstone_ring_version`,
|
||||
`turnstone_ring_rebalance_total`, `turnstone_ring_migrations_total`.
|
||||
`turnstone_router_membership_size`, `turnstone_router_refresh_total`.
|
||||
|
||||
@@ -1,35 +1,27 @@
|
||||
# Consistent Hash Ring — Reference Design
|
||||
|
||||
**Status**: Reference (not currently in the hot path)
|
||||
**Date**: 2026-03-30
|
||||
**Status**: Reference — alternative routing strategy
|
||||
|
||||
## Overview
|
||||
Live routing uses **rendezvous (HRW) hashing** in
|
||||
`turnstone/core/rendezvous.py` and `turnstone/console/router.py`. This
|
||||
document captures a vnode-ring approach as a reference for future
|
||||
evaluation if the cluster outgrows rendezvous's O(N)-per-route
|
||||
characteristic.
|
||||
|
||||
This document describes a consistent hash ring algorithm evaluated during
|
||||
the design of the direct HTTP transport routing system. The current
|
||||
implementation uses weight-proportional bucket assignment with a
|
||||
donor/recipient rebalancing algorithm (see the routing section of
|
||||
[../architecture.md](../architecture.md)).
|
||||
The consistent hash ring is documented here as a reference for future
|
||||
scalability work — if the cluster grows beyond the point where the
|
||||
weight-proportional approach is sufficient, the ring provides a
|
||||
proven alternative with stronger stability guarantees.
|
||||
The FNV-1a-32 hash function specified below is bit-identical to the
|
||||
hash used by the live rendezvous implementation; cross-language clients
|
||||
can rely on these test vectors.
|
||||
|
||||
## When to consider the ring approach
|
||||
## When the ring approach becomes interesting
|
||||
|
||||
The current weight-proportional seeding + donor/recipient rebalancer works
|
||||
well when:
|
||||
- Cluster size is moderate (< 50 nodes)
|
||||
- Nodes join/leave infrequently
|
||||
- The rebalancer runs centrally (in the console)
|
||||
The vnode ring becomes preferable to rendezvous hashing when:
|
||||
|
||||
The consistent hash ring becomes advantageous when:
|
||||
- Cluster size grows large (50+ nodes) and frequent membership changes
|
||||
cause the donor/recipient algorithm to churn
|
||||
- Decentralized routing is needed (each node computes the ring locally,
|
||||
no central console required)
|
||||
- Cross-language determinism is important (multiple implementations must
|
||||
agree on the same assignment without sharing state)
|
||||
- Cluster size grows large (50+ nodes) and the per-route O(N) hash
|
||||
computation becomes visible against downstream HTTP cost.
|
||||
- Decentralised routing is needed (each node computes the ring locally,
|
||||
no central console required).
|
||||
- A precomputed flat-array lookup is desired so the routing hot path
|
||||
avoids hashing entirely.
|
||||
|
||||
## Algorithm
|
||||
|
||||
@@ -134,16 +126,17 @@ class HashRing:
|
||||
# Precompute all 65536 bucket assignments
|
||||
```
|
||||
|
||||
## Comparison with current approach
|
||||
## Comparison with rendezvous (HRW) hashing
|
||||
|
||||
| Aspect | Weight-proportional (current) | Consistent hash ring |
|
||||
|--------|------------------------------|---------------------|
|
||||
| Seeding | Exact weight split, deterministic | Hash-based, ~3% variance |
|
||||
| Node addition | Donor/recipient moves only excess | Ring moves ~1/N buckets |
|
||||
| Node removal | Dead buckets → most underloaded | Ring redistributes to clockwise neighbors |
|
||||
| Cross-node churn | Zero (only donor→recipient) | Zero (ring stability guarantee) |
|
||||
| Decentralized | No (needs central rebalancer) | Yes (each node computes locally) |
|
||||
| Complexity | Simple weight arithmetic | Virtual node construction + bisect |
|
||||
| Aspect | Rendezvous (live) | Consistent hash ring (this doc) |
|
||||
|--------|-------------------|---------------------------------|
|
||||
| Per-route cost | O(N) hash computes | O(log V) bisect against precomputed array |
|
||||
| Seeding | None — pure function | Build vnode array on every membership change |
|
||||
| Node addition | Pure function moves ~1/N keys | Ring moves ~1/N buckets |
|
||||
| Node removal | Surviving nodes' keys unchanged | Surviving nodes' buckets unchanged |
|
||||
| Decentralised | Yes — pure function over services | Yes — each node computes locally |
|
||||
| Persistent state | None | None on the hot path; precomputed array in memory |
|
||||
| Complexity | ~20 LOC | Virtual-node construction + bisect |
|
||||
|
||||
## Test vectors
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ An MCP server that exposes tools for executing commands across a Turnstone clust
|
||||
|
||||
This server uses the Turnstone console SDK (`TurnstoneConsole`) for node discovery and routing, and `TurnstoneServer` for per-node SSE streaming. The dispatch flow for each command is:
|
||||
|
||||
1. **Route** — `TurnstoneConsole.route_create_workstream(target_node=..., auto_approve=True)` creates a workstream pinned to the target node via the console's hash-ring routing proxy, returning `ws_id` and `node_url`.
|
||||
1. **Route** — `TurnstoneConsole.route_create_workstream(target_node=..., auto_approve=True)` creates a workstream pinned to the target node via the console's rendezvous routing proxy, returning `ws_id` and `node_url`.
|
||||
2. **Execute** — `TurnstoneServer(node_url, token=...)` connects directly to the node's SSE stream using the same `TURNSTONE_API_TOKEN`. `send_and_wait(prompt, ws_id)` runs the command and the raw bash output is captured from the `ToolResultEvent` — bypassing the costly "agent reads output then re-generates output as completion tokens" round-trip.
|
||||
3. **Cleanup** — `TurnstoneConsole.route_close(ws_id)` closes the workstream.
|
||||
|
||||
|
||||
@@ -4286,7 +4286,7 @@
|
||||
},
|
||||
"/v1/api/route/workstreams/new": {
|
||||
"post": {
|
||||
"summary": "Create workstream via hash-ring routing proxy",
|
||||
"summary": "Create workstream via rendezvous routing proxy",
|
||||
"operationId": "v1_api_route_workstreams_new_post",
|
||||
"tags": [
|
||||
"Routing"
|
||||
|
||||
@@ -146,7 +146,7 @@ export class TurnstoneConsole extends BaseClient {
|
||||
// -- Routing proxy --------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Create a workstream via the console hash-ring router.
|
||||
* Create a workstream via the console rendezvous router.
|
||||
*
|
||||
* When `attachments` is non-empty the request is sent as
|
||||
* multipart/form-data and the console routes via `?ws_id=<hex>`
|
||||
|
||||
@@ -37,61 +37,22 @@ class TestRecordRoute:
|
||||
assert "turnstone_router_request_duration_seconds_sum" in text
|
||||
|
||||
|
||||
class TestRingInfo:
|
||||
"""Ring membership and version gauges."""
|
||||
class TestRouterInfo:
|
||||
"""Live-membership gauge + refresh counter."""
|
||||
|
||||
def test_defaults_zero(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
text = m.generate_text()
|
||||
assert "turnstone_ring_membership_size 0" in text
|
||||
assert "turnstone_ring_version 0" in text
|
||||
assert "turnstone_router_membership_size 0" in text
|
||||
assert "turnstone_router_refresh_total 0" in text
|
||||
|
||||
def test_set_ring_info(self) -> None:
|
||||
def test_set_router_info(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
m.set_ring_info(3, 7)
|
||||
m.set_router_info(3, 7)
|
||||
|
||||
text = m.generate_text()
|
||||
assert "turnstone_ring_membership_size 3" in text
|
||||
assert "turnstone_ring_version 7" in text
|
||||
|
||||
|
||||
class TestRebalance:
|
||||
"""Rebalance and migration counters."""
|
||||
|
||||
def test_noop(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
m.record_rebalance("noop")
|
||||
|
||||
text = m.generate_text()
|
||||
assert 'turnstone_ring_rebalance_total{result="noop"} 1' in text
|
||||
|
||||
def test_seeded(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
m.record_rebalance("seeded")
|
||||
|
||||
text = m.generate_text()
|
||||
assert 'turnstone_ring_rebalance_total{result="seeded"} 1' in text
|
||||
|
||||
def test_rebalanced(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
m.record_rebalance("rebalanced")
|
||||
m.record_rebalance("rebalanced")
|
||||
|
||||
text = m.generate_text()
|
||||
assert 'turnstone_ring_rebalance_total{result="rebalanced"} 2' in text
|
||||
|
||||
def test_migrations(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
m.record_migrations(5)
|
||||
m.record_migrations(3)
|
||||
|
||||
text = m.generate_text()
|
||||
assert "turnstone_ring_migrations_total 8" in text
|
||||
|
||||
def test_migrations_default_zero(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
text = m.generate_text()
|
||||
assert "turnstone_ring_migrations_total 0" in text
|
||||
assert "turnstone_router_membership_size 3" in text
|
||||
assert "turnstone_router_refresh_total 7" in text
|
||||
|
||||
|
||||
class TestGenerateText:
|
||||
@@ -103,10 +64,8 @@ class TestGenerateText:
|
||||
expected = [
|
||||
"turnstone_router_requests_total",
|
||||
"turnstone_router_request_duration_seconds",
|
||||
"turnstone_ring_membership_size",
|
||||
"turnstone_ring_version",
|
||||
"turnstone_ring_rebalance_total",
|
||||
"turnstone_ring_migrations_total",
|
||||
"turnstone_router_membership_size",
|
||||
"turnstone_router_refresh_total",
|
||||
]
|
||||
for name in expected:
|
||||
assert name in text, f"Missing metric: {name}"
|
||||
@@ -116,8 +75,8 @@ class TestGenerateText:
|
||||
text = m.generate_text()
|
||||
assert "# HELP turnstone_router_requests_total" in text
|
||||
assert "# TYPE turnstone_router_requests_total counter" in text
|
||||
assert "# HELP turnstone_ring_membership_size" in text
|
||||
assert "# TYPE turnstone_ring_membership_size gauge" in text
|
||||
assert "# HELP turnstone_router_membership_size" in text
|
||||
assert "# TYPE turnstone_router_membership_size gauge" in text
|
||||
|
||||
def test_ends_with_newline(self) -> None:
|
||||
m = ConsoleMetrics()
|
||||
@@ -125,24 +84,16 @@ class TestGenerateText:
|
||||
assert text.endswith("\n")
|
||||
|
||||
def test_combined_scenario(self) -> None:
|
||||
"""Full scenario: routes, ring info, rebalances, migrations."""
|
||||
"""Full scenario: routes + router info."""
|
||||
m = ConsoleMetrics()
|
||||
m.record_route("create", 200, 0.1)
|
||||
m.record_route("send", 200, 0.05)
|
||||
m.record_route("send", 502, 1.2)
|
||||
m.set_ring_info(3, 12)
|
||||
m.record_rebalance("seeded")
|
||||
m.record_rebalance("noop")
|
||||
m.record_rebalance("rebalanced")
|
||||
m.record_migrations(4)
|
||||
m.set_router_info(3, 12)
|
||||
|
||||
text = m.generate_text()
|
||||
assert 'turnstone_router_requests_total{method="create",status="2xx"} 1' in text
|
||||
assert 'turnstone_router_requests_total{method="send",status="2xx"} 1' in text
|
||||
assert 'turnstone_router_requests_total{method="send",status="5xx"} 1' in text
|
||||
assert "turnstone_ring_membership_size 3" in text
|
||||
assert "turnstone_ring_version 12" in text
|
||||
assert 'turnstone_ring_rebalance_total{result="noop"} 1' in text
|
||||
assert 'turnstone_ring_rebalance_total{result="rebalanced"} 1' in text
|
||||
assert 'turnstone_ring_rebalance_total{result="seeded"} 1' in text
|
||||
assert "turnstone_ring_migrations_total 4" in text
|
||||
assert "turnstone_router_membership_size 3" in text
|
||||
assert "turnstone_router_refresh_total 12" in text
|
||||
|
||||
+178
-192
@@ -1,17 +1,13 @@
|
||||
"""Tests for turnstone.console.router."""
|
||||
"""Tests for turnstone.console.router (rendezvous routing)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
import secrets
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.console.router import ConsoleRouter, NodeRef
|
||||
from turnstone.core.hash_ring import RING_SIZE, NoAvailableNodeError
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fake storage
|
||||
# ---------------------------------------------------------------------------
|
||||
from turnstone.core.rendezvous import NoAvailableNodeError
|
||||
|
||||
|
||||
class FakeStorage:
|
||||
@@ -19,26 +15,14 @@ class FakeStorage:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.services: list[dict[str, str]] = []
|
||||
self.buckets: list[dict[str, Any]] = []
|
||||
self.overrides: list[dict[str, str]] = []
|
||||
self.settings: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def list_services(self, service_type: str, max_age_seconds: int = 120) -> list[dict[str, str]]:
|
||||
return list(self.services)
|
||||
|
||||
def list_ring_buckets(self) -> list[dict[str, Any]]:
|
||||
return list(self.buckets)
|
||||
|
||||
def list_workstream_overrides(self) -> list[dict[str, str]]:
|
||||
return list(self.overrides)
|
||||
|
||||
def get_system_setting(self, key: str, node_id: str = "") -> dict[str, Any] | None:
|
||||
return self.settings.get(key)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
NODE_A = {"service_id": "node-a", "url": "http://a:8080", "metadata": "{}"}
|
||||
NODE_B = {"service_id": "node-b", "url": "http://b:8080", "metadata": "{}"}
|
||||
@@ -50,260 +34,262 @@ def _make_router(storage: FakeStorage | None = None) -> tuple[ConsoleRouter, Fak
|
||||
return ConsoleRouter(s), s # type: ignore[arg-type]
|
||||
|
||||
|
||||
def _ws_id_for_bucket(bucket: int) -> str:
|
||||
"""Build a 32-char hex ws_id whose first 4 chars encode *bucket*."""
|
||||
return f"{bucket:04x}" + "0" * 28
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestRouteBasic
|
||||
# ---------------------------------------------------------------------------
|
||||
def _random_ws_id() -> str:
|
||||
return secrets.token_hex(16)
|
||||
|
||||
|
||||
class TestRouteBasic:
|
||||
"""Basic routing through the bucket cache."""
|
||||
|
||||
def test_route_returns_correct_node(self) -> None:
|
||||
def test_route_returns_a_live_node(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||
storage.buckets = [
|
||||
{"bucket": 0x0000, "node_id": "node-a"},
|
||||
{"bucket": 0x0001, "node_id": "node-b"},
|
||||
{"bucket": 0x0002, "node_id": "node-c"},
|
||||
]
|
||||
router.refresh_cache()
|
||||
|
||||
assert router.route(_ws_id_for_bucket(0x0000)) == NodeRef("node-a", "http://a:8080")
|
||||
assert router.route(_ws_id_for_bucket(0x0001)) == NodeRef("node-b", "http://b:8080")
|
||||
assert router.route(_ws_id_for_bucket(0x0002)) == NodeRef("node-c", "http://c:8080")
|
||||
ref = router.route(_random_ws_id())
|
||||
assert ref.node_id in {"node-a", "node-b", "node-c"}
|
||||
|
||||
def test_route_is_deterministic_for_same_ws_id(self) -> None:
|
||||
"""Same ws_id + same membership → same target every time."""
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||
router.refresh_cache()
|
||||
|
||||
ws_id = _random_ws_id()
|
||||
first = router.route(ws_id)
|
||||
for _ in range(50):
|
||||
assert router.route(ws_id) == first
|
||||
|
||||
def test_route_override_priority(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A, NODE_B]
|
||||
storage.buckets = [{"bucket": 0x0000, "node_id": "node-a"}]
|
||||
ws_id = _ws_id_for_bucket(0x0000)
|
||||
ws_id = _random_ws_id()
|
||||
storage.overrides = [{"ws_id": ws_id, "node_id": "node-b"}]
|
||||
router.refresh_cache()
|
||||
|
||||
# Override wins over bucket assignment
|
||||
# Override wins regardless of HRW score.
|
||||
assert router.route(ws_id) == NodeRef("node-b", "http://b:8080")
|
||||
|
||||
def test_route_empty_cache_raises(self) -> None:
|
||||
def test_route_empty_membership_raises(self) -> None:
|
||||
router, _ = _make_router()
|
||||
with pytest.raises(NoAvailableNodeError):
|
||||
router.route(_random_ws_id())
|
||||
|
||||
with pytest.raises(NoAvailableNodeError, match="not assigned"):
|
||||
router.route(_ws_id_for_bucket(0x0000))
|
||||
def test_route_empty_ws_id_raises(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
router.refresh_cache()
|
||||
with pytest.raises(NoAvailableNodeError, match="empty"):
|
||||
router.route("")
|
||||
|
||||
def test_route_url_convenience(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0x0010, "node_id": "node-a"}]
|
||||
router.refresh_cache()
|
||||
assert router.route_url(_random_ws_id()) == "http://a:8080"
|
||||
|
||||
|
||||
class TestMembershipConvergence:
|
||||
"""Rendezvous gives the minimal-moves property; pin it."""
|
||||
|
||||
def test_node_join_only_steals_some_keys(self) -> None:
|
||||
"""Adding a 4th node moves ~1/4 of keys to it; the other 3
|
||||
nodes' kept keys are unchanged."""
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||
router.refresh_cache()
|
||||
|
||||
assert router.route_url(_ws_id_for_bucket(0x0010)) == "http://a:8080"
|
||||
sample = [_random_ws_id() for _ in range(2000)]
|
||||
before = {ws: router.route(ws).node_id for ws in sample}
|
||||
|
||||
storage.services = [
|
||||
NODE_A,
|
||||
NODE_B,
|
||||
NODE_C,
|
||||
{"service_id": "node-d", "url": "http://d:8080", "metadata": "{}"},
|
||||
]
|
||||
router.refresh_cache()
|
||||
after = {ws: router.route(ws).node_id for ws in sample}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestRefreshCache
|
||||
# ---------------------------------------------------------------------------
|
||||
moved = sum(1 for ws in sample if before[ws] != after[ws])
|
||||
moved_to_new = sum(1 for ws in sample if after[ws] == "node-d")
|
||||
# Every move must be onto the new node — no churn between
|
||||
# existing nodes.
|
||||
assert moved == moved_to_new
|
||||
# Should be roughly 1/4 of keys; allow a wide band for variance.
|
||||
assert 0.15 < moved / len(sample) < 0.35
|
||||
|
||||
|
||||
class TestRefreshCache:
|
||||
"""Cache loading from storage."""
|
||||
|
||||
def test_refresh_loads_from_storage(self) -> None:
|
||||
def test_node_leave_only_redistributes_dead_node_keys(self) -> None:
|
||||
"""Removing node-a sends node-a's keys to b/c only; keys that
|
||||
were on b/c stay put."""
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 100, "node_id": "node-a"}]
|
||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||
router.refresh_cache()
|
||||
|
||||
ref = router.route(_ws_id_for_bucket(100))
|
||||
assert ref.node_id == "node-a"
|
||||
sample = [_random_ws_id() for _ in range(2000)]
|
||||
before = {ws: router.route(ws).node_id for ws in sample}
|
||||
|
||||
def test_refresh_handles_dead_nodes(self) -> None:
|
||||
storage.services = [NODE_B, NODE_C]
|
||||
router.refresh_cache()
|
||||
after = {ws: router.route(ws).node_id for ws in sample}
|
||||
|
||||
for ws in sample:
|
||||
if before[ws] in ("node-b", "node-c"):
|
||||
assert after[ws] == before[ws], (
|
||||
f"key {ws} moved from {before[ws]} to {after[ws]} "
|
||||
"even though its old owner is still live"
|
||||
)
|
||||
else: # was on node-a
|
||||
assert after[ws] in ("node-b", "node-c")
|
||||
|
||||
|
||||
class TestWeights:
|
||||
def test_weight_2_node_gets_more_keys_than_weight_1(self) -> None:
|
||||
router, storage = _make_router()
|
||||
# node-b is in buckets but not in services (dead/expired)
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [
|
||||
{"bucket": 0x0000, "node_id": "node-a"},
|
||||
{"bucket": 0x0001, "node_id": "node-b"},
|
||||
storage.services = [
|
||||
{"service_id": "node-a", "url": "http://a:8080", "metadata": '{"weight": 2}'},
|
||||
{"service_id": "node-b", "url": "http://b:8080", "metadata": '{"weight": 1}'},
|
||||
]
|
||||
router.refresh_cache()
|
||||
|
||||
assert router.route(_ws_id_for_bucket(0x0000)).node_id == "node-a"
|
||||
with pytest.raises(NoAvailableNodeError):
|
||||
router.route(_ws_id_for_bucket(0x0001))
|
||||
sample = [_random_ws_id() for _ in range(5000)]
|
||||
on_a = sum(1 for ws in sample if router.route(ws).node_id == "node-a")
|
||||
# Heavier node should win clearly more than half; exact ratio
|
||||
# depends on the simple hash×weight formulation but a/b > 1.4
|
||||
# for weight 2:1 across 5k samples is reliable.
|
||||
assert on_a / len(sample) > 0.55
|
||||
|
||||
def test_refresh_returns_true_on_change(self) -> None:
|
||||
def test_invalid_metadata_falls_back_to_weight_1(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
||||
|
||||
assert router.refresh_cache() is True
|
||||
|
||||
def test_refresh_returns_false_on_no_change(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
||||
|
||||
storage.services = [
|
||||
{"service_id": "node-a", "url": "http://a:8080", "metadata": "not json"},
|
||||
]
|
||||
router.refresh_cache()
|
||||
assert router.refresh_cache() is False
|
||||
# Just confirms it doesn't blow up.
|
||||
router.route(_random_ws_id())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestCheckVersion
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCheckVersion:
|
||||
"""Version-gated refresh."""
|
||||
|
||||
def test_version_change_triggers_refresh(self) -> None:
|
||||
class TestRefreshLifecycle:
|
||||
def test_refresh_cache_publishes_new_membership_immediately(self) -> None:
|
||||
"""refresh_cache() reloads on the calling thread — the next
|
||||
route() sees the new membership without any further trigger."""
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
||||
storage.settings["rebalancer_version"] = {"value": "1"}
|
||||
router.refresh_cache()
|
||||
assert router.node_count() == 1
|
||||
|
||||
assert router.check_version() is True
|
||||
assert router.is_ready()
|
||||
storage.services = [NODE_A, NODE_B]
|
||||
router.refresh_cache()
|
||||
assert router.node_count() == 2
|
||||
|
||||
def test_concurrent_refresh_returns_false_on_lock_contention(self) -> None:
|
||||
"""refresh_cache uses a non-blocking lock acquire — if another
|
||||
thread is already refreshing, the second caller bails so the
|
||||
in-flight refresh's result is the one that publishes."""
|
||||
|
||||
def test_same_version_skips(self) -> None:
|
||||
router, storage = _make_router()
|
||||
# Default version is 0; setting absent also means 0
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
||||
|
||||
# First call: version=0 matches self._version=0 -> no refresh
|
||||
assert router.check_version() is False
|
||||
assert not router.is_ready() # cache was never loaded
|
||||
with router._refresh_lock:
|
||||
# Lock held by this thread → the call below can't acquire.
|
||||
assert router.refresh_cache() is False
|
||||
|
||||
def test_force_refresh_blocks_until_in_flight_refresh_releases(self) -> None:
|
||||
"""force_refresh acquires the refresh lock blocking — used by the
|
||||
404-retry path to guarantee a fresh view even under contention."""
|
||||
import threading
|
||||
|
||||
def test_version_none_treated_as_zero(self) -> None:
|
||||
router, storage = _make_router()
|
||||
# settings dict is empty -> get_system_setting returns None
|
||||
assert router.check_version() is False
|
||||
storage.services = [NODE_A]
|
||||
|
||||
# Hold the refresh lock from another thread.
|
||||
lock_held = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestGenerateWsId
|
||||
# ---------------------------------------------------------------------------
|
||||
def hold_lock() -> None:
|
||||
with router._refresh_lock:
|
||||
lock_held.set()
|
||||
release.wait(timeout=2)
|
||||
|
||||
holder = threading.Thread(target=hold_lock, daemon=True)
|
||||
holder.start()
|
||||
assert lock_held.wait(timeout=1)
|
||||
|
||||
# force_refresh should block, not bail.
|
||||
result_box: list[bool] = []
|
||||
|
||||
def call_force() -> None:
|
||||
result_box.append(router.force_refresh())
|
||||
|
||||
caller = threading.Thread(target=call_force, daemon=True)
|
||||
caller.start()
|
||||
caller.join(timeout=0.2)
|
||||
assert caller.is_alive(), "force_refresh returned without acquiring lock"
|
||||
|
||||
release.set()
|
||||
holder.join(timeout=1)
|
||||
caller.join(timeout=1)
|
||||
assert not caller.is_alive()
|
||||
# Membership changed from empty → 1 live node.
|
||||
assert result_box == [True]
|
||||
assert router.node_count() == 1
|
||||
|
||||
def test_force_refresh_always_reloads(self) -> None:
|
||||
"""force_refresh skips the non-blocking-lock bail and always
|
||||
publishes a fresh view — back-to-back calls each pick up the
|
||||
latest storage state."""
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
router.force_refresh()
|
||||
assert router.node_count() == 1
|
||||
|
||||
storage.services = [NODE_A, NODE_B]
|
||||
router.force_refresh()
|
||||
assert router.node_count() == 2
|
||||
|
||||
def test_version_is_monotonic_across_refreshes(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
router.refresh_cache()
|
||||
v1 = router.version
|
||||
router.refresh_cache()
|
||||
v2 = router.version
|
||||
assert v2 > v1
|
||||
router.force_refresh()
|
||||
assert router.version > v2
|
||||
|
||||
|
||||
class TestGenerateWsId:
|
||||
"""Workstream ID generation targeting a specific node."""
|
||||
|
||||
def test_generates_routable_id(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A, NODE_B]
|
||||
storage.buckets = [
|
||||
{"bucket": 0x00FF, "node_id": "node-a"},
|
||||
{"bucket": 0x0100, "node_id": "node-b"},
|
||||
]
|
||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||
router.refresh_cache()
|
||||
|
||||
ws_id = router.generate_ws_id_for_node("node-a")
|
||||
ws_id = router.generate_ws_id_for_node("node-b")
|
||||
assert len(ws_id) == 32
|
||||
assert router.route(ws_id).node_id == "node-a"
|
||||
assert router.route(ws_id).node_id == "node-b"
|
||||
|
||||
def test_unknown_node_raises(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
||||
router.refresh_cache()
|
||||
|
||||
with pytest.raises(NoAvailableNodeError, match="node-z"):
|
||||
router.generate_ws_id_for_node("node-z")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestIsReady
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIsReady:
|
||||
"""Readiness checks."""
|
||||
|
||||
def test_false_when_empty(self) -> None:
|
||||
router, _ = _make_router()
|
||||
assert router.is_ready() is False
|
||||
|
||||
def test_true_after_refresh(self) -> None:
|
||||
def test_true_after_membership_loads(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A]
|
||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
||||
router.refresh_cache()
|
||||
|
||||
assert router.is_ready() is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestPopulateFromAssignments
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPopulateFromAssignments:
|
||||
"""Direct cache population without DB round-trip."""
|
||||
|
||||
def test_populate_makes_router_ready(self) -> None:
|
||||
router, _ = _make_router()
|
||||
assignments = [(b, "node-a") for b in range(RING_SIZE)]
|
||||
nodes = {"node-a": NodeRef("node-a", "http://a:8080")}
|
||||
router.populate_from_assignments(assignments, nodes)
|
||||
|
||||
assert router.is_ready()
|
||||
assert router.node_count() == 1
|
||||
assert router.route(_ws_id_for_bucket(0)).node_id == "node-a"
|
||||
|
||||
def test_populate_multi_node(self) -> None:
|
||||
router, _ = _make_router()
|
||||
assignments = [(0, "node-a"), (1, "node-b"), (2, "node-a")]
|
||||
nodes = {
|
||||
"node-a": NodeRef("node-a", "http://a:8080"),
|
||||
"node-b": NodeRef("node-b", "http://b:8080"),
|
||||
}
|
||||
router.populate_from_assignments(assignments, nodes)
|
||||
|
||||
assert router.route(_ws_id_for_bucket(0)).node_id == "node-a"
|
||||
assert router.route(_ws_id_for_bucket(1)).node_id == "node-b"
|
||||
assert router.route(_ws_id_for_bucket(2)).node_id == "node-a"
|
||||
|
||||
def test_populate_loads_overrides_from_db(self) -> None:
|
||||
router, storage = _make_router()
|
||||
ws_id = _ws_id_for_bucket(0)
|
||||
storage.overrides = [{"ws_id": ws_id, "node_id": "node-b"}]
|
||||
nodes = {
|
||||
"node-a": NodeRef("node-a", "http://a:8080"),
|
||||
"node-b": NodeRef("node-b", "http://b:8080"),
|
||||
}
|
||||
router.populate_from_assignments([(0, "node-a")], nodes)
|
||||
|
||||
# Override should route bucket 0 to node-b despite assignment to node-a
|
||||
assert router.route(ws_id) == NodeRef("node-b", "http://b:8080")
|
||||
|
||||
def test_populate_no_overrides_when_table_empty(self) -> None:
|
||||
router, storage = _make_router()
|
||||
# No overrides in storage
|
||||
router.populate_from_assignments(
|
||||
[(0, "node-a")],
|
||||
{"node-a": NodeRef("node-a", "http://a:8080")},
|
||||
)
|
||||
assert len(router._overrides) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestNodeCount
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestNodeCount:
|
||||
"""Distinct node counting."""
|
||||
|
||||
def test_count_distinct_nodes(self) -> None:
|
||||
def test_count_matches_live_services(self) -> None:
|
||||
router, storage = _make_router()
|
||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||
# Spread all 65536 buckets across 3 nodes
|
||||
storage.buckets = [
|
||||
{"bucket": b, "node_id": f"node-{['a', 'b', 'c'][b % 3]}"} for b in range(RING_SIZE)
|
||||
]
|
||||
router.refresh_cache()
|
||||
|
||||
assert router.node_count() == 3
|
||||
|
||||
@@ -11,7 +11,7 @@ from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.console.collector import ClusterCollector
|
||||
from turnstone.console.router import ConsoleRouter, NodeRef
|
||||
from turnstone.core.hash_ring import NoAvailableNodeError
|
||||
from turnstone.core.rendezvous import NoAvailableNodeError
|
||||
|
||||
# Shared test auth — JWT-based
|
||||
_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
@@ -102,7 +102,7 @@ def _wire_proxy(app: Any, mock_post: MagicMock | None = None) -> None:
|
||||
|
||||
|
||||
class TestRouteCreate:
|
||||
"""POST /v1/api/route/workstreams/new — create via hash-ring routing."""
|
||||
"""POST /v1/api/route/workstreams/new — create via rendezvous routing."""
|
||||
|
||||
@pytest.fixture()
|
||||
def client(self):
|
||||
@@ -178,17 +178,17 @@ class TestRouteCreate:
|
||||
router.generate_ws_id_for_node.assert_called_with("node-c")
|
||||
client.close()
|
||||
|
||||
def test_route_create_routing_strategy_hash_ring(self, client):
|
||||
def test_route_create_routing_strategy_rendezvous(self, client):
|
||||
"""Default fan-out (no resume_ws / no target_node) reports
|
||||
routing_strategy='hash_ring' so the coordinator's spawn tool can
|
||||
explain why the node was chosen."""
|
||||
routing_strategy='rendezvous' so the coordinator's spawn tool
|
||||
can explain why the node was chosen."""
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test-ws"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["routing_strategy"] == "hash_ring"
|
||||
assert resp.json()["routing_strategy"] == "rendezvous"
|
||||
|
||||
def test_route_create_routing_strategy_target_node(self):
|
||||
router = _make_mock_router()
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
"""Tests for turnstone.core.hash_ring."""
|
||||
|
||||
from turnstone.core.hash_ring import bucket_of
|
||||
|
||||
|
||||
class TestBucketOf:
|
||||
def test_known_vectors(self):
|
||||
assert bucket_of("a3f1" + "0" * 28) == 0xA3F1
|
||||
assert bucket_of("0000" + "a" * 28) == 0
|
||||
assert bucket_of("ffff" + "b" * 28) == 65535
|
||||
|
||||
def test_hex_prefix(self):
|
||||
# Only the first 4 hex chars matter — the rest is ignored.
|
||||
assert bucket_of("abcd0000") == bucket_of("abcdffff")
|
||||
assert bucket_of("abcd0000") == 0xABCD
|
||||
@@ -1,174 +0,0 @@
|
||||
"""Tests for the hash ring routing storage methods."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class TestHashRingBuckets:
|
||||
def test_list_empty(self, storage):
|
||||
assert storage.list_ring_buckets() == []
|
||||
|
||||
def test_seed_and_list(self, storage):
|
||||
storage.seed_ring_buckets([(0, "node-a"), (1, "node-b"), (2, "node-a")])
|
||||
rows = storage.list_ring_buckets()
|
||||
assert len(rows) == 3
|
||||
assert rows[0] == {"bucket": 0, "node_id": "node-a"}
|
||||
assert rows[1] == {"bucket": 1, "node_id": "node-b"}
|
||||
assert rows[2] == {"bucket": 2, "node_id": "node-a"}
|
||||
|
||||
def test_seed_idempotent(self, storage):
|
||||
storage.seed_ring_buckets([(0, "node-a"), (1, "node-b")])
|
||||
# Re-seed with conflicting assignment: should keep original
|
||||
storage.seed_ring_buckets([(0, "node-x"), (2, "node-c")])
|
||||
rows = storage.list_ring_buckets()
|
||||
by_bucket = {r["bucket"]: r["node_id"] for r in rows}
|
||||
assert by_bucket[0] == "node-a" # original preserved
|
||||
assert by_bucket[1] == "node-b"
|
||||
assert by_bucket[2] == "node-c" # new bucket added
|
||||
|
||||
def test_assign_buckets(self, storage):
|
||||
storage.seed_ring_buckets([(0, "node-a"), (1, "node-a"), (2, "node-b")])
|
||||
storage.assign_buckets([0, 1], "node-c")
|
||||
rows = storage.list_ring_buckets()
|
||||
by_bucket = {r["bucket"]: r["node_id"] for r in rows}
|
||||
assert by_bucket[0] == "node-c"
|
||||
assert by_bucket[1] == "node-c"
|
||||
assert by_bucket[2] == "node-b"
|
||||
|
||||
def test_assign_returns_count(self, storage):
|
||||
storage.seed_ring_buckets([(0, "node-a"), (1, "node-a")])
|
||||
count = storage.assign_buckets([0, 1], "node-b")
|
||||
assert count == 2
|
||||
# Empty list returns 0
|
||||
assert storage.assign_buckets([], "node-x") == 0
|
||||
|
||||
def test_assign_large_list_exceeds_chunk_size(self, storage):
|
||||
"""Regression: lists larger than chunk_size must not hit param limits."""
|
||||
n = 1200 # exceeds SQLite chunk_size (500) and exercises multi-chunk path
|
||||
storage.seed_ring_buckets([(i, "node-a") for i in range(n)])
|
||||
count = storage.assign_buckets(list(range(n)), "node-b")
|
||||
assert count == n
|
||||
rows = storage.list_ring_buckets()
|
||||
assert all(r["node_id"] == "node-b" for r in rows)
|
||||
|
||||
def test_assign_deduplicates_input(self, storage):
|
||||
"""Duplicates in the input list should not inflate rowcount."""
|
||||
storage.seed_ring_buckets([(0, "node-a"), (1, "node-a")])
|
||||
count = storage.assign_buckets([0, 1, 0, 1, 0], "node-b")
|
||||
assert count == 2
|
||||
|
||||
|
||||
class TestBucketStats:
|
||||
def test_increment_creates_row(self, storage):
|
||||
storage.increment_bucket_count(42)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert len(stats) == 1
|
||||
assert stats[0]["bucket"] == 42
|
||||
assert stats[0]["ws_count"] == 1
|
||||
assert stats[0]["active_count"] == 0
|
||||
|
||||
def test_increment_active(self, storage):
|
||||
storage.increment_bucket_count(10, active=True)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert stats[0]["ws_count"] == 1
|
||||
assert stats[0]["active_count"] == 1
|
||||
# Increment again without active
|
||||
storage.increment_bucket_count(10)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert stats[0]["ws_count"] == 2
|
||||
assert stats[0]["active_count"] == 1
|
||||
|
||||
def test_decrement(self, storage):
|
||||
storage.increment_bucket_count(5, active=True)
|
||||
storage.increment_bucket_count(5, active=True)
|
||||
storage.decrement_bucket_count(5, active=True)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert stats[0]["ws_count"] == 1
|
||||
assert stats[0]["active_count"] == 1
|
||||
|
||||
def test_decrement_clamps_at_zero(self, storage):
|
||||
storage.increment_bucket_count(7)
|
||||
storage.decrement_bucket_count(7)
|
||||
storage.decrement_bucket_count(7) # already at 0
|
||||
stats = storage.list_bucket_stats()
|
||||
# ws_count is 0, so should not appear (filter ws_count > 0)
|
||||
assert len(stats) == 0
|
||||
|
||||
def test_adjust_active_only(self, storage):
|
||||
storage.increment_bucket_count(20, active=True)
|
||||
storage.increment_bucket_count(20, active=True)
|
||||
# Decrease active without changing ws_count
|
||||
storage.adjust_bucket_active(20, -1)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert stats[0]["ws_count"] == 2
|
||||
assert stats[0]["active_count"] == 1
|
||||
# Clamp at zero
|
||||
storage.adjust_bucket_active(20, -5)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert stats[0]["active_count"] == 0
|
||||
|
||||
def test_list_sparse(self, storage):
|
||||
storage.increment_bucket_count(100)
|
||||
storage.increment_bucket_count(200)
|
||||
storage.increment_bucket_count(300)
|
||||
# Decrement 200 to zero
|
||||
storage.decrement_bucket_count(200)
|
||||
stats = storage.list_bucket_stats()
|
||||
buckets = [s["bucket"] for s in stats]
|
||||
assert 100 in buckets
|
||||
assert 200 not in buckets
|
||||
assert 300 in buckets
|
||||
|
||||
def test_set_bucket_stat_creates(self, storage):
|
||||
"""set_bucket_stat upserts a new row."""
|
||||
storage.set_bucket_stat(42, 5, 2)
|
||||
stats = storage.list_bucket_stats()
|
||||
row = next(s for s in stats if s["bucket"] == 42)
|
||||
assert row["ws_count"] == 5
|
||||
assert row["active_count"] == 2
|
||||
|
||||
def test_set_bucket_stat_overwrites(self, storage):
|
||||
"""set_bucket_stat overwrites existing values."""
|
||||
storage.set_bucket_stat(42, 10, 3)
|
||||
storage.set_bucket_stat(42, 2, 0)
|
||||
stats = storage.list_bucket_stats()
|
||||
row = next(s for s in stats if s["bucket"] == 42)
|
||||
assert row["ws_count"] == 2
|
||||
assert row["active_count"] == 0
|
||||
|
||||
def test_set_bucket_stat_zero_removes_from_sparse(self, storage):
|
||||
"""Setting ws_count=0 means list_bucket_stats excludes it (sparse)."""
|
||||
storage.set_bucket_stat(42, 5, 1)
|
||||
storage.set_bucket_stat(42, 0, 0)
|
||||
stats = storage.list_bucket_stats()
|
||||
assert not any(s["bucket"] == 42 for s in stats)
|
||||
|
||||
|
||||
class TestWorkstreamOverrides:
|
||||
def test_set_and_list(self, storage):
|
||||
storage.set_workstream_override("ws-001", "node-a", reason="affinity")
|
||||
overrides = storage.list_workstream_overrides()
|
||||
assert len(overrides) == 1
|
||||
assert overrides[0]["ws_id"] == "ws-001"
|
||||
assert overrides[0]["node_id"] == "node-a"
|
||||
assert overrides[0]["reason"] == "affinity"
|
||||
|
||||
def test_upsert(self, storage):
|
||||
storage.set_workstream_override("ws-002", "node-a")
|
||||
storage.set_workstream_override("ws-002", "node-b", reason="migration")
|
||||
overrides = storage.list_workstream_overrides()
|
||||
assert len(overrides) == 1
|
||||
assert overrides[0]["node_id"] == "node-b"
|
||||
assert overrides[0]["reason"] == "migration"
|
||||
|
||||
def test_delete(self, storage):
|
||||
storage.set_workstream_override("ws-003", "node-a")
|
||||
result = storage.delete_workstream_override("ws-003")
|
||||
assert result is True
|
||||
assert storage.list_workstream_overrides() == []
|
||||
|
||||
def test_delete_nonexistent(self, storage):
|
||||
result = storage.delete_workstream_override("ws-nope")
|
||||
assert result is False
|
||||
|
||||
def test_list_empty(self, storage):
|
||||
assert storage.list_workstream_overrides() == []
|
||||
@@ -1,561 +0,0 @@
|
||||
"""Tests for turnstone.console.rebalancer."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.console.rebalancer import Rebalancer
|
||||
from turnstone.core.hash_ring import RING_SIZE
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path):
|
||||
"""Fresh SQLite backend for each test."""
|
||||
return SQLiteBackend(str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
def _register_nodes(storage: SQLiteBackend, count: int, *, weight: int = 1) -> None:
|
||||
"""Register *count* server nodes in the services table."""
|
||||
for i in range(count):
|
||||
meta = json.dumps({"weight": weight, "started": "2026-01-01T00:00:00Z"})
|
||||
storage.register_service("server", f"node-{i}", f"http://node-{i}:8080", metadata=meta)
|
||||
|
||||
|
||||
def _register_weighted_nodes(storage: SQLiteBackend, weights: dict[str, int]) -> None:
|
||||
"""Register nodes with specific weights."""
|
||||
for node_id, w in weights.items():
|
||||
meta = json.dumps({"weight": w, "started": "2026-01-01T00:00:00Z"})
|
||||
storage.register_service("server", node_id, f"http://{node_id}:8080", metadata=meta)
|
||||
|
||||
|
||||
def _get_version(storage: SQLiteBackend) -> int:
|
||||
"""Read the rebalancer_version from system_settings."""
|
||||
raw = storage.get_system_setting("rebalancer_version", node_id="")
|
||||
if raw is None:
|
||||
return 0
|
||||
try:
|
||||
return int(json.loads(raw.get("value", "0")))
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
class TestFirstRunSeed:
|
||||
def test_first_run_seeds_ring(self, storage):
|
||||
"""Empty assignment table + 2 nodes -> seed all 65536 rows."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage)
|
||||
result = rb.rebalance_once()
|
||||
|
||||
assert result.seeded is True
|
||||
assert result.noop is False
|
||||
assert result.nodes == 2
|
||||
|
||||
buckets = storage.list_ring_buckets()
|
||||
assert len(buckets) == RING_SIZE
|
||||
|
||||
# All buckets should be assigned to one of the two nodes
|
||||
node_ids = {b["node_id"] for b in buckets}
|
||||
assert node_ids == {"node-0", "node-1"}
|
||||
|
||||
|
||||
class TestSeedPopulatesRouter:
|
||||
def test_seed_populates_router_directly(self, storage):
|
||||
"""On first seed, the router cache is populated without a DB read-back."""
|
||||
from turnstone.console.router import ConsoleRouter
|
||||
|
||||
_register_nodes(storage, 2)
|
||||
router = ConsoleRouter(storage)
|
||||
assert not router.is_ready()
|
||||
|
||||
rb = Rebalancer(storage=storage, router=router)
|
||||
result = rb.rebalance_once()
|
||||
|
||||
assert result.seeded is True
|
||||
assert router.is_ready()
|
||||
assert router.node_count() == 2
|
||||
|
||||
# Routing should work for any valid ws_id
|
||||
ws_id = "0000" + "a" * 28
|
||||
ref = router.route(ws_id)
|
||||
assert ref.node_id in {"node-0", "node-1"}
|
||||
|
||||
|
||||
class TestIdempotent:
|
||||
def test_second_run_is_noop(self, storage):
|
||||
"""Running rebalance twice with same membership produces noop on second pass."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage)
|
||||
|
||||
r1 = rb.rebalance_once()
|
||||
assert r1.seeded is True
|
||||
|
||||
r2 = rb.rebalance_once()
|
||||
assert r2.noop is True
|
||||
assert r2.moves == 0
|
||||
|
||||
|
||||
class TestNewNodeRebalances:
|
||||
def test_adding_node_moves_buckets(self, storage):
|
||||
"""Seed with 2 nodes, add 3rd -> some buckets move to the new node."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage, threshold=0.01)
|
||||
rb.rebalance_once() # seed
|
||||
|
||||
# Verify only 2 nodes initially
|
||||
buckets_before = storage.list_ring_buckets()
|
||||
nodes_before = {b["node_id"] for b in buckets_before}
|
||||
assert nodes_before == {"node-0", "node-1"}
|
||||
|
||||
# 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()
|
||||
assert result.noop is False
|
||||
assert result.moves > 0
|
||||
assert result.nodes == 3
|
||||
|
||||
# Verify all three nodes have buckets
|
||||
buckets_after = storage.list_ring_buckets()
|
||||
nodes_after = {b["node_id"] for b in buckets_after}
|
||||
assert "node-2" in nodes_after
|
||||
|
||||
|
||||
class TestDeadNodeReassigned:
|
||||
def test_dead_node_buckets_move_to_survivors(self, storage):
|
||||
"""Seed with 3 nodes, deregister one -> its buckets move to survivors."""
|
||||
_register_nodes(storage, 3)
|
||||
rb = Rebalancer(storage=storage, threshold=0.01)
|
||||
rb.rebalance_once() # seed
|
||||
|
||||
# Verify node-2 has some buckets
|
||||
buckets = storage.list_ring_buckets()
|
||||
node2_count = sum(1 for b in buckets if b["node_id"] == "node-2")
|
||||
assert node2_count > 0
|
||||
|
||||
# Deregister node-2
|
||||
storage.deregister_service("server", "node-2")
|
||||
|
||||
result = rb.rebalance_once()
|
||||
assert result.noop is False
|
||||
assert result.moves > 0
|
||||
|
||||
# Verify no buckets assigned to dead node
|
||||
buckets_after = storage.list_ring_buckets()
|
||||
nodes_after = {b["node_id"] for b in buckets_after}
|
||||
assert "node-2" not in nodes_after
|
||||
|
||||
|
||||
class TestSingleNodeNoop:
|
||||
def test_single_node_already_assigned_is_noop(self, storage):
|
||||
"""1 node with all buckets assigned -> noop."""
|
||||
_register_nodes(storage, 1)
|
||||
rb = Rebalancer(storage=storage)
|
||||
|
||||
# Seed with single node
|
||||
rb.rebalance_once()
|
||||
|
||||
# Second run should be noop
|
||||
result = rb.rebalance_once()
|
||||
assert result.noop is True
|
||||
|
||||
|
||||
class TestWeightedDistribution:
|
||||
def test_weight_2_gets_more_buckets(self, storage):
|
||||
"""Node with weight=2 gets roughly 2x the buckets of weight=1."""
|
||||
_register_weighted_nodes(storage, {"heavy": 2, "light": 1})
|
||||
rb = Rebalancer(storage=storage, vnodes_per_unit=150)
|
||||
rb.rebalance_once() # seed
|
||||
|
||||
buckets = storage.list_ring_buckets()
|
||||
heavy_count = sum(1 for b in buckets if b["node_id"] == "heavy")
|
||||
light_count = sum(1 for b in buckets if b["node_id"] == "light")
|
||||
|
||||
# heavy should have roughly 2/3 of total, light roughly 1/3
|
||||
# Allow 10% tolerance
|
||||
expected_heavy = RING_SIZE * 2 // 3
|
||||
assert abs(heavy_count - expected_heavy) < RING_SIZE * 0.10
|
||||
assert heavy_count > light_count
|
||||
|
||||
|
||||
class TestVersionIncremented:
|
||||
def test_version_bumps_on_seed(self, storage):
|
||||
"""Verify rebalancer_version increments after seed."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage)
|
||||
|
||||
v0 = _get_version(storage)
|
||||
assert v0 == 0
|
||||
|
||||
rb.rebalance_once()
|
||||
v1 = _get_version(storage)
|
||||
assert v1 == 1
|
||||
|
||||
def test_version_bumps_on_rebalance(self, storage):
|
||||
"""Version bumps on actual moves, not on noops."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage, threshold=0.01)
|
||||
rb.rebalance_once() # seed: version -> 1
|
||||
|
||||
# Noop: version stays at 1
|
||||
rb.rebalance_once()
|
||||
assert _get_version(storage) == 1
|
||||
|
||||
# Add node: version -> 2
|
||||
meta = json.dumps({"weight": 1, "started": "2026-01-01T00:00:00Z"})
|
||||
storage.register_service("server", "node-2", "http://node-2:8080", metadata=meta)
|
||||
rb.rebalance_once()
|
||||
assert _get_version(storage) == 2
|
||||
|
||||
|
||||
class TestReconcileStats:
|
||||
def test_bucket_stats_corrected(self, storage):
|
||||
"""Create workstreams in DB, verify bucket_stats are reconciled."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage)
|
||||
rb.rebalance_once() # seed
|
||||
|
||||
# Create some workstreams — ws_id starts with hex bucket
|
||||
# Bucket 0x0000 = 0, bucket 0x0001 = 1
|
||||
storage.register_workstream("0000" + "a" * 28, state="idle")
|
||||
storage.register_workstream("0000" + "b" * 28, state="running")
|
||||
storage.register_workstream("0001" + "c" * 28, state="idle")
|
||||
|
||||
# Set bogus stats that will be corrected
|
||||
storage.increment_bucket_count(0) # says 1, should be 2
|
||||
storage.increment_bucket_count(5) # says 1, should be 0
|
||||
|
||||
rb._reconcile_bucket_stats()
|
||||
|
||||
stats = storage.list_bucket_stats()
|
||||
stats_map = {s["bucket"]: s for s in stats}
|
||||
|
||||
# Bucket 0 should have 2 ws, 1 active (running)
|
||||
assert stats_map[0]["ws_count"] == 2
|
||||
assert stats_map[0]["active_count"] == 1
|
||||
|
||||
# Bucket 1 should have 1 ws, 0 active
|
||||
assert stats_map[1]["ws_count"] == 1
|
||||
assert stats_map[1]["active_count"] == 0
|
||||
|
||||
# Bucket 5 should have been removed (ws_count=0)
|
||||
assert 5 not in stats_map
|
||||
|
||||
|
||||
class TestTransferPriorityEmptyFirst:
|
||||
def test_empty_buckets_moved_before_occupied(self, storage):
|
||||
"""Verify the sort key puts empty buckets before occupied ones.
|
||||
|
||||
Rather than asserting specific bucket assignments (which depend on
|
||||
hash ring placement), we verify the sorting invariant directly by
|
||||
checking that moves with zero occupancy come before occupied ones
|
||||
in the internal ordering.
|
||||
"""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage, threshold=0.01)
|
||||
rb.rebalance_once() # seed
|
||||
|
||||
# Create workstreams in a few buckets owned by node-0
|
||||
buckets = storage.list_ring_buckets()
|
||||
node0_buckets = [b["bucket"] for b in buckets if b["node_id"] == "node-0"]
|
||||
|
||||
occupied = set()
|
||||
for b in node0_buckets[:3]:
|
||||
ws_id = f"{b:04x}" + "d" * 28
|
||||
storage.register_workstream(ws_id, state="running")
|
||||
storage.increment_bucket_count(b, active=True)
|
||||
occupied.add(b)
|
||||
|
||||
# Reconcile stats so the rebalancer sees them
|
||||
rb._reconcile_bucket_stats()
|
||||
|
||||
# Read stats to verify ordering assumptions
|
||||
stats = storage.list_bucket_stats()
|
||||
stats_map = {s["bucket"]: (s["ws_count"], s["active_count"]) for s in stats}
|
||||
|
||||
# The sort key is (active_count, ws_count) — occupied buckets
|
||||
# must sort AFTER empty buckets
|
||||
for b in occupied:
|
||||
assert stats_map[b][0] > 0 # ws_count > 0
|
||||
assert stats_map[b][1] > 0 # active_count > 0
|
||||
|
||||
# Empty buckets have (0, 0) which sorts before (1, 1)
|
||||
assert (0, 0) < (1, 1)
|
||||
|
||||
|
||||
class TestLeaderElection:
|
||||
def test_two_rebalancers_one_runs(self, storage):
|
||||
"""Two rebalancers compete — only one acquires the lock."""
|
||||
_register_nodes(storage, 2)
|
||||
rb1 = Rebalancer(storage=storage)
|
||||
rb2 = Rebalancer(storage=storage)
|
||||
|
||||
# rb1 acquires the lock
|
||||
assert rb1._try_acquire_lock() is True
|
||||
|
||||
# rb2 cannot acquire (lock is fresh)
|
||||
assert rb2._try_acquire_lock() is False
|
||||
|
||||
# rb1 releases
|
||||
rb1._release_lock()
|
||||
|
||||
# Now rb2 can acquire
|
||||
assert rb2._try_acquire_lock() is True
|
||||
rb2._release_lock()
|
||||
|
||||
|
||||
class TestZeroNodes:
|
||||
def test_no_nodes_returns_noop(self, storage):
|
||||
"""Zero live nodes -> noop result."""
|
||||
rb = Rebalancer(storage=storage)
|
||||
result = rb.rebalance_once()
|
||||
assert result.noop is True
|
||||
assert result.nodes == 0
|
||||
|
||||
|
||||
class TestStartStop:
|
||||
def test_start_stop_lifecycle(self, storage):
|
||||
"""Verify start/stop lifecycle doesn't hang or crash."""
|
||||
_register_nodes(storage, 1)
|
||||
rb = Rebalancer(storage=storage, interval=1)
|
||||
rb.start()
|
||||
assert rb._thread is not None
|
||||
assert rb._thread.is_alive()
|
||||
rb.stop()
|
||||
assert not rb._thread.is_alive()
|
||||
|
||||
def test_trigger_wakes_thread(self, storage):
|
||||
"""Verify trigger() causes an immediate pass."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage, interval=3600) # long interval
|
||||
rb.start()
|
||||
try:
|
||||
rb.trigger()
|
||||
# Give it a moment to process
|
||||
rb._stop_event.wait(timeout=2)
|
||||
finally:
|
||||
rb.stop()
|
||||
# After trigger, the ring should be seeded
|
||||
assert len(storage.list_ring_buckets()) == RING_SIZE
|
||||
|
||||
|
||||
class TestGetStatus:
|
||||
def test_status_before_any_run(self, storage):
|
||||
"""Status returns version=0 and no last_result before any run."""
|
||||
rb = Rebalancer(storage=storage)
|
||||
status = rb.get_status()
|
||||
assert status["version"] == 0
|
||||
assert status["last_result"] is None
|
||||
|
||||
def test_status_after_seed(self, storage):
|
||||
"""Status reflects the seed run when result is stored."""
|
||||
_register_nodes(storage, 2)
|
||||
rb = Rebalancer(storage=storage)
|
||||
result = rb.rebalance_once()
|
||||
# The loop normally sets _last_result; simulate that here
|
||||
rb._last_result = result
|
||||
status = rb.get_status()
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Tests for turnstone.core.rendezvous (HRW routing primitive)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.rendezvous import NoAvailableNodeError, NodeRef, fnv1a_32, select, select_all
|
||||
|
||||
|
||||
class TestFnv1aVectors:
|
||||
"""Pin the FNV-1a-32 implementation against the documented test
|
||||
vectors so cross-language readers (Go, TS) stay in sync."""
|
||||
|
||||
def test_empty_input(self) -> None:
|
||||
assert fnv1a_32(b"") == 0x811C9DC5 # basis
|
||||
|
||||
def test_foobar(self) -> None:
|
||||
assert fnv1a_32(b"foobar") == 0xBF9CF968
|
||||
|
||||
def test_single_byte(self) -> None:
|
||||
# Hand-computed: (basis ^ 0x61) * prime, masked to 32 bits.
|
||||
expected = ((0x811C9DC5 ^ 0x61) * 0x01000193) & 0xFFFFFFFF
|
||||
assert fnv1a_32(b"a") == expected
|
||||
|
||||
|
||||
class TestSelect:
|
||||
def test_empty_node_list_raises(self) -> None:
|
||||
with pytest.raises(NoAvailableNodeError):
|
||||
select("any-key", [])
|
||||
|
||||
def test_single_node_always_wins(self) -> None:
|
||||
only = NodeRef("solo", "http://solo")
|
||||
for key in ("a", "b", "00ff" + "0" * 28):
|
||||
assert select(key, [only]) is only
|
||||
|
||||
def test_deterministic(self) -> None:
|
||||
nodes = [NodeRef(f"n{i}", f"http://n{i}") for i in range(5)]
|
||||
key = "deadbeef" * 4
|
||||
first = select(key, nodes)
|
||||
for _ in range(20):
|
||||
assert select(key, nodes) is first
|
||||
|
||||
def test_independent_of_node_list_order(self) -> None:
|
||||
nodes = [NodeRef(f"n{i}", f"http://n{i}") for i in range(5)]
|
||||
key = "feedface" * 4
|
||||
forward = select(key, nodes)
|
||||
backward = select(key, list(reversed(nodes)))
|
||||
assert forward.node_id == backward.node_id
|
||||
|
||||
def test_distribution_roughly_uniform(self) -> None:
|
||||
nodes = [NodeRef(f"n{i}", f"http://n{i}") for i in range(4)]
|
||||
counts = {n.node_id: 0 for n in nodes}
|
||||
# Use sequential keys — 32 hex chars is what the router actually
|
||||
# passes in. Sequential isn't a problem because FNV-1a smears.
|
||||
for i in range(4000):
|
||||
key = f"{i:08x}" + "0" * 24
|
||||
counts[select(key, nodes).node_id] += 1
|
||||
# Each node should win ~25% (1000); allow ±15% drift.
|
||||
for c in counts.values():
|
||||
assert 850 < c < 1150, counts
|
||||
|
||||
|
||||
class TestMinimalMoves:
|
||||
def test_join_only_moves_to_new_node(self) -> None:
|
||||
old = [NodeRef(f"n{i}", f"http://n{i}") for i in range(3)]
|
||||
new = [*old, NodeRef("n3", "http://n3")]
|
||||
moved_correctly = 0
|
||||
moved_incorrectly = 0
|
||||
for i in range(2000):
|
||||
key = f"{i:08x}" + "0" * 24
|
||||
before = select(key, old).node_id
|
||||
after = select(key, new).node_id
|
||||
if before == after:
|
||||
continue
|
||||
if after == "n3":
|
||||
moved_correctly += 1
|
||||
else:
|
||||
moved_incorrectly += 1
|
||||
# Strict invariant: a join must never move a key between two
|
||||
# surviving nodes.
|
||||
assert moved_incorrectly == 0
|
||||
# Sanity: some keys did move.
|
||||
assert moved_correctly > 0
|
||||
|
||||
def test_leave_does_not_disturb_surviving_nodes(self) -> None:
|
||||
old = [NodeRef(f"n{i}", f"http://n{i}") for i in range(4)]
|
||||
new = old[:-1] # n3 leaves
|
||||
for i in range(2000):
|
||||
key = f"{i:08x}" + "0" * 24
|
||||
before = select(key, old).node_id
|
||||
after = select(key, new).node_id
|
||||
if before == "n3":
|
||||
# Must rehome to a survivor.
|
||||
assert after in {"n0", "n1", "n2"}
|
||||
else:
|
||||
# Must not move.
|
||||
assert after == before
|
||||
|
||||
|
||||
class TestWeights:
|
||||
def test_higher_weight_wins_more_often(self) -> None:
|
||||
nodes = [
|
||||
NodeRef("light", "http://l", weight=1),
|
||||
NodeRef("heavy", "http://h", weight=4),
|
||||
]
|
||||
on_heavy = 0
|
||||
for i in range(5000):
|
||||
key = f"{i:08x}" + "0" * 24
|
||||
if select(key, nodes).node_id == "heavy":
|
||||
on_heavy += 1
|
||||
# Heavy gets clearly more than half; tolerance for the simple
|
||||
# hash×weight formulation is wide.
|
||||
assert on_heavy / 5000 > 0.65
|
||||
|
||||
def test_zero_weight_clamped_to_one(self) -> None:
|
||||
# A weight-0 node still participates as if weight 1 — defended
|
||||
# at both NodeRef construction and _score(). Use the public
|
||||
# surface to sanity check.
|
||||
nodes = [
|
||||
NodeRef("a", "http://a", weight=0),
|
||||
NodeRef("b", "http://b", weight=0),
|
||||
]
|
||||
# Just confirms it doesn't divide-by-zero or score to 0.
|
||||
winner = select("any-key", nodes)
|
||||
assert winner.node_id in {"a", "b"}
|
||||
|
||||
|
||||
class TestSelectAll:
|
||||
def test_returns_all_nodes_in_score_order(self) -> None:
|
||||
nodes = [NodeRef(f"n{i}", f"http://n{i}") for i in range(4)]
|
||||
ranked = select_all("some-key", nodes)
|
||||
assert len(ranked) == 4
|
||||
assert {n.node_id for n in ranked} == {"n0", "n1", "n2", "n3"}
|
||||
# Top of the ranked list matches the single-select winner.
|
||||
assert ranked[0] is select("some-key", nodes)
|
||||
|
||||
def test_empty_list_returns_empty(self) -> None:
|
||||
assert select_all("any-key", []) == []
|
||||
@@ -1074,7 +1074,7 @@ CONSOLE_ENDPOINTS: list[EndpointSpec] = [
|
||||
EndpointSpec(
|
||||
"/v1/api/route/workstreams/new",
|
||||
"POST",
|
||||
"Create workstream via hash-ring routing proxy",
|
||||
"Create workstream via rendezvous routing proxy",
|
||||
response_model=RouteCreateResponse,
|
||||
error_codes=[400, 503],
|
||||
tags=["Routing"],
|
||||
|
||||
@@ -427,16 +427,17 @@ class ClusterCollector:
|
||||
for nid in lost_nodes:
|
||||
asyncio.run_coroutine_threadsafe(self._stop_node(nid), self._sse_loop)
|
||||
|
||||
# Notify the routing layer so it can refresh its hash-ring cache
|
||||
# when the rebalancer has published a new version.
|
||||
# Drive the router's cache from the collector's discovery
|
||||
# thread so the async route() handlers stay pure-in-memory.
|
||||
# The unconditional refresh also picks up admin-written
|
||||
# workstream_overrides between membership events.
|
||||
if self._router is not None:
|
||||
try:
|
||||
self._router.check_version()
|
||||
self._router.refresh_cache()
|
||||
except Exception:
|
||||
log.debug("Router version check failed", exc_info=True)
|
||||
# Update ring gauge metrics after version check
|
||||
log.debug("Router refresh failed", exc_info=True)
|
||||
if self._console_metrics is not None:
|
||||
self._console_metrics.set_ring_info(
|
||||
self._console_metrics.set_router_info(
|
||||
self._router.node_count(),
|
||||
self._router.version,
|
||||
)
|
||||
|
||||
@@ -57,7 +57,7 @@ class CoordinatorManager:
|
||||
|
||||
# Pseudo-node id persisted on coordinator rows so ``workstreams.node_id``
|
||||
# stays non-NULL and list / audit surfaces can distinguish coordinators
|
||||
# from real-node workstreams. The hash-ring router treats it as an
|
||||
# from real-node workstreams. Routed through the router as an
|
||||
# unroutable sentinel — coordinators never land on real nodes.
|
||||
# Bound from ``ClusterCollector.CONSOLE_PSEUDO_NODE_ID`` so the two
|
||||
# literals can't drift (the collector's eviction + query filters
|
||||
|
||||
@@ -8,10 +8,10 @@ from collections import defaultdict
|
||||
|
||||
|
||||
class ConsoleMetrics:
|
||||
"""Collects console routing and ring metrics in Prometheus text exposition format.
|
||||
"""Collects console routing and membership metrics in Prometheus text exposition format.
|
||||
|
||||
Lighter-weight than the server's MetricsCollector — tracks only router
|
||||
request counters, ring membership gauges, and rebalancer activity.
|
||||
Lighter-weight than the server's MetricsCollector — tracks router
|
||||
request counters and live-membership gauges.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
@@ -19,10 +19,8 @@ class ConsoleMetrics:
|
||||
self._router_requests: dict[tuple[str, str], int] = defaultdict(int)
|
||||
self._router_duration_sum: dict[str, float] = defaultdict(float)
|
||||
self._router_duration_count: dict[str, int] = defaultdict(int)
|
||||
self._ring_membership: int = 0
|
||||
self._ring_version: int = 0
|
||||
self._rebalance_total: dict[str, int] = defaultdict(int)
|
||||
self._migrations_total: int = 0
|
||||
self._router_membership: int = 0
|
||||
self._router_refresh_count: int = 0
|
||||
self._start_time: float = time.monotonic()
|
||||
|
||||
def record_route(self, method: str, status: int, duration: float) -> None:
|
||||
@@ -33,21 +31,11 @@ class ConsoleMetrics:
|
||||
self._router_duration_sum[method] += duration
|
||||
self._router_duration_count[method] += 1
|
||||
|
||||
def set_ring_info(self, membership: int, version: int) -> None:
|
||||
"""Update the current ring membership size and version."""
|
||||
def set_router_info(self, membership: int, refresh_count: int) -> None:
|
||||
"""Update current live-node count + the router's refresh counter."""
|
||||
with self._lock:
|
||||
self._ring_membership = membership
|
||||
self._ring_version = version
|
||||
|
||||
def record_rebalance(self, result: str) -> None:
|
||||
"""Record a rebalance pass outcome (noop/seeded/rebalanced)."""
|
||||
with self._lock:
|
||||
self._rebalance_total[result] += 1
|
||||
|
||||
def record_migrations(self, count: int) -> None:
|
||||
"""Record eager migration count from a rebalance pass."""
|
||||
with self._lock:
|
||||
self._migrations_total += count
|
||||
self._router_membership = membership
|
||||
self._router_refresh_count = refresh_count
|
||||
|
||||
def generate_text(self) -> str:
|
||||
"""Return Prometheus text exposition format (v0.0.4)."""
|
||||
@@ -57,10 +45,8 @@ class ConsoleMetrics:
|
||||
router_requests = dict(self._router_requests)
|
||||
duration_sum = dict(self._router_duration_sum)
|
||||
duration_count = dict(self._router_duration_count)
|
||||
ring_membership = self._ring_membership
|
||||
ring_version = self._ring_version
|
||||
rebalance_total = dict(self._rebalance_total)
|
||||
migrations_total = self._migrations_total
|
||||
router_membership = self._router_membership
|
||||
router_refresh_count = self._router_refresh_count
|
||||
|
||||
# turnstone_router_requests_total
|
||||
lines.append("# HELP turnstone_router_requests_total Console-routed requests")
|
||||
@@ -85,26 +71,17 @@ class ConsoleMetrics:
|
||||
f" {duration_count[method]}"
|
||||
)
|
||||
|
||||
# turnstone_ring_membership_size
|
||||
lines.append("# HELP turnstone_ring_membership_size Current ring node count")
|
||||
lines.append("# TYPE turnstone_ring_membership_size gauge")
|
||||
lines.append(f"turnstone_ring_membership_size {ring_membership}")
|
||||
# turnstone_router_membership_size
|
||||
lines.append("# HELP turnstone_router_membership_size Current live-node count")
|
||||
lines.append("# TYPE turnstone_router_membership_size gauge")
|
||||
lines.append(f"turnstone_router_membership_size {router_membership}")
|
||||
|
||||
# turnstone_ring_version
|
||||
lines.append("# HELP turnstone_ring_version Current ring version")
|
||||
lines.append("# TYPE turnstone_ring_version gauge")
|
||||
lines.append(f"turnstone_ring_version {ring_version}")
|
||||
|
||||
# turnstone_ring_rebalance_total
|
||||
lines.append("# HELP turnstone_ring_rebalance_total Rebalancer runs by result")
|
||||
lines.append("# TYPE turnstone_ring_rebalance_total counter")
|
||||
for result, count in sorted(rebalance_total.items()):
|
||||
lines.append(f'turnstone_ring_rebalance_total{{result="{result}"}} {count}')
|
||||
|
||||
# turnstone_ring_migrations_total
|
||||
lines.append("# HELP turnstone_ring_migrations_total Workstream migrations from rebalancer")
|
||||
lines.append("# TYPE turnstone_ring_migrations_total counter")
|
||||
lines.append(f"turnstone_ring_migrations_total {migrations_total}")
|
||||
# turnstone_router_refresh_total — bumped on every successful
|
||||
# cache refresh. A flat counter under churn means the
|
||||
# collector's discovery loop is stuck.
|
||||
lines.append("# HELP turnstone_router_refresh_total Router cache refresh counter")
|
||||
lines.append("# TYPE turnstone_router_refresh_total counter")
|
||||
lines.append(f"turnstone_router_refresh_total {router_refresh_count}")
|
||||
|
||||
lines.append("") # trailing newline
|
||||
return "\n".join(lines)
|
||||
|
||||
@@ -1,627 +0,0 @@
|
||||
"""Hash ring rebalancer — maintains bucket-to-node assignments.
|
||||
|
||||
Runs as a daemon thread inside the console process, following the same
|
||||
lifecycle pattern as ClusterCollector and TaskScheduler.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import structlog
|
||||
|
||||
from turnstone.core.hash_ring import RING_SIZE, RingNode, bucket_of
|
||||
from turnstone.core.storage._registry import StorageUnavailableError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from turnstone.console.collector import ClusterCollector
|
||||
from turnstone.console.metrics import ConsoleMetrics
|
||||
from turnstone.console.router import ConsoleRouter
|
||||
from turnstone.core.storage._protocol import StorageBackend
|
||||
|
||||
log = structlog.get_logger(__name__)
|
||||
|
||||
# States considered "active" for bucket stat reconciliation
|
||||
_ACTIVE_STATES = frozenset({"running", "thinking", "attention"})
|
||||
|
||||
|
||||
@dataclass
|
||||
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
|
||||
seeded: bool = False
|
||||
noop: bool = True
|
||||
|
||||
|
||||
class Rebalancer:
|
||||
"""Background daemon thread that maintains hash ring bucket assignments.
|
||||
|
||||
Uses the same lifecycle pattern as TaskScheduler: daemon thread, DB-based
|
||||
leader lock, periodic wake or event-driven trigger.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage: StorageBackend,
|
||||
router: ConsoleRouter | None = None,
|
||||
collector: ClusterCollector | None = None,
|
||||
console_metrics: ConsoleMetrics | None = None,
|
||||
interval: int = 60,
|
||||
threshold: float = 0.10,
|
||||
vnodes_per_unit: int = 150,
|
||||
lock_ttl: int = 120,
|
||||
eager_migrate: bool = False,
|
||||
api_token: str = "",
|
||||
token_manager: Any = None,
|
||||
) -> None:
|
||||
self._storage = storage
|
||||
self._router = router
|
||||
self._collector = collector
|
||||
self._console_metrics = console_metrics
|
||||
self._interval = interval
|
||||
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._token_manager = token_manager
|
||||
self._stop_event = threading.Event()
|
||||
self._trigger_event = threading.Event()
|
||||
self._thread: threading.Thread | None = None
|
||||
self._lock_owner = uuid.uuid4().hex
|
||||
self._last_result: RebalanceResult | None = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the rebalancer daemon thread."""
|
||||
self._stop_event.clear()
|
||||
self._trigger_event.clear()
|
||||
self._thread = threading.Thread(target=self._loop, daemon=True, name="rebalancer")
|
||||
self._thread.start()
|
||||
log.info("rebalancer.started", interval=self._interval)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the rebalancer and wait for the thread to finish."""
|
||||
self._stop_event.set()
|
||||
self._trigger_event.set() # wake the thread so it exits promptly
|
||||
if self._thread is not None:
|
||||
self._thread.join(timeout=5)
|
||||
log.info("rebalancer.stopped")
|
||||
|
||||
def trigger(self) -> None:
|
||||
"""Wake the rebalancer for an immediate check."""
|
||||
self._trigger_event.set()
|
||||
|
||||
def get_status(self) -> dict[str, Any]:
|
||||
"""Return current rebalancer status for the admin API."""
|
||||
raw = self._storage.get_system_setting("rebalancer_version", node_id="")
|
||||
version = 0
|
||||
if raw is not None:
|
||||
with contextlib.suppress(json.JSONDecodeError, TypeError, ValueError):
|
||||
version = int(json.loads(raw.get("value", "0")))
|
||||
result: dict[str, Any] = {
|
||||
"version": version,
|
||||
"is_leader": False,
|
||||
"last_result": None,
|
||||
}
|
||||
if self._last_result is not None:
|
||||
lr = self._last_result
|
||||
result["last_result"] = {
|
||||
"moves": lr.moves,
|
||||
"trigger": lr.trigger,
|
||||
"duration_ms": lr.duration_ms,
|
||||
"nodes": lr.nodes,
|
||||
"seeded": lr.seeded,
|
||||
"noop": lr.noop,
|
||||
}
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Main loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _loop(self) -> None:
|
||||
"""Main rebalancer loop — sleep or wait for trigger, then rebalance."""
|
||||
while not self._stop_event.is_set():
|
||||
self._trigger_event.wait(timeout=self._interval)
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
trigger = "triggered" if self._trigger_event.is_set() else "periodic"
|
||||
self._trigger_event.clear()
|
||||
if not self._try_acquire_lock():
|
||||
continue
|
||||
try:
|
||||
result = self.rebalance_once(trigger=trigger)
|
||||
self._last_result = result
|
||||
self._record_result_metrics(result)
|
||||
except StorageUnavailableError:
|
||||
pass # already logged by storage layer
|
||||
except Exception:
|
||||
log.exception("rebalancer.error")
|
||||
finally:
|
||||
self._release_lock()
|
||||
|
||||
def _record_result_metrics(self, result: RebalanceResult) -> None:
|
||||
"""Push rebalance result counters to the console metrics collector."""
|
||||
if self._console_metrics is None:
|
||||
return
|
||||
if result.seeded:
|
||||
self._console_metrics.record_rebalance("seeded")
|
||||
elif not result.noop:
|
||||
self._console_metrics.record_rebalance("rebalanced")
|
||||
else:
|
||||
self._console_metrics.record_rebalance("noop")
|
||||
if result.migrations > 0:
|
||||
self._console_metrics.record_migrations(result.migrations)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Leader lock (same pattern as TaskScheduler)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _try_acquire_lock(self) -> bool:
|
||||
"""Try to acquire the rebalancer lock via system_settings.
|
||||
|
||||
Uses a row with key ``rebalancer_lock``. The value is a JSON
|
||||
object ``{"owner": "<id>", "acquired": "<iso>"}``. Another
|
||||
instance's lock is considered expired when its timestamp is
|
||||
older than ``_lock_ttl`` seconds.
|
||||
|
||||
To reduce the TOCTOU window of a read-then-write approach, this
|
||||
method writes unconditionally and reads back to verify ownership.
|
||||
If two rebalancers race, one write wins and the loser sees the
|
||||
winner's value on read-back.
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
now_str = now.strftime("%Y-%m-%dT%H:%M:%S")
|
||||
|
||||
existing = self._storage.get_system_setting("rebalancer_lock")
|
||||
if existing is not None:
|
||||
try:
|
||||
lock_data = json.loads(existing.get("value", "{}"))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
lock_data = {}
|
||||
owner = lock_data.get("owner", "")
|
||||
acquired_str = lock_data.get("acquired", "")
|
||||
if owner != self._lock_owner and acquired_str:
|
||||
try:
|
||||
acquired_dt = datetime.strptime(acquired_str, "%Y-%m-%dT%H:%M:%S").replace(
|
||||
tzinfo=UTC
|
||||
)
|
||||
if (now - acquired_dt).total_seconds() < self._lock_ttl:
|
||||
return False # Another instance holds a valid lock
|
||||
except ValueError:
|
||||
pass # Malformed timestamp — take the lock
|
||||
|
||||
lock_value = json.dumps({"owner": self._lock_owner, "acquired": now_str})
|
||||
self._storage.upsert_system_setting("rebalancer_lock", lock_value)
|
||||
stored = self._storage.get_system_setting("rebalancer_lock")
|
||||
if stored is not None:
|
||||
try:
|
||||
data = json.loads(stored.get("value", "{}"))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return False
|
||||
return bool(data.get("owner") == self._lock_owner)
|
||||
return False
|
||||
|
||||
def _release_lock(self) -> None:
|
||||
"""Release the rebalancer lock if we still own it."""
|
||||
existing = self._storage.get_system_setting("rebalancer_lock")
|
||||
if existing is not None:
|
||||
try:
|
||||
lock_data = json.loads(existing.get("value", "{}"))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
lock_data = {}
|
||||
if lock_data.get("owner") == self._lock_owner:
|
||||
self._storage.delete_system_setting("rebalancer_lock")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Rebalance algorithm
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def rebalance_once(self, trigger: str = "periodic") -> RebalanceResult:
|
||||
"""Execute a single rebalance pass.
|
||||
|
||||
Returns a RebalanceResult describing what happened.
|
||||
"""
|
||||
t0 = time.monotonic()
|
||||
result = RebalanceResult(trigger=trigger)
|
||||
|
||||
# 1. Read live server nodes
|
||||
nodes_raw = self._storage.list_services("server", max_age_seconds=120)
|
||||
if not nodes_raw:
|
||||
result.duration_ms = (time.monotonic() - t0) * 1000
|
||||
return result
|
||||
|
||||
ring_nodes = _build_ring_nodes(nodes_raw)
|
||||
result.nodes = len(ring_nodes)
|
||||
|
||||
# 2. Read current bucket assignments
|
||||
current_rows = self._storage.list_ring_buckets()
|
||||
|
||||
# 3. If table is empty — first run, seed all 65536 buckets
|
||||
if not current_rows:
|
||||
assignments = _weight_based_assignments(ring_nodes)
|
||||
self._storage.seed_ring_buckets(assignments)
|
||||
new_version = self._bump_version()
|
||||
# Populate router cache directly from computed assignments
|
||||
# to avoid reading 65 536 rows back from DB.
|
||||
if self._router is not None:
|
||||
from turnstone.console.router import NodeRef
|
||||
|
||||
node_refs = {n.node_id: NodeRef(n.node_id, n.url) for n in ring_nodes}
|
||||
self._router.populate_from_assignments(assignments, node_refs, version=new_version)
|
||||
result.seeded = True
|
||||
result.noop = False
|
||||
result.duration_ms = (time.monotonic() - t0) * 1000
|
||||
log.info(
|
||||
"rebalancer.seeded",
|
||||
nodes=len(ring_nodes),
|
||||
buckets=len(assignments),
|
||||
)
|
||||
return result
|
||||
|
||||
# 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. Reconcile bucket_stats before computing transfer costs
|
||||
self._reconcile_bucket_stats()
|
||||
stats_rows = self._storage.list_bucket_stats()
|
||||
stats_map: dict[int, tuple[int, int]] = {}
|
||||
for s in stats_rows:
|
||||
stats_map[s["bucket"]] = (s["ws_count"], s["active_count"])
|
||||
|
||||
# 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)
|
||||
|
||||
# 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
|
||||
|
||||
# 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)
|
||||
|
||||
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
|
||||
return result
|
||||
|
||||
# 11. Execute moves: group by target node
|
||||
by_target: dict[str, list[int]] = defaultdict(list)
|
||||
for bucket, _from, to in filtered_moves:
|
||||
by_target[to].append(bucket)
|
||||
|
||||
total_moved = 0
|
||||
for target_node_id, bucket_list in by_target.items():
|
||||
total_moved += self._storage.assign_buckets(bucket_list, target_node_id)
|
||||
|
||||
# 12. Bump version and refresh cache
|
||||
self._bump_version()
|
||||
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),
|
||||
)
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _bump_version(self) -> int:
|
||||
"""Increment the rebalancer_version counter in system_settings.
|
||||
|
||||
Returns the new version number.
|
||||
|
||||
The read-then-write is safe because this method is only called while
|
||||
the leader lock is held (``_try_acquire_lock`` succeeded). Concurrent
|
||||
writers are prevented by the lock, so no CAS or timestamp trick is
|
||||
needed.
|
||||
"""
|
||||
raw = self._storage.get_system_setting("rebalancer_version", node_id="")
|
||||
version = 0
|
||||
if raw is not None:
|
||||
with contextlib.suppress(json.JSONDecodeError, TypeError, ValueError):
|
||||
version = int(json.loads(raw.get("value", "0")))
|
||||
new_version = version + 1
|
||||
self._storage.upsert_system_setting(
|
||||
"rebalancer_version", json.dumps(new_version), node_id=""
|
||||
)
|
||||
return new_version
|
||||
|
||||
def _reconcile_bucket_stats(self) -> None:
|
||||
"""Reconcile bucket_stats against actual workstream table data.
|
||||
|
||||
Self-heals counter drift from server crashes (a crashed server
|
||||
can't decrement its counters).
|
||||
"""
|
||||
ws_data = self._storage.list_workstream_routing_data()
|
||||
|
||||
# Compute actual per-bucket counts
|
||||
actual: dict[int, tuple[int, int]] = {} # bucket -> (ws_count, active_count)
|
||||
for ws_id, state in ws_data:
|
||||
if len(ws_id) < 4:
|
||||
continue
|
||||
bucket = bucket_of(ws_id)
|
||||
ws_count, active_count = actual.get(bucket, (0, 0))
|
||||
ws_count += 1
|
||||
if state in _ACTIVE_STATES:
|
||||
active_count += 1
|
||||
actual[bucket] = (ws_count, active_count)
|
||||
|
||||
# Load current stats
|
||||
stats_rows = self._storage.list_bucket_stats()
|
||||
stored: dict[int, tuple[int, int]] = {}
|
||||
for s in stats_rows:
|
||||
stored[s["bucket"]] = (s["ws_count"], s["active_count"])
|
||||
|
||||
# All buckets that appear in either set
|
||||
all_buckets = set(actual.keys()) | set(stored.keys())
|
||||
|
||||
for bucket in all_buckets:
|
||||
act = actual.get(bucket, (0, 0))
|
||||
sto = stored.get(bucket, (0, 0))
|
||||
if act != sto:
|
||||
# Pass current stored values to avoid re-querying the DB
|
||||
self._reset_bucket_stat(
|
||||
bucket,
|
||||
act[0],
|
||||
act[1],
|
||||
current_ws=sto[0],
|
||||
current_active=sto[1],
|
||||
)
|
||||
|
||||
def _reset_bucket_stat(
|
||||
self,
|
||||
bucket: int,
|
||||
ws_count: int,
|
||||
active_count: int,
|
||||
current_ws: int = 0,
|
||||
current_active: int = 0,
|
||||
) -> None:
|
||||
"""Reset a bucket_stats row to exact values via single upsert."""
|
||||
if (ws_count, active_count) == (current_ws, current_active):
|
||||
return # no change
|
||||
self._storage.set_bucket_stat(bucket, ws_count, active_count)
|
||||
|
||||
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._token_manager is not None:
|
||||
headers["Authorization"] = f"Bearer {self._token_manager.token}"
|
||||
elif 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 _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)
|
||||
# Compute per-node counts using the same int() + remainder distribution
|
||||
# as rebalance_once step 7, so seeding is a guaranteed noop on first rebalance.
|
||||
sorted_nodes = sorted(nodes, key=lambda n: n.node_id)
|
||||
counts: dict[str, int] = {}
|
||||
assigned = 0
|
||||
for n in sorted_nodes:
|
||||
c = int((n.weight / total_weight) * RING_SIZE)
|
||||
counts[n.node_id] = c
|
||||
assigned += c
|
||||
# Distribute remainder to heaviest nodes (same as rebalance_once step 7)
|
||||
remainder_pool = sorted(counts, key=lambda nid: counts[nid], reverse=True)
|
||||
for i in range(RING_SIZE - assigned):
|
||||
counts[remainder_pool[i % len(remainder_pool)]] += 1
|
||||
|
||||
assignments: list[tuple[int, str]] = []
|
||||
bucket = 0
|
||||
for node in sorted_nodes:
|
||||
for _ in range(counts[node.node_id]):
|
||||
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] = []
|
||||
for svc in services:
|
||||
meta_str = svc.get("metadata", "{}")
|
||||
try:
|
||||
meta = json.loads(meta_str)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
meta = {}
|
||||
weight = int(meta.get("weight", 1))
|
||||
if weight < 1:
|
||||
weight = 1
|
||||
nodes.append(RingNode(node_id=svc["service_id"], url=svc["url"], weight=weight))
|
||||
return nodes
|
||||
+149
-127
@@ -1,150 +1,151 @@
|
||||
"""Console routing layer — routes workstream requests to server nodes.
|
||||
|
||||
Maintains an in-memory flat array of 65536 bucket->NodeRef entries populated
|
||||
from the hash_ring_buckets table. Routing is O(1): cache[int(ws_id[:4], 16)].
|
||||
Uses rendezvous (HRW) hashing over the live ``services`` table. The
|
||||
routing function is a pure function of ``(ws_id, live_nodes)``: every
|
||||
reader given the same membership list produces the same answer, and
|
||||
``services.last_heartbeat`` is the single source of truth for both
|
||||
liveness and routing.
|
||||
|
||||
**Cache ownership**: the router's cache is push-driven by the
|
||||
collector's background discovery thread. ``route()`` and ``is_ready()``
|
||||
are pure in-memory lookups — they do not touch storage on the hot path.
|
||||
The collector calls ``refresh_cache()`` on every discovery tick and
|
||||
again immediately on observed membership changes (node_joined /
|
||||
node_lost). ``force_refresh()`` exists for the 404-retry path;
|
||||
callers must wrap it in ``asyncio.to_thread`` when invoking from an
|
||||
async handler so its DB read doesn't stall the event loop.
|
||||
|
||||
Per-route cost is O(N) hash computes — microseconds at typical cluster
|
||||
sizes, dwarfed by every downstream HTTP round-trip.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import secrets
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from turnstone.core.hash_ring import RING_SIZE, NoAvailableNodeError
|
||||
from turnstone.core.rendezvous import NoAvailableNodeError, NodeRef, select
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from turnstone.core.storage._protocol import StorageBackend
|
||||
|
||||
log = logging.getLogger("turnstone.console.router")
|
||||
# Brute-force attempt cap for ``generate_ws_id_for_node``. Expected
|
||||
# attempts for a weight-w_t target in a cluster with total weight W is
|
||||
# W/w_t (the target wins w_t/W of keys). At typical scale (N≤50,
|
||||
# weights ∈ {1..4}) the worst case is ~200 attempts; the cap is well
|
||||
# above that to absorb pathologically-skewed configurations without
|
||||
# spurious failures.
|
||||
_GENERATE_ATTEMPT_CAP = 65_536
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NodeRef:
|
||||
"""A server node that can receive proxied requests."""
|
||||
def _parse_weight(metadata_json: str) -> int:
|
||||
"""Pull the ``weight`` key out of a service-registry metadata blob.
|
||||
|
||||
node_id: str
|
||||
url: str
|
||||
A single corrupt row must not abort the cache refresh — fall back to
|
||||
weight=1 on any parse / type / value error, including JSON shapes
|
||||
that aren't dicts (``null``, lists, scalars).
|
||||
"""
|
||||
try:
|
||||
meta = json.loads(metadata_json or "{}")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return 1
|
||||
if not isinstance(meta, dict):
|
||||
return 1
|
||||
try:
|
||||
weight = int(meta.get("weight", 1))
|
||||
except (TypeError, ValueError):
|
||||
return 1
|
||||
return max(weight, 1)
|
||||
|
||||
|
||||
class ConsoleRouter:
|
||||
"""Routes workstream operations to the correct server node.
|
||||
|
||||
Maintains an in-memory flat array of 65536 bucket->NodeRef entries,
|
||||
populated from the hash_ring_buckets table. All routing is a
|
||||
single O(1) array lookup: ``cache[int(ws_id[:4], 16)]``.
|
||||
Thread-safe. All state mutation goes through ``_lock``; lookups
|
||||
snapshot the node list under the lock and run the rendezvous select
|
||||
outside it (the select is a pure function over an immutable list).
|
||||
"""
|
||||
|
||||
def __init__(self, storage: StorageBackend) -> None:
|
||||
self._storage = storage
|
||||
self._cache: list[NodeRef | None] = [None] * RING_SIZE
|
||||
self._lock = threading.Lock()
|
||||
self._nodes: list[NodeRef] = []
|
||||
self._overrides: dict[str, NodeRef] = {}
|
||||
self._version: int = 0
|
||||
self._refresh_lock = threading.Lock()
|
||||
# Monotonic counter bumped on every successful refresh — used by
|
||||
# the metrics gauge. Strictly increasing so dashboards can
|
||||
# detect when membership stops being refreshed.
|
||||
self._refresh_counter: int = 0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Cache management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def refresh_cache(self) -> bool:
|
||||
"""Reload the assignment cache from DB.
|
||||
"""Reload live-node list + overrides from storage.
|
||||
|
||||
Thread-safe: if another thread is already refreshing, this call
|
||||
returns False immediately (the other thread's refresh will apply).
|
||||
Returns True if the cache changed compared to the previous load.
|
||||
returns False immediately (the in-flight refresh will publish
|
||||
the latest state). Returns True if the membership changed
|
||||
compared to the previous load.
|
||||
|
||||
Called by the collector's discovery thread on every tick — never
|
||||
invoke from an async event-loop handler (this is a blocking DB
|
||||
read). Use ``force_refresh`` if you need a guaranteed-fresh
|
||||
view, and wrap that in ``asyncio.to_thread``.
|
||||
"""
|
||||
if not self._refresh_lock.acquire(blocking=False):
|
||||
return False # another thread is refreshing
|
||||
return False
|
||||
try:
|
||||
return self._refresh_cache_locked()
|
||||
return self._refresh_locked()
|
||||
finally:
|
||||
self._refresh_lock.release()
|
||||
|
||||
def _refresh_cache_locked(self) -> bool:
|
||||
"""Inner refresh — must be called with _refresh_lock held."""
|
||||
# Load node URLs from services table
|
||||
members = self._storage.list_services("server", max_age_seconds=120)
|
||||
nodes: dict[str, NodeRef] = {
|
||||
m["service_id"]: NodeRef(m["service_id"], m["url"]) for m in members
|
||||
}
|
||||
def force_refresh(self) -> bool:
|
||||
"""Refresh now, blocking if another refresh is in progress.
|
||||
|
||||
# Load bucket assignments into flat array
|
||||
buckets = self._storage.list_ring_buckets()
|
||||
new_cache: list[NodeRef | None] = [None] * RING_SIZE
|
||||
for row in buckets:
|
||||
ref = nodes.get(row["node_id"])
|
||||
if ref is not None:
|
||||
new_cache[row["bucket"]] = ref
|
||||
Used by the 404-retry path in the routing proxy when ``route()``
|
||||
sent the request to a node that doesn't have the workstream —
|
||||
the retry needs a guaranteed-fresh view of membership +
|
||||
overrides before giving up.
|
||||
|
||||
# Load per-workstream overrides (pinned workstreams)
|
||||
overrides = self._storage.list_workstream_overrides()
|
||||
new_overrides: dict[str, NodeRef] = {}
|
||||
for row in overrides:
|
||||
ref = nodes.get(row["node_id"])
|
||||
if ref is not None:
|
||||
new_overrides[row["ws_id"]] = ref
|
||||
|
||||
changed = new_cache != self._cache or new_overrides != self._overrides
|
||||
|
||||
# Atomic swap
|
||||
self._overrides = new_overrides
|
||||
self._cache = new_cache
|
||||
|
||||
return changed
|
||||
|
||||
def populate_from_assignments(
|
||||
self,
|
||||
assignments: list[tuple[int, str]],
|
||||
nodes: dict[str, NodeRef],
|
||||
*,
|
||||
version: int = 0,
|
||||
) -> None:
|
||||
"""Populate cache directly from computed assignments (no DB round-trip).
|
||||
|
||||
Used during initial seed to avoid a read-back of 65 536 rows.
|
||||
Overrides are loaded from DB since they may exist from a prior run
|
||||
(e.g. table was cleared but overrides survive). Setting *version*
|
||||
prevents ``check_version()`` from triggering an immediate refresh.
|
||||
Async callers must wrap this in ``asyncio.to_thread`` — the
|
||||
method takes a blocking lock and issues storage queries.
|
||||
"""
|
||||
new_cache: list[NodeRef | None] = [None] * RING_SIZE
|
||||
for bucket, node_id in assignments:
|
||||
ref = nodes.get(node_id)
|
||||
if ref is not None:
|
||||
new_cache[bucket] = ref
|
||||
|
||||
overrides = self._storage.list_workstream_overrides()
|
||||
new_overrides: dict[str, NodeRef] = {}
|
||||
for row in overrides:
|
||||
ref = nodes.get(row["node_id"])
|
||||
if ref is not None:
|
||||
new_overrides[row["ws_id"]] = ref
|
||||
|
||||
with self._refresh_lock:
|
||||
self._cache = new_cache
|
||||
return self._refresh_locked()
|
||||
|
||||
def _refresh_locked(self) -> bool:
|
||||
services = self._storage.list_services("server", max_age_seconds=120)
|
||||
new_nodes = sorted(
|
||||
(
|
||||
NodeRef(
|
||||
node_id=s["service_id"],
|
||||
url=s["url"],
|
||||
weight=_parse_weight(s.get("metadata", "{}")),
|
||||
)
|
||||
for s in services
|
||||
if s.get("service_id") and s.get("url")
|
||||
),
|
||||
key=lambda n: n.node_id,
|
||||
)
|
||||
nodes_by_id = {n.node_id: n for n in new_nodes}
|
||||
|
||||
overrides_rows = self._storage.list_workstream_overrides()
|
||||
new_overrides: dict[str, NodeRef] = {}
|
||||
for row in overrides_rows:
|
||||
ref = nodes_by_id.get(row["node_id"])
|
||||
if ref is not None:
|
||||
new_overrides[row["ws_id"]] = ref
|
||||
|
||||
with self._lock:
|
||||
changed = new_nodes != self._nodes or new_overrides != self._overrides
|
||||
self._nodes = new_nodes
|
||||
self._overrides = new_overrides
|
||||
self._version = version
|
||||
|
||||
def check_version(self) -> bool:
|
||||
"""Poll the rebalancer version and refresh if it changed.
|
||||
|
||||
Returns True if a refresh was triggered.
|
||||
"""
|
||||
setting = self._storage.get_system_setting("rebalancer_version", node_id="")
|
||||
if setting is not None:
|
||||
try:
|
||||
version = int(json.loads(setting.get("value", "0")))
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
version = 0
|
||||
else:
|
||||
version = 0
|
||||
|
||||
if version != self._version:
|
||||
self.refresh_cache()
|
||||
self._version = version
|
||||
return True
|
||||
return False
|
||||
self._refresh_counter += 1
|
||||
return changed
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Routing
|
||||
@@ -155,21 +156,21 @@ class ConsoleRouter:
|
||||
|
||||
Priority:
|
||||
1. Per-workstream override (pinned to a specific node)
|
||||
2. Bucket assignment (first 4 hex chars -> array index)
|
||||
2. Rendezvous (HRW) selection over the live-node list
|
||||
|
||||
Pure in-memory lookup — does not touch storage. Cache freshness
|
||||
is the collector's responsibility (see module docstring).
|
||||
"""
|
||||
ref = self._overrides.get(ws_id)
|
||||
if ref is not None:
|
||||
return ref
|
||||
if len(ws_id) < 4:
|
||||
raise NoAvailableNodeError(f"invalid ws_id: {ws_id!r}")
|
||||
try:
|
||||
bucket = int(ws_id[:4], 16)
|
||||
except ValueError:
|
||||
raise NoAvailableNodeError(f"invalid ws_id prefix: {ws_id[:4]!r}") from None
|
||||
ref = self._cache[bucket]
|
||||
if ref is None:
|
||||
raise NoAvailableNodeError(f"bucket {bucket} not assigned")
|
||||
return ref
|
||||
with self._lock:
|
||||
ref = self._overrides.get(ws_id)
|
||||
if ref is not None:
|
||||
return ref
|
||||
nodes = self._nodes # snapshot — list is replaced wholesale on refresh
|
||||
if not nodes:
|
||||
raise NoAvailableNodeError("no live nodes")
|
||||
if not ws_id:
|
||||
raise NoAvailableNodeError("invalid ws_id: empty")
|
||||
return select(ws_id, nodes)
|
||||
|
||||
def route_url(self, ws_id: str) -> str:
|
||||
"""Convenience — return just the URL for the target node."""
|
||||
@@ -180,29 +181,50 @@ class ConsoleRouter:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def is_ready(self) -> bool:
|
||||
"""Return True if at least one bucket is assigned."""
|
||||
return any(ref is not None for ref in self._cache)
|
||||
"""True if the router knows about at least one live node."""
|
||||
with self._lock:
|
||||
return bool(self._nodes)
|
||||
|
||||
def node_count(self) -> int:
|
||||
"""Number of distinct live nodes in the current view."""
|
||||
with self._lock:
|
||||
return len(self._nodes)
|
||||
|
||||
@property
|
||||
def version(self) -> int:
|
||||
"""The last seen rebalancer version."""
|
||||
return self._version
|
||||
"""Monotonic counter bumped on every successful cache refresh.
|
||||
|
||||
def node_count(self) -> int:
|
||||
"""Count distinct nodes present in the cache."""
|
||||
return len({ref.node_id for ref in self._cache if ref is not None})
|
||||
Surfaced by the collector's ``set_ring_info`` gauge so a
|
||||
dashboard can detect when membership stops being refreshed.
|
||||
Strictly increasing across the process lifetime.
|
||||
"""
|
||||
with self._lock:
|
||||
return self._refresh_counter
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Workstream ID generation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def generate_ws_id_for_node(self, node_id: str) -> str:
|
||||
"""Generate a routable workstream ID targeting *node_id*.
|
||||
"""Generate a 32-hex-char ws_id where rendezvous selects *node_id*.
|
||||
|
||||
The first 4 hex chars encode a bucket owned by the node; the
|
||||
remaining 28 hex chars are random (32 chars total).
|
||||
Brute-force loop: pick a random candidate, check whether HRW
|
||||
picks the target. Expected attempts ≈ ``W/w_t`` where ``W`` is
|
||||
total cluster weight and ``w_t`` is the target's weight. Cap
|
||||
at ``_GENERATE_ATTEMPT_CAP`` to bound worst case for skewed
|
||||
configurations.
|
||||
"""
|
||||
for bucket, ref in enumerate(self._cache):
|
||||
if ref is not None and ref.node_id == node_id:
|
||||
return f"{bucket:04x}" + secrets.token_hex(14)
|
||||
raise NoAvailableNodeError(f"no bucket assigned to node {node_id!r}")
|
||||
with self._lock:
|
||||
nodes = list(self._nodes)
|
||||
if not nodes:
|
||||
raise NoAvailableNodeError(f"no live node {node_id!r}")
|
||||
if not any(n.node_id == node_id for n in nodes):
|
||||
raise NoAvailableNodeError(f"no live node {node_id!r}")
|
||||
|
||||
for _ in range(_GENERATE_ATTEMPT_CAP):
|
||||
candidate = secrets.token_hex(16)
|
||||
if select(candidate, nodes).node_id == node_id:
|
||||
return candidate
|
||||
raise NoAvailableNodeError(
|
||||
f"could not generate ws_id targeting {node_id!r} after {_GENERATE_ATTEMPT_CAP} attempts"
|
||||
)
|
||||
|
||||
+28
-101
@@ -55,7 +55,7 @@ from turnstone.core.auth import (
|
||||
jwt_version_slot,
|
||||
require_permission,
|
||||
)
|
||||
from turnstone.core.hash_ring import NoAvailableNodeError
|
||||
from turnstone.core.rendezvous import NoAvailableNodeError
|
||||
from turnstone.core.skill_kind import SkillKind
|
||||
from turnstone.core.web_helpers import (
|
||||
read_json_or_400,
|
||||
@@ -1371,7 +1371,7 @@ async def create_workstream(request: Request) -> JSONResponse:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Route handlers — workstream routing proxy (hash-ring)
|
||||
# Route handlers — workstream routing proxy (rendezvous)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -1386,7 +1386,7 @@ def _record_route(
|
||||
|
||||
|
||||
async def route_create(request: Request) -> Response:
|
||||
"""POST /v1/api/route/workstreams/new — create via hash-ring routing.
|
||||
"""POST /v1/api/route/workstreams/new — create via rendezvous routing.
|
||||
|
||||
Accepts both `application/json` and `multipart/form-data`. Multipart
|
||||
callers must include ``?ws_id=<hex>`` in the URL query string so the
|
||||
@@ -1397,9 +1397,9 @@ async def route_create(request: Request) -> Response:
|
||||
router: ConsoleRouter | None = request.app.state.router
|
||||
ring_ready = router is not None and router.is_ready()
|
||||
if not ring_ready:
|
||||
# Ring not yet populated (rebalancer hasn't run or is disabled).
|
||||
# Try a one-shot refresh before giving up — the rebalancer may
|
||||
# have written buckets since the last collector poll.
|
||||
# Router cache empty — the collector hasn't published a
|
||||
# services list yet. One-shot refresh off the event loop
|
||||
# before giving up.
|
||||
if router is not None:
|
||||
await asyncio.to_thread(router.refresh_cache)
|
||||
ring_ready = router.is_ready()
|
||||
@@ -1426,7 +1426,7 @@ async def route_create(request: Request) -> Response:
|
||||
# Routing strategy is surfaced on the response so callers (the
|
||||
# coordinator's spawn_workstream tool especially) can explain why a
|
||||
# given node was chosen. Set on every branch below.
|
||||
routing_strategy = "hash_ring"
|
||||
routing_strategy = "rendezvous"
|
||||
|
||||
if is_multipart:
|
||||
# Multipart: caller must pass ws_id as a query param so we can
|
||||
@@ -1501,7 +1501,9 @@ async def route_create(request: Request) -> Response:
|
||||
ref = router.route(body["resume_ws"])
|
||||
routing_strategy = "resume"
|
||||
elif body.get("target_node"):
|
||||
ws_id = router.generate_ws_id_for_node(body["target_node"])
|
||||
# Brute-force HRW search can take up to _GENERATE_ATTEMPT_CAP
|
||||
# iterations for skewed weights; off the event loop.
|
||||
ws_id = await asyncio.to_thread(router.generate_ws_id_for_node, body["target_node"])
|
||||
body["ws_id"] = ws_id
|
||||
ref = router.route(ws_id)
|
||||
pin = True
|
||||
@@ -1595,12 +1597,12 @@ async def route_create(request: Request) -> Response:
|
||||
audit_ws_id = body.get("ws_id") or body.get("resume_ws", "") or ""
|
||||
# Return the storage-authoritative node_id so subsequent
|
||||
# inspect / list calls agree on the binding. ``ref.node_id`` is
|
||||
# the hash-ring target AT SPAWN TIME — stale once the
|
||||
# rebalancer runs or a node comes/goes, and the node's own
|
||||
# create handler is the source of truth for what node_id got
|
||||
# persisted on the workstream row. Fall back to ref.node_id
|
||||
# only when the storage lookup fails, matching the previous
|
||||
# behaviour so this change is strictly additive.
|
||||
# the rendezvous target AT SPAWN TIME — stale once membership
|
||||
# changes — and the node's own create handler is the source of
|
||||
# truth for what node_id got persisted on the workstream row.
|
||||
# Fall back to ref.node_id only when the storage lookup fails,
|
||||
# matching the previous behaviour so this change is strictly
|
||||
# additive.
|
||||
bound_node_id = ref.node_id
|
||||
storage = getattr(request.app.state, "auth_storage", None)
|
||||
if storage is not None and audit_ws_id:
|
||||
@@ -1633,7 +1635,7 @@ async def route_create(request: Request) -> Response:
|
||||
|
||||
|
||||
async def route_attachment_proxy(request: Request) -> Response:
|
||||
"""Proxy ws-id-keyed attachment endpoints through the hash-ring router.
|
||||
"""Proxy ws-id-keyed attachment endpoints through the router.
|
||||
|
||||
Handles all four shapes mounted under
|
||||
``/v1/api/route/workstreams/{ws_id}/attachments[/...]``:
|
||||
@@ -1846,19 +1848,16 @@ async def route_proxy(request: Request) -> Response:
|
||||
|
||||
# Transparent retry on 404 (at most once):
|
||||
#
|
||||
# The bucket-routed node doesn't have the workstream. Refresh the
|
||||
# cache (reloads overrides + bucket assignments from DB) and re-route.
|
||||
# If the route changed (e.g., a local-create override was added since
|
||||
# the last cache load), retry on the new node. If the route is the
|
||||
# same, return the 404 as-is — no loop, no scan.
|
||||
# The rendezvous-selected node doesn't have the workstream. Refresh
|
||||
# membership + overrides and re-route. If the route changed (e.g., a
|
||||
# local-create override was added since the last cache load, or a
|
||||
# node has joined / dropped), retry on the new node. If the route
|
||||
# is the same, return the 404 as-is — no loop, no scan.
|
||||
if resp.status_code == 404:
|
||||
# Blocking refresh — wait for any in-progress refresh to finish
|
||||
# so the retry uses the latest data, not stale cache.
|
||||
router._refresh_lock.acquire()
|
||||
try:
|
||||
router._refresh_cache_locked()
|
||||
finally:
|
||||
router._refresh_lock.release()
|
||||
# Off the event loop — force_refresh takes a blocking lock and
|
||||
# issues two storage queries. Coalesces internally so a 404
|
||||
# stampede after a node churn doesn't N×-multiply DB reads.
|
||||
await asyncio.to_thread(router.force_refresh)
|
||||
try:
|
||||
new_ref = router.route(ws_id)
|
||||
except (NoAvailableNodeError, ValueError):
|
||||
@@ -3764,7 +3763,7 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
audit_exec = ThreadPoolExecutor(max_workers=4, thread_name_prefix="coord-audit")
|
||||
app.state.audit_executor = audit_exec
|
||||
_set_audit_executor(audit_exec)
|
||||
# Populate hash-ring routing cache if a router is configured
|
||||
# Populate the router's services cache if a router is configured
|
||||
_router: ConsoleRouter | None = getattr(app.state, "router", None)
|
||||
if _router is not None:
|
||||
_router.refresh_cache()
|
||||
@@ -3779,10 +3778,6 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
scheduler = getattr(app.state, "scheduler", None)
|
||||
if scheduler is not None:
|
||||
scheduler.start()
|
||||
# Start rebalancer if configured
|
||||
_rebalancer = getattr(app.state, "rebalancer", None)
|
||||
if _rebalancer is not None:
|
||||
_rebalancer.start()
|
||||
# OIDC discovery (if configured)
|
||||
oidc_config = app.state.oidc_config
|
||||
if oidc_config.enabled:
|
||||
@@ -3982,9 +3977,6 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
await tls_mgr.stop_renewal()
|
||||
if scheduler is not None:
|
||||
scheduler.stop()
|
||||
_rebalancer = getattr(app.state, "rebalancer", None)
|
||||
if _rebalancer is not None:
|
||||
_rebalancer.stop()
|
||||
coord_mgr_shutdown = getattr(app.state, "coord_mgr", None)
|
||||
if coord_mgr_shutdown is not None:
|
||||
try:
|
||||
@@ -9912,35 +9904,6 @@ async def admin_delete_node_metadata_key(request: Request) -> JSONResponse:
|
||||
return JSONResponse({"ok": True})
|
||||
|
||||
|
||||
async def admin_ring_status(request: Request) -> JSONResponse:
|
||||
"""GET /v1/api/admin/ring/status — hash ring rebalancer status."""
|
||||
from turnstone.core.auth import require_permission
|
||||
|
||||
err = require_permission(request, "admin.settings")
|
||||
if err:
|
||||
return err
|
||||
rebalancer = getattr(request.app.state, "rebalancer", None)
|
||||
if rebalancer is None:
|
||||
return JSONResponse({"enabled": False})
|
||||
return JSONResponse({"enabled": True, **rebalancer.get_status()})
|
||||
|
||||
|
||||
async def admin_ring_rebalance(request: Request) -> JSONResponse:
|
||||
"""POST /v1/api/admin/ring/rebalance — trigger an immediate rebalance."""
|
||||
from turnstone.core.auth import require_permission
|
||||
|
||||
err = require_permission(request, "admin.settings")
|
||||
if err:
|
||||
return err
|
||||
rebalancer = getattr(request.app.state, "rebalancer", None)
|
||||
if rebalancer is None:
|
||||
return JSONResponse(
|
||||
{"status": "error", "reason": "rebalancer not enabled"}, status_code=503
|
||||
)
|
||||
rebalancer.trigger()
|
||||
return JSONResponse({"status": "ok"})
|
||||
|
||||
|
||||
async def tls_ca_status(request: Request) -> JSONResponse:
|
||||
"""GET /v1/api/admin/tls/ca — CA status."""
|
||||
from turnstone.core.auth import require_permission
|
||||
@@ -10089,7 +10052,6 @@ def create_app(
|
||||
tls_manager: Any = None,
|
||||
console_url: str = "",
|
||||
router: ConsoleRouter | None = None,
|
||||
rebalancer: Any = None,
|
||||
console_metrics: ConsoleMetrics | None = None,
|
||||
) -> Starlette:
|
||||
"""Build the Starlette ASGI application for the console dashboard."""
|
||||
@@ -10112,7 +10074,7 @@ def create_app(
|
||||
Route("/api/cluster/node/{node_id}", cluster_node_detail),
|
||||
Route("/api/cluster/snapshot", cluster_snapshot),
|
||||
Route("/api/cluster/events", cluster_events_sse),
|
||||
# Workstream routing (proxy to server nodes via hash ring)
|
||||
# Workstream routing (rendezvous proxy to server nodes)
|
||||
Route("/api/route/workstreams/new", route_create, methods=["POST"]),
|
||||
Route("/api/route/send", route_proxy, methods=["POST"]),
|
||||
Route("/api/route/approve", route_proxy, methods=["POST"]),
|
||||
@@ -10564,13 +10526,6 @@ def create_app(
|
||||
admin_set_node_metadata,
|
||||
methods=["PUT"],
|
||||
),
|
||||
# Hash ring
|
||||
Route("/api/admin/ring/status", admin_ring_status),
|
||||
Route(
|
||||
"/api/admin/ring/rebalance",
|
||||
admin_ring_rebalance,
|
||||
methods=["POST"],
|
||||
),
|
||||
# TLS / ACME
|
||||
Route("/api/admin/tls/ca", tls_ca_status),
|
||||
Route("/api/admin/tls/ca.pem", tls_ca_cert),
|
||||
@@ -10630,7 +10585,6 @@ def create_app(
|
||||
app.state.console_url = console_url
|
||||
app.state.tls_manager = tls_manager
|
||||
app.state.router = router
|
||||
app.state.rebalancer = rebalancer
|
||||
app.state.console_metrics = console_metrics or ConsoleMetrics()
|
||||
|
||||
# Mount ACME responder whenever a TLS manager is configured.
|
||||
@@ -10861,32 +10815,6 @@ def main() -> None:
|
||||
except Exception:
|
||||
log.debug("Failed to sync TLS state to ConfigStore", exc_info=True)
|
||||
|
||||
# Rebalancer — create if enabled in ConfigStore
|
||||
rebalancer = None
|
||||
if auth_storage:
|
||||
try:
|
||||
from turnstone.core.config_store import ConfigStore
|
||||
|
||||
_rcs = ConfigStore(auth_storage)
|
||||
if _rcs.get("rebalancer.enabled"):
|
||||
from turnstone.console.rebalancer import Rebalancer
|
||||
|
||||
rebalancer = Rebalancer(
|
||||
storage=auth_storage,
|
||||
router=router,
|
||||
collector=collector,
|
||||
console_metrics=console_metrics,
|
||||
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="",
|
||||
token_manager=proxy_token_mgr,
|
||||
)
|
||||
log.info("rebalancer.configured")
|
||||
except Exception:
|
||||
log.warning("Failed to configure rebalancer", exc_info=True)
|
||||
|
||||
app = create_app(
|
||||
collector=collector,
|
||||
jwt_secret=jwt_secret,
|
||||
@@ -10896,7 +10824,6 @@ def main() -> None:
|
||||
tls_manager=tls_mgr,
|
||||
console_url=console_url,
|
||||
router=router,
|
||||
rebalancer=rebalancer,
|
||||
console_metrics=console_metrics,
|
||||
)
|
||||
|
||||
|
||||
@@ -195,7 +195,6 @@ APPROVE_PATHS: frozenset[str] = frozenset(
|
||||
"/api/_internal/config-reload",
|
||||
"/api/_internal/mcp-reload",
|
||||
"/api/_internal/model-reload",
|
||||
"/api/_internal/migrate",
|
||||
}
|
||||
)
|
||||
ADMIN_PREFIX = "/api/admin/"
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
"""Hash ring routing primitives.
|
||||
|
||||
Constants, helpers, and data types for the bucket-based routing system.
|
||||
The consistent hash ring algorithm itself is documented in
|
||||
``docs/design/consistent-hash-ring.md`` as a reference design for future
|
||||
scalability work. The current rebalancer uses weight-proportional
|
||||
distribution instead (see ``turnstone/console/rebalancer.py``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
RING_SIZE = 65536 # 16-bit bucket space
|
||||
|
||||
|
||||
def bucket_of(ws_id: str) -> int:
|
||||
"""Extract the bucket from a workstream UUID.
|
||||
|
||||
ws_ids are hex strings (``secrets.token_hex``). The first 4 hex
|
||||
characters give us 16 bits = 65536 buckets. Since UUIDs are random,
|
||||
this is already uniformly distributed — no hash function needed.
|
||||
"""
|
||||
return int(ws_id[:4], 16)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RingNode:
|
||||
"""A physical node participating in the cluster."""
|
||||
|
||||
node_id: str
|
||||
url: str
|
||||
weight: int = 1
|
||||
|
||||
|
||||
class NoAvailableNodeError(Exception):
|
||||
"""Raised when routing fails (no nodes registered, bucket not assigned)."""
|
||||
@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from turnstone.core.hash_ring import bucket_of as _bucket_of
|
||||
from turnstone.core.log import get_logger
|
||||
from turnstone.core.storage import get_storage
|
||||
from turnstone.core.workstream import WorkstreamKind
|
||||
@@ -301,33 +300,6 @@ def update_workstream_state(ws_id: str, state: str) -> None:
|
||||
log.warning("Failed to update workstream state ws=%s state=%s", ws_id, state, exc_info=True)
|
||||
|
||||
|
||||
# -- Hash ring bucket counts --------------------------------------------------
|
||||
|
||||
|
||||
def increment_bucket_count(ws_id: str, active: bool = False) -> None:
|
||||
"""Fire-and-forget bucket count increment."""
|
||||
try:
|
||||
get_storage().increment_bucket_count(_bucket_of(ws_id), active)
|
||||
except Exception:
|
||||
log.warning("bucket count increment failed for %s", ws_id[:8], exc_info=True)
|
||||
|
||||
|
||||
def decrement_bucket_count(ws_id: str, active: bool = False) -> None:
|
||||
"""Fire-and-forget bucket count decrement."""
|
||||
try:
|
||||
get_storage().decrement_bucket_count(_bucket_of(ws_id), active)
|
||||
except Exception:
|
||||
log.warning("bucket count decrement failed for %s", ws_id[:8], exc_info=True)
|
||||
|
||||
|
||||
def adjust_bucket_active(ws_id: str, delta: int) -> None:
|
||||
"""Fire-and-forget active count adjustment."""
|
||||
try:
|
||||
get_storage().adjust_bucket_active(_bucket_of(ws_id), delta)
|
||||
except Exception:
|
||||
log.warning("bucket active adjust failed for %s", ws_id[:8], exc_info=True)
|
||||
|
||||
|
||||
def delete_workstream_override(ws_id: str) -> None:
|
||||
"""Fire-and-forget override deletion."""
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Rendezvous (Highest Random Weight, HRW) routing.
|
||||
|
||||
Pure function over ``(ws_id, live_nodes)`` → ``NodeRef``. Liveness is
|
||||
sourced from the ``services`` table; routing requires no extra
|
||||
persistent state.
|
||||
|
||||
Properties (per the standard HRW result):
|
||||
|
||||
- **Determinism** — every reader given the same membership list produces
|
||||
the same answer, no coordination required.
|
||||
- **Minimal moves** — when a node joins, only the keys for which it now
|
||||
scores highest move to it (~``1/N`` for equal weights). When a node
|
||||
leaves, only the keys it was winning fail over, distributed across
|
||||
the remaining nodes proportionally to their weight. Other keys do
|
||||
not move.
|
||||
- **Cross-language compatible** — the FNV-1a hash spec matches
|
||||
``docs/design/consistent-hash-ring.md`` so Go or TypeScript clients
|
||||
can compute the same routes.
|
||||
|
||||
Cost: O(N) per lookup where N is the live-node count, dominated by
|
||||
N FNV-1a hash computes. Negligible compared with any downstream HTTP
|
||||
round-trip.
|
||||
|
||||
Weighting: hash value is multiplied by the node weight rather than
|
||||
using the Skeena ``-weight / ln(uniform)`` formulation. The simpler
|
||||
form gives indistinguishable distribution at typical weight ranges
|
||||
(1-3) and avoids floating-point comparisons that would be a
|
||||
cross-language portability hazard. Switch to Skeena if a node ever
|
||||
ships with weight ≥ 10 alongside weight-1 peers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
_FNV_BASIS = 0x811C9DC5
|
||||
_FNV_PRIME = 0x01000193
|
||||
_MASK_32 = 0xFFFFFFFF
|
||||
|
||||
|
||||
def fnv1a_32(data: bytes) -> int:
|
||||
"""FNV-1a 32-bit hash.
|
||||
|
||||
Reference impl from ``docs/design/consistent-hash-ring.md`` — kept
|
||||
bit-identical so cross-language clients can compute the same routes.
|
||||
Test vectors:
|
||||
|
||||
>>> hex(fnv1a_32(b""))
|
||||
'0x811c9dc5'
|
||||
>>> hex(fnv1a_32(b"foobar"))
|
||||
'0xbf9cf968'
|
||||
"""
|
||||
h = _FNV_BASIS
|
||||
for b in data:
|
||||
h ^= b
|
||||
h = (h * _FNV_PRIME) & _MASK_32
|
||||
return h
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NodeRef:
|
||||
"""A live node eligible for routing."""
|
||||
|
||||
node_id: str
|
||||
url: str
|
||||
weight: int = 1
|
||||
|
||||
|
||||
class NoAvailableNodeError(Exception):
|
||||
"""Raised when the live-node list is empty."""
|
||||
|
||||
|
||||
def _score(node_id: str, key: str, weight: int) -> int:
|
||||
"""HRW score — higher wins.
|
||||
|
||||
Hash key is ``"{node_id}\\x00{key}"`` — the NUL separator prevents
|
||||
boundary collisions (e.g. ``("ab", "cd")`` vs ``("a", "bcd")``).
|
||||
"""
|
||||
payload = f"{node_id}\x00{key}".encode()
|
||||
return fnv1a_32(payload) * max(weight, 1)
|
||||
|
||||
|
||||
def select(key: str, nodes: list[NodeRef]) -> NodeRef:
|
||||
"""Pick the highest-scoring node for *key*.
|
||||
|
||||
Tie-breaks on ``node_id`` lexicographically descending (the
|
||||
higher-sorted ``node_id`` wins) so behavior stays deterministic if
|
||||
two nodes happen to score identically — vanishingly rare with
|
||||
32-bit hashes but worth pinning for tests.
|
||||
"""
|
||||
if not nodes:
|
||||
raise NoAvailableNodeError("no live nodes")
|
||||
return max(
|
||||
nodes,
|
||||
key=lambda n: (_score(n.node_id, key, n.weight), n.node_id),
|
||||
)
|
||||
|
||||
|
||||
def select_all(key: str, nodes: list[NodeRef]) -> list[NodeRef]:
|
||||
"""Return *all* nodes ranked by score, highest first.
|
||||
|
||||
Tie-break matches ``select`` — lexicographically descending on
|
||||
``node_id`` so the top of the list always equals ``select(...)``.
|
||||
Used by callers that want a pre-computed fail-over list — e.g. a
|
||||
retry policy that, on connect failure to the primary, falls through
|
||||
to the second-highest scorer without re-running selection.
|
||||
"""
|
||||
if not nodes:
|
||||
return []
|
||||
return sorted(
|
||||
nodes,
|
||||
key=lambda n: (_score(n.node_id, key, n.weight), n.node_id),
|
||||
reverse=True,
|
||||
)
|
||||
@@ -612,75 +612,18 @@ def _build_registry() -> dict[str, SettingDef]:
|
||||
"directory) to get a publicly trusted certificate for the console's HTTPS endpoint. "
|
||||
"Leave empty to self-issue from the internal CA (use when behind a reverse proxy).",
|
||||
),
|
||||
# -- ring ---------------------------------------------------------------
|
||||
SettingDef(
|
||||
"ring.vnodes_per_unit",
|
||||
"int",
|
||||
150,
|
||||
"Virtual nodes per unit weight on the hash ring",
|
||||
"ring",
|
||||
min_value=10,
|
||||
max_value=1000,
|
||||
help="Controls the granularity of the consistent hash ring. Higher values give a "
|
||||
"more uniform distribution of buckets to nodes at the cost of slightly more memory. "
|
||||
"Each physical node gets weight * vnodes_per_unit virtual positions on the ring.",
|
||||
),
|
||||
# -- rebalancer ---------------------------------------------------------
|
||||
SettingDef(
|
||||
"rebalancer.enabled",
|
||||
"bool",
|
||||
True,
|
||||
"Enable the hash ring rebalancer daemon",
|
||||
"rebalancer",
|
||||
help="When enabled, the console runs a background thread that monitors cluster "
|
||||
"membership and automatically redistributes hash ring buckets when nodes join "
|
||||
"or leave. Required for the channel gateway and multi-node routing.",
|
||||
),
|
||||
SettingDef(
|
||||
"rebalancer.interval",
|
||||
"int",
|
||||
60,
|
||||
"Rebalancer check interval in seconds",
|
||||
"rebalancer",
|
||||
min_value=10,
|
||||
max_value=3600,
|
||||
help="How often the rebalancer wakes up to check whether bucket assignments "
|
||||
"need updating. It also wakes immediately on membership changes.",
|
||||
),
|
||||
SettingDef(
|
||||
"rebalancer.threshold",
|
||||
"float",
|
||||
0.10,
|
||||
"Imbalance threshold before rebalancing (0.0\u20131.0)",
|
||||
"rebalancer",
|
||||
min_value=0.01,
|
||||
max_value=0.50,
|
||||
help="Minimum deviation from the ideal distribution before buckets are moved. "
|
||||
"A value of 0.10 means 10% deviation triggers rebalancing. Lower values keep "
|
||||
"the cluster more balanced but cause more frequent bucket moves.",
|
||||
),
|
||||
SettingDef(
|
||||
"rebalancer.eager_migrate",
|
||||
"bool",
|
||||
False,
|
||||
"Eagerly migrate active workstreams after rebalance",
|
||||
"rebalancer",
|
||||
help="When enabled, the rebalancer asks source nodes to evict workstreams whose "
|
||||
"buckets have been reassigned. When disabled (default), workstreams migrate lazily "
|
||||
"on the next request — the old copy is eventually evicted by idle timeout.",
|
||||
),
|
||||
# -- node ---------------------------------------------------------------
|
||||
SettingDef(
|
||||
"node.weight",
|
||||
"int",
|
||||
1,
|
||||
"Node weight for hash ring distribution",
|
||||
"Node weight for rendezvous routing",
|
||||
"node",
|
||||
min_value=1,
|
||||
max_value=100,
|
||||
help="Relative capacity of this server node. A node with weight 2 receives "
|
||||
"roughly twice as many bucket assignments (and therefore workstreams) as a "
|
||||
"node with weight 1.",
|
||||
help="Relative capacity of this server node. A node with weight 2 wins "
|
||||
"roughly twice as many ws_id rendezvous selections (and therefore "
|
||||
"receives twice as many workstreams) as a node with weight 1.",
|
||||
),
|
||||
# -- coordinator --------------------------------------------------------
|
||||
SettingDef(
|
||||
|
||||
@@ -16,11 +16,9 @@ from turnstone.core.log import get_logger
|
||||
from turnstone.core.storage._schema import (
|
||||
api_tokens,
|
||||
audit_events,
|
||||
bucket_stats,
|
||||
channel_routes,
|
||||
channel_users,
|
||||
conversations,
|
||||
hash_ring_buckets,
|
||||
heuristic_rules,
|
||||
intent_verdicts,
|
||||
mcp_servers,
|
||||
@@ -1851,122 +1849,7 @@ class PostgreSQLBackend:
|
||||
rows = conn.execute(stmt).fetchall()
|
||||
return {r[0] for r in rows}
|
||||
|
||||
# -- Hash ring routing -----------------------------------------------------
|
||||
|
||||
def list_ring_buckets(self) -> list[dict[str, Any]]:
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
sa.select(hash_ring_buckets).order_by(hash_ring_buckets.c.bucket)
|
||||
).fetchall()
|
||||
return [dict(r._mapping) for r in rows]
|
||||
|
||||
def seed_ring_buckets(self, assignments: list[tuple[int, str]]) -> None:
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
chunk_size = 16_000 # 2 params/row × 16k = 32k, within psycopg 65 535 limit
|
||||
with self._conn() as conn:
|
||||
for i in range(0, len(assignments), chunk_size):
|
||||
chunk = assignments[i : i + chunk_size]
|
||||
stmt = pg_insert(hash_ring_buckets).values(
|
||||
[{"bucket": b, "node_id": n} for b, n in chunk]
|
||||
)
|
||||
stmt = stmt.on_conflict_do_nothing(index_elements=[hash_ring_buckets.c.bucket])
|
||||
conn.execute(stmt)
|
||||
conn.commit()
|
||||
|
||||
def assign_buckets(self, buckets: list[int], node_id: str) -> int:
|
||||
if not buckets:
|
||||
return 0
|
||||
# De-duplicate so rowcount stays accurate across chunks.
|
||||
buckets = list(dict.fromkeys(buckets))
|
||||
# psycopg limits query parameters to 65 535; chunk to stay well under.
|
||||
chunk_size = 10_000
|
||||
total = 0
|
||||
with self._conn() as conn:
|
||||
for i in range(0, len(buckets), chunk_size):
|
||||
chunk = buckets[i : i + chunk_size]
|
||||
result = conn.execute(
|
||||
sa.update(hash_ring_buckets)
|
||||
.where(hash_ring_buckets.c.bucket.in_(chunk))
|
||||
.values(node_id=node_id)
|
||||
)
|
||||
total += result.rowcount
|
||||
conn.commit()
|
||||
return total
|
||||
|
||||
def increment_bucket_count(self, bucket: int, active: bool = False) -> None:
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
with self._conn() as conn:
|
||||
stmt = pg_insert(bucket_stats).values(
|
||||
bucket=bucket,
|
||||
ws_count=1,
|
||||
active_count=1 if active else 0,
|
||||
)
|
||||
set_: dict[str, Any] = {"ws_count": bucket_stats.c.ws_count + 1}
|
||||
if active:
|
||||
set_["active_count"] = bucket_stats.c.active_count + 1
|
||||
stmt = stmt.on_conflict_do_update(index_elements=[bucket_stats.c.bucket], set_=set_)
|
||||
conn.execute(stmt)
|
||||
conn.commit()
|
||||
|
||||
def decrement_bucket_count(self, bucket: int, active: bool = False) -> None:
|
||||
vals: dict[str, Any] = {
|
||||
"ws_count": sa.case(
|
||||
(bucket_stats.c.ws_count > 0, bucket_stats.c.ws_count - 1),
|
||||
else_=0,
|
||||
)
|
||||
}
|
||||
if active:
|
||||
vals["active_count"] = sa.case(
|
||||
(bucket_stats.c.active_count > 0, bucket_stats.c.active_count - 1),
|
||||
else_=0,
|
||||
)
|
||||
with self._conn() as conn:
|
||||
conn.execute(
|
||||
sa.update(bucket_stats).where(bucket_stats.c.bucket == bucket).values(**vals)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def adjust_bucket_active(self, bucket: int, delta: int) -> None:
|
||||
with self._conn() as conn:
|
||||
conn.execute(
|
||||
sa.update(bucket_stats)
|
||||
.where(bucket_stats.c.bucket == bucket)
|
||||
.values(
|
||||
active_count=sa.case(
|
||||
(
|
||||
bucket_stats.c.active_count + sa.literal(delta) >= 0,
|
||||
bucket_stats.c.active_count + sa.literal(delta),
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def list_bucket_stats(self) -> list[dict[str, Any]]:
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
sa.select(bucket_stats)
|
||||
.where(bucket_stats.c.ws_count > 0)
|
||||
.order_by(bucket_stats.c.bucket)
|
||||
).fetchall()
|
||||
return [dict(r._mapping) for r in rows]
|
||||
|
||||
def set_bucket_stat(self, bucket: int, ws_count: int, active_count: int) -> None:
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
stmt = pg_insert(bucket_stats).values(
|
||||
bucket=bucket, ws_count=ws_count, active_count=active_count
|
||||
)
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
index_elements=[bucket_stats.c.bucket],
|
||||
set_={"ws_count": ws_count, "active_count": active_count},
|
||||
)
|
||||
with self._conn() as conn:
|
||||
conn.execute(stmt)
|
||||
conn.commit()
|
||||
# -- Routing overrides -----------------------------------------------------
|
||||
|
||||
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
@@ -1998,16 +1881,6 @@ class PostgreSQLBackend:
|
||||
).fetchall()
|
||||
return [dict(r._mapping) for r in rows]
|
||||
|
||||
def list_workstream_routing_data(self) -> list[tuple[str, str]]:
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
sa.text(
|
||||
"SELECT w.ws_id, w.state FROM workstreams w "
|
||||
"WHERE w.ws_id NOT IN (SELECT ws_id FROM workstream_overrides)"
|
||||
)
|
||||
).fetchall()
|
||||
return [(r[0], r[1]) for r in rows]
|
||||
|
||||
# -- Roles -----------------------------------------------------------------
|
||||
|
||||
def create_role(
|
||||
|
||||
@@ -787,39 +787,7 @@ class StorageBackend(Protocol):
|
||||
"""Return node_ids where ALL key=value filters match (exact match)."""
|
||||
...
|
||||
|
||||
# -- Hash ring routing ---
|
||||
|
||||
def list_ring_buckets(self) -> list[dict[str, Any]]:
|
||||
"""Return all rows from hash_ring_buckets. Empty if not seeded."""
|
||||
...
|
||||
|
||||
def seed_ring_buckets(self, assignments: list[tuple[int, str]]) -> None:
|
||||
"""Insert (bucket, node_id) rows. Idempotent (ON CONFLICT DO NOTHING)."""
|
||||
...
|
||||
|
||||
def assign_buckets(self, buckets: list[int], node_id: str) -> int:
|
||||
"""Reassign buckets to a node. Returns rows updated."""
|
||||
...
|
||||
|
||||
def increment_bucket_count(self, bucket: int, active: bool = False) -> None:
|
||||
"""Increment ws_count (and active_count if active) in bucket_stats. Upserts."""
|
||||
...
|
||||
|
||||
def decrement_bucket_count(self, bucket: int, active: bool = False) -> None:
|
||||
"""Decrement ws_count (and active_count if active). Clamps at zero."""
|
||||
...
|
||||
|
||||
def adjust_bucket_active(self, bucket: int, delta: int) -> None:
|
||||
"""Adjust active_count only (not ws_count). For state transitions."""
|
||||
...
|
||||
|
||||
def list_bucket_stats(self) -> list[dict[str, Any]]:
|
||||
"""Return all bucket_stats rows with ws_count > 0."""
|
||||
...
|
||||
|
||||
def set_bucket_stat(self, bucket: int, ws_count: int, active_count: int) -> None:
|
||||
"""Set bucket_stats to exact values. Upserts the row."""
|
||||
...
|
||||
# -- Routing overrides ---
|
||||
|
||||
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
||||
"""Pin a workstream to a specific node. Upserts."""
|
||||
@@ -833,10 +801,6 @@ class StorageBackend(Protocol):
|
||||
"""Return all overrides."""
|
||||
...
|
||||
|
||||
def list_workstream_routing_data(self) -> list[tuple[str, str]]:
|
||||
"""Return (ws_id, state) for all non-override workstreams."""
|
||||
...
|
||||
|
||||
# -- Roles (RBAC) ----------------------------------------------------------
|
||||
|
||||
def create_role(
|
||||
|
||||
@@ -262,26 +262,9 @@ node_metadata = sa.Table(
|
||||
sa.Index("idx_node_metadata_key", node_metadata.c.key)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Hash ring routing tables
|
||||
# Routing — per-workstream pinning overrides
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
hash_ring_buckets = sa.Table(
|
||||
"hash_ring_buckets",
|
||||
metadata,
|
||||
sa.Column("bucket", sa.Integer, primary_key=True),
|
||||
sa.Column("node_id", sa.Text, nullable=False),
|
||||
)
|
||||
|
||||
sa.Index("idx_ring_buckets_node", hash_ring_buckets.c.node_id)
|
||||
|
||||
bucket_stats = sa.Table(
|
||||
"bucket_stats",
|
||||
metadata,
|
||||
sa.Column("bucket", sa.Integer, primary_key=True),
|
||||
sa.Column("ws_count", sa.Integer, nullable=False, server_default="0"),
|
||||
sa.Column("active_count", sa.Integer, nullable=False, server_default="0"),
|
||||
)
|
||||
|
||||
workstream_overrides = sa.Table(
|
||||
"workstream_overrides",
|
||||
metadata,
|
||||
|
||||
@@ -16,11 +16,9 @@ from turnstone.core.log import get_logger
|
||||
from turnstone.core.storage._schema import (
|
||||
api_tokens,
|
||||
audit_events,
|
||||
bucket_stats,
|
||||
channel_routes,
|
||||
channel_users,
|
||||
conversations,
|
||||
hash_ring_buckets,
|
||||
heuristic_rules,
|
||||
intent_verdicts,
|
||||
mcp_servers,
|
||||
@@ -1965,122 +1963,7 @@ class SQLiteBackend:
|
||||
rows = conn.execute(stmt).fetchall()
|
||||
return {r[0] for r in rows}
|
||||
|
||||
# -- Hash ring routing -----------------------------------------------------
|
||||
|
||||
def list_ring_buckets(self) -> list[dict[str, Any]]:
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
sa.select(hash_ring_buckets).order_by(hash_ring_buckets.c.bucket)
|
||||
).fetchall()
|
||||
return [dict(r._mapping) for r in rows]
|
||||
|
||||
def seed_ring_buckets(self, assignments: list[tuple[int, str]]) -> None:
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
chunk_size = 8_000 # 2 params/row × 8k = 16k, within SQLite 3.32+ limit (32 766)
|
||||
with self._conn() as conn:
|
||||
for i in range(0, len(assignments), chunk_size):
|
||||
chunk = assignments[i : i + chunk_size]
|
||||
stmt = sqlite_insert(hash_ring_buckets).values(
|
||||
[{"bucket": b, "node_id": n} for b, n in chunk]
|
||||
)
|
||||
stmt = stmt.on_conflict_do_nothing(index_elements=["bucket"])
|
||||
conn.execute(stmt)
|
||||
conn.commit()
|
||||
|
||||
def assign_buckets(self, buckets: list[int], node_id: str) -> int:
|
||||
if not buckets:
|
||||
return 0
|
||||
# De-duplicate so rowcount stays accurate across chunks.
|
||||
buckets = list(dict.fromkeys(buckets))
|
||||
# SQLite default SQLITE_MAX_VARIABLE_NUMBER is 999; chunk conservatively.
|
||||
chunk_size = 500
|
||||
total = 0
|
||||
with self._conn() as conn:
|
||||
for i in range(0, len(buckets), chunk_size):
|
||||
chunk = buckets[i : i + chunk_size]
|
||||
result = conn.execute(
|
||||
sa.update(hash_ring_buckets)
|
||||
.where(hash_ring_buckets.c.bucket.in_(chunk))
|
||||
.values(node_id=node_id)
|
||||
)
|
||||
total += result.rowcount
|
||||
conn.commit()
|
||||
return total
|
||||
|
||||
def increment_bucket_count(self, bucket: int, active: bool = False) -> None:
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
with self._conn() as conn:
|
||||
stmt = sqlite_insert(bucket_stats).values(
|
||||
bucket=bucket,
|
||||
ws_count=1,
|
||||
active_count=1 if active else 0,
|
||||
)
|
||||
set_ = {"ws_count": bucket_stats.c.ws_count + 1}
|
||||
if active:
|
||||
set_["active_count"] = bucket_stats.c.active_count + 1
|
||||
stmt = stmt.on_conflict_do_update(index_elements=["bucket"], set_=set_)
|
||||
conn.execute(stmt)
|
||||
conn.commit()
|
||||
|
||||
def decrement_bucket_count(self, bucket: int, active: bool = False) -> None:
|
||||
vals: dict[str, Any] = {
|
||||
"ws_count": sa.case(
|
||||
(bucket_stats.c.ws_count > 0, bucket_stats.c.ws_count - 1),
|
||||
else_=0,
|
||||
)
|
||||
}
|
||||
if active:
|
||||
vals["active_count"] = sa.case(
|
||||
(bucket_stats.c.active_count > 0, bucket_stats.c.active_count - 1),
|
||||
else_=0,
|
||||
)
|
||||
with self._conn() as conn:
|
||||
conn.execute(
|
||||
sa.update(bucket_stats).where(bucket_stats.c.bucket == bucket).values(**vals)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def adjust_bucket_active(self, bucket: int, delta: int) -> None:
|
||||
with self._conn() as conn:
|
||||
conn.execute(
|
||||
sa.update(bucket_stats)
|
||||
.where(bucket_stats.c.bucket == bucket)
|
||||
.values(
|
||||
active_count=sa.case(
|
||||
(
|
||||
bucket_stats.c.active_count + sa.literal(delta) >= 0,
|
||||
bucket_stats.c.active_count + sa.literal(delta),
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def list_bucket_stats(self) -> list[dict[str, Any]]:
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
sa.select(bucket_stats)
|
||||
.where(bucket_stats.c.ws_count > 0)
|
||||
.order_by(bucket_stats.c.bucket)
|
||||
).fetchall()
|
||||
return [dict(r._mapping) for r in rows]
|
||||
|
||||
def set_bucket_stat(self, bucket: int, ws_count: int, active_count: int) -> None:
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
stmt = sqlite_insert(bucket_stats).values(
|
||||
bucket=bucket, ws_count=ws_count, active_count=active_count
|
||||
)
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
index_elements=["bucket"],
|
||||
set_={"ws_count": ws_count, "active_count": active_count},
|
||||
)
|
||||
with self._conn() as conn:
|
||||
conn.execute(stmt)
|
||||
conn.commit()
|
||||
# -- Routing overrides -----------------------------------------------------
|
||||
|
||||
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
@@ -2112,16 +1995,6 @@ class SQLiteBackend:
|
||||
).fetchall()
|
||||
return [dict(r._mapping) for r in rows]
|
||||
|
||||
def list_workstream_routing_data(self) -> list[tuple[str, str]]:
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
sa.text(
|
||||
"SELECT w.ws_id, w.state FROM workstreams w "
|
||||
"WHERE w.ws_id NOT IN (SELECT ws_id FROM workstream_overrides)"
|
||||
)
|
||||
).fetchall()
|
||||
return [(r[0], r[1]) for r in rows]
|
||||
|
||||
# -- Roles -----------------------------------------------------------------
|
||||
|
||||
def create_role(
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Drop hash ring bucket tables.
|
||||
|
||||
Routing is now rendezvous (HRW) hashing over the live ``services``
|
||||
table; ``hash_ring_buckets`` and ``bucket_stats`` are no longer read
|
||||
or written. ``workstream_overrides`` stays — manual per-ws pinning
|
||||
still takes priority over the rendezvous select.
|
||||
|
||||
Also clears the ``rebalancer_version`` and ``rebalancer_lock`` rows in
|
||||
``system_settings`` so they don't linger as orphan keys.
|
||||
|
||||
Revision ID: 046
|
||||
Revises: 045
|
||||
Create Date: 2026-04-18
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "046"
|
||||
down_revision = "045"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.drop_index("idx_ring_buckets_node", table_name="hash_ring_buckets")
|
||||
op.drop_table("hash_ring_buckets")
|
||||
op.drop_table("bucket_stats")
|
||||
bind = op.get_bind()
|
||||
bind.execute(
|
||||
sa.text(
|
||||
"DELETE FROM system_settings WHERE key IN ('rebalancer_version', 'rebalancer_lock')"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.create_table(
|
||||
"hash_ring_buckets",
|
||||
sa.Column("bucket", sa.Integer, primary_key=True),
|
||||
sa.Column("node_id", sa.Text, nullable=False),
|
||||
)
|
||||
op.create_index("idx_ring_buckets_node", "hash_ring_buckets", ["node_id"])
|
||||
op.create_table(
|
||||
"bucket_stats",
|
||||
sa.Column("bucket", sa.Integer, primary_key=True),
|
||||
sa.Column("ws_count", sa.Integer, nullable=False, server_default="0"),
|
||||
sa.Column("active_count", sa.Integer, nullable=False, server_default="0"),
|
||||
)
|
||||
@@ -93,9 +93,6 @@ class WorkstreamState(enum.Enum):
|
||||
ERROR = "error" # last operation failed
|
||||
|
||||
|
||||
_ACTIVE_STATES = {"running", "thinking", "attention"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Workstream dataclass
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -251,11 +248,9 @@ class WorkstreamManager:
|
||||
if first_evicted is not None:
|
||||
self._cleanup_ui(first_evicted)
|
||||
self._last_evicted = first_evicted
|
||||
from turnstone.core.memory import decrement_bucket_count as _dbc1
|
||||
from turnstone.core.memory import delete_workstream_override as _dwo1
|
||||
from turnstone.core.metrics import metrics as _m1
|
||||
|
||||
_dbc1(first_evicted.id, active=False) # evicted ws is always idle
|
||||
_dwo1(first_evicted.id)
|
||||
_m1.record_eviction()
|
||||
|
||||
@@ -320,7 +315,7 @@ class WorkstreamManager:
|
||||
self._active_id = ws.id
|
||||
|
||||
# Persist to storage only after successful insertion
|
||||
from turnstone.core.memory import increment_bucket_count, register_workstream
|
||||
from turnstone.core.memory import register_workstream
|
||||
|
||||
register_workstream(
|
||||
ws.id,
|
||||
@@ -332,17 +327,14 @@ class WorkstreamManager:
|
||||
kind=kind,
|
||||
parent_ws_id=ws.parent_ws_id,
|
||||
)
|
||||
increment_bucket_count(ws.id)
|
||||
|
||||
# Cleanup second-phase eviction outside the lock.
|
||||
if second_evicted is not None:
|
||||
self._cleanup_ui(second_evicted)
|
||||
self._last_evicted = second_evicted
|
||||
from turnstone.core.memory import decrement_bucket_count as _dbc2
|
||||
from turnstone.core.memory import delete_workstream_override as _dwo2
|
||||
from turnstone.core.metrics import metrics as _m2
|
||||
|
||||
_dbc2(second_evicted.id, active=False)
|
||||
_dwo2(second_evicted.id)
|
||||
_m2.record_eviction()
|
||||
return ws
|
||||
@@ -414,20 +406,17 @@ class WorkstreamManager:
|
||||
ws = self._workstreams.pop(ws_id, None)
|
||||
if ws is None:
|
||||
return False
|
||||
was_active = ws.state.value in _ACTIVE_STATES
|
||||
self._order.remove(ws_id)
|
||||
if self._active_id == ws_id:
|
||||
self._active_id = self._order[0]
|
||||
# Unblock any waiting approval/plan events so worker thread can exit
|
||||
self._cleanup_ui(ws)
|
||||
from turnstone.core.memory import (
|
||||
decrement_bucket_count,
|
||||
delete_workstream_override,
|
||||
update_workstream_state,
|
||||
)
|
||||
|
||||
update_workstream_state(ws_id, "closed")
|
||||
decrement_bucket_count(ws_id, active=was_active)
|
||||
delete_workstream_override(ws_id)
|
||||
return True
|
||||
|
||||
@@ -489,18 +478,12 @@ class WorkstreamManager:
|
||||
ws = self._workstreams.get(ws_id)
|
||||
if ws:
|
||||
with ws._lock:
|
||||
old_active = ws.state.value in _ACTIVE_STATES
|
||||
ws.state = state
|
||||
ws.last_active = time.monotonic()
|
||||
ws.error_message = error_msg
|
||||
from turnstone.core.memory import adjust_bucket_active, update_workstream_state
|
||||
from turnstone.core.memory import update_workstream_state
|
||||
|
||||
update_workstream_state(ws_id, state.value)
|
||||
new_active = state.value in _ACTIVE_STATES
|
||||
if old_active and not new_active:
|
||||
adjust_bucket_active(ws_id, -1)
|
||||
elif new_active and not old_active:
|
||||
adjust_bucket_active(ws_id, 1)
|
||||
if self._on_state_change:
|
||||
self._on_state_change(ws_id, state)
|
||||
|
||||
|
||||
@@ -4224,30 +4224,6 @@ def internal_model_status(request: Request) -> JSONResponse:
|
||||
return JSONResponse({"models": models})
|
||||
|
||||
|
||||
# -- internal workstream migration -------------------------------------------
|
||||
|
||||
|
||||
async def internal_migrate(request: Request) -> JSONResponse:
|
||||
"""POST /v1/api/_internal/migrate — evict a workstream for rebalancer migration."""
|
||||
from turnstone.core.web_helpers import read_json_or_400
|
||||
|
||||
body = await read_json_or_400(request)
|
||||
if isinstance(body, JSONResponse):
|
||||
return body
|
||||
ws_id = body.get("ws_id", "")
|
||||
if not ws_id:
|
||||
return JSONResponse({"status": "error", "reason": "ws_id required"}, status_code=400)
|
||||
mgr: WorkstreamManager = request.app.state.workstreams
|
||||
if mgr.get(ws_id) is None:
|
||||
return JSONResponse({"status": "not_found", "ws_id": ws_id}, status_code=404)
|
||||
if not mgr.close(ws_id):
|
||||
return JSONResponse(
|
||||
{"status": "refused", "reason": "last_workstream", "ws_id": ws_id},
|
||||
status_code=409,
|
||||
)
|
||||
return JSONResponse({"status": "ok", "ws_id": ws_id})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Global SSE fan-out
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -4705,7 +4681,6 @@ def create_app(
|
||||
methods=["POST"],
|
||||
),
|
||||
Route("/api/_internal/model-status", internal_model_status),
|
||||
Route("/api/_internal/migrate", internal_migrate, methods=["POST"]),
|
||||
],
|
||||
),
|
||||
Route("/health", health),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "inspect_workstream",
|
||||
"description": "Read the persisted state of a workstream you previously spawned: state, title, skill, timestamps, and the last N messages. Auto-approved — safe, read-only. A `live` block from the owning node is merged when available (current tokens, activity, pending_approval). The `node_id` here is the CURRENT binding (storage-authoritative) — if the rebalancer migrated the workstream after a prior spawn_workstream call, this read reflects the new node. When the workstream was closed with a reason, that string surfaces as `close_reason`. Provider-native content blocks are stripped by default to keep the response compact; pass include_provider_content=true if you need the full fidelity payload for replay tooling.",
|
||||
"description": "Read the persisted state of a workstream you previously spawned: state, title, skill, timestamps, and the last N messages. Auto-approved — safe, read-only. A `live` block from the owning node is merged when available (current tokens, activity, pending_approval). The `node_id` here is the spawn-time binding from storage; the actively-routed owner can differ if cluster membership changed since spawn (rendezvous re-derives per call). The `live` block, when present, reflects the node that currently holds the in-memory state. When the workstream was closed with a reason, that string surfaces as `close_reason`. Provider-native content blocks are stripped by default to keep the response compact; pass include_provider_content=true if you need the full fidelity payload for replay tooling.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "list_nodes",
|
||||
"description": "List ACTIVE server nodes in the cluster with their metadata. Auto-approved — safe, read-only. By default only nodes with a fresh service-registry heartbeat (within the last 120s) are returned — stale registrations (decommissioned containers, migrated hosts) are filtered out so `target_node` suggestions from this call actually route. Each node carries two metadata sources: (a) AUTO-populated on every node startup — keys arch, cpu_count, fqdn, hostname, os, os_release, python (always present); (b) USER-supplied via the console Nodes admin tab — keys like capability, region, tenant, role (deployment-specific). Pass arbitrary key=value filters to narrow the result; ALL filters must match (AND semantics). Pair with target_node on spawn_workstream to pin a child workstream to a node that matches your capability requirements. CAVEAT: the heartbeat is a 120s sliding window, so a node returned here can drop out before a follow-up `spawn_workstream(target_node=…)` lands — the spawn returns \"No available node for routing\" in that race. If you can tolerate landing on any node, omit `target_node` and let the hash ring pick from the still-healthy set; if a specific node is required, retry the spawn after re-listing. Internal network detail (the `interfaces` key — container IPs and interface names) is stripped by default because routing decisions should use capability/region/role tags, not IPs; set include_network_detail=true only if you specifically need it for debugging. Set include_inactive=true to surface stale registrations for troubleshooting — those nodes will reject `target_node` pinning.",
|
||||
"description": "List ACTIVE server nodes in the cluster with their metadata. Auto-approved — safe, read-only. By default only nodes with a fresh service-registry heartbeat (within the last 120s) are returned — stale registrations (decommissioned containers, migrated hosts) are filtered out so `target_node` suggestions from this call actually route. Each node carries two metadata sources: (a) AUTO-populated on every node startup — keys arch, cpu_count, fqdn, hostname, os, os_release, python (always present); (b) USER-supplied via the console Nodes admin tab — keys like capability, region, tenant, role (deployment-specific). Pass arbitrary key=value filters to narrow the result; ALL filters must match (AND semantics). Pair with target_node on spawn_workstream to pin a child workstream to a node that matches your capability requirements. CAVEAT: the heartbeat is a 120s sliding window, so a node returned here can drop out before a follow-up `spawn_workstream(target_node=…)` lands — the spawn returns \"No available node for routing\" in that race. If you can tolerate landing on any node, omit `target_node` and let rendezvous pick deterministically from the still-healthy set; if a specific node is required, retry the spawn after re-listing. Internal network detail (the `interfaces` key — container IPs and interface names) is stripped by default because routing decisions should use capability/region/role tags, not IPs; set include_network_detail=true only if you specifically need it for debugging. Set include_inactive=true to surface stale registrations for troubleshooting — those nodes will reject `target_node` pinning.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "spawn_workstream",
|
||||
"description": "Create a new child workstream and optionally dispatch an initial message. Use this to kick off a focused sub-task on a separate workstream (different skill, model, or node) while your coordinator remains in charge of orchestration. The child runs independently; use send_to_workstream / inspect_workstream / close_workstream to drive and observe it. The returned `node_id` is a POINT-IN-TIME snapshot of the binding at spawn — the cluster rebalancer can migrate a workstream to a different node later, so don't cache the value for long-running callbacks; re-read with inspect_workstream when you need the current binding.",
|
||||
"description": "Create a new child workstream and optionally dispatch an initial message. Use this to kick off a focused sub-task on a separate workstream (different skill, model, or node) while your coordinator remains in charge of orchestration. The child runs independently; use send_to_workstream / inspect_workstream / close_workstream to drive and observe it. The returned `node_id` is the routing target at spawn time — subsequent operations on the workstream re-route through rendezvous over the current live-node set, so a node join or drop after spawn can shift the active owner. Conversation state is shared via storage; the new owner lazily rehydrates. Don't cache `node_id` for long-running callbacks; re-read with inspect_workstream when you need the current binding.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
@@ -22,7 +22,7 @@
|
||||
},
|
||||
"target_node": {
|
||||
"type": "string",
|
||||
"description": "Optional node_id hint. When provided, the routing proxy pins the workstream to that node via generate_ws_id_for_node. Otherwise the hash ring picks. CAVEAT: pinning is a hard constraint — if the named node has dropped out of the 120s service-registry heartbeat window between your `list_nodes` call and this spawn, the spawn fails with \"No available node for routing\" rather than falling back. Omit `target_node` if you can tolerate landing on any node; pin only when the workload truly requires that specific node's capabilities/region/tenancy."
|
||||
"description": "Optional node_id hint. When provided, the routing proxy pins the workstream to that node by generating a ws_id that rendezvous-hashes onto it. Otherwise rendezvous picks deterministically from the live-node set. CAVEAT: pinning is a hard constraint — if the named node has dropped out of the 120s service-registry heartbeat window between your `list_nodes` call and this spawn, the spawn fails with \"No available node for routing\" rather than falling back. Omit `target_node` if you can tolerate landing on any node; pin only when the workload truly requires that specific node's capabilities/region/tenancy."
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
|
||||
Reference in New Issue
Block a user