Replication: leaders, followers and quorums

postgres/streaming_lab.py

"""Streaming replication on PostgreSQL 18.6, measured: a primary and a standby whose link adds 5 ms in each direction
(a 10 ms round trip). What each synchronous_commit level adds to a commit, how far behind a read on the standby is,
and what a crash of the primary loses with asynchronous and with synchronous commits."""

import statistics

from pglab import Lab, wait_until

DELAY_MS = 5


def pair(lab, tag, sync):
    """A primary with one table and a streaming standby named standby1 behind a 5 ms link."""
    primary = lab.server(f"p{tag}")
    primary.start()
    primary.sql("CREATE TABLE ledger (id int PRIMARY KEY, note text)")
    link = lab.link(f"l{tag}", primary, DELAY_MS)
    standby = lab.standby(f"s{tag}", primary, primary_at=link.dir, app="standby1")
    standby.start()
    wait_until(lambda: primary.sql("SELECT state FROM pg_stat_replication").value == "streaming")
    if sync:
        primary.sql("ALTER SYSTEM SET synchronous_standby_names = 'standby1'")
        primary.sql("SELECT pg_reload_conf()")
        wait_until(lambda: primary.sql("SELECT sync_state FROM pg_stat_replication").value == "sync")
    return primary, standby, link


def visible(standby, row_id):
    return standby.sql(f"SELECT count(*) FROM ledger WHERE id = {row_id}").value == "1"


with Lab() as lab:
    primary, standby, link = pair(lab, 1, sync=True)
    row = primary.sql("SELECT application_name, state, sync_state FROM pg_stat_replication").rows[0]
    print(f"standby: {row[0]}, state {row[1]}, sync_state {row[2]}")
    print(f"the link adds {DELAY_MS} ms each way, so a round trip takes {2 * DELAY_MS} ms")
    print()

    print("1. What a commit costs, median of 30 single-row commits")
    next_id = 1
    with primary.connect() as c:
        for level in ("off", "local", "remote_write", "on", "remote_apply"):
            c.query(f"SET synchronous_commit = {level}")
            times = []
            for _ in range(30):
                times.append(c.timed(f"INSERT INTO ledger VALUES ({next_id}, 'level {level}')"))
                next_id += 1
            print(f"   {level:<13} {statistics.median(times):5.1f} ms")
    print()

    print("2. Is the new row on the standby straight after the commit returns? (10 commits per level)")

    def reads(level):
        global next_id
        found = 0
        with primary.connect() as c, standby.connect() as r:
            c.query(f"SET synchronous_commit = {level}")
            for _ in range(10):
                c.query(f"INSERT INTO ledger VALUES ({next_id}, 'read check')")
                found += r.value(f"SELECT count(*) FROM ledger WHERE id = {next_id}") == "1"
                next_id += 1
        return found

    idle = {level: reads(level) for level in ("local", "on", "remote_apply")}
    standby.configure(recovery_min_apply_delay="50ms")
    standby.restart()
    wait_until(lambda: primary.sql("SELECT sync_state FROM pg_stat_replication").value == "sync")
    late = {level: reads(level) for level in ("local", "on", "remote_apply")}
    print("   level          standby applies at once   standby applies each commit 50 ms late")
    for level in idle:
        print(f"   {level:<13}  {idle[level]:>2} of 10                  {late[level]:>2} of 10")
    with primary.connect() as c:
        c.query("SET synchronous_commit = remote_apply")
        slow = statistics.median(c.timed(f"INSERT INTO ledger VALUES ({next_id + i}, 'slow')") for i in range(5))
        next_id += 5
    print(f"   with the late standby, a remote_apply commit took {slow:.0f} ms (median of 5)")
    standby.configure(recovery_min_apply_delay="0")
    standby.restart()
    wait_until(lambda: primary.sql("SELECT sync_state FROM pg_stat_replication").value == "sync")
    with primary.connect() as c, standby.connect() as r:
        c.query("SET synchronous_commit = local")
        waits = []
        for _ in range(20):
            c.query(f"INSERT INTO ledger VALUES ({next_id}, 'lag check')")
            _, waited = wait_until(lambda: r.value(f"SELECT count(*) FROM ledger WHERE id = {next_id}") == "1")
            waits.append(waited * 1000)
            next_id += 1
    print(f"   an asynchronous commit reached readers on the standby after {statistics.median(waits):.1f} ms"
          " (median of 20)")
    print()

    print("3. The primary crashes while the link is down")
    for sync in (False, True):
        primary, standby, link = pair(lab, 2 if not sync else 3, sync)
        mode = "synchronous (on)" if sync else "asynchronous (local)"
        with primary.connect() as c:
            c.query("SET synchronous_commit = " + ("on" if sync else "local"))
            for i in range(1, 6):
                c.query(f"INSERT INTO ledger VALUES ({i}, 'before the outage')")
            wait_until(lambda: visible(standby, 5))
            link.stall()
            acknowledged = 0
            for i in range(6, 26):
                c.send(f"INSERT INTO ledger VALUES ({i}, 'during the outage')")
                try:
                    c.wait(timeout=1.0)
                    acknowledged += 1
                except TimeoutError:
                    print(f"   {mode}: commit {i} got no answer within 1 s, so the application does not count it")
                    break
            link.cut()
            primary.stop("immediate")
        standby.promote()
        survived = int(standby.sql("SELECT count(*) FROM ledger WHERE id > 5").value)
        print(f"   {mode}: {acknowledged} commits acknowledged during the outage, {survived} of them on the promoted"
              f" standby: {acknowledged - survived} acknowledged commits lost")
        primary.start()
        extra = int(primary.sql("SELECT count(*) FROM ledger WHERE id > 5").value)
        print(f"   the old primary, restarted on its own, holds {extra} row(s) the new primary does not have")
        primary.stop("immediate")

