From 631903e5c2f0df9d8b55bcb037dd3c691a28edc2 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Wed, 17 Jun 2026 12:44:58 +0300 Subject: [PATCH 01/17] Add SLO checks with a SQLAlchemy read/write workload Introduce a parallel read/write SLO workload built on the ydb_sqlalchemy dialect (SQLAlchemy Core and ORM modes) and wire it into ydb-slo-action via a label-gated GitHub workflow. - tests/slo: workload runner, Dockerfile, entrypoint, requirements, README - .github/workflows/slo.yml: build current+baseline images, run init@v2 and publish report@v2 on PRs labelled "SLO" --- .dockerignore | 13 ++ .github/workflows/slo.yml | 137 +++++++++++++++ tests/slo/Dockerfile | 39 +++++ tests/slo/README.md | 78 +++++++++ tests/slo/docker-entrypoint.sh | 34 ++++ tests/slo/requirements.txt | 9 + tests/slo/src/__main__.py | 12 ++ tests/slo/src/generator.py | 53 ++++++ tests/slo/src/metrics.py | 237 ++++++++++++++++++++++++++ tests/slo/src/models.py | 94 +++++++++++ tests/slo/src/options.py | 67 ++++++++ tests/slo/src/workload.py | 299 +++++++++++++++++++++++++++++++++ 12 files changed, 1072 insertions(+) create mode 100644 .dockerignore create mode 100644 .github/workflows/slo.yml create mode 100644 tests/slo/Dockerfile create mode 100644 tests/slo/README.md create mode 100644 tests/slo/docker-entrypoint.sh create mode 100644 tests/slo/requirements.txt create mode 100644 tests/slo/src/__main__.py create mode 100644 tests/slo/src/generator.py create mode 100644 tests/slo/src/metrics.py create mode 100644 tests/slo/src/models.py create mode 100644 tests/slo/src/options.py create mode 100644 tests/slo/src/workload.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..56ef81c --- /dev/null +++ b/.dockerignore @@ -0,0 +1,13 @@ +.git +.venv +venv +dist +build +*.egg-info +**/__pycache__ +**/*.pyc +.pytest_cache +docs/.build +.vscode +.idea +.DS_Store diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml new file mode 100644 index 0000000..28111fe --- /dev/null +++ b/.github/workflows/slo.yml @@ -0,0 +1,137 @@ +name: SLO + +on: + pull_request: + types: [opened, reopened, synchronize, labeled] + +permissions: + contents: read + pull-requests: write + checks: write + +jobs: + workload: + if: contains(github.event.pull_request.labels.*.name, 'SLO') + + name: SLO workload (${{ matrix.workload.name }}) + runs-on: ubuntu-latest + + strategy: + fail-fast: false + matrix: + workload: + - name: core + command: "--read-rps 500 --write-rps 50" + - name: orm + command: "--read-rps 500 --write-rps 50" + + concurrency: + group: slo-${{ github.ref }}-${{ matrix.workload.name }} + cancel-in-progress: true + + steps: + - name: Install tooling + run: | + set -euxo pipefail + YQ_VERSION=v4.48.2 + BUILDX_VERSION=0.30.1 + COMPOSE_VERSION=2.40.3 + + sudo curl -fLo /usr/local/bin/yq \ + "https://github.com/mikefarah/yq/releases/download/${YQ_VERSION}/yq_linux_amd64" + sudo chmod +x /usr/local/bin/yq + + sudo mkdir -p /usr/local/lib/docker/cli-plugins + + sudo curl -fLo /usr/local/lib/docker/cli-plugins/docker-buildx \ + "https://github.com/docker/buildx/releases/download/v${BUILDX_VERSION}/buildx-v${BUILDX_VERSION}.linux-amd64" + sudo chmod +x /usr/local/lib/docker/cli-plugins/docker-buildx + + sudo curl -fLo /usr/local/lib/docker/cli-plugins/docker-compose \ + "https://github.com/docker/compose/releases/download/v${COMPOSE_VERSION}/docker-compose-linux-x86_64" + sudo chmod +x /usr/local/lib/docker/cli-plugins/docker-compose + + yq --version + docker --version + docker buildx version + docker compose version + + - name: Checkout current dialect version + uses: actions/checkout@v5 + with: + path: current + fetch-depth: 0 + + - name: Determine baseline commit + id: baseline + working-directory: current + run: | + set -euo pipefail + BASELINE=$(git merge-base HEAD origin/main) + echo "sha=${BASELINE}" >> "$GITHUB_OUTPUT" + if [ "$(git rev-parse origin/main)" = "${BASELINE}" ]; then + echo "ref=main" >> "$GITHUB_OUTPUT" + else + echo "ref=${BASELINE:0:7}" >> "$GITHUB_OUTPUT" + fi + + - name: Checkout baseline dialect version + uses: actions/checkout@v5 + with: + ref: ${{ steps.baseline.outputs.sha }} + path: baseline + fetch-depth: 1 + + - name: Build workload images (current + baseline) + run: | + set -euxo pipefail + + # Current: PR dialect + PR workload runner. + docker build \ + -f "$GITHUB_WORKSPACE/current/tests/slo/Dockerfile" \ + -t "ydb-slo-current" \ + "$GITHUB_WORKSPACE/current" + + # Baseline: baseline dialect with the CURRENT workload runner so the + # runner-side contract (entrypoint, metric names) is identical for both. + rm -rf "$GITHUB_WORKSPACE/baseline/tests/slo" + cp -r "$GITHUB_WORKSPACE/current/tests/slo" "$GITHUB_WORKSPACE/baseline/tests/slo" + docker build \ + -f "$GITHUB_WORKSPACE/baseline/tests/slo/Dockerfile" \ + -t "ydb-slo-baseline" \ + "$GITHUB_WORKSPACE/baseline" + + - name: Run SLO workload + uses: ydb-platform/ydb-slo-action/init@v2 + timeout-minutes: 30 + with: + github_issue: ${{ github.event.pull_request.number }} + github_token: ${{ secrets.GITHUB_TOKEN }} + workload_name: ${{ matrix.workload.name }} + workload_duration: "120" + workload_current_ref: ${{ github.head_ref || github.ref_name }} + workload_current_image: ydb-slo-current + workload_current_command: ${{ matrix.workload.command }} + workload_baseline_ref: ${{ steps.baseline.outputs.ref }} + workload_baseline_image: ydb-slo-baseline + workload_baseline_command: ${{ matrix.workload.command }} + # Trim the cluster (drop database-3/4/5) to fit a GitHub-hosted runner. + disable_compose_profiles: extra-nodes + + report: + # Runs in the same workflow run as `workload`, so it sees the freshly + # uploaded metrics artifacts and can compare current vs baseline and gate + # the PR. Same-repo PRs have the pull-requests:write token this needs. + if: ${{ always() && contains(github.event.pull_request.labels.*.name, 'SLO') }} + + name: SLO report + needs: workload + runs-on: ubuntu-latest + + steps: + - name: Publish SLO report + uses: ydb-platform/ydb-slo-action/report@v2 + with: + github_issue: ${{ github.event.pull_request.number }} + github_token: ${{ secrets.GITHUB_TOKEN }} + github_run_id: ${{ github.run_id }} diff --git a/tests/slo/Dockerfile b/tests/slo/Dockerfile new file mode 100644 index 0000000..5276efd --- /dev/null +++ b/tests/slo/Dockerfile @@ -0,0 +1,39 @@ +# syntax=docker/dockerfile:1 +# +# Packages the SQLAlchemy SLO workload runner for ydb-slo-action v2. +# +# Build context is the repository root. The action builds this image twice — for +# the `current` (PR) and `baseline` (merge-base) versions of the dialect — and +# runs both in parallel against the same YDB cluster, distinguishing them by the +# WORKLOAD_REF metric label (see tests/slo/docker-entrypoint.sh). +# +# Notes: +# - OpenTelemetry 1.39.x requires Python >= 3.9. +# - gcc/libc6-dev are needed to build hdrhistogram from sdist. + +FROM python:3.11-slim + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends gcc libc6-dev \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /src + +# 1. ydb-sqlalchemy dialect (current or baseline source tree). +COPY setup.py pyproject.toml requirements.txt README.md ./ +COPY ydb_sqlalchemy/ ydb_sqlalchemy/ +RUN pip install --no-cache-dir . + +# 2. SLO runner deps. +COPY tests/slo/requirements.txt tests/slo/requirements.txt +RUN pip install --no-cache-dir -r tests/slo/requirements.txt + +# 3. Workload source + entrypoint. +COPY tests/slo/src /src/tests/slo/src +COPY tests/slo/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh +RUN chmod +x /usr/local/bin/docker-entrypoint.sh + +ENTRYPOINT ["/usr/local/bin/docker-entrypoint.sh"] diff --git a/tests/slo/README.md b/tests/slo/README.md new file mode 100644 index 0000000..1bbc518 --- /dev/null +++ b/tests/slo/README.md @@ -0,0 +1,78 @@ +# SQLAlchemy SLO workload + +A load generator that exercises the `ydb-sqlalchemy` dialect under +[`ydb-platform/ydb-slo-action`](https://github.com/ydb-platform/ydb-slo-action). +It keeps reader and writer threads busy against a single key/value table and +reports latency / throughput / availability to Prometheus over OTLP, so the +action can compare the current (PR) dialect against a baseline and gate the PR +on regressions. + +## What it does + +* **read** — `SELECT * FROM WHERE object_id = :id` for a random id. +* **write** — `UPSERT INTO
(...) VALUES (...)` for a fresh id. + +Both run in parallel from dedicated thread pools. Every operation is wrapped in +an idempotent retry loop, so transient errors injected by the action's chaos +layer become latency instead of availability drops. + +Two execution modes (selected by `WORKLOAD_NAME` / `--mode`): + +| mode | read path | write path | +|-------|---------------------------------|---------------------------------------------| +| `core`| `Connection.execute(select())` | `Connection.execute(upsert())` | +| `orm` | `Session.get(KeyValueRow, id)` | `Session.execute(upsert())` + `commit()` | + +## Layout + +``` +tests/slo/ +├── Dockerfile # image used by ydb-slo-action (build context = repo root) +├── docker-entrypoint.sh # create (idempotent) then run, honouring injected env +├── requirements.txt # hdrhistogram + opentelemetry +└── src/ + ├── __main__.py # entrypoint: python ./tests/slo/src ... + ├── options.py # argparse (create / run / cleanup) + ├── models.py # engine factory, table + imperatively-mapped ORM entity + ├── generator.py # row payload generator + ├── metrics.py # OTLP metrics (names match the action's metrics.yaml) + └── workload.py # rate limiting, retries, parallel read/write jobs +``` + +## CLI + +```bash +# Create the table and fill it with initial rows. +python ./tests/slo/src create grpc://localhost:2136 /local --mode core + +# Run the parallel read/write load for 60s. +python ./tests/slo/src run grpc://localhost:2136 /local \ + --mode core --time 60 --read-rps 500 --write-rps 50 \ + --otlp-endpoint http://localhost:9090/api/v1/otlp/v1/metrics + +# Drop the table. +python ./tests/slo/src cleanup grpc://localhost:2136 /local +``` + +When run by the action, `YDB_ENDPOINT`, `YDB_DATABASE`, `WORKLOAD_DURATION`, +`WORKLOAD_NAME`, `WORKLOAD_REF` and `OTEL_EXPORTER_OTLP_METRICS_ENDPOINT` are +injected automatically; the metrics endpoint is picked up from the environment. + +## Metrics + +Emitted via OTLP/HTTP and consumed by the action's default `metrics.yaml`: + +| Prometheus name | type | labels | +|-------------------------------------|---------|------------------------------------------| +| `sdk_operations_total` | counter | `ref`, `operation_type`, `operation_status` | +| `sdk_operations_success_total` | counter | `ref`, `operation_type` | +| `sdk_operations_failure_total` | counter | `ref`, `operation_type` | +| `sdk_retry_attempts_total` | counter | `ref`, `operation_type` | +| `sdk_operation_latency_p{50,95,99}_seconds` | gauge | `ref`, `operation_type`, `operation_status` | + +## CI + +`.github/workflows/slo.yml` runs the workload on pull requests labelled `SLO`. +It builds the current and baseline images, hands them to `ydb-slo-action/init`, +then publishes a comparison report with `ydb-slo-action/report`. The cluster is +trimmed to fit a GitHub-hosted runner via `disable_compose_profiles: extra-nodes`. diff --git a/tests/slo/docker-entrypoint.sh b/tests/slo/docker-entrypoint.sh new file mode 100644 index 0000000..9b926fb --- /dev/null +++ b/tests/slo/docker-entrypoint.sh @@ -0,0 +1,34 @@ +#!/bin/sh +# Workload entrypoint used by ydb-slo-action v2. +# +# The action launches the same image for both `current` and `baseline` workloads +# in parallel; both prepare the schema (idempotent) and then run the load. +# +# Env vars injected by the action: +# WORKLOAD_NAME core | orm (selects the SQLAlchemy layer) +# WORKLOAD_REF current / (used as the `ref` metric label) +# WORKLOAD_DURATION run duration in seconds +# YDB_ENDPOINT grpc://ydb:2136 +# YDB_DATABASE /Root/testdb +# OTEL_EXPORTER_OTLP_METRICS_ENDPOINT Prometheus OTLP receiver +# +# Anything passed after the script name is appended to the `run` command — this +# is how tuning flags from `workload_current_command` (e.g. --read-rps) arrive. + +set -e + +ENDPOINT="${YDB_ENDPOINT:-grpc://localhost:2136}" +DATABASE="${YDB_DATABASE:-/local}" +DURATION="${WORKLOAD_DURATION:-60}" +MODE="${WORKLOAD_NAME:-core}" + +echo "SLO workload: mode=${MODE} ref=${WORKLOAD_REF:-current} endpoint=${ENDPOINT} db=${DATABASE} duration=${DURATION}s" + +# Tolerate a parallel container winning the create race. +python ./tests/slo/src create "$ENDPOINT" "$DATABASE" --mode "$MODE" \ + || echo "WARN: create exited non-zero (treated as already-prepared)" >&2 + +exec python ./tests/slo/src run "$ENDPOINT" "$DATABASE" \ + --mode "$MODE" \ + --time "$DURATION" \ + "$@" diff --git a/tests/slo/requirements.txt b/tests/slo/requirements.txt new file mode 100644 index 0000000..f7ace45 --- /dev/null +++ b/tests/slo/requirements.txt @@ -0,0 +1,9 @@ +# Extra deps for the SLO workload runner. +# SQLAlchemy, ydb and ydb-dbapi come from installing the ydb-sqlalchemy dialect. + +hdrhistogram==0.10.3 + +# OpenTelemetry (OTLP/HTTP exporter). 1.39.x requires Python >= 3.9. +opentelemetry-api==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-exporter-otlp-proto-http==1.39.1 diff --git a/tests/slo/src/__main__.py b/tests/slo/src/__main__.py new file mode 100644 index 0000000..ef9959c --- /dev/null +++ b/tests/slo/src/__main__.py @@ -0,0 +1,12 @@ +import logging + +from options import parse_options +from workload import run_from_args + +if __name__ == "__main__": + args = parse_options() + logging.basicConfig( + level=logging.DEBUG if args.debug else logging.INFO, + format="%(asctime)s %(levelname)-8s %(threadName)s %(message)s", + ) + run_from_args(args) diff --git a/tests/slo/src/generator.py b/tests/slo/src/generator.py new file mode 100644 index 0000000..ad3f929 --- /dev/null +++ b/tests/slo/src/generator.py @@ -0,0 +1,53 @@ +"""Row payload generator shared by the create/write paths.""" + +import random +import string +import threading +from dataclasses import dataclass +from datetime import datetime, timezone + +MAX_UINT64 = 2**64 - 1 + + +def _random_string(min_len: int, max_len: int) -> str: + length = random.randint(min_len, max_len) + return "".join(random.choices(string.ascii_lowercase, k=length)) + + +@dataclass +class Row: + object_id: int + payload_str: str + payload_double: float + payload_timestamp: datetime + + def as_params(self) -> dict: + return { + "object_id": self.object_id, + "payload_str": self.payload_str, + "payload_double": self.payload_double, + "payload_timestamp": self.payload_timestamp, + } + + +class RowGenerator: + """Thread-safe generator of rows with monotonically increasing ids.""" + + def __init__(self, start_id: int = 0) -> None: + self._id = start_id + self._lock = threading.Lock() + + def _next_id(self) -> int: + with self._lock: + self._id += 1 + if self._id >= MAX_UINT64: + self._id = 1 + return self._id + + def get(self) -> Row: + return Row( + object_id=self._next_id(), + payload_str=_random_string(20, 40), + payload_double=random.random(), + payload_timestamp=datetime.now(timezone.utc), + ) diff --git a/tests/slo/src/metrics.py b/tests/slo/src/metrics.py new file mode 100644 index 0000000..aac95c3 --- /dev/null +++ b/tests/slo/src/metrics.py @@ -0,0 +1,237 @@ +""" +OTLP metrics for the SQLAlchemy SLO workload. + +The instrument names here are chosen so that, once exported through the +Prometheus OTLP receiver, they line up with the queries shipped in +ydb-slo-action's default ``metrics.yaml``: + + * ``sdk_operations_total`` (counter, labels: ref, operation_type, operation_status) + * ``sdk_operations_success_total`` (counter) + * ``sdk_operations_failure_total`` (counter) + * ``sdk_retry_attempts_total`` (counter, labels: ref, operation_type) + * ``sdk_operation_latency_p50_seconds`` / ``_p95_`` / ``_p99_`` (gauge) + +Latency percentiles are computed client-side per push window via HdrHistogram +and emitted as gauges; the histogram is reset after every push so each sample +describes only the most recent window. +""" + +import logging +import threading +import time +from abc import ABC, abstractmethod +from contextlib import contextmanager +from os import environ +from typing import Optional + +OP_TYPE_READ, OP_TYPE_WRITE = "read", "write" +OP_STATUS_SUCCESS, OP_STATUS_FAILURE = "success", "error" + +REF = environ.get("WORKLOAD_REF") or environ.get("REF") or "current" +WORKLOAD = environ.get("WORKLOAD_NAME") or environ.get("WORKLOAD") or "core" + +logger = logging.getLogger(__name__) + + +def _sdk_version() -> str: + try: + from importlib.metadata import version + + return version("ydb-sqlalchemy") + except Exception: + return "0.0.0" + + +class BaseMetrics(ABC): + @abstractmethod + def start(self, op_type: str) -> float: ... + + @abstractmethod + def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Optional[BaseException] = None) -> None: ... + + @abstractmethod + def push(self) -> None: ... + + @abstractmethod + def reset(self) -> None: ... + + @contextmanager + def measure(self, op_type: str): + start_ts = self.start(op_type) + error = None + try: + yield self + except Exception as err: + error = err + raise + finally: + self.stop(op_type, start_ts, error=error) + + +class DummyMetrics(BaseMetrics): + """No-op metrics used when no OTLP endpoint is configured (local runs).""" + + def start(self, op_type: str) -> float: + return time.time() + + def stop(self, op_type, start_time, attempts=1, error=None) -> None: + return None + + def push(self) -> None: + return None + + def reset(self) -> None: + return None + + +class OtlpMetrics(BaseMetrics): + _HDR_MIN_US = 1 + _HDR_MAX_US = 60_000_000 # 60s + _HDR_SIG_FIGS = 3 + _PERCENTILES = (("p50", 50.0), ("p95", 95.0), ("p99", 99.0)) + + def __init__(self, otlp_metrics_endpoint: str): + from hdrh.histogram import HdrHistogram + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + from opentelemetry.sdk.resources import Resource + + self._HdrHistogram = HdrHistogram + + resource = Resource.create( + { + "service.name": f"workload-{WORKLOAD}", + "service.instance.id": environ.get("SLO_INSTANCE_ID", f"{REF}-{WORKLOAD}"), + "ref": REF, + "sdk": "ydb-sqlalchemy", + "sdk_version": _sdk_version(), + "workload": WORKLOAD, + } + ) + + exporter = OTLPMetricExporter(endpoint=otlp_metrics_endpoint) + reader = PeriodicExportingMetricReader(exporter) + self._provider = MeterProvider(resource=resource, metric_readers=[reader]) + self._meter = self._provider.get_meter("ydb-sqlalchemy-slo") + + self._errors = self._meter.create_counter( + name="sdk.errors.total", + description="Total number of errors, categorized by error type.", + ) + self._operations_total = self._meter.create_counter( + name="sdk.operations.total", + description="Total number of operations attempted.", + ) + self._operations_success_total = self._meter.create_counter( + name="sdk.operations.success.total", + description="Total number of successful operations.", + ) + self._operations_failure_total = self._meter.create_counter( + name="sdk.operations.failure.total", + description="Total number of failed operations.", + ) + self._retry_attempts_total = self._meter.create_counter( + name="sdk.retry.attempts.total", + description="Total number of attempts (including the first one).", + ) + self._pending = self._meter.create_up_down_counter( + name="sdk.pending.operations", + description="Current number of in-flight operations.", + ) + self._latency_gauges = { + name: self._meter.create_gauge( + name=f"sdk.operation.latency.{name}.seconds", + unit="s", + description=f"Operation latency {name} over the last push window.", + ) + for name, _ in self._PERCENTILES + } + + self._lock = threading.Lock() + self._hdr: dict = {} + + def _get_hdr(self, op_type: str, op_status: str): + key = (op_type, op_status) + hist = self._hdr.get(key) + if hist is None: + hist = self._HdrHistogram(self._HDR_MIN_US, self._HDR_MAX_US, self._HDR_SIG_FIGS) + self._hdr[key] = hist + return hist + + def start(self, op_type: str) -> float: + self._pending.add(1, attributes={"ref": REF, "operation_type": op_type}) + return time.time() + + def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Optional[BaseException] = None) -> None: + duration = time.time() - start_time + duration_us = min(max(int(duration * 1_000_000), self._HDR_MIN_US), self._HDR_MAX_US) + + op_status = OP_STATUS_SUCCESS if error is None else OP_STATUS_FAILURE + base_attrs = {"ref": REF, "operation_type": op_type} + op_attrs = {**base_attrs, "operation_status": op_status} + + self._retry_attempts_total.add(int(attempts), attributes=base_attrs) + self._pending.add(-1, attributes=base_attrs) + self._operations_total.add(1, attributes=op_attrs) + + if error is not None: + self._errors.add(1, attributes={**base_attrs, "error_type": type(error).__name__}) + self._operations_failure_total.add(1, attributes=base_attrs) + else: + self._operations_success_total.add(1, attributes=base_attrs) + + with self._lock: + self._get_hdr(op_type, op_status).record_value(duration_us) + + def push(self) -> None: + with self._lock: + for (op_type, op_status), hist in self._hdr.items(): + if hist.get_total_count() == 0: + continue + attrs = {"ref": REF, "operation_type": op_type, "operation_status": op_status} + for name, percentile in self._PERCENTILES: + value_s = hist.get_value_at_percentile(percentile) / 1_000_000 + self._latency_gauges[name].set(value_s, attributes=attrs) + for hist in self._hdr.values(): + hist.reset() + self._provider.force_flush() + + def reset(self) -> None: + with self._lock: + for hist in self._hdr.values(): + hist.reset() + self._provider.force_flush() + + +def _resolve_metrics_endpoint(cli_endpoint: Optional[str]) -> str: + """ + Resolution order: + 1. OTEL_EXPORTER_OTLP_METRICS_ENDPOINT (used as-is) + 2. OTEL_EXPORTER_OTLP_ENDPOINT + /v1/metrics suffix + 3. CLI --otlp-endpoint + """ + metrics_env = environ.get("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", "").strip() + if metrics_env: + return metrics_env + + base_env = environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "").strip() + if base_env: + base = base_env.rstrip("/") + return base if base.endswith("/v1/metrics") else f"{base}/v1/metrics" + + return (cli_endpoint or "").strip() + + +def create_metrics(otlp_endpoint: Optional[str]) -> BaseMetrics: + endpoint = _resolve_metrics_endpoint(otlp_endpoint) + if not endpoint: + logger.info("Metrics disabled (no OTLP endpoint); using DummyMetrics") + return DummyMetrics() + + logger.info("Exporting metrics via OTLP to: %s", endpoint) + try: + return OtlpMetrics(endpoint) + except Exception: + logger.exception("Failed to init OTLP metrics; falling back to DummyMetrics") + return DummyMetrics() diff --git a/tests/slo/src/models.py b/tests/slo/src/models.py new file mode 100644 index 0000000..f584f7a --- /dev/null +++ b/tests/slo/src/models.py @@ -0,0 +1,94 @@ +""" +SQLAlchemy plumbing for the SLO workload: engine factory, the key/value table +definition and an imperatively-mapped ORM class for the ``orm`` workload mode. +""" + +import re + +import sqlalchemy as sa +from sqlalchemy.engine import URL +from sqlalchemy.orm import registry + +from ydb_sqlalchemy import types as ydb_types + +_ENDPOINT_RE = re.compile(r"^(?Pgrpcs?|grpc)://(?P[^:/]+):(?P\d+)") + + +def build_url(endpoint: str, database: str) -> URL: + """ + Turn a YDB endpoint (``grpc://ydb:2136``) and database path (``/Root/testdb``) + into a ``yql+ydb`` SQLAlchemy URL. + """ + match = _ENDPOINT_RE.match(endpoint.strip()) + if not match: + raise ValueError(f"Cannot parse YDB endpoint: {endpoint!r}") + + scheme = match.group("scheme") + host = match.group("host") + port = int(match.group("port")) + + # The dialect re-adds the leading slash to the database, so strip it here to + # avoid a doubled slash in the rendered URL. + db = database.strip().lstrip("/") + + query = {"protocol": "grpcs"} if scheme == "grpcs" else {} + return URL.create("yql+ydb", host=host, port=port, database=db, query=query) + + +def build_engine(endpoint: str, database: str, pool_size: int) -> sa.engine.Engine: + url = build_url(endpoint, database) + # pool_pre_ping lets the pool transparently discard connections that broke + # while a node was down (chaos), instead of handing a dead one to a worker. + return sa.create_engine( + url, + pool_size=pool_size, + max_overflow=max(4, pool_size), + pool_pre_ping=True, + pool_recycle=300, + ) + + +def build_table( + metadata: sa.MetaData, + table_name: str, + *, + min_partitions: int = 6, + max_partitions: int = 100, + partition_size_mb: int = 100, +) -> sa.Table: + return sa.Table( + table_name, + metadata, + sa.Column("object_id", ydb_types.UInt64, primary_key=True), + sa.Column("payload_str", sa.Unicode), + sa.Column("payload_double", sa.Float), + sa.Column("payload_timestamp", sa.TIMESTAMP), + ydb_auto_partitioning_by_size=True, + ydb_auto_partitioning_by_load=True, + ydb_auto_partitioning_min_partitions_count=min_partitions, + ydb_auto_partitioning_max_partitions_count=max_partitions, + ydb_auto_partitioning_partition_size_mb=partition_size_mb, + ) + + +class KeyValueRow: + """Plain class mapped imperatively onto the workload table for ``orm`` mode.""" + + def __init__(self, object_id, payload_str=None, payload_double=None, payload_timestamp=None): + self.object_id = object_id + self.payload_str = payload_str + self.payload_double = payload_double + self.payload_timestamp = payload_timestamp + + +_mapper_registry = registry() +_mapped = False + + +def ensure_mapped(table: sa.Table) -> type: + """Map :class:`KeyValueRow` onto ``table`` exactly once.""" + global _mapped + if not _mapped: + _mapper_registry.map_imperatively(KeyValueRow, table) + _mapped = True + return KeyValueRow diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py new file mode 100644 index 0000000..05302bb --- /dev/null +++ b/tests/slo/src/options.py @@ -0,0 +1,67 @@ +import argparse +from os import environ + +_DEFAULT_MODE = (environ.get("WORKLOAD_NAME") or "core").strip().lower() +if _DEFAULT_MODE not in ("core", "orm"): + _DEFAULT_MODE = "core" + + +def _add_common(parser): + parser.add_argument("endpoint", help="YDB endpoint, e.g. grpc://localhost:2136") + parser.add_argument("database", help="YDB database path, e.g. /local") + parser.add_argument("-t", "--table-name", default="slo_sqlalchemy", help="Workload table name") + parser.add_argument( + "--mode", + choices=("core", "orm"), + default=_DEFAULT_MODE, + help="SQLAlchemy layer to exercise (defaults to WORKLOAD_NAME)", + ) + parser.add_argument("--debug", action="store_true", help="Enable debug logging") + + +def _add_create(subparsers): + parser = subparsers.add_parser("create", help="Create the table and fill it with initial rows") + _add_common(parser) + parser.add_argument("-c", "--initial-data-count", default=1000, type=int, help="Initial row count") + parser.add_argument("--batch-size", default=100, type=int, help="Rows per insert batch") + parser.add_argument("-p-min", "--min-partitions-count", default=6, type=int, help="Min partitions") + parser.add_argument("-p-max", "--max-partitions-count", default=100, type=int, help="Max partitions") + parser.add_argument("-p-size", "--partition-size", default=100, type=int, help="Partition size [MB]") + parser.add_argument("--write-timeout", default=20000, type=int, help="Write timeout [ms]") + parser.add_argument("--max-retries", default=30, type=int, help="Max attempts per operation") + + +def _add_run(subparsers): + parser = subparsers.add_parser("run", help="Run the parallel read/write SLO load") + _add_common(parser) + parser.add_argument("--read-rps", default=500, type=int, help="Target read RPS") + parser.add_argument("--write-rps", default=50, type=int, help="Target write RPS") + parser.add_argument("--read-threads", default=8, type=int, help="Reader threads") + parser.add_argument("--write-threads", default=4, type=int, help="Writer threads") + parser.add_argument("--read-timeout", default=20000, type=int, help="Read timeout [ms]") + parser.add_argument("--write-timeout", default=20000, type=int, help="Write timeout [ms]") + parser.add_argument("--initial-data-count", default=1000, type=int, help="Fallback id space when table is empty") + parser.add_argument("--time", default=60, type=int, help="Run duration [s]") + parser.add_argument("--shutdown-time", default=10, type=int, help="Graceful shutdown time [s]") + parser.add_argument("--max-retries", default=30, type=int, help="Max attempts per operation") + parser.add_argument( + "--otlp-endpoint", + default="", + type=str, + help="OTLP metrics endpoint; empty disables metrics unless OTEL_* env vars are set", + ) + parser.add_argument("--report-period", default=1000, type=int, help="Metrics push period [ms]") + + +def _add_cleanup(subparsers): + parser = subparsers.add_parser("cleanup", help="Drop the workload table") + _add_common(parser) + + +def parse_options(): + parser = argparse.ArgumentParser(description="YDB SQLAlchemy SLO workload") + subparsers = parser.add_subparsers(dest="command", required=True) + _add_create(subparsers) + _add_run(subparsers) + _add_cleanup(subparsers) + return parser.parse_args() diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py new file mode 100644 index 0000000..f632bbd --- /dev/null +++ b/tests/slo/src/workload.py @@ -0,0 +1,299 @@ +""" +Core of the SQLAlchemy SLO workload. + +A single workload run keeps several reader and writer threads busy against one +YDB table through the ``ydb_sqlalchemy`` dialect, while a metrics thread ships +latency/availability samples to Prometheus over OTLP. Two execution modes are +supported, selected by ``WORKLOAD_NAME`` / ``--mode``: + + * ``core`` — SQLAlchemy Core (``Connection.execute``) + * ``orm`` — SQLAlchemy ORM (``Session`` + imperatively mapped entity) + +Every operation is wrapped in an idempotent retry loop so that transient +failures injected by ydb-slo-action's chaos layer turn into latency rather than +into availability drops. +""" + +import logging +import random +import threading +import time + +import sqlalchemy as sa +from sqlalchemy.orm import Session + +from ydb_sqlalchemy import upsert as ydb_upsert + +from generator import RowGenerator +from metrics import OP_TYPE_READ, OP_TYPE_WRITE, create_metrics +from models import build_engine, build_table, ensure_mapped + +logger = logging.getLogger(__name__) + + +class SyncRateLimiter: + """Thread-safe limiter enforcing a minimum interval between permits.""" + + def __init__(self, min_interval_s: float) -> None: + self._min_interval_s = max(0.0, float(min_interval_s)) + self._lock = threading.Lock() + self._next_allowed_ts = 0.0 + + def __enter__(self): + if self._min_interval_s <= 0.0: + return self + while True: + with self._lock: + now = time.monotonic() + if now >= self._next_allowed_ts: + self._next_allowed_ts = now + self._min_interval_s + return self + sleep_for = self._next_allowed_ts - now + if sleep_for > 0: + time.sleep(sleep_for) + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + +def _limiter_for_rps(rps: int) -> SyncRateLimiter: + return SyncRateLimiter(0.0 if rps <= 0 else 1.0 / rps) + + +def retry_call(fn, *, deadline: float, max_attempts: int): + """ + Run ``fn`` until it succeeds, the deadline passes or ``max_attempts`` is hit. + + Returns ``(attempts, error)`` where ``error`` is ``None`` on success. Reads + and UPSERT writes are idempotent, so retrying any failure is safe. + """ + attempt = 0 + last_error = None + while True: + attempt += 1 + try: + fn() + return attempt, None + except Exception as err: # noqa: BLE001 - idempotent ops, retry everything + last_error = err + + if attempt >= max_attempts or time.monotonic() >= deadline: + return attempt, last_error + + backoff = min(1.0, 0.02 * (2 ** min(attempt, 6))) + time.sleep(backoff * (0.5 + random.random())) + + +class Workload: + def __init__(self, args): + self.args = args + self.mode = getattr(args, "mode", "core") + self.table_name = args.table_name + + # -- schema lifecycle --------------------------------------------------- + + def create(self): + engine = build_engine(self.args.endpoint, self.args.database, pool_size=4) + try: + metadata = sa.MetaData() + table = build_table( + metadata, + self.table_name, + min_partitions=self.args.min_partitions_count, + max_partitions=self.args.max_partitions_count, + partition_size_mb=self.args.partition_size, + ) + + # current + baseline containers create the table in parallel; the + # loser of the race sees "already exists" and we just move on. + try: + table.create(engine, checkfirst=True) + logger.info("Table %s is ready", self.table_name) + except Exception: + logger.warning("Create table %s reported an error (assuming it exists)", self.table_name, exc_info=True) + + self._fill_initial_data(engine, table) + finally: + engine.dispose() + + def _fill_initial_data(self, engine, table): + total = self.args.initial_data_count + batch_size = self.args.batch_size + generator = RowGenerator(start_id=0) + logger.info("Filling %s with %s initial rows", self.table_name, total) + + inserted = 0 + while inserted < total: + size = min(batch_size, total - inserted) + batch = [generator.get().as_params() for _ in range(size)] + deadline = time.monotonic() + self.args.write_timeout / 1000 + attempts, error = retry_call( + lambda b=batch: self._upsert_batch(engine, table, b), + deadline=deadline, + max_attempts=self.args.max_retries, + ) + if error is not None: + raise error + inserted += size + + logger.info("Inserted %s rows into %s", inserted, self.table_name) + + @staticmethod + def _upsert_batch(engine, table, batch): + with engine.begin() as conn: + conn.execute(ydb_upsert(table), batch) + + def cleanup(self): + engine = build_engine(self.args.endpoint, self.args.database, pool_size=2) + try: + metadata = sa.MetaData() + table = build_table(metadata, self.table_name) + table.drop(engine, checkfirst=True) + logger.info("Dropped table %s", self.table_name) + finally: + engine.dispose() + + # -- load --------------------------------------------------------------- + + def run(self): + args = self.args + metrics = create_metrics(args.otlp_endpoint) + pool_size = args.read_threads + args.write_threads + 2 + engine = build_engine(args.endpoint, args.database, pool_size=pool_size) + + metadata = sa.MetaData() + table = build_table(metadata, self.table_name) + if self.mode == "orm": + ensure_mapped(table) + + max_id = self._max_id(engine, table) + logger.info("Starting '%s' SLO load on %s (max_id=%s)", self.mode, self.table_name, max_id) + + read_stmt = sa.select(table).where(table.c.object_id == sa.bindparam("object_id")) + read_limiter = _limiter_for_rps(args.read_rps) + write_limiter = _limiter_for_rps(args.write_rps) + row_generator = RowGenerator(start_id=max_id) + end_time = time.monotonic() + args.time + + threads = [] + for i in range(args.read_threads): + threads.append( + threading.Thread( + name=f"slo_read_{i}", + target=self._reader_loop, + args=(engine, table, read_stmt, metrics, read_limiter, max_id, end_time), + ) + ) + for i in range(args.write_threads): + threads.append( + threading.Thread( + name=f"slo_write_{i}", + target=self._writer_loop, + args=(engine, table, metrics, write_limiter, row_generator, end_time), + ) + ) + metrics_thread = threading.Thread( + name="slo_metrics", + target=self._metrics_loop, + args=(metrics, end_time, args.report_period), + ) + + for t in threads: + t.start() + metrics_thread.start() + + for t in threads: + t.join() + metrics_thread.join() + + metrics.push() + metrics.reset() + engine.dispose() + logger.info("Finished '%s' SLO load", self.mode) + + def _max_id(self, engine, table) -> int: + try: + with engine.connect() as conn: + value = conn.execute(sa.select(sa.func.max(table.c.object_id))).scalar() + return int(value) if value else self.args.initial_data_count + except Exception: + logger.warning("Could not read max(object_id); falling back to initial-data-count", exc_info=True) + return max(1, self.args.initial_data_count) + + # -- per-operation work ------------------------------------------------- + + def _reader_loop(self, engine, table, read_stmt, metrics, limiter, max_id, end_time): + model = ensure_mapped(table) if self.mode == "orm" else None + timeout_s = self.args.read_timeout / 1000 + + while time.monotonic() < end_time: + with limiter: + object_id = random.randint(1, max(1, max_id)) + + if self.mode == "orm": + + def do_read(oid=object_id): + with Session(engine) as session: + session.get(model, oid) + + else: + + def do_read(oid=object_id): + with engine.connect() as conn: + conn.execute(read_stmt, {"object_id": oid}).fetchall() + + start = metrics.start(OP_TYPE_READ) + attempts, error = retry_call( + do_read, + deadline=time.monotonic() + timeout_s, + max_attempts=self.args.max_retries, + ) + metrics.stop(OP_TYPE_READ, start, attempts=attempts, error=error) + + def _writer_loop(self, engine, table, metrics, limiter, row_generator, end_time): + timeout_s = self.args.write_timeout / 1000 + + while time.monotonic() < end_time: + with limiter: + params = row_generator.get().as_params() + + if self.mode == "orm": + + def do_write(p=params): + with Session(engine) as session: + session.execute(ydb_upsert(table).values(**p)) + session.commit() + + else: + + def do_write(p=params): + with engine.begin() as conn: + conn.execute(ydb_upsert(table).values(**p)) + + start = metrics.start(OP_TYPE_WRITE) + attempts, error = retry_call( + do_write, + deadline=time.monotonic() + timeout_s, + max_attempts=self.args.max_retries, + ) + metrics.stop(OP_TYPE_WRITE, start, attempts=attempts, error=error) + + @staticmethod + def _metrics_loop(metrics, end_time, report_period_ms): + limiter = SyncRateLimiter(max(1, int(report_period_ms)) / 1000.0) + while time.monotonic() < end_time: + with limiter: + metrics.push() + + +def run_from_args(args): + workload = Workload(args) + command = args.command + if command == "create": + workload.create() + elif command == "run": + workload.run() + elif command == "cleanup": + workload.cleanup() + else: + raise ValueError(f"Unknown command: {command}") From 5e90c6bc7880dad0ec33554bfd52748976603fab Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Wed, 17 Jun 2026 12:48:01 +0300 Subject: [PATCH 02/17] SLO: create baseline tests/ parent dir before copying runner --- .github/workflows/slo.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 28111fe..69ff5d0 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -94,7 +94,9 @@ jobs: # Baseline: baseline dialect with the CURRENT workload runner so the # runner-side contract (entrypoint, metric names) is identical for both. + # main has no tests/ dir, so create the parent before copying the runner. rm -rf "$GITHUB_WORKSPACE/baseline/tests/slo" + mkdir -p "$GITHUB_WORKSPACE/baseline/tests" cp -r "$GITHUB_WORKSPACE/current/tests/slo" "$GITHUB_WORKSPACE/baseline/tests/slo" docker build \ -f "$GITHUB_WORKSPACE/baseline/tests/slo/Dockerfile" \ From 56e319fd31281aff849b20a770bb626a91178247 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Wed, 17 Jun 2026 14:18:04 +0300 Subject: [PATCH 03/17] Move dialect tests under tests/ to drop the test/ vs tests/ split The dialect integration tests now live in tests/integration/ alongside tests/slo/, so the repo no longer has both a test/ and a tests/ directory. tox.ini (lint + dialect pytest paths) and setup.cfg (profile_file) are updated accordingly. --- setup.cfg | 2 +- {test => tests/integration}/__init__.py | 0 {test => tests/integration}/conftest.py | 0 {test => tests/integration}/test_core.py | 0 {test => tests/integration}/test_inspect.py | 0 {test => tests/integration}/test_orm.py | 0 {test => tests/integration}/test_suite.py | 0 tox.ini | 12 ++++++------ 8 files changed, 7 insertions(+), 7 deletions(-) rename {test => tests/integration}/__init__.py (100%) rename {test => tests/integration}/conftest.py (100%) rename {test => tests/integration}/test_core.py (100%) rename {test => tests/integration}/test_inspect.py (100%) rename {test => tests/integration}/test_orm.py (100%) rename {test => tests/integration}/test_suite.py (100%) diff --git a/setup.cfg b/setup.cfg index 650a160..5fec16e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,7 +3,7 @@ addopts= --tb native -v -r fxX -p no:warnings [sqla_testing] requirement_cls=ydb_sqlalchemy.sqlalchemy.requirements:Requirements -profile_file=test/profiles.txt +profile_file=tests/integration/profiles.txt [db] default=yql+ydb://localhost:2136/local diff --git a/test/__init__.py b/tests/integration/__init__.py similarity index 100% rename from test/__init__.py rename to tests/integration/__init__.py diff --git a/test/conftest.py b/tests/integration/conftest.py similarity index 100% rename from test/conftest.py rename to tests/integration/conftest.py diff --git a/test/test_core.py b/tests/integration/test_core.py similarity index 100% rename from test/test_core.py rename to tests/integration/test_core.py diff --git a/test/test_inspect.py b/tests/integration/test_inspect.py similarity index 100% rename from test/test_inspect.py rename to tests/integration/test_inspect.py diff --git a/test/test_orm.py b/tests/integration/test_orm.py similarity index 100% rename from test/test_orm.py rename to tests/integration/test_orm.py diff --git a/test/test_suite.py b/tests/integration/test_suite.py similarity index 100% rename from test/test_suite.py rename to tests/integration/test_suite.py diff --git a/tox.ini b/tox.ini index 111242a..3b6b2bc 100644 --- a/tox.ini +++ b/tox.ini @@ -25,7 +25,7 @@ ignore_errors = True commands = docker-compose up -d python {toxinidir}/wait_container_ready.py - pytest -v test --dbdriver ydb --dbdriver ydb_async + pytest -v tests/integration --dbdriver ydb --dbdriver ydb_async pytest -v ydb_sqlalchemy docker-compose down -v @@ -33,7 +33,7 @@ commands = commands = docker-compose up -d python {toxinidir}/wait_container_ready.py - pytest -v test --dbdriver ydb --dbdriver ydb_async + pytest -v tests/integration --dbdriver ydb --dbdriver ydb_async docker-compose down -v [testenv:test-unit] @@ -51,22 +51,22 @@ commands = [testenv:black] skip_install = true commands = - black --diff --check ydb_sqlalchemy examples/basic_example test + black --diff --check ydb_sqlalchemy examples/basic_example tests/integration [testenv:black-format] skip_install = true commands = - black ydb_sqlalchemy examples/basic_example test + black ydb_sqlalchemy examples/basic_example tests/integration [testenv:isort] skip_install = true commands = - isort ydb_sqlalchemy examples/basic_example test + isort ydb_sqlalchemy examples/basic_example tests/integration [testenv:style] ignore_errors = True commands = - flake8 ydb_sqlalchemy examples/basic_example test + flake8 ydb_sqlalchemy examples/basic_example tests/integration [flake8] show-source = true From fb3cfe238ce5238adbc1a4b2d0d3a81ad4caa9f7 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Wed, 17 Jun 2026 14:44:16 +0300 Subject: [PATCH 04/17] SLO: drop app-level retries, rely on autocommit ydb-dbapi retries The dialect runs in AUTOCOMMIT, so each single-statement read/write already goes through the YDB SDK's retry_operation_sync inside ydb-dbapi. The workload now performs one attempt per operation and records any surfaced exception as a real SLO failure, instead of a broad app-level retry loop that masked non-retryable errors. Removes the now-unused timeout/max-retries flags. --- tests/slo/README.md | 8 +++-- tests/slo/src/options.py | 5 --- tests/slo/src/workload.py | 67 +++++++++++---------------------------- 3 files changed, 24 insertions(+), 56 deletions(-) diff --git a/tests/slo/README.md b/tests/slo/README.md index 1bbc518..a8fb742 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -12,9 +12,11 @@ on regressions. * **read** — `SELECT * FROM
WHERE object_id = :id` for a random id. * **write** — `UPSERT INTO
(...) VALUES (...)` for a fresh id. -Both run in parallel from dedicated thread pools. Every operation is wrapped in -an idempotent retry loop, so transient errors injected by the action's chaos -layer become latency instead of availability drops. +Both run in parallel from dedicated thread pools. Each operation is a single +autocommit statement — there is no app-level retry. The dialect runs in +AUTOCOMMIT by default, so every `execute` already goes through the YDB SDK's +`retry_operation_sync` inside `ydb-dbapi`; any exception that still surfaces is +recorded as a real SLO failure. Two execution modes (selected by `WORKLOAD_NAME` / `--mode`): diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py index 05302bb..1a23bf2 100644 --- a/tests/slo/src/options.py +++ b/tests/slo/src/options.py @@ -27,8 +27,6 @@ def _add_create(subparsers): parser.add_argument("-p-min", "--min-partitions-count", default=6, type=int, help="Min partitions") parser.add_argument("-p-max", "--max-partitions-count", default=100, type=int, help="Max partitions") parser.add_argument("-p-size", "--partition-size", default=100, type=int, help="Partition size [MB]") - parser.add_argument("--write-timeout", default=20000, type=int, help="Write timeout [ms]") - parser.add_argument("--max-retries", default=30, type=int, help="Max attempts per operation") def _add_run(subparsers): @@ -38,12 +36,9 @@ def _add_run(subparsers): parser.add_argument("--write-rps", default=50, type=int, help="Target write RPS") parser.add_argument("--read-threads", default=8, type=int, help="Reader threads") parser.add_argument("--write-threads", default=4, type=int, help="Writer threads") - parser.add_argument("--read-timeout", default=20000, type=int, help="Read timeout [ms]") - parser.add_argument("--write-timeout", default=20000, type=int, help="Write timeout [ms]") parser.add_argument("--initial-data-count", default=1000, type=int, help="Fallback id space when table is empty") parser.add_argument("--time", default=60, type=int, help="Run duration [s]") parser.add_argument("--shutdown-time", default=10, type=int, help="Graceful shutdown time [s]") - parser.add_argument("--max-retries", default=30, type=int, help="Max attempts per operation") parser.add_argument( "--otlp-endpoint", default="", diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py index f632bbd..34062c7 100644 --- a/tests/slo/src/workload.py +++ b/tests/slo/src/workload.py @@ -9,9 +9,10 @@ * ``core`` — SQLAlchemy Core (``Connection.execute``) * ``orm`` — SQLAlchemy ORM (``Session`` + imperatively mapped entity) -Every operation is wrapped in an idempotent retry loop so that transient -failures injected by ydb-slo-action's chaos layer turn into latency rather than -into availability drops. +Each operation is a single autocommit statement. The dialect runs in AUTOCOMMIT +by default, so every ``execute`` already goes through the YDB SDK's +``retry_operation_sync`` inside ydb-dbapi — there is no app-level retry here, and +a surfaced exception is recorded as a genuine SLO failure. """ import logging @@ -60,28 +61,20 @@ def _limiter_for_rps(rps: int) -> SyncRateLimiter: return SyncRateLimiter(0.0 if rps <= 0 else 1.0 / rps) -def retry_call(fn, *, deadline: float, max_attempts: int): - """ - Run ``fn`` until it succeeds, the deadline passes or ``max_attempts`` is hit. +def _measure(metrics, op_type, fn): + """Run a single operation, timing it and recording success/failure. - Returns ``(attempts, error)`` where ``error`` is ``None`` on success. Reads - and UPSERT writes are idempotent, so retrying any failure is safe. + No app-level retry: in AUTOCOMMIT the ydb-dbapi layer already retries + transient YDB errors, so any exception that reaches here is a real failure. + Exceptions are recorded (not re-raised) to keep the worker thread alive. """ - attempt = 0 - last_error = None - while True: - attempt += 1 - try: - fn() - return attempt, None - except Exception as err: # noqa: BLE001 - idempotent ops, retry everything - last_error = err - - if attempt >= max_attempts or time.monotonic() >= deadline: - return attempt, last_error - - backoff = min(1.0, 0.02 * (2 ** min(attempt, 6))) - time.sleep(backoff * (0.5 + random.random())) + start = metrics.start(op_type) + error = None + try: + fn() + except Exception as err: # noqa: BLE001 - count as a failed SLO op, keep going + error = err + metrics.stop(op_type, start, attempts=1, error=error) class Workload: @@ -126,14 +119,7 @@ def _fill_initial_data(self, engine, table): while inserted < total: size = min(batch_size, total - inserted) batch = [generator.get().as_params() for _ in range(size)] - deadline = time.monotonic() + self.args.write_timeout / 1000 - attempts, error = retry_call( - lambda b=batch: self._upsert_batch(engine, table, b), - deadline=deadline, - max_attempts=self.args.max_retries, - ) - if error is not None: - raise error + self._upsert_batch(engine, table, batch) inserted += size logger.info("Inserted %s rows into %s", inserted, self.table_name) @@ -224,7 +210,6 @@ def _max_id(self, engine, table) -> int: def _reader_loop(self, engine, table, read_stmt, metrics, limiter, max_id, end_time): model = ensure_mapped(table) if self.mode == "orm" else None - timeout_s = self.args.read_timeout / 1000 while time.monotonic() < end_time: with limiter: @@ -242,17 +227,9 @@ def do_read(oid=object_id): with engine.connect() as conn: conn.execute(read_stmt, {"object_id": oid}).fetchall() - start = metrics.start(OP_TYPE_READ) - attempts, error = retry_call( - do_read, - deadline=time.monotonic() + timeout_s, - max_attempts=self.args.max_retries, - ) - metrics.stop(OP_TYPE_READ, start, attempts=attempts, error=error) + _measure(metrics, OP_TYPE_READ, do_read) def _writer_loop(self, engine, table, metrics, limiter, row_generator, end_time): - timeout_s = self.args.write_timeout / 1000 - while time.monotonic() < end_time: with limiter: params = row_generator.get().as_params() @@ -270,13 +247,7 @@ def do_write(p=params): with engine.begin() as conn: conn.execute(ydb_upsert(table).values(**p)) - start = metrics.start(OP_TYPE_WRITE) - attempts, error = retry_call( - do_write, - deadline=time.monotonic() + timeout_s, - max_attempts=self.args.max_retries, - ) - metrics.stop(OP_TYPE_WRITE, start, attempts=attempts, error=error) + _measure(metrics, OP_TYPE_WRITE, do_write) @staticmethod def _metrics_loop(metrics, end_time, report_period_ms): From 86f8e2902c3a0eff26895cf84d4801d8c34b538e Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Wed, 17 Jun 2026 14:45:54 +0300 Subject: [PATCH 05/17] SLO: match the python-sdk reference config (600s, 1000/100 rps) Align workload_duration and read/write RPS with ydb-python-sdk's tests/slo workflow. extra-nodes stays disabled to fit a GitHub-hosted runner. --- .github/workflows/slo.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 69ff5d0..20156c5 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -21,9 +21,9 @@ jobs: matrix: workload: - name: core - command: "--read-rps 500 --write-rps 50" + command: "--read-rps 1000 --write-rps 100" - name: orm - command: "--read-rps 500 --write-rps 50" + command: "--read-rps 1000 --write-rps 100" concurrency: group: slo-${{ github.ref }}-${{ matrix.workload.name }} @@ -110,7 +110,7 @@ jobs: github_issue: ${{ github.event.pull_request.number }} github_token: ${{ secrets.GITHUB_TOKEN }} workload_name: ${{ matrix.workload.name }} - workload_duration: "120" + workload_duration: "600" workload_current_ref: ${{ github.head_ref || github.ref_name }} workload_current_image: ydb-slo-current workload_current_command: ${{ matrix.workload.command }} From f866e5dceedfa692eaef2327e7de542515eee3b8 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Wed, 17 Jun 2026 14:50:25 +0300 Subject: [PATCH 06/17] SLO: match the python-sdk reference config (self-hosted runner, full cluster, 600s, 1000/100 rps) Run the workload job on the large-runner-sqlalchemy self-hosted runner with the full YDB cluster (all compose profiles), and align workload_duration and read/write RPS with ydb-python-sdk's tests/slo workflow. The report job stays on ubuntu-latest. --- .github/workflows/slo.yml | 4 +--- tests/slo/README.md | 5 +++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 20156c5..66cb2f6 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -14,7 +14,7 @@ jobs: if: contains(github.event.pull_request.labels.*.name, 'SLO') name: SLO workload (${{ matrix.workload.name }}) - runs-on: ubuntu-latest + runs-on: "large-runner-sqlalchemy" strategy: fail-fast: false @@ -117,8 +117,6 @@ jobs: workload_baseline_ref: ${{ steps.baseline.outputs.ref }} workload_baseline_image: ydb-slo-baseline workload_baseline_command: ${{ matrix.workload.command }} - # Trim the cluster (drop database-3/4/5) to fit a GitHub-hosted runner. - disable_compose_profiles: extra-nodes report: # Runs in the same workflow run as `workload`, so it sees the freshly diff --git a/tests/slo/README.md b/tests/slo/README.md index a8fb742..5e4b4a7 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -76,5 +76,6 @@ Emitted via OTLP/HTTP and consumed by the action's default `metrics.yaml`: `.github/workflows/slo.yml` runs the workload on pull requests labelled `SLO`. It builds the current and baseline images, hands them to `ydb-slo-action/init`, -then publishes a comparison report with `ydb-slo-action/report`. The cluster is -trimmed to fit a GitHub-hosted runner via `disable_compose_profiles: extra-nodes`. +then publishes a comparison report with `ydb-slo-action/report`. The workload +job runs on the `large-runner-sqlalchemy` self-hosted runner with the full YDB +cluster; the report job runs on `ubuntu-latest`. From 7ede3a02cda14803da581d83c7d51c9db20ea73d Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 18 Jun 2026 13:09:54 +0300 Subject: [PATCH 07/17] SLO: make tooling install resilient to transient network errors A curl (28) SSL connection timeout while downloading docker-compose failed the whole job. Wrap the yq/buildx/compose downloads in a retrying curl (--retry --retry-all-errors with connect/max timeouts) so a transient network blip retries instead of failing the run. --- .github/workflows/slo.yml | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 66cb2f6..2488725 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -37,17 +37,26 @@ jobs: BUILDX_VERSION=0.30.1 COMPOSE_VERSION=2.40.3 - sudo curl -fLo /usr/local/bin/yq \ + # Resilient download: retry transient network errors (e.g. curl 28 + # SSL connection timeout) instead of failing the whole job. + dl() { + sudo curl -fL \ + --retry 5 --retry-all-errors --retry-delay 3 \ + --connect-timeout 30 --max-time 300 \ + -o "$1" "$2" + } + + dl /usr/local/bin/yq \ "https://github.com/mikefarah/yq/releases/download/${YQ_VERSION}/yq_linux_amd64" sudo chmod +x /usr/local/bin/yq sudo mkdir -p /usr/local/lib/docker/cli-plugins - sudo curl -fLo /usr/local/lib/docker/cli-plugins/docker-buildx \ + dl /usr/local/lib/docker/cli-plugins/docker-buildx \ "https://github.com/docker/buildx/releases/download/v${BUILDX_VERSION}/buildx-v${BUILDX_VERSION}.linux-amd64" sudo chmod +x /usr/local/lib/docker/cli-plugins/docker-buildx - sudo curl -fLo /usr/local/lib/docker/cli-plugins/docker-compose \ + dl /usr/local/lib/docker/cli-plugins/docker-compose \ "https://github.com/docker/compose/releases/download/v${COMPOSE_VERSION}/docker-compose-linux-x86_64" sudo chmod +x /usr/local/lib/docker/cli-plugins/docker-compose From c2bf6fdeba9b3a216a294169a1b7cf68a6fe2195 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 18 Jun 2026 13:37:34 +0300 Subject: [PATCH 08/17] SLO: remove the SLO label after a run Mirror the ydb-python-sdk reference behaviour: SLO is opt-in per label, so drop the label once the run finishes (any conclusion) to avoid re-running the heavy suite on every push. Re-add the label to trigger another run. --- .github/workflows/slo.yml | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 2488725..7a75413 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -144,3 +144,23 @@ jobs: github_issue: ${{ github.event.pull_request.number }} github_token: ${{ secrets.GITHUB_TOKEN }} github_run_id: ${{ github.run_id }} + + remove-label: + # SLO runs once per label: drop the "SLO" label after the run so ordinary + # pushes don't re-trigger the (heavy) SLO suite. Re-add the label to re-run. + if: ${{ always() && contains(github.event.pull_request.labels.*.name, 'SLO') }} + + name: Remove SLO label + needs: [workload, report] + runs-on: ubuntu-latest + permissions: + pull-requests: write + + steps: + - name: Remove SLO label + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + gh api -X DELETE \ + "repos/${{ github.repository }}/issues/${{ github.event.pull_request.number }}/labels/SLO" \ + && echo "Removed SLO label" || echo "SLO label already absent" From 4ab593b46d8dc1534b6628af6725fc53d0dcfbc9 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 18 Jun 2026 15:16:52 +0300 Subject: [PATCH 09/17] SLO: trigger only when the SLO label is added Switch the trigger to pull_request: [labeled] and gate every job on github.event.label.name == 'SLO', so the suite runs exactly when the label is attached and not on ordinary pushes. Matches the django-ydb-backend setup. --- .github/workflows/slo.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 7a75413..c627905 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -2,7 +2,7 @@ name: SLO on: pull_request: - types: [opened, reopened, synchronize, labeled] + types: [labeled] permissions: contents: read @@ -11,7 +11,7 @@ permissions: jobs: workload: - if: contains(github.event.pull_request.labels.*.name, 'SLO') + if: github.event.label.name == 'SLO' name: SLO workload (${{ matrix.workload.name }}) runs-on: "large-runner-sqlalchemy" @@ -131,7 +131,7 @@ jobs: # Runs in the same workflow run as `workload`, so it sees the freshly # uploaded metrics artifacts and can compare current vs baseline and gate # the PR. Same-repo PRs have the pull-requests:write token this needs. - if: ${{ always() && contains(github.event.pull_request.labels.*.name, 'SLO') }} + if: ${{ always() && github.event.label.name == 'SLO' }} name: SLO report needs: workload @@ -148,7 +148,7 @@ jobs: remove-label: # SLO runs once per label: drop the "SLO" label after the run so ordinary # pushes don't re-trigger the (heavy) SLO suite. Re-add the label to re-run. - if: ${{ always() && contains(github.event.pull_request.labels.*.name, 'SLO') }} + if: ${{ always() && github.event.label.name == 'SLO' }} name: Remove SLO label needs: [workload, report] From 120d9f0173a246dff3901956cdd5d0fbb712014c Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 18 Jun 2026 15:16:52 +0300 Subject: [PATCH 10/17] SLO: drop pool_pre_ping and use idiomatic ORM access pool_pre_ping added a SELECT 1 round-trip on every checkout (~+3ms/op, measured); ydb-dbapi already retries transient errors and re-acquires sessions, so it is unnecessary. ORM mode now uses a sessionmaker with session.get() for reads and session.add()+commit() for writes (a real unit-of-work INSERT) instead of a Core upsert through the session; ORM inserts use a random id to stay collision-free and avoid a hot last partition. --- tests/slo/README.md | 28 ++++++++++++++++------------ tests/slo/src/generator.py | 15 +++++++++++++++ tests/slo/src/models.py | 7 +++---- tests/slo/src/workload.py | 32 +++++++++++++++++--------------- 4 files changed, 51 insertions(+), 31 deletions(-) diff --git a/tests/slo/README.md b/tests/slo/README.md index 5e4b4a7..0f6bca9 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -9,21 +9,25 @@ on regressions. ## What it does -* **read** — `SELECT * FROM
WHERE object_id = :id` for a random id. -* **write** — `UPSERT INTO
(...) VALUES (...)` for a fresh id. +A reader looks up a row by primary key; a writer adds a new row. Both run in +parallel from dedicated thread pools. -Both run in parallel from dedicated thread pools. Each operation is a single -autocommit statement — there is no app-level retry. The dialect runs in -AUTOCOMMIT by default, so every `execute` already goes through the YDB SDK's -`retry_operation_sync` inside `ydb-dbapi`; any exception that still surfaces is -recorded as a real SLO failure. +Two execution modes (selected by `WORKLOAD_NAME` / `--mode`) exercise the two +ways real applications use the dialect: -Two execution modes (selected by `WORKLOAD_NAME` / `--mode`): +| mode | read path (point lookup) | write path | +|-------|---------------------------------|--------------------------------------------------| +| `core`| `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | +| `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | -| mode | read path | write path | -|-------|---------------------------------|---------------------------------------------| -| `core`| `Connection.execute(select())` | `Connection.execute(upsert())` | -| `orm` | `Session.get(KeyValueRow, id)` | `Session.execute(upsert())` + `commit()` | +Each operation is a single autocommit statement — there is no app-level retry. +The dialect runs in AUTOCOMMIT by default, so every `execute` already goes +through the YDB SDK's `retry_operation_sync` inside `ydb-dbapi`; any exception +that still surfaces is recorded as a real SLO failure. + +`core` writes use a monotonic id and UPSERT (idempotent, safe for the shared +table); `orm` writes use a random id so each `session.add` INSERT is collision-free +and avoids a hot last partition. ## Layout diff --git a/tests/slo/src/generator.py b/tests/slo/src/generator.py index ad3f929..9137b71 100644 --- a/tests/slo/src/generator.py +++ b/tests/slo/src/generator.py @@ -51,3 +51,18 @@ def get(self) -> Row: payload_double=random.random(), payload_timestamp=datetime.now(timezone.utc), ) + + +def random_row() -> Row: + """A row with a random 63-bit id. + + Used for ORM inserts: a random id keeps each ``session.add`` collision-free + without cross-thread coordination, and avoids the hot last partition that a + monotonically increasing primary key would create in YDB. + """ + return Row( + object_id=random.getrandbits(63) + 1, + payload_str=_random_string(20, 40), + payload_double=random.random(), + payload_timestamp=datetime.now(timezone.utc), + ) diff --git a/tests/slo/src/models.py b/tests/slo/src/models.py index f584f7a..63d13d3 100644 --- a/tests/slo/src/models.py +++ b/tests/slo/src/models.py @@ -37,14 +37,13 @@ def build_url(endpoint: str, database: str) -> URL: def build_engine(endpoint: str, database: str, pool_size: int) -> sa.engine.Engine: url = build_url(endpoint, database) - # pool_pre_ping lets the pool transparently discard connections that broke - # while a node was down (chaos), instead of handing a dead one to a worker. + # No pool_pre_ping: it would add a SELECT 1 round-trip on every checkout. + # ydb-dbapi already retries transient errors and re-acquires sessions from + # its internal session pool, so broken connections recover without a ping. return sa.create_engine( url, pool_size=pool_size, max_overflow=max(4, pool_size), - pool_pre_ping=True, - pool_recycle=300, ) diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py index 34062c7..5ec28c4 100644 --- a/tests/slo/src/workload.py +++ b/tests/slo/src/workload.py @@ -21,13 +21,13 @@ import time import sqlalchemy as sa -from sqlalchemy.orm import Session +from sqlalchemy.orm import sessionmaker from ydb_sqlalchemy import upsert as ydb_upsert -from generator import RowGenerator +from generator import RowGenerator, random_row from metrics import OP_TYPE_READ, OP_TYPE_WRITE, create_metrics -from models import build_engine, build_table, ensure_mapped +from models import KeyValueRow, build_engine, build_table, ensure_mapped logger = logging.getLogger(__name__) @@ -82,6 +82,7 @@ def __init__(self, args): self.args = args self.mode = getattr(args, "mode", "core") self.table_name = args.table_name + self._session_factory = None # -- schema lifecycle --------------------------------------------------- @@ -151,6 +152,7 @@ def run(self): table = build_table(metadata, self.table_name) if self.mode == "orm": ensure_mapped(table) + self._session_factory = sessionmaker(engine) max_id = self._max_id(engine, table) logger.info("Starting '%s' SLO load on %s (max_id=%s)", self.mode, self.table_name, max_id) @@ -167,7 +169,7 @@ def run(self): threading.Thread( name=f"slo_read_{i}", target=self._reader_loop, - args=(engine, table, read_stmt, metrics, read_limiter, max_id, end_time), + args=(engine, read_stmt, metrics, read_limiter, max_id, end_time), ) ) for i in range(args.write_threads): @@ -208,18 +210,16 @@ def _max_id(self, engine, table) -> int: # -- per-operation work ------------------------------------------------- - def _reader_loop(self, engine, table, read_stmt, metrics, limiter, max_id, end_time): - model = ensure_mapped(table) if self.mode == "orm" else None - + def _reader_loop(self, engine, read_stmt, metrics, limiter, max_id, end_time): while time.monotonic() < end_time: with limiter: object_id = random.randint(1, max(1, max_id)) if self.mode == "orm": - + # Typical ORM read: fetch one mapped entity by primary key. def do_read(oid=object_id): - with Session(engine) as session: - session.get(model, oid) + with self._session_factory() as session: + session.get(KeyValueRow, oid) else: @@ -232,16 +232,18 @@ def do_read(oid=object_id): def _writer_loop(self, engine, table, metrics, limiter, row_generator, end_time): while time.monotonic() < end_time: with limiter: - params = row_generator.get().as_params() - if self.mode == "orm": + # Typical ORM write: add a new mapped object and commit the + # unit of work (a real INSERT through the session). + row = random_row() - def do_write(p=params): - with Session(engine) as session: - session.execute(ydb_upsert(table).values(**p)) + def do_write(r=row): + with self._session_factory() as session: + session.add(KeyValueRow(**r.as_params())) session.commit() else: + params = row_generator.get().as_params() def do_write(p=params): with engine.begin() as conn: From b212c7b08bfc221c671d7d8891e7f9b44bd7825d Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 18 Jun 2026 17:39:56 +0300 Subject: [PATCH 11/17] SLO: add a shared YDB session pool workload variant New 'shared' mode runs the same Core read/write path as 'core' but builds the engine with a single ydb.QuerySessionPool shared across all connections (connect_args={'ydb_session_pool': ...}), instead of every pooled DBAPI connection creating its own driver and session pool. Added as a third matrix entry to compare session-pool strategies. Local 12-thread read benchmark: ~+14% throughput and ~-30% p95/p99 vs per-connection pools. --- .github/workflows/slo.yml | 2 ++ tests/slo/README.md | 13 +++++++++---- tests/slo/src/models.py | 31 +++++++++++++++++++++++++++++++ tests/slo/src/options.py | 8 +++++--- tests/slo/src/workload.py | 13 +++++++++++-- 5 files changed, 58 insertions(+), 9 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index c627905..72c664e 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -24,6 +24,8 @@ jobs: command: "--read-rps 1000 --write-rps 100" - name: orm command: "--read-rps 1000 --write-rps 100" + - name: shared + command: "--read-rps 1000 --write-rps 100" concurrency: group: slo-${{ github.ref }}-${{ matrix.workload.name }} diff --git a/tests/slo/README.md b/tests/slo/README.md index 0f6bca9..1ad7fc5 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -15,10 +15,15 @@ parallel from dedicated thread pools. Two execution modes (selected by `WORKLOAD_NAME` / `--mode`) exercise the two ways real applications use the dialect: -| mode | read path (point lookup) | write path | -|-------|---------------------------------|--------------------------------------------------| -| `core`| `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | -| `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | +| mode | read path (point lookup) | write path | +|---------|---------------------------------|--------------------------------------------------| +| `core` | `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | +| `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | +| `shared`| same as `core` | same as `core` | + +`shared` uses the same Core access as `core` but builds the engine with one +`ydb.QuerySessionPool` shared across all connections (`connect_args={"ydb_session_pool": ...}`), +instead of each pooled connection spinning up its own driver and session pool. Each operation is a single autocommit statement — there is no app-level retry. The dialect runs in AUTOCOMMIT by default, so every `execute` already goes diff --git a/tests/slo/src/models.py b/tests/slo/src/models.py index 63d13d3..c229206 100644 --- a/tests/slo/src/models.py +++ b/tests/slo/src/models.py @@ -6,6 +6,7 @@ import re import sqlalchemy as sa +import ydb from sqlalchemy.engine import URL from sqlalchemy.orm import registry @@ -47,6 +48,36 @@ def build_engine(endpoint: str, database: str, pool_size: int) -> sa.engine.Engi ) +def build_shared_pool_engine(endpoint: str, database: str, pool_size: int): + """Engine where all connections share ONE ``ydb.QuerySessionPool`` and driver. + + By default every pooled DBAPI connection builds its own ``ydb.Driver`` plus a + ``QuerySessionPool``; passing a shared pool via ``connect_args`` means a + single driver and a single set of YDB sessions back the whole engine. + + Returns ``(engine, session_pool, driver)`` so the caller can stop the pool + and driver when the workload is done (disposing the engine leaves the shared + pool/driver running, since they are owned externally). + """ + match = _ENDPOINT_RE.match(endpoint.strip()) + if not match: + raise ValueError(f"Cannot parse YDB endpoint: {endpoint!r}") + ydb_endpoint = f"{match.group('scheme')}://{match.group('host')}:{match.group('port')}" + ydb_database = "/" + database.strip().lstrip("/") + + driver = ydb.Driver(endpoint=ydb_endpoint, database=ydb_database) + driver.wait(timeout=30, fail_fast=True) + session_pool = ydb.QuerySessionPool(driver, size=pool_size) + + engine = sa.create_engine( + build_url(endpoint, database), + pool_size=pool_size, + max_overflow=max(4, pool_size), + connect_args={"ydb_session_pool": session_pool}, + ) + return engine, session_pool, driver + + def build_table( metadata: sa.MetaData, table_name: str, diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py index 1a23bf2..08a0a1e 100644 --- a/tests/slo/src/options.py +++ b/tests/slo/src/options.py @@ -1,8 +1,9 @@ import argparse from os import environ +_MODES = ("core", "orm", "shared") _DEFAULT_MODE = (environ.get("WORKLOAD_NAME") or "core").strip().lower() -if _DEFAULT_MODE not in ("core", "orm"): +if _DEFAULT_MODE not in _MODES: _DEFAULT_MODE = "core" @@ -12,9 +13,10 @@ def _add_common(parser): parser.add_argument("-t", "--table-name", default="slo_sqlalchemy", help="Workload table name") parser.add_argument( "--mode", - choices=("core", "orm"), + choices=_MODES, default=_DEFAULT_MODE, - help="SQLAlchemy layer to exercise (defaults to WORKLOAD_NAME)", + help="Access pattern: core (Connection), orm (Session), " + "shared (Connection over one shared YDB session pool). Defaults to WORKLOAD_NAME.", ) parser.add_argument("--debug", action="store_true", help="Enable debug logging") diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py index 5ec28c4..0c44cf4 100644 --- a/tests/slo/src/workload.py +++ b/tests/slo/src/workload.py @@ -27,7 +27,7 @@ from generator import RowGenerator, random_row from metrics import OP_TYPE_READ, OP_TYPE_WRITE, create_metrics -from models import KeyValueRow, build_engine, build_table, ensure_mapped +from models import KeyValueRow, build_engine, build_shared_pool_engine, build_table, ensure_mapped logger = logging.getLogger(__name__) @@ -146,7 +146,12 @@ def run(self): args = self.args metrics = create_metrics(args.otlp_endpoint) pool_size = args.read_threads + args.write_threads + 2 - engine = build_engine(args.endpoint, args.database, pool_size=pool_size) + shared_handles = None + if self.mode == "shared": + engine, session_pool, driver = build_shared_pool_engine(args.endpoint, args.database, pool_size) + shared_handles = (session_pool, driver) + else: + engine = build_engine(args.endpoint, args.database, pool_size=pool_size) metadata = sa.MetaData() table = build_table(metadata, self.table_name) @@ -197,6 +202,10 @@ def run(self): metrics.push() metrics.reset() engine.dispose() + if shared_handles: + session_pool, driver = shared_handles + session_pool.stop() + driver.stop() logger.info("Finished '%s' SLO load", self.mode) def _max_id(self, engine, table) -> int: From 7eaa569d48cf1ea97946c8bc2d7ca9d299d2dd72 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 19 Jun 2026 11:55:23 +0300 Subject: [PATCH 12/17] SLO: remove the shared session pool workload variant The shared mode ran the same Core operations as core, only with a shared ydb.QuerySessionPool, so it added no distinct load signal; its benefit is resource efficiency at high connection counts, which this SLO does not exercise and cannot measure cleanly across separate matrix runners. The shared-pool feature stays covered by the dialect unit tests. Drops a third full-cluster run from every SLO invocation. --- .github/workflows/slo.yml | 2 -- tests/slo/README.md | 5 ----- tests/slo/src/models.py | 31 ------------------------------- tests/slo/src/options.py | 5 ++--- tests/slo/src/workload.py | 13 ++----------- 5 files changed, 4 insertions(+), 52 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 72c664e..c627905 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -24,8 +24,6 @@ jobs: command: "--read-rps 1000 --write-rps 100" - name: orm command: "--read-rps 1000 --write-rps 100" - - name: shared - command: "--read-rps 1000 --write-rps 100" concurrency: group: slo-${{ github.ref }}-${{ matrix.workload.name }} diff --git a/tests/slo/README.md b/tests/slo/README.md index 1ad7fc5..7fe3e76 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -19,11 +19,6 @@ ways real applications use the dialect: |---------|---------------------------------|--------------------------------------------------| | `core` | `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | | `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | -| `shared`| same as `core` | same as `core` | - -`shared` uses the same Core access as `core` but builds the engine with one -`ydb.QuerySessionPool` shared across all connections (`connect_args={"ydb_session_pool": ...}`), -instead of each pooled connection spinning up its own driver and session pool. Each operation is a single autocommit statement — there is no app-level retry. The dialect runs in AUTOCOMMIT by default, so every `execute` already goes diff --git a/tests/slo/src/models.py b/tests/slo/src/models.py index c229206..63d13d3 100644 --- a/tests/slo/src/models.py +++ b/tests/slo/src/models.py @@ -6,7 +6,6 @@ import re import sqlalchemy as sa -import ydb from sqlalchemy.engine import URL from sqlalchemy.orm import registry @@ -48,36 +47,6 @@ def build_engine(endpoint: str, database: str, pool_size: int) -> sa.engine.Engi ) -def build_shared_pool_engine(endpoint: str, database: str, pool_size: int): - """Engine where all connections share ONE ``ydb.QuerySessionPool`` and driver. - - By default every pooled DBAPI connection builds its own ``ydb.Driver`` plus a - ``QuerySessionPool``; passing a shared pool via ``connect_args`` means a - single driver and a single set of YDB sessions back the whole engine. - - Returns ``(engine, session_pool, driver)`` so the caller can stop the pool - and driver when the workload is done (disposing the engine leaves the shared - pool/driver running, since they are owned externally). - """ - match = _ENDPOINT_RE.match(endpoint.strip()) - if not match: - raise ValueError(f"Cannot parse YDB endpoint: {endpoint!r}") - ydb_endpoint = f"{match.group('scheme')}://{match.group('host')}:{match.group('port')}" - ydb_database = "/" + database.strip().lstrip("/") - - driver = ydb.Driver(endpoint=ydb_endpoint, database=ydb_database) - driver.wait(timeout=30, fail_fast=True) - session_pool = ydb.QuerySessionPool(driver, size=pool_size) - - engine = sa.create_engine( - build_url(endpoint, database), - pool_size=pool_size, - max_overflow=max(4, pool_size), - connect_args={"ydb_session_pool": session_pool}, - ) - return engine, session_pool, driver - - def build_table( metadata: sa.MetaData, table_name: str, diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py index 08a0a1e..a6ba844 100644 --- a/tests/slo/src/options.py +++ b/tests/slo/src/options.py @@ -1,7 +1,7 @@ import argparse from os import environ -_MODES = ("core", "orm", "shared") +_MODES = ("core", "orm") _DEFAULT_MODE = (environ.get("WORKLOAD_NAME") or "core").strip().lower() if _DEFAULT_MODE not in _MODES: _DEFAULT_MODE = "core" @@ -15,8 +15,7 @@ def _add_common(parser): "--mode", choices=_MODES, default=_DEFAULT_MODE, - help="Access pattern: core (Connection), orm (Session), " - "shared (Connection over one shared YDB session pool). Defaults to WORKLOAD_NAME.", + help="Access pattern: core (Connection) or orm (Session). Defaults to WORKLOAD_NAME.", ) parser.add_argument("--debug", action="store_true", help="Enable debug logging") diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py index 0c44cf4..5ec28c4 100644 --- a/tests/slo/src/workload.py +++ b/tests/slo/src/workload.py @@ -27,7 +27,7 @@ from generator import RowGenerator, random_row from metrics import OP_TYPE_READ, OP_TYPE_WRITE, create_metrics -from models import KeyValueRow, build_engine, build_shared_pool_engine, build_table, ensure_mapped +from models import KeyValueRow, build_engine, build_table, ensure_mapped logger = logging.getLogger(__name__) @@ -146,12 +146,7 @@ def run(self): args = self.args metrics = create_metrics(args.otlp_endpoint) pool_size = args.read_threads + args.write_threads + 2 - shared_handles = None - if self.mode == "shared": - engine, session_pool, driver = build_shared_pool_engine(args.endpoint, args.database, pool_size) - shared_handles = (session_pool, driver) - else: - engine = build_engine(args.endpoint, args.database, pool_size=pool_size) + engine = build_engine(args.endpoint, args.database, pool_size=pool_size) metadata = sa.MetaData() table = build_table(metadata, self.table_name) @@ -202,10 +197,6 @@ def run(self): metrics.push() metrics.reset() engine.dispose() - if shared_handles: - session_pool, driver = shared_handles - session_pool.stop() - driver.stop() logger.info("Finished '%s' SLO load", self.mode) def _max_id(self, engine, table) -> int: From c2153239fdeec4e638e83849c1ca99999e681288 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 19 Jun 2026 12:33:10 +0300 Subject: [PATCH 13/17] SLO: stop emitting retry_attempts to avoid spurious threshold failures With no app-level retry, attempts is always 1, so sdk_retry_attempts_total just mirrored sdk_operations_total. The action's *_retry_attempts metric (increase(retry) - increase(ops)) was then pure PromQL noise around zero, and comparing that noise current-vs-baseline produced arbitrary swings (43.8% warning one run, 100% critical the next). Drop the counter so the metric is empty and no longer gates. --- tests/slo/README.md | 5 ++++- tests/slo/src/metrics.py | 11 +++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/slo/README.md b/tests/slo/README.md index 7fe3e76..4e609fe 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -73,9 +73,12 @@ Emitted via OTLP/HTTP and consumed by the action's default `metrics.yaml`: | `sdk_operations_total` | counter | `ref`, `operation_type`, `operation_status` | | `sdk_operations_success_total` | counter | `ref`, `operation_type` | | `sdk_operations_failure_total` | counter | `ref`, `operation_type` | -| `sdk_retry_attempts_total` | counter | `ref`, `operation_type` | | `sdk_operation_latency_p{50,95,99}_seconds` | gauge | `ref`, `operation_type`, `operation_status` | +The action's default `*_retry_attempts` metrics stay empty: this workload has no +app-level retry (transient errors are retried inside `ydb-dbapi`), so it does not +emit `sdk_retry_attempts_total`. + ## CI `.github/workflows/slo.yml` runs the workload on pull requests labelled `SLO`. diff --git a/tests/slo/src/metrics.py b/tests/slo/src/metrics.py index aac95c3..70bcdf1 100644 --- a/tests/slo/src/metrics.py +++ b/tests/slo/src/metrics.py @@ -8,7 +8,6 @@ * ``sdk_operations_total`` (counter, labels: ref, operation_type, operation_status) * ``sdk_operations_success_total`` (counter) * ``sdk_operations_failure_total`` (counter) - * ``sdk_retry_attempts_total`` (counter, labels: ref, operation_type) * ``sdk_operation_latency_p50_seconds`` / ``_p95_`` / ``_p99_`` (gauge) Latency percentiles are computed client-side per push window via HdrHistogram @@ -131,10 +130,11 @@ def __init__(self, otlp_metrics_endpoint: str): name="sdk.operations.failure.total", description="Total number of failed operations.", ) - self._retry_attempts_total = self._meter.create_counter( - name="sdk.retry.attempts.total", - description="Total number of attempts (including the first one).", - ) + # No sdk.retry.attempts.total: this workload has no app-level retry, so + # every operation is a single attempt and the counter would just mirror + # sdk.operations.total. The action's *_retry_attempts metric (which is + # increase(retry) - increase(ops)) would then be pure PromQL noise around + # zero and produce spurious threshold violations, so we don't emit it. self._pending = self._meter.create_up_down_counter( name="sdk.pending.operations", description="Current number of in-flight operations.", @@ -171,7 +171,6 @@ def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Option base_attrs = {"ref": REF, "operation_type": op_type} op_attrs = {**base_attrs, "operation_status": op_status} - self._retry_attempts_total.add(int(attempts), attributes=base_attrs) self._pending.add(-1, attributes=base_attrs) self._operations_total.add(1, attributes=op_attrs) From b1fe4b09acc830af3668fe644e057c5240bb31a7 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 19 Jun 2026 12:36:45 +0300 Subject: [PATCH 14/17] Revert "SLO: stop emitting retry_attempts to avoid spurious threshold failures" This reverts commit c2153239fdeec4e638e83849c1ca99999e681288. --- tests/slo/README.md | 5 +---- tests/slo/src/metrics.py | 11 ++++++----- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/tests/slo/README.md b/tests/slo/README.md index 4e609fe..7fe3e76 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -73,12 +73,9 @@ Emitted via OTLP/HTTP and consumed by the action's default `metrics.yaml`: | `sdk_operations_total` | counter | `ref`, `operation_type`, `operation_status` | | `sdk_operations_success_total` | counter | `ref`, `operation_type` | | `sdk_operations_failure_total` | counter | `ref`, `operation_type` | +| `sdk_retry_attempts_total` | counter | `ref`, `operation_type` | | `sdk_operation_latency_p{50,95,99}_seconds` | gauge | `ref`, `operation_type`, `operation_status` | -The action's default `*_retry_attempts` metrics stay empty: this workload has no -app-level retry (transient errors are retried inside `ydb-dbapi`), so it does not -emit `sdk_retry_attempts_total`. - ## CI `.github/workflows/slo.yml` runs the workload on pull requests labelled `SLO`. diff --git a/tests/slo/src/metrics.py b/tests/slo/src/metrics.py index 70bcdf1..aac95c3 100644 --- a/tests/slo/src/metrics.py +++ b/tests/slo/src/metrics.py @@ -8,6 +8,7 @@ * ``sdk_operations_total`` (counter, labels: ref, operation_type, operation_status) * ``sdk_operations_success_total`` (counter) * ``sdk_operations_failure_total`` (counter) + * ``sdk_retry_attempts_total`` (counter, labels: ref, operation_type) * ``sdk_operation_latency_p50_seconds`` / ``_p95_`` / ``_p99_`` (gauge) Latency percentiles are computed client-side per push window via HdrHistogram @@ -130,11 +131,10 @@ def __init__(self, otlp_metrics_endpoint: str): name="sdk.operations.failure.total", description="Total number of failed operations.", ) - # No sdk.retry.attempts.total: this workload has no app-level retry, so - # every operation is a single attempt and the counter would just mirror - # sdk.operations.total. The action's *_retry_attempts metric (which is - # increase(retry) - increase(ops)) would then be pure PromQL noise around - # zero and produce spurious threshold violations, so we don't emit it. + self._retry_attempts_total = self._meter.create_counter( + name="sdk.retry.attempts.total", + description="Total number of attempts (including the first one).", + ) self._pending = self._meter.create_up_down_counter( name="sdk.pending.operations", description="Current number of in-flight operations.", @@ -171,6 +171,7 @@ def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Option base_attrs = {"ref": REF, "operation_type": op_type} op_attrs = {**base_attrs, "operation_status": op_status} + self._retry_attempts_total.add(int(attempts), attributes=base_attrs) self._pending.add(-1, attributes=base_attrs) self._operations_total.add(1, attributes=op_attrs) From 014e7439358499407d72751dd5eeea1d7e7128a0 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 19 Jun 2026 13:49:28 +0300 Subject: [PATCH 15/17] Add retry_ydb_operation helper for SQLAlchemy operations Unwrap a SQLAlchemy DBAPIError to the underlying ydb.Error and run the operation under the SDK retry policy (retry_ydb_operation / _async, plus a retry_ydb decorator for sync and async functions). Parameters are max_retries and idempotent; no ydb objects are exposed. Lives in the sqlalchemy subpackage with unit tests and docs. --- docs/advanced.rst | 49 ++++++++++ ydb_sqlalchemy/__init__.py | 1 + ydb_sqlalchemy/sqlalchemy/__init__.py | 5 + ydb_sqlalchemy/sqlalchemy/retry.py | 120 ++++++++++++++++++++++++ ydb_sqlalchemy/sqlalchemy/test_retry.py | 98 +++++++++++++++++++ 5 files changed, 273 insertions(+) create mode 100644 ydb_sqlalchemy/sqlalchemy/retry.py create mode 100644 ydb_sqlalchemy/sqlalchemy/test_retry.py diff --git a/docs/advanced.rst b/docs/advanced.rst index 9df13b5..35d801d 100644 --- a/docs/advanced.rst +++ b/docs/advanced.rst @@ -40,3 +40,52 @@ Pass ``_add_declare_for_yql_stmt_vars=True`` to :func:`sqlalchemy.create_engine` ) with engine.connect() as conn: conn.execute(sa.text("SELECT :id"), {"id": 1}) # runs as "DECLARE `$id` as Int64;\nSELECT $id" with param + +Retrying operations +------------------- + +YDB returns retryable errors (``Unavailable``, ``Overloaded``, ``Aborted``, +``BadSession``, ...) that the SDK knows how to retry. SQLAlchemy, however, wraps +the underlying driver error in :class:`sqlalchemy.exc.DBAPIError`, which the SDK +retry logic does not recognise. The helpers in ``ydb_sqlalchemy`` translate the +SQLAlchemy error back to the original ``ydb.Error`` and run the operation under +the SDK's retry policy: + +.. code-block:: python + + import sqlalchemy as sa + from ydb_sqlalchemy import retry_ydb_operation + + engine = sa.create_engine("yql+ydb://localhost:2136/local") + + def read_modify_write(): + with engine.begin() as conn: + value = conn.execute(sa.text("SELECT v FROM t WHERE id = 1")).scalar() + conn.execute(sa.text("UPDATE t SET v = :v WHERE id = 1"), {"v": value + 1}) + + # idempotent=True also retries ambiguous errors; set it only when re-running + # the whole callable is safe. + retry_ydb_operation(read_modify_write, max_retries=10, idempotent=True) + +The callable is re-run from scratch on every attempt, so it must open its own +connection/transaction and not rely on earlier state. Only ``ydb.Error`` is +retried; any other error propagates unchanged. + +There is also a decorator form that works on both sync and async functions, and +an async call form :func:`ydb_sqlalchemy.retry_ydb_operation_async`: + +.. code-block:: python + + from ydb_sqlalchemy import retry_ydb + + @retry_ydb(idempotent=True) + def read_modify_write(): + ... + +.. note:: + + In the default ``AUTOCOMMIT`` isolation level each statement is its own + transaction and is already retried inside ``ydb-dbapi``, so the helper is + mainly useful for **interactive transactions** (a multi-statement + read-modify-write under ``SERIALIZABLE``/``SNAPSHOT``), where the whole + transaction must be retried as a unit. diff --git a/ydb_sqlalchemy/__init__.py b/ydb_sqlalchemy/__init__.py index 55ade24..6e2bc08 100644 --- a/ydb_sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/__init__.py @@ -1,4 +1,5 @@ from ._version import VERSION # noqa: F401 from ydb_dbapi import IsolationLevel # noqa: F401 from .sqlalchemy import Upsert, types, upsert # noqa: F401 +from .sqlalchemy import retry_ydb, retry_ydb_operation, retry_ydb_operation_async # noqa: F401 import ydb_dbapi as dbapi diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index affed36..73198cd 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -21,6 +21,11 @@ import ydb_dbapi from ydb_sqlalchemy.sqlalchemy.dbapi_adapter import AdaptedAsyncConnection from ydb_sqlalchemy.sqlalchemy.dml import Upsert +from ydb_sqlalchemy.sqlalchemy.retry import ( # noqa: F401 + retry_ydb, + retry_ydb_operation, + retry_ydb_operation_async, +) from ydb_sqlalchemy.sqlalchemy.compiler import YqlCompiler, YqlDDLCompiler, YqlIdentifierPreparer, YqlTypeCompiler diff --git a/ydb_sqlalchemy/sqlalchemy/retry.py b/ydb_sqlalchemy/sqlalchemy/retry.py new file mode 100644 index 0000000..c812092 --- /dev/null +++ b/ydb_sqlalchemy/sqlalchemy/retry.py @@ -0,0 +1,120 @@ +"""Retry SQLAlchemy operations that fail with retryable YDB errors. + +SQLAlchemy wraps the underlying ydb-dbapi error in ``sqlalchemy.exc.DBAPIError``, +while YDB's retry logic only recognises ``ydb.Error``. These helpers unwrap the +original YDB error and re-raise it so the SDK can apply its retryable-error +classification (Unavailable, Overloaded, Aborted, BadSession, ...) and back-off. +Only YDB errors are retried; any other error propagates unchanged. + +The retried callable is run from scratch on every attempt, so it must be safe to +re-run (open its own connection/session, not depend on earlier state). + +Call form (sync and async):: + + rows = retry_ydb_operation(read, idempotent=True) + rows = await retry_ydb_operation_async(read, idempotent=True) + +Decorator form (works on both sync and async functions):: + + @retry_ydb(idempotent=True) + def read(): + with engine.connect() as conn: + return conn.execute(stmt).fetchall() +""" + +import functools +import inspect +from typing import Awaitable, Callable, Optional, TypeVar + +import sqlalchemy.exc +import ydb + +T = TypeVar("T") + +DEFAULT_MAX_RETRIES = 10 + + +def _retry_settings(max_retries: int, idempotent: bool) -> ydb.RetrySettings: + return ydb.RetrySettings(max_retries=max_retries, idempotent=idempotent) + + +def _unwrap_ydb_error(exc: BaseException) -> Optional[ydb.Error]: + """Return the ``ydb.Error`` wrapped inside a SQLAlchemy / ydb-dbapi error, if any.""" + original = getattr(getattr(exc, "orig", None), "original_error", None) + return original if isinstance(original, ydb.Error) else None + + +def retry_ydb_operation( + callee: Callable[[], T], + *, + max_retries: int = DEFAULT_MAX_RETRIES, + idempotent: bool = False, +) -> T: + """Run ``callee``, retrying transient YDB errors up to ``max_retries`` times. + + Set ``idempotent=True`` only when re-running ``callee`` is safe, so that + ambiguous errors are retried too. + """ + + def attempt() -> T: + try: + return callee() + except sqlalchemy.exc.DBAPIError as exc: + ydb_error = _unwrap_ydb_error(exc) + if ydb_error is not None: + raise ydb_error from exc + raise + + return ydb.retry_operation_sync(attempt, _retry_settings(max_retries, idempotent)) + + +async def retry_ydb_operation_async( + callee: Callable[[], Awaitable[T]], + *, + max_retries: int = DEFAULT_MAX_RETRIES, + idempotent: bool = False, +) -> T: + """Async counterpart of :func:`retry_ydb_operation`.""" + + async def attempt() -> T: + try: + return await callee() + except sqlalchemy.exc.DBAPIError as exc: + ydb_error = _unwrap_ydb_error(exc) + if ydb_error is not None: + raise ydb_error from exc + raise + + return await ydb.retry_operation_async(attempt, _retry_settings(max_retries, idempotent)) + + +def retry_ydb(*, max_retries: int = DEFAULT_MAX_RETRIES, idempotent: bool = False): + """Decorator that retries transient YDB errors raised by the wrapped function. + + Works on both sync and async functions:: + + @retry_ydb(idempotent=True) + def read(): ... + + @retry_ydb(idempotent=True) + async def read(): ... + """ + + def decorator(func): + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + return await retry_ydb_operation_async( + lambda: func(*args, **kwargs), max_retries=max_retries, idempotent=idempotent + ) + + return async_wrapper + + @functools.wraps(func) + def sync_wrapper(*args, **kwargs): + return retry_ydb_operation(lambda: func(*args, **kwargs), max_retries=max_retries, idempotent=idempotent) + + return sync_wrapper + + return decorator diff --git a/ydb_sqlalchemy/sqlalchemy/test_retry.py b/ydb_sqlalchemy/sqlalchemy/test_retry.py new file mode 100644 index 0000000..f77d5a1 --- /dev/null +++ b/ydb_sqlalchemy/sqlalchemy/test_retry.py @@ -0,0 +1,98 @@ +import asyncio + +import pytest +import sqlalchemy.exc +import ydb +import ydb_dbapi + +from .retry import _unwrap_ydb_error, retry_ydb, retry_ydb_operation, retry_ydb_operation_async + + +def _sa_dbapi_error(ydb_issue): + """Wrap a ydb issue the way the real stack does: ydb.Error -> ydb-dbapi -> SQLAlchemy.""" + dbapi_error = ydb_dbapi.OperationalError(str(ydb_issue), original_error=ydb_issue) + return sqlalchemy.exc.OperationalError("SELECT 1", {}, dbapi_error) + + +def test_unwrap_ydb_error_returns_original(): + issue = ydb.issues.Unavailable("node is down") + assert _unwrap_ydb_error(_sa_dbapi_error(issue)) is issue + + +def test_unwrap_ydb_error_none_for_plain_exception(): + assert _unwrap_ydb_error(ValueError("boom")) is None + + +def test_retry_ydb_operation_passes_through_success(): + assert retry_ydb_operation(lambda: 42) == 42 + + +def test_retry_ydb_operation_retries_transient_then_succeeds(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + def callee(): + calls.append(1) + if len(calls) < 3: + raise _sa_dbapi_error(issue) + return "ok" + + assert retry_ydb_operation(callee, max_retries=5) == "ok" + assert len(calls) == 3 # two transient failures, then success + + +def test_retry_ydb_operation_does_not_retry_non_retryable(): + issue = ydb.issues.BadRequest("bad query") + calls = [] + + def callee(): + calls.append(1) + raise _sa_dbapi_error(issue) + + with pytest.raises(ydb.issues.BadRequest): + retry_ydb_operation(callee, max_retries=5) + assert len(calls) == 1 # not retried + + +def test_retry_ydb_decorator_sync(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + @retry_ydb(max_retries=5) + def op(): + calls.append(1) + if len(calls) < 3: + raise _sa_dbapi_error(issue) + return "ok" + + assert op() == "ok" + assert len(calls) == 3 + + +def test_retry_ydb_decorator_async(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + @retry_ydb(max_retries=5) + async def op(): + calls.append(1) + if len(calls) < 3: + raise _sa_dbapi_error(issue) + return "ok" + + assert asyncio.run(op()) == "ok" + assert len(calls) == 3 + + +def test_retry_ydb_operation_async_call_form(): + issue = ydb.issues.Unavailable("node is down") + calls = [] + + async def callee(): + calls.append(1) + if len(calls) < 2: + raise _sa_dbapi_error(issue) + return "ok" + + assert asyncio.run(retry_ydb_operation_async(callee, max_retries=5)) == "ok" + assert len(calls) == 2 From 4c1c0febe4ce8596a9185881968539da4aee4ce0 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 19 Jun 2026 13:49:28 +0300 Subject: [PATCH 16/17] SLO: add interactive-transaction (tx) workload and count real retries New tx mode runs read-modify-write in a SERIALIZABLE interactive transaction wrapped in retry_ydb_operation (where ydb-dbapi does not retry on its own). core/orm stay autocommit with no app-level retry. The retry counter now emits retries (attempts beyond the first), so it is exactly 0 for autocommit instead of PromQL noise. --- tests/slo/README.md | 36 +++++++++-------- tests/slo/src/metrics.py | 7 +++- tests/slo/src/options.py | 5 ++- tests/slo/src/workload.py | 84 ++++++++++++++++++++++++++++++--------- 4 files changed, 94 insertions(+), 38 deletions(-) diff --git a/tests/slo/README.md b/tests/slo/README.md index 7fe3e76..8600979 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -9,25 +9,29 @@ on regressions. ## What it does -A reader looks up a row by primary key; a writer adds a new row. Both run in -parallel from dedicated thread pools. - -Two execution modes (selected by `WORKLOAD_NAME` / `--mode`) exercise the two -ways real applications use the dialect: - -| mode | read path (point lookup) | write path | -|---------|---------------------------------|--------------------------------------------------| -| `core` | `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | -| `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | - -Each operation is a single autocommit statement — there is no app-level retry. -The dialect runs in AUTOCOMMIT by default, so every `execute` already goes -through the YDB SDK's `retry_operation_sync` inside `ydb-dbapi`; any exception -that still surfaces is recorded as a real SLO failure. +Reader and writer threads run in parallel from dedicated pools. Three execution +modes (selected by `WORKLOAD_NAME` / `--mode`) exercise the ways real +applications use the dialect: + +| mode | read path | write path | +|--------|-----------------------------------|--------------------------------------------------| +| `core` | `Connection.execute(select())` | `Connection.execute(upsert())` (bulk KV upsert) | +| `orm` | `Session.get(KeyValueRow, id)` | `Session.add(KeyValueRow(...))` + `commit()` (ORM insert) | +| `tx` | `select()` in a SERIALIZABLE tx | read-modify-write (`select()` + `upsert()`) in a SERIALIZABLE tx | + +`core` and `orm` issue **single autocommit statements**, which ydb-dbapi already +retries internally, so no app-level retry is used. `tx` runs **interactive +transactions**, which ydb-dbapi does *not* retry on its own, so each transaction +is wrapped in `ydb_sqlalchemy.retry_ydb_operation` — the dialect helper that +translates a SQLAlchemy error back to the underlying `ydb.Error` and retries the +transient ones; the workload counts the attempts, so `sdk_retry_attempts_total` +reflects real retries in that mode. `core` writes use a monotonic id and UPSERT (idempotent, safe for the shared table); `orm` writes use a random id so each `session.add` INSERT is collision-free -and avoids a hot last partition. +and avoids a hot last partition; `tx` writes increment an existing row, which is +why the read-modify-write must be one transaction (and why it is retried with +`idempotent=False`). ## Layout diff --git a/tests/slo/src/metrics.py b/tests/slo/src/metrics.py index aac95c3..f184dc8 100644 --- a/tests/slo/src/metrics.py +++ b/tests/slo/src/metrics.py @@ -133,7 +133,7 @@ def __init__(self, otlp_metrics_endpoint: str): ) self._retry_attempts_total = self._meter.create_counter( name="sdk.retry.attempts.total", - description="Total number of attempts (including the first one).", + description="Total number of retries (attempts beyond the first one).", ) self._pending = self._meter.create_up_down_counter( name="sdk.pending.operations", @@ -171,7 +171,10 @@ def stop(self, op_type: str, start_time: float, attempts: int = 1, error: Option base_attrs = {"ref": REF, "operation_type": op_type} op_attrs = {**base_attrs, "operation_status": op_status} - self._retry_attempts_total.add(int(attempts), attributes=base_attrs) + # Count retries (attempts beyond the first), not total attempts: this is + # exactly 0 when nothing was retried, so the *_retry_attempts metric is a + # clean 0 for autocommit workloads instead of PromQL noise around zero. + self._retry_attempts_total.add(max(int(attempts) - 1, 0), attributes=base_attrs) self._pending.add(-1, attributes=base_attrs) self._operations_total.add(1, attributes=op_attrs) diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py index a6ba844..013431a 100644 --- a/tests/slo/src/options.py +++ b/tests/slo/src/options.py @@ -1,7 +1,7 @@ import argparse from os import environ -_MODES = ("core", "orm") +_MODES = ("core", "orm", "tx") _DEFAULT_MODE = (environ.get("WORKLOAD_NAME") or "core").strip().lower() if _DEFAULT_MODE not in _MODES: _DEFAULT_MODE = "core" @@ -15,7 +15,8 @@ def _add_common(parser): "--mode", choices=_MODES, default=_DEFAULT_MODE, - help="Access pattern: core (Connection) or orm (Session). Defaults to WORKLOAD_NAME.", + help="Access pattern: core (Connection), orm (Session) or " + "tx (interactive read-modify-write transactions). Defaults to WORKLOAD_NAME.", ) parser.add_argument("--debug", action="store_true", help="Enable debug logging") diff --git a/tests/slo/src/workload.py b/tests/slo/src/workload.py index 5ec28c4..0f64d22 100644 --- a/tests/slo/src/workload.py +++ b/tests/slo/src/workload.py @@ -9,10 +9,10 @@ * ``core`` — SQLAlchemy Core (``Connection.execute``) * ``orm`` — SQLAlchemy ORM (``Session`` + imperatively mapped entity) -Each operation is a single autocommit statement. The dialect runs in AUTOCOMMIT -by default, so every ``execute`` already goes through the YDB SDK's -``retry_operation_sync`` inside ydb-dbapi — there is no app-level retry here, and -a surfaced exception is recorded as a genuine SLO failure. +Each operation is wrapped in ``ydb_sqlalchemy.retry_operation``, which unwraps a +SQLAlchemy error back to the underlying ``ydb.Error`` and retries the transient +ones through the SDK's retry policy. We count the attempts (so ``retry_attempts`` +is a real signal) and record the final exception, if any, as an SLO failure. """ import logging @@ -23,6 +23,7 @@ import sqlalchemy as sa from sqlalchemy.orm import sessionmaker +from ydb_sqlalchemy import retry_ydb_operation from ydb_sqlalchemy import upsert as ydb_upsert from generator import RowGenerator, random_row @@ -61,20 +62,34 @@ def _limiter_for_rps(rps: int) -> SyncRateLimiter: return SyncRateLimiter(0.0 if rps <= 0 else 1.0 / rps) -def _measure(metrics, op_type, fn): - """Run a single operation, timing it and recording success/failure. +def _measure(metrics, op_type, fn, *, retry=False, idempotent=False): + """Time ``fn``, count its attempts and record the outcome. - No app-level retry: in AUTOCOMMIT the ydb-dbapi layer already retries - transient YDB errors, so any exception that reaches here is a real failure. - Exceptions are recorded (not re-raised) to keep the worker thread alive. + Autocommit single statements (``core``/``orm``) are already retried inside + ydb-dbapi, so they run once (``retry=False``). Interactive transactions + (``tx``) are not, so they run under ``ydb_sqlalchemy.retry_ydb_operation``, + which translates the SQLAlchemy error back to the underlying ``ydb.Error`` + and retries the transient ones — and we report how many attempts it took, so + ``retry_attempts`` is a real signal. The final exception, if any, is recorded + as an SLO failure but not re-raised, so the worker thread keeps going. """ start = metrics.start(op_type) + attempts = 0 error = None + + def counted(): + nonlocal attempts + attempts += 1 + return fn() + try: - fn() - except Exception as err: # noqa: BLE001 - count as a failed SLO op, keep going + if retry: + retry_ydb_operation(counted, idempotent=idempotent) + else: + counted() + except Exception as err: # noqa: BLE001 - record as a failed SLO op, keep going error = err - metrics.stop(op_type, start, attempts=1, error=error) + metrics.stop(op_type, start, attempts=max(attempts, 1), error=error) class Workload: @@ -177,7 +192,7 @@ def run(self): threading.Thread( name=f"slo_write_{i}", target=self._writer_loop, - args=(engine, table, metrics, write_limiter, row_generator, end_time), + args=(engine, table, metrics, write_limiter, row_generator, max_id, end_time), ) ) metrics_thread = threading.Thread( @@ -221,15 +236,26 @@ def do_read(oid=object_id): with self._session_factory() as session: session.get(KeyValueRow, oid) - else: + _measure(metrics, OP_TYPE_READ, do_read) + + elif self.mode == "tx": + # Read inside an interactive (SERIALIZABLE) transaction. + def do_read(oid=object_id): + with engine.connect().execution_options(isolation_level="SERIALIZABLE") as conn: + with conn.begin(): + conn.execute(read_stmt, {"object_id": oid}).fetchall() + + _measure(metrics, OP_TYPE_READ, do_read, retry=True, idempotent=True) + + else: # core def do_read(oid=object_id): with engine.connect() as conn: conn.execute(read_stmt, {"object_id": oid}).fetchall() - _measure(metrics, OP_TYPE_READ, do_read) + _measure(metrics, OP_TYPE_READ, do_read) - def _writer_loop(self, engine, table, metrics, limiter, row_generator, end_time): + def _writer_loop(self, engine, table, metrics, limiter, row_generator, max_id, end_time): while time.monotonic() < end_time: with limiter: if self.mode == "orm": @@ -242,14 +268,36 @@ def do_write(r=row): session.add(KeyValueRow(**r.as_params())) session.commit() - else: + _measure(metrics, OP_TYPE_WRITE, do_write) + + elif self.mode == "tx": + # Read-modify-write in one interactive (SERIALIZABLE) transaction: + # exactly where the dialect's retry helper is needed (ydb-dbapi + # does not retry an interactive transaction on its own). + object_id = random.randint(1, max(1, max_id)) + + def do_write(oid=object_id): + with engine.connect().execution_options(isolation_level="SERIALIZABLE") as conn: + with conn.begin(): + current = conn.execute( + sa.select(table.c.payload_double).where(table.c.object_id == oid) + ).scalar() + conn.execute( + ydb_upsert(table).values(object_id=oid, payload_double=(current or 0.0) + 1.0) + ) + + # An increment is not idempotent, so only the pre-commit transient + # errors (the SDK's default, idempotent=False) are retried. + _measure(metrics, OP_TYPE_WRITE, do_write, retry=True, idempotent=False) + + else: # core params = row_generator.get().as_params() def do_write(p=params): with engine.begin() as conn: conn.execute(ydb_upsert(table).values(**p)) - _measure(metrics, OP_TYPE_WRITE, do_write) + _measure(metrics, OP_TYPE_WRITE, do_write) @staticmethod def _metrics_loop(metrics, end_time, report_period_ms): From 30776b7e03818bf7c11bd963088fe9db827e269d Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 19 Jun 2026 13:49:28 +0300 Subject: [PATCH 17/17] SLO: count retries directly and set per-scenario thresholds Switch to ydb-slo-action@feat/per-scenario-thresholds. Override the *_retry_attempts metric to count the retry counter directly (the new emission is retries, not total attempts). Make retry_attempts informational for the tx scenario only; core/orm keep the strict default (their value is structurally 0). --- .github/workflows/slo.yml | 32 ++++++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index c627905..519d028 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -22,8 +22,24 @@ jobs: workload: - name: core command: "--read-rps 1000 --write-rps 100" + thresholds: "" - name: orm command: "--read-rps 1000 --write-rps 100" + thresholds: "" + - name: tx + command: "--read-rps 1000 --write-rps 100" + # Interactive transactions legitimately retry (chaos / lock conflicts), + # so retry_attempts is informational for this scenario only. + thresholds: | + metrics: + - name: read_retry_attempts + direction: lower_is_better + warning_change_percent: 100000000 + critical_change_percent: 100000000 + - name: write_retry_attempts + direction: lower_is_better + warning_change_percent: 100000000 + critical_change_percent: 100000000 concurrency: group: slo-${{ github.ref }}-${{ matrix.workload.name }} @@ -113,7 +129,7 @@ jobs: "$GITHUB_WORKSPACE/baseline" - name: Run SLO workload - uses: ydb-platform/ydb-slo-action/init@v2 + uses: ydb-platform/ydb-slo-action/init@feat/per-scenario-thresholds timeout-minutes: 30 with: github_issue: ${{ github.event.pull_request.number }} @@ -126,6 +142,18 @@ jobs: workload_baseline_ref: ${{ steps.baseline.outputs.ref }} workload_baseline_image: ydb-slo-baseline workload_baseline_command: ${{ matrix.workload.command }} + # The workload emits only real retries (attempts beyond the first), so + # count them directly instead of the default "attempts - operations" + # (which would read 0 for every scenario under the new emission). + metrics_yaml: | + metrics: + - name: read_retry_attempts + query: | + sum by(ref) (increase(sdk_retry_attempts_total{operation_type="read"}[5s])) + - name: write_retry_attempts + query: | + sum by(ref) (increase(sdk_retry_attempts_total{operation_type="write"}[5s])) + thresholds_yaml: ${{ matrix.workload.thresholds }} report: # Runs in the same workflow run as `workload`, so it sees the freshly @@ -139,7 +167,7 @@ jobs: steps: - name: Publish SLO report - uses: ydb-platform/ydb-slo-action/report@v2 + uses: ydb-platform/ydb-slo-action/report@feat/per-scenario-thresholds with: github_issue: ${{ github.event.pull_request.number }} github_token: ${{ secrets.GITHUB_TOKEN }}