From 5cdeb8774a9bc0268b1998898c1eda3b507eaab7 Mon Sep 17 00:00:00 2001 From: Deathcharge Date: Sun, 2 Aug 2026 07:56:00 -0400 Subject: [PATCH 1/2] feat(postgres): enforce distributed connection capacity --- CHANGELOG.md | 2 + ROADMAP.md | 3 +- docs/MULTI_INSTANCE_ARCHITECTURE.md | 6 +- samsarix_chat_engine/postgres.py | 63 +++- samsarix_chat_engine/postgres_connections.py | 316 +++++++++++++++++++ samsarix_chat_engine/postgres_store.py | 4 +- tests/conftest.py | 1 + tests/test_postgres_connections.py | 229 ++++++++++++++ tests/test_postgres_store.py | 2 +- 9 files changed, 619 insertions(+), 7 deletions(-) create mode 100644 samsarix_chat_engine/postgres_connections.py create mode 100644 tests/test_postgres_connections.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ff90f30..a9003be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ This project follows semantic versioning while it is in alpha: minor versions ma - Internal PostgreSQL schema v3 and a complete `ChatStorage` implementation for rooms, bounded messages, Unicode-normalized search, moderation, audit history, monotonic read state, stable spooled exports, explicit retention, and cross-instance idempotency/capacity enforcement. - Transactional PostgreSQL webhook outbox with stable delivery IDs, database-time scheduling, expiring worker-owned claims, `SKIP LOCKED` work selection, crash recovery, bounded terminal-history pruning, and operator replay. - Internal per-process PostgreSQL realtime relay with durable cursors, ordered polling/replay, post-dispatch acknowledgement, archive/ban socket teardown, and lease-loss fencing. Polling is the correctness path; a future `LISTEN`/`NOTIFY` listener may only reduce latency. +- PostgreSQL schema v4 and an internal connection registry with database-time socket leases, atomic deployment-wide and per-room capacity, owner-bound renewal/release, archived-room rejection, and crashed-process reclamation. ### Security and operations @@ -22,6 +23,7 @@ This project follows semantic versioning while it is in alpha: minor versions ma - PostgreSQL event and webhook envelopes remain bounded at 512 KiB so every valid 100,000-character message, including four-byte Unicode, fits without turning a valid domain write into a coordination failure. - PostgreSQL webhook claims can be acknowledged only by the live lease owner. Expired claims are safely redelivered with the same ID; receivers must still deduplicate because delivery is at least once. - A PostgreSQL relay that loses its database lease closes every local socket before renewing the same durable cursor. Unsupported internal event types are not forwarded onto the public WebSocket protocol. +- Expired process IDs discard their stale socket rows before re-registration, preventing a restarted replica from reviving phantom occupancy. Archived and expired reservations stop consuming capacity even before physical cleanup. ## 0.12.0 — 2026-08-02 diff --git a/ROADMAP.md b/ROADMAP.md index 097b576..61758d3 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -85,8 +85,9 @@ This makes the supported topology repeatable without implying that a container m - [x] define a storage-neutral application boundary without changing the single-instance default; - [x] implement the internal PostgreSQL authoritative store, ordered event log, read state, stable exports, retention, and leased webhook outbox; - [x] implement a cursor-backed per-process realtime relay with ordered replay and lease-loss socket fencing; +- [x] implement PostgreSQL-owned expiring connection leases with atomic global/per-room caps and crash reclamation; - [ ] expose guarded PostgreSQL configuration only after the remaining topology gates are proven; -- [ ] implement and test cross-worker fan-out, presence, distributed rate limits, and reconnect recovery; +- [ ] implement and test presence, distributed rate limits, and reconnect recovery, then wire the proven cross-worker fan-out path into the application; - [ ] run sustained load/soak and reconnect-storm tests and publish measured limits; - [ ] add OpenTelemetry hooks only when an operator needs them, with telemetry disabled by default. diff --git a/docs/MULTI_INSTANCE_ARCHITECTURE.md b/docs/MULTI_INSTANCE_ARCHITECTURE.md index eecd18d..22d7932 100644 --- a/docs/MULTI_INSTANCE_ARCHITECTURE.md +++ b/docs/MULTI_INSTANCE_ARCHITECTURE.md @@ -16,6 +16,8 @@ The initial event-log implementation serializes sequence allocation with a trans Schema v3 uses PostgreSQL database time for room/message/moderation ordering, read cursors, webhook due times and leases, and retention boundaries. The internal store now implements the full storage protocol, including monotonic subject-scoped read state, transactionally stable bounded-memory exports, explicit retention, and a leased transactional webhook outbox. Transaction-scoped capacity locks currently serialize room creation, message mutation/retention, bounded audit insertion, and webhook capacity changes across replicas. Deletion and retention cancel unsent sensitive webhook payloads and scrub message bodies from older durable event and terminal-webhook envelopes before commit. These conservative global locks make correctness inspectable first; the load gate must measure their throughput before v0.13 receives a scale claim. +Schema v4 adds one database-owned lease per admitted WebSocket. A transaction locks the live owning instance and usable room, serializes the capacity decision, removes stale reservations, and atomically checks deployment-wide and per-room caps before inserting. Renewal and release require the owning instance ID. Counts exclude expired socket leases, expired owners, and archived rooms; cleanup returns those stale rows for future presence convergence. Re-registering an expired stable instance ID deletes its former socket rows before renewing the process lease, so a crash/restart cannot revive phantom occupancy. All replicas must use identical capacity settings; public application admission and presence wiring remain release gates. + No Redis dependency is planned for the first supported topology. Redis Pub/Sub is at-most-once, while Streams introduce a second durable system whose commit cannot be atomic with the authoritative database without an additional outbox relay. PostgreSQL already supplies transactions, row locks, advisory locks, `SKIP LOCKED`, and commit-coupled notifications needed by this product's current scale boundary. ## Why multi-process SQLite is rejected @@ -43,7 +45,7 @@ An instance that falls behind the retained event window must close its local soc PostgreSQL-backed deployments require all of the following: - **Migrations:** one transaction-scoped advisory lock serializes schema inspection and migration. A newer unsupported schema fails closed. -- **Connection capacity:** expiring per-socket leases enforce deployment-wide and per-room caps. Heartbeats renew leases; crashed-instance leases expire. +- **Connection capacity:** implemented schema-v4 per-socket leases enforce deployment-wide and per-room caps with database-time expiry and owner-bound renewal/release. Heartbeats will renew leases when application wiring lands; crashed-instance and archived-room rows are already excluded and reclaimable. - **Presence:** join/leave transitions derive from connection leases. A leader-elected sweeper emits bounded expiry transitions for crashed instances. Presence remains best effort and carries no authorization meaning. - **Rate limits:** atomic time-bucket counters enforce deployment-wide subject/client limits. Database time defines bucket boundaries so host clock skew cannot multiply quotas. - **Typing:** transition events use the durable event path but expire automatically and are not retained as chat history or audit content. @@ -71,7 +73,7 @@ v0.13 cannot claim multi-instance support until CI proves: - two or more real app processes share one PostgreSQL database and deliver create/update/delete events exactly once to each connected test socket under normal operation; - a listener disconnect/reconnect replays committed event rows without relying on `NOTIFY` delivery (the internal polling relay and cursor replay case are implemented; the real-process listener gate remains); - concurrent idempotent message creation returns one authoritative message; -- global and per-room connection caps plus rate limits hold across processes; +- global and per-room connection caps plus rate limits hold across processes (the storage-level concurrent connection-cap case is implemented; real-process admission and rate-limit gates remain); - archive and ban actions close matching sockets on every process; - two webhook workers never hold the same live claim, and a killed worker's claim is recovered (the storage-level lease/recovery case is implemented; the killed-process gate remains); - migration concurrency is serialized and newer schemas fail closed; diff --git a/samsarix_chat_engine/postgres.py b/samsarix_chat_engine/postgres.py index 05f212a..c50e931 100644 --- a/samsarix_chat_engine/postgres.py +++ b/samsarix_chat_engine/postgres.py @@ -23,7 +23,7 @@ from psycopg.types.json import Jsonb from psycopg_pool import AsyncConnectionPool, PoolTimeout -POSTGRES_SCHEMA_VERSION = 3 +POSTGRES_SCHEMA_VERSION = 4 POSTGRES_MIGRATION_LOCK_ID = 7_495_346_927_831_819_041 POSTGRES_EVENT_SEQUENCE_LOCK_ID = 7_495_346_927_831_819_042 REALTIME_CHANNEL = "samsarix_realtime_v1" @@ -198,6 +198,24 @@ async def register_instance(self, instance_id: str, *, lease_seconds: int) -> in _validate_instance(instance_id, lease_seconds) async with self.transaction() as connection: + cursor = await connection.execute( + """ + SELECT lease_expires_at <= clock_timestamp() + FROM public.samsarix_instance_cursors + WHERE instance_id = %s + FOR UPDATE + """, + (instance_id,), + ) + existing = await cursor.fetchone() + if existing is not None and bool(existing[0]): + # A stable instance ID may be reused after a crash. Its old + # sockets no longer exist and must not become live again when + # the owner lease is renewed. + await connection.execute( + "DELETE FROM public.samsarix_connection_leases WHERE instance_id = %s", + (instance_id,), + ) cursor = await connection.execute( """ INSERT INTO public.samsarix_instance_cursors ( @@ -411,6 +429,49 @@ async def _initialize_schema(self) -> None: ) """ ) + await connection.execute( + """ + CREATE TABLE IF NOT EXISTS public.samsarix_connection_leases ( + connection_id TEXT PRIMARY KEY CHECK ( + char_length(connection_id) BETWEEN 1 AND 128 + AND connection_id ~ '^[A-Za-z0-9][A-Za-z0-9._:-]*$' + ), + instance_id TEXT NOT NULL REFERENCES public.samsarix_instance_cursors(instance_id) + ON DELETE CASCADE, + room_id TEXT NOT NULL REFERENCES public.samsarix_rooms(id) ON DELETE CASCADE, + username TEXT NOT NULL CHECK (char_length(username) BETWEEN 1 AND 64), + subject TEXT CHECK (subject IS NULL OR char_length(subject) BETWEEN 1 AND 64), + lease_expires_at TIMESTAMPTZ NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT clock_timestamp(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT clock_timestamp(), + CHECK (lease_expires_at > created_at) + ) + """ + ) + await connection.execute( + """ + CREATE INDEX IF NOT EXISTS samsarix_connection_leases_expiry + ON public.samsarix_connection_leases (lease_expires_at, connection_id) + """ + ) + await connection.execute( + """ + CREATE INDEX IF NOT EXISTS samsarix_connection_leases_room + ON public.samsarix_connection_leases (room_id, lease_expires_at, connection_id) + """ + ) + await connection.execute( + """ + CREATE INDEX IF NOT EXISTS samsarix_connection_leases_instance + ON public.samsarix_connection_leases (instance_id, lease_expires_at, connection_id) + """ + ) + await connection.execute( + """ + CREATE INDEX IF NOT EXISTS samsarix_connection_leases_member + ON public.samsarix_connection_leases (room_id, subject, lease_expires_at) + """ + ) await connection.execute( """ CREATE TABLE IF NOT EXISTS public.samsarix_messages ( diff --git a/samsarix_chat_engine/postgres_connections.py b/samsarix_chat_engine/postgres_connections.py new file mode 100644 index 0000000..dda5ed8 --- /dev/null +++ b/samsarix_chat_engine/postgres_connections.py @@ -0,0 +1,316 @@ +# Copyright (c) 2026 Samsarix LLC +# SPDX-License-Identifier: MPL-2.0 +"""PostgreSQL-owned expiring WebSocket connection leases.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from typing import Any, cast + +from psycopg import AsyncConnection +from psycopg.errors import UniqueViolation + +from .postgres import ( + _ROOM_ID_PATTERN, + InstanceLeaseError, + PostgresFoundation, + _validate_instance_id, +) + +POSTGRES_CONNECTION_CAP_LOCK_ID = 7_495_346_927_831_819_047 +_CONNECTION_ID_MAX_CHARS = 128 +_USERNAME_MAX_CHARS = 64 +_SUBJECT_MAX_CHARS = 64 + + +class ConnectionLeaseError(RuntimeError): + """Raised when a connection lease is invalid, missing, expired, or owned elsewhere.""" + + +class ConnectionRoomUnavailableError(ConnectionLeaseError): + """Raised when a connection's room is missing or archived.""" + + +@dataclass(frozen=True, slots=True) +class ConnectionLease: + """One process-owned, expiring connection reservation.""" + + connection_id: str + instance_id: str + room_id: str + username: str + subject: str | None + lease_expires_at: datetime + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class ConnectionCounts: + """Live global and selected-room occupancy at one database instant.""" + + total: int + room: int + + +class PostgresConnectionRegistry: + """Atomically reserve global and per-room connection capacity across processes.""" + + def __init__( + self, + foundation: PostgresFoundation, + *, + max_connections: int, + max_connections_per_room: int, + lease_seconds: int = 30, + ) -> None: + if max_connections < 1: + raise ValueError("PostgreSQL connection capacity must be positive") + if max_connections_per_room < 1: + raise ValueError("PostgreSQL per-room connection capacity must be positive") + if max_connections_per_room > max_connections: + raise ValueError("per-room connection capacity cannot exceed global capacity") + if not 3 <= lease_seconds <= 300: + raise ValueError("PostgreSQL connection lease must be between 3 and 300 seconds") + self.foundation = foundation + self.max_connections = max_connections + self.max_connections_per_room = max_connections_per_room + self.lease_seconds = lease_seconds + + async def try_acquire( + self, + *, + connection_id: str, + instance_id: str, + room_id: str, + username: str, + subject: str | None, + ) -> ConnectionLease | None: + """Reserve capacity, returning ``None`` when either configured cap is full.""" + + _validate_connection(connection_id, instance_id, room_id, username, subject) + async with self.foundation.transaction() as connection: + cursor = await connection.execute( + """ + SELECT 1 + FROM public.samsarix_instance_cursors + WHERE instance_id = %s AND lease_expires_at > clock_timestamp() + FOR SHARE + """, + (instance_id,), + ) + if await cursor.fetchone() is None: + raise InstanceLeaseError("instance lease is missing or expired") + + cursor = await connection.execute( + """ + SELECT archived_at + FROM public.samsarix_rooms + WHERE id = %s + FOR SHARE + """, + (room_id,), + ) + room = await cursor.fetchone() + if room is None or room[0] is not None: + raise ConnectionRoomUnavailableError("connection room is missing or archived") + + await connection.execute("SELECT pg_advisory_xact_lock(%s)", (POSTGRES_CONNECTION_CAP_LOCK_ID,)) + await _delete_expired(connection) + + cursor = await connection.execute( + """ + SELECT + COUNT(*)::BIGINT, + COUNT(*) FILTER (WHERE lease.room_id = %s)::BIGINT + FROM public.samsarix_connection_leases AS lease + JOIN public.samsarix_instance_cursors AS owner + ON owner.instance_id = lease.instance_id + JOIN public.samsarix_rooms AS room + ON room.id = lease.room_id + WHERE lease.lease_expires_at > clock_timestamp() + AND owner.lease_expires_at > clock_timestamp() + AND room.archived_at IS NULL + """, + (room_id,), + ) + counts = await cursor.fetchone() + total = int(counts[0]) if counts is not None else 0 + room_total = int(counts[1]) if counts is not None else 0 + if total >= self.max_connections or room_total >= self.max_connections_per_room: + return None + + try: + cursor = await connection.execute( + """ + INSERT INTO public.samsarix_connection_leases ( + connection_id, instance_id, room_id, username, subject, lease_expires_at + ) + SELECT + %s, %s, %s, %s, %s, + clock_timestamp() + make_interval(secs => %s) + WHERE EXISTS ( + SELECT 1 + FROM public.samsarix_instance_cursors + WHERE instance_id = %s + AND lease_expires_at > clock_timestamp() + ) + RETURNING + connection_id, instance_id, room_id, username, subject, + lease_expires_at, created_at + """, + ( + connection_id, + instance_id, + room_id, + username, + subject, + self.lease_seconds, + instance_id, + ), + ) + except UniqueViolation: + raise ConnectionLeaseError("connection ID is already leased") from None + row = await cursor.fetchone() + if row is None: + raise InstanceLeaseError("instance lease expired before connection reservation") + return _lease_from_row(row) + + async def renew(self, *, connection_id: str, instance_id: str) -> datetime: + """Extend one live connection only while its owner and room remain usable.""" + + _validate_connection_id(connection_id) + _validate_instance_id(instance_id) + async with self.foundation.transaction() as connection: + cursor = await connection.execute( + """ + UPDATE public.samsarix_connection_leases AS lease + SET lease_expires_at = clock_timestamp() + make_interval(secs => %s), + updated_at = clock_timestamp() + FROM public.samsarix_instance_cursors AS owner, + public.samsarix_rooms AS room + WHERE lease.connection_id = %s + AND lease.instance_id = %s + AND lease.lease_expires_at > clock_timestamp() + AND owner.instance_id = lease.instance_id + AND owner.lease_expires_at > clock_timestamp() + AND room.id = lease.room_id + AND room.archived_at IS NULL + RETURNING lease.lease_expires_at + """, + (self.lease_seconds, connection_id, instance_id), + ) + row = await cursor.fetchone() + if row is None: + raise ConnectionLeaseError("connection lease is missing, expired, or unavailable") + return cast(datetime, row[0]) + + async def release(self, *, connection_id: str, instance_id: str) -> bool: + """Release an owned connection reservation; repeated release is harmless.""" + + _validate_connection_id(connection_id) + _validate_instance_id(instance_id) + async with self.foundation.transaction() as connection: + cursor = await connection.execute( + """ + DELETE FROM public.samsarix_connection_leases + WHERE connection_id = %s AND instance_id = %s + """, + (connection_id, instance_id), + ) + return cursor.rowcount == 1 + + async def counts(self, *, room_id: str) -> ConnectionCounts: + """Return live occupancy without allowing stale rows to inflate it.""" + + if not _ROOM_ID_PATTERN.fullmatch(room_id): + raise ValueError("invalid room ID") + async with self.foundation.transaction() as connection: + cursor = await connection.execute( + """ + SELECT + COUNT(*)::BIGINT, + COUNT(*) FILTER (WHERE lease.room_id = %s)::BIGINT + FROM public.samsarix_connection_leases AS lease + JOIN public.samsarix_instance_cursors AS owner + ON owner.instance_id = lease.instance_id + JOIN public.samsarix_rooms AS room + ON room.id = lease.room_id + WHERE lease.lease_expires_at > clock_timestamp() + AND owner.lease_expires_at > clock_timestamp() + AND room.archived_at IS NULL + """, + (room_id,), + ) + row = await cursor.fetchone() + return ConnectionCounts(total=int(row[0]), room=int(row[1])) if row is not None else ConnectionCounts(0, 0) + + async def reap_expired(self) -> list[ConnectionLease]: + """Delete and return rows whose socket/owner expired or whose room was archived.""" + + async with self.foundation.transaction() as connection: + await connection.execute("SELECT pg_advisory_xact_lock(%s)", (POSTGRES_CONNECTION_CAP_LOCK_ID,)) + rows = await _delete_expired(connection) + return [_lease_from_row(row) for row in rows] + + +async def _delete_expired(connection: AsyncConnection[tuple[Any, ...]]) -> list[tuple[Any, ...]]: + cursor = await connection.execute( + """ + DELETE FROM public.samsarix_connection_leases AS lease + WHERE lease.lease_expires_at <= clock_timestamp() + OR EXISTS ( + SELECT 1 + FROM public.samsarix_instance_cursors AS owner + WHERE owner.instance_id = lease.instance_id + AND owner.lease_expires_at <= clock_timestamp() + ) + OR EXISTS ( + SELECT 1 + FROM public.samsarix_rooms AS room + WHERE room.id = lease.room_id + AND room.archived_at IS NOT NULL + ) + RETURNING + lease.connection_id, lease.instance_id, lease.room_id, lease.username, + lease.subject, lease.lease_expires_at, lease.created_at + """ + ) + return list(await cursor.fetchall()) + + +def _lease_from_row(row: tuple[Any, ...]) -> ConnectionLease: + return ConnectionLease( + connection_id=str(row[0]), + instance_id=str(row[1]), + room_id=str(row[2]), + username=str(row[3]), + subject=None if row[4] is None else str(row[4]), + lease_expires_at=cast(datetime, row[5]), + created_at=cast(datetime, row[6]), + ) + + +def _validate_connection( + connection_id: str, + instance_id: str, + room_id: str, + username: str, + subject: str | None, +) -> None: + _validate_connection_id(connection_id) + _validate_instance_id(instance_id) + if not _ROOM_ID_PATTERN.fullmatch(room_id): + raise ValueError("invalid room ID") + if not 1 <= len(username) <= _USERNAME_MAX_CHARS: + raise ValueError("username must be between 1 and 64 characters") + if subject is not None and not 1 <= len(subject) <= _SUBJECT_MAX_CHARS: + raise ValueError("subject must be between 1 and 64 characters") + + +def _validate_connection_id(connection_id: str) -> None: + if not 1 <= len(connection_id) <= _CONNECTION_ID_MAX_CHARS or not connection_id[0].isalnum(): + raise ValueError("invalid connection ID") + allowed = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._:-" + if any(character not in allowed for character in connection_id): + raise ValueError("invalid connection ID") diff --git a/samsarix_chat_engine/postgres_store.py b/samsarix_chat_engine/postgres_store.py index 51e3573..05f4792 100644 --- a/samsarix_chat_engine/postgres_store.py +++ b/samsarix_chat_engine/postgres_store.py @@ -26,7 +26,7 @@ RoomCreate, WebhookDelivery, ) -from .postgres import PostgresFoundation, PostgresFoundationError +from .postgres import POSTGRES_SCHEMA_VERSION, PostgresFoundation, PostgresFoundationError from .store import ( ChatStorage, InvalidAuditCursorError, @@ -148,7 +148,7 @@ async def close(self) -> None: async def check_ready(self) -> bool: try: - return await self.foundation.schema_version() >= 3 + return await self.foundation.schema_version() >= POSTGRES_SCHEMA_VERSION except PostgresFoundationError: return False diff --git a/tests/conftest.py b/tests/conftest.py index 594a1a2..41c2331 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -67,6 +67,7 @@ async def _reset_postgres_test_database(conninfo: str) -> None: row = await cursor.fetchone() if row is None or row[0] != "samsarix_test": raise RuntimeError("live PostgreSQL tests require the dedicated samsarix_test database") + await connection.execute("DROP TABLE IF EXISTS public.samsarix_connection_leases") await connection.execute("DROP TABLE IF EXISTS public.samsarix_instance_cursors") await connection.execute("DROP TABLE IF EXISTS public.samsarix_realtime_events") await connection.execute("DROP TABLE IF EXISTS public.samsarix_room_read_states") diff --git a/tests/test_postgres_connections.py b/tests/test_postgres_connections.py new file mode 100644 index 0000000..30f64f4 --- /dev/null +++ b/tests/test_postgres_connections.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026 Samsarix LLC +# SPDX-License-Identifier: MPL-2.0 +"""Live tests for PostgreSQL-owned, crash-reclaimable connection capacity.""" + +from __future__ import annotations + +import asyncio + +import pytest + +pytest.importorskip("psycopg") + +from samsarix_chat_engine.models import RoomCreate # noqa: E402 +from samsarix_chat_engine.postgres import InstanceLeaseError, PostgresFoundation # noqa: E402 +from samsarix_chat_engine.postgres_connections import ( # noqa: E402 + ConnectionCounts, + ConnectionLeaseError, + ConnectionRoomUnavailableError, + PostgresConnectionRegistry, +) +from samsarix_chat_engine.postgres_store import PostgresChatStore # noqa: E402 + +pytestmark = pytest.mark.postgres + + +def _store(conninfo: str) -> PostgresChatStore: + return PostgresChatStore( + conninfo, + max_rooms=10, + max_stored_messages=20, + max_stored_messages_per_room=10, + ) + + +@pytest.mark.asyncio +async def test_concurrent_registries_enforce_exact_global_capacity(clean_postgres_database: str) -> None: + first = _store(clean_postgres_database) + second_foundation = PostgresFoundation(clean_postgres_database) + await first.initialize() + await second_foundation.open() + try: + await first.create_room(RoomCreate(id="alpha", name="Alpha")) + await first.create_room(RoomCreate(id="beta", name="Beta")) + await first.foundation.register_instance("node-a", lease_seconds=30) + await second_foundation.register_instance("node-b", lease_seconds=30) + first_registry = PostgresConnectionRegistry( + first.foundation, + max_connections=4, + max_connections_per_room=3, + ) + second_registry = PostgresConnectionRegistry( + second_foundation, + max_connections=4, + max_connections_per_room=3, + ) + + acquisitions = await asyncio.gather( + *( + (first_registry if index % 2 == 0 else second_registry).try_acquire( + connection_id=f"socket-{index}", + instance_id="node-a" if index % 2 == 0 else "node-b", + room_id="alpha" if index % 3 else "beta", + username=f"user-{index}", + subject=f"subject-{index}", + ) + for index in range(12) + ) + ) + + assert sum(lease is not None for lease in acquisitions) == 4 + assert await first_registry.counts(room_id="alpha") == await second_registry.counts(room_id="alpha") + assert (await first_registry.counts(room_id="alpha")).total == 4 + finally: + await asyncio.gather(first.close(), second_foundation.close()) + + +@pytest.mark.asyncio +async def test_concurrent_registries_enforce_exact_per_room_capacity(clean_postgres_database: str) -> None: + store = _store(clean_postgres_database) + await store.initialize() + try: + await store.create_room(RoomCreate(id="general", name="General")) + await store.foundation.register_instance("node-a", lease_seconds=30) + registry = PostgresConnectionRegistry( + store.foundation, + max_connections=10, + max_connections_per_room=2, + ) + + acquisitions = await asyncio.gather( + *( + registry.try_acquire( + connection_id=f"room-socket-{index}", + instance_id="node-a", + room_id="general", + username=f"user-{index}", + subject=None, + ) + for index in range(8) + ) + ) + + assert sum(lease is not None for lease in acquisitions) == 2 + assert (await registry.counts(room_id="general")).room == 2 + finally: + await store.close() + + +@pytest.mark.asyncio +async def test_renew_release_and_room_lifecycle_fail_closed(clean_postgres_database: str) -> None: + store = _store(clean_postgres_database) + await store.initialize() + try: + await store.create_room(RoomCreate(id="general", name="General")) + await store.foundation.register_instance("node-a", lease_seconds=30) + await store.foundation.register_instance("node-b", lease_seconds=30) + registry = PostgresConnectionRegistry( + store.foundation, + max_connections=5, + max_connections_per_room=5, + ) + lease = await registry.try_acquire( + connection_id="socket-one", + instance_id="node-a", + room_id="general", + username="Andrew", + subject="user-1", + ) + assert lease is not None + assert await registry.release(connection_id="socket-one", instance_id="node-b") is False + assert await registry.renew(connection_id="socket-one", instance_id="node-a") > lease.lease_expires_at + + await store.set_room_state("general", archived=True, frozen=None, actor="operator") + with pytest.raises(ConnectionLeaseError, match="unavailable"): + await registry.renew(connection_id="socket-one", instance_id="node-a") + assert await registry.counts(room_id="general") == ConnectionCounts(0, 0) + assert [item.connection_id for item in await registry.reap_expired()] == ["socket-one"] + with pytest.raises(ConnectionRoomUnavailableError, match="archived"): + await registry.try_acquire( + connection_id="socket-two", + instance_id="node-a", + room_id="general", + username="Andrew", + subject=None, + ) + finally: + await store.close() + + +@pytest.mark.asyncio +async def test_expired_socket_and_crashed_instance_occupancy_are_reclaimed( + clean_postgres_database: str, +) -> None: + store = _store(clean_postgres_database) + await store.initialize() + try: + await store.create_room(RoomCreate(id="general", name="General")) + await store.foundation.register_instance("node-a", lease_seconds=30) + registry = PostgresConnectionRegistry( + store.foundation, + max_connections=1, + max_connections_per_room=1, + ) + assert await registry.try_acquire( + connection_id="socket-old", + instance_id="node-a", + room_id="general", + username="Andrew", + subject=None, + ) + async with store.foundation.transaction() as connection: + await connection.execute( + """ + UPDATE public.samsarix_connection_leases + SET lease_expires_at = clock_timestamp() - interval '1 second' + WHERE connection_id = 'socket-old' + """ + ) + assert await registry.try_acquire( + connection_id="socket-new", + instance_id="node-a", + room_id="general", + username="Andrew", + subject=None, + ) + + async with store.foundation.transaction() as connection: + await connection.execute( + """ + UPDATE public.samsarix_instance_cursors + SET lease_expires_at = clock_timestamp() - interval '1 second' + WHERE instance_id = 'node-a' + """ + ) + assert await registry.counts(room_id="general") == ConnectionCounts(0, 0) + with pytest.raises(InstanceLeaseError, match="expired"): + await registry.try_acquire( + connection_id="socket-blocked", + instance_id="node-a", + room_id="general", + username="Andrew", + subject=None, + ) + + await store.foundation.register_instance("node-a", lease_seconds=30) + assert await registry.try_acquire( + connection_id="socket-after-restart", + instance_id="node-a", + room_id="general", + username="Andrew", + subject=None, + ) + assert await registry.release(connection_id="socket-after-restart", instance_id="node-a") + assert not await registry.release(connection_id="socket-after-restart", instance_id="node-a") + finally: + await store.close() + + +def test_registry_rejects_unsafe_configuration_and_identifiers() -> None: + foundation = PostgresFoundation("postgresql://unused") + with pytest.raises(ValueError, match="cannot exceed"): + PostgresConnectionRegistry(foundation, max_connections=2, max_connections_per_room=3) + with pytest.raises(ValueError, match="between 3 and 300"): + PostgresConnectionRegistry( + foundation, + max_connections=2, + max_connections_per_room=2, + lease_seconds=2, + ) diff --git a/tests/test_postgres_store.py b/tests/test_postgres_store.py index 15b97fa..7caf342 100644 --- a/tests/test_postgres_store.py +++ b/tests/test_postgres_store.py @@ -96,7 +96,7 @@ async def test_schema_v2_migrates_transactionally_and_widens_event_payloads( service = _store(clean_postgres_database) await service.initialize() try: - assert await service.foundation.schema_version() == POSTGRES_SCHEMA_VERSION == 3 + assert await service.foundation.schema_version() == POSTGRES_SCHEMA_VERSION == 4 assert await service.check_ready() assert await service.list_rooms() == [] async with service.foundation.transaction() as connection: From 8e7d9be6108ffc0214a35149b037ef54d2af83a4 Mon Sep 17 00:00:00 2001 From: Deathcharge Date: Sun, 2 Aug 2026 07:58:30 -0400 Subject: [PATCH 2/2] test(postgres): keep forced expiry schema-valid --- tests/test_postgres_connections.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_postgres_connections.py b/tests/test_postgres_connections.py index 30f64f4..fb1b918 100644 --- a/tests/test_postgres_connections.py +++ b/tests/test_postgres_connections.py @@ -172,7 +172,8 @@ async def test_expired_socket_and_crashed_instance_occupancy_are_reclaimed( await connection.execute( """ UPDATE public.samsarix_connection_leases - SET lease_expires_at = clock_timestamp() - interval '1 second' + SET created_at = clock_timestamp() - interval '2 seconds', + lease_expires_at = clock_timestamp() - interval '1 second' WHERE connection_id = 'socket-old' """ )