SLO for query service

Author: vgvoleg

tests/slo/src/jobs.py

@@ -3,7 +3,7 @@
 import logging
 import dataclasses
 from random import randint
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 from ratelimiter import RateLimiter
 
 import threading
@@ -36,7 +36,7 @@
 
 @dataclasses.dataclass
 class RequestParams:
-    pool: ydb.SessionPool
+    pool: Union[ydb.SessionPool, ydb.QuerySessionPool]
     query: str
     params: dict
     metrics: Metrics
@@ -56,7 +56,7 @@ def transaction(session):
 
         result = session.transaction().execute(
             params.query,
-            params.params,
+            parameters=params.params,
             commit_tx=True,
             settings=params.request_settings,
         )
@@ -135,6 +135,61 @@ def run_read_jobs(args, driver, tb_name, max_id, metrics):
     return futures
 
 
+def run_reads_query(driver, query, max_id, metrics, limiter, runtime, timeout):
+    start_time = time.time()
+
+    logger.info("Start read workload")
+
+    request_settings = ydb.BaseRequestSettings().with_timeout(timeout)
+    retry_setting = ydb.RetrySettings(
+        idempotent=True,
+        max_session_acquire_timeout=timeout,
+    )
+
+    with ydb.QuerySessionPool(driver) as pool:
+        logger.info("Session pool for read requests created")
+
+        while time.time() - start_time < runtime:
+            params = {"$object_id": randint(1, max_id)}
+            with limiter:
+
+                def check_result(result):
+                    res = next(result)
+                    assert res[0].rows[0]
+
+                params = RequestParams(
+                    pool=pool,
+                    query=query,
+                    params=params,
+                    metrics=metrics,
+                    labels=(JOB_READ_LABEL,),
+                    request_settings=request_settings,
+                    retry_settings=retry_setting,
+                    check_result_cb=check_result,
+                )
+                execute_query(params)
+
+        logger.info("Stop read workload")
+
+
+def run_read_jobs_query(args, driver, tb_name, max_id, metrics):
+    logger.info("Start read jobs for query service")
+
+    read_q = READ_QUERY_TEMPLATE.format(tb_name)
+
+    read_limiter = RateLimiter(max_calls=args.read_rps, period=1)
+    futures = []
+    for _ in range(args.read_threads):
+        future = threading.Thread(
+            name="slo_run_read",
+            target=run_reads_query,
+            args=(driver, read_q, max_id, metrics, read_limiter, args.time, args.read_timeout / 1000),
+        )
+        future.start()
+        futures.append(future)
+    return futures
+
+
 def run_writes(driver, query, row_generator, metrics, limiter, runtime, timeout):
     start_time = time.time()
 
@@ -157,6 +212,11 @@ def run_writes(driver, query, row_generator, metrics, limiter, runtime, timeout)
                 "$payload_double": row.payload_double,
                 "$payload_timestamp": row.payload_timestamp,
             }
+
+            def check_result(result):
+                with result:
+                    pass
+
             with limiter:
                 params = RequestParams(
                     pool=pool,
@@ -166,6 +226,7 @@ def run_writes(driver, query, row_generator, metrics, limiter, runtime, timeout)
                     labels=(JOB_WRITE_LABEL,),
                     request_settings=request_settings,
                     retry_settings=retry_setting,
+                    check_result_cb=check_result,
                 )
                 execute_query(params)
 
@@ -194,6 +255,63 @@ def run_write_jobs(args, driver, tb_name, max_id, metrics):
     return futures
 
 
+def run_writes_query(driver, query, row_generator, metrics, limiter, runtime, timeout):
+    start_time = time.time()
+
+    logger.info("Start write workload")
+
+    request_settings = ydb.BaseRequestSettings().with_timeout(timeout)
+    retry_setting = ydb.RetrySettings(
+        idempotent=True,
+        max_session_acquire_timeout=timeout,
+    )
+
+    with ydb.QuerySessionPool(driver) as pool:
+        logger.info("Session pool for read requests created")
+
+        while time.time() - start_time < runtime:
+            row = row_generator.get()
+            params = {
+                "$object_id": (row.object_id, ydb.PrimitiveType.Int64),
+                "$payload_str": (row.payload_str, ydb.PrimitiveType.Utf8),
+                "$payload_double": (row.payload_double, ydb.PrimitiveType.Double),
+                "$payload_timestamp": (row.payload_timestamp, ydb.PrimitiveType.Timestamp),
+            }
+            with limiter:
+                params = RequestParams(
+                    pool=pool,
+                    query=query,
+                    params=params,
+                    metrics=metrics,
+                    labels=(JOB_WRITE_LABEL,),
+                    request_settings=request_settings,
+                    retry_settings=retry_setting,
+                )
+                execute_query(params)
+
+        logger.info("Stop write workload")
+
+
+def run_write_jobs_query(args, driver, tb_name, max_id, metrics):
+    logger.info("Start write jobs for query service")
+
+    write_q = WRITE_QUERY_TEMPLATE.format(tb_name)
+
+    write_limiter = RateLimiter(max_calls=args.write_rps, period=1)
+    row_generator = RowGenerator(max_id)
+
+    futures = []
+    for _ in range(args.write_threads):
+        future = threading.Thread(
+            name="slo_run_write",
+            target=run_writes_query,
+            args=(driver, write_q, row_generator, metrics, write_limiter, args.time, args.write_timeout / 1000),
+        )
+        future.start()
+        futures.append(future)
+    return futures
+
+
 def push_metric(limiter, runtime, metrics):
     start_time = time.time()
     logger.info("Start push metrics")

tests/slo/src/runner.py

@@ -7,7 +7,13 @@
 import concurrent.futures
 from concurrent.futures import ThreadPoolExecutor
 
-from jobs import run_read_jobs, run_write_jobs, run_metric_job
+from jobs import (
+    run_read_jobs,
+    run_write_jobs,
+    run_read_jobs_query,
+    run_write_jobs_query,
+    run_metric_job,
+)
 from metrics import Metrics, SDK_SERVICE_NAME
 
 logger = logging.getLogger(__name__)
@@ -93,8 +99,8 @@ def run_slo(args, driver, tb_name):
         )
     elif SDK_SERVICE_NAME == "query-service":
         futures = (
-            *run_read_jobs(args, driver, tb_name, max_id, metrics),
-            *run_write_jobs(args, driver, tb_name, max_id, metrics),
+            *run_read_jobs_query(args, driver, tb_name, max_id, metrics),
+            *run_write_jobs_query(args, driver, tb_name, max_id, metrics),
             run_metric_job(args, metrics),
         )
     else:

ydb/query/transaction.py

@@ -370,6 +370,7 @@ def execute(
         exec_mode: Optional[base.QueryExecMode] = None,
         parameters: Optional[dict] = None,
         concurrent_result_sets: Optional[bool] = False,
+        settings: dict = None,  # TODO: temporary hack
     ) -> base.SyncResponseContextIterator:
         """WARNING: This API is experimental and could be changed.