The action sets these variables on the workload container: - *
Metrics emitted (matching the contract from - * {@code ydb-platform/ydb-slo-action}): - *
Every metric carries the {@code ref} label so the report action can
- * separate current and baseline series.
- */
public final class Metrics implements AutoCloseable {
public enum OperationType {
@@ -83,7 +60,7 @@ public String label() {
private static final AttributeKey Defaults match the SLO workloads in the Go and JavaScript SDKs so the
- * three runs are comparable. JCommander annotations let the operator override
- * any field from the command line, e.g.
- * {@code --read-rps 500 --write-rps 50}.
- */
@SuppressWarnings("FieldMayBeFinal")
public final class KvWorkloadParams {
@@ -62,11 +54,29 @@ public final class KvWorkloadParams {
private int maxPartitionCount = 1_000;
@Parameter(
- names = {"--duration"},
+ names = {"--duration", "--time"},
description = "Run duration in seconds (overrides WORKLOAD_DURATION when > 0)"
)
private int durationSeconds = 0;
+ @Parameter(
+ names = {"--shutdown-time"},
+ description = "Extra seconds, on top of --duration, given to in-flight ops before force-shutdown"
+ )
+ private int shutdownTimeSeconds = 30;
+
+ @Parameter(
+ names = {"--max-attempts"},
+ description = "Maximum total attempts per operation (initial + retries)"
+ )
+ private int maxAttempts = 10;
+
+ @Parameter(
+ names = {"--max-workers"},
+ description = "Hard cap on the number of worker threads per operation type"
+ )
+ private int maxWorkers = 64;
+
public int readRps() {
return readRps;
}
@@ -99,12 +109,8 @@ public int maxPartitionCount() {
return maxPartitionCount;
}
- /**
- * Effective run duration. If the CLI flag was omitted (left at 0), falls
- * back to the value supplied via the {@code WORKLOAD_DURATION} environment
- * variable through {@code Config}.
- * @return Effective run duration value
- */
+
+
public int durationSeconds() {
return durationSeconds;
}
@@ -112,4 +118,16 @@ public int durationSeconds() {
public void setDurationSeconds(int durationSeconds) {
this.durationSeconds = durationSeconds;
}
+
+ public int shutdownTimeSeconds() {
+ return shutdownTimeSeconds;
+ }
+
+ public int maxAttempts() {
+ return maxAttempts;
+ }
+
+ public int maxWorkers() {
+ return maxWorkers;
+ }
}
diff --git a/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/OpOutcome.java b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/OpOutcome.java
new file mode 100644
index 0000000..34bda8a
--- /dev/null
+++ b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/OpOutcome.java
@@ -0,0 +1,37 @@
+package tech.ydb.slo.core.kv;
+
+public final class OpOutcome {
+ private final boolean success;
+ private final int retryAttempts;
+ private final String errorKind;
+
+ private OpOutcome(boolean success, int retryAttempts, String errorKind) {
+ this.success = success;
+ this.retryAttempts = Math.max(0, retryAttempts);
+ this.errorKind = errorKind;
+ }
+
+
+
+ public static OpOutcome success(int retryAttempts) {
+ return new OpOutcome(true, retryAttempts, null);
+ }
+
+
+
+ public static OpOutcome error(int retryAttempts, String errorKind) {
+ return new OpOutcome(false, retryAttempts, errorKind);
+ }
+
+ public boolean isSuccess() {
+ return success;
+ }
+
+ public int retryAttempts() {
+ return retryAttempts;
+ }
+
+ public String errorKind() {
+ return errorKind;
+ }
+}
diff --git a/slo/src/main/java/tech/ydb/slo/kv/Row.java b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/Row.java
similarity index 60%
rename from slo/src/main/java/tech/ydb/slo/kv/Row.java
rename to slo-workload/core/src/main/java/tech/ydb/slo/core/kv/Row.java
index 9f1c292..1d10f08 100644
--- a/slo/src/main/java/tech/ydb/slo/kv/Row.java
+++ b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/Row.java
@@ -1,24 +1,7 @@
-package tech.ydb.slo.kv;
+package tech.ydb.slo.core.kv;
import java.time.Instant;
-/**
- * A single row of the KV workload table.
- *
- * The schema mirrors the one used by SLO workloads in other YDB SDKs
- * (Go, JavaScript) so reports across SDKs are comparable:
- * The {@code hash} column is computed by YDB at insert time via
- * {@code Digest::NumericHash($id)}, so we don't carry it on the client.
- */
public final class Row {
private final long id;
private final String payloadStr;
diff --git a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/RowGenerator.java b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/RowGenerator.java
similarity index 71%
rename from slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/RowGenerator.java
rename to slo-workload/core/src/main/java/tech/ydb/slo/core/kv/RowGenerator.java
index 04fdb9d..60765b0 100644
--- a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/RowGenerator.java
+++ b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/RowGenerator.java
@@ -1,17 +1,10 @@
-package tech.ydb.slo.kv;
+package tech.ydb.slo.core.kv;
import java.time.Instant;
import java.util.Base64;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
-/**
- * Generates rows for the KV workload.
- *
- * Each row gets a monotonically increasing {@code id} and a random payload.
- * The format mirrors the SLO workloads in the Go and JS SDKs so the resulting
- * tables are comparable.
- */
public final class RowGenerator {
private static final int MIN_PAYLOAD_LENGTH = 20;
private static final int MAX_PAYLOAD_LENGTH = 40;
@@ -22,20 +15,15 @@ public RowGenerator(long startId) {
this.nextId = new AtomicLong(startId);
}
- /**
- * Generates a new row with a fresh monotonically increasing id.
- * @return a new row
- */
+
+
public Row generate() {
long id = nextId.getAndIncrement();
return generate(id);
}
- /**
- * Generates a row with an explicit id (used during prefill to control IDs).
- * @param id row id
- * @return a new row
- */
+
+
public static Row generate(long id) {
long payloadHash = ThreadLocalRandom.current().nextLong();
double payloadDouble = ThreadLocalRandom.current().nextDouble();
@@ -52,4 +40,13 @@ private static String randomPayloadString() {
ThreadLocalRandom.current().nextBytes(bytes);
return Base64.getEncoder().withoutPadding().encodeToString(bytes);
}
+
+
+
+ public static long numericHash(long id) {
+ long z = id + 0x9E3779B97F4A7C15L;
+ z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L;
+ z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL;
+ return z ^ (z >>> 31);
+ }
}
diff --git a/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/WorkloadRunner.java b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/WorkloadRunner.java
new file mode 100644
index 0000000..b3a5c69
--- /dev/null
+++ b/slo-workload/core/src/main/java/tech/ydb/slo/core/kv/WorkloadRunner.java
@@ -0,0 +1,309 @@
+package tech.ydb.slo.core.kv;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import com.google.common.util.concurrent.RateLimiter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import tech.ydb.slo.core.Metrics;
+
+public final class WorkloadRunner {
+ private static final Logger logger = LoggerFactory.getLogger(WorkloadRunner.class);
+
+
+
+ private static final double PREFILL_SUCCESS_THRESHOLD = 0.5;
+
+ private final KvClient client;
+ private final Metrics metrics;
+ private final KvWorkloadParams params;
+ private final String tablePath;
+ private final RowGenerator generator;
+
+ public WorkloadRunner(KvClient client, Metrics metrics, KvWorkloadParams params, String tablePath) {
+ this.client = client;
+ this.metrics = metrics;
+ this.params = params;
+ this.tablePath = tablePath;
+ this.generator = new RowGenerator(params.prefillCount());
+ }
+
+
+
+ public void setup() throws Exception {
+ logger.info("creating table {}", tablePath);
+ client.createTable(params, tablePath);
+ logger.info("table {} created", tablePath);
+
+ if (params.prefillCount() <= 0) {
+ logger.info("prefill count <= 0, skipping prefill");
+ return;
+ }
+
+ logger.info("prefilling {} rows into {}", params.prefillCount(), tablePath);
+ int parallelism = Math.min(params.maxWorkers(), Math.max(1, params.prefillCount()));
+ ExecutorService prefillPool = Executors.newFixedThreadPool(
+ parallelism, namedThreadFactory("slo-prefill-")
+ );
+ AtomicLong nextId = new AtomicLong(0);
+ AtomicInteger failed = new AtomicInteger();
+ AtomicInteger sessionOpenFailures = new AtomicInteger();
+ try {
+ for (int w = 0; w < parallelism; w++) {
+ prefillPool.execute(() -> {
+ try (KvSession session = client.openSession()) {
+ long id;
+ while ((id = nextId.getAndIncrement()) < params.prefillCount()) {
+ OpOutcome outcome = session.write(
+ RowGenerator.generate(id), params.writeTimeoutMs());
+ if (!outcome.isSuccess()) {
+ int f = failed.incrementAndGet();
+ if (f <= 5) {
+ logger.warn("prefill row {} failed: {}", id, outcome.errorKind());
+ }
+ }
+ }
+ } catch (Exception e) {
+
+
+
+
+
+ sessionOpenFailures.incrementAndGet();
+ long firstUnclaimed = nextId.getAndSet(params.prefillCount());
+ if (firstUnclaimed < params.prefillCount()) {
+ failed.addAndGet((int) (params.prefillCount() - firstUnclaimed));
+ }
+ logger.error("prefill worker failed to open session: {}", e.toString());
+ }
+ });
+ }
+ } finally {
+ prefillPool.shutdown();
+ if (!prefillPool.awaitTermination(5, TimeUnit.MINUTES)) {
+ prefillPool.shutdownNow();
+ }
+ }
+
+ int total = params.prefillCount();
+ int failedCount = failed.get();
+ int succeeded = total - failedCount;
+ if (sessionOpenFailures.get() == parallelism) {
+ throw new IllegalStateException(
+ "all " + parallelism + " prefill workers failed to open a session — "
+ + "check YDB connectivity and credentials"
+ );
+ }
+ if (succeeded < total * PREFILL_SUCCESS_THRESHOLD) {
+ throw new IllegalStateException(
+ "prefill completed with " + failedCount + " failed rows out of " + total
+ + " (success rate < " + (int) (PREFILL_SUCCESS_THRESHOLD * 100)
+ + "%); reads would target an empty key-space, refusing to run"
+ );
+ }
+ if (failedCount > 0) {
+ logger.warn("prefill completed with {} failed rows out of {}", failedCount, total);
+ } else {
+ logger.info("prefill completed");
+ }
+ }
+
+
+
+ public void run() throws InterruptedException {
+ long durationSeconds = params.durationSeconds();
+ long endNanos = durationSeconds > 0
+ ? System.nanoTime() + TimeUnit.SECONDS.toNanos(durationSeconds)
+ : Long.MAX_VALUE;
+
+
+
+
+ AtomicLong writesIssued = new AtomicLong();
+
+ int readWorkers = workerCount(params.readRps());
+ int writeWorkers = workerCount(params.writeRps());
+
+ if (readWorkers == 0 && writeWorkers == 0) {
+ logger.warn("both read and write RPS are <= 0, run phase has nothing to do");
+ return;
+ }
+
+ ExecutorService readPool = null;
+ ExecutorService writePool = null;
+ try {
+ if (readWorkers > 0) {
+ readPool = Executors.newFixedThreadPool(readWorkers, namedThreadFactory("slo-read-"));
+ RateLimiter readLimiter = RateLimiter.create(params.readRps());
+ for (int i = 0; i < readWorkers; i++) {
+ readPool.execute(() -> readWorkerLoop(endNanos, readLimiter, writesIssued));
+ }
+ } else {
+ logger.info("read RPS <= 0, skipping read workers");
+ }
+
+ if (writeWorkers > 0) {
+ writePool = Executors.newFixedThreadPool(writeWorkers, namedThreadFactory("slo-write-"));
+ RateLimiter writeLimiter = RateLimiter.create(params.writeRps());
+ for (int i = 0; i < writeWorkers; i++) {
+ writePool.execute(() -> writeWorkerLoop(endNanos, writeLimiter, writesIssued));
+ }
+ } else {
+ logger.info("write RPS <= 0, skipping write workers");
+ }
+
+
+ long graceNanos = TimeUnit.SECONDS.toNanos(params.shutdownTimeSeconds());
+ long waitNanos = durationSeconds > 0
+ ? Math.max(0L, endNanos - System.nanoTime()) + graceNanos
+ : Long.MAX_VALUE;
+
+ if (readPool != null) {
+ readPool.shutdown();
+ }
+ if (writePool != null) {
+ writePool.shutdown();
+ }
+
+ if (readPool != null) {
+ long started = System.nanoTime();
+ if (!readPool.awaitTermination(waitNanos, TimeUnit.NANOSECONDS)) {
+ logger.warn("read pool did not drain within deadline, forcing shutdown");
+ readPool.shutdownNow();
+ }
+ waitNanos = Math.max(0L, waitNanos - (System.nanoTime() - started));
+ }
+ if (writePool != null) {
+ if (!writePool.awaitTermination(waitNanos, TimeUnit.NANOSECONDS)) {
+ logger.warn("write pool did not drain within deadline, forcing shutdown");
+ writePool.shutdownNow();
+ }
+ }
+ } finally {
+ forceShutdown(readPool, "read pool");
+ forceShutdown(writePool, "write pool");
+ }
+ }
+
+
+
+ public void teardown() {
+ logger.info("dropping table {}", tablePath);
+ client.dropTable(tablePath);
+ }
+
+
+
+ private void readWorkerLoop(long endNanos, RateLimiter limiter, AtomicLong writesIssued) {
+ try (KvSession session = client.openSession()) {
+ while (!Thread.currentThread().isInterrupted()) {
+ long remaining = endNanos - System.nanoTime();
+ if (remaining <= 0) {
+ return;
+ }
+
+
+
+ if (!limiter.tryAcquire(remaining, TimeUnit.NANOSECONDS)) {
+ return;
+ }
+ try {
+ readOnce(session, writesIssued.get());
+ } catch (Throwable t) {
+ logger.warn("read op threw unexpectedly: {}", t.toString());
+ }
+ }
+ } catch (Exception e) {
+ logger.warn("read worker failed to open session: {}", e.toString());
+ }
+ }
+
+ private void writeWorkerLoop(long endNanos, RateLimiter limiter, AtomicLong writesIssued) {
+ try (KvSession session = client.openSession()) {
+ while (!Thread.currentThread().isInterrupted()) {
+ long remaining = endNanos - System.nanoTime();
+ if (remaining <= 0) {
+ return;
+ }
+ if (!limiter.tryAcquire(remaining, TimeUnit.NANOSECONDS)) {
+ return;
+ }
+ try {
+ writeOnce(session, generator.generate());
+ writesIssued.incrementAndGet();
+ } catch (Throwable t) {
+ logger.warn("write op threw unexpectedly: {}", t.toString());
+ }
+ }
+ } catch (Exception e) {
+ logger.warn("write worker failed to open session: {}", e.toString());
+ }
+ }
+
+
+
+ private void readOnce(KvSession session, long writesObserved) {
+ long upperBound = Math.max(1L, params.prefillCount() + writesObserved);
+ long id = ThreadLocalRandom.current().nextLong(upperBound);
+
+ Metrics.Span span = metrics.startOperation(Metrics.OperationType.READ);
+ OpOutcome outcome = session.read(id, params.readTimeoutMs());
+ if (outcome.isSuccess()) {
+ span.finishSuccess(outcome.retryAttempts());
+ } else {
+ span.finishError(outcome.retryAttempts(), outcome.errorKind());
+ logger.debug("read {} failed: {}", id, outcome.errorKind());
+ }
+ }
+
+ private void writeOnce(KvSession session, Row row) {
+ Metrics.Span span = metrics.startOperation(Metrics.OperationType.WRITE);
+ OpOutcome outcome = session.write(row, params.writeTimeoutMs());
+ if (outcome.isSuccess()) {
+ span.finishSuccess(outcome.retryAttempts());
+ } else {
+ span.finishError(outcome.retryAttempts(), outcome.errorKind());
+ logger.debug("write {} failed: {}", row.id(), outcome.errorKind());
+ }
+ }
+
+
+
+ private int workerCount(int rps) {
+ if (rps <= 0) {
+ return 0;
+ }
+ return Math.min(params.maxWorkers(), Math.max(1, rps));
+ }
+
+ private static ThreadFactory namedThreadFactory(String prefix) {
+ AtomicInteger counter = new AtomicInteger();
+ return r -> {
+ Thread t = new Thread(r, prefix + counter.getAndIncrement());
+ t.setDaemon(true);
+ return t;
+ };
+ }
+
+ private static void forceShutdown(ExecutorService pool, String name) {
+ if (pool == null || pool.isTerminated()) {
+ return;
+ }
+ logger.warn("{} still active in cleanup, forcing shutdown", name);
+ pool.shutdownNow();
+ try {
+ if (!pool.awaitTermination(5, TimeUnit.SECONDS)) {
+ logger.warn("{} did not terminate after shutdownNow", name);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+}
diff --git a/slo-workload/jdbc/Dockerfile b/slo-workload/jdbc/Dockerfile
index 943be48..1331818 100644
--- a/slo-workload/jdbc/Dockerfile
+++ b/slo-workload/jdbc/Dockerfile
@@ -5,7 +5,9 @@
# and run parameters from environment variables and pushes OTLP metrics to the
# endpoint the action provides.
#
-# Build context: the `ydb-java-examples` repository root.
+# Build context: the `ydb-java-examples` repository root. For ydb-jdbc-driver
+# CI the context may also contain `./ydb-jdbc-driver`; when present, the driver
+# is installed from source and the workload is pinned to that exact version.
#
# Optional build args:
# MAVEN_IMAGE Builder image. Defaults to `maven:3.9-eclipse-temurin-17`.
@@ -23,9 +25,26 @@ COPY . /src
ARG YDB_JDBC_VERSION=""
-# Pin the JDBC driver version under test when provided, then build only the
-# workload module (and the parent context it needs).
-RUN if [ -n "${YDB_JDBC_VERSION}" ]; then \
+# Install the JDBC driver from source when the driver checkout is present in
+# the build context. Otherwise use YDB_JDBC_VERSION when provided.
+RUN if [ -d /src/ydb-jdbc-driver ]; then \
+ cd /src/ydb-jdbc-driver && \
+ mvn -B -q \
+ -DskipTests \
+ -Dmaven.javadoc.skip=true \
+ -Dmaven.source.skip=true \
+ -Dgpg.skip=true \
+ install && \
+ mvn -B -q help:evaluate -Dexpression=project.version -DforceStdout > /tmp/ydb-jdbc.version && \
+ YDB_JDBC_VERSION="$(cat /tmp/ydb-jdbc.version)" && \
+ cd /src && \
+ echo "Pinning ydb-jdbc-driver to source-built ${YDB_JDBC_VERSION}" && \
+ mvn -B -q versions:set-property \
+ -Dproperty=ydb.jdbc.version \
+ -DnewVersion="${YDB_JDBC_VERSION}" \
+ -DgenerateBackupPoms=false \
+ -pl slo-workload ; \
+ elif [ -n "${YDB_JDBC_VERSION}" ]; then \
echo "Pinning ydb-jdbc-driver to ${YDB_JDBC_VERSION}" && \
mvn -B -q versions:set-property \
-Dproperty=ydb.jdbc.version \
diff --git a/slo-workload/jdbc/README.md b/slo-workload/jdbc/README.md
index 19ee337..f28d8c9 100644
--- a/slo-workload/jdbc/README.md
+++ b/slo-workload/jdbc/README.md
@@ -40,14 +40,21 @@ relying on server-side YQL builtins inside parameterized statements.
## Retries
-Operations are retried with exponential backoff (up to 10 attempts). An error
-is considered retryable when the driver throws a `SQLRecoverableException` or
-`SQLTransientException` (which covers the driver's
-`YdbRetryableException`, `YdbConditionallyRetryableException`,
-`YdbUnavailbaleException` and `YdbTimeoutException`). The number of retries is
-recorded in `sdk_retry_attempts_total`, and the failure reason is reported via
-the `error_kind` label on `sdk_errors_total` (using the YDB status code when
-available).
+Reads (`SELECT`) and writes (`UPSERT`) are both idempotent, so retries decide
+based on the YDB status code: a `YdbStatusable` whose
+`StatusCode.isRetryable(true)` is true is retried (this covers ABORTED,
+OVERLOADED, UNAVAILABLE, BAD_SESSION, SESSION_BUSY, UNDETERMINED on idempotent
+operations). Anything else falls back to the JDBC marker types
+`SQLRecoverableException` / `SQLTransientException`. The retry attempt cap is
+configurable via `--max-attempts` (default 10) and backoff is capped at 1s.
+Connection-level errors (`SQLRecoverableException`,
+`SQLTransientConnectionException`, `SQLNonTransientConnectionException`)
+invalidate the worker's cached connection before the next attempt opens a
+fresh one.
+
+The number of retries is recorded in `sdk_retry_attempts_total`, and the
+failure reason is reported via the `error_kind` label on `sdk_errors_total`
+(using the YDB status code when available).
## Files
@@ -58,18 +65,18 @@ jdbc/
├── README.md
└── src/main/
├── java/tech/ydb/slo/
- │ ├── Config.java Reads env vars, resolves the JDBC URL
- │ ├── Main.java Entry point
- │ ├── Metrics.java OTLP metrics + HDR histograms
- │ └── kv/
- │ ├── KvWorkload.java Setup/run/teardown loop over JDBC
- │ ├── KvWorkloadParams.java JCommander-bound CLI flags
- │ ├── Row.java Row data class
- │ └── RowGenerator.java Random payload generator
+ │ ├── Main.java Entry point, loads the JDBC driver
+ │ └── jdbc/
+ │ └── JdbcKvClient.java KvClient: connection lifecycle + retry
└── resources/
└── log4j2.xml Console logging config
```
+The shared harness (`Config`, `Metrics`, `KvWorkloadParams`, `Row`,
+`RowGenerator`, `WorkloadRunner`, `Launcher`) lives in
+[`../core`](../core/src/main/java/tech/ydb/slo/core), so every implementation
+emits the same metric contract.
+
## Building and running locally
```bash
diff --git a/slo-workload/jdbc/pom.xml b/slo-workload/jdbc/pom.xml
index c4b174e..20475df 100644
--- a/slo-workload/jdbc/pom.xml
+++ b/slo-workload/jdbc/pom.xml
@@ -11,55 +11,22 @@
The action sets these variables on the workload container:
- * Because the component under test here is the JDBC driver, the
- * YDB connection is expressed as a JDBC URL ({@code jdbc:ydb:...}). The URL is
- * resolved in this order: {@code YDB_JDBC_URL} (used verbatim), then
- * {@code YDB_CONNECTION_STRING} (prefixed with {@code jdbc:ydb:}), then
- * {@code YDB_ENDPOINT} + {@code YDB_DATABASE}.
- */
-public final class Config {
- private final String jdbcUrl;
- private final String token;
- private final String ref;
- private final String workloadName;
- private final int durationSeconds;
- private final String otlpEndpoint;
-
- private Config(
- String jdbcUrl,
- String token,
- String ref,
- String workloadName,
- int durationSeconds,
- String otlpEndpoint
- ) {
- this.jdbcUrl = jdbcUrl;
- this.token = token;
- this.ref = ref;
- this.workloadName = workloadName;
- this.durationSeconds = durationSeconds;
- this.otlpEndpoint = otlpEndpoint;
- }
-
- public String jdbcUrl() {
- return jdbcUrl;
- }
-
- public String token() {
- return token;
- }
-
- public String ref() {
- return ref;
- }
-
- public String workloadName() {
- return workloadName;
- }
-
- public int durationSeconds() {
- return durationSeconds;
- }
-
- public String otlpEndpoint() {
- return otlpEndpoint;
- }
-
- /**
- * Loads configuration from environment variables.
- *
- * @return configuration instance
- * @throws IllegalStateException if required variables are missing or invalid
- */
- public static Config fromEnv() {
- String jdbcUrl = resolveJdbcUrl();
- if (jdbcUrl == null || jdbcUrl.isEmpty()) {
- throw new IllegalStateException(
- "YDB connection is not configured: set YDB_JDBC_URL, "
- + "YDB_CONNECTION_STRING or YDB_ENDPOINT + YDB_DATABASE"
- );
- }
-
- String token = envOrDefault("YDB_TOKEN", "");
- String ref = envOrDefault("WORKLOAD_REF", "unknown");
- String workloadName = envOrDefault("WORKLOAD_NAME", "java-slo-jdbc-workload");
- int durationSeconds = parseInt(envOrDefault("WORKLOAD_DURATION", "600"), 600);
- String otlpEndpoint = envOrDefault("OTEL_EXPORTER_OTLP_ENDPOINT", "");
-
- return new Config(jdbcUrl, token, ref, workloadName, durationSeconds, otlpEndpoint);
- }
-
- private static String resolveJdbcUrl() {
- String explicit = System.getenv("YDB_JDBC_URL");
- if (explicit != null && !explicit.isEmpty()) {
- return explicit;
- }
-
- String connectionString = System.getenv("YDB_CONNECTION_STRING");
- if (connectionString != null && !connectionString.isEmpty()) {
- return toJdbcUrl(connectionString);
- }
-
- String endpoint = System.getenv("YDB_ENDPOINT");
- String database = System.getenv("YDB_DATABASE");
- if (endpoint == null || endpoint.isEmpty() || database == null || database.isEmpty()) {
- return null;
- }
- return toJdbcUrl(composeConnectionString(endpoint, database));
- }
-
- /**
- * Turns a YDB connection string ({@code grpc://host:port/database}) into a
- * JDBC URL understood by the YDB JDBC driver. If the value already starts
- * with {@code jdbc:}, it is returned unchanged.
- */
- private static String toJdbcUrl(String connectionString) {
- if (connectionString.startsWith("jdbc:")) {
- return connectionString;
- }
- return "jdbc:ydb:" + connectionString;
- }
-
- private static String composeConnectionString(String endpoint, String database) {
- // Compose a connection string in the form grpc://host:port/database.
- if (endpoint.endsWith("/") && database.startsWith("/")) {
- return endpoint + database.substring(1);
- }
- if (!endpoint.endsWith("/") && !database.startsWith("/")) {
- return endpoint + "/" + database;
- }
- return endpoint + database;
- }
-
- private static String envOrDefault(String name, String defaultValue) {
- String value = System.getenv(name);
- return (value == null || value.isEmpty()) ? defaultValue : value;
- }
-
- private static int parseInt(String value, int defaultValue) {
- try {
- return Integer.parseInt(value);
- } catch (NumberFormatException e) {
- return defaultValue;
- }
- }
-}
diff --git a/slo-workload/jdbc/src/main/java/tech/ydb/slo/Main.java b/slo-workload/jdbc/src/main/java/tech/ydb/slo/Main.java
index 997f8b9..b370983 100644
--- a/slo-workload/jdbc/src/main/java/tech/ydb/slo/Main.java
+++ b/slo-workload/jdbc/src/main/java/tech/ydb/slo/Main.java
@@ -1,160 +1,30 @@
package tech.ydb.slo;
-import java.util.Properties;
+import tech.ydb.slo.core.Launcher;
+import tech.ydb.slo.jdbc.JdbcKvClient;
-import com.beust.jcommander.JCommander;
-import com.beust.jcommander.ParameterException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import tech.ydb.slo.kv.KvWorkload;
-import tech.ydb.slo.kv.KvWorkloadParams;
-
-/**
- * Entry point of the JDBC SLO workload.
- *
- * Reads connection details and run parameters from environment variables
- * (see {@link Config}), parses workload-specific flags from the command line
- * (see {@link KvWorkloadParams}), and runs the KV workload phases — setup,
- * run, teardown — pushing metrics to the OTLP endpoint configured by the YDB
- * SLO action runtime.
- *
- * Exit codes:
- * The workload creates a partitioned table, prefills it with rows, and then
- * runs read and write loops at fixed RPS for the configured duration. Each
- * operation is timed and retried with exponential backoff; the outcome is
- * recorded into {@link Metrics} so the SLO action can compare current and
- * baseline runs.
- *
- * Schema and queries mirror the KV workloads in the Go and JavaScript SDKs
- * so the produced metrics are directly comparable across SDKs. Unlike the
- * query-service workload, the primary-key {@code hash} column is derived on
- * the client (see {@link #numericHash(long)}) instead of via the server-side
- * {@code Digest::NumericHash}, which keeps the parameterized JDBC statements
- * free of type-inference ambiguity.
- *
- * Concurrency model: each operation type (read / write) gets a dedicated
- * thread pool sized to the configured RPS. Every worker thread owns its own
- * JDBC {@link Connection} (the YDB driver's connections are not thread-safe),
- * pulls a permit from a shared Guava {@link RateLimiter}, and executes the
- * operation inline. There is no separate driver thread and no work queue.
- */
-public final class KvWorkload {
- private static final Logger logger = LoggerFactory.getLogger(KvWorkload.class);
-
- private static final String CREATE_TABLE_QUERY_TEMPLATE = ""
- + "CREATE TABLE IF NOT EXISTS `%s` ("
- + " hash Uint64,"
- + " id Uint64,"
- + " payload_str Utf8,"
- + " payload_double Double,"
- + " payload_timestamp Timestamp,"
- + " payload_hash Uint64,"
- + " PRIMARY KEY (hash, id)"
- + ") WITH ("
- + " UNIFORM_PARTITIONS = %d,"
- + " AUTO_PARTITIONING_BY_SIZE = ENABLED,"
- + " AUTO_PARTITIONING_PARTITION_SIZE_MB = %d,"
- + " AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = %d,"
- + " AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = %d"
- + ")";
-
- private static final String DROP_TABLE_QUERY_TEMPLATE = "DROP TABLE `%s`";
-
- private static final String WRITE_QUERY_TEMPLATE = ""
- + "UPSERT INTO `%s` ("
- + " hash, id, payload_str, payload_double, payload_timestamp, payload_hash"
- + ") VALUES (?, ?, ?, ?, ?, ?)";
-
- private static final String READ_QUERY_TEMPLATE = ""
- + "SELECT id, payload_str, payload_double, payload_timestamp, payload_hash"
- + " FROM `%s`"
- + " WHERE id = ? AND hash = ?";
-
- /*
- * Hard cap on the number of worker threads spawned for a single operation
- * type. The SLO targets a few hundred RPS in CI; allowing more workers
- * than this just wastes threads on JIT-warmup contention without
- * improving throughput.
- */
- private static final int MAX_WORKERS = 64;
-
- /*
- * Maximum number of attempts (initial + retries) per operation before it
- * is recorded as a failure. Mirrors the order of magnitude of the
- * query-service SessionRetryContext default.
- */
- private static final int MAX_ATTEMPTS = 10;
-
- private static final long INITIAL_BACKOFF_MS = 10L;
- private static final long MAX_BACKOFF_MS = 1_000L;
-
- /*
- * Extra time, on top of the workload duration, given to worker pools to
- * complete their last in-flight operations before {@link #run()} forces
- * shutdown.
- */
- private static final long SHUTDOWN_GRACE_SECONDS = 30L;
-
- private final String jdbcUrl;
- private final Properties connectionProperties;
- private final Metrics metrics;
- private final KvWorkloadParams params;
- private final String tablePath;
- private final RowGenerator generator;
-
- public KvWorkload(
- String jdbcUrl,
- Properties connectionProperties,
- Metrics metrics,
- KvWorkloadParams params,
- String tablePath
- ) {
- this.jdbcUrl = jdbcUrl;
- this.connectionProperties = connectionProperties;
- this.metrics = metrics;
- this.params = params;
- this.tablePath = tablePath;
- this.generator = new RowGenerator(params.prefillCount());
- }
-
- /*
- * Creates the table (if missing) and prefills it with
- * {@code params.prefillCount()} rows using a bounded pool of worker
- * connections.
- */
- public void setup() throws InterruptedException, SQLException {
- logger.info("creating table {}", tablePath);
- try (Connection conn = openConnection();
- Statement stmt = conn.createStatement()) {
- stmt.execute(String.format(
- CREATE_TABLE_QUERY_TEMPLATE,
- tablePath,
- params.minPartitionCount(),
- params.partitionSizeMb(),
- params.minPartitionCount(),
- params.maxPartitionCount()
- ));
- }
- logger.info("table {} created", tablePath);
-
- if (params.prefillCount() <= 0) {
- logger.info("prefill count <= 0, skipping prefill");
- return;
- }
-
- logger.info("prefilling {} rows into {}", params.prefillCount(), tablePath);
- int parallelism = Math.min(MAX_WORKERS, Math.max(1, params.prefillCount()));
- ExecutorService prefillPool = Executors.newFixedThreadPool(
- parallelism, namedThreadFactory("slo-prefill-")
- );
- AtomicLong nextId = new AtomicLong(0);
- AtomicInteger failed = new AtomicInteger();
- try {
- for (int w = 0; w < parallelism; w++) {
- prefillPool.execute(() -> {
- try (WorkerConnection wc = new WorkerConnection()) {
- long id;
- while ((id = nextId.getAndIncrement()) < params.prefillCount()) {
- SQLException err = writeWithRetry(wc, RowGenerator.generate(id),
- params.writeTimeoutMs(), null);
- if (err != null) {
- int f = failed.incrementAndGet();
- if (f <= 5) {
- logger.warn("prefill row {} failed: {}", id, err.toString());
- }
- }
- }
- }
- });
- }
- } finally {
- prefillPool.shutdown();
- if (!prefillPool.awaitTermination(5, TimeUnit.MINUTES)) {
- prefillPool.shutdownNow();
- }
- }
- if (failed.get() > 0) {
- logger.warn("prefill completed with {} failed rows out of {}",
- failed.get(), params.prefillCount());
- } else {
- logger.info("prefill completed");
- }
- }
-
- /*
- * Runs the workload until the configured deadline or thread interruption.
- */
- public void run() throws InterruptedException {
- long durationSeconds = params.durationSeconds();
- long endNanos = durationSeconds > 0
- ? System.nanoTime() + TimeUnit.SECONDS.toNanos(durationSeconds)
- : Long.MAX_VALUE;
-
- // Track how many writes have completed so reads target a key-space
- // that's actually been populated. The generator was constructed with
- // nextId = prefillCount, so writes pick up where prefill left off.
- AtomicLong writesIssued = new AtomicLong();
-
- int readWorkers = workerCount(params.readRps());
- int writeWorkers = workerCount(params.writeRps());
-
- if (readWorkers == 0 && writeWorkers == 0) {
- logger.warn("both read and write RPS are <= 0, run phase has nothing to do");
- return;
- }
-
- ExecutorService readPool = null;
- ExecutorService writePool = null;
- try {
- if (readWorkers > 0) {
- readPool = Executors.newFixedThreadPool(readWorkers, namedThreadFactory("slo-read-"));
- RateLimiter readLimiter = RateLimiter.create(params.readRps());
- for (int i = 0; i < readWorkers; i++) {
- readPool.execute(() -> readWorkerLoop(endNanos, readLimiter, writesIssued));
- }
- } else {
- logger.info("read RPS <= 0, skipping read workers");
- }
-
- if (writeWorkers > 0) {
- writePool = Executors.newFixedThreadPool(writeWorkers, namedThreadFactory("slo-write-"));
- RateLimiter writeLimiter = RateLimiter.create(params.writeRps());
- for (int i = 0; i < writeWorkers; i++) {
- writePool.execute(() -> writeWorkerLoop(endNanos, writeLimiter, writesIssued));
- }
- } else {
- logger.info("write RPS <= 0, skipping write workers");
- }
-
- // Wait for workers to drain naturally as they hit the deadline.
- long graceNanos = TimeUnit.SECONDS.toNanos(SHUTDOWN_GRACE_SECONDS);
- long waitNanos = durationSeconds > 0
- ? Math.max(0L, endNanos - System.nanoTime()) + graceNanos
- : Long.MAX_VALUE;
-
- if (readPool != null) {
- readPool.shutdown();
- }
- if (writePool != null) {
- writePool.shutdown();
- }
-
- if (readPool != null) {
- long started = System.nanoTime();
- if (!readPool.awaitTermination(waitNanos, TimeUnit.NANOSECONDS)) {
- logger.warn("read pool did not drain within deadline, forcing shutdown");
- readPool.shutdownNow();
- }
- waitNanos = Math.max(0L, waitNanos - (System.nanoTime() - started));
- }
- if (writePool != null) {
- if (!writePool.awaitTermination(waitNanos, TimeUnit.NANOSECONDS)) {
- logger.warn("write pool did not drain within deadline, forcing shutdown");
- writePool.shutdownNow();
- }
- }
- } finally {
- forceShutdown(readPool, "read pool");
- forceShutdown(writePool, "write pool");
- }
- }
-
- /*
- * Drops the workload table. Called from the {@code finally} block in
- * {@code Main} so the database is left clean even on failure.
- */
- public void teardown() {
- logger.info("dropping table {}", tablePath);
- try (Connection conn = openConnection();
- Statement stmt = conn.createStatement()) {
- stmt.execute(String.format(DROP_TABLE_QUERY_TEMPLATE, tablePath));
- logger.info("table {} dropped", tablePath);
- } catch (SQLException e) {
- logger.warn("failed to drop table {}: {}", tablePath, e.toString());
- }
- }
-
- // --- worker loops ------------------------------------------------------
-
- private void readWorkerLoop(long endNanos, RateLimiter limiter, AtomicLong writesIssued) {
- try (WorkerConnection wc = new WorkerConnection()) {
- while (System.nanoTime() < endNanos && !Thread.currentThread().isInterrupted()) {
- limiter.acquire();
- try {
- readOnce(wc, writesIssued.get());
- } catch (Throwable t) {
- logger.warn("read op threw unexpectedly: {}", t.toString());
- }
- }
- }
- }
-
- private void writeWorkerLoop(long endNanos, RateLimiter limiter, AtomicLong writesIssued) {
- try (WorkerConnection wc = new WorkerConnection()) {
- while (System.nanoTime() < endNanos && !Thread.currentThread().isInterrupted()) {
- limiter.acquire();
- try {
- writeOnce(wc, generator.generate());
- writesIssued.incrementAndGet();
- } catch (Throwable t) {
- logger.warn("write op threw unexpectedly: {}", t.toString());
- }
- }
- }
- }
-
- // --- single operations -------------------------------------------------
-
- private void readOnce(WorkerConnection wc, long writesObserved) {
- long upperBound = Math.max(1L, params.prefillCount() + writesObserved);
- long id = ThreadLocalRandom.current().nextLong(upperBound);
- long hash = numericHash(id);
-
- Metrics.Span span = metrics.startOperation(Metrics.OperationType.READ);
- int attempts = 0;
- SQLException last = null;
- while (attempts < MAX_ATTEMPTS) {
- attempts++;
- try {
- wc.read(id, hash, timeoutSeconds(params.readTimeoutMs()));
- span.finishSuccess(attempts - 1);
- return;
- } catch (SQLException e) {
- last = e;
- if (!isRetryable(e) || attempts >= MAX_ATTEMPTS) {
- break;
- }
- wc.invalidateOnConnectionError(e);
- backoff(attempts);
- }
- }
- span.finishError(attempts - 1, classifyError(last));
- logger.debug("read {} failed: {}", id, last == null ? "?" : last.toString());
- }
-
- private void writeOnce(WorkerConnection wc, Row row) {
- Metrics.Span span = metrics.startOperation(Metrics.OperationType.WRITE);
- int[] attemptsOut = new int[1];
- SQLException err = writeWithRetry(wc, row, params.writeTimeoutMs(), attemptsOut);
- if (err == null) {
- span.finishSuccess(attemptsOut[0] - 1);
- } else {
- span.finishError(Math.max(0, attemptsOut[0] - 1), classifyError(err));
- logger.debug("write {} failed: {}", row.id(), err.toString());
- }
- }
-
- /*
- * Writes a single row with retry. When {@code attemptsOut} is non-null, the
- * total number of attempts is written to its first element. Returns
- * {@code null} on success or the last {@link SQLException} on failure.
- * Used both by the run phase (with metrics handled by the caller) and
- * prefill (silent).
- */
- private SQLException writeWithRetry(WorkerConnection wc, Row row, int timeoutMs, int[] attemptsOut) {
- long hash = numericHash(row.id());
- int attempts = 0;
- SQLException last = null;
- while (attempts < MAX_ATTEMPTS) {
- attempts++;
- try {
- wc.write(row, hash, timeoutSeconds(timeoutMs));
- if (attemptsOut != null) {
- attemptsOut[0] = attempts;
- }
- return null;
- } catch (SQLException e) {
- last = e;
- if (!isRetryable(e) || attempts >= MAX_ATTEMPTS) {
- break;
- }
- wc.invalidateOnConnectionError(e);
- backoff(attempts);
- }
- }
- if (attemptsOut != null) {
- attemptsOut[0] = attempts;
- }
- return last;
- }
-
- // --- helpers -----------------------------------------------------------
-
- private Connection openConnection() throws SQLException {
- return DriverManager.getConnection(jdbcUrl, connectionProperties);
- }
-
- private static int timeoutSeconds(int timeoutMs) {
- return Math.max(1, (timeoutMs + 999) / 1000);
- }
-
- private static boolean isRetryable(SQLException e) {
- return e instanceof SQLRecoverableException || e instanceof SQLTransientException;
- }
-
- private static boolean isConnectionError(SQLException e) {
- return e instanceof SQLRecoverableException || e instanceof SQLTransientConnectionException;
- }
-
- private static String classifyError(SQLException e) {
- if (e == null) {
- return "unknown";
- }
- if (e instanceof YdbStatusable) {
- try {
- return "ydb/" + ((YdbStatusable) e).getStatus().getCode().name().toLowerCase();
- } catch (RuntimeException ignored) {
- // fall through to the generic classification
- }
- }
- return e.getClass().getSimpleName().toLowerCase();
- }
-
- private static void backoff(int attempt) {
- long delay = Math.min(MAX_BACKOFF_MS, INITIAL_BACKOFF_MS * (1L << Math.min(attempt - 1, 20)));
- try {
- Thread.sleep(delay);
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- }
- }
-
- /**
- * Derives the primary-key {@code hash} column from {@code id} using a
- * SplitMix64-style mix. Reads and writes both call this so they always
- * target the same key. The exact function does not need to match YQL's
- * {@code Digest::NumericHash}; it only needs to be deterministic and
- * well distributed across partitions.
- * @param id row id
- * @return derived hash value
- */
- private static long numericHash(long id) {
- long z = id + 0x9E3779B97F4A7C15L;
- z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L;
- z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL;
- return z ^ (z >>> 31);
- }
-
- private static int workerCount(int rps) {
- if (rps <= 0) {
- return 0;
- }
- return Math.min(MAX_WORKERS, Math.max(1, rps));
- }
-
- private static ThreadFactory namedThreadFactory(String prefix) {
- AtomicInteger counter = new AtomicInteger();
- return r -> {
- Thread t = new Thread(r, prefix + counter.getAndIncrement());
- t.setDaemon(true);
- return t;
- };
- }
-
- private static void forceShutdown(ExecutorService pool, String name) {
- if (pool == null || pool.isTerminated()) {
- return;
- }
- logger.warn("{} still active in cleanup, forcing shutdown", name);
- pool.shutdownNow();
- try {
- if (!pool.awaitTermination(5, TimeUnit.SECONDS)) {
- logger.warn("{} did not terminate after shutdownNow", name);
- }
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- }
- }
-
- /**
- * A JDBC connection owned by a single worker thread, together with lazily
- * prepared read/write statements. On a connection-level error the holder
- * is invalidated so the next operation transparently reconnects.
- */
- private final class WorkerConnection implements AutoCloseable {
- private Connection connection;
- private PreparedStatement readStmt;
- private PreparedStatement writeStmt;
-
- private Connection connection() throws SQLException {
- if (connection == null || connection.isClosed()) {
- connection = openConnection();
- readStmt = null;
- writeStmt = null;
- }
- return connection;
- }
-
- private PreparedStatement readStmt() throws SQLException {
- Connection conn = connection();
- if (readStmt == null) {
- readStmt = conn.prepareStatement(String.format(READ_QUERY_TEMPLATE, tablePath));
- }
- return readStmt;
- }
-
- private PreparedStatement writeStmt() throws SQLException {
- Connection conn = connection();
- if (writeStmt == null) {
- writeStmt = conn.prepareStatement(String.format(WRITE_QUERY_TEMPLATE, tablePath));
- }
- return writeStmt;
- }
-
- void read(long id, long hash, int timeoutSeconds) throws SQLException {
- PreparedStatement stmt = readStmt();
- stmt.setQueryTimeout(timeoutSeconds);
- stmt.setLong(1, id);
- stmt.setLong(2, hash);
- try (ResultSet rs = stmt.executeQuery()) {
- // Touch the result set so we exercise the deserialization path.
- while (rs.next()) {
- rs.getLong("id");
- }
- }
- }
-
- void write(Row row, long hash, int timeoutSeconds) throws SQLException {
- PreparedStatement stmt = writeStmt();
- stmt.setQueryTimeout(timeoutSeconds);
- stmt.setLong(1, hash);
- stmt.setLong(2, row.id());
- stmt.setString(3, row.payloadStr());
- stmt.setDouble(4, row.payloadDouble());
- stmt.setTimestamp(5, Timestamp.from(row.payloadTimestamp()));
- stmt.setLong(6, row.payloadHash());
- stmt.executeUpdate();
- }
-
- void invalidateOnConnectionError(SQLException e) {
- if (isConnectionError(e)) {
- close();
- }
- }
-
- @Override
- public void close() {
- closeQuietly(readStmt);
- closeQuietly(writeStmt);
- closeQuietly(connection);
- readStmt = null;
- writeStmt = null;
- connection = null;
- }
-
- private void closeQuietly(AutoCloseable closeable) {
- if (closeable == null) {
- return;
- }
- try {
- closeable.close();
- } catch (Exception ignored) {
- // best-effort cleanup
- }
- }
- }
-}
diff --git a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/KvWorkloadParams.java b/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/KvWorkloadParams.java
deleted file mode 100644
index 41f2156..0000000
--- a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/KvWorkloadParams.java
+++ /dev/null
@@ -1,115 +0,0 @@
-package tech.ydb.slo.kv;
-
-import com.beust.jcommander.Parameter;
-
-/**
- * Tunable parameters for the KV workload.
- *
- * Defaults match the SLO workloads in the Go and JavaScript SDKs so the
- * runs are comparable. JCommander annotations let the operator override
- * any field from the command line, e.g.
- * {@code --read-rps 500 --write-rps 50}.
- */
-@SuppressWarnings("FieldMayBeFinal")
-public final class KvWorkloadParams {
-
- @Parameter(
- names = {"--read-rps"},
- description = "Target read operations per second"
- )
- private int readRps = 1000;
-
- @Parameter(
- names = {"--write-rps"},
- description = "Target write operations per second"
- )
- private int writeRps = 100;
-
- @Parameter(
- names = {"--read-timeout-ms"},
- description = "Per-attempt read timeout in milliseconds"
- )
- private int readTimeoutMs = 10_000;
-
- @Parameter(
- names = {"--write-timeout-ms"},
- description = "Per-attempt write timeout in milliseconds"
- )
- private int writeTimeoutMs = 10_000;
-
- @Parameter(
- names = {"--prefill-count"},
- description = "Number of rows to prefill before the run phase"
- )
- private int prefillCount = 1_000;
-
- @Parameter(
- names = {"--partition-size"},
- description = "Auto-partitioning partition size in MB"
- )
- private int partitionSizeMb = 1;
-
- @Parameter(
- names = {"--min-partition-count"},
- description = "Minimum number of table partitions"
- )
- private int minPartitionCount = 6;
-
- @Parameter(
- names = {"--max-partition-count"},
- description = "Maximum number of table partitions"
- )
- private int maxPartitionCount = 1_000;
-
- @Parameter(
- names = {"--duration"},
- description = "Run duration in seconds (overrides WORKLOAD_DURATION when > 0)"
- )
- private int durationSeconds = 0;
-
- public int readRps() {
- return readRps;
- }
-
- public int writeRps() {
- return writeRps;
- }
-
- public int readTimeoutMs() {
- return readTimeoutMs;
- }
-
- public int writeTimeoutMs() {
- return writeTimeoutMs;
- }
-
- public int prefillCount() {
- return prefillCount;
- }
-
- public int partitionSizeMb() {
- return partitionSizeMb;
- }
-
- public int minPartitionCount() {
- return minPartitionCount;
- }
-
- public int maxPartitionCount() {
- return maxPartitionCount;
- }
-
- /**
- * Effective run duration. If the CLI flag was omitted (left at 0), falls
- * back to the value supplied via the {@code WORKLOAD_DURATION} environment
- * variable through {@code Config}.
- * @return Effective run duration value
- */
- public int durationSeconds() {
- return durationSeconds;
- }
-
- public void setDurationSeconds(int durationSeconds) {
- this.durationSeconds = durationSeconds;
- }
-}
diff --git a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/Row.java b/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/Row.java
deleted file mode 100644
index a53ba89..0000000
--- a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/Row.java
+++ /dev/null
@@ -1,64 +0,0 @@
-package tech.ydb.slo.kv;
-
-import java.time.Instant;
-
-/**
- * A single row of the KV workload table.
- *
- * The schema mirrors the one used by SLO workloads in other YDB SDKs
- * (Go, JavaScript):
- * The {@code hash} primary-key column is derived from {@code id}; the JDBC
- * workload computes it on the client (see {@code KvWorkload}) so reads and
- * writes target the same key without relying on server-side YQL builtins in
- * parameterized statements.
- */
-public final class Row {
- private final long id;
- private final String payloadStr;
- private final double payloadDouble;
- private final Instant payloadTimestamp;
- private final long payloadHash;
-
- public Row(
- long id,
- String payloadStr,
- double payloadDouble,
- Instant payloadTimestamp,
- long payloadHash
- ) {
- this.id = id;
- this.payloadStr = payloadStr;
- this.payloadDouble = payloadDouble;
- this.payloadTimestamp = payloadTimestamp;
- this.payloadHash = payloadHash;
- }
-
- public long id() {
- return id;
- }
-
- public String payloadStr() {
- return payloadStr;
- }
-
- public double payloadDouble() {
- return payloadDouble;
- }
-
- public Instant payloadTimestamp() {
- return payloadTimestamp;
- }
-
- public long payloadHash() {
- return payloadHash;
- }
-}
diff --git a/slo-workload/pom.xml b/slo-workload/pom.xml
index b89ae23..a43b80f 100644
--- a/slo-workload/pom.xml
+++ b/slo-workload/pom.xml
@@ -17,29 +17,35 @@
Reads connection details and run parameters from environment variables
- * (see {@link Config}), parses workload-specific flags from the command line
- * (see {@link KvWorkloadParams}), and runs the KV workload phases — setup,
- * run, teardown — pushing metrics to the OTLP endpoint configured by the YDB
- * SLO action runtime.
- *
- * Exit codes:
- * Metrics emitted (matching the contract from
- * {@code ydb-platform/ydb-slo-action}):
- * Every metric carries the {@code ref} label so the report action can
- * separate current and baseline series.
- */
-public final class Metrics implements AutoCloseable {
-
- public enum OperationType {
- READ("read"),
- WRITE("write");
-
- private final String label;
-
- OperationType(String label) {
- this.label = label;
- }
-
- public String label() {
- return label;
- }
- }
-
- public enum OperationStatus {
- SUCCESS("success"),
- ERROR("error");
-
- private final String label;
-
- OperationStatus(String label) {
- this.label = label;
- }
-
- public String label() {
- return label;
- }
- }
-
- private static final AttributeKey The workload creates a partitioned table, prefills it with rows, and then
- * runs read and write loops at fixed RPS for the configured duration. Each
- * operation is timed and retried via {@link SessionRetryContext}; the outcome
- * is recorded into {@link Metrics} so the SLO action can compare current and
- * baseline runs.
- *
- * Schema and queries mirror the KV workloads in the Go and JavaScript SDKs
- * so the produced metrics are directly comparable across SDKs.
- *
- * Concurrency model: each operation type (read / write) gets a dedicated
- * thread pool sized to the configured RPS. Every worker thread pulls a permit
- * from a shared Guava {@link RateLimiter} and executes the operation inline.
- * There is no separate driver thread and no work queue, which removes the
- * unbounded backlog risk under chaos and keeps the worker count proportional
- * to the actual concurrency budget.
- */
-public final class KvWorkload {
- private static final Logger logger = LoggerFactory.getLogger(KvWorkload.class);
-
- private static final String CREATE_TABLE_QUERY_TEMPLATE = ""
- + "CREATE TABLE IF NOT EXISTS `%s` ("
- + " hash Uint64,"
- + " id Uint64,"
- + " payload_str Utf8,"
- + " payload_double Double,"
- + " payload_timestamp Timestamp,"
- + " payload_hash Uint64,"
- + " PRIMARY KEY (hash, id)"
- + ") WITH ("
- + " UNIFORM_PARTITIONS = %d,"
- + " AUTO_PARTITIONING_BY_SIZE = ENABLED,"
- + " AUTO_PARTITIONING_PARTITION_SIZE_MB = %d,"
- + " AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = %d,"
- + " AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = %d"
- + ")";
-
- private static final String DROP_TABLE_QUERY_TEMPLATE = "DROP TABLE `%s`";
-
- private static final String WRITE_QUERY_TEMPLATE = ""
- + "DECLARE $id AS Uint64;"
- + "DECLARE $payload_str AS Utf8;"
- + "DECLARE $payload_double AS Double;"
- + "DECLARE $payload_timestamp AS Timestamp;"
- + "DECLARE $payload_hash AS Uint64;"
- + "UPSERT INTO `%s` ("
- + " id, hash, payload_str, payload_double, payload_timestamp, payload_hash"
- + ") VALUES ("
- + " $id,"
- + " Digest::NumericHash($id),"
- + " $payload_str,"
- + " $payload_double,"
- + " $payload_timestamp,"
- + " $payload_hash"
- + ");";
-
- private static final String READ_QUERY_TEMPLATE = ""
- + "DECLARE $id AS Uint64;"
- + "SELECT id, payload_str, payload_double, payload_timestamp, payload_hash"
- + " FROM `%s`"
- + " WHERE id = $id AND hash = Digest::NumericHash($id);";
-
- /*
- * Hard cap on the number of worker threads spawned for a single operation
- * type. The SLO targets a few hundred RPS in CI; allowing more workers
- * than this just wastes threads on JIT-warmup contention without
- * improving throughput.
- */
- private static final int MAX_WORKERS = 64;
-
- /*
- * Extra time, on top of the workload duration, given to worker pools to
- * complete their last in-flight operations before {@link #run()} forces
- * shutdown. Picked to be larger than the default per-attempt timeout so
- * a request that started just before the deadline can finish cleanly.
- */
- private static final long SHUTDOWN_GRACE_SECONDS = 30L;
-
- private final SessionRetryContext retryCtx;
- private final Metrics metrics;
- private final KvWorkloadParams params;
- private final String tablePath;
-
- private final RowGenerator generator;
-
- public KvWorkload(QueryClient queryClient, Metrics metrics, KvWorkloadParams params, String tablePath) {
- this.retryCtx = SessionRetryContext.create(queryClient).build();
- this.metrics = metrics;
- this.params = params;
- this.tablePath = tablePath;
- this.generator = new RowGenerator(params.prefillCount());
- }
-
- /*
- * Creates the table (if missing) and prefills it with
- * {@code params.prefillCount()} rows. Prefill uses a fixed-size thread pool
- * so we don't open thousands of sessions in parallel on slow runners.
- */
- public void setup() throws InterruptedException {
- logger.info("creating table {}", tablePath);
- Status createStatus = retryCtx.supplyResult(session ->
- session.createQuery(
- String.format(
- CREATE_TABLE_QUERY_TEMPLATE,
- tablePath,
- params.minPartitionCount(),
- params.partitionSizeMb(),
- params.minPartitionCount(),
- params.maxPartitionCount()
- ),
- TxMode.NONE
- ).execute()
- ).join().getStatus();
- createStatus.expectSuccess("failed to create table " + tablePath);
- logger.info("table {} created", tablePath);
-
- logger.info("prefilling {} rows into {}", params.prefillCount(), tablePath);
- int parallelism = Math.min(MAX_WORKERS, Math.max(1, Math.min(params.prefillCount(), MAX_WORKERS)));
- ExecutorService prefillPool = Executors.newFixedThreadPool(
- parallelism, namedThreadFactory("slo-prefill-")
- );
- try {
- List Read and write workers run concurrently on dedicated thread pools.
- * Each worker pulls a permit from its rate limiter and executes the
- * operation inline, so there is no shared work queue and no driver
- * thread. Sub-zero RPS disables the corresponding loop entirely (useful
- * for write-only or read-only smoke tests).
- */
- public void run() throws InterruptedException {
- long durationSeconds = params.durationSeconds();
- long endNanos = durationSeconds > 0
- ? System.nanoTime() + TimeUnit.SECONDS.toNanos(durationSeconds)
- : Long.MAX_VALUE;
-
- // Track how many writes have completed so reads target a key-space
- // that's actually been populated. The generator itself was
- // constructed with nextId = prefillCount, so writes pick up where
- // prefill left off.
- AtomicLong writesIssued = new AtomicLong();
-
- int readWorkers = workerCount(params.readRps());
- int writeWorkers = workerCount(params.writeRps());
-
- if (readWorkers == 0 && writeWorkers == 0) {
- logger.warn("both read and write RPS are <= 0, run phase has nothing to do");
- return;
- }
-
- ExecutorService readPool = null;
- ExecutorService writePool = null;
- try {
- if (readWorkers > 0) {
- readPool = Executors.newFixedThreadPool(readWorkers, namedThreadFactory("slo-read-"));
- RateLimiter readLimiter = RateLimiter.create(params.readRps());
- for (int i = 0; i < readWorkers; i++) {
- readPool.execute(() -> workerLoop(
- endNanos, readLimiter,
- () -> readOnce(writesIssued.get()),
- "read"
- ));
- }
- } else {
- logger.info("read RPS <= 0, skipping read workers");
- }
-
- if (writeWorkers > 0) {
- writePool = Executors.newFixedThreadPool(writeWorkers, namedThreadFactory("slo-write-"));
- RateLimiter writeLimiter = RateLimiter.create(params.writeRps());
- for (int i = 0; i < writeWorkers; i++) {
- writePool.execute(() -> workerLoop(
- endNanos, writeLimiter,
- () -> {
- writeOnce(generator.generate());
- writesIssued.incrementAndGet();
- },
- "write"
- ));
- }
- } else {
- logger.info("write RPS <= 0, skipping write workers");
- }
-
- // Wait for workers to drain naturally as they hit the deadline.
- // shutdown() lets in-flight ops finish; awaitTermination caps the
- // wait at duration + grace so the run phase can't hang past the
- // configured budget. Workers are stopped via shutdownNow() in
- // the finally block if they exceed the grace window.
- long graceNanos = TimeUnit.SECONDS.toNanos(SHUTDOWN_GRACE_SECONDS);
- long waitNanos = durationSeconds > 0
- ? Math.max(0L, endNanos - System.nanoTime()) + graceNanos
- : Long.MAX_VALUE;
-
- if (readPool != null) {
- readPool.shutdown();
- }
- if (writePool != null) {
- writePool.shutdown();
- }
-
- long readWaitNanos = waitNanos;
- if (readPool != null) {
- long started = System.nanoTime();
- if (!readPool.awaitTermination(readWaitNanos, TimeUnit.NANOSECONDS)) {
- logger.warn("read pool did not drain within deadline, forcing shutdown");
- readPool.shutdownNow();
- }
- waitNanos = Math.max(0L, waitNanos - (System.nanoTime() - started));
- }
- if (writePool != null) {
- if (!writePool.awaitTermination(waitNanos, TimeUnit.NANOSECONDS)) {
- logger.warn("write pool did not drain within deadline, forcing shutdown");
- writePool.shutdownNow();
- }
- }
- } finally {
- forceShutdown(readPool, "read pool");
- forceShutdown(writePool, "write pool");
- }
- }
-
- /*
- * Drops the workload table. Called from the {@code finally} block in
- * {@code Main} so the database is left clean even on failure.
- */
- public void teardown() {
- logger.info("dropping table {}", tablePath);
- Status status = retryCtx.supplyResult(session ->
- session.createQuery(
- String.format(DROP_TABLE_QUERY_TEMPLATE, tablePath),
- TxMode.NONE
- ).execute()
- ).join().getStatus();
- if (!status.isSuccess()) {
- logger.warn("failed to drop table {}: {}", tablePath, status);
- } else {
- logger.info("table {} dropped", tablePath);
- }
- }
-
- // --- internals ---------------------------------------------------------
-
- /*
- * Loops on a single worker thread until the deadline or interruption,
- * pacing each iteration through the shared rate limiter and running the
- * operation inline. No work queue is involved — backpressure comes
- * naturally from the limiter blocking the worker.
- */
- private void workerLoop(long endNanos, RateLimiter limiter, Runnable singleOp, String name) {
- while (System.nanoTime() < endNanos && !Thread.currentThread().isInterrupted()) {
- limiter.acquire();
- try {
- singleOp.run();
- } catch (Throwable t) {
- logger.warn("{} op threw unexpectedly: {}", name, t.toString());
- }
- }
- }
-
- /*
- * Computes the number of worker threads for a given RPS target.
- * Returns 0 for non-positive RPS so the caller skips the loop entirely.
- */
- private static int workerCount(int rps) {
- if (rps <= 0) {
- return 0;
- }
- return Math.min(MAX_WORKERS, Math.max(1, rps));
- }
-
- /*
- * Picks a random id in [0, keyspaceUpper) and reads it back from the table.
- * Reads target only ids known to exist (the prefilled range plus rows
- * written so far during this run), so a successful read always returns
- * a row and exercises the deserialization path.
- */
- private void readOnce(long writesObserved) {
- long upperBound = Math.max(1L, params.prefillCount() + writesObserved);
- long id = ThreadLocalRandom.current().nextLong(upperBound);
-
- Metrics.Span span = metrics.startOperation(Metrics.OperationType.READ);
- AtomicInteger attempts = new AtomicInteger();
- ExecuteQuerySettings settings = ExecuteQuerySettings.newBuilder()
- .withRequestTimeout(Duration.ofMillis(params.readTimeoutMs()))
- .build();
-
- Result Each row gets a monotonically increasing {@code id} and a random payload.
- * The {@code hash} column is computed server-side via
- * {@code Digest::NumericHash($id)} at insert time, so it is not carried on
- * the client. The format mirrors the SLO workloads in the Go and JS SDKs so
- * the resulting tables are interchangeable.
- */
-public final class RowGenerator {
- private static final int MIN_PAYLOAD_LENGTH = 20;
- private static final int MAX_PAYLOAD_LENGTH = 40;
-
- private final AtomicLong nextId;
-
- public RowGenerator(long startId) {
- this.nextId = new AtomicLong(startId);
- }
-
- /**
- * Generates a new row with a fresh monotonically increasing id.
- * @return a new row
- */
- public Row generate() {
- long id = nextId.getAndIncrement();
- return generate(id);
- }
-
- /**
- * Generates a row with an explicit id (used during prefill to control IDs).
- * @param id
- * @return a new row
- */
- public static Row generate(long id) {
- long payloadHash = ThreadLocalRandom.current().nextLong();
- double payloadDouble = ThreadLocalRandom.current().nextDouble();
- String payloadStr = randomPayloadString();
- Instant payloadTimestamp = Instant.now();
-
- return new Row(id, payloadStr, payloadDouble, payloadTimestamp, payloadHash);
- }
-
- private static String randomPayloadString() {
- int length = MIN_PAYLOAD_LENGTH
- + ThreadLocalRandom.current().nextInt(MAX_PAYLOAD_LENGTH - MIN_PAYLOAD_LENGTH + 1);
- byte[] bytes = new byte[length];
- ThreadLocalRandom.current().nextBytes(bytes);
- return Base64.getEncoder().withoutPadding().encodeToString(bytes);
- }
-}
- * hash Uint64 (primary key, computed server-side via Digest::NumericHash(id))
- * id Uint64 (primary key)
- * payload_str Utf8
- * payload_double Double
- * payload_timestamp Timestamp
- * payload_hash Uint64
- *
- *
- *
- *
- *
- *
- *
- */
public final class Main {
- private static final Logger logger = LoggerFactory.getLogger(Main.class);
-
private static final String YDB_DRIVER_CLASS = "tech.ydb.jdbc.YdbDriver";
private Main() {
- // utility class
+
}
public static void main(String[] args) {
- Config config;
try {
- config = Config.fromEnv();
- } catch (IllegalStateException e) {
- logger.error("invalid environment configuration: {}", e.getMessage());
- System.exit(2);
- return;
- }
-
- KvWorkloadParams params = new KvWorkloadParams();
- try {
- JCommander.newBuilder()
- .programName("ydb-slo-jdbc-workload")
- .acceptUnknownOptions(true)
- .addObject(params)
- .build()
- .parse(args);
- } catch (ParameterException e) {
- logger.error("invalid CLI arguments: {}", e.getMessage());
- System.exit(2);
- return;
- }
-
- // CLI duration takes precedence over WORKLOAD_DURATION when supplied.
- if (params.durationSeconds() <= 0) {
- params.setDurationSeconds(config.durationSeconds());
- }
-
- try {
- // The driver auto-registers via the JDBC SPI, but loading it
- // explicitly fails fast with a clear message if it's missing.
Class.forName(YDB_DRIVER_CLASS);
} catch (ClassNotFoundException e) {
- logger.error("YDB JDBC driver not found on classpath: {}", YDB_DRIVER_CLASS);
+ org.slf4j.LoggerFactory.getLogger(Main.class)
+ .error("YDB JDBC driver not found on classpath: {}", YDB_DRIVER_CLASS);
System.exit(1);
return;
}
- logger.info("starting SLO workload: name={}, ref={}, duration={}s, readRps={}, writeRps={}, url={}",
- config.workloadName(),
- config.ref(),
- params.durationSeconds(),
- params.readRps(),
- params.writeRps(),
- config.jdbcUrl());
-
- // The table path embeds workload name and ref so concurrent runs of
- // the current and baseline images don't step on each other. Both
- // components are sanitized: WORKLOAD_NAME comes from the action input
- // and is normally already safe, but we don't trust user input to be
- // a valid YDB identifier.
- String tablePath = sanitize(config.workloadName()) + "_" + sanitize(config.ref());
-
- Properties connectionProperties = new Properties();
- if (config.token() != null && !config.token().isEmpty()) {
- connectionProperties.setProperty("token", config.token());
- }
-
- int exitCode = 0;
- Metrics metrics = Metrics.create(config);
- KvWorkload workload = new KvWorkload(
- config.jdbcUrl(), connectionProperties, metrics, params, tablePath
+ Launcher.launch(
+ "ydb-slo-jdbc-workload",
+ "java-jdbc-kv",
+ args,
+ (config, params, tablePath) -> new JdbcKvClient(config, params, tablePath)
);
-
- try {
- workload.setup();
- workload.run();
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- logger.warn("workload interrupted");
- exitCode = 1;
- } catch (Throwable t) {
- logger.error("workload failed", t);
- exitCode = 1;
- } finally {
- try {
- workload.teardown();
- } catch (Throwable t) {
- logger.warn("teardown failed", t);
- }
-
- try {
- metrics.flush();
- } catch (Throwable t) {
- logger.warn("metrics flush failed", t);
- }
-
- closeQuietly(metrics, "metrics");
- }
-
- System.exit(exitCode);
- }
-
- private static void closeQuietly(AutoCloseable closeable, String name) {
- if (closeable == null) {
- return;
- }
- try {
- closeable.close();
- } catch (Throwable t) {
- logger.warn("failed to close {}: {}", name, t.toString());
- }
- }
-
- /**
- * Replaces characters that aren't valid in YDB table names with underscores.
- * Refs from CI may include slashes ({@code release/1.2}) or dots, which
- * the action permits in metrics labels but YDB rejects in table paths.
- */
- private static String sanitize(String value) {
- StringBuilder sb = new StringBuilder(value.length());
- for (int i = 0; i < value.length(); i++) {
- char c = value.charAt(i);
- if (Character.isLetterOrDigit(c) || c == '_') {
- sb.append(c);
- } else {
- sb.append('_');
- }
- }
- return sb.toString();
}
}
diff --git a/slo-workload/jdbc/src/main/java/tech/ydb/slo/jdbc/JdbcKvClient.java b/slo-workload/jdbc/src/main/java/tech/ydb/slo/jdbc/JdbcKvClient.java
new file mode 100644
index 0000000..970ff60
--- /dev/null
+++ b/slo-workload/jdbc/src/main/java/tech/ydb/slo/jdbc/JdbcKvClient.java
@@ -0,0 +1,268 @@
+package tech.ydb.slo.jdbc;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.SQLNonTransientConnectionException;
+import java.sql.SQLRecoverableException;
+import java.sql.SQLTransientConnectionException;
+import java.sql.SQLTransientException;
+import java.sql.Statement;
+import java.sql.Timestamp;
+import java.util.Properties;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import tech.ydb.jdbc.exception.YdbStatusable;
+import tech.ydb.slo.core.Config;
+import tech.ydb.slo.core.kv.KvClient;
+import tech.ydb.slo.core.kv.KvSchema;
+import tech.ydb.slo.core.kv.KvSession;
+import tech.ydb.slo.core.kv.KvWorkloadParams;
+import tech.ydb.slo.core.kv.OpOutcome;
+import tech.ydb.slo.core.kv.Row;
+import tech.ydb.slo.core.kv.RowGenerator;
+
+public final class JdbcKvClient implements KvClient {
+ private static final Logger logger = LoggerFactory.getLogger(JdbcKvClient.class);
+
+ private static final long INITIAL_BACKOFF_MS = 10L;
+ private static final long MAX_BACKOFF_MS = 1_000L;
+
+ private final String jdbcUrl;
+ private final Properties connectionProperties;
+ private final String tablePath;
+ private final int maxAttempts;
+
+ public JdbcKvClient(Config config, KvWorkloadParams params, String tablePath) {
+ this.jdbcUrl = config.jdbcUrl();
+ this.tablePath = tablePath;
+ this.maxAttempts = Math.max(1, params.maxAttempts());
+ this.connectionProperties = new Properties();
+ if (config.token() != null && !config.token().isEmpty()) {
+ connectionProperties.setProperty("token", config.token());
+ }
+ }
+
+ @Override
+ public void createTable(KvWorkloadParams params, String tablePath) throws SQLException {
+ try (Connection conn = openConnection();
+ Statement stmt = conn.createStatement()) {
+ stmt.execute(String.format(
+ KvSchema.CREATE_TABLE_TEMPLATE,
+ tablePath,
+ params.minPartitionCount(),
+ params.partitionSizeMb(),
+ params.minPartitionCount(),
+ params.maxPartitionCount()
+ ));
+ }
+ }
+
+ @Override
+ public void dropTable(String tablePath) {
+ try (Connection conn = openConnection();
+ Statement stmt = conn.createStatement()) {
+ stmt.execute(String.format(KvSchema.DROP_TABLE_TEMPLATE, tablePath));
+ } catch (SQLException e) {
+ logger.warn("failed to drop table {}: {}", tablePath, e.toString());
+ }
+ }
+
+ @Override
+ public KvSession openSession() throws SQLException {
+ return new JdbcKvSession();
+ }
+
+ private Connection openConnection() throws SQLException {
+ return DriverManager.getConnection(jdbcUrl, connectionProperties);
+ }
+
+ private final class JdbcKvSession implements KvSession {
+ private Connection connection;
+ private PreparedStatement readStmt;
+ private PreparedStatement writeStmt;
+
+ @Override
+ public OpOutcome read(long id, int timeoutMs) {
+ long hash = RowGenerator.numericHash(id);
+ int attempts = 0;
+ SQLException last = null;
+ while (attempts < maxAttempts) {
+ attempts++;
+ try {
+ read(id, hash, timeoutSeconds(timeoutMs));
+ return OpOutcome.success(attempts - 1);
+ } catch (SQLException e) {
+ last = e;
+ if (!isRetryable(e) || attempts >= maxAttempts) {
+ break;
+ }
+ invalidateOnConnectionError(e);
+ if (!backoff(attempts)) {
+
+
+
+ break;
+ }
+ }
+ }
+ return OpOutcome.error(attempts - 1, classifyError(last));
+ }
+
+ @Override
+ public OpOutcome write(Row row, int timeoutMs) {
+ long hash = RowGenerator.numericHash(row.id());
+ int attempts = 0;
+ SQLException last = null;
+ while (attempts < maxAttempts) {
+ attempts++;
+ try {
+ write(row, hash, timeoutSeconds(timeoutMs));
+ return OpOutcome.success(attempts - 1);
+ } catch (SQLException e) {
+ last = e;
+ if (!isRetryable(e) || attempts >= maxAttempts) {
+ break;
+ }
+ invalidateOnConnectionError(e);
+ if (!backoff(attempts)) {
+ break;
+ }
+ }
+ }
+ return OpOutcome.error(attempts - 1, classifyError(last));
+ }
+
+ @Override
+ public void close() {
+ closeQuietly(readStmt);
+ closeQuietly(writeStmt);
+ closeQuietly(connection);
+ readStmt = null;
+ writeStmt = null;
+ connection = null;
+ }
+
+ private Connection connection() throws SQLException {
+ if (connection == null || connection.isClosed()) {
+ connection = openConnection();
+ readStmt = null;
+ writeStmt = null;
+ }
+ return connection;
+ }
+
+ private PreparedStatement readStmt() throws SQLException {
+ Connection conn = connection();
+ if (readStmt == null) {
+ readStmt = conn.prepareStatement(String.format(KvSchema.SELECT_TEMPLATE, tablePath));
+ }
+ return readStmt;
+ }
+
+ private PreparedStatement writeStmt() throws SQLException {
+ Connection conn = connection();
+ if (writeStmt == null) {
+ writeStmt = conn.prepareStatement(String.format(KvSchema.UPSERT_TEMPLATE, tablePath));
+ }
+ return writeStmt;
+ }
+
+ private void read(long id, long hash, int timeoutSeconds) throws SQLException {
+ PreparedStatement stmt = readStmt();
+ stmt.setQueryTimeout(timeoutSeconds);
+ stmt.setLong(1, id);
+ stmt.setLong(2, hash);
+ try (ResultSet rs = stmt.executeQuery()) {
+ while (rs.next()) {
+ rs.getLong("id");
+ }
+ }
+ }
+
+ private void write(Row row, long hash, int timeoutSeconds) throws SQLException {
+ PreparedStatement stmt = writeStmt();
+ stmt.setQueryTimeout(timeoutSeconds);
+ stmt.setLong(1, hash);
+ stmt.setLong(2, row.id());
+ stmt.setString(3, row.payloadStr());
+ stmt.setDouble(4, row.payloadDouble());
+ stmt.setTimestamp(5, Timestamp.from(row.payloadTimestamp()));
+ stmt.setLong(6, row.payloadHash());
+ stmt.executeUpdate();
+ }
+
+ private void invalidateOnConnectionError(SQLException e) {
+ if (isConnectionError(e)) {
+ close();
+ }
+ }
+ }
+
+ private static int timeoutSeconds(int timeoutMs) {
+ return Math.max(1, (timeoutMs + 999) / 1000);
+ }
+
+
+
+ private static boolean isRetryable(SQLException e) {
+ if (e instanceof YdbStatusable) {
+ try {
+ return ((YdbStatusable) e).getStatus().getCode().isRetryable(true);
+ } catch (RuntimeException ignored) {
+
+ }
+ }
+ return e instanceof SQLRecoverableException || e instanceof SQLTransientException;
+ }
+
+ private static boolean isConnectionError(SQLException e) {
+ return e instanceof SQLRecoverableException
+ || e instanceof SQLTransientConnectionException
+ || e instanceof SQLNonTransientConnectionException;
+ }
+
+ private static String classifyError(SQLException e) {
+ if (e == null) {
+ return "unknown";
+ }
+ if (e instanceof YdbStatusable) {
+ try {
+ return "ydb/" + ((YdbStatusable) e).getStatus().getCode().name().toLowerCase();
+ } catch (RuntimeException ignored) {
+
+ }
+ }
+ return e.getClass().getSimpleName().toLowerCase();
+ }
+
+
+
+ private static boolean backoff(int attempt) {
+ long delay = Math.min(MAX_BACKOFF_MS, INITIAL_BACKOFF_MS * (1L << Math.min(attempt - 1, 20)));
+ try {
+ Thread.sleep(delay);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ private static void closeQuietly(AutoCloseable closeable) {
+ if (closeable == null) {
+ return;
+ }
+ try {
+ closeable.close();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ } catch (Exception ignored) {
+
+ }
+ }
+}
diff --git a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/KvWorkload.java b/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/KvWorkload.java
deleted file mode 100644
index dd74c32..0000000
--- a/slo-workload/jdbc/src/main/java/tech/ydb/slo/kv/KvWorkload.java
+++ /dev/null
@@ -1,566 +0,0 @@
-package tech.ydb.slo.kv;
-
-import java.sql.Connection;
-import java.sql.DriverManager;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.SQLRecoverableException;
-import java.sql.SQLTransientConnectionException;
-import java.sql.SQLTransientException;
-import java.sql.Statement;
-import java.sql.Timestamp;
-import java.util.Properties;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.ThreadLocalRandom;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
-
-import com.google.common.util.concurrent.RateLimiter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import tech.ydb.jdbc.exception.YdbStatusable;
-import tech.ydb.slo.Metrics;
-
-/**
- * Key-value workload for the SLO test, driving the YDB JDBC driver.
- *
- *
- * hash Uint64 (primary key, derived from id)
- * id Uint64 (primary key)
- * payload_str Utf8
- * payload_double Double
- * payload_timestamp Timestamp
- * payload_hash Uint64
- *
- *
- *
- *
- */
-public final class Main {
- private static final Logger logger = LoggerFactory.getLogger(Main.class);
-
- private Main() {
- // utility class
- }
-
- public static void main(String[] args) {
- Config config;
- try {
- config = Config.fromEnv();
- } catch (IllegalStateException e) {
- logger.error("invalid environment configuration: {}", e.getMessage());
- System.exit(2);
- return;
- }
-
- KvWorkloadParams params = new KvWorkloadParams();
- try {
- JCommander.newBuilder()
- .programName("ydb-slo-workload")
- .acceptUnknownOptions(true)
- .addObject(params)
- .build()
- .parse(args);
- } catch (ParameterException e) {
- logger.error("invalid CLI arguments: {}", e.getMessage());
- System.exit(2);
- return;
- }
-
- // CLI duration takes precedence over WORKLOAD_DURATION when supplied.
- if (params.durationSeconds() <= 0) {
- params.setDurationSeconds(config.durationSeconds());
- }
-
- logger.info("starting SLO workload: name={}, ref={}, duration={}s, readRps={}, writeRps={}",
- config.workloadName(),
- config.ref(),
- params.durationSeconds(),
- params.readRps(),
- params.writeRps());
-
- // The table path embeds workload name and ref so concurrent runs of
- // the current and baseline images don't step on each other. Both
- // components are sanitized: WORKLOAD_NAME comes from the action input
- // and is normally already safe, but we don't trust user input to be
- // a valid YDB identifier.
- String tablePath = sanitize(config.workloadName()) + "_" + sanitize(config.ref());
-
- int exitCode = 0;
- Metrics metrics = Metrics.create(config);
- AuthProvider provider = NopAuthProvider.INSTANCE;
- if (config.token() != null && !config.token().isEmpty()) {
- provider = new TokenAuthProvider(config.token());
- }
- GrpcTransport transport = GrpcTransport.forConnectionString(config.connectionString())
- .withAuthProvider(provider)
- .build();
- QueryClient queryClient = QueryClient.newClient(transport).build();
-
- KvWorkload workload = new KvWorkload(queryClient, metrics, params, tablePath);
-
- try {
- workload.setup();
- workload.run();
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- logger.warn("workload interrupted");
- exitCode = 1;
- } catch (Throwable t) {
- logger.error("workload failed", t);
- exitCode = 1;
- } finally {
- try {
- workload.teardown();
- } catch (Throwable t) {
- logger.warn("teardown failed", t);
- }
-
- try {
- metrics.flush();
- } catch (Throwable t) {
- logger.warn("metrics flush failed", t);
- }
-
- closeQuietly(metrics, "metrics");
- closeQuietly(queryClient, "query client");
- closeQuietly(transport, "transport");
- }
-
- System.exit(exitCode);
- }
-
- private static void closeQuietly(AutoCloseable closeable, String name) {
- if (closeable == null) {
- return;
- }
- try {
- closeable.close();
- } catch (Throwable t) {
- logger.warn("failed to close {}: {}", name, t.toString());
- }
- }
-
- /**
- * Replaces characters that aren't valid in YDB table names with underscores.
- * Refs from CI may include slashes ({@code release/1.2}) or dots, which
- * the action permits in metrics labels but YDB rejects in table paths.
- */
- private static String sanitize(String value) {
- StringBuilder sb = new StringBuilder(value.length());
- for (int i = 0; i < value.length(); i++) {
- char c = value.charAt(i);
- if (Character.isLetterOrDigit(c) || c == '_') {
- sb.append(c);
- } else {
- sb.append('_');
- }
- }
- return sb.toString();
- }
-}
diff --git a/slo/src/main/java/tech/ydb/slo/Metrics.java b/slo/src/main/java/tech/ydb/slo/Metrics.java
deleted file mode 100644
index 767a747..0000000
--- a/slo/src/main/java/tech/ydb/slo/Metrics.java
+++ /dev/null
@@ -1,381 +0,0 @@
-package tech.ydb.slo;
-
-import java.time.Duration;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.TimeUnit;
-
-import io.opentelemetry.api.common.AttributeKey;
-import io.opentelemetry.api.common.Attributes;
-import io.opentelemetry.api.metrics.LongCounter;
-import io.opentelemetry.api.metrics.LongUpDownCounter;
-import io.opentelemetry.api.metrics.Meter;
-import io.opentelemetry.api.metrics.ObservableDoubleMeasurement;
-import io.opentelemetry.exporter.otlp.http.metrics.OtlpHttpMetricExporter;
-import io.opentelemetry.sdk.metrics.SdkMeterProvider;
-import io.opentelemetry.sdk.metrics.SdkMeterProviderBuilder;
-import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader;
-import io.opentelemetry.sdk.resources.Resource;
-import org.HdrHistogram.AtomicHistogram;
-import org.HdrHistogram.Histogram;
-
-/**
- * Collects and pushes SLO workload metrics to the OTLP endpoint configured by
- * the YDB SLO action runtime.
- *
- *
- *
- *
- *