orchestrata_ultimate.semantic_graph.ha_controller

Reference › orchestrata_ultimate › semantic_graph › ha_controller

A high-availability controller for Memgraph Community replication, the engine of the shared semantic graph (memgraph), which has no built-in HA. The controller monitors a set of data instances, enforces a single MAIN role, fails over automatically and reintegrates recovered nodes as replicas. It also serves an HTTP API so clients can discover the current MAIN URI without hardcoding it. It is deliberately a single-controller design that persists its state to SQLite so it survives restarts; its docstring notes that a three-node Raft controller cluster could follow. The command-line wrapper is ha_controller_cli, and main_discovery mentions it as the legacy discovery route.

Three dataclasses model the cluster: the frozen NodeConfig (name, bolt URI, replication host), NodeState (runtime health of one instance) and ClusterStatus (current MAIN, epoch and node states, returned by the status endpoint). MemgraphHAController takes the node list, a state database path, a poll interval (2 seconds by default), a replication mode (ASYNC by default) and a stall threshold. run loads state (_load_state restores the current MAIN from SQLite or bootstraps from the first node) and starts the monitoring loop. Each _poll_cycle updates node health (_update_node_health, _check_node), reconciles roles (_reconcile_roles), performs failover when MAIN is unhealthy (_perform_failover) and idempotently registers healthy non-MAIN nodes as replicas (_ensure_replicas_registered), reading SHOW REPLICAS data and tracking per-cycle replication progress to detect a frozen but ready replica. The HTTP surface is a FastAPI app exposed through the app property.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

NodeConfig (dataclass)

Decorators: dataclass(frozen=True) · Line 37

Configuration for one Memgraph data instance.

NodeState (dataclass)

Decorators: dataclass · Line 49

Runtime state of one Memgraph data instance.

ClusterStatus (dataclass)

Decorators: dataclass · Line 60

Snapshot returned by the HTTP status endpoint.

MemgraphHAController

Line 99

Monitor and manage a Memgraph replication cluster.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/semantic_graph/ha_controller.py

"""High-availability controller for Memgraph Community replication.

The controller monitors a set of Memgraph data instances, enforces a single
MAIN role, performs automatic failover, and reintegrates recovered nodes as
REPLICAs (autofailback). It exposes an HTTP API so clients can discover the
current MAIN URI without hardcoding it.

This is intentionally a single-controller design. It is the source of truth for
the cluster and persists its state to SQLite so it survives restarts. A future
iteration could turn this into a 3-node Raft controller cluster.
"""

from __future__ import annotations

import logging
import sqlite3
import time
from dataclasses import dataclass
from typing import Any

import anyio
import anyio.to_thread
from fastapi import FastAPI

logger = logging.getLogger(__name__)

# SHOW REPLICAS data_info statuses that mean "healthy, leave it alone": Memgraph
# reports "ready" (in sync) plus "replicating"/"recovery" (actively catching up).
# Anything else (e.g. "invalid") — or no entry at all — means the replica must be
# (re)registered.
_HEALTHY_REPLICA_STATUSES: frozenset[str] = frozenset(
    {"ready", "replicating", "recovery"}
)


@dataclass(frozen=True)
class NodeConfig:
    """Configuration for one Memgraph data instance."""

    name: str
    bolt_uri: str
    replication_host: str
    replication_port: int = 10000
    username: str = ""
    password: str = ""


@dataclass
class NodeState:
    """Runtime state of one Memgraph data instance."""

    config: NodeConfig
    role: str = "unknown"
    healthy: bool = False
    last_seen: float = 0.0
    error: str | None = None


@dataclass
class ClusterStatus:
    """Snapshot returned by the HTTP status endpoint."""

    current_main: str | None
    epoch: int
    nodes: list[NodeState]


