193 lines
9 KiB
Python
193 lines
9 KiB
Python
#!/usr/bin/env python3
|
|
"""Проверка локальной пропускной способности записи PostgreSQL."""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import math
|
|
import os
|
|
import statistics
|
|
import time
|
|
import urllib.parse
|
|
from uuid import uuid4
|
|
|
|
import asyncpg
|
|
|
|
|
|
def local_database_url(value: str) -> str:
|
|
normalized = value.replace("postgresql+asyncpg://", "postgresql://", 1)
|
|
parsed = urllib.parse.urlsplit(normalized)
|
|
if parsed.scheme not in {"postgres", "postgresql"}:
|
|
raise ValueError("поддерживается только PostgreSQL")
|
|
if parsed.hostname not in {"127.0.0.1", "localhost", "::1"}:
|
|
raise ValueError("предохранитель: benchmark разрешён только для локальной БД")
|
|
return normalized
|
|
|
|
|
|
def percentile(values: list[float], percent: float) -> float:
|
|
ordered = sorted(values)
|
|
return ordered[max(0, math.ceil(percent / 100 * len(ordered)) - 1)]
|
|
|
|
|
|
async def run(args: argparse.Namespace) -> int:
|
|
if not 100 <= args.operations <= 10_000:
|
|
raise ValueError("operations должно быть от 100 до 10000")
|
|
if not 1 <= args.concurrency <= 100:
|
|
raise ValueError("concurrency должно быть от 1 до 100")
|
|
if args.workload == "temporary" and args.concurrency != 1:
|
|
raise ValueError("temporary workload поддерживает только concurrency=1")
|
|
database_url = local_database_url(args.database_url)
|
|
worker_count = min(args.concurrency, args.operations)
|
|
pool = (
|
|
await asyncpg.create_pool(
|
|
database_url, min_size=worker_count, max_size=worker_count,
|
|
command_timeout=30,
|
|
)
|
|
if args.workload == "audit" and worker_count > 1 else None
|
|
)
|
|
connection = await asyncpg.connect(database_url, timeout=10) if pool is None else None
|
|
latencies: list[float] = []
|
|
run_marker = None
|
|
cleanup_rows = 0
|
|
try:
|
|
if args.workload == "temporary":
|
|
assert connection is not None
|
|
async with connection.transaction():
|
|
await connection.execute("""
|
|
CREATE TEMPORARY TABLE lct_write_benchmark (
|
|
sequence_no integer PRIMARY KEY,
|
|
written_at timestamptz NOT NULL DEFAULT now(),
|
|
payload text NOT NULL
|
|
) ON COMMIT DROP
|
|
""")
|
|
started = time.perf_counter()
|
|
for sequence_no in range(args.operations):
|
|
operation_started = time.perf_counter()
|
|
await connection.execute(
|
|
"INSERT INTO lct_write_benchmark(sequence_no, payload) VALUES($1, $2)",
|
|
sequence_no, f"local-load-{sequence_no:05d}",
|
|
)
|
|
latencies.append((time.perf_counter() - operation_started) * 1000)
|
|
elapsed = time.perf_counter() - started
|
|
count = await connection.fetchval("SELECT count(*) FROM lct_write_benchmark")
|
|
else:
|
|
if not args.allow_audit_writes:
|
|
raise ValueError("для workload=audit требуется явный --allow-audit-writes")
|
|
assert pool is not None or connection is not None
|
|
fetchval = pool.fetchval if pool is not None else connection.fetchval
|
|
if await fetchval("SELECT to_regclass('public.audit_log')") is None:
|
|
raise ValueError("таблица public.audit_log отсутствует; сначала примените миграции")
|
|
# This opt-in mode measures durable product-table writes. Every
|
|
# INSERT autocommits and exercises the real audit index; exact
|
|
# run-scoped rows are removed in finally, even after a partial run.
|
|
run_marker = f"load-benchmark:{uuid4().hex}"
|
|
started = time.perf_counter()
|
|
|
|
async def write_partition(worker_index: int) -> None:
|
|
if pool is None:
|
|
assert connection is not None
|
|
worker_connection = connection
|
|
for sequence_no in range(worker_index, args.operations, worker_count):
|
|
operation_started = time.perf_counter()
|
|
await worker_connection.execute(
|
|
"""INSERT INTO public.audit_log(id, actor, role, action, object_id, detail)
|
|
VALUES($1, $2, 'system', 'benchmark.write', $3, $4)""",
|
|
uuid4(), "load-benchmark", run_marker, f"sequence={sequence_no}",
|
|
)
|
|
latencies.append((time.perf_counter() - operation_started) * 1000)
|
|
else:
|
|
async with pool.acquire() as worker_connection:
|
|
for sequence_no in range(worker_index, args.operations, worker_count):
|
|
operation_started = time.perf_counter()
|
|
await worker_connection.execute(
|
|
"""INSERT INTO public.audit_log(id, actor, role, action, object_id, detail)
|
|
VALUES($1, $2, 'system', 'benchmark.write', $3, $4)""",
|
|
uuid4(), "load-benchmark", run_marker, f"sequence={sequence_no}",
|
|
)
|
|
latencies.append((time.perf_counter() - operation_started) * 1000)
|
|
|
|
await asyncio.gather(*(write_partition(worker) for worker in range(worker_count)))
|
|
elapsed = time.perf_counter() - started
|
|
count = await fetchval(
|
|
"SELECT count(*) FROM public.audit_log WHERE object_id = $1", run_marker
|
|
)
|
|
finally:
|
|
try:
|
|
if run_marker is not None:
|
|
if pool is not None:
|
|
cleanup_rows = await pool.fetchval(
|
|
"WITH removed AS (DELETE FROM public.audit_log WHERE object_id = $1 RETURNING 1) "
|
|
"SELECT count(*) FROM removed", run_marker,
|
|
)
|
|
elif connection is not None:
|
|
cleanup_rows = await connection.fetchval(
|
|
"WITH removed AS (DELETE FROM public.audit_log WHERE object_id = $1 RETURNING 1) "
|
|
"SELECT count(*) FROM removed", run_marker,
|
|
)
|
|
finally:
|
|
if pool is not None:
|
|
await pool.close()
|
|
elif connection is not None:
|
|
await connection.close()
|
|
|
|
throughput = args.operations / elapsed
|
|
result = {
|
|
"database": "local PostgreSQL",
|
|
"workload": args.workload,
|
|
"operations_requested": args.operations,
|
|
"operations_written": count,
|
|
"cleanup_rows_deleted": cleanup_rows,
|
|
"concurrency": worker_count,
|
|
"elapsed_seconds": elapsed,
|
|
"operations_per_second": throughput,
|
|
"operation_latency_ms": {
|
|
"min": min(latencies),
|
|
"median": statistics.median(latencies),
|
|
"p95": percentile(latencies, 95),
|
|
"max": max(latencies),
|
|
},
|
|
"pass_100_writes_per_second": (
|
|
count == args.operations
|
|
and throughput >= 100
|
|
and (args.workload != "audit" or cleanup_rows == args.operations)
|
|
),
|
|
"scope": (
|
|
f"autocommitted INSERTs over {worker_count} local asyncpg connection(s) into the "
|
|
"indexed product audit_log table; this run's rows are deleted in finally; "
|
|
"not target 6-core/32-GB hardware"
|
|
if args.workload == "audit" else
|
|
"sequential INSERT statements over one local asyncpg connection into a temporary "
|
|
"table; table is dropped on commit; not target 6-core/32-GB hardware"
|
|
),
|
|
}
|
|
rendered = json.dumps(result, ensure_ascii=False, indent=2)
|
|
print(rendered)
|
|
if args.output:
|
|
with open(args.output, "w", encoding="utf-8") as stream:
|
|
stream.write(rendered + "\n")
|
|
return 0 if result["pass_100_writes_per_second"] else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--database-url",
|
|
default=os.environ.get("DATABASE_URL", "postgresql://lct:lct@127.0.0.1:5432/lct"),
|
|
)
|
|
parser.add_argument("--operations", type=int, default=1000)
|
|
parser.add_argument("--concurrency", type=int, default=1,
|
|
help="number of concurrent PostgreSQL clients (audit workload only)")
|
|
parser.add_argument(
|
|
"--workload", choices=("temporary", "audit"), default="temporary",
|
|
help="temporary benchmark table (default) or real indexed audit_log writes",
|
|
)
|
|
parser.add_argument(
|
|
"--allow-audit-writes", action="store_true",
|
|
help="required opt-in: workload=audit inserts then removes rows in audit_log",
|
|
)
|
|
parser.add_argument("--output", help="сохранить JSON-отчёт")
|
|
try:
|
|
raise SystemExit(asyncio.run(run(parser.parse_args())))
|
|
except Exception as exc:
|
|
print(f"database load failed: {type(exc).__name__}: {exc}")
|
|
raise SystemExit(2)
|