diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 0037e3dc..161cbb78 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -33,6 +33,8 @@ jobs: command: "--read-rps 1000 --write-rps 100" - name: async-query command: "--read-rps 1000 --write-rps 100" + - name: async-query-rmw + command: "--write-threads 4 --counter-count 2" - name: sync-topic command: "--write-rps 200 --write-threads 8 --read-threads 8" metrics_yaml_path: sdk-current/tests/slo/metrics-topic.yaml diff --git a/tests/slo/RMW.md b/tests/slo/RMW.md new file mode 100644 index 00000000..eebf5598 --- /dev/null +++ b/tests/slo/RMW.md @@ -0,0 +1,63 @@ +# Concurrent transactional RMW + +`async-query-rmw` is a separate SLO scenario. It uses only the asynchronous +Query service adapter (`ydb.aio.QuerySessionPool` and +`QuerySession.transaction(ydb.QuerySerializableReadWrite())`). Sync Query, +Table service and Topic are not RMW adapters. + +Use the existing image entrypoint with `WORKLOAD_NAME=async-query-rmw`, or run: + +```sh +python tests/slo/src rmw-run grpc://localhost:2136 /local \ + --time 60 --write-threads 4 --counter-count 2 \ + --otlp-endpoint http://localhost:9090/api/v1/otlp/v1/metrics +``` + +The command creates two tables named `slo_rmw__counters` +and `slo_rmw__ledger` inside the selected database, initializes +zero-valued counters, runs, validates and drops only tables whose creation +succeeded. There is no +separate startup process, create command or cross-process namespace to carry. +Current, baseline and overlapping runs use different tables even with the same +ref. The existing consuming workflow overlays the same SLO harness on the +baseline checkout while retaining each checkout's SDK and its own Dockerfile. + +Each round admits one operation per writer against a shared counter. Rounds +alternate counters. All admitted calls settle before the next round; this +prevents a fast writer from starving peers that are retrying. IDs are allocated +before the SDK retry callback, and each operation has an independently retained +counter, delta and payload. A serializable transaction reads that ID's ledger +record first. An exact replay is a no-op; a conflicting replay fails. A new +operation reads the counter, increments its value and version, inserts the +ledger record and commits atomically. Streams are consumed before commit and +each callback owns its session and transaction. + +Admission stops on duration, SIGINT, SIGTERM or a final error. It does not cancel +admitted SDK calls. After they settle, a new serializable transaction reads all +counters and ledger rows. Validation requires the exact admitted ID/counter/ +delta/payload set, no duplicate IDs, every expected counter, and exact per-counter +value and version. This detects sum-preserving lost writes and a coherently +missing operation as well as corrupted data. The first final transaction error +wins over validation, metrics and cleanup errors. + +Write/read request timeout defaults remain 20/10 seconds, with the same +idempotent retry settings and session-acquisition budgets as the existing async +Query workload. SDK retry limits, backoff and driver deadlines are not changed. +`--time` bounds admission, not settlement; in-flight calls retain their SDK +operation budgets. Existing `read`/`write` operation and retry-attempt metrics +are used unchanged. A write metric represents one logical operation. Its +technical attempts include session checkout failures before the transaction +body, using the existing public SDK retry helper with the same settings and +pooled session lifecycle. The independent final check is a read operation. + +Deterministic regression coverage: + +```sh +python -m pytest -q tests/table/test_tx.py::TestSloRmwData +``` + +The tests execute production Apply/Run/Validate through the actual SDK +retry helper and a transactional store seam. Their controlled conflicts and +post-commit error simulate retry/acknowledgement behavior; they are not proof of +a real transport-level ambiguous commit. Native database and CI receipts belong +to the release verification, not to this deterministic test claim. diff --git a/tests/slo/docker-entrypoint.sh b/tests/slo/docker-entrypoint.sh index 55d933ec..3b13e53b 100755 --- a/tests/slo/docker-entrypoint.sh +++ b/tests/slo/docker-entrypoint.sh @@ -5,7 +5,7 @@ # in parallel; both must be able to schema-prepare and then run. # # Inputs come from the env vars injected by the action: -# WORKLOAD_NAME sync-table | sync-query | async-query | sync-topic | async-topic +# WORKLOAD_NAME sync-table | sync-query | async-query | async-query-rmw | sync-topic | async-topic # WORKLOAD_DURATION run duration in seconds # YDB_ENDPOINT grpc://ydb:2136 # YDB_DATABASE /Root/testdb @@ -21,6 +21,7 @@ set -e case "${WORKLOAD_NAME:-sync-query}" in sync-table|sync-query|async-query) PREFIX=table ;; + async-query-rmw) PREFIX=rmw ;; topic|sync-topic|async-topic) PREFIX=topic ;; *) echo "Unknown WORKLOAD_NAME: ${WORKLOAD_NAME}" >&2 @@ -44,10 +45,13 @@ if [ "$PREFIX" = "topic" ]; then EXTRA_ARGS="--path ${DATABASE%/}/slo_topic_${SAFE_REF}" fi +# RMW owns its process-unique schema inside run, including cleanup. # Schema prep is idempotent at the SDK level for topics; for tables, a parallel # baseline container may race and fail with "already exists" — tolerate it. -python ./tests/slo/src "${PREFIX}-create" "$ENDPOINT" "$DATABASE" $EXTRA_ARGS \ - || echo "WARN: ${PREFIX}-create exited non-zero (treated as already-prepared)" >&2 +if [ "$PREFIX" != "rmw" ]; then + python ./tests/slo/src "${PREFIX}-create" "$ENDPOINT" "$DATABASE" $EXTRA_ARGS \ + || echo "WARN: ${PREFIX}-create exited non-zero (treated as already-prepared)" >&2 +fi exec python ./tests/slo/src \ "${PREFIX}-run" "$ENDPOINT" "$DATABASE" \ diff --git a/tests/slo/src/jobs/rmw_jobs.py b/tests/slo/src/jobs/rmw_jobs.py new file mode 100644 index 00000000..dc290074 --- /dev/null +++ b/tests/slo/src/jobs/rmw_jobs.py @@ -0,0 +1,228 @@ +import asyncio +import logging +from dataclasses import dataclass + +import ydb + +from core.metrics import OP_TYPE_READ, OP_TYPE_WRITE + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class Operation: + op_id: str + counter_id: int + delta: int + + def row(self): + return dict( + op_id=self.op_id, + counter_id=self.counter_id, + delta=self.delta, + payload=f"{self.op_id}:{self.counter_id}:{self.delta}", + ) + + +async def read_rows(tx, query, settings, parameters=None): + rows = [] + async with await tx.execute(query, parameters=parameters, settings=settings) as stream: + async for part in stream: + rows.extend(part.rows) + return rows + + +class AsyncRmwJobManager: + def __init__(self, pool, args, metrics, counters, ledger): + self.pool, self.args, self.metrics = pool, args, metrics + self.counters, self.ledger = counters, ledger + self.stop = asyncio.Event() + self.expected = {} + self.failure = None + self.diagnostic_failure = None + self.attempts = 0 + self.replays = 0 + self.write_settings = ydb.BaseRequestSettings().with_timeout(args.write_timeout / 1000) + self.read_settings = ydb.BaseRequestSettings().with_timeout(args.read_timeout / 1000) + self.retry_settings = ydb.RetrySettings(idempotent=True, max_session_acquire_timeout=args.write_timeout / 1000) + + async def apply(self, operation): + attempts = 0 + error = None + labels = (OP_TYPE_WRITE,) + started = self.metrics.start(labels) + + async def transaction(): + nonlocal attempts + attempts += 1 + async with self.pool.checkout(timeout=self.retry_settings.max_session_acquire_timeout) as session: + async with session.transaction(ydb.QuerySerializableReadWrite()) as tx: + rows = await read_rows( + tx, + f"SELECT op_id, counter_id, delta, payload FROM `{self.ledger}` WHERE op_id = $op_id;", + self.write_settings, + {"$op_id": (operation.op_id, ydb.PrimitiveType.Utf8)}, + ) + if rows: + if rows != [operation.row()]: + raise RuntimeError(f"RMW replay mismatch: {operation.op_id}: {rows!r}") + await tx.commit(settings=self.write_settings) + self.replays += 1 + return + rows = await read_rows( + tx, + f"SELECT counter_id, value, version FROM `{self.counters}` WHERE counter_id = $counter_id;", + self.write_settings, + {"$counter_id": (operation.counter_id, ydb.PrimitiveType.Uint64)}, + ) + if len(rows) != 1 or rows[0]["counter_id"] != operation.counter_id: + raise RuntimeError(f"RMW counter missing: {operation.counter_id}: {rows!r}") + values = dict( + operation.row(), value=rows[0]["value"] + operation.delta, version=rows[0]["version"] + 1 + ) + parameters = { + f"${key}": ( + value, + ydb.PrimitiveType.Utf8 if isinstance(value, str) else ydb.PrimitiveType.Uint64, + ) + for key, value in values.items() + } + await read_rows( + tx, + f"UPSERT INTO `{self.counters}` (counter_id, value, version) VALUES ($counter_id, $value, $version);" + f" INSERT INTO `{self.ledger}` (op_id, counter_id, delta, payload)" + " VALUES ($op_id, $counter_id, $delta, $payload);", + self.write_settings, + parameters, + ) + await tx.commit(settings=self.write_settings) + + try: + await ydb.retry_operation_async(transaction, self.retry_settings) + except BaseException as exc: + error = exc + if self.failure is None: + self.failure = exc + self.stop.set() + raise + finally: + self.attempts += attempts + self._finish_metric(labels, started, attempts, error) + + def _finish_metric(self, labels, started, attempts, error): + try: + self.metrics.stop(labels, started, attempts=attempts, error=error) + except BaseException: + if error is None: + raise + logger.exception("Cannot finish RMW metrics after operation failure") + + async def validate(self): + async def snapshot(session): + async with session.transaction(ydb.QuerySerializableReadWrite()) as tx: + counters = await read_rows( + tx, f"SELECT counter_id, value, version FROM `{self.counters}`;", self.read_settings + ) + ledger = await read_rows( + tx, f"SELECT op_id, counter_id, delta, payload FROM `{self.ledger}`;", self.read_settings + ) + await tx.commit(settings=self.read_settings) + return counters, ledger + + error = None + started = self.metrics.start((OP_TYPE_READ,)) + attempts = 0 + + retry_settings = ydb.RetrySettings(idempotent=True, max_session_acquire_timeout=self.args.read_timeout / 1000) + + async def read(): + nonlocal attempts + attempts += 1 + async with self.pool.checkout(timeout=retry_settings.max_session_acquire_timeout) as session: + return await snapshot(session) + + try: + counters, ledger = await ydb.retry_operation_async(read, retry_settings) + observed = {} + for row in ledger: + op_id = row["op_id"] + if op_id in observed: + raise RuntimeError(f"RMW duplicate ledger ID: {op_id}") + observed[op_id] = dict(row) + expected = {key: operation.row() for key, operation in self.expected.items()} + if observed != expected: + raise RuntimeError( + f"RMW ledger mismatch: expected {len(expected)} exact operations, observed {len(observed)}" + ) + totals = {key: dict(counter_id=key, value=0, version=0) for key in range(self.args.counter_count)} + for operation in self.expected.values(): + totals[operation.counter_id]["value"] += operation.delta + totals[operation.counter_id]["version"] += 1 + if len(counters) != len(totals) or {row["counter_id"]: dict(row) for row in counters} != totals: + raise RuntimeError(f"RMW counter mismatch: expected {totals!r}, observed {counters!r}") + logger.info( + "Validated RMW counters=%s operations=%s attempts=%s replays=%s tables=%s,%s", + len(totals), + len(expected), + self.attempts, + self.replays, + self.counters, + self.ledger, + ) + except BaseException as exc: + error = exc + raise + finally: + self._finish_metric((OP_TYPE_READ,), started, attempts, error) + + async def _apply_operation(self, operation): + try: + await self.apply(operation) + except BaseException as exc: + if self.diagnostic_failure is None: + self.diagnostic_failure = exc + self.stop.set() + + async def _metrics_sender(self): + try: + while not self.stop.is_set(): + try: + await asyncio.wait_for(self.stop.wait(), self.args.report_period / 1000) + except asyncio.TimeoutError: + await asyncio.get_running_loop().run_in_executor(None, self.metrics.push) + except BaseException as exc: + self.stop.set() + return exc + return None + + async def run_tests(self): + timer = asyncio.get_running_loop().call_later(self.args.time, self.stop.set) + sender = asyncio.create_task(self._metrics_sender()) + sequence = 0 + try: + while not self.stop.is_set(): + sequence += 1 + operations = [ + Operation(f"{writer}:{sequence}", (sequence - 1) % self.args.counter_count, 1 + writer % 3) + for writer in range(self.args.write_threads) + ] + self.expected.update((operation.op_id, operation) for operation in operations) + # Bound admission to one contended round; never cancel an admitted SDK call. + await asyncio.gather(*(self._apply_operation(operation) for operation in operations)) + finally: + self.stop.set() + timer.cancel() + metric_failure = await sender + try: + await self.validate() + except BaseException: + if self.failure is None and self.diagnostic_failure is None and metric_failure is None: + raise + logger.exception("RMW final validation failed after workload failure") + if self.failure is not None: + raise self.failure + if self.diagnostic_failure is not None: + raise self.diagnostic_failure + if metric_failure is not None: + raise metric_failure + self.metrics.push() diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py index 200c3c55..66cb18e0 100644 --- a/tests/slo/src/options.py +++ b/tests/slo/src/options.py @@ -97,6 +97,33 @@ def make_table_cleanup_parser(subparsers): add_common_options(table_cleanup_parser) +def make_rmw_run_parser(subparsers): + parser = subparsers.add_parser("rmw-run", help="Run concurrent transactional RMW using async Query service") + parser.add_argument("endpoint", help="YDB endpoint") + parser.add_argument("db", help="YDB database name") + parser.add_argument("--debug", action="store_true", help="Enable debug logging") + + def positive(value): + result = int(value) + if result < 1: + raise argparse.ArgumentTypeError("must be positive") + return result + + parser.add_argument( + "--time", default=10, type=positive, help="Admission duration [s]; admitted calls settle afterwards" + ) + parser.add_argument("--write-threads", default=4, type=positive, help="Concurrent writers sharing counters") + parser.add_argument("--counter-count", default=2, type=positive, help="Contended counters") + parser.add_argument("--write-timeout", default=20000, type=positive, help="Write request timeout [ms]") + parser.add_argument("--read-timeout", default=10000, type=positive, help="Validation request timeout [ms]") + parser.add_argument("--report-period", default=1000, type=positive, help="Metrics push period [ms]") + parser.add_argument( + "--otlp-endpoint", + default="http://localhost:9090/api/v1/otlp/v1/metrics", + help="OTLP metrics endpoint; empty to disable", + ) + + def make_topic_create_parser(subparsers): topic_create_parser = subparsers.add_parser("topic-create", help="Create topic with consumer") add_common_options(topic_create_parser) @@ -169,6 +196,7 @@ def get_root_parser(): make_table_create_parser(subparsers) make_table_run_parser(subparsers) make_table_cleanup_parser(subparsers) + make_rmw_run_parser(subparsers) make_topic_create_parser(subparsers) make_topic_run_parser(subparsers) diff --git a/tests/slo/src/root_runner.py b/tests/slo/src/root_runner.py index fb7117c4..76883d1e 100644 --- a/tests/slo/src/root_runner.py +++ b/tests/slo/src/root_runner.py @@ -7,6 +7,7 @@ from core.metrics import WORKLOAD from runners.topic_runner import TopicRunner from runners.table_runner import TableRunner +from runners.rmw_runner import RmwRunner from runners.base import BaseRunner logger = logging.getLogger(__name__) @@ -14,8 +15,8 @@ def _is_async_workload(args) -> bool: """Async mode is driven by the workload label (``async-*``); the ``--async`` - CLI flag is kept as a manual override.""" - return getattr(args, "async", False) or str(WORKLOAD).startswith("async-") + CLI flag is kept as a manual override. RMW has only an async Query adapter.""" + return args.subcommand == "rmw-run" or getattr(args, "async", False) or str(WORKLOAD).startswith("async-") class SLORunner: @@ -108,6 +109,7 @@ def create_runner() -> SLORunner: runner = SLORunner() runner.register_runner("table", TableRunner) runner.register_runner("topic", TopicRunner) + runner.register_runner("rmw", RmwRunner) return runner diff --git a/tests/slo/src/runners/rmw_runner.py b/tests/slo/src/runners/rmw_runner.py new file mode 100644 index 00000000..faa9fe3d --- /dev/null +++ b/tests/slo/src/runners/rmw_runner.py @@ -0,0 +1,75 @@ +import asyncio +import signal +from uuid import uuid4 + +import ydb.aio + +from core.metrics import create_metrics +from jobs.rmw_jobs import AsyncRmwJobManager +from .base import BaseRunner + + +class RmwRunner(BaseRunner): + @property + def prefix(self): + return "rmw" + + def create(self, args): + raise ValueError("rmw-run owns its schema lifecycle") + + def cleanup(self, args): + raise ValueError("rmw-run owns its schema lifecycle") + + def run(self, args): + raise ValueError("RMW supports only the async Query adapter") + + async def run_async(self, args): + metrics = create_metrics(args.otlp_endpoint) + namespace = f"{args.db.rstrip('/')}/slo_rmw_{uuid4().hex}" + counters, ledger = f"{namespace}_counters", f"{namespace}_ledger" + self.logger.info("Creating RMW tables: %s, %s", counters, ledger) + failure = None + try: + async with ydb.aio.QuerySessionPool(self.driver) as pool: + owned = [] + handlers = {} + loop = asyncio.get_running_loop() + try: + for table, schema in ( + (counters, "counter_id Uint64, value Uint64, version Uint64, PRIMARY KEY(counter_id)"), + (ledger, "op_id Utf8, counter_id Uint64, delta Uint64, payload Utf8, PRIMARY KEY(op_id)"), + ): + await pool.execute_with_retries(f"CREATE TABLE `{table}` ({schema});") + owned.append(table) + rows = ", ".join(f"({key}u, 0u, 0u)" for key in range(args.counter_count)) + await pool.execute_with_retries( + f"UPSERT INTO `{counters}` (counter_id, value, version) VALUES {rows};" + ) + manager = AsyncRmwJobManager(pool, args, metrics, counters, ledger) + for sig in (signal.SIGINT, signal.SIGTERM): + handlers[sig] = signal.getsignal(sig) + loop.add_signal_handler(sig, manager.stop.set) + self.logger.info("Starting RMW writers=%s counters=%s", args.write_threads, args.counter_count) + await manager.run_tests() + metrics.reset() + except BaseException as exc: + failure = exc + raise + finally: + for sig, handler in handlers.items(): + loop.remove_signal_handler(sig) + signal.signal(sig, handler) + for table in reversed(owned): + try: + await pool.execute_with_retries(f"DROP TABLE `{table}`;") + self.logger.info("RMW table dropped: %s", table) + except BaseException as exc: + self.logger.exception("Cannot drop RMW table %s", table) + if failure is None: + failure = exc + if failure is not None: + raise failure + except BaseException as exc: + if failure is not None and exc is not failure: + raise failure from exc + raise diff --git a/tests/table/test_tx.py b/tests/table/test_tx.py index ffdbc4eb..e01893e5 100644 --- a/tests/table/test_tx.py +++ b/tests/table/test_tx.py @@ -1,4 +1,11 @@ +import asyncio +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + import pytest +import ydb.aio from contextlib import suppress import ydb.iam @@ -379,3 +386,385 @@ def owner(): future.result(timeout=10) assert len(tables) == len(set(tables)) == 2 assert sorted(call.args[0] for call in session.drop_table.call_args_list) == sorted(tables) + + +def _rmw_manager(rmw, store): + args = SimpleNamespace( + write_timeout=20000, read_timeout=10000, counter_count=2, time=3600, write_threads=4, report_period=1000 + ) + return rmw.AsyncRmwJobManager(store, args, Mock(), "/local/owned_counters", "/local/owned_ledger") + + +class RmwStore: + # A serializable transaction seam, not evidence of real SDK/server recovery. + def __init__(self): + self.counters = {key: dict(counter_id=key, value=0, version=0) for key in range(2)} + self.ledger = {} + self.revision = 0 + self.commits = 0 + self.before_commit = None + self.after_commit = None + self.on_counter_read = None + self.extra_ledger = [] + self.extra_counters = [] + self.read_failure = None + self.sessions = [] + + def checkout(self, timeout=None): + store = self + + class Session: + async def __aenter__(self): + store.sessions.append(self) + return self + + async def __aexit__(self, *args): + pass + + def transaction(self, mode): + assert isinstance(mode, ydb.QuerySerializableReadWrite) + return RmwTransaction(store) + + return Session() + + +class RmwTransaction: + def __init__(self, store): + self.store = store + self.counters = deepcopy(store.counters) + self.ledger = deepcopy(store.ledger) + self.revision = store.revision + self.wrote = False + self.operation = None + self.finished = False + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + async def execute(self, query, parameters=None, settings=None): + assert not self.finished + assert settings.timeout in (10, 20) + parameters = {key[1:]: value[0] for key, value in (parameters or {}).items()} + if query.startswith("UPSERT"): + self.wrote = True + self.operation = parameters["op_id"] + self.counters[parameters["counter_id"]] = { + key: parameters[key] for key in ("counter_id", "value", "version") + } + assert parameters["op_id"] not in self.ledger + self.ledger[parameters["op_id"]] = { + key: parameters[key] for key in ("op_id", "counter_id", "delta", "payload") + } + rows = [] + elif "WHERE op_id" in query: + row = self.ledger.get(parameters["op_id"]) + rows = [row] if row is not None else [] + elif "WHERE counter_id" in query: + row = self.counters.get(parameters["counter_id"]) + rows = [row] if row is not None else [] + if self.store.on_counter_read: + await self.store.on_counter_read() + else: + if self.store.read_failure: + raise self.store.read_failure + if "counter_id, value, version" in query: + rows = list(self.counters.values()) + self.store.extra_counters + else: + rows = list(self.ledger.values()) + self.store.extra_ledger + + class Stream: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + def __aiter__(self): + async def parts(): + yield SimpleNamespace(rows=deepcopy(rows)) + + return parts() + + return Stream() + + async def commit(self, settings=None): + if self.wrote: + if self.store.before_commit: + await self.store.before_commit(self) + if self.revision != self.store.revision: + raise ydb.Aborted("serializable conflict") + self.store.counters, self.store.ledger = self.counters, self.ledger + self.store.revision += 1 + self.store.commits += 1 + self.finished = True + if self.wrote and self.store.after_commit: + await self.store.after_commit(self) + + +class TestSloRmwData: + @pytest.fixture + def rmw(self, monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "tests/slo/src")) + from jobs import rmw_jobs + + return rmw_jobs + + @pytest.fixture + def store(self): + return RmwStore() + + @pytest.mark.asyncio + @pytest.mark.parametrize("ambiguous", [False, True]) + async def test_apply_uses_same_id_after_sdk_retry_and_deduplicates(self, rmw, store, ambiguous): + workload = _rmw_manager(rmw, store) + operation = rmw.Operation("writer:1", 0, 3) + workload.expected[operation.op_id] = operation + injected = False + + async def failure(tx): + nonlocal injected + if not injected: + injected = True + raise ydb.BadSession("synthetic lost commit acknowledgement") if ambiguous else ydb.Aborted("conflict") + + if ambiguous: + store.after_commit = failure + else: + store.before_commit = failure + await workload.apply(operation) + await workload.apply(operation) + await workload.validate() + assert store.commits == 1 + assert store.counters[0] == dict(counter_id=0, value=3, version=1) + assert store.ledger == {operation.op_id: operation.row()} + assert workload.attempts == 3 + assert workload.replays == (2 if ambiguous else 1) + assert len({id(session) for session in store.sessions}) == len(store.sessions) + assert workload.retry_settings.max_retries == ydb.RetrySettings().max_retries + assert workload.retry_settings.idempotent + assert workload.retry_settings.max_session_acquire_timeout == 20 + assert workload.metrics.stop.call_args.kwargs["attempts"] == 1 + + @pytest.mark.asyncio + async def test_run_stops_admission_and_settles_contending_calls_before_validate(self, rmw, store): + workload = _rmw_manager(rmw, store) + admitted = 0 + barrier = asyncio.Event() + + async def first_reads(): + nonlocal admitted + if admitted < workload.args.write_threads: + admitted += 1 + if admitted == workload.args.write_threads: + workload.stop.set() + barrier.set() + await asyncio.wait_for(barrier.wait(), 5) + + store.on_counter_read = first_reads + await asyncio.wait_for(workload.run_tests(), 5) + assert len(workload.expected) == len(store.ledger) == store.commits == 4 + assert workload.attempts > store.commits + assert store.counters[0] == dict(counter_id=0, value=7, version=4) + assert store.counters[1] == dict(counter_id=1, value=0, version=0) + assert workload.metrics.stop.call_args.args[0] == ("read",) + + @pytest.mark.asyncio + async def test_original_final_error_wins_and_other_admitted_call_settles(self, rmw, store): + workload = _rmw_manager(rmw, store) + workload.args.write_threads = 2 + both = asyncio.Event() + terminal = ydb.BadRequest("original final SDK error") + started = 0 + + async def first_reads(): + nonlocal started + started += 1 + if started == 2: + both.set() + await asyncio.wait_for(both.wait(), 5) + + async def fail_one(tx): + if tx.operation.startswith("0:"): + store.read_failure = ydb.BadRequest("diagnostic failure") + raise terminal + await asyncio.wait_for(workload.stop.wait(), 5) + + def metric_stop(*args, **kwargs): + if kwargs["error"] is terminal: + raise RuntimeError("metric diagnostic failure") + + workload.metrics.stop.side_effect = metric_stop + store.on_counter_read = first_reads + store.before_commit = fail_one + with pytest.raises(ydb.BadRequest) as failure: + await asyncio.wait_for(workload.run_tests(), 5) + assert failure.value is terminal + assert store.commits == 1 + assert set(store.ledger) == {"1:1"} + assert set(workload.expected) == {"0:1", "1:1"} + + @pytest.mark.asyncio + async def test_sdk_failure_wins_over_an_earlier_success_metric_failure(self, rmw, store): + workload = _rmw_manager(rmw, store) + workload.args.write_threads = 2 + both = asyncio.Event() + terminal = ydb.BadRequest("original final SDK error") + started = 0 + + async def first_reads(): + nonlocal started + started += 1 + if started == 2: + both.set() + await asyncio.wait_for(both.wait(), 5) + + async def fail_second(tx): + if tx.operation.startswith("1:"): + await asyncio.wait_for(workload.stop.wait(), 5) + raise terminal + + def metric_stop(labels, *args, **kwargs): + if labels == ("write",) and kwargs["error"] is None: + raise RuntimeError("earlier metrics failure") + + workload.metrics.stop.side_effect = metric_stop + store.on_counter_read = first_reads + store.before_commit = fail_second + with pytest.raises(ydb.BadRequest) as failure: + await asyncio.wait_for(workload.run_tests(), 5) + assert failure.value is terminal + assert store.commits == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "corruption", + [ + "sum_preserving", + "coherent_missing", + "duplicate", + "payload", + "version", + "counter_id", + "delta", + "extra", + "missing_counter", + "duplicate_counter", + ], + ) + async def test_independent_validation_rejects_exact_data_corruption(self, rmw, store, corruption): + workload = _rmw_manager(rmw, store) + for key in range(2): + operation = rmw.Operation(f"writer:{key}", key, 2) + workload.expected[operation.op_id] = operation + await workload.apply(operation) + if corruption == "sum_preserving": + store.counters[0]["value"] -= 1 + store.counters[1]["value"] += 1 + elif corruption == "coherent_missing": + del store.ledger["writer:0"] + store.counters[0].update(value=0, version=0) + elif corruption == "duplicate": + store.extra_ledger = [store.ledger["writer:0"]] + elif corruption in ("payload", "counter_id", "delta"): + store.ledger["writer:0"][corruption] = "bad" if corruption == "payload" else 17 + elif corruption == "version": + store.counters[0]["version"] += 1 + elif corruption == "extra": + store.ledger["extra"] = rmw.Operation("extra", 0, 0).row() + elif corruption == "missing_counter": + del store.counters[0] + elif corruption == "duplicate_counter": + store.extra_counters = [store.counters[0]] + with pytest.raises(RuntimeError): + await workload.validate() + + @pytest.mark.asyncio + async def test_apply_rejects_nonidentical_replay_without_an_increment(self, rmw, store): + workload = _rmw_manager(rmw, store) + operation = rmw.Operation("same-id", 0, 1) + await workload.apply(operation) + with pytest.raises(RuntimeError): + await workload.apply(rmw.Operation("same-id", 1, 2)) + assert store.commits == 1 + assert store.counters[1]["value"] == 0 + + def test_rmw_cli_uses_existing_async_lifecycle_and_rejects_invalid_inputs(self, rmw): + from options import get_root_parser + from root_runner import _is_async_workload, create_runner + from runners.rmw_runner import RmwRunner + + parser = get_root_parser() + args = parser.parse_args(["rmw-run", "grpc://localhost:13936", "/local"]) + assert _is_async_workload(args) + assert create_runner().runners["rmw"] is RmwRunner + for option in ( + "--time", + "--write-threads", + "--counter-count", + "--write-timeout", + "--read-timeout", + "--report-period", + ): + with pytest.raises(SystemExit) as failure: + parser.parse_args(["rmw-run", "grpc://localhost:13936", "/local", option, "0"]) + assert failure.value.code == 2 + + @pytest.mark.asyncio + async def test_counts_sdk_checkout_retry_before_transaction_body(self, rmw, store, monkeypatch): + workload = _rmw_manager(rmw, store) + operation = rmw.Operation("checkout:1", 0, 1) + workload.expected[operation.op_id] = operation + original = store.checkout + acquisitions = 0 + + class FailedCheckout: + async def __aenter__(self): + raise ydb.Aborted("checkout failed before the transaction callback") + + async def __aexit__(self, *args): + raise AssertionError("failed checkout was never entered") + + def checkout(timeout=None): + nonlocal acquisitions + acquisitions += 1 + return FailedCheckout() if acquisitions == 1 else original(timeout) + + monkeypatch.setattr(store, "checkout", checkout) + await workload.apply(operation) + assert acquisitions == 2 + assert workload.attempts == 2 + assert workload.metrics.stop.call_args.kwargs["attempts"] == 2 + assert store.counters[0] == dict(counter_id=0, value=1, version=1) + assert store.ledger == {operation.op_id: operation.row()} + + @pytest.mark.asyncio + @pytest.mark.parametrize("failed_create", [1, 2]) + async def test_setup_failure_drops_only_successfully_created_tables(self, rmw, monkeypatch, failed_create): + from runners import rmw_runner + + failure = ydb.BadRequest("CREATE failed without ownership") + statements = [] + + class Pool: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + async def execute_with_retries(self, query): + statements.append(query) + if query.startswith("CREATE") and len(statements) == failed_create: + raise failure + + monkeypatch.setattr(rmw_runner, "create_metrics", lambda _: Mock()) + monkeypatch.setattr(rmw_runner.ydb.aio, "QuerySessionPool", lambda _: Pool()) + with pytest.raises(ydb.BadRequest) as caught: + await rmw_runner.RmwRunner().run_async(SimpleNamespace(otlp_endpoint="", db="/local")) + assert caught.value is failure + drops = [query for query in statements if query.startswith("DROP")] + owned = statements[0].split("`")[1] + assert drops == ([f"DROP TABLE `{owned}`;"] if failed_create == 2 else [])