standby: standby1, state streaming, sync_state sync
the link adds 5 ms each way, so a round trip takes 10 ms

1. What a commit costs, median of 30 single-row commits
   off             0.2 ms
   local           0.1 ms
   remote_write   13.0 ms
   on             12.9 ms
   remote_apply   12.9 ms

2. Is the new row on the standby straight after the commit returns? (10 commits per level)
   level          standby applies at once   standby applies each commit 50 ms late
   local           0 of 10                   0 of 10
   on             10 of 10                   0 of 10
   remote_apply   10 of 10                  10 of 10
   with the late standby, a remote_apply commit took 58 ms (median of 5)
   an asynchronous commit reached readers on the standby after 6.8 ms (median of 20)

3. The primary crashes while the link is down
   asynchronous (local): 20 commits acknowledged during the outage, 0 of them on the promoted standby: 20 acknowledged commits lost
   the old primary, restarted on its own, holds 20 row(s) the new primary does not have
   synchronous (on): commit 6 got no answer within 1 s, so the application does not count it
   synchronous (on): 0 commits acknowledged during the outage, 0 of them on the promoted standby: 0 acknowledged commits lost
   the old primary, restarted on its own, holds 1 row(s) the new primary does not have

postgres/pglab.py

"""A small PostgreSQL lab for the examples of this lesson: throwaway PostgreSQL 18.6 servers, a client that speaks the
PostgreSQL wire protocol directly (so the examples can time a single commit without starting a program per query),
standbys made with pg_basebackup, and a link between two servers that adds a fixed delay in each direction or stalls.

Every server lives in a new folder under the run's temporary folder (TMPDIR), listens only on a Unix socket in that
folder (never on a network port), and is stopped and removed at the end. The PostgreSQL 18.6 programs are the ones
installed by mise next to the Python that runs this file. fsync is off: these servers never need to survive a power
cut, so the times the examples print come from replication, not from the disk.
"""

import os
import shutil
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time
from collections import deque
from pathlib import Path

PG_VERSION = "18.6"
# .../installs/python/3.14.8/bin/python3 -> .../installs
INSTALLS = Path(sys.executable).resolve().parents[3]
BIN = INSTALLS / "postgres" / PG_VERSION / "bin"


def _run(args, env=None, check=True):
    done = subprocess.run([str(a) for a in args], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
                          env=env)
    if check and done.returncode != 0:
        sys.exit(f"{Path(str(args[0])).name} failed:\n{done.stdout}")
    return done


