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:
@@ -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.
|
**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 |
|
| Component | Purpose |
|
||||||
|-----------|---------|
|
|-----------|---------|
|
||||||
|
|||||||
@@ -2121,15 +2121,15 @@ turnstone_tool_calls_total{tool="read_file"} 3
|
|||||||
## Console Routing Proxy Endpoints
|
## Console Routing Proxy Endpoints
|
||||||
|
|
||||||
These endpoints are served by the console (`turnstone-console`) and proxy
|
These endpoints are served by the console (`turnstone-console`) and proxy
|
||||||
requests to the correct server node via the hash ring bucket cache. In
|
requests to the correct server node via rendezvous (HRW) hashing over the
|
||||||
multi-node deployments, clients (SDK, channel gateway) talk to the console
|
live service registry. In multi-node deployments, clients (SDK, channel
|
||||||
instead of individual server nodes.
|
gateway) talk to the console instead of individual server nodes.
|
||||||
|
|
||||||
### `POST /v1/api/route/workstreams/new`
|
### `POST /v1/api/route/workstreams/new`
|
||||||
|
|
||||||
Create a workstream via hash-ring routing. The console generates the `ws_id`,
|
Create a workstream via rendezvous routing. The console generates the `ws_id`,
|
||||||
routes to the assigned node, and includes `node_url` in the response for
|
routes to the rendezvous-selected node, and includes `node_url` in the
|
||||||
direct SSE connections.
|
response for direct SSE connections.
|
||||||
|
|
||||||
### `POST /v1/api/route/send`
|
### `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:
|
Prometheus metrics for the console routing layer. Includes:
|
||||||
`turnstone_router_requests_total`, `turnstone_router_request_duration_seconds`,
|
`turnstone_router_requests_total`, `turnstone_router_request_duration_seconds`,
|
||||||
`turnstone_ring_membership_size`, `turnstone_ring_version`,
|
`turnstone_router_membership_size`, `turnstone_router_refresh_total`.
|
||||||
`turnstone_ring_rebalance_total`, `turnstone_ring_migrations_total`.
|
|
||||||
|
|||||||
@@ -1,35 +1,27 @@
|
|||||||
# Consistent Hash Ring — Reference Design
|
# Consistent Hash Ring — Reference Design
|
||||||
|
|
||||||
**Status**: Reference (not currently in the hot path)
|
**Status**: Reference — alternative routing strategy
|
||||||
**Date**: 2026-03-30
|
|
||||||
|
|
||||||
## 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 FNV-1a-32 hash function specified below is bit-identical to the
|
||||||
the design of the direct HTTP transport routing system. The current
|
hash used by the live rendezvous implementation; cross-language clients
|
||||||
implementation uses weight-proportional bucket assignment with a
|
can rely on these test vectors.
|
||||||
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.
|
|
||||||
|
|
||||||
## When to consider the ring approach
|
## When the ring approach becomes interesting
|
||||||
|
|
||||||
The current weight-proportional seeding + donor/recipient rebalancer works
|
The vnode ring becomes preferable to rendezvous hashing when:
|
||||||
well when:
|
|
||||||
- Cluster size is moderate (< 50 nodes)
|
|
||||||
- Nodes join/leave infrequently
|
|
||||||
- The rebalancer runs centrally (in the console)
|
|
||||||
|
|
||||||
The consistent hash ring becomes advantageous when:
|
- Cluster size grows large (50+ nodes) and the per-route O(N) hash
|
||||||
- Cluster size grows large (50+ nodes) and frequent membership changes
|
computation becomes visible against downstream HTTP cost.
|
||||||
cause the donor/recipient algorithm to churn
|
- Decentralised routing is needed (each node computes the ring locally,
|
||||||
- Decentralized routing is needed (each node computes the ring locally,
|
no central console required).
|
||||||
no central console required)
|
- A precomputed flat-array lookup is desired so the routing hot path
|
||||||
- Cross-language determinism is important (multiple implementations must
|
avoids hashing entirely.
|
||||||
agree on the same assignment without sharing state)
|
|
||||||
|
|
||||||
## Algorithm
|
## Algorithm
|
||||||
|
|
||||||
@@ -134,16 +126,17 @@ class HashRing:
|
|||||||
# Precompute all 65536 bucket assignments
|
# Precompute all 65536 bucket assignments
|
||||||
```
|
```
|
||||||
|
|
||||||
## Comparison with current approach
|
## Comparison with rendezvous (HRW) hashing
|
||||||
|
|
||||||
| Aspect | Weight-proportional (current) | Consistent hash ring |
|
| Aspect | Rendezvous (live) | Consistent hash ring (this doc) |
|
||||||
|--------|------------------------------|---------------------|
|
|--------|-------------------|---------------------------------|
|
||||||
| Seeding | Exact weight split, deterministic | Hash-based, ~3% variance |
|
| Per-route cost | O(N) hash computes | O(log V) bisect against precomputed array |
|
||||||
| Node addition | Donor/recipient moves only excess | Ring moves ~1/N buckets |
|
| Seeding | None — pure function | Build vnode array on every membership change |
|
||||||
| Node removal | Dead buckets → most underloaded | Ring redistributes to clockwise neighbors |
|
| Node addition | Pure function moves ~1/N keys | Ring moves ~1/N buckets |
|
||||||
| Cross-node churn | Zero (only donor→recipient) | Zero (ring stability guarantee) |
|
| Node removal | Surviving nodes' keys unchanged | Surviving nodes' buckets unchanged |
|
||||||
| Decentralized | No (needs central rebalancer) | Yes (each node computes locally) |
|
| Decentralised | Yes — pure function over services | Yes — each node computes locally |
|
||||||
| Complexity | Simple weight arithmetic | Virtual node construction + bisect |
|
| Persistent state | None | None on the hot path; precomputed array in memory |
|
||||||
|
| Complexity | ~20 LOC | Virtual-node construction + bisect |
|
||||||
|
|
||||||
## Test vectors
|
## 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:
|
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.
|
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.
|
3. **Cleanup** — `TurnstoneConsole.route_close(ws_id)` closes the workstream.
|
||||||
|
|
||||||
|
|||||||
@@ -4286,7 +4286,7 @@
|
|||||||
},
|
},
|
||||||
"/v1/api/route/workstreams/new": {
|
"/v1/api/route/workstreams/new": {
|
||||||
"post": {
|
"post": {
|
||||||
"summary": "Create workstream via hash-ring routing proxy",
|
"summary": "Create workstream via rendezvous routing proxy",
|
||||||
"operationId": "v1_api_route_workstreams_new_post",
|
"operationId": "v1_api_route_workstreams_new_post",
|
||||||
"tags": [
|
"tags": [
|
||||||
"Routing"
|
"Routing"
|
||||||
|
|||||||
@@ -146,7 +146,7 @@ export class TurnstoneConsole extends BaseClient {
|
|||||||
// -- Routing proxy --------------------------------------------------------
|
// -- 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
|
* When `attachments` is non-empty the request is sent as
|
||||||
* multipart/form-data and the console routes via `?ws_id=<hex>`
|
* 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
|
assert "turnstone_router_request_duration_seconds_sum" in text
|
||||||
|
|
||||||
|
|
||||||
class TestRingInfo:
|
class TestRouterInfo:
|
||||||
"""Ring membership and version gauges."""
|
"""Live-membership gauge + refresh counter."""
|
||||||
|
|
||||||
def test_defaults_zero(self) -> None:
|
def test_defaults_zero(self) -> None:
|
||||||
m = ConsoleMetrics()
|
m = ConsoleMetrics()
|
||||||
text = m.generate_text()
|
text = m.generate_text()
|
||||||
assert "turnstone_ring_membership_size 0" in text
|
assert "turnstone_router_membership_size 0" in text
|
||||||
assert "turnstone_ring_version 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 = ConsoleMetrics()
|
||||||
m.set_ring_info(3, 7)
|
m.set_router_info(3, 7)
|
||||||
|
|
||||||
text = m.generate_text()
|
text = m.generate_text()
|
||||||
assert "turnstone_ring_membership_size 3" in text
|
assert "turnstone_router_membership_size 3" in text
|
||||||
assert "turnstone_ring_version 7" in text
|
assert "turnstone_router_refresh_total 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
|
|
||||||
|
|
||||||
|
|
||||||
class TestGenerateText:
|
class TestGenerateText:
|
||||||
@@ -103,10 +64,8 @@ class TestGenerateText:
|
|||||||
expected = [
|
expected = [
|
||||||
"turnstone_router_requests_total",
|
"turnstone_router_requests_total",
|
||||||
"turnstone_router_request_duration_seconds",
|
"turnstone_router_request_duration_seconds",
|
||||||
"turnstone_ring_membership_size",
|
"turnstone_router_membership_size",
|
||||||
"turnstone_ring_version",
|
"turnstone_router_refresh_total",
|
||||||
"turnstone_ring_rebalance_total",
|
|
||||||
"turnstone_ring_migrations_total",
|
|
||||||
]
|
]
|
||||||
for name in expected:
|
for name in expected:
|
||||||
assert name in text, f"Missing metric: {name}"
|
assert name in text, f"Missing metric: {name}"
|
||||||
@@ -116,8 +75,8 @@ class TestGenerateText:
|
|||||||
text = m.generate_text()
|
text = m.generate_text()
|
||||||
assert "# HELP turnstone_router_requests_total" in text
|
assert "# HELP turnstone_router_requests_total" in text
|
||||||
assert "# TYPE turnstone_router_requests_total counter" in text
|
assert "# TYPE turnstone_router_requests_total counter" in text
|
||||||
assert "# HELP turnstone_ring_membership_size" in text
|
assert "# HELP turnstone_router_membership_size" in text
|
||||||
assert "# TYPE turnstone_ring_membership_size gauge" in text
|
assert "# TYPE turnstone_router_membership_size gauge" in text
|
||||||
|
|
||||||
def test_ends_with_newline(self) -> None:
|
def test_ends_with_newline(self) -> None:
|
||||||
m = ConsoleMetrics()
|
m = ConsoleMetrics()
|
||||||
@@ -125,24 +84,16 @@ class TestGenerateText:
|
|||||||
assert text.endswith("\n")
|
assert text.endswith("\n")
|
||||||
|
|
||||||
def test_combined_scenario(self) -> None:
|
def test_combined_scenario(self) -> None:
|
||||||
"""Full scenario: routes, ring info, rebalances, migrations."""
|
"""Full scenario: routes + router info."""
|
||||||
m = ConsoleMetrics()
|
m = ConsoleMetrics()
|
||||||
m.record_route("create", 200, 0.1)
|
m.record_route("create", 200, 0.1)
|
||||||
m.record_route("send", 200, 0.05)
|
m.record_route("send", 200, 0.05)
|
||||||
m.record_route("send", 502, 1.2)
|
m.record_route("send", 502, 1.2)
|
||||||
m.set_ring_info(3, 12)
|
m.set_router_info(3, 12)
|
||||||
m.record_rebalance("seeded")
|
|
||||||
m.record_rebalance("noop")
|
|
||||||
m.record_rebalance("rebalanced")
|
|
||||||
m.record_migrations(4)
|
|
||||||
|
|
||||||
text = m.generate_text()
|
text = m.generate_text()
|
||||||
assert 'turnstone_router_requests_total{method="create",status="2xx"} 1' in 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="2xx"} 1' in text
|
||||||
assert 'turnstone_router_requests_total{method="send",status="5xx"} 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_router_membership_size 3" in text
|
||||||
assert "turnstone_ring_version 12" in text
|
assert "turnstone_router_refresh_total 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
|
|
||||||
|
|||||||
+178
-192
@@ -1,17 +1,13 @@
|
|||||||
"""Tests for turnstone.console.router."""
|
"""Tests for turnstone.console.router (rendezvous routing)."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Any
|
import secrets
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from turnstone.console.router import ConsoleRouter, NodeRef
|
from turnstone.console.router import ConsoleRouter, NodeRef
|
||||||
from turnstone.core.hash_ring import RING_SIZE, NoAvailableNodeError
|
from turnstone.core.rendezvous import NoAvailableNodeError
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Fake storage
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class FakeStorage:
|
class FakeStorage:
|
||||||
@@ -19,26 +15,14 @@ class FakeStorage:
|
|||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.services: list[dict[str, str]] = []
|
self.services: list[dict[str, str]] = []
|
||||||
self.buckets: list[dict[str, Any]] = []
|
|
||||||
self.overrides: list[dict[str, str]] = []
|
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]]:
|
def list_services(self, service_type: str, max_age_seconds: int = 120) -> list[dict[str, str]]:
|
||||||
return list(self.services)
|
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]]:
|
def list_workstream_overrides(self) -> list[dict[str, str]]:
|
||||||
return list(self.overrides)
|
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_A = {"service_id": "node-a", "url": "http://a:8080", "metadata": "{}"}
|
||||||
NODE_B = {"service_id": "node-b", "url": "http://b: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]
|
return ConsoleRouter(s), s # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
def _ws_id_for_bucket(bucket: int) -> str:
|
def _random_ws_id() -> str:
|
||||||
"""Build a 32-char hex ws_id whose first 4 chars encode *bucket*."""
|
return secrets.token_hex(16)
|
||||||
return f"{bucket:04x}" + "0" * 28
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# TestRouteBasic
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class TestRouteBasic:
|
class TestRouteBasic:
|
||||||
"""Basic routing through the bucket cache."""
|
def test_route_returns_a_live_node(self) -> None:
|
||||||
|
|
||||||
def test_route_returns_correct_node(self) -> None:
|
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
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()
|
router.refresh_cache()
|
||||||
|
|
||||||
assert router.route(_ws_id_for_bucket(0x0000)) == NodeRef("node-a", "http://a:8080")
|
ref = router.route(_random_ws_id())
|
||||||
assert router.route(_ws_id_for_bucket(0x0001)) == NodeRef("node-b", "http://b:8080")
|
assert ref.node_id in {"node-a", "node-b", "node-c"}
|
||||||
assert router.route(_ws_id_for_bucket(0x0002)) == NodeRef("node-c", "http://c:8080")
|
|
||||||
|
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:
|
def test_route_override_priority(self) -> None:
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A, NODE_B]
|
storage.services = [NODE_A, NODE_B]
|
||||||
storage.buckets = [{"bucket": 0x0000, "node_id": "node-a"}]
|
ws_id = _random_ws_id()
|
||||||
ws_id = _ws_id_for_bucket(0x0000)
|
|
||||||
storage.overrides = [{"ws_id": ws_id, "node_id": "node-b"}]
|
storage.overrides = [{"ws_id": ws_id, "node_id": "node-b"}]
|
||||||
router.refresh_cache()
|
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")
|
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()
|
router, _ = _make_router()
|
||||||
|
with pytest.raises(NoAvailableNodeError):
|
||||||
|
router.route(_random_ws_id())
|
||||||
|
|
||||||
with pytest.raises(NoAvailableNodeError, match="not assigned"):
|
def test_route_empty_ws_id_raises(self) -> None:
|
||||||
router.route(_ws_id_for_bucket(0x0000))
|
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:
|
def test_route_url_convenience(self) -> None:
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A]
|
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()
|
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}
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
moved = sum(1 for ws in sample if before[ws] != after[ws])
|
||||||
# TestRefreshCache
|
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
|
||||||
|
|
||||||
|
def test_node_leave_only_redistributes_dead_node_keys(self) -> None:
|
||||||
class TestRefreshCache:
|
"""Removing node-a sends node-a's keys to b/c only; keys that
|
||||||
"""Cache loading from storage."""
|
were on b/c stay put."""
|
||||||
|
|
||||||
def test_refresh_loads_from_storage(self) -> None:
|
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A]
|
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||||
storage.buckets = [{"bucket": 100, "node_id": "node-a"}]
|
|
||||||
router.refresh_cache()
|
router.refresh_cache()
|
||||||
|
|
||||||
ref = router.route(_ws_id_for_bucket(100))
|
sample = [_random_ws_id() for _ in range(2000)]
|
||||||
assert ref.node_id == "node-a"
|
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()
|
router, storage = _make_router()
|
||||||
# node-b is in buckets but not in services (dead/expired)
|
storage.services = [
|
||||||
storage.services = [NODE_A]
|
{"service_id": "node-a", "url": "http://a:8080", "metadata": '{"weight": 2}'},
|
||||||
storage.buckets = [
|
{"service_id": "node-b", "url": "http://b:8080", "metadata": '{"weight": 1}'},
|
||||||
{"bucket": 0x0000, "node_id": "node-a"},
|
|
||||||
{"bucket": 0x0001, "node_id": "node-b"},
|
|
||||||
]
|
]
|
||||||
router.refresh_cache()
|
router.refresh_cache()
|
||||||
|
|
||||||
assert router.route(_ws_id_for_bucket(0x0000)).node_id == "node-a"
|
sample = [_random_ws_id() for _ in range(5000)]
|
||||||
with pytest.raises(NoAvailableNodeError):
|
on_a = sum(1 for ws in sample if router.route(ws).node_id == "node-a")
|
||||||
router.route(_ws_id_for_bucket(0x0001))
|
# 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()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A]
|
storage.services = [
|
||||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
{"service_id": "node-a", "url": "http://a:8080", "metadata": "not json"},
|
||||||
|
]
|
||||||
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"}]
|
|
||||||
|
|
||||||
router.refresh_cache()
|
router.refresh_cache()
|
||||||
assert router.refresh_cache() is False
|
# Just confirms it doesn't blow up.
|
||||||
|
router.route(_random_ws_id())
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
class TestRefreshLifecycle:
|
||||||
# TestCheckVersion
|
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."""
|
||||||
|
|
||||||
class TestCheckVersion:
|
|
||||||
"""Version-gated refresh."""
|
|
||||||
|
|
||||||
def test_version_change_triggers_refresh(self) -> None:
|
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A]
|
storage.services = [NODE_A]
|
||||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
router.refresh_cache()
|
||||||
storage.settings["rebalancer_version"] = {"value": "1"}
|
assert router.node_count() == 1
|
||||||
|
|
||||||
assert router.check_version() is True
|
storage.services = [NODE_A, NODE_B]
|
||||||
assert router.is_ready()
|
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()
|
router, storage = _make_router()
|
||||||
# Default version is 0; setting absent also means 0
|
|
||||||
storage.services = [NODE_A]
|
storage.services = [NODE_A]
|
||||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
|
||||||
|
|
||||||
# First call: version=0 matches self._version=0 -> no refresh
|
with router._refresh_lock:
|
||||||
assert router.check_version() is False
|
# Lock held by this thread → the call below can't acquire.
|
||||||
assert not router.is_ready() # cache was never loaded
|
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()
|
router, storage = _make_router()
|
||||||
# settings dict is empty -> get_system_setting returns None
|
storage.services = [NODE_A]
|
||||||
assert router.check_version() is False
|
|
||||||
|
|
||||||
|
# Hold the refresh lock from another thread.
|
||||||
|
lock_held = threading.Event()
|
||||||
|
release = threading.Event()
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
def hold_lock() -> None:
|
||||||
# TestGenerateWsId
|
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:
|
class TestGenerateWsId:
|
||||||
"""Workstream ID generation targeting a specific node."""
|
|
||||||
|
|
||||||
def test_generates_routable_id(self) -> None:
|
def test_generates_routable_id(self) -> None:
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A, NODE_B]
|
storage.services = [NODE_A, NODE_B, NODE_C]
|
||||||
storage.buckets = [
|
|
||||||
{"bucket": 0x00FF, "node_id": "node-a"},
|
|
||||||
{"bucket": 0x0100, "node_id": "node-b"},
|
|
||||||
]
|
|
||||||
router.refresh_cache()
|
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 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:
|
def test_unknown_node_raises(self) -> None:
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A]
|
storage.services = [NODE_A]
|
||||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
|
||||||
router.refresh_cache()
|
router.refresh_cache()
|
||||||
|
|
||||||
with pytest.raises(NoAvailableNodeError, match="node-z"):
|
with pytest.raises(NoAvailableNodeError, match="node-z"):
|
||||||
router.generate_ws_id_for_node("node-z")
|
router.generate_ws_id_for_node("node-z")
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# TestIsReady
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class TestIsReady:
|
class TestIsReady:
|
||||||
"""Readiness checks."""
|
|
||||||
|
|
||||||
def test_false_when_empty(self) -> None:
|
def test_false_when_empty(self) -> None:
|
||||||
router, _ = _make_router()
|
router, _ = _make_router()
|
||||||
assert router.is_ready() is False
|
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()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A]
|
storage.services = [NODE_A]
|
||||||
storage.buckets = [{"bucket": 0, "node_id": "node-a"}]
|
|
||||||
router.refresh_cache()
|
router.refresh_cache()
|
||||||
|
|
||||||
assert router.is_ready() is True
|
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:
|
class TestNodeCount:
|
||||||
"""Distinct node counting."""
|
def test_count_matches_live_services(self) -> None:
|
||||||
|
|
||||||
def test_count_distinct_nodes(self) -> None:
|
|
||||||
router, storage = _make_router()
|
router, storage = _make_router()
|
||||||
storage.services = [NODE_A, NODE_B, NODE_C]
|
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()
|
router.refresh_cache()
|
||||||
|
|
||||||
assert router.node_count() == 3
|
assert router.node_count() == 3
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ from starlette.testclient import TestClient
|
|||||||
|
|
||||||
from turnstone.console.collector import ClusterCollector
|
from turnstone.console.collector import ClusterCollector
|
||||||
from turnstone.console.router import ConsoleRouter, NodeRef
|
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
|
# Shared test auth — JWT-based
|
||||||
_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"
|
_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:
|
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()
|
@pytest.fixture()
|
||||||
def client(self):
|
def client(self):
|
||||||
@@ -178,17 +178,17 @@ class TestRouteCreate:
|
|||||||
router.generate_ws_id_for_node.assert_called_with("node-c")
|
router.generate_ws_id_for_node.assert_called_with("node-c")
|
||||||
client.close()
|
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
|
"""Default fan-out (no resume_ws / no target_node) reports
|
||||||
routing_strategy='hash_ring' so the coordinator's spawn tool can
|
routing_strategy='rendezvous' so the coordinator's spawn tool
|
||||||
explain why the node was chosen."""
|
can explain why the node was chosen."""
|
||||||
resp = client.post(
|
resp = client.post(
|
||||||
"/v1/api/route/workstreams/new",
|
"/v1/api/route/workstreams/new",
|
||||||
json={"name": "test-ws"},
|
json={"name": "test-ws"},
|
||||||
headers=_TEST_AUTH_HEADERS,
|
headers=_TEST_AUTH_HEADERS,
|
||||||
)
|
)
|
||||||
assert resp.status_code == 200
|
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):
|
def test_route_create_routing_strategy_target_node(self):
|
||||||
router = _make_mock_router()
|
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(
|
EndpointSpec(
|
||||||
"/v1/api/route/workstreams/new",
|
"/v1/api/route/workstreams/new",
|
||||||
"POST",
|
"POST",
|
||||||
"Create workstream via hash-ring routing proxy",
|
"Create workstream via rendezvous routing proxy",
|
||||||
response_model=RouteCreateResponse,
|
response_model=RouteCreateResponse,
|
||||||
error_codes=[400, 503],
|
error_codes=[400, 503],
|
||||||
tags=["Routing"],
|
tags=["Routing"],
|
||||||
|
|||||||
@@ -427,16 +427,17 @@ class ClusterCollector:
|
|||||||
for nid in lost_nodes:
|
for nid in lost_nodes:
|
||||||
asyncio.run_coroutine_threadsafe(self._stop_node(nid), self._sse_loop)
|
asyncio.run_coroutine_threadsafe(self._stop_node(nid), self._sse_loop)
|
||||||
|
|
||||||
# Notify the routing layer so it can refresh its hash-ring cache
|
# Drive the router's cache from the collector's discovery
|
||||||
# when the rebalancer has published a new version.
|
# 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:
|
if self._router is not None:
|
||||||
try:
|
try:
|
||||||
self._router.check_version()
|
self._router.refresh_cache()
|
||||||
except Exception:
|
except Exception:
|
||||||
log.debug("Router version check failed", exc_info=True)
|
log.debug("Router refresh failed", exc_info=True)
|
||||||
# Update ring gauge metrics after version check
|
|
||||||
if self._console_metrics is not None:
|
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.node_count(),
|
||||||
self._router.version,
|
self._router.version,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -57,7 +57,7 @@ class CoordinatorManager:
|
|||||||
|
|
||||||
# Pseudo-node id persisted on coordinator rows so ``workstreams.node_id``
|
# Pseudo-node id persisted on coordinator rows so ``workstreams.node_id``
|
||||||
# stays non-NULL and list / audit surfaces can distinguish coordinators
|
# 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.
|
# unroutable sentinel — coordinators never land on real nodes.
|
||||||
# Bound from ``ClusterCollector.CONSOLE_PSEUDO_NODE_ID`` so the two
|
# Bound from ``ClusterCollector.CONSOLE_PSEUDO_NODE_ID`` so the two
|
||||||
# literals can't drift (the collector's eviction + query filters
|
# literals can't drift (the collector's eviction + query filters
|
||||||
|
|||||||
@@ -8,10 +8,10 @@ from collections import defaultdict
|
|||||||
|
|
||||||
|
|
||||||
class ConsoleMetrics:
|
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
|
Lighter-weight than the server's MetricsCollector — tracks router
|
||||||
request counters, ring membership gauges, and rebalancer activity.
|
request counters and live-membership gauges.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
@@ -19,10 +19,8 @@ class ConsoleMetrics:
|
|||||||
self._router_requests: dict[tuple[str, str], int] = defaultdict(int)
|
self._router_requests: dict[tuple[str, str], int] = defaultdict(int)
|
||||||
self._router_duration_sum: dict[str, float] = defaultdict(float)
|
self._router_duration_sum: dict[str, float] = defaultdict(float)
|
||||||
self._router_duration_count: dict[str, int] = defaultdict(int)
|
self._router_duration_count: dict[str, int] = defaultdict(int)
|
||||||
self._ring_membership: int = 0
|
self._router_membership: int = 0
|
||||||
self._ring_version: int = 0
|
self._router_refresh_count: int = 0
|
||||||
self._rebalance_total: dict[str, int] = defaultdict(int)
|
|
||||||
self._migrations_total: int = 0
|
|
||||||
self._start_time: float = time.monotonic()
|
self._start_time: float = time.monotonic()
|
||||||
|
|
||||||
def record_route(self, method: str, status: int, duration: float) -> None:
|
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_sum[method] += duration
|
||||||
self._router_duration_count[method] += 1
|
self._router_duration_count[method] += 1
|
||||||
|
|
||||||
def set_ring_info(self, membership: int, version: int) -> None:
|
def set_router_info(self, membership: int, refresh_count: int) -> None:
|
||||||
"""Update the current ring membership size and version."""
|
"""Update current live-node count + the router's refresh counter."""
|
||||||
with self._lock:
|
with self._lock:
|
||||||
self._ring_membership = membership
|
self._router_membership = membership
|
||||||
self._ring_version = version
|
self._router_refresh_count = refresh_count
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
def generate_text(self) -> str:
|
def generate_text(self) -> str:
|
||||||
"""Return Prometheus text exposition format (v0.0.4)."""
|
"""Return Prometheus text exposition format (v0.0.4)."""
|
||||||
@@ -57,10 +45,8 @@ class ConsoleMetrics:
|
|||||||
router_requests = dict(self._router_requests)
|
router_requests = dict(self._router_requests)
|
||||||
duration_sum = dict(self._router_duration_sum)
|
duration_sum = dict(self._router_duration_sum)
|
||||||
duration_count = dict(self._router_duration_count)
|
duration_count = dict(self._router_duration_count)
|
||||||
ring_membership = self._ring_membership
|
router_membership = self._router_membership
|
||||||
ring_version = self._ring_version
|
router_refresh_count = self._router_refresh_count
|
||||||
rebalance_total = dict(self._rebalance_total)
|
|
||||||
migrations_total = self._migrations_total
|
|
||||||
|
|
||||||
# turnstone_router_requests_total
|
# turnstone_router_requests_total
|
||||||
lines.append("# HELP turnstone_router_requests_total Console-routed requests")
|
lines.append("# HELP turnstone_router_requests_total Console-routed requests")
|
||||||
@@ -85,26 +71,17 @@ class ConsoleMetrics:
|
|||||||
f" {duration_count[method]}"
|
f" {duration_count[method]}"
|
||||||
)
|
)
|
||||||
|
|
||||||
# turnstone_ring_membership_size
|
# turnstone_router_membership_size
|
||||||
lines.append("# HELP turnstone_ring_membership_size Current ring node count")
|
lines.append("# HELP turnstone_router_membership_size Current live-node count")
|
||||||
lines.append("# TYPE turnstone_ring_membership_size gauge")
|
lines.append("# TYPE turnstone_router_membership_size gauge")
|
||||||
lines.append(f"turnstone_ring_membership_size {ring_membership}")
|
lines.append(f"turnstone_router_membership_size {router_membership}")
|
||||||
|
|
||||||
# turnstone_ring_version
|
# turnstone_router_refresh_total — bumped on every successful
|
||||||
lines.append("# HELP turnstone_ring_version Current ring version")
|
# cache refresh. A flat counter under churn means the
|
||||||
lines.append("# TYPE turnstone_ring_version gauge")
|
# collector's discovery loop is stuck.
|
||||||
lines.append(f"turnstone_ring_version {ring_version}")
|
lines.append("# HELP turnstone_router_refresh_total Router cache refresh counter")
|
||||||
|
lines.append("# TYPE turnstone_router_refresh_total counter")
|
||||||
# turnstone_ring_rebalance_total
|
lines.append(f"turnstone_router_refresh_total {router_refresh_count}")
|
||||||
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}")
|
|
||||||
|
|
||||||
lines.append("") # trailing newline
|
lines.append("") # trailing newline
|
||||||
return "\n".join(lines)
|
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.
|
"""Console routing layer — routes workstream requests to server nodes.
|
||||||
|
|
||||||
Maintains an in-memory flat array of 65536 bucket->NodeRef entries populated
|
Uses rendezvous (HRW) hashing over the live ``services`` table. The
|
||||||
from the hash_ring_buckets table. Routing is O(1): cache[int(ws_id[:4], 16)].
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
import logging
|
|
||||||
import secrets
|
import secrets
|
||||||
import threading
|
import threading
|
||||||
from dataclasses import dataclass
|
|
||||||
from typing import TYPE_CHECKING
|
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:
|
if TYPE_CHECKING:
|
||||||
from turnstone.core.storage._protocol import StorageBackend
|
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)
|
def _parse_weight(metadata_json: str) -> int:
|
||||||
class NodeRef:
|
"""Pull the ``weight`` key out of a service-registry metadata blob.
|
||||||
"""A server node that can receive proxied requests."""
|
|
||||||
|
|
||||||
node_id: str
|
A single corrupt row must not abort the cache refresh — fall back to
|
||||||
url: str
|
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:
|
class ConsoleRouter:
|
||||||
"""Routes workstream operations to the correct server node.
|
"""Routes workstream operations to the correct server node.
|
||||||
|
|
||||||
Maintains an in-memory flat array of 65536 bucket->NodeRef entries,
|
Thread-safe. All state mutation goes through ``_lock``; lookups
|
||||||
populated from the hash_ring_buckets table. All routing is a
|
snapshot the node list under the lock and run the rendezvous select
|
||||||
single O(1) array lookup: ``cache[int(ws_id[:4], 16)]``.
|
outside it (the select is a pure function over an immutable list).
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, storage: StorageBackend) -> None:
|
def __init__(self, storage: StorageBackend) -> None:
|
||||||
self._storage = storage
|
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._overrides: dict[str, NodeRef] = {}
|
||||||
self._version: int = 0
|
|
||||||
self._refresh_lock = threading.Lock()
|
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
|
# Cache management
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def refresh_cache(self) -> bool:
|
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
|
Thread-safe: if another thread is already refreshing, this call
|
||||||
returns False immediately (the other thread's refresh will apply).
|
returns False immediately (the in-flight refresh will publish
|
||||||
Returns True if the cache changed compared to the previous load.
|
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):
|
if not self._refresh_lock.acquire(blocking=False):
|
||||||
return False # another thread is refreshing
|
return False
|
||||||
try:
|
try:
|
||||||
return self._refresh_cache_locked()
|
return self._refresh_locked()
|
||||||
finally:
|
finally:
|
||||||
self._refresh_lock.release()
|
self._refresh_lock.release()
|
||||||
|
|
||||||
def _refresh_cache_locked(self) -> bool:
|
def force_refresh(self) -> bool:
|
||||||
"""Inner refresh — must be called with _refresh_lock held."""
|
"""Refresh now, blocking if another refresh is in progress.
|
||||||
# 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
|
|
||||||
}
|
|
||||||
|
|
||||||
# Load bucket assignments into flat array
|
Used by the 404-retry path in the routing proxy when ``route()``
|
||||||
buckets = self._storage.list_ring_buckets()
|
sent the request to a node that doesn't have the workstream —
|
||||||
new_cache: list[NodeRef | None] = [None] * RING_SIZE
|
the retry needs a guaranteed-fresh view of membership +
|
||||||
for row in buckets:
|
overrides before giving up.
|
||||||
ref = nodes.get(row["node_id"])
|
|
||||||
if ref is not None:
|
|
||||||
new_cache[row["bucket"]] = ref
|
|
||||||
|
|
||||||
# Load per-workstream overrides (pinned workstreams)
|
Async callers must wrap this in ``asyncio.to_thread`` — the
|
||||||
overrides = self._storage.list_workstream_overrides()
|
method takes a blocking lock and issues storage queries.
|
||||||
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.
|
|
||||||
"""
|
"""
|
||||||
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:
|
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._overrides = new_overrides
|
||||||
self._version = version
|
self._refresh_counter += 1
|
||||||
|
return changed
|
||||||
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
|
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
# Routing
|
# Routing
|
||||||
@@ -155,21 +156,21 @@ class ConsoleRouter:
|
|||||||
|
|
||||||
Priority:
|
Priority:
|
||||||
1. Per-workstream override (pinned to a specific node)
|
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)
|
with self._lock:
|
||||||
if ref is not None:
|
ref = self._overrides.get(ws_id)
|
||||||
return ref
|
if ref is not None:
|
||||||
if len(ws_id) < 4:
|
return ref
|
||||||
raise NoAvailableNodeError(f"invalid ws_id: {ws_id!r}")
|
nodes = self._nodes # snapshot — list is replaced wholesale on refresh
|
||||||
try:
|
if not nodes:
|
||||||
bucket = int(ws_id[:4], 16)
|
raise NoAvailableNodeError("no live nodes")
|
||||||
except ValueError:
|
if not ws_id:
|
||||||
raise NoAvailableNodeError(f"invalid ws_id prefix: {ws_id[:4]!r}") from None
|
raise NoAvailableNodeError("invalid ws_id: empty")
|
||||||
ref = self._cache[bucket]
|
return select(ws_id, nodes)
|
||||||
if ref is None:
|
|
||||||
raise NoAvailableNodeError(f"bucket {bucket} not assigned")
|
|
||||||
return ref
|
|
||||||
|
|
||||||
def route_url(self, ws_id: str) -> str:
|
def route_url(self, ws_id: str) -> str:
|
||||||
"""Convenience — return just the URL for the target node."""
|
"""Convenience — return just the URL for the target node."""
|
||||||
@@ -180,29 +181,50 @@ class ConsoleRouter:
|
|||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def is_ready(self) -> bool:
|
def is_ready(self) -> bool:
|
||||||
"""Return True if at least one bucket is assigned."""
|
"""True if the router knows about at least one live node."""
|
||||||
return any(ref is not None for ref in self._cache)
|
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
|
@property
|
||||||
def version(self) -> int:
|
def version(self) -> int:
|
||||||
"""The last seen rebalancer version."""
|
"""Monotonic counter bumped on every successful cache refresh.
|
||||||
return self._version
|
|
||||||
|
|
||||||
def node_count(self) -> int:
|
Surfaced by the collector's ``set_ring_info`` gauge so a
|
||||||
"""Count distinct nodes present in the cache."""
|
dashboard can detect when membership stops being refreshed.
|
||||||
return len({ref.node_id for ref in self._cache if ref is not None})
|
Strictly increasing across the process lifetime.
|
||||||
|
"""
|
||||||
|
with self._lock:
|
||||||
|
return self._refresh_counter
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
# Workstream ID generation
|
# Workstream ID generation
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def generate_ws_id_for_node(self, node_id: str) -> str:
|
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
|
Brute-force loop: pick a random candidate, check whether HRW
|
||||||
remaining 28 hex chars are random (32 chars total).
|
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):
|
with self._lock:
|
||||||
if ref is not None and ref.node_id == node_id:
|
nodes = list(self._nodes)
|
||||||
return f"{bucket:04x}" + secrets.token_hex(14)
|
if not nodes:
|
||||||
raise NoAvailableNodeError(f"no bucket assigned to node {node_id!r}")
|
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,
|
jwt_version_slot,
|
||||||
require_permission,
|
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.skill_kind import SkillKind
|
||||||
from turnstone.core.web_helpers import (
|
from turnstone.core.web_helpers import (
|
||||||
read_json_or_400,
|
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:
|
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
|
Accepts both `application/json` and `multipart/form-data`. Multipart
|
||||||
callers must include ``?ws_id=<hex>`` in the URL query string so the
|
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
|
router: ConsoleRouter | None = request.app.state.router
|
||||||
ring_ready = router is not None and router.is_ready()
|
ring_ready = router is not None and router.is_ready()
|
||||||
if not ring_ready:
|
if not ring_ready:
|
||||||
# Ring not yet populated (rebalancer hasn't run or is disabled).
|
# Router cache empty — the collector hasn't published a
|
||||||
# Try a one-shot refresh before giving up — the rebalancer may
|
# services list yet. One-shot refresh off the event loop
|
||||||
# have written buckets since the last collector poll.
|
# before giving up.
|
||||||
if router is not None:
|
if router is not None:
|
||||||
await asyncio.to_thread(router.refresh_cache)
|
await asyncio.to_thread(router.refresh_cache)
|
||||||
ring_ready = router.is_ready()
|
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
|
# Routing strategy is surfaced on the response so callers (the
|
||||||
# coordinator's spawn_workstream tool especially) can explain why a
|
# coordinator's spawn_workstream tool especially) can explain why a
|
||||||
# given node was chosen. Set on every branch below.
|
# given node was chosen. Set on every branch below.
|
||||||
routing_strategy = "hash_ring"
|
routing_strategy = "rendezvous"
|
||||||
|
|
||||||
if is_multipart:
|
if is_multipart:
|
||||||
# Multipart: caller must pass ws_id as a query param so we can
|
# 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"])
|
ref = router.route(body["resume_ws"])
|
||||||
routing_strategy = "resume"
|
routing_strategy = "resume"
|
||||||
elif body.get("target_node"):
|
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
|
body["ws_id"] = ws_id
|
||||||
ref = router.route(ws_id)
|
ref = router.route(ws_id)
|
||||||
pin = True
|
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 ""
|
audit_ws_id = body.get("ws_id") or body.get("resume_ws", "") or ""
|
||||||
# Return the storage-authoritative node_id so subsequent
|
# Return the storage-authoritative node_id so subsequent
|
||||||
# inspect / list calls agree on the binding. ``ref.node_id`` is
|
# inspect / list calls agree on the binding. ``ref.node_id`` is
|
||||||
# the hash-ring target AT SPAWN TIME — stale once the
|
# the rendezvous target AT SPAWN TIME — stale once membership
|
||||||
# rebalancer runs or a node comes/goes, and the node's own
|
# changes — and the node's own create handler is the source of
|
||||||
# create handler is the source of truth for what node_id got
|
# truth for what node_id got persisted on the workstream row.
|
||||||
# persisted on the workstream row. Fall back to ref.node_id
|
# Fall back to ref.node_id only when the storage lookup fails,
|
||||||
# only when the storage lookup fails, matching the previous
|
# matching the previous behaviour so this change is strictly
|
||||||
# behaviour so this change is strictly additive.
|
# additive.
|
||||||
bound_node_id = ref.node_id
|
bound_node_id = ref.node_id
|
||||||
storage = getattr(request.app.state, "auth_storage", None)
|
storage = getattr(request.app.state, "auth_storage", None)
|
||||||
if storage is not None and audit_ws_id:
|
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:
|
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
|
Handles all four shapes mounted under
|
||||||
``/v1/api/route/workstreams/{ws_id}/attachments[/...]``:
|
``/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):
|
# Transparent retry on 404 (at most once):
|
||||||
#
|
#
|
||||||
# The bucket-routed node doesn't have the workstream. Refresh the
|
# The rendezvous-selected node doesn't have the workstream. Refresh
|
||||||
# cache (reloads overrides + bucket assignments from DB) and re-route.
|
# membership + overrides and re-route. If the route changed (e.g., a
|
||||||
# If the route changed (e.g., a local-create override was added since
|
# local-create override was added since the last cache load, or a
|
||||||
# the last cache load), retry on the new node. If the route is the
|
# node has joined / dropped), retry on the new node. If the route
|
||||||
# same, return the 404 as-is — no loop, no scan.
|
# is the same, return the 404 as-is — no loop, no scan.
|
||||||
if resp.status_code == 404:
|
if resp.status_code == 404:
|
||||||
# Blocking refresh — wait for any in-progress refresh to finish
|
# Off the event loop — force_refresh takes a blocking lock and
|
||||||
# so the retry uses the latest data, not stale cache.
|
# issues two storage queries. Coalesces internally so a 404
|
||||||
router._refresh_lock.acquire()
|
# stampede after a node churn doesn't N×-multiply DB reads.
|
||||||
try:
|
await asyncio.to_thread(router.force_refresh)
|
||||||
router._refresh_cache_locked()
|
|
||||||
finally:
|
|
||||||
router._refresh_lock.release()
|
|
||||||
try:
|
try:
|
||||||
new_ref = router.route(ws_id)
|
new_ref = router.route(ws_id)
|
||||||
except (NoAvailableNodeError, ValueError):
|
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")
|
audit_exec = ThreadPoolExecutor(max_workers=4, thread_name_prefix="coord-audit")
|
||||||
app.state.audit_executor = audit_exec
|
app.state.audit_executor = audit_exec
|
||||||
_set_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)
|
_router: ConsoleRouter | None = getattr(app.state, "router", None)
|
||||||
if _router is not None:
|
if _router is not None:
|
||||||
_router.refresh_cache()
|
_router.refresh_cache()
|
||||||
@@ -3779,10 +3778,6 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
|||||||
scheduler = getattr(app.state, "scheduler", None)
|
scheduler = getattr(app.state, "scheduler", None)
|
||||||
if scheduler is not None:
|
if scheduler is not None:
|
||||||
scheduler.start()
|
scheduler.start()
|
||||||
# Start rebalancer if configured
|
|
||||||
_rebalancer = getattr(app.state, "rebalancer", None)
|
|
||||||
if _rebalancer is not None:
|
|
||||||
_rebalancer.start()
|
|
||||||
# OIDC discovery (if configured)
|
# OIDC discovery (if configured)
|
||||||
oidc_config = app.state.oidc_config
|
oidc_config = app.state.oidc_config
|
||||||
if oidc_config.enabled:
|
if oidc_config.enabled:
|
||||||
@@ -3982,9 +3977,6 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
|||||||
await tls_mgr.stop_renewal()
|
await tls_mgr.stop_renewal()
|
||||||
if scheduler is not None:
|
if scheduler is not None:
|
||||||
scheduler.stop()
|
scheduler.stop()
|
||||||
_rebalancer = getattr(app.state, "rebalancer", None)
|
|
||||||
if _rebalancer is not None:
|
|
||||||
_rebalancer.stop()
|
|
||||||
coord_mgr_shutdown = getattr(app.state, "coord_mgr", None)
|
coord_mgr_shutdown = getattr(app.state, "coord_mgr", None)
|
||||||
if coord_mgr_shutdown is not None:
|
if coord_mgr_shutdown is not None:
|
||||||
try:
|
try:
|
||||||
@@ -9912,35 +9904,6 @@ async def admin_delete_node_metadata_key(request: Request) -> JSONResponse:
|
|||||||
return JSONResponse({"ok": True})
|
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:
|
async def tls_ca_status(request: Request) -> JSONResponse:
|
||||||
"""GET /v1/api/admin/tls/ca — CA status."""
|
"""GET /v1/api/admin/tls/ca — CA status."""
|
||||||
from turnstone.core.auth import require_permission
|
from turnstone.core.auth import require_permission
|
||||||
@@ -10089,7 +10052,6 @@ def create_app(
|
|||||||
tls_manager: Any = None,
|
tls_manager: Any = None,
|
||||||
console_url: str = "",
|
console_url: str = "",
|
||||||
router: ConsoleRouter | None = None,
|
router: ConsoleRouter | None = None,
|
||||||
rebalancer: Any = None,
|
|
||||||
console_metrics: ConsoleMetrics | None = None,
|
console_metrics: ConsoleMetrics | None = None,
|
||||||
) -> Starlette:
|
) -> Starlette:
|
||||||
"""Build the Starlette ASGI application for the console dashboard."""
|
"""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/node/{node_id}", cluster_node_detail),
|
||||||
Route("/api/cluster/snapshot", cluster_snapshot),
|
Route("/api/cluster/snapshot", cluster_snapshot),
|
||||||
Route("/api/cluster/events", cluster_events_sse),
|
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/workstreams/new", route_create, methods=["POST"]),
|
||||||
Route("/api/route/send", route_proxy, methods=["POST"]),
|
Route("/api/route/send", route_proxy, methods=["POST"]),
|
||||||
Route("/api/route/approve", route_proxy, methods=["POST"]),
|
Route("/api/route/approve", route_proxy, methods=["POST"]),
|
||||||
@@ -10564,13 +10526,6 @@ def create_app(
|
|||||||
admin_set_node_metadata,
|
admin_set_node_metadata,
|
||||||
methods=["PUT"],
|
methods=["PUT"],
|
||||||
),
|
),
|
||||||
# Hash ring
|
|
||||||
Route("/api/admin/ring/status", admin_ring_status),
|
|
||||||
Route(
|
|
||||||
"/api/admin/ring/rebalance",
|
|
||||||
admin_ring_rebalance,
|
|
||||||
methods=["POST"],
|
|
||||||
),
|
|
||||||
# TLS / ACME
|
# TLS / ACME
|
||||||
Route("/api/admin/tls/ca", tls_ca_status),
|
Route("/api/admin/tls/ca", tls_ca_status),
|
||||||
Route("/api/admin/tls/ca.pem", tls_ca_cert),
|
Route("/api/admin/tls/ca.pem", tls_ca_cert),
|
||||||
@@ -10630,7 +10585,6 @@ def create_app(
|
|||||||
app.state.console_url = console_url
|
app.state.console_url = console_url
|
||||||
app.state.tls_manager = tls_manager
|
app.state.tls_manager = tls_manager
|
||||||
app.state.router = router
|
app.state.router = router
|
||||||
app.state.rebalancer = rebalancer
|
|
||||||
app.state.console_metrics = console_metrics or ConsoleMetrics()
|
app.state.console_metrics = console_metrics or ConsoleMetrics()
|
||||||
|
|
||||||
# Mount ACME responder whenever a TLS manager is configured.
|
# Mount ACME responder whenever a TLS manager is configured.
|
||||||
@@ -10861,32 +10815,6 @@ def main() -> None:
|
|||||||
except Exception:
|
except Exception:
|
||||||
log.debug("Failed to sync TLS state to ConfigStore", exc_info=True)
|
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(
|
app = create_app(
|
||||||
collector=collector,
|
collector=collector,
|
||||||
jwt_secret=jwt_secret,
|
jwt_secret=jwt_secret,
|
||||||
@@ -10896,7 +10824,6 @@ def main() -> None:
|
|||||||
tls_manager=tls_mgr,
|
tls_manager=tls_mgr,
|
||||||
console_url=console_url,
|
console_url=console_url,
|
||||||
router=router,
|
router=router,
|
||||||
rebalancer=rebalancer,
|
|
||||||
console_metrics=console_metrics,
|
console_metrics=console_metrics,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -195,7 +195,6 @@ APPROVE_PATHS: frozenset[str] = frozenset(
|
|||||||
"/api/_internal/config-reload",
|
"/api/_internal/config-reload",
|
||||||
"/api/_internal/mcp-reload",
|
"/api/_internal/mcp-reload",
|
||||||
"/api/_internal/model-reload",
|
"/api/_internal/model-reload",
|
||||||
"/api/_internal/migrate",
|
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
ADMIN_PREFIX = "/api/admin/"
|
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
|
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.log import get_logger
|
||||||
from turnstone.core.storage import get_storage
|
from turnstone.core.storage import get_storage
|
||||||
from turnstone.core.workstream import WorkstreamKind
|
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)
|
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:
|
def delete_workstream_override(ws_id: str) -> None:
|
||||||
"""Fire-and-forget override deletion."""
|
"""Fire-and-forget override deletion."""
|
||||||
try:
|
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. "
|
"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).",
|
"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 ---------------------------------------------------------------
|
# -- node ---------------------------------------------------------------
|
||||||
SettingDef(
|
SettingDef(
|
||||||
"node.weight",
|
"node.weight",
|
||||||
"int",
|
"int",
|
||||||
1,
|
1,
|
||||||
"Node weight for hash ring distribution",
|
"Node weight for rendezvous routing",
|
||||||
"node",
|
"node",
|
||||||
min_value=1,
|
min_value=1,
|
||||||
max_value=100,
|
max_value=100,
|
||||||
help="Relative capacity of this server node. A node with weight 2 receives "
|
help="Relative capacity of this server node. A node with weight 2 wins "
|
||||||
"roughly twice as many bucket assignments (and therefore workstreams) as a "
|
"roughly twice as many ws_id rendezvous selections (and therefore "
|
||||||
"node with weight 1.",
|
"receives twice as many workstreams) as a node with weight 1.",
|
||||||
),
|
),
|
||||||
# -- coordinator --------------------------------------------------------
|
# -- coordinator --------------------------------------------------------
|
||||||
SettingDef(
|
SettingDef(
|
||||||
|
|||||||
@@ -16,11 +16,9 @@ from turnstone.core.log import get_logger
|
|||||||
from turnstone.core.storage._schema import (
|
from turnstone.core.storage._schema import (
|
||||||
api_tokens,
|
api_tokens,
|
||||||
audit_events,
|
audit_events,
|
||||||
bucket_stats,
|
|
||||||
channel_routes,
|
channel_routes,
|
||||||
channel_users,
|
channel_users,
|
||||||
conversations,
|
conversations,
|
||||||
hash_ring_buckets,
|
|
||||||
heuristic_rules,
|
heuristic_rules,
|
||||||
intent_verdicts,
|
intent_verdicts,
|
||||||
mcp_servers,
|
mcp_servers,
|
||||||
@@ -1851,122 +1849,7 @@ class PostgreSQLBackend:
|
|||||||
rows = conn.execute(stmt).fetchall()
|
rows = conn.execute(stmt).fetchall()
|
||||||
return {r[0] for r in rows}
|
return {r[0] for r in rows}
|
||||||
|
|
||||||
# -- Hash ring routing -----------------------------------------------------
|
# -- Routing overrides -----------------------------------------------------
|
||||||
|
|
||||||
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()
|
|
||||||
|
|
||||||
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
||||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||||
@@ -1998,16 +1881,6 @@ class PostgreSQLBackend:
|
|||||||
).fetchall()
|
).fetchall()
|
||||||
return [dict(r._mapping) for r in rows]
|
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 -----------------------------------------------------------------
|
# -- Roles -----------------------------------------------------------------
|
||||||
|
|
||||||
def create_role(
|
def create_role(
|
||||||
|
|||||||
@@ -787,39 +787,7 @@ class StorageBackend(Protocol):
|
|||||||
"""Return node_ids where ALL key=value filters match (exact match)."""
|
"""Return node_ids where ALL key=value filters match (exact match)."""
|
||||||
...
|
...
|
||||||
|
|
||||||
# -- Hash ring routing ---
|
# -- Routing overrides ---
|
||||||
|
|
||||||
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."""
|
|
||||||
...
|
|
||||||
|
|
||||||
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
||||||
"""Pin a workstream to a specific node. Upserts."""
|
"""Pin a workstream to a specific node. Upserts."""
|
||||||
@@ -833,10 +801,6 @@ class StorageBackend(Protocol):
|
|||||||
"""Return all overrides."""
|
"""Return all overrides."""
|
||||||
...
|
...
|
||||||
|
|
||||||
def list_workstream_routing_data(self) -> list[tuple[str, str]]:
|
|
||||||
"""Return (ws_id, state) for all non-override workstreams."""
|
|
||||||
...
|
|
||||||
|
|
||||||
# -- Roles (RBAC) ----------------------------------------------------------
|
# -- Roles (RBAC) ----------------------------------------------------------
|
||||||
|
|
||||||
def create_role(
|
def create_role(
|
||||||
|
|||||||
@@ -262,26 +262,9 @@ node_metadata = sa.Table(
|
|||||||
sa.Index("idx_node_metadata_key", node_metadata.c.key)
|
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 = sa.Table(
|
||||||
"workstream_overrides",
|
"workstream_overrides",
|
||||||
metadata,
|
metadata,
|
||||||
|
|||||||
@@ -16,11 +16,9 @@ from turnstone.core.log import get_logger
|
|||||||
from turnstone.core.storage._schema import (
|
from turnstone.core.storage._schema import (
|
||||||
api_tokens,
|
api_tokens,
|
||||||
audit_events,
|
audit_events,
|
||||||
bucket_stats,
|
|
||||||
channel_routes,
|
channel_routes,
|
||||||
channel_users,
|
channel_users,
|
||||||
conversations,
|
conversations,
|
||||||
hash_ring_buckets,
|
|
||||||
heuristic_rules,
|
heuristic_rules,
|
||||||
intent_verdicts,
|
intent_verdicts,
|
||||||
mcp_servers,
|
mcp_servers,
|
||||||
@@ -1965,122 +1963,7 @@ class SQLiteBackend:
|
|||||||
rows = conn.execute(stmt).fetchall()
|
rows = conn.execute(stmt).fetchall()
|
||||||
return {r[0] for r in rows}
|
return {r[0] for r in rows}
|
||||||
|
|
||||||
# -- Hash ring routing -----------------------------------------------------
|
# -- Routing overrides -----------------------------------------------------
|
||||||
|
|
||||||
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()
|
|
||||||
|
|
||||||
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
def set_workstream_override(self, ws_id: str, node_id: str, reason: str = "targeted") -> None:
|
||||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||||
@@ -2112,16 +1995,6 @@ class SQLiteBackend:
|
|||||||
).fetchall()
|
).fetchall()
|
||||||
return [dict(r._mapping) for r in rows]
|
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 -----------------------------------------------------------------
|
# -- Roles -----------------------------------------------------------------
|
||||||
|
|
||||||
def create_role(
|
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
|
ERROR = "error" # last operation failed
|
||||||
|
|
||||||
|
|
||||||
_ACTIVE_STATES = {"running", "thinking", "attention"}
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Workstream dataclass
|
# Workstream dataclass
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -251,11 +248,9 @@ class WorkstreamManager:
|
|||||||
if first_evicted is not None:
|
if first_evicted is not None:
|
||||||
self._cleanup_ui(first_evicted)
|
self._cleanup_ui(first_evicted)
|
||||||
self._last_evicted = 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.memory import delete_workstream_override as _dwo1
|
||||||
from turnstone.core.metrics import metrics as _m1
|
from turnstone.core.metrics import metrics as _m1
|
||||||
|
|
||||||
_dbc1(first_evicted.id, active=False) # evicted ws is always idle
|
|
||||||
_dwo1(first_evicted.id)
|
_dwo1(first_evicted.id)
|
||||||
_m1.record_eviction()
|
_m1.record_eviction()
|
||||||
|
|
||||||
@@ -320,7 +315,7 @@ class WorkstreamManager:
|
|||||||
self._active_id = ws.id
|
self._active_id = ws.id
|
||||||
|
|
||||||
# Persist to storage only after successful insertion
|
# 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(
|
register_workstream(
|
||||||
ws.id,
|
ws.id,
|
||||||
@@ -332,17 +327,14 @@ class WorkstreamManager:
|
|||||||
kind=kind,
|
kind=kind,
|
||||||
parent_ws_id=ws.parent_ws_id,
|
parent_ws_id=ws.parent_ws_id,
|
||||||
)
|
)
|
||||||
increment_bucket_count(ws.id)
|
|
||||||
|
|
||||||
# Cleanup second-phase eviction outside the lock.
|
# Cleanup second-phase eviction outside the lock.
|
||||||
if second_evicted is not None:
|
if second_evicted is not None:
|
||||||
self._cleanup_ui(second_evicted)
|
self._cleanup_ui(second_evicted)
|
||||||
self._last_evicted = 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.memory import delete_workstream_override as _dwo2
|
||||||
from turnstone.core.metrics import metrics as _m2
|
from turnstone.core.metrics import metrics as _m2
|
||||||
|
|
||||||
_dbc2(second_evicted.id, active=False)
|
|
||||||
_dwo2(second_evicted.id)
|
_dwo2(second_evicted.id)
|
||||||
_m2.record_eviction()
|
_m2.record_eviction()
|
||||||
return ws
|
return ws
|
||||||
@@ -414,20 +406,17 @@ class WorkstreamManager:
|
|||||||
ws = self._workstreams.pop(ws_id, None)
|
ws = self._workstreams.pop(ws_id, None)
|
||||||
if ws is None:
|
if ws is None:
|
||||||
return False
|
return False
|
||||||
was_active = ws.state.value in _ACTIVE_STATES
|
|
||||||
self._order.remove(ws_id)
|
self._order.remove(ws_id)
|
||||||
if self._active_id == ws_id:
|
if self._active_id == ws_id:
|
||||||
self._active_id = self._order[0]
|
self._active_id = self._order[0]
|
||||||
# Unblock any waiting approval/plan events so worker thread can exit
|
# Unblock any waiting approval/plan events so worker thread can exit
|
||||||
self._cleanup_ui(ws)
|
self._cleanup_ui(ws)
|
||||||
from turnstone.core.memory import (
|
from turnstone.core.memory import (
|
||||||
decrement_bucket_count,
|
|
||||||
delete_workstream_override,
|
delete_workstream_override,
|
||||||
update_workstream_state,
|
update_workstream_state,
|
||||||
)
|
)
|
||||||
|
|
||||||
update_workstream_state(ws_id, "closed")
|
update_workstream_state(ws_id, "closed")
|
||||||
decrement_bucket_count(ws_id, active=was_active)
|
|
||||||
delete_workstream_override(ws_id)
|
delete_workstream_override(ws_id)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
@@ -489,18 +478,12 @@ class WorkstreamManager:
|
|||||||
ws = self._workstreams.get(ws_id)
|
ws = self._workstreams.get(ws_id)
|
||||||
if ws:
|
if ws:
|
||||||
with ws._lock:
|
with ws._lock:
|
||||||
old_active = ws.state.value in _ACTIVE_STATES
|
|
||||||
ws.state = state
|
ws.state = state
|
||||||
ws.last_active = time.monotonic()
|
ws.last_active = time.monotonic()
|
||||||
ws.error_message = error_msg
|
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)
|
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:
|
if self._on_state_change:
|
||||||
self._on_state_change(ws_id, state)
|
self._on_state_change(ws_id, state)
|
||||||
|
|
||||||
|
|||||||
@@ -4224,30 +4224,6 @@ def internal_model_status(request: Request) -> JSONResponse:
|
|||||||
return JSONResponse({"models": models})
|
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
|
# Global SSE fan-out
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -4705,7 +4681,6 @@ def create_app(
|
|||||||
methods=["POST"],
|
methods=["POST"],
|
||||||
),
|
),
|
||||||
Route("/api/_internal/model-status", internal_model_status),
|
Route("/api/_internal/model-status", internal_model_status),
|
||||||
Route("/api/_internal/migrate", internal_migrate, methods=["POST"]),
|
|
||||||
],
|
],
|
||||||
),
|
),
|
||||||
Route("/health", health),
|
Route("/health", health),
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "inspect_workstream",
|
"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": {
|
"parameters": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "list_nodes",
|
"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": {
|
"parameters": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "spawn_workstream",
|
"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": {
|
"parameters": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
@@ -22,7 +22,7 @@
|
|||||||
},
|
},
|
||||||
"target_node": {
|
"target_node": {
|
||||||
"type": "string",
|
"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": []
|
"required": []
|
||||||
|
|||||||
Reference in New Issue
Block a user