enable topic slo

Author: vgvoleg

.github/workflows/slo.yml

@@ -57,6 +57,24 @@ jobs:
               --read-timeout 1000 \
               --write-timeout 1000
             cleanup-args: grpc://localhost:2135 /Root/testdb
+          - prefix: topic
+            workload: topic-basic
+            create-args: |
+              grpc://localhost:2135 /Root/testdb \
+              --path /Root/testdb/slo_topic \
+              --partitions-count 10
+            run-args: |
+              grpc://localhost:2135 /Root/testdb \
+              --path /Root/testdb/slo_topic \
+              --partitions-count 10
+              --prom-pgw localhost:9091 \
+              --report-period 250 \
+              --time ${{inputs.slo_workload_duration_seconds || 600}} \
+              --read-rps ${{inputs.slo_workload_read_max_rps || 100}} \
+              --write-rps ${{inputs.slo_workload_write_max_rps || 100}} \
+              --read-timeout 5000 \
+              --write-timeout 5000
+            cleanup-args: grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic
 
 
     concurrency:

tests/slo/TOPIC_SLO.md

@@ -97,8 +97,8 @@ python -m src topic-run grpc://localhost:2135 /local \
    - `cleanup_topic()` - очистка топика
 
 2. **metrics.py** - расширен для поддержки топик метрик:
-   - `OP_TYPE_TOPIC_WRITE` - метрики записи
-   - `OP_TYPE_TOPIC_READ` - метрики чтения
+   - `OP_TYPE_WRITE` - метрики записи
+   - `OP_TYPE_READ` - метрики чтения
 
 3. **options.py** - добавлена команда `topic-run` с параметрами для топиков
 

tests/slo/src/core/metrics.py

@@ -10,7 +10,6 @@
 from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram, push_to_gateway  # noqa: E402
 
 OP_TYPE_READ, OP_TYPE_WRITE = "read", "write"
-OP_TYPE_TOPIC_READ, OP_TYPE_TOPIC_WRITE = "topic_read", "topic_write"
 OP_STATUS_SUCCESS, OP_STATUS_FAILURE = "success", "err"
 
 REF = environ.get("REF", "main")

tests/slo/src/jobs/base.py

@@ -4,12 +4,14 @@
 import time
 from ratelimiter import RateLimiter
 
+import ydb
+
 logger = logging.getLogger(__name__)
 
 
 class BaseJobManager(ABC):
     def __init__(self, driver, args, metrics):
-        self.driver = driver
+        self.driver: ydb.Driver = driver
         self.args = args
         self.metrics = metrics
 

tests/slo/src/jobs/topic_jobs.py

@@ -5,7 +5,7 @@
 from ratelimiter import RateLimiter
 
 from .base import BaseJobManager
-from core.metrics import OP_TYPE_TOPIC_READ, OP_TYPE_TOPIC_WRITE
+from core.metrics import OP_TYPE_READ, OP_TYPE_WRITE
 
 logger = logging.getLogger(__name__)
 
@@ -24,14 +24,14 @@ def run_tests(self):
     def _run_topic_write_jobs(self):
         logger.info("Start topic write jobs")
 
-        write_limiter = RateLimiter(max_calls=self.args.topic_write_rps, period=1)
+        write_limiter = RateLimiter(max_calls=self.args.write_rps, period=1)
 
         futures = []
-        for i in range(self.args.topic_write_threads):
+        for i in range(self.args.write_threads):
             future = threading.Thread(
                 name=f"slo_topic_write_{i}",
                 target=self._run_topic_writes,
-                args=(write_limiter,),
+                args=(write_limiter, i, ),
             )
             future.start()
             futures.append(future)
@@ -41,29 +41,30 @@ def _run_topic_write_jobs(self):
     def _run_topic_read_jobs(self):
         logger.info("Start topic read jobs")
 
-        read_limiter = RateLimiter(max_calls=self.args.topic_read_rps, period=1)
+        read_limiter = RateLimiter(max_calls=self.args.read_rps, period=1)
 
         futures = []
-        for i in range(self.args.topic_read_threads):
+        for i in range(self.args.read_threads):
             future = threading.Thread(
                 name=f"slo_topic_read_{i}",
                 target=self._run_topic_reads,
-                args=(read_limiter,),
+                args=(read_limiter, ),
             )
             future.start()
             futures.append(future)
 
         return futures
 