class Lab:
    """A folder for the lab's servers and links, removed when the lab closes."""

    def __init__(self):
        if not (BIN / "postgres").exists():
            sys.exit(f"PostgreSQL {PG_VERSION} was not found next to this Python")
        self.base = Path(tempfile.mkdtemp(prefix="pg"))
        self.servers = []
        self.links = []

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        for link in self.links:
            link.close()
        for s in reversed(self.servers):
            s.stop("immediate", quiet=True)
        shutil.rmtree(self.base, ignore_errors=True)

    def server(self, name, settings=None):
        """A new, empty server (initdb), not started yet."""
        s = Server(self.base / name, settings)
        s.initdb()
        self.servers.append(s)
        return s

    def standby(self, name, source, primary_at=None, app=None, settings=None):
        """A standby of `source`, copied with pg_basebackup; it streams from `primary_at` (a folder whose socket leads
        to the primary: the primary's own, or a Link's) under the application_name `app`."""
        s = Server(self.base / name, settings)
        s.dir.mkdir(parents=True)
        _run([BIN / "pg_basebackup", "-h", source.dir, "-U", "learner", "-D", s.data, "-X", "stream", "-c", "fast",
              "--no-sync"])
        conninfo = f"host={primary_at or source.dir} port=5432 user=learner application_name={app or name}"
        with open(s.data / "postgresql.auto.conf", "a") as f:
            f.write(f"primary_conninfo = '{conninfo}'\n")
        (s.data / "standby.signal").write_text("")
        self.servers.append(s)
        return s

    def link(self, name, to, delay_ms):
        """A link to server `to` whose socket is in a folder of its own, delaying each direction by delay_ms."""
        l = Link(self.base / name, to.dir, delay_ms / 1000)
        self.links.append(l)
        return l


class Server:
    def __init__(self, folder, settings=None):
        self.dir = folder
        self.name = folder.name
        self.data = folder / "data"
        self.settings = {"listen_addresses": "''", "unix_socket_directories": f"'{folder}'", "fsync": "off",
                         "full_page_writes": "off", "synchronous_commit": "on", **(settings or {})}

    def initdb(self):
        self.dir.mkdir(parents=True)
        _run([BIN / "initdb", "-D", self.data, "-U", "learner", "--auth=trust", "--no-instructions", "--no-sync",
              "-E", "UTF8", "--locale=C"])

    def configure(self, **settings):
        """Settings for the next start (or a reload)."""
        for k, v in settings.items():
            self.settings[k] = v

    def _options(self):
        return " ".join(f"-c {k}={v}" for k, v in self.settings.items())

    def start(self):
        done = _run([BIN / "pg_ctl", "-D", self.data, "-o", self._options(), "-l", self.dir / "log", "-w", "start"],
                    check=False)
        if done.returncode != 0:
            sys.exit(f"the server did not start:\n{(self.dir / 'log').read_text()}")

    def stop(self, mode="fast", quiet=False):
        if (self.data / "postmaster.pid").exists():
            _run([BIN / "pg_ctl", "-D", self.data, "-m", mode, "-w", "stop"], check=not quiet)

    def restart(self):
        self.stop("fast")
        self.start()

    def promote(self):
        _run([BIN / "pg_ctl", "-D", self.data, "-w", "promote"])

    def connect(self, app="client", db="postgres"):
        return Conn(self.dir, app=app, db=db)

    def sql(self, statement, db="postgres"):
        """Runs one statement on a connection of its own; returns the Result."""
        with self.connect(db=db) as c:
            return c.query(statement)

    def psql(self, statement, db="postgres"):
        """What psql prints for one statement (for output that looks like a terminal session)."""
        env = {**os.environ, "PGHOST": str(self.dir), "PGUSER": "learner", "PGDATABASE": db}
        return _run([BIN / "psql", "-X", "-c", statement], env=env, check=False).stdout.rstrip("\n")


class PgError(Exception):
    def __init__(self, fields):
        self.fields = fields
        self.code = fields.get("C", "")
        super().__init__(f"{fields.get('S', 'ERROR')}:  {fields.get('M', '')}")


class Result:
    def __init__(self):
        self.columns = []
        self.rows = []
        self.tag = None
        self.error = None
        self.notices = []
        self.status = None

    @property
    def value(self):
        """The first column of the first row (or None)."""
        return self.rows[0][0] if self.rows else None


