Add a Postgres proxy that counts round trips, and use it in the bench

`scripts-dev/pg_roundtrip_proxy.py` sits between Synapse and Postgres. It
counts round trips, a client message after a server message on the same
connection, can add latency to each, and records the protocol messages
and SQL of each round trip. SIGUSR2 resets it and SIGUSR1 writes the
counts and log to a file.

`createroom_bench.py --proxy-delay-ms N` runs the server through it and
reports the round trips per request beside the timings. With a delay,
the extra wall time over a run without one shows how many of those round
trips the request waits for in turn. A new `e2ee_invite` scenario is the
room the E2EE monitor creates on jki.re, and `--scenarios` picks a
subset. `persist_bench.Server` honours `PGPORT`.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PVLsNh6akt4ifEizY4D1ta

scripts-dev/createroom_bench.py

@@ -29,7 +29,10 @@ the server on psycopg2 instead of the native driver.
 import argparse
 import json
 import os
+import shutil
+import signal
 import statistics
+import subprocess
 import sys
 import tempfile
 import time
@@ -41,6 +44,19 @@ import persist_bench as pb  # noqa: E402
 
 SCENARIOS: dict[str, dict[str, Any]] = {
     "plain": {},
+    # The room the E2EE monitor on jki.re creates every minute.
+    "e2ee_invite": {
+        "preset": "trusted_private_chat",
+        "is_direct": True,
+        "invite": ["@invitee:bench"],
+        "initial_state": [
+            {
+                "type": "m.room.encryption",
+                "state_key": "",
+                "content": {"algorithm": "m.megolm.v1.aes-sha2"},
+            }
+        ],
+    },
     "private_invite": {"preset": "private_chat", "invite": ["@invitee:bench"]},
     "public_named": {"preset": "public_chat", "name": "Bench room", "topic": "A topic"},
     "initial_state": {
@@ -71,8 +87,77 @@ SCENARIOS: dict[str, dict[str, Any]] = {
 
 CREATE_ROOM = "/_matrix/client/v3/createRoom"
 
+# The local port `pg_roundtrip_proxy.py` listens on, with `--proxy-delay-ms`.
+PROXY_PORT = 15432
 
-def bench_level(level: str, rooms: int, rust_driver: bool) -> list[dict[str, Any]]:
+
+class RoundTripProxy:
+    """`pg_roundtrip_proxy.py` between the server and Postgres, counting the
+    round trips of each request."""
+
+    def __init__(self, delay_ms: float) -> None:
+        self.stats_dir = tempfile.mkdtemp(prefix="pg-roundtrips-")
+        self.stats = os.path.join(self.stats_dir, "stats.json")
+        target = (
+            f"{os.environ.get('PGHOST', 'localhost')}:"
+            f"{os.environ.get('PGPORT', '5432')}"
+        )
+        self.process = subprocess.Popen(
+            [
+                sys.executable,
+                os.path.join(
+                    os.path.dirname(os.path.abspath(__file__)), "pg_roundtrip_proxy.py"
+                ),
+                f"--listen={PROXY_PORT}",
+                f"--target={target}",
+                f"--delay-ms={delay_ms}",
+                f"--stats={self.stats}",
+            ]
+        )
+        time.sleep(1)
+        self.check_running()
+        os.environ["PGHOST"] = "127.0.0.1"
+        os.environ["PGPORT"] = str(PROXY_PORT)
+
+    def check_running(self) -> None:
+        if self.process.poll() is not None:
+            raise RuntimeError(
+                f"pg_roundtrip_proxy.py exited with {self.process.returncode}"
+            )
+
+    def reset(self) -> None:
+        self.process.send_signal(signal.SIGUSR2)
+        time.sleep(0.05)
+
+    def turns(self) -> int:
+        """The round trips since the last `reset`."""
+        # The proxy replaces the file atomically, so it is complete once it
+        # is gone and back.
+        if os.path.exists(self.stats):
+            os.unlink(self.stats)
+        self.process.send_signal(signal.SIGUSR1)
+        deadline = time.monotonic() + 10
+        while not os.path.exists(self.stats):
+            self.check_running()
+            if time.monotonic() > deadline:
+                raise RuntimeError("pg_roundtrip_proxy.py did not write its stats")
+            time.sleep(0.01)
+        with open(self.stats) as f:
+            return json.load(f)["turns"]
+
+    def stop(self) -> None:
+        self.process.terminate()
+        self.process.wait()
+        shutil.rmtree(self.stats_dir, ignore_errors=True)
+
+
+def bench_level(
+    level: str,
+    rooms: int,
+    rust_driver: bool,
+    scenarios: list[str],
+    proxy: RoundTripProxy | None = None,
+) -> list[dict[str, Any]]:
     """Benchmark a server at `level`: a `rust_event_persistence` level,
     optionally followed by `+rooms` for `rust_room_creation: on`."""
     results = []
@@ -88,15 +173,23 @@ def bench_level(level: str, rooms: int, rust_driver: bool) -> list[dict[str, Any
         server.start()
         token = pb.register("creator")
         pb.register("invitee")
-        for name, body in SCENARIOS.items():
+        for name in scenarios:
+            body = SCENARIOS[name]
             for _ in range(5):
                 pb.ok("POST", CREATE_ROOM, body, token)
             before = server.sample()
             times = []
+            turns = []
             for _ in range(rooms):
+                if proxy is not None:
+                    # Let the last room's background work finish first.
+                    time.sleep(0.5)
+                    proxy.reset()
                 t0 = time.perf_counter()
                 pb.ok("POST", CREATE_ROOM, body, token)
                 times.append(time.perf_counter() - t0)
+                if proxy is not None:
+                    turns.append(proxy.turns())
             after = server.sample()
             row = {
                 "level": level,
@@ -110,6 +203,10 @@ def bench_level(level: str, rooms: int, rust_driver: bool) -> list[dict[str, Any
                 * (after.process_cpu - before.process_cpu)
                 / rooms,
             }
+            if turns:
+                # Background work from the request counts too, if it
+                # overlaps.
+                row["round_trips"] = statistics.median(turns)
             results.append(row)
             print(row, flush=True)
     finally:
@@ -127,22 +224,50 @@ def main() -> None:
     )
     parser.add_argument("--rooms", type=int, default=40)
     parser.add_argument("--psycopg2", action="store_true")
+    parser.add_argument(
+        "--scenarios",
+        default=",".join(SCENARIOS),
+        help="comma-separated scenarios to run, from: " + ", ".join(SCENARIOS),
+    )
+    parser.add_argument(
+        "--proxy-delay-ms",
+        type=float,
+        help="run Postgres connections through pg_roundtrip_proxy.py, adding this"
+        " much latency to each round trip (0 to only count), and report the round"
+        " trips per request. The extra wall time over a run without the proxy,"
+        " divided by the delay, is the number of round trips the request waits"
+        " for in turn.",
+    )
     parser.add_argument("--json", help="also write the results to this file")
     args = parser.parse_args()
 
+    proxy = None
+    if args.proxy_delay_ms is not None:
+        proxy = RoundTripProxy(args.proxy_delay_ms)
+
     results = []
-    for level in args.levels.split(","):
-        results += bench_level(level, args.rooms, rust_driver=not args.psycopg2)
+    try:
+        for level in args.levels.split(","):
+            results += bench_level(
+                level,
+                args.rooms,
+                rust_driver=not args.psycopg2,
+                scenarios=args.scenarios.split(","),
+                proxy=proxy,
+            )
+    finally:
+        if proxy is not None:
+            proxy.stop()
 
     print(
         f"\n{'scenario':<16} {'level':<12} {'mean ms':>8} {'p50 ms':>8}"
-        f" {'reactor ms':>11} {'process ms':>11}"
+        f" {'reactor ms':>11} {'process ms':>11} {'round trips':>12}"
     )
     for r in results:
         print(
             f"{r['scenario']:<16} {r['level']:<12} {r['mean_ms']:>8.1f}"
             f" {r['p50_ms']:>8.1f} {r['reactor_cpu_ms']:>11.1f}"
-            f" {r['process_cpu_ms']:>11.1f}"
+            f" {r['process_cpu_ms']:>11.1f} {r.get('round_trips', ''):>12}"
         )
     if args.json:
         with open(args.json, "w") as f:

scripts-dev/persist_bench.py

@@ -166,6 +166,7 @@ class Server:
                 "args": {
                     "database": self.database,
                     "host": os.environ.get("PGHOST", "localhost"),
+                    "port": int(os.environ.get("PGPORT", "5432")),
                     "user": os.environ.get("PGUSER", "postgres"),
                     "password": os.environ.get("PGPASSWORD", ""),
                     "cp_min": 5,

scripts-dev/pg_roundtrip_proxy.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2026 Element Creations Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+#
+
+"""A TCP proxy in front of Postgres that counts round trips and can add latency.
+
+A round trip, or "turn", is a client message that follows a server message on
+the same connection: the client had to wait for the server before sending it.
+Pipelined messages count as one turn.
+
+    poetry run python scripts-dev/pg_roundtrip_proxy.py --listen 15432 \\
+        --target postgres:5432 --delay-ms 1 --stats /tmp/stats.json
+
+Signals:
+  - SIGUSR2 resets the counters and starts recording a log of turns.
+  - SIGUSR1 writes the counters, and the log since the last reset, to the
+    `--stats` file as JSON: `turns`, `per_connection` and `log`, a list of
+    `[connection, summary]` where the summary is the protocol message types
+    and the SQL of the turn's first statement.
+
+Run Synapse against the proxy, reset before a request, dump after it. Adding a
+delay and comparing the request's wall time with and without it shows how many
+of the turns are on the request's critical path.
+"""
+
+import argparse
+import asyncio
+import json
+import os
+import signal
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Conn:
+    """What the proxy knows about one client connection."""
+
+    turns: int = 0
+    # Whether the last message seen went from the client to the server.
+    client_spoke_last: bool = False
+    # Prepared statement names to their SQL, from the Parse messages seen.
+    statements: dict[bytes, str] = field(default_factory=dict)
+    # Client bytes not yet parsed: the start of a message split across reads.
+    pending: bytes = b""
+    # Whether the startup packet, which has no type byte, has been seen.
+    started: bool = False
+
+
+class Proxy:
+    def __init__(self, target: tuple[str, int], delay: float, stats_path: str):
+        self.target = target
+        self.delay = delay
+        self.stats_path = stats_path
+        self.turns = 0
+        self.conns: dict[int, Conn] = {}
+        # The turns since the last reset, or None before the first reset.
+        self.log: list[tuple[int, str]] | None = None
+
+    def describe(self, data: bytes, conn: Conn) -> str:
+        """The frontend message types completed by `data`, and the SQL of the
+        first Parse, Bind (by statement name) or simple Query among them.
+
+        Messages can be split across reads, so a partial one is kept for the
+        next read."""
+        buf = conn.pending + data
+        kinds: list[str] = []
+        sql = ""
+        i = 0
+        while not conn.started and i + 8 <= len(buf):
+            # The startup packet, and an SSL or GSS encryption request before
+            # it, are a length then a code, with no type byte.
+            length = int.from_bytes(buf[i : i + 4], "big")
+            if i + length > len(buf):
+                break
+            code = int.from_bytes(buf[i + 4 : i + 8], "big")
+            if code not in (80877103, 80877104):
+                conn.started = True
+            i += length
+        while conn.started and i + 5 <= len(buf):
+            kind = chr(buf[i])
+            length = int.from_bytes(buf[i + 1 : i + 5], "big")
+            if i + 1 + length > len(buf):
+                break
+            body = buf[i + 5 : i + 1 + length]
+            kinds.append(kind)
+            parts = body.split(b"\0")
+            if kind == "P" and len(parts) > 1:
+                conn.statements[parts[0]] = parts[1].decode(errors="replace")
+                sql = sql or conn.statements[parts[0]]
+            elif kind == "B" and not sql and len(parts) > 1:
+                sql = conn.statements.get(parts[1], "?")
+            elif kind == "Q" and not sql:
+                sql = parts[0].decode(errors="replace")
+            i += 1 + length
+        conn.pending = buf[i:]
+        return "".join(kinds)[:16] + " " + " ".join(sql.split())[:120]
+
+    async def pipe(
+        self,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+        conn_id: int,
+        from_client: bool,
+    ) -> None:
+        conn = self.conns[conn_id]
+        try:
+            while data := await reader.read(65536):
+                if from_client:
+                    # Parse every message, to learn the statement names.
+                    # Counting must go on whatever the parser makes of it.
+                    try:
+                        summary = self.describe(data, conn)
+                    except Exception:
+                        conn.pending = b""
+                        summary = "? unparsed"
+                    if not conn.client_spoke_last:
+                        self.turns += 1
+                        conn.turns += 1
+                        if self.log is not None:
+                            self.log.append((conn_id, summary))
+                        if self.delay:
+                            await asyncio.sleep(self.delay)
+                    conn.client_spoke_last = True
+                else:
+                    conn.client_spoke_last = False
+                writer.write(data)
+                await writer.drain()
+        except ConnectionError:
+            pass
+        finally:
+            writer.close()
+
+    async def handle(
+        self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter
+    ) -> None:
+        try:
+            server_reader, server_writer = await asyncio.open_connection(*self.target)
+        except OSError:
+            client_writer.close()
+            raise
+        conn_id = len(self.conns)
+        self.conns[conn_id] = Conn()
+        await asyncio.gather(
+            self.pipe(client_reader, server_writer, conn_id, True),
+            self.pipe(server_reader, client_writer, conn_id, False),
+        )
+
+    def dump(self) -> None:
+        # Written to a temporary file and renamed, so a reader never sees it
+        # half written.
+        tmp_path = self.stats_path + ".tmp"
+        with open(tmp_path, "w") as f:
+            json.dump(
+                {
+                    "turns": self.turns,
+                    "per_connection": {
+                        conn_id: conn.turns
+                        for conn_id, conn in self.conns.items()
+                        if conn.turns
+                    },
+                    "log": self.log or [],
+                },
+                f,
+            )
+        os.replace(tmp_path, self.stats_path)
+
+    def reset(self) -> None:
+        self.turns = 0
+        for conn in self.conns.values():
+            conn.turns = 0
+        self.log = []
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--listen", type=int, required=True, help="local port")
+    parser.add_argument("--target", required=True, help="Postgres HOST:PORT")
+    parser.add_argument(
+        "--delay-ms", type=float, default=0, help="latency to add to each turn"
+    )
+    parser.add_argument("--stats", required=True, help="file SIGUSR1 writes to")
+    args = parser.parse_args()
+
+    host, _, port = args.target.rpartition(":")
+    proxy = Proxy((host, int(port)), args.delay_ms / 1000, args.stats)
+
+    loop = asyncio.get_running_loop()
+    loop.add_signal_handler(signal.SIGUSR1, proxy.dump)
+    loop.add_signal_handler(signal.SIGUSR2, proxy.reset)
+    server = await asyncio.start_server(proxy.handle, "127.0.0.1", args.listen)
+    async with server:
+        await server.serve_forever()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())