stackbox 0.1.0__py3-none-any.whl

stackbox/ledger.py

@@ -0,0 +1,387 @@
+"""Durable record of every sandbox the gateway placed and every provisioning attempt.
+
+Intent is written *before* an upstream create is sent, so a crash or an ambiguous
+timeout never leaves a billable sandbox the gateway has no record of.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+# Sandbox states
+PROVISIONING = "provisioning"  # create sent (or about to be); outcome not yet known
+UNKNOWN = (
+    "unknown"  # create outcome ambiguous; must be reconciled before reuse of the slot
+)
+ACTIVE = "active"
+FAILED = "failed"
+DELETED = "deleted"
+EXPIRED = "expired"  # upstream no longer has it (timeout or external deletion)
+
+HOLDING_STATES = (PROVISIONING, UNKNOWN, ACTIVE)
+
+# Attempt outcomes
+OK = "ok"
+FAILED_CLEAR = "failed"  # upstream definitely did not create a sandbox
+AMBIGUOUS = "ambiguous"  # may or may not exist upstream
+ADOPTED = "adopted"  # ambiguous attempt later found to exist upstream
+ABSENT = "absent"  # ambiguous attempt later confirmed absent
+SATURATED = (
+    "saturated"  # backend refused for quota/capacity; nothing created, not a fault
+)
+DISCARDED = (
+    "discarded"  # ambiguous attempt found upstream but deleted instead of adopted
+)
+CLIENT_ERROR = "client_error"  # request rejected as invalid; not retried elsewhere
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS sandboxes (
+    gw_id TEXT PRIMARY KEY,
+    backend TEXT,
+    upstream_id TEXT,
+    state TEXT NOT NULL,
+    requested_at REAL NOT NULL,
+    ready_at REAL,
+    ended_at REAL,
+    attempts INTEGER NOT NULL DEFAULT 0,
+    last_error TEXT,
+    cpu REAL,
+    memory_gib REAL,
+    pinned_backend TEXT,
+    client_metadata TEXT,
+    image TEXT,
+    backend_data TEXT
+);
+CREATE INDEX IF NOT EXISTS sandboxes_backend_state ON sandboxes(backend, state);
+CREATE INDEX IF NOT EXISTS sandboxes_upstream ON sandboxes(backend, upstream_id);
+CREATE TABLE IF NOT EXISTS attempts (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    gw_id TEXT NOT NULL,
+    backend TEXT NOT NULL,
+    started_at REAL NOT NULL,
+    finished_at REAL,
+    outcome TEXT,
+    http_status INTEGER,
+    error TEXT
+);
+CREATE INDEX IF NOT EXISTS attempts_gw ON attempts(gw_id);
+CREATE TABLE IF NOT EXISTS images (
+    backend TEXT NOT NULL,
+    image TEXT NOT NULL,
+    status TEXT NOT NULL,
+    ref TEXT,
+    error TEXT,
+    detail TEXT,
+    updated_at REAL NOT NULL,
+    PRIMARY KEY (backend, image)
+);
+CREATE TABLE IF NOT EXISTS kv (
+    key TEXT PRIMARY KEY,
+    value TEXT NOT NULL
+);
+"""
+
+# Columns added after the first release; applied to existing databases on open.
+_MIGRATIONS = [
+    "ALTER TABLE sandboxes ADD COLUMN image TEXT",
+    "ALTER TABLE sandboxes ADD COLUMN backend_data TEXT",
+]
+
+# Image preparation states
+IMAGE_BUILDING = "building"
+IMAGE_READY = "ready"
+IMAGE_FAILED = "failed"
+
+
+@dataclass
+class SandboxRecord:
+    gw_id: str
+    backend: str | None
+    upstream_id: str | None
+    state: str
+    requested_at: float
+    ready_at: float | None
+    ended_at: float | None
+    attempts: int
+    last_error: str | None
+    cpu: float | None
+    memory_gib: float | None
+    pinned_backend: str | None
+    client_metadata: dict[str, str]
+    image: str | None = None
+    backend_data: dict[str, Any] | None = None
+
+    @classmethod
+    def from_row(cls, row: sqlite3.Row) -> SandboxRecord:
+        d = dict(row)
+        d["client_metadata"] = json.loads(d["client_metadata"] or "{}")
+        d["backend_data"] = (
+            json.loads(d["backend_data"]) if d.get("backend_data") else None
+        )
+        return cls(**d)
+
+
+@dataclass
+class ImageRecord:
+    backend: str
+    image: str
+    status: str
+    ref: str | None
+    error: str | None
+    detail: dict[str, Any] | None
+    updated_at: float
+
+
+class Ledger:
+    def __init__(self, path: str):
+        if path != ":memory:":
+            Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
+            path = str(Path(path).expanduser())
+        self._conn = sqlite3.connect(
+            path, check_same_thread=False, isolation_level=None
+        )
+        self._conn.row_factory = sqlite3.Row
+        self._conn.execute("PRAGMA journal_mode=WAL")
+        self._conn.executescript(_SCHEMA)
+        for stmt in _MIGRATIONS:
+            try:
+                self._conn.execute(stmt)
+            except sqlite3.OperationalError:
+                pass  # column already exists
+        self._lock = threading.Lock()
+
+    def _exec(self, sql: str, params: tuple[Any, ...] = ()) -> sqlite3.Cursor:
+        with self._lock:
+            return self._conn.execute(sql, params)
+
+    # --- sandboxes -------------------------------------------------------
+    def create_intent(
+        self,
+        gw_id: str,
+        *,
+        cpu: float | None,
+        memory_gib: float | None,
+        pinned_backend: str | None,
+        client_metadata: dict[str, str],
+        image: str | None = None,
+    ) -> None:
+        self._exec(
+            "INSERT INTO sandboxes (gw_id, state, requested_at, cpu, memory_gib, pinned_backend,"
+            " client_metadata, image) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+            (
+                gw_id,
+                PROVISIONING,
+                time.time(),
+                cpu,
+                memory_gib,
+                pinned_backend,
+                json.dumps(client_metadata),
+                image,
+            ),
+        )
+
+    def claim_backend(self, gw_id: str, backend: str, max_active: int) -> bool:
+        """Atomically assign gw_id to backend if the backend has a free slot."""
+        with self._lock:
+            self._conn.execute("BEGIN IMMEDIATE")
+            try:
+                (held,) = self._conn.execute(
+                    f"SELECT COUNT(*) FROM sandboxes WHERE backend = ? AND state IN ({','.join('?' * len(HOLDING_STATES))})",
+                    (backend, *HOLDING_STATES),
+                ).fetchone()
+                if held >= max_active:
+                    self._conn.execute("ROLLBACK")
+                    return False
+                self._conn.execute(
+                    "UPDATE sandboxes SET backend = ?, state = ?, attempts = attempts + 1 WHERE gw_id = ?",
+                    (backend, PROVISIONING, gw_id),
+                )
+                self._conn.execute("COMMIT")
+                return True
+            except BaseException:
+                self._conn.execute("ROLLBACK")
+                raise
+
+    def release_backend(self, gw_id: str, error: str | None) -> None:
+        """Provisioning on the current backend definitely failed; free the slot for retry elsewhere."""
+        self._exec(
+            "UPDATE sandboxes SET backend = NULL, state = ?, last_error = ? WHERE gw_id = ?",
+            (PROVISIONING, error, gw_id),
+        )
+
+    def mark_active(
+        self,
+        gw_id: str,
+        upstream_id: str,
+        *,
+        cpu: float | None = None,
+        memory_gib: float | None = None,
+        backend_data: dict[str, Any] | None = None,
+    ) -> None:
+        """Record the sandbox as running. cpu/memory, when given, are what the backend
+        actually allocated (they replace the requested values for cost accounting)."""
+        self._exec(
+            "UPDATE sandboxes SET upstream_id = ?, state = ?, ready_at = ?,"
+            " cpu = COALESCE(?, cpu), memory_gib = COALESCE(?, memory_gib),"
+            " backend_data = COALESCE(?, backend_data) WHERE gw_id = ?",
+            (
+                upstream_id,
+                ACTIVE,
+                time.time(),
+                cpu,
+                memory_gib,
+                json.dumps(backend_data) if backend_data is not None else None,
+                gw_id,
+            ),
+        )
+
+    def set_state(self, gw_id: str, state: str, error: str | None = None) -> None:
+        ended = time.time() if state in (FAILED, DELETED, EXPIRED) else None
+        self._exec(
+            "UPDATE sandboxes SET state = ?, last_error = COALESCE(?, last_error),"
+            " ended_at = COALESCE(?, ended_at) WHERE gw_id = ?",
+            (state, error, ended, gw_id),
+        )
+
+    def get(self, gw_id: str) -> SandboxRecord | None:
+        row = self._exec("SELECT * FROM sandboxes WHERE gw_id = ?", (gw_id,)).fetchone()
+        return SandboxRecord.from_row(row) if row else None
+
+    def by_upstream(self, backend: str, upstream_id: str) -> SandboxRecord | None:
+        row = self._exec(
+            "SELECT * FROM sandboxes WHERE backend = ? AND upstream_id = ?",
+            (backend, upstream_id),
+        ).fetchone()
+        return SandboxRecord.from_row(row) if row else None
+
+    def in_states(self, *states: str) -> list[SandboxRecord]:
+        rows = self._exec(
+            f"SELECT * FROM sandboxes WHERE state IN ({','.join('?' * len(states))})",
+            states,
+        ).fetchall()
+        return [SandboxRecord.from_row(r) for r in rows]
+
+    def held_count(self, backend: str) -> int:
+        (n,) = self._exec(
+            f"SELECT COUNT(*) FROM sandboxes WHERE backend = ? AND state IN ({','.join('?' * len(HOLDING_STATES))})",
+            (backend, *HOLDING_STATES),
+        ).fetchone()
+        return n
+
+    def active_count(self, backend: str) -> int:
+        (n,) = self._exec(
+            "SELECT COUNT(*) FROM sandboxes WHERE backend = ? AND state = ?",
+            (backend, ACTIVE),
+        ).fetchone()
+        return n
+
+    def all(self) -> list[SandboxRecord]:
+        return [
+            SandboxRecord.from_row(r)
+            for r in self._exec("SELECT * FROM sandboxes").fetchall()
+        ]
+
+    # --- attempts --------------------------------------------------------
+    def start_attempt(self, gw_id: str, backend: str) -> int:
+        cur = self._exec(
+            "INSERT INTO attempts (gw_id, backend, started_at) VALUES (?, ?, ?)",
+            (gw_id, backend, time.time()),
+        )
+        return int(cur.lastrowid or 0)
+
+    def finish_attempt(
+        self,
+        attempt_id: int,
+        outcome: str,
+        http_status: int | None = None,
+        error: str | None = None,
+    ) -> None:
+        self._exec(
+            "UPDATE attempts SET finished_at = ?, outcome = ?, http_status = ?, error = ? WHERE id = ?",
+            (time.time(), outcome, http_status, error, attempt_id),
+        )
+
+    def resolve_ambiguous_attempt(self, gw_id: str, backend: str, outcome: str) -> None:
+        self._exec(
+            "UPDATE attempts SET outcome = ? WHERE gw_id = ? AND backend = ? AND outcome = ?",
+            (outcome, gw_id, backend, AMBIGUOUS),
+        )
+
+    def attempts(self) -> list[dict[str, Any]]:
+        return [
+            dict(r) for r in self._exec("SELECT * FROM attempts ORDER BY id").fetchall()
+        ]
+
+    def startup_samples(self) -> list[tuple[str, str | None, float]]:
+        """(backend, image, seconds) for every successful create, oldest first."""
+        rows = self._exec(
+            "SELECT a.backend, s.image, a.finished_at - a.started_at FROM attempts a"
+            " JOIN sandboxes s ON s.gw_id = a.gw_id"
+            " WHERE a.outcome = ? AND a.finished_at IS NOT NULL ORDER BY a.finished_at",
+            (OK,),
+        ).fetchall()
+        return [(r[0], r[1], float(r[2])) for r in rows]
+
+    # --- prepared images ---------------------------------------------------
+    def set_image(
+        self,
+        backend: str,
+        image: str,
+        status: str,
+        *,
+        ref: str | None = None,
+        error: str | None = None,
+        detail: dict[str, Any] | None = None,
+    ) -> None:
+        self._exec(
+            "INSERT INTO images (backend, image, status, ref, error, detail, updated_at)"
+            " VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT (backend, image) DO UPDATE SET"
+            " status = excluded.status, ref = COALESCE(excluded.ref, images.ref),"
+            " error = excluded.error, detail = COALESCE(excluded.detail, images.detail),"
+            " updated_at = excluded.updated_at",
+            (
+                backend,
+                image,
+                status,
+                ref,
+                error,
+                json.dumps(detail) if detail is not None else None,
+                time.time(),
+            ),
+        )
+
+    def get_image(self, backend: str, image: str) -> ImageRecord | None:
+        row = self._exec(
+            "SELECT * FROM images WHERE backend = ? AND image = ?", (backend, image)
+        ).fetchone()
+        if row is None:
+            return None
+        d = dict(row)
+        d["detail"] = json.loads(d["detail"]) if d["detail"] else None
+        return ImageRecord(**d)
+
+    def images(self) -> list[dict[str, Any]]:
+        return [dict(r) for r in self._exec("SELECT * FROM images").fetchall()]
+
+    # --- key/value ---------------------------------------------------------
+    def get_or_create_secret(self, key: str, factory: Any) -> str:
+        with self._lock:
+            row = self._conn.execute(
+                "SELECT value FROM kv WHERE key = ?", (key,)
+            ).fetchone()
+            if row:
+                return str(row[0])
+            value = factory()
+            self._conn.execute(
+                "INSERT OR IGNORE INTO kv (key, value) VALUES (?, ?)", (key, value)
+            )
+            (value,) = self._conn.execute(
+                "SELECT value FROM kv WHERE key = ?", (key,)
+            ).fetchone()
+            return str(value)