Add concurrent transactional Query RMW SLO workload

.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

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_<process-unique UUID>_counters`
+and `slo_rmw_<same UUID>_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.

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" \

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()

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)