def probe_coordinator_for_main(bolt_uri: str) -> str | None:
    """Ask one Memgraph coordinator (``SHOW INSTANCES``) for the current MAIN
    data-instance's bolt URI, or ``None`` (unreachable, not a coordinator, or no
    MAIN reported — all best-effort).

    ``SHOW INSTANCES`` is a Raft admin query only a *coordinator* answers over a
    direct (non-routed) ``bolt://`` connection. This is the single raw-driver
    coordinator probe for MAIN discovery; it lives here — the allowlisted
    Memgraph HA seam — so clients like :mod:`semantic_graph.main_discovery` need
    no raw graph access of their own (D13 graph-seam rule).
    """
    from orchestrata_ultimate.core.graph_factory import NO_AUTH, create_graph

    try:
        driver = create_graph(uri=bolt_uri, auth=NO_AUTH, connection_timeout=5)
    except Exception:  # noqa: BLE001 - probing is inherently best-effort
        return None
    try:
        with driver.session() as session:
            for record in session.run("SHOW INSTANCES"):
                if record.get("role") == "main":
                    bolt_server = record.get("bolt_server")
                    if bolt_server:
                        return f"bolt://{bolt_server}"
    except Exception:  # noqa: BLE001 - probing is inherently best-effort
        return None
    finally:
        driver.close()
    return None


