workload replay: reduce disk usage & reuse reference connections for DBs

misc/python/materialize/workload_replay/data.py

@@ -37,6 +37,7 @@
     ingest_webhook,
     make_kafka_producer,
     produce_json,
+    reuse_reference_connections,
 )
 from materialize.workload_replay.util import (
     get_kafka_topic,
@@ -345,6 +346,7 @@ def create_initial_data_external(
 
     # Lazily resolved ports — only looked up when a source of that type is seen.
     _ports: dict[str, int] = {}
+    postgres_conns: dict[str, dict[str, Any]] = {}
 
     def port(service: str) -> int:
         if service not in _ports:
@@ -384,6 +386,7 @@ def port(service: str) -> int:
                                 "password": "postgres",
                                 "dbname": ref_db,
                             }
+                            postgres_conns[ref_db] = conn
                             _submit_chunks(
                                 pool,
                                 _copy_chunk,
@@ -459,6 +462,12 @@ def port(service: str) -> int:
             )
             _await_futures(futures, totals)
 
+    # Reduces disk usage: otherwise autovacuum's WAL is retained by source slots.
+    for dbname, conn_params in postgres_conns.items():
+        print(f"Vacuuming Postgres database {dbname}")
+        with psycopg.connect(**conn_params, autocommit=True) as conn:
+            conn.execute("VACUUM (FREEZE)")
+
     for child, source, cols, num_rows, pretty_name in sequential:
         data_columns = [
             Column(
@@ -612,6 +621,7 @@ def continuous_ingestion_webhook(
                         f"({batch_size} every {period_s:.2f}s)"
                     )
 
+                    @reuse_reference_connections()
                     def continuous_ingestion_source(
                         source: dict[str, Any],
                         pretty_name: str,
@@ -697,6 +707,7 @@ def continuous_ingestion_source(
                             f"({batch_size} every {period_s:.2f}s)"
                         )
 
+                        @reuse_reference_connections()
                         def continuous_ingestion_child(
                             source: dict[str, Any],
                             child: dict[str, Any],

misc/python/materialize/workload_replay/ingest.py

@@ -17,8 +17,11 @@
 import json
 import random
 import re
+import threading
 import time
 import urllib.parse
+from collections.abc import Callable, Iterator
+from contextlib import ExitStack, contextmanager
 from functools import cache
 from textwrap import dedent
 from typing import Any
@@ -48,6 +51,47 @@
     get_sql_server_reference_db_schema_table,
 )
 
+# Per-row reconnects can overload reference database authentication.
+_reference_connections = threading.local()
+
+
+@contextmanager
+def reuse_reference_connections() -> Iterator[None]:
+    """Reuse reference connections in this thread and close them on exit."""
+    previous = getattr(_reference_connections, "scope", None)
+    with ExitStack() as stack:
+        _reference_connections.scope = (stack, {})
+        try:
+            yield
+        finally:
+            _reference_connections.scope = previous
+
+
+def _connection_open(conn: Any) -> bool:
+    return not getattr(conn, "closed", False) and getattr(conn, "open", True)
+
+
+@contextmanager
+def _reference_connection(
+    key: tuple[str, int, str], connect: Callable[[], Any]
+) -> Iterator[Any]:
+    with ExitStack() as temporary:
+        stack, connections = getattr(_reference_connections, "scope", None) or (
+            temporary,
+            {},
+        )
+        conn = connections.get(key)
+        if conn is None or not _connection_open(conn):
+            conn = connections[key] = connect()
+
+            def close() -> None:
+                # PyMySQL's close raises if a protocol error already closed the socket.
+                if _connection_open(conn):
+                    conn.close()
+
+            stack.callback(close)
+        yield conn
+
 
 def delivery_report(err: KafkaError | None, msg: Any) -> None:
     """Kafka delivery report callback."""
@@ -411,56 +455,61 @@ def ingest(
         ref_database, ref_schema, ref_table = get_postgres_reference_db_schema_table(
             child
         )
-        conn = psycopg.connect(
-            host="127.0.0.1",
-            port=c.default_port("postgres"),
-            user="postgres",
-            password="postgres",
-            dbname=ref_database,
-        )
-        conn.autocommit = True
-
-        col_names = [col.name for col in columns]
-
-        with conn.cursor() as cur:
-            copy_stmt = SQL("COPY {}.{} ({}) FROM STDIN").format(
-                Identifier(ref_schema),
-                Identifier(ref_table),
-                SQL(", ").join(map(Identifier, col_names)),
-            )
+        port = c.default_port("postgres")
+        with _reference_connection(
+            ("postgres", port, ref_database),
+            lambda: psycopg.connect(
+                host="127.0.0.1",
+                port=port,
+                user="postgres",
+                password="postgres",
+                dbname=ref_database,
+                autocommit=True,
+            ),
+        ) as conn:
+            col_names = [col.name for col in columns]
+
+            with conn.cursor() as cur:
+                copy_stmt = SQL("COPY {}.{} ({}) FROM STDIN").format(
+                    Identifier(ref_schema),
+                    Identifier(ref_table),
+                    SQL(", ").join(map(Identifier, col_names)),
+                )
 
-            with cur.copy(copy_stmt) as copy:
-                for _ in range(num_rows):
-                    row = [col.value(rng, in_query=False) for col in columns]
-                    copy.write_row(row)
+                with cur.copy(copy_stmt) as copy:
+                    for _ in range(num_rows):
+                        row = [col.value(rng, in_query=False) for col in columns]
+                        copy.write_row(row)
 
     elif source["type"] == "mysql":
         ref_database, ref_table = get_mysql_reference_db_table(child)
 
-        conn = pymysql.connect(
-            host="127.0.0.1",
-            user="root",
-            password=MySql.DEFAULT_ROOT_PASSWORD,
-            database=ref_database,
-            port=c.default_port("mysql"),
-            autocommit=False,
-        )
-
-        value_funcs = [col.value for col in columns]
-        rows_sql = []
-        for _ in range(num_rows):
-            row = [fn(rng) for fn in value_funcs]
-            rows_sql.append("(" + ", ".join(row) + ")")
-
-        stmt = f"INSERT INTO {ref_table} VALUES " + ", ".join(rows_sql)
-
-        with conn.cursor() as cur:
-            cur.execute(stmt)
-        # The connection is opened with autocommit=False, so without this commit
-        # the implicit InnoDB transaction is discarded on close and no rows are
-        # ingested.
-        conn.commit()
-        conn.close()
+        port = c.default_port("mysql")
+        with _reference_connection(
+            ("mysql", port, ref_database),
+            lambda: pymysql.connect(
+                host="127.0.0.1",
+                user="root",
+                password=MySql.DEFAULT_ROOT_PASSWORD,
+                database=ref_database,
+                port=port,
+                autocommit=False,
+            ),
+        ) as conn:
+            value_funcs = [col.value for col in columns]
+            rows_sql = []
+            for _ in range(num_rows):
+                row = [fn(rng) for fn in value_funcs]
+                rows_sql.append("(" + ", ".join(row) + ")")
+
+            stmt = f"INSERT INTO {ref_table} VALUES " + ", ".join(rows_sql)
+
+            with conn.cursor() as cur:
+                cur.execute(stmt)
+            # The connection is opened with autocommit=False, so without this commit
+            # the implicit InnoDB transaction is discarded on close and no rows are
+            # ingested.
+            conn.commit()
 
     elif source["type"] == "kafka":
         topic = get_kafka_topic(source)

misc/python/materialize/workload_replay/replay_test.py

@@ -12,8 +12,16 @@
 from __future__ import annotations
 
 import random
+import threading
+from unittest.mock import MagicMock, Mock
+
+import psycopg
+import pymysql
+import pytest
 
 from materialize.cli.mz_workload_capture import parse_mz_list
+from materialize.workload_replay import data
+from materialize.workload_replay import ingest as ingest_module
 from materialize.workload_replay.column import Column
 from materialize.workload_replay.replay import pg_params_to_psycopg
 
@@ -67,3 +75,92 @@ def test_kafka_value_is_raw_and_never_a_default_expr() -> None:
     values = {col.kafka_value(random.Random(seed)) for seed in range(200)}
     assert "'x'::text" not in values
     assert all(not v.startswith("'") for v in values if isinstance(v, str))
+
+
+@pytest.mark.parametrize("source_type", ["postgres", "mysql"])
+@pytest.mark.parametrize("fail_write", [False, True])
+def test_continuous_ingestion_reuses_and_closes_connection(
+    monkeypatch: pytest.MonkeyPatch,
+    source_type: str,
+    fail_write: bool,
+) -> None:
+    stop = threading.Event()
+    c = Mock()
+    c.default_port.return_value = 5432
+    conn = MagicMock(closed=False, open=True)
+    connect = Mock(return_value=conn)
+    monkeypatch.setattr(ingest_module.psycopg, "connect", connect)
+    monkeypatch.setattr(ingest_module.pymysql, "connect", connect)
+    batches = 0
+    write_error = (
+        psycopg.OperationalError
+        if source_type == "postgres"
+        else pymysql.err.OperationalError
+    )
+
+    def cursor():
+        nonlocal batches
+        batches += 1
+        if batches == 3:
+            stop.set()
+            if fail_write:
+                if source_type == "mysql":
+                    conn.open = False
+                    conn.close.side_effect = pymysql.err.Error("Already closed")
+                raise write_error("injected write failure")
+        return MagicMock()
+
+    conn.cursor.side_effect = cursor
+    reference = "db.s.t" if source_type == "postgres" else "db.t"
+    child = {
+        "type": "table",
+        "create_sql": f"CREATE TABLE t FROM SOURCE src (REFERENCE = {reference})",
+        "messages_second": 10000,
+        "columns": [
+            {"name": "c", "type": "integer", "nullable": False, "default": None}
+        ],
+    }
+    source = {"type": source_type, "children": {"t": child}}
+    workload = {"databases": {"db": {"s": {"sources": {"source": source}}}}}
+    threads = data.create_ingestions(
+        c, workload, stop, 1, False, {"total": 0, "slow": 0, "failed": 0}
+    )
+    thread = threads[0]
+    thread.start()
+    try:
+        if fail_write:
+            with pytest.raises(write_error, match="injected write failure"):
+                thread.join(timeout=5)
+        else:
+            thread.join(timeout=5)
+        assert not thread.is_alive()
+    finally:
+        stop.set()
+        if thread.is_alive():
+            thread.join(timeout=5)
+    assert batches == 3
+    assert connect.call_count == 1
+    if source_type == "mysql" and fail_write:
+        conn.close.assert_not_called()
+    else:
+        conn.close.assert_called_once_with()
+
+
+@pytest.mark.parametrize("source_type", ["postgres", "mysql"])
+def test_reference_connection_replaces_closed_connection(source_type: str) -> None:
+    old = MagicMock(closed=False, open=True)
+    new = MagicMock(closed=False, open=True)
+    connect = Mock(side_effect=[old, new])
+    key = (source_type, 5432, "db")
+    with ingest_module.reuse_reference_connections():
+        with ingest_module._reference_connection(key, connect) as conn:
+            assert conn is old
+        if source_type == "postgres":
+            old.closed = True
+        else:
+            old.open = False
+        with ingest_module._reference_connection(key, connect) as conn:
+            assert conn is new
+    assert connect.call_count == 2
+    old.close.assert_not_called()
+    new.close.assert_called_once_with()