diff --git a/scripts-dev/createroom_bench.py b/scripts-dev/createroom_bench.py index 05bed3d4a8..c0a58307e4 100755 --- a/scripts-dev/createroom_bench.py +++ b/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: diff --git a/scripts-dev/persist_bench.py b/scripts-dev/persist_bench.py index 46c06cd207..eb5732e7f8 100755 --- a/scripts-dev/persist_bench.py +++ b/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, diff --git a/scripts-dev/pg_roundtrip_proxy.py b/scripts-dev/pg_roundtrip_proxy.py new file mode 100755 index 0000000000..f01031bfc5 --- /dev/null +++ b/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: +# . +# + +"""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())