-    def _run_topic_writes(self, limiter):
+    def _run_topic_writes(self, limiter, partition_id = None):
         start_time = time.time()
         logger.info("Start topic write workload")
 
-        writer = self.driver.topic_client.writer(self.args.topic_path, codec=ydb.TopicCodec.GZIP)
-        logger.info("Topic writer created")
-
-        try:
-            write_session = writer.__enter__()
+        with self.driver.topic_client.writer(
+            self.args.path,
+            codec=ydb.TopicCodec.GZIP,
+            partition_id=partition_id,
+        ) as writer:
+            logger.info("Topic writer created")
 
             message_count = 0
             while time.time() - start_time < self.args.time:
@@ -72,47 +73,44 @@ def _run_topic_writes(self, limiter):
 
                     content = f"message_{message_count}_{threading.current_thread().name}".encode("utf-8")
 
-                    if len(content) < self.args.topic_message_size:
-                        content += b"x" * (self.args.topic_message_size - len(content))
+                    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_TOPIC_WRITE,))
+                    ts = self.metrics.start((OP_TYPE_WRITE,))
                     try:
-                        write_session.write(message)
-                        self.metrics.stop((OP_TYPE_TOPIC_WRITE,), ts)
+                        writer.write(message)
+                        if message_count % 100 == 0:
+                            writer.flush(timeout=self.args.write_timeout / 1000)
+
+                        self.metrics.stop((OP_TYPE_WRITE,), ts)
                     except Exception as e:
-                        self.metrics.stop((OP_TYPE_TOPIC_WRITE,), ts, error=e)
+                        self.metrics.stop((OP_TYPE_WRITE,), ts, error=e)
                         logger.error("Write error: %s", e)
 
-        finally:
-            writer.__exit__(None, None, None)
-
         logger.info("Stop topic write workload")
 
     def _run_topic_reads(self, limiter):
         start_time = time.time()
         logger.info("Start topic read workload")
 
-        reader = self.driver.topic_client.reader(self.args.topic_consumer, self.args.topic_path)
-        logger.info("Topic reader created")
-
-        try:
-            read_session = reader.__enter__()
+        with self.driver.topic_client.reader(
+            self.args.consumer,
+            self.args.path,
+        ) as reader:
+            logger.info("Topic reader created")
 
             while time.time() - start_time < self.args.time:
                 with limiter:
-                    ts = self.metrics.start((OP_TYPE_TOPIC_READ,))
+                    ts = self.metrics.start((OP_TYPE_READ,))
                     try:
-                        batch = read_session.receive_message(timeout=self.args.topic_read_timeout / 1000)
-                        if batch is not None:
-                            read_session.commit_offset(batch.batches[-1].message_offset_end)
-                        self.metrics.stop((OP_TYPE_TOPIC_READ,), ts)
+                        msg = reader.receive_message(timeout=self.args.read_timeout / 1000)
+                        if msg is not None:
+                            reader.commit(msg)
+                        self.metrics.stop((OP_TYPE_READ,), ts)
                     except Exception as e:
-                        self.metrics.stop((OP_TYPE_TOPIC_READ,), ts, error=e)
+                        self.metrics.stop((OP_TYPE_READ,), ts, error=e)
                         logger.debug("Read timeout or error: %s", e)
 
-        finally:
-            reader.__exit__(None, None, None)
-
         logger.info("Stop topic read workload")

tests/slo/src/options.py

@@ -61,46 +61,41 @@ def make_table_cleanup_parser(subparsers):
     add_common_options(table_cleanup_parser)
 
 
+def make_topic_create_parser(subparsers):
+    topic_create_parser = subparsers.add_parser("topic-create", help="Create topic with consumer")
+    add_common_options(topic_create_parser)
+
+    topic_create_parser.add_argument("--path", default="/local/slo_topic", type=str, help="Topic path")
+    topic_create_parser.add_argument("--consumer", default="slo_consumer", type=str, help="Topic consumer name")
+    topic_create_parser.add_argument("--partitions-count", default=1, type=int, help="Partition count")
+
+
 def make_topic_run_parser(subparsers):
-    """Создает парсер для команды topic-run - запуск SLO тестов для топиков"""
     topic_parser = subparsers.add_parser("topic-run", help="Run topic SLO workload")
     add_common_options(topic_parser)
 
