viewer stress test (#26479)

Author: adameat

ydb/tests/stress/viewer/__main__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+import argparse
+from ydb.tests.stress.viewer.workload import Workload
+
+if __name__ == '__main__':
+    text = """\033[92mViewer workload\x1b[0m"""
+    parser = argparse.ArgumentParser(description=text, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('--mon_endpoint', default='http://localhost:8765', help="An endpoint to be used")
+    parser.add_argument('--database', default=None, help='A database to connect')
+    parser.add_argument('--duration', default=10 ** 9, type=lambda x: int(x), help='A duration of workload in seconds.')
+    args = parser.parse_args()
+    with Workload(args.mon_endpoint, args.database, args.duration) as workload:
+        workload.start()

ydb/tests/stress/viewer/tests/test_workload.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+import pytest
+import os
+import yatest
+from ydb.tests.library.stress.fixtures import StressFixture
+
+
+class TestYdbWorkload(StressFixture):
+    @pytest.fixture(autouse=True, scope="function")
+    def setup(self):
+        yield from self.setup_cluster()
+
+    def test(self):
+        if not self.mon_endpoint:
+            self.mon_endpoint = "http://localhost:8765"
+        cmd = [
+            yatest.common.binary_path(os.getenv("YDB_TEST_PATH")),
+            "--mon_endpoint", self.mon_endpoint,
+            "--database", self.database,
+            "--duration", "60",
+        ]
+        yatest.common.execute(cmd, wait=True)

ydb/tests/stress/viewer/tests/ya.make

@@ -0,0 +1,25 @@
+PY3TEST()
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/ydbd_dep.inc)
+ENV(YDB_TEST_PATH="ydb/tests/stress/viewer/viewer")
+
+TEST_SRCS(
+    test_workload.py
+)
+
+IF (SANITIZER_TYPE)
+    REQUIREMENTS(ram:32)
+ENDIF()
+
+SIZE(MEDIUM)
+
+DEPENDS(
+    ydb/tests/stress/viewer
+)
+
+PEERDIR(
+    ydb/tests/library
+    ydb/tests/library/stress
+    ydb/tests/stress/viewer/workload
+)
+
+END()

ydb/tests/stress/viewer/workload/__init__.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+import time
+import threading
+import queue
+import traceback
+import requests
+from urllib.parse import urljoin
+from concurrent.futures import ThreadPoolExecutor
+
+class Workload:
+    def __init__(self, mon_endpoint, database, duration):
+        self.mon_endpoint = mon_endpoint
+        self.database = database
+        self.round_size = 1000
+        self.duration = duration
+        self.delayed_events = queue.Queue()
+        self.pool_semaphore = threading.BoundedSemaphore(value=100)
+        self.worker_exception = []
+        self.pool = None
+
+    def wrapper(self, f, *args, **kwargs):
+        try:
+            if len(self.worker_exception) == 0:
+                result = f(*args, **kwargs)
+                if result.status_code != 200:
+                    raise AssertionError(f"Bad response code {result.status_code}, text: {result.text}")
+        except Exception:
+            with self.lock:
+                self.worker_exception.append(traceback.format_exc())
+        finally:
+            self.pool_semaphore.release()
+
+    def call_viewer_api_post(self, url, body=None, headers=None):
+        if body is None:
+            body = {}
+        if self.database:
+            body["database"] = self.database
+        return requests.post(self.mon_endpoint + url, json=body, headers=headers)
+
+    def loop(self):
+        started_at = time.time()
+
+        while time.time() - started_at < self.duration:
+            yield self.call_viewer_api_post, "/viewer/capabilities"
+            yield self.call_viewer_api_post, "/viewer/whoami"
+            yield self.call_viewer_api_post, "/viewer/nodelist"
+            yield self.call_viewer_api_post, "/viewer/nodes?fields_required=all"
+            yield self.call_viewer_api_post, "/viewer/groups?fields_required=all"
+            if self.database:
+                yield self.call_viewer_api_post, "/viewer/describe", {"path": self.database}
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.pool:
+            self.pool.shutdown(wait=True)
+            self.pool = None
+
+    def start(self):
+        self.pool = ThreadPoolExecutor(max_workers=100)
+        for call in self.loop():
+            if len(self.worker_exception) == 0:
+                self.pool_semaphore.acquire()
+                func, *args = call
+                self.pool.submit(self.wrapper, func, *args)
+            else:
+                raise AssertionError(f"Worker exceptions {self.worker_exception}")
+        self.pool.shutdown(wait=True)
+        self.pool = None

ydb/tests/stress/viewer/workload/ya.make

@@ -0,0 +1,11 @@
+PY3_LIBRARY()
+
+PY_SRCS(
+    __init__.py
+)
+
+PEERDIR(
+    contrib/python/requests
+)
+
+END()

ydb/tests/stress/viewer/ya.make

@@ -0,0 +1,15 @@
+PY3_PROGRAM(viewer)
+
+PY_SRCS(
+    __main__.py
+)
+
+PEERDIR(
+    ydb/tests/stress/viewer/workload
+)
+
+END()
+
+RECURSE_FOR_TESTS(
+    tests
+)