class Conn:
    """One connection, speaking the PostgreSQL frontend/backend protocol (version 3.0) over the Unix socket."""

    def __init__(self, folder, port=5432, app="client", db="postgres"):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(str(Path(folder) / f".s.PGSQL.{port}"))
        self.buf = b""
        params = {"user": "learner", "database": db, "application_name": app, "client_encoding": "UTF8"}
        body = struct.pack("!i", 196608) + b"".join(k.encode() + b"\0" + v.encode() + b"\0" for k, v in params.items())
        body += b"\0"
        self.sock.sendall(struct.pack("!i", len(body) + 4) + body)
        self.pending = 0
        result = self._collect()
        if result.error:
            raise PgError(result.error)

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def close(self):
        try:
            self.sock.sendall(b"X" + struct.pack("!i", 4))
        except OSError:
            pass
        self.sock.close()

    def _read(self, n):
        while len(self.buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise ConnectionError("the server closed the connection")
            self.buf += chunk
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def _message(self):
        kind = self._read(1)
        (length,) = struct.unpack("!i", self._read(4))
        return kind, self._read(length - 4)

    @staticmethod
    def _fields(payload):
        out = {}
        for part in payload.split(b"\0"):
            if part:
                out[chr(part[0])] = part[1:].decode()
        return out

    def _collect(self):
        """Reads messages up to ReadyForQuery."""
        r = Result()
        while True:
            kind, payload = self._message()
            if kind == b"T":
                (n,) = struct.unpack("!h", payload[:2])
                pos, cols = 2, []
                for _ in range(n):
                    end = payload.index(b"\0", pos)
                    cols.append(payload[pos:end].decode())
                    pos = end + 1 + 18
                r.columns = cols
            elif kind == b"D":
                (n,) = struct.unpack("!h", payload[:2])
                pos, row = 2, []
                for _ in range(n):
                    (size,) = struct.unpack("!i", payload[pos:pos + 4])
                    pos += 4
                    if size < 0:
                        row.append(None)
                    else:
                        row.append(payload[pos:pos + size].decode())
                        pos += size
                r.rows.append(tuple(row))
            elif kind == b"C":
                r.tag = payload.rstrip(b"\0").decode()
            elif kind == b"E":
                r.error = self._fields(payload)
            elif kind == b"N":
                r.notices.append(self._fields(payload))
            elif kind == b"Z":
                r.status = payload.decode()
                return r

    def send(self, sql):
        """Sends a query without waiting for its result (wait() reads it)."""
        q = sql.encode() + b"\0"
        self.sock.sendall(b"Q" + struct.pack("!i", len(q) + 4) + q)
        self.pending += 1

    def wait(self, timeout=None):
        """The result of the query sent last; raises TimeoutError when nothing complete arrives within timeout."""
        self.sock.settimeout(timeout)
        try:
            result = self._collect()
        except (socket.timeout, TimeoutError):
            raise TimeoutError from None
        finally:
            self.sock.settimeout(None)
        self.pending -= 1
        return result

    def query(self, sql, check=True):
        self.send(sql)
        r = self.wait()
        if check and r.error:
            raise PgError(r.error)
        return r

    def value(self, sql):
        return self.query(sql).value

    def timed(self, sql):
        """Runs a query and returns how long it took, in milliseconds."""
        t = time.perf_counter()
        self.query(sql)
        return (time.perf_counter() - t) * 1000


class Link:
    """A Unix-socket link to a server that delays every byte by `delay` seconds in each direction, or holds it back
    while stalled. A standby that connects through it behaves as if the primary were that far away."""

    def __init__(self, folder, target_dir, delay):
        folder.mkdir(parents=True)
        self.target = str(Path(target_dir) / ".s.PGSQL.5432")
        self.delay = delay
        self.stalled = threading.Event()
        self.closed = threading.Event()
        self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.listener.bind(str(folder / ".s.PGSQL.5432"))
        self.listener.listen(8)
        self.dir = folder
        self.socks = []
        threading.Thread(target=self._accept, daemon=True).start()

    def stall(self):
        """Holds back everything in both directions, as a network that stops delivering would."""
        self.stalled.set()

    def resume(self):
        self.stalled.clear()

    def cut(self):
        """Breaks every connection through the link; what it was holding back is never delivered."""
        self.closed.set()
        for s in self.socks:
            try:
                s.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
        self.close()

    def close(self):
        self.closed.set()
        for s in [self.listener, *self.socks]:
            try:
                s.close()
            except OSError:
                pass

    def _accept(self):
        while not self.closed.is_set():
            try:
                client, _ = self.listener.accept()
            except OSError:
                return
            server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            server.connect(self.target)
            self.socks += [client, server]
            for a, b in ((client, server), (server, client)):
                q, cv = deque(), threading.Condition()
                threading.Thread(target=self._read, args=(a, q, cv), daemon=True).start()
                threading.Thread(target=self._write, args=(b, q, cv), daemon=True).start()

    def _read(self, src, q, cv):
        while True:
            try:
                data = src.recv(65536)
            except OSError:
                data = b""
            with cv:
                q.append((time.monotonic() + self.delay, data))
                cv.notify()
            if not data:
                return

    def _write(self, dst, q, cv):
        while True:
            with cv:
                while not q:
                    cv.wait()
                due, data = q[0]
            if self.closed.is_set():
                return
            now = time.monotonic()
            if now < due or self.stalled.is_set():
                time.sleep(max(due - now, 0.0005))
                continue
            with cv:
                q.popleft()
            if not data:
                try:
                    dst.shutdown(socket.SHUT_WR)
                except OSError:
                    pass
                return
            try:
                dst.sendall(data)
            except OSError:
                return


def wait_until(check, timeout=10.0, step=0.0005):
    """Polls check() until it returns something true; returns (value, seconds waited)."""
    start = time.perf_counter()
    while True:
        v = check()
        if v:
            return v, time.perf_counter() - start
        if time.perf_counter() - start > timeout:
            sys.exit("timed out waiting")
        time.sleep(step)

postgres/logical_lab.py

"""Logical replication on PostgreSQL 18.6: a publisher and a subscriber that is an ordinary, writable database.
Rows travel as changes to named tables; a table without a key cannot publish updates, and a column added on the
publisher alone stops replication until the subscriber has it too."""

from pglab import Lab, wait_until


def show(server, statement, shown=None):
    """Prints a statement and what psql printed for it, as a terminal session would."""
    print(f"{server.name}=> {shown or statement}")
    print(server.psql(statement))
    print()


with Lab() as lab:
    pub = lab.server("pub", {"wal_level": "logical"})
    sub = lab.server("sub", {"wal_retrieve_retry_interval": "200ms"})
    pub.name, sub.name = "publisher", "subscriber"
    pub.start()
    sub.start()
    for s in (pub, sub):
        s.sql("CREATE TABLE accounts (id int PRIMARY KEY, holder text, balance int)")
        s.sql("CREATE TABLE visits (page text, seen int)")
    pub.sql("INSERT INTO accounts VALUES (1, 'Asha', 500), (2, 'Ravi', 300)")

    show(pub, "CREATE PUBLICATION shop FOR TABLE accounts, visits;")
    subscribe = "CREATE SUBSCRIPTION shop_sub CONNECTION 'host={} user=learner dbname=postgres' PUBLICATION shop;"
    # The publisher's socket folder is a temporary path of this computer: the session shows a placeholder instead.
    show(sub, subscribe.format(pub.dir), shown=subscribe.format("<publisher socket folder>"))
    wait_until(lambda: sub.sql("SELECT count(*) FROM accounts").value == "2")
    show(sub, "SELECT * FROM accounts ORDER BY id;")

    show(pub, "UPDATE accounts SET balance = balance - 100 WHERE id = 1;")
    wait_until(lambda: sub.sql("SELECT balance FROM accounts WHERE id = 1").value == "400")
    show(sub, "SELECT * FROM accounts ORDER BY id;")

    show(pub, "INSERT INTO visits VALUES ('/home', 1);")
    show(pub, "UPDATE visits SET seen = 2;")

    show(sub, "CREATE TABLE notes (account_id int, note text);")
    show(sub, "CREATE INDEX accounts_holder ON accounts (holder);")

    show(pub, "ALTER TABLE accounts ADD COLUMN branch text;")
    show(pub, "INSERT INTO accounts VALUES (3, 'Meena', 900, 'Pune');")
    wait_until(lambda: sub.sql("SELECT apply_error_count > 0 FROM pg_stat_subscription_stats").value == "t")
    log = (sub.dir / "log").read_text().splitlines()
    error = next(line for line in log if "missing replicated column" in line)
    print("the subscriber's log says:")
    print("  " + error[error.index("ERROR:"):])
    print()
    show(sub, "SELECT count(*) FROM accounts;")
    show(sub, "ALTER TABLE accounts ADD COLUMN branch text;")
    wait_until(lambda: sub.sql("SELECT count(*) FROM accounts").value == "3")
    show(sub, "SELECT * FROM accounts ORDER BY id;")

publisher=> CREATE PUBLICATION shop FOR TABLE accounts, visits;
CREATE PUBLICATION

subscriber=> CREATE SUBSCRIPTION shop_sub CONNECTION 'host=<publisher socket folder> user=learner dbname=postgres' PUBLICATION shop;
NOTICE:  created replication slot "shop_sub" on publisher
CREATE SUBSCRIPTION

subscriber=> SELECT * FROM accounts ORDER BY id;
 id | holder | balance 
----+--------+---------
  1 | Asha   |     500
  2 | Ravi   |     300
(2 rows)

publisher=> UPDATE accounts SET balance = balance - 100 WHERE id = 1;
UPDATE 1

subscriber=> SELECT * FROM accounts ORDER BY id;
 id | holder | balance 
----+--------+---------
  1 | Asha   |     400
  2 | Ravi   |     300
(2 rows)

publisher=> INSERT INTO visits VALUES ('/home', 1);
INSERT 0 1

publisher=> UPDATE visits SET seen = 2;
ERROR:  cannot update table "visits" because it does not have a replica identity and publishes updates
HINT:  To enable updating the table, set REPLICA IDENTITY using ALTER TABLE.

subscriber=> CREATE TABLE notes (account_id int, note text);
CREATE TABLE

subscriber=> CREATE INDEX accounts_holder ON accounts (holder);
CREATE INDEX

publisher=> ALTER TABLE accounts ADD COLUMN branch text;
ALTER TABLE

publisher=> INSERT INTO accounts VALUES (3, 'Meena', 900, 'Pune');
INSERT 0 1

the subscriber's log says:
  ERROR:  logical replication target relation "public.accounts" is missing replicated column: "branch"

subscriber=> SELECT count(*) FROM accounts;
 count 
-------
     2
(1 row)

subscriber=> ALTER TABLE accounts ADD COLUMN branch text;
ALTER TABLE

subscriber=> SELECT * FROM accounts ORDER BY id;
 id | holder | balance | branch 
----+--------+---------+--------
  1 | Asha   |     400 | 
  2 | Ravi   |     300 | 
  3 | Meena  |     900 | Pune
(3 rows)

postgres/pglab.py

"""A small PostgreSQL lab for the examples of this lesson: throwaway PostgreSQL 18.6 servers, a client that speaks the
PostgreSQL wire protocol directly (so the examples can time a single commit without starting a program per query),
standbys made with pg_basebackup, and a link between two servers that adds a fixed delay in each direction or stalls.

Every server lives in a new folder under the run's temporary folder (TMPDIR), listens only on a Unix socket in that
folder (never on a network port), and is stopped and removed at the end. The PostgreSQL 18.6 programs are the ones
installed by mise next to the Python that runs this file. fsync is off: these servers never need to survive a power
cut, so the times the examples print come from replication, not from the disk.
"""

import os
import shutil
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time
from collections import deque
from pathlib import Path

PG_VERSION = "18.6"
# .../installs/python/3.14.8/bin/python3 -> .../installs
INSTALLS = Path(sys.executable).resolve().parents[3]
BIN = INSTALLS / "postgres" / PG_VERSION / "bin"


def _run(args, env=None, check=True):
    done = subprocess.run([str(a) for a in args], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
                          env=env)
    if check and done.returncode != 0:
        sys.exit(f"{Path(str(args[0])).name} failed:\n{done.stdout}")
    return done


class Lab:
    """A folder for the lab's servers and links, removed when the lab closes."""

    def __init__(self):
        if not (BIN / "postgres").exists():
            sys.exit(f"PostgreSQL {PG_VERSION} was not found next to this Python")
        self.base = Path(tempfile.mkdtemp(prefix="pg"))
        self.servers = []
        self.links = []

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        for link in self.links:
            link.close()
        for s in reversed(self.servers):
            s.stop("immediate", quiet=True)
        shutil.rmtree(self.base, ignore_errors=True)

    def server(self, name, settings=None):
        """A new, empty server (initdb), not started yet."""
        s = Server(self.base / name, settings)
        s.initdb()
        self.servers.append(s)
        return s

    def standby(self, name, source, primary_at=None, app=None, settings=None):
        """A standby of `source`, copied with pg_basebackup; it streams from `primary_at` (a folder whose socket leads
        to the primary: the primary's own, or a Link's) under the application_name `app`."""
        s = Server(self.base / name, settings)
        s.dir.mkdir(parents=True)
        _run([BIN / "pg_basebackup", "-h", source.dir, "-U", "learner", "-D", s.data, "-X", "stream", "-c", "fast",
              "--no-sync"])
        conninfo = f"host={primary_at or source.dir} port=5432 user=learner application_name={app or name}"
        with open(s.data / "postgresql.auto.conf", "a") as f:
            f.write(f"primary_conninfo = '{conninfo}'\n")
        (s.data / "standby.signal").write_text("")
        self.servers.append(s)
        return s

    def link(self, name, to, delay_ms):
        """A link to server `to` whose socket is in a folder of its own, delaying each direction by delay_ms."""
        l = Link(self.base / name, to.dir, delay_ms / 1000)
        self.links.append(l)
        return l


class Server:
    def __init__(self, folder, settings=None):
        self.dir = folder
        self.name = folder.name
        self.data = folder / "data"
        self.settings = {"listen_addresses": "''", "unix_socket_directories": f"'{folder}'", "fsync": "off",
                         "full_page_writes": "off", "synchronous_commit": "on", **(settings or {})}

    def initdb(self):
        self.dir.mkdir(parents=True)
        _run([BIN / "initdb", "-D", self.data, "-U", "learner", "--auth=trust", "--no-instructions", "--no-sync",
              "-E", "UTF8", "--locale=C"])

    def configure(self, **settings):
        """Settings for the next start (or a reload)."""
        for k, v in settings.items():
            self.settings[k] = v

    def _options(self):
        return " ".join(f"-c {k}={v}" for k, v in self.settings.items())

    def start(self):
        done = _run([BIN / "pg_ctl", "-D", self.data, "-o", self._options(), "-l", self.dir / "log", "-w", "start"],
                    check=False)
        if done.returncode != 0:
            sys.exit(f"the server did not start:\n{(self.dir / 'log').read_text()}")

    def stop(self, mode="fast", quiet=False):
        if (self.data / "postmaster.pid").exists():
            _run([BIN / "pg_ctl", "-D", self.data, "-m", mode, "-w", "stop"], check=not quiet)

    def restart(self):
        self.stop("fast")
        self.start()

    def promote(self):
        _run([BIN / "pg_ctl", "-D", self.data, "-w", "promote"])

    def connect(self, app="client", db="postgres"):
        return Conn(self.dir, app=app, db=db)

    def sql(self, statement, db="postgres"):
        """Runs one statement on a connection of its own; returns the Result."""
        with self.connect(db=db) as c:
            return c.query(statement)

    def psql(self, statement, db="postgres"):
        """What psql prints for one statement (for output that looks like a terminal session)."""
        env = {**os.environ, "PGHOST": str(self.dir), "PGUSER": "learner", "PGDATABASE": db}
        return _run([BIN / "psql", "-X", "-c", statement], env=env, check=False).stdout.rstrip("\n")


class PgError(Exception):
    def __init__(self, fields):
        self.fields = fields
        self.code = fields.get("C", "")
        super().__init__(f"{fields.get('S', 'ERROR')}:  {fields.get('M', '')}")


class Result:
    def __init__(self):
        self.columns = []
        self.rows = []
        self.tag = None
        self.error = None
        self.notices = []
        self.status = None

    @property
    def value(self):
        """The first column of the first row (or None)."""
        return self.rows[0][0] if self.rows else None


class Conn:
    """One connection, speaking the PostgreSQL frontend/backend protocol (version 3.0) over the Unix socket."""

    def __init__(self, folder, port=5432, app="client", db="postgres"):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(str(Path(folder) / f".s.PGSQL.{port}"))
        self.buf = b""
        params = {"user": "learner", "database": db, "application_name": app, "client_encoding": "UTF8"}
        body = struct.pack("!i", 196608) + b"".join(k.encode() + b"\0" + v.encode() + b"\0" for k, v in params.items())
        body += b"\0"
        self.sock.sendall(struct.pack("!i", len(body) + 4) + body)
        self.pending = 0
        result = self._collect()
        if result.error:
            raise PgError(result.error)

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def close(self):
        try:
            self.sock.sendall(b"X" + struct.pack("!i", 4))
        except OSError:
            pass
        self.sock.close()

    def _read(self, n):
        while len(self.buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise ConnectionError("the server closed the connection")
            self.buf += chunk
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def _message(self):
        kind = self._read(1)
        (length,) = struct.unpack("!i", self._read(4))
        return kind, self._read(length - 4)

    @staticmethod
    def _fields(payload):
        out = {}
        for part in payload.split(b"\0"):
            if part:
                out[chr(part[0])] = part[1:].decode()
        return out

    def _collect(self):
        """Reads messages up to ReadyForQuery."""
        r = Result()
        while True:
            kind, payload = self._message()
            if kind == b"T":
                (n,) = struct.unpack("!h", payload[:2])
                pos, cols = 2, []
                for _ in range(n):
                    end = payload.index(b"\0", pos)
                    cols.append(payload[pos:end].decode())
                    pos = end + 1 + 18
                r.columns = cols
            elif kind == b"D":
                (n,) = struct.unpack("!h", payload[:2])
                pos, row = 2, []
                for _ in range(n):
                    (size,) = struct.unpack("!i", payload[pos:pos + 4])
                    pos += 4
                    if size < 0:
                        row.append(None)
                    else:
                        row.append(payload[pos:pos + size].decode())
                        pos += size
                r.rows.append(tuple(row))
            elif kind == b"C":
                r.tag = payload.rstrip(b"\0").decode()
            elif kind == b"E":
                r.error = self._fields(payload)
            elif kind == b"N":
                r.notices.append(self._fields(payload))
            elif kind == b"Z":
                r.status = payload.decode()
                return r

    def send(self, sql):
        """Sends a query without waiting for its result (wait() reads it)."""
        q = sql.encode() + b"\0"
        self.sock.sendall(b"Q" + struct.pack("!i", len(q) + 4) + q)
        self.pending += 1

    def wait(self, timeout=None):
        """The result of the query sent last; raises TimeoutError when nothing complete arrives within timeout."""
        self.sock.settimeout(timeout)
        try:
            result = self._collect()
        except (socket.timeout, TimeoutError):
            raise TimeoutError from None
        finally:
            self.sock.settimeout(None)
        self.pending -= 1
        return result

    def query(self, sql, check=True):
        self.send(sql)
        r = self.wait()
        if check and r.error:
            raise PgError(r.error)
        return r

    def value(self, sql):
        return self.query(sql).value

    def timed(self, sql):
        """Runs a query and returns how long it took, in milliseconds."""
        t = time.perf_counter()
        self.query(sql)
        return (time.perf_counter() - t) * 1000


class Link:
    """A Unix-socket link to a server that delays every byte by `delay` seconds in each direction, or holds it back
    while stalled. A standby that connects through it behaves as if the primary were that far away."""

    def __init__(self, folder, target_dir, delay):
        folder.mkdir(parents=True)
        self.target = str(Path(target_dir) / ".s.PGSQL.5432")
        self.delay = delay
        self.stalled = threading.Event()
        self.closed = threading.Event()
        self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.listener.bind(str(folder / ".s.PGSQL.5432"))
        self.listener.listen(8)
        self.dir = folder
        self.socks = []
        threading.Thread(target=self._accept, daemon=True).start()

    def stall(self):
        """Holds back everything in both directions, as a network that stops delivering would."""
        self.stalled.set()

    def resume(self):
        self.stalled.clear()

    def cut(self):
        """Breaks every connection through the link; what it was holding back is never delivered."""
        self.closed.set()
        for s in self.socks:
            try:
                s.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
        self.close()

    def close(self):
        self.closed.set()
        for s in [self.listener, *self.socks]:
            try:
                s.close()
            except OSError:
                pass

    def _accept(self):
        while not self.closed.is_set():
            try:
                client, _ = self.listener.accept()
            except OSError:
                return
            server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            server.connect(self.target)
            self.socks += [client, server]
            for a, b in ((client, server), (server, client)):
                q, cv = deque(), threading.Condition()
                threading.Thread(target=self._read, args=(a, q, cv), daemon=True).start()
                threading.Thread(target=self._write, args=(b, q, cv), daemon=True).start()

    def _read(self, src, q, cv):
        while True:
            try:
                data = src.recv(65536)
            except OSError:
                data = b""
            with cv:
                q.append((time.monotonic() + self.delay, data))
                cv.notify()
            if not data:
                return

    def _write(self, dst, q, cv):
        while True:
            with cv:
                while not q:
                    cv.wait()
                due, data = q[0]
            if self.closed.is_set():
                return
            now = time.monotonic()
            if now < due or self.stalled.is_set():
                time.sleep(max(due - now, 0.0005))
                continue
            with cv:
                q.popleft()
            if not data:
                try:
                    dst.shutdown(socket.SHUT_WR)
                except OSError:
                    pass
                return
            try:
                dst.sendall(data)
            except OSError:
                return


def wait_until(check, timeout=10.0, step=0.0005):
    """Polls check() until it returns something true; returns (value, seconds waited)."""
    start = time.perf_counter()
    while True:
        v = check()
        if v:
            return v, time.perf_counter() - start
        if time.perf_counter() - start > timeout:
            sys.exit("timed out waiting")
        time.sleep(step)