-    topic_parser.add_argument("--topic-read-rps", default=50, type=int, help="Topic read request rps")
-    topic_parser.add_argument("--topic-read-timeout", default=5000, type=int, help="Topic read timeout [ms]")
-    topic_parser.add_argument("--topic-write-rps", default=20, type=int, help="Topic write request rps")
-    topic_parser.add_argument("--topic-write-timeout", default=10000, type=int, help="Topic write timeout [ms]")
-    topic_parser.add_argument("--topic-read-threads", default=1, type=int, help="Number of threads for topic reading")
-    topic_parser.add_argument("--topic-write-threads", default=1, type=int, help="Number of threads for topic writing")
-    topic_parser.add_argument("--topic-path", default="/local/slo_topic", type=str, help="Topic path")
-    topic_parser.add_argument("--topic-consumer", default="slo_consumer", type=str, help="Topic consumer name")
-    topic_parser.add_argument("--topic-message-size", default=100, type=int, help="Topic message size in bytes")
-    topic_parser.add_argument("--topic-min-partitions", default=1, type=int, help="Minimum active partitions")
-    topic_parser.add_argument("--topic-max-partitions", default=10, type=int, help="Maximum active partitions")
-    topic_parser.add_argument("--topic-retention-hours", default=24, type=int, help="Retention period in hours")
-
-    topic_parser.add_argument("--time", default=60, type=int, help="Time to run in seconds")
+    topic_parser.add_argument("--path", default="/local/slo_topic", type=str, help="Topic path")
+    topic_parser.add_argument("--consumer", default="slo_consumer", type=str, help="Topic consumer name")
+    topic_parser.add_argument("--partitions-count", default=1, type=int, help="Partition count")
+    topic_parser.add_argument("--read-rps", default=100, type=int, help="Topic read request rps")
+    topic_parser.add_argument("--read-timeout", default=5000, type=int, help="Topic read timeout [ms]")
+    topic_parser.add_argument("--write-rps", default=100, type=int, help="Topic write request rps")
+    topic_parser.add_argument("--write-timeout", default=5000, type=int, help="Topic write timeout [ms]")
+    topic_parser.add_argument("--read-threads", default=1, type=int, help="Number of threads for topic reading")
+    topic_parser.add_argument("--write-threads", default=1, type=int, help="Number of threads for topic writing")
+    topic_parser.add_argument("--message-size", default=100, type=int, help="Topic message size in bytes")
+
+    topic_parser.add_argument("--time", default=10, type=int, help="Time to run in seconds")
     topic_parser.add_argument("--shutdown-time", default=10, type=int, help="Graceful shutdown time in seconds")
-    topic_parser.add_argument("--prom-pgw", default="", type=str, help="Prometheus push gateway (empty to disable)")
+    topic_parser.add_argument("--prom-pgw", default="localhost:9091", type=str, help="Prometheus push gateway (empty to disable)")
     topic_parser.add_argument("--report-period", default=1000, type=int, help="Prometheus push period in [ms]")
 
 
-def make_topic_create_parser(subparsers):
-    topic_create_parser = subparsers.add_parser("topic-create", help="Create topic with consumer")
-    add_common_options(topic_create_parser)
-
-    topic_create_parser.add_argument("--topic-path", default="/local/slo_topic", type=str, help="Topic path")
-    topic_create_parser.add_argument("--topic-consumer", default="slo_consumer", type=str, help="Topic consumer name")
-    topic_create_parser.add_argument("--topic-min-partitions", default=1, type=int, help="Minimum active partitions")
-    topic_create_parser.add_argument("--topic-max-partitions", default=10, type=int, help="Maximum active partitions")
-    topic_create_parser.add_argument("--topic-retention-hours", default=24, type=int, help="Retention period in hours")
-
-
 def make_topic_cleanup_parser(subparsers):
     topic_cleanup_parser = subparsers.add_parser("topic-cleanup", help="Drop topic")
     add_common_options(topic_cleanup_parser)
 
-    topic_cleanup_parser.add_argument("--topic-path", default="/local/slo_topic", type=str, help="Topic path")
+    topic_cleanup_parser.add_argument("--path", default="/local/slo_topic", type=str, help="Topic path")
 
 
 def get_root_parser():