class MemgraphHAController:
    """Monitor and manage a Memgraph replication cluster."""

    def __init__(
        self,
        nodes: list[NodeConfig],
        db_path: str = "/var/lib/memgraph-ha/state.db",
        poll_interval: float = 2.0,
        replication_mode: str = "ASYNC",
        stall_cycles_threshold: int = 10,
    ) -> None:
        self._nodes = {node.name: NodeState(config=node) for node in nodes}
        self._db_path = db_path
        self._poll_interval = poll_interval
        self._replication_mode = replication_mode.upper()
        self._stall_cycles_threshold = stall_cycles_threshold
        # Per-replica (last_seen_ts, consecutive_frozen_cycles) — detects
        # replicas that report 'ready' while their replication timestamp is
        # frozen with data outstanding (post-reboot incident 2026-07-08).
        # See _note_replica_progress.
        self._replica_progress: dict[str, tuple[int | None, int]] = {}
        self._persisted_main: str = ""
        self._epoch: int = 0
        self._app = FastAPI(title="Memgraph HA Controller")
        self._setup_routes()

    @property
    def app(self) -> FastAPI:
        return self._app

    def _setup_routes(self) -> None:
        @self._app.get("/status")
        async def status() -> dict[str, Any]:
            return self._status_dict()

        @self._app.get("/main")
        async def main() -> dict[str, Any]:
            current = self._current_main()
            if current is None:
                return {"main": None, "uri": None}
            return {"main": current.config.name, "uri": current.config.bolt_uri}

        @self._app.post("/failover")
        async def failover() -> dict[str, Any]:
            await self._perform_failover()
            return self._status_dict()

    async def run(self) -> None:
        """Initialize state and start the monitoring loop."""
        await anyio.to_thread.run_sync(self._init_db)
        await self._load_state()
        while True:
            try:
                await self._poll_cycle()
            except Exception as exc:
                logger.exception("Polling cycle failed: %s", exc)
            await anyio.sleep(self._poll_interval)

    # ------------------------------------------------------------------
    # State persistence
    # ------------------------------------------------------------------

    def _init_db(self) -> None:
        parent = self._db_path.rsplit("/", 1)[0]
        if parent:
            import os

            os.makedirs(parent, exist_ok=True)
        with sqlite3.connect(self._db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS cluster_state (
                    id INTEGER PRIMARY KEY CHECK (id = 1),
                    current_main TEXT NOT NULL,
                    epoch INTEGER NOT NULL,
                    updated_at REAL NOT NULL
                )
                """)
            conn.commit()

    async def _load_state(self) -> None:
        """Restore current_main from SQLite, or bootstrap from first node."""

        def _read() -> tuple[str, int] | None:
            with sqlite3.connect(self._db_path) as conn:
                row = conn.execute(
                    "SELECT current_main, epoch FROM cluster_state WHERE id = 1"
                ).fetchone()
                if row is None:
                    return None
                return row[0], row[1]

        state = await anyio.to_thread.run_sync(_read)
        if state is not None:
            current_main, epoch = state
            logger.info("Restored cluster state: main=%s epoch=%d", current_main, epoch)
            self._persisted_main = current_main
            self._epoch = epoch
        else:
            first = next(iter(self._nodes.values())).config.name
            logger.info("Bootstrapping cluster with main=%s", first)
            self._persisted_main = first
            self._epoch = 1
            await self._save_state()

    async def _save_state(self) -> None:
        def _write() -> None:
            with sqlite3.connect(self._db_path) as conn:
                conn.execute(
                    """
                    INSERT INTO cluster_state (id, current_main, epoch, updated_at)
                    VALUES (1, ?, ?, ?)
                    ON CONFLICT(id) DO UPDATE SET
                        current_main = excluded.current_main,
                        epoch = excluded.epoch,
                        updated_at = excluded.updated_at
                    """,
                    (self._persisted_main, self._epoch, time.time()),
                )
                conn.commit()

        await anyio.to_thread.run_sync(_write)

    # ------------------------------------------------------------------
    # Polling and reconciliation
    # ------------------------------------------------------------------

    async def _poll_cycle(self) -> None:
        await self._update_node_health()
        await self._reconcile_roles()

    async def _update_node_health(self) -> None:
        async with anyio.create_task_group() as tg:
            for state in self._nodes.values():
                tg.start_soon(self._check_node, state)

    async def _check_node(self, state: NodeState) -> None:
        try:
            role = await self._query_role(state.config)
            state.role = role
            state.healthy = True
            state.last_seen = time.time()
            state.error = None
        except Exception as exc:
            state.healthy = False
            state.error = str(exc)
            logger.debug("Node %s health check failed: %s", state.config.name, exc)

    async def _reconcile_roles(self) -> None:
        healthy_nodes = [s for s in self._nodes.values() if s.healthy]
        if len(healthy_nodes) < 2:
            logger.critical(
                "Only %d healthy node(s) (%s); role changes skipped (split-brain guard). "
                "The cluster cannot self-heal until at least 2 nodes are reachable.",
                len(healthy_nodes),
                ", ".join(s.config.name for s in healthy_nodes),
            )
            return

        current_main_state = self._nodes.get(self._persisted_main)

        # If the persisted MAIN is healthy and actually MAIN, ensure replicas are registered.
        if (
            current_main_state is not None
            and current_main_state.healthy
            and current_main_state.role == "main"
        ):
            await self._ensure_replicas_registered(current_main_state)
            await self._demote_stray_mains(current_main_state)
            return

        # MAIN is down or demoted. Perform failover.
        await self._perform_failover()

    async def _perform_failover(self) -> None:
        healthy = [s for s in self._nodes.values() if s.healthy]
        if len(healthy) < 2:
            logger.critical(
                "Cannot failover: only %d healthy node(s) (%s). "
                "Restart at least one more node.",
                len(healthy),
                ", ".join(s.config.name for s in healthy),
            )
            return

        # Prefer a node that already reports as replica.
        candidates = [s for s in healthy if s.role == "replica"]
        if not candidates:
            candidates = [s for s in healthy if s.role in ("main", "unknown")]
        if not candidates:
            logger.warning("No failover candidate available")
            return

        new_main = candidates[0]
        logger.warning(
            "Failing over from %s to %s", self._persisted_main, new_main.config.name
        )

        # Promote the candidate to MAIN.
        await self._set_role(new_main.config, "main")
        new_main.role = "main"

        # Update persisted state.
        self._persisted_main = new_main.config.name
        self._epoch += 1
        await self._save_state()

        # Register remaining healthy nodes as replicas of the new MAIN.
        await self._ensure_replicas_registered(new_main)
        await self._demote_stray_mains(new_main)

    async def _ensure_replicas_registered(self, main_state: NodeState) -> None:
        """Idempotently register every healthy non-MAIN node as a replica.

        Non-destructive: a replica already registered AND in a healthy status is
        left ALONE — never DROP+re-REGISTERed. Blindly tearing down a working
        registration every reconcile cycle (the previous behaviour) is exactly how
        the cluster once ended up with zero live replicas — one cycle DROPped the
        working replicas and the following REGISTER failed, stranding them. Each
        replica is handled in isolation so one failure cannot abort the others or
        bubble up and kill the reconcile loop (``_register_replica`` never raises).
        """
        try:
            live = await self._show_replicas(main_state.config)
        except Exception as exc:
            logger.warning(
                "SHOW REPLICAS on %s failed; will (re)register defensively: %s",
                main_state.config.name,
                exc,
            )
            live = {}
        for state in self._nodes.values():
            if state.config.name == main_state.config.name or not state.healthy:
                continue
            if state.role != "replica":
                try:
                    await self._set_role(state.config, "replica")
                    state.role = "replica"
                except Exception as exc:
                    logger.warning(
                        "Setting REPLICA role on %s failed: %s", state.config.name, exc
                    )
                    continue
            data_info = live.get(state.config.name)
            if self._replica_is_healthy(data_info):
                if not self._note_replica_progress(state.config.name, data_info):
                    continue  # registered, healthy, progressing — do not disturb
                logger.warning(
                    "Replica %s reports 'ready' but its replication timestamp "
                    "has been frozen with data outstanding for %d consecutive "
                    "cycles — forcing re-registration",
                    state.config.name,
                    self._stall_cycles_threshold,
                )
            else:
                self._replica_progress.pop(state.config.name, None)
            await self._register_replica(
                main_state.config,
                state.config,
                drop_first=state.config.name in live,
            )

    def _note_replica_progress(
        self, name: str, data_info: dict[str, Any] | None
    ) -> bool:
        """Track per-cycle replication progress; True when frozen-but-ready.

        The post-reboot incident (2026-07-08) left both replicas reporting
        ``status='ready'`` while their replication timestamp sat frozen 131
        transactions behind MAIN — a state ``_replica_is_healthy`` cannot
        see, so reconcile never repaired it. A replica counts as *frozen*
        when, for ``stall_cycles_threshold`` consecutive poll cycles, it is
        'ready' AND has data outstanding (``behind != 0``) AND its ``ts``
        never advances. Genuine catch-up ('recovery' status) is exempt —
        only the pathological ready-but-stuck state triggers.

        Returns True at the threshold and resets the counter so the forced
        re-registration gets a full window before firing again.
        """
        info: dict[str, Any] | None = None
        if data_info:
            info = next((v for v in data_info.values() if isinstance(v, dict)), None)
        if info is None or info.get("status") != "ready":
            self._replica_progress.pop(name, None)
            return False
        behind = info.get("behind")
        ts = info.get("ts")
        if not isinstance(behind, int) or behind == 0 or not isinstance(ts, int):
            self._replica_progress.pop(name, None)
            return False
        last_ts, cycles = self._replica_progress.get(name, (None, 0))
        cycles = cycles + 1 if last_ts == ts else 0
        self._replica_progress[name] = (ts, cycles)
        if cycles >= self._stall_cycles_threshold:
            del self._replica_progress[name]
            return True
        return False

    @staticmethod
    def _replica_is_healthy(data_info: dict[str, Any] | None) -> bool:
        """True iff a SHOW REPLICAS ``data_info`` reports a healthy/syncing status.

        Pure decision function (no I/O). ``data_info`` is the column Memgraph
        returns per replica — ``{"memgraph": {"status": "ready"|"recovery"|...}}``
        — or ``None`` when the replica is not registered at all. A missing entry
        or any non-healthy status means it must be (re)registered.
        """
        if not data_info:
            return False
        for info in data_info.values():
            if not isinstance(info, dict):
                return False
            if info.get("status") not in _HEALTHY_REPLICA_STATUSES:
                return False
        return True

    async def _show_replicas(
        self, main_config: NodeConfig
    ) -> dict[str, dict[str, Any]]:
        """Return ``{replica_name: data_info}`` for replicas registered on the MAIN."""

        def _run() -> dict[str, dict[str, Any]]:
            from orchestrata_ultimate.core.graph_factory import NO_AUTH, create_graph

            driver = create_graph(uri=main_config.bolt_uri, auth=NO_AUTH)
            try:
                with driver.session() as session:
                    out: dict[str, dict[str, Any]] = {}
                    for record in session.run("SHOW REPLICAS"):
                        name = record.get("name")
                        if name:
                            out[str(name)] = record.get("data_info") or {}
                    return out
            finally:
                driver.close()

        return await anyio.to_thread.run_sync(_run)

    async def _demote_stray_mains(self, main_state: NodeState) -> None:
        """Demote any healthy node that thinks it is MAIN but is not the current MAIN."""
        for state in self._nodes.values():
            if state.config.name == main_state.config.name:
                continue
            if state.healthy and state.role == "main":
                logger.warning(
                    "Demoting stray main %s (current main is %s)",
                    state.config.name,
                    main_state.config.name,
                )
                await self._set_role(state.config, "replica")
                state.role = (
                    "replica"  # pragma: no cover - requires a real stray-main demotion
                )
                await self._register_replica(main_state.config, state.config)

    # ------------------------------------------------------------------
    # Bolt helpers
    # ------------------------------------------------------------------

    async def _query_role(self, config: NodeConfig) -> str:
        def _run() -> str:
            from orchestrata_ultimate.core.graph_factory import NO_AUTH, create_graph

            driver = create_graph(uri=config.bolt_uri, auth=NO_AUTH)
            try:
                with driver.session() as session:
                    result = session.run("SHOW REPLICATION ROLE")
                    record = result.single()
                    if record is None:
                        return "unknown"  # pragma: no cover - Memgraph always returns a role
                    return str(record[0]).lower()
            finally:
                driver.close()

        return await anyio.to_thread.run_sync(_run)

    async def _set_role(self, config: NodeConfig, role: str) -> None:
        def _run() -> None:
            from orchestrata_ultimate.core.graph_factory import NO_AUTH, create_graph

            driver = create_graph(uri=config.bolt_uri, auth=NO_AUTH)
            try:
                with driver.session() as session:
                    if role not in ("main", "replica"):
                        raise ValueError(f"Unknown role: {role}")
                    # Memgraph raises on a no-op SET REPLICATION ROLE (e.g.
                    # "TO MAIN" on a node that already IS main), so skip when
                    # the node is already in the target role. Before the
                    # consume() below, that failure was silently discarded and
                    # the tests only passed because of it — a REPLICA set-role
                    # failure would likewise have gone unnoticed in a real
                    # failover, with the controller believing it succeeded.
                    record = session.run("SHOW REPLICATION ROLE").single()
                    current = str(record["replication role"]).lower() if record else ""
                    if current == role:
                        return
                    if role == "main":
                        session.run("SET REPLICATION ROLE TO MAIN").consume()
                    else:
                        session.run(
                            f"SET REPLICATION ROLE TO REPLICA WITH PORT {config.replication_port}"
                        ).consume()
            finally:
                driver.close()

        await anyio.to_thread.run_sync(_run)

    async def _register_replica(
        self,
        main_config: NodeConfig,
        replica_config: NodeConfig,
        *,
        drop_first: bool = True,
    ) -> bool:
        """(Re)register ``replica_config`` on the MAIN. Never raises; returns success.

        ``drop_first`` DROPs any existing registration before re-adding — needed
        when an already-registered replica is in a bad state. For a fresh
        registration it is skipped, avoiding a spurious "failed to unregister"
        error against a replica that was never registered. A failed REGISTER is
        logged and swallowed so it can be retried next reconcile cycle instead of
        killing the loop (the un-wrapped REGISTER here previously did just that).
        """

        def _run() -> None:
            from orchestrata_ultimate.core.graph_factory import NO_AUTH, create_graph

            driver = create_graph(uri=main_config.bolt_uri, auth=NO_AUTH)
            try:
                with driver.session() as session:
                    if drop_first:
                        # gate-ok: absence is the expected outcome of this best-effort pre-clean; a genuinely wrong state still fails loudly at the REGISTER REPLICA below
                        try:
                            # consume() so a server-side failure raises here —
                            # an unconsumed result is discarded on session
                            # close and the except below never fires.
                            session.run(
                                f"DROP REPLICA `{replica_config.name}`"
                            ).consume()
                        except Exception as exc:
                            logger.debug(
                                "DROP REPLICA %s failed (may not exist): %s",
                                replica_config.name,
                                exc,
                            )
                    session.run(
                        f"REGISTER REPLICA `{replica_config.name}` "
                        f"{self._replication_mode} TO "
                        f"'{replica_config.replication_host}:{replica_config.replication_port}'"
                    ).consume()
            finally:
                driver.close()

        try:
            await anyio.to_thread.run_sync(_run)
        except Exception as exc:
            logger.warning(
                "REGISTER REPLICA %s on %s failed (will retry next cycle): %s",
                replica_config.name,
                main_config.name,
                exc,
            )
            return False
        logger.info(
            "Registered replica %s on %s", replica_config.name, main_config.name
        )
        return True

    # ------------------------------------------------------------------
    # Status helpers
    # ------------------------------------------------------------------

    def _current_main(self) -> NodeState | None:
        return self._nodes.get(self._persisted_main)

    def _status_dict(self) -> dict[str, Any]:
        return {
            "current_main": self._persisted_main,
            "epoch": self._epoch,
            "nodes": [
                {
                    "name": s.config.name,
                    "bolt_uri": s.config.bolt_uri,
                    "role": s.role,
                    "healthy": s.healthy,
                    "last_seen": s.last_seen,
                    "error": s.error,
                }
                for s in self._nodes.values()
            ],
        }