async slo jobs for local testing

Author: vgvoleg

tests/slo/requirements.txt

@@ -1,4 +1,5 @@
 requests==2.28.2
 ratelimiter==1.2.0.post0
+aiolimiter==1.1.0
 prometheus-client==0.17.0
 quantile-estimator==0.1.2

tests/slo/slo_runner.sh

@@ -3,6 +3,7 @@ docker compose -f playground/configs/compose.yaml up -d --wait
 
 ../../.venv/bin/python ./src topic-create grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic
 
-../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --time 10
+# ../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --write-rps 1 --time 600 --async
 
-../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --write-threads 0 --read-rps 1 --debug --time 600
+../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --time 10
+../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --write-threads 0 --read-rps 1 --time 600 

tests/slo/src/jobs/async_topic_jobs.py

@@ -0,0 +1,130 @@
+import asyncio
+import ydb.aio
+import time
+import logging
+from aiolimiter import AsyncLimiter
+
+from .base import BaseJobManager
+from core.metrics import OP_TYPE_READ, OP_TYPE_WRITE
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncTopicJobManager(BaseJobManager):
+    def __init__(self, driver, args, metrics):
+        super().__init__(driver, args, metrics)
+        self.driver: ydb.aio.Driver = driver
+
+    async def run_tests(self):
+        tasks = [
+            *await self._run_topic_write_jobs(),
+            *await self._run_topic_read_jobs(),
+            *self._run_metric_job(),
+        ]
+
+        await asyncio.gather(*tasks)
+
+    async def _run_topic_write_jobs(self):
+        logger.info("Start async topic write jobs")
+
+        write_limiter = AsyncLimiter(max_rate=self.args.write_rps, time_period=1)
+
+        tasks = []
+        for i in range(self.args.write_threads):
+            task = asyncio.create_task(self._run_topic_writes(write_limiter, i), name=f"slo_topic_write_{i}")
+            tasks.append(task)
+
+        return tasks
+
+    async def _run_topic_read_jobs(self):
+        logger.info("Start async topic read jobs")
+
+        read_limiter = AsyncLimiter(max_rate=self.args.read_rps, time_period=1)
+
+        tasks = []
+        for i in range(self.args.read_threads):
+            task = asyncio.create_task(self._run_topic_reads(read_limiter), name=f"slo_topic_read_{i}")
+            tasks.append(task)
+
+        return tasks
+
+    async def _run_topic_writes(self, limiter, partition_id=None):
+        start_time = time.time()
+        logger.info("Start async topic write workload")
+
+        async with self.driver.topic_client.writer(
+            self.args.path,
+            codec=ydb.TopicCodec.GZIP,
+            partition_id=partition_id,
+        ) as writer:
+            logger.info("Async topic writer created")
+
+            message_count = 0
+            while time.time() - start_time < self.args.time:
+                async with limiter:
+                    message_count += 1
+
+                    content = f"message_{message_count}_{asyncio.current_task().get_name()}".encode("utf-8")
+
+                    if len(content) < self.args.message_size:
+                        content += b"x" * (self.args.message_size - len(content))
+
+                    message = ydb.TopicWriterMessage(data=content)
+
+                    ts = self.metrics.start((OP_TYPE_WRITE,))
+                    try:
+                        await writer.write_with_ack(message)
+                        logger.info("Write message: %s", content)
+                        self.metrics.stop((OP_TYPE_WRITE,), ts)
+                    except Exception as e:
+                        self.metrics.stop((OP_TYPE_WRITE,), ts, error=e)
+                        logger.error("Write error: %s", e)
+
+        logger.info("Stop async topic write workload")
+
+    async def _run_topic_reads(self, limiter):
+        start_time = time.time()
+        logger.info("Start async topic read workload")
+
+        async with self.driver.topic_client.reader(
+            self.args.path,
+            self.args.consumer,
+        ) as reader:
+            logger.info("Async topic reader created")
+
+            while time.time() - start_time < self.args.time:
+                async with limiter:
+                    ts = self.metrics.start((OP_TYPE_READ,))
+                    try:
+                        msg = await reader.receive_message()
+                        if msg is not None:
+                            logger.info("Read message: %s", msg.data.decode())
+                            await reader.commit_with_ack(msg)
+
+                        self.metrics.stop((OP_TYPE_READ,), ts)
+                    except Exception as e:
+                        self.metrics.stop((OP_TYPE_READ,), ts, error=e)
+                        logger.error("Read error: %s", e)
+
+        logger.info("Stop async topic read workload")
+
+    def _run_metric_job(self):
+        if not self.args.prom_pgw:
+            return []
+
+        # Create async task for metrics
+        task = asyncio.create_task(self._async_metric_sender(self.args.time), name="slo_metrics_sender")
+        return [task]
+
+    async def _async_metric_sender(self, runtime):
+        start_time = time.time()
+        logger.info("Start push metrics (async)")
+
+        limiter = AsyncLimiter(max_rate=10**6 // self.args.report_period, time_period=1)
+
+        while time.time() - start_time < runtime:
+            async with limiter:
+                # Call sync metrics.push() in executor to avoid blocking
+                await asyncio.get_event_loop().run_in_executor(None, self.metrics.push)
+
+        logger.info("Stop push metrics (async)")

tests/slo/src/options.py

@@ -6,6 +6,7 @@ def add_common_options(parser):
     parser.add_argument("db", help="YDB database name")
     parser.add_argument("-t", "--table-name", default="key_value", help="Table name")
     parser.add_argument("--debug", action="store_true", help="Enable debug logging")
+    parser.add_argument("--async", action="store_true", help="Use async mode for operations")
 
 
 def make_table_create_parser(subparsers):

tests/slo/src/root_runner.py

@@ -1,4 +1,6 @@
+import asyncio
 import ydb
+import ydb.aio
 import logging
 from typing import Dict
 
@@ -26,6 +28,15 @@ def run_command(self, args):
             raise ValueError(f"Unknown prefix: {prefix}. Available: {list(self.runners.keys())}")
 
         runner_instance = self.runners[prefix]()
+
+        # Check if async mode is requested and command is 'run'
+        if getattr(args, "async", False) and command == "run":
+            asyncio.run(self._run_async_command(args, runner_instance, command))
+        else:
+            self._run_sync_command(args, runner_instance, command)
+
+    def _run_sync_command(self, args, runner_instance, command):
+        """Run command in synchronous mode"""
         driver_config = ydb.DriverConfig(
             args.endpoint,
             database=args.db,
@@ -43,13 +54,33 @@ def run_command(self, args):
                 elif command == "cleanup":
                     runner_instance.cleanup(args)
                 else:
-                    raise RuntimeError(f"Unknown command {command} for prefix {prefix}")
+                    raise RuntimeError(f"Unknown command {command} for prefix {runner_instance.prefix}")
             except BaseException:
                 logger.exception("Something went wrong")
                 raise
             finally:
                 driver.stop(timeout=getattr(args, "shutdown_time", 10))
 
+    async def _run_async_command(self, args, runner_instance, command):
+        """Run command in asynchronous mode"""
+        driver_config = ydb.DriverConfig(
+            args.endpoint,
+            database=args.db,
+            grpc_keep_alive_timeout=5000,
+        )
+
+        async with ydb.aio.Driver(driver_config) as driver:
+            await driver.wait(timeout=300)
+            try:
+                runner_instance.set_driver(driver)
+                if command == "run":
+                    await runner_instance.run_async(args)
+                else:
+                    raise RuntimeError(f"Async mode only supports 'run' command, got '{command}'")
+            except BaseException:
+                logger.exception("Something went wrong in async mode")
+                raise
+
 
 def create_runner() -> SLORunner:
     runner = SLORunner()

tests/slo/src/runners/base.py

@@ -24,6 +24,9 @@ def create(self, args):
     def run(self, args):
         pass
 
+    async def run_async(self, args):
+        raise NotImplementedError(f"Async mode not supported for {self.prefix}")
+
     @abstractmethod
     def cleanup(self, args):
         pass

tests/slo/src/runners/topic_runner.py

@@ -1,8 +1,10 @@
 import time
 import ydb
+import ydb.aio
 
 from .base import BaseRunner
 from jobs.topic_jobs import TopicJobManager
+from jobs.async_topic_jobs import AsyncTopicJobManager
 from core.metrics import create_metrics
 
 
@@ -76,6 +78,21 @@ def run(self, args):
         if hasattr(metrics, "reset"):
             metrics.reset()
 
+    async def run_async(self, args):
+        """Async version of topic SLO tests using ydb.aio.Driver"""
+        metrics = create_metrics(args.prom_pgw)
+
+        self.logger.info("Starting async topic SLO tests")
+
+        # Use async driver for topic operations
+        job_manager = AsyncTopicJobManager(self.driver, args, metrics)
+        await job_manager.run_tests()
+
+        self.logger.info("Async topic SLO tests completed")
+
+        if hasattr(metrics, "reset"):
+            metrics.reset()
+
     def cleanup(self, args):
         self.logger.info("Cleaning up topic: %s", args.path)
 

ydb/_grpc/grpcwrapper/common_utils.py

@@ -165,7 +165,7 @@ def __init__(self, convert_server_grpc_to_wrapper):
         self._connection_state = "new"
         self._stream_call = None
         self._wait_executor = None
-        self._on_disconnected_lambda = None
+        self._on_disconnected_callback = None
 
         self._stream_settings: BaseRequestSettings = (
             BaseRequestSettings()
@@ -197,32 +197,32 @@ def _clean_executor(self, wait: bool):
 
     async def _start_asyncio_driver(self, driver: DriverIO, stub, method):
         requests_iterator = QueueToIteratorAsyncIO(self.from_client_grpc)
-        stream_call, on_disconnected_lambda = await driver(
+        stream_call, on_disconnected_callback = await driver(
             requests_iterator,
             stub,
             method,
             settings=self._stream_settings,
-            include_disconnected_lambda_to_result=True,
+            include_disconnected_callback_to_result=True,
         )
         self._stream_call = stream_call
-        self._on_disconnected_lambda = on_disconnected_lambda
+        self._on_disconnected_callback = on_disconnected_callback
         self.from_server_grpc = stream_call.__aiter__()
 
     async def _start_sync_driver(self, driver: Driver, stub, method):
         requests_iterator = AsyncQueueToSyncIteratorAsyncIO(self.from_client_grpc)
         self._wait_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 
-        stream_call, on_disconnected_lambda = await to_thread(
+        stream_call, on_disconnected_callback = await to_thread(
             driver,
             requests_iterator,
             stub,
             method,
             executor=self._wait_executor,
             settings=self._stream_settings,
-            include_disconnected_lambda_to_result=True,
+            include_disconnected_callback_to_result=True,
         )
         self._stream_call = stream_call
-        self._on_disconnected_lambda = on_disconnected_lambda
+        self._on_disconnected_callback = on_disconnected_callback
         self.from_server_grpc = SyncToAsyncIterator(stream_call.__iter__(), self._wait_executor)
 
     async def receive(self, timeout: Optional[int] = None) -> Any:
@@ -238,8 +238,8 @@ async def get_response():
                 grpc_message = await asyncio.wait_for(get_response(), timeout)
 
         except (grpc.RpcError, grpc.aio.AioRpcError) as e:
-            if self._on_disconnected_lambda:
-                coro = self._on_disconnected_lambda()
+            if self._on_disconnected_callback:
+                coro = self._on_disconnected_callback()
                 if asyncio.iscoroutine(coro):
                     await coro
 

ydb/aio/connection.py

@@ -235,7 +235,7 @@ async def close(self, grace: float = 30):
         been terminated are cancelled. If grace is None, this method will wait until all tasks are finished.
         :return: None
         """
-        logger.info("Closing channel for endpoint %s", self.endpoint)
+        logger.debug("Closing channel for endpoint %s", self.endpoint)
 
         self.closing = True
 

ydb/aio/pool.py

@@ -259,7 +259,7 @@ async def __call__(
         wrap_args=(),
         preferred_endpoint=None,
         fast_fail=False,
-        include_disconnected_lambda_to_result=False,
+        include_disconnected_callback_to_result=False,
     ):
         if self._stopped:
             raise issues.Error("Driver was stopped")
@@ -283,6 +283,6 @@ async def __call__(
             on_disconnected,
         )
 
-        if include_disconnected_lambda_to_result:
+        if include_disconnected_callback_to_result:
             return res, on_disconnected
         return res