orchestrata_ultimate.kernel.performance.cluster

Reference › orchestrata_ultimate › kernel › performance › cluster

The clustering primitives of the performance package, SPEC-001 story US-010 (REQ-001-US-010-024 to -026). They let several kernel nodes share sessions and volks and elect one coordinator. manager composes them; the package initialiser re-exports them.

• ConsistentHash places each node on a hash ring with a number of virtual replicas. add_node and remove_node change the ring, get_node returns the node for a key, get_nodes returns several distinct nodes and get_all_nodes lists them.
• LeaderElection elects a leader. With a Redis client, try_acquire_leadership sets a leader key only if absent with an expiry (30 seconds by default) and, if the key is already held by this node, renews it; without Redis, the first node becomes leader. release_leadership, is_leader, get_leader, on_become_leader and on_become_follower complete it. Redis errors are logged and count as not acquired.
• NodeHealth is the record of one node's health.
• ClusterManager runs a heartbeat loop at an interval (5 seconds by default), checks node health against a failure threshold, and on failure removes the node from the ring and calls _rebalance. register_node, update_node_health, get_node_for_session, get_node_for_volk, get_healthy_nodes and get_cluster_stats serve callers, and on_node_added, on_node_removed and on_rebalance register callbacks whose errors are logged.

At a glance

Class diagram

Dependency diagram

API

Classes

ConsistentHash

Decorators: traces(spec='001', story='US-010', reqs=['REQ-001-US-010-024']) · Line 27

Consistent hashing for routing (REQ-001-US-010-024).

LeaderElection

Decorators: traces(spec='001', story='US-010', reqs=['REQ-001-US-010-025']) · Line 137

Leader election for coordination (REQ-001-US-010-025).

NodeHealth (dataclass)

Decorators: traces(spec='001', story='US-010', reqs=['REQ-001-US-010-026']), dataclass · Line 288

Health status of a kernel node (REQ-001-US-010-026).

ClusterManager

Decorators: traces(spec='001', story='US-010', reqs=['REQ-001-US-010-021', 'REQ-001-US-010-026', 'REQ-001-US-010-027']) · Line 309

Manages cluster of kernel instances (REQ-001-US-010-021, REQ-001-US-010-026, REQ-001-US-010-027).

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/kernel/performance/cluster.py

# mypy: ignore-errors
"""Performance subsystem: cluster."""

import asyncio
import bisect
import hashlib
import threading
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import datetime
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.kernel.performance._common import (
    CLUSTER_DEFAULT_HEARTBEAT_INTERVAL_SECONDS,
    CONSISTENT_HASH_DEFAULT_REPLICAS,
    _logger,
    _utc_now,
)


@traces(
    spec="001",
    story="US-010",
    reqs=["REQ-001-US-010-024"],
)
class ConsistentHash:
    """Consistent hashing for routing (REQ-001-US-010-024).

    Provides stable routing even when nodes change.
    """

    def __init__(self, replicas: int = CONSISTENT_HASH_DEFAULT_REPLICAS) -> None:
        """Initialize consistent hash ring.

        Args:
            replicas: Virtual nodes per real node
        """
        self._replicas = replicas
        self._ring: list[tuple[int, str]] = []  # (hash, node_id)
        self._nodes: set[str] = set()
        self._lock = threading.Lock()

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-024"])
    def _hash(self, key: str) -> int:
        """Hash a key to ring position."""
        return int(hashlib.md5(key.encode()).hexdigest(), 16)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-024"])
    def add_node(self, node_id: str) -> None:
        """Add node to ring.

        Args:
            node_id: Node identifier
        """
        with self._lock:
            if node_id in self._nodes:
                return

            self._nodes.add(node_id)
            for i in range(self._replicas):
                key = f"{node_id}:{i}"
                h = self._hash(key)
                bisect.insort(self._ring, (h, node_id))

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-024"])
    def remove_node(self, node_id: str) -> None:
        """Remove node from ring.

        Args:
            node_id: Node identifier
        """
        with self._lock:
            if node_id not in self._nodes:
                return

            self._nodes.discard(node_id)
            self._ring = [(h, n) for h, n in self._ring if n != node_id]

    @traces(
        spec="001", story="US-010", reqs=["REQ-001-US-010-022", "REQ-001-US-010-023"]
    )
    def get_node(self, key: str) -> str | None:
        """Get node for key (REQ-001-US-010-022, REQ-001-US-010-023 - sharding).

        Args:
            key: Key to route (session_id or volk_alias)

        Returns:
            Node identifier or None if no nodes
        """
        with self._lock:
            if not self._ring:
                return None

            h = self._hash(key)
            idx = bisect.bisect_left(self._ring, (h, ""))
            if idx >= len(self._ring):
                idx = 0
            return self._ring[idx][1]

    def get_nodes(self, key: str, count: int) -> list[str]:
        """Get multiple nodes for key (for replication).

        Args:
            key: Key to route
            count: Number of nodes to return

        Returns:
            List of node identifiers
        """
        with self._lock:
            if not self._ring or count <= 0:
                return []

            h = self._hash(key)
            idx = bisect.bisect_left(self._ring, (h, ""))

            result = []
            seen = set()
            for i in range(len(self._ring)):
                node = self._ring[(idx + i) % len(self._ring)][1]
                if node not in seen:
                    seen.add(node)
                    result.append(node)
                    if len(result) >= count:
                        break
            return result

    def get_all_nodes(self) -> list[str]:
        """Get all nodes."""
        with self._lock:
            return list(self._nodes)


@traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
class LeaderElection:
    """Leader election for coordination (REQ-001-US-010-025).

    Simple leader election using Redis or local fallback.
    """

    def __init__(
        self,
        node_id: str,
        ttl_seconds: float = 30.0,
        redis_client: Any | None = None,
        leader_key: str = "orchestrata:leader",
    ) -> None:
        """Initialize leader election.

        Args:
            node_id: This node's identifier
            ttl_seconds: Leadership TTL
            redis_client: Optional Redis client for distributed election
            leader_key: Redis key for leader election (default: orchestrata:leader)
        """
        self._node_id = node_id
        self._ttl_seconds = ttl_seconds
        self._redis = redis_client
        self._leader_key = leader_key

        self._is_leader = False
        self._current_leader: str | None = None
        self._leadership_acquired_at: datetime | None = None
        self._lock = asyncio.Lock()

        self._on_leader_callbacks: list[Callable[[], Awaitable[None]]] = []
        self._on_follower_callbacks: list[Callable[[], Awaitable[None]]] = []

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    async def try_acquire_leadership(self) -> bool:
        """Try to become leader.

        Returns:
            True if this node is now leader
        """
        async with self._lock:
            if self._redis:
                # Use Redis SETNX for distributed lock
                try:
                    key = self._leader_key
                    acquired = await self._redis.set(
                        key,
                        self._node_id,
                        nx=True,
                        ex=int(self._ttl_seconds),
                    )
                    if acquired:
                        self._is_leader = True
                        self._current_leader = self._node_id
                        self._leadership_acquired_at = _utc_now()
                        await self._notify_leader()
                        return True

                    # Check who is leader
                    self._current_leader = await self._redis.get(key)
                    if self._current_leader == self._node_id:
                        # We're still leader, renew
                        self._is_leader = True
                        await self._redis.expire(key, int(self._ttl_seconds))
                        return True

                    return False
                except Exception as e:
                    _logger.error(f"Leader election error: {e}")
                    return False
            else:
                # Local fallback - first node wins
                if self._current_leader is None:
                    self._is_leader = True
                    self._current_leader = self._node_id
                    self._leadership_acquired_at = _utc_now()
                    await self._notify_leader()
                    return True
                return self._is_leader

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    async def release_leadership(self) -> None:
        """Release leadership."""
        async with self._lock:
            if not self._is_leader:
                return

            if self._redis:
                try:
                    key = self._leader_key
                    current = await self._redis.get(key)
                    if current == self._node_id:
                        await self._redis.delete(key)
                except Exception as exc:
                    _logger.warning(
                        "Failed to release leader key from Redis: %s",
                        exc,
                        exc_info=True,
                    )

            self._is_leader = False
            self._current_leader = None
            self._leadership_acquired_at = None
            await self._notify_follower()

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    async def _notify_leader(self) -> None:
        """Notify callbacks of becoming leader."""
        for callback in self._on_leader_callbacks:
            try:
                await callback()
            except Exception as exc:
                _logger.warning(
                    "Leader notification callback failed: %s", exc, exc_info=True
                )

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    async def _notify_follower(self) -> None:
        """Notify callbacks of becoming follower."""
        for callback in self._on_follower_callbacks:
            try:
                await callback()
            except Exception as exc:
                _logger.warning(
                    "Follower notification callback failed: %s", exc, exc_info=True
                )

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    def is_leader(self) -> bool:
        """Check if this node is leader."""
        return self._is_leader

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    def get_leader(self) -> str | None:
        """Get current leader node ID."""
        return self._current_leader

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    def on_become_leader(self, callback: Callable[[], Awaitable[None]]) -> None:
        """Register callback for becoming leader."""
        self._on_leader_callbacks.append(callback)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-025"])
    def on_become_follower(self, callback: Callable[[], Awaitable[None]]) -> None:
        """Register callback for becoming follower."""
        self._on_follower_callbacks.append(callback)


@traces(spec="001", story="US-010", reqs=["REQ-001-US-010-026"])
@dataclass
class NodeHealth:
    """Health status of a kernel node (REQ-001-US-010-026)."""

    node_id: str
    is_healthy: bool
    last_heartbeat: datetime
    load: float
    active_sessions: int
    error_count: int = 0
    consecutive_failures: int = 0


@traces(
    spec="001",
    story="US-010",
    reqs=[
        "REQ-001-US-010-021",
        "REQ-001-US-010-026",
        "REQ-001-US-010-027",
    ],
)
class ClusterManager:
    """Manages cluster of kernel instances (REQ-001-US-010-021, REQ-001-US-010-026, REQ-001-US-010-027).

    Handles:
    - Node registration and discovery
    - Health monitoring
    - Failure detection and handling
    - Load rebalancing
    """

    def __init__(
        self,
        node_id: str,
        heartbeat_interval: float = CLUSTER_DEFAULT_HEARTBEAT_INTERVAL_SECONDS,
        failure_threshold: int = 3,
    ) -> None:
        """Initialize cluster manager.

        Args:
            node_id: This node's identifier
            heartbeat_interval: Seconds between heartbeats
            failure_threshold: Consecutive failures before marking unhealthy
        """
        self._node_id = node_id
        self._heartbeat_interval = heartbeat_interval
        self._failure_threshold = failure_threshold

        self._nodes: dict[str, NodeHealth] = {}
        self._consistent_hash = ConsistentHash()
        self._lock = asyncio.Lock()

        self._heartbeat_task: asyncio.Task | None = None
        self._running = False

        self._on_node_added: list[Callable[[str], Awaitable[None]]] = []
        self._on_node_removed: list[Callable[[str], Awaitable[None]]] = []
        self._on_rebalance: list[Callable[[], Awaitable[None]]] = []

    async def start(self) -> None:
        """Start cluster manager."""
        self._running = True
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())

        # Register self
        await self.register_node(
            self._node_id,
            NodeHealth(
                node_id=self._node_id,
                is_healthy=True,
                last_heartbeat=_utc_now(),
                load=0.0,
                active_sessions=0,
            ),
        )

    async def stop(self) -> None:
        """Stop cluster manager."""
        self._running = False
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                _logger.debug("heartbeat task cancelled")

    async def _heartbeat_loop(self) -> None:
        """Periodic heartbeat and health check loop."""
        while self._running:
            try:
                await asyncio.sleep(self._heartbeat_interval)
                await self._check_node_health()
            except asyncio.CancelledError:
                break
            except Exception as e:
                _logger.error(f"Heartbeat error: {e}")

    async def _check_node_health(self) -> None:
        """Check health of all nodes."""
        now = _utc_now()
        failed_nodes = []

        async with self._lock:
            for node_id, health in self._nodes.items():
                if node_id == self._node_id:
                    # Update self
                    health.last_heartbeat = now
                    continue

                # Check if node is stale
                elapsed = (now - health.last_heartbeat).total_seconds()
                if elapsed > self._heartbeat_interval * self._failure_threshold:
                    health.consecutive_failures += 1
                    if health.consecutive_failures >= self._failure_threshold:
                        health.is_healthy = False
                        failed_nodes.append(node_id)

        # Handle failed nodes
        for node_id in failed_nodes:
            await self._handle_node_failure(node_id)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-026"])
    async def _handle_node_failure(self, node_id: str) -> None:
        """Handle node failure (REQ-001-US-010-026).

        Args:
            node_id: Failed node identifier
        """
        _logger.warning(f"Node {node_id} failed, triggering rebalance")

        async with self._lock:
            self._consistent_hash.remove_node(node_id)

        # Notify callbacks
        for callback in self._on_node_removed:
            try:
                await callback(node_id)
            except Exception as exc:
                _logger.warning(
                    "Node-removed callback failed for node %s: %s",
                    node_id,
                    exc,
                    exc_info=True,
                )

        # Trigger rebalance
        await self._rebalance()

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-027"])
    async def _rebalance(self) -> None:
        """Rebalance load across nodes (REQ-001-US-010-027).

        Called when nodes are added/removed.
        """
        for callback in self._on_rebalance:
            try:
                await callback()
            except Exception as exc:
                _logger.warning("Rebalance callback failed: %s", exc, exc_info=True)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-021"])
    async def register_node(self, node_id: str, health: NodeHealth) -> None:
        """Register a node in the cluster.

        Args:
            node_id: Node identifier
            health: Initial health status
        """
        async with self._lock:
            is_new = node_id not in self._nodes
            self._nodes[node_id] = health
            self._consistent_hash.add_node(node_id)

        if is_new:
            for callback in self._on_node_added:
                try:
                    await callback(node_id)
                except Exception as exc:
                    _logger.warning(
                        "Node-added callback failed for node %s: %s",
                        node_id,
                        exc,
                        exc_info=True,
                    )

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-026"])
    async def update_node_health(self, node_id: str, health: NodeHealth) -> None:
        """Update node health status.

        Args:
            node_id: Node identifier
            health: Updated health status
        """
        async with self._lock:
            if node_id in self._nodes:
                old_health = self._nodes[node_id]
                self._nodes[node_id] = health
                if health.is_healthy and not old_health.is_healthy:
                    # Node recovered
                    self._consistent_hash.add_node(node_id)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-022"])
    async def get_node_for_session(self, session_id: str) -> str | None:
        """Get node for session (REQ-001-US-010-022).

        Args:
            session_id: Session identifier

        Returns:
            Node identifier
        """
        async with self._lock:
            return self._consistent_hash.get_node(session_id)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-023"])
    async def get_node_for_volk(self, volk_alias: str) -> str | None:
        """Get node for Volk (REQ-001-US-010-023).

        Args:
            volk_alias: Volk alias

        Returns:
            Node identifier
        """
        async with self._lock:
            return self._consistent_hash.get_node(volk_alias)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-026"])
    async def get_healthy_nodes(self) -> list[str]:
        """Get list of healthy nodes."""
        async with self._lock:
            return [
                node_id for node_id, health in self._nodes.items() if health.is_healthy
            ]

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-029"])
    async def get_cluster_stats(self) -> dict[str, Any]:
        """Get cluster statistics."""
        async with self._lock:
            healthy = sum(1 for h in self._nodes.values() if h.is_healthy)
            total_load = sum(h.load for h in self._nodes.values())
            total_sessions = sum(h.active_sessions for h in self._nodes.values())

            return {
                "total_nodes": len(self._nodes),
                "healthy_nodes": healthy,
                "total_load": total_load,
                "avg_load": total_load / len(self._nodes) if self._nodes else 0,
                "total_sessions": total_sessions,
                "nodes": {
                    node_id: {
                        "is_healthy": h.is_healthy,
                        "load": h.load,
                        "active_sessions": h.active_sessions,
                        "error_count": h.error_count,
                    }
                    for node_id, h in self._nodes.items()
                },
            }

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-021"])
    def on_node_added(self, callback: Callable[[str], Awaitable[None]]) -> None:
        """Register callback for node addition."""
        self._on_node_added.append(callback)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-026"])
    def on_node_removed(self, callback: Callable[[str], Awaitable[None]]) -> None:
        """Register callback for node removal."""
        self._on_node_removed.append(callback)

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-027"])
    def on_rebalance(self, callback: Callable[[], Awaitable[None]]) -> None:
        """Register callback for rebalance events."""
        self._on_rebalance.append(callback)