diff --git a/misc/python/materialize/workload_replay/data.py b/misc/python/materialize/workload_replay/data.py index 9eb153e592a27..bb0fa45060e7a 100644 --- a/misc/python/materialize/workload_replay/data.py +++ b/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], diff --git a/misc/python/materialize/workload_replay/ingest.py b/misc/python/materialize/workload_replay/ingest.py index f2afc9fb4a142..62a19f3b8b648 100644 --- a/misc/python/materialize/workload_replay/ingest.py +++ b/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) diff --git a/misc/python/materialize/workload_replay/replay_test.py b/misc/python/materialize/workload_replay/replay_test.py index d3e3441fd065c..1b80129726cb3 100644 --- a/misc/python/materialize/workload_replay/replay_test.py +++ b/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()