From fae3e0152215babdd1c1e4b8c33105330eba7048 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <me@cschwarz.com>
Date: Thu, 9 Nov 2023 15:58:16 +0000
Subject: [PATCH] WIP: performance test that uses the getpage benchmark

---
 setup_bench_repo_dir.bash                  |  41 +++++++
 test_runner/fixtures/neon_fixtures.py      |  31 +++++-
 test_runner/performance/test_pageserver.py | 122 +++++++++++++++++++++
 3 files changed, 188 insertions(+), 6 deletions(-)
 create mode 100644 setup_bench_repo_dir.bash
 create mode 100644 test_runner/performance/test_pageserver.py

diff --git a/setup_bench_repo_dir.bash b/setup_bench_repo_dir.bash
new file mode 100644
index 0000000000..cf1dbfd819
--- /dev/null
+++ b/setup_bench_repo_dir.bash
@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+if [ "$(cat /sys/class/block/nvme1n1/device/model)" != "Amazon EC2 NVMe Instance Storage        " ]; then
+    echo "nvme1n1 is not Amazon EC2 NVMe Instance Storage: '$(cat /sys/class/block/nvme1n1/device/model)'"
+    exit 1
+fi
+
+rmdir bench_repo_dir || true
+
+sudo mkfs.ext4 -E lazy_itable_init=0,lazy_journal_init=0  /dev/nvme1n1
+
+sudo mount /dev/nvme1n1 /mnt
+sudo chown -R "$(id -u)":"$(id -g)" /mnt
+
+mkdir /mnt/bench_repo_dir
+mkdir bench_repo_dir
+sudo mount --bind /mnt/bench_repo_dir bench_repo_dir
+
+mkdir /mnt/test_output
+
+mkdir /mnt/many_tenants
+
+echo run the following commands
+
+cat <<EOF
+    # test suite run
+    export TEST_OUTPUT="/mnt/test_output"
+    DEFAULT_PG_VERSION=15 BUILD_TYPE=release ./scripts/pytest test_runner/performance/test_pageserver.py
+
+    # for interactive use
+    export NEON_REPO_DIR="$(readlink -f ./bench_repo_dir)/repo"
+    cargo build_testing --release
+    ./target/release/neon_local init
+    # ... create tenant, seed it using pgbench
+    # then duplicate the tenant using
+    # poetry run python3 ./test_runner/duplicate_tenant.py TENANT_ID 200 8
+EOF
+
+
diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index 7dfdd9274d..c5c28faad1 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -772,13 +772,10 @@ class NeonEnv:
         self.initial_tenant = config.initial_tenant
         self.initial_timeline = config.initial_timeline
 
+        self.control_plane_api: Optional[str] = None
+        self.attachment_service: Optional[NeonAttachmentService] = None
         if config.enable_generations:
-            attachment_service_port = self.port_distributor.get_port()
-            self.control_plane_api: Optional[str] = f"http://127.0.0.1:{attachment_service_port}"
-            self.attachment_service: Optional[NeonAttachmentService] = NeonAttachmentService(self)
-        else:
-            self.control_plane_api = None
-            self.attachment_service = None
+            self.enable_generations()
 
         # Create a config file corresponding to the options
         cfg: Dict[str, Any] = {
@@ -847,6 +844,18 @@ class NeonEnv:
         log.info(f"Config: {cfg}")
         self.neon_cli.init(cfg)
 
+    def enable_generations(self, start=False):
+        if not start:
+            # TODO: assert that we haven't `self.start()`ed yet
+            pass
+        assert self.control_plane_api is None
+        assert self.attachment_service is None
+        attachment_service_port = self.port_distributor.get_port()
+        self.control_plane_api = f"http://127.0.0.1:{attachment_service_port}"
+        self.attachment_service = NeonAttachmentService(self)
+        if start:
+            self.attachment_service.start()
+
     def start(self):
         # Start up broker, pageserver and all safekeepers
         self.broker.try_start()
@@ -1580,6 +1589,16 @@ class Pagectl(AbstractNeonCli):
         parsed = json.loads(res.stdout)
         return IndexPartDump.from_json(parsed)
 
+# class GetpageBenchLibpq(AbstractNeonCli):
+#     """
+#     A typed wrapper around the `getpage_bench_libpq` CLI.
+#     """
+#
+#     COMMAND = "getpage_bench_libpq"
+#
+#     def run(self):
+#         pass
+
 
 class NeonAttachmentService:
     def __init__(self, env: NeonEnv):
diff --git a/test_runner/performance/test_pageserver.py b/test_runner/performance/test_pageserver.py
new file mode 100644
index 0000000000..4a1c3190aa
--- /dev/null
+++ b/test_runner/performance/test_pageserver.py
@@ -0,0 +1,122 @@
+import json
+import shutil
+import subprocess
+from pathlib import Path
+from typing import List
+
+from fixtures.benchmark_fixture import NeonBenchmarker
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnvBuilder, PgBin, last_flush_lsn_upload
+from fixtures.pageserver.utils import wait_until_tenant_active
+from fixtures.remote_storage import LocalFsStorage, RemoteStorageKind
+from fixtures.types import TenantId
+
+
+def test_getpage_throughput(
+    neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenchmarker, pg_bin: PgBin
+):
+    neon_env_builder.enable_generations = True
+    neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS)
+    env = neon_env_builder.init_start()
+
+    remote_storage = env.pageserver_remote_storage
+    assert isinstance(remote_storage, LocalFsStorage)
+
+    ps_http = env.pageserver.http_client()
+
+    # clean up the useless default tenant
+    ps_http.tenant_delete(env.initial_tenant)
+
+    # create our template tenant
+    tenant_config_mgmt_api = {
+        "gc_period": "0s",
+        "checkpoint_timeout": "3650 day",
+        "compaction_period": "20 s",
+        "compaction_threshold": 10,
+        "compaction_target_size": 134217728,
+        "checkpoint_distance": 268435456,
+        "image_creation_threshold": 3,
+    }
+    tenant_config_cli = {k: str(v) for k, v in tenant_config_mgmt_api.items()}
+
+    template_tenant, template_timeline = env.neon_cli.create_tenant(conf=tenant_config_cli)
+    template_tenant_gen = int(ps_http.tenant_status(template_tenant)["generation"])
+    with env.endpoints.create_start("main", tenant_id=template_tenant) as ep:
+        pg_bin.run_capture(["pgbench", "-i", "-s50", ep.connstr()])
+        last_flush_lsn_upload(env, ep, template_tenant, template_timeline)
+    ps_http.tenant_detach(template_tenant)
+
+    # stop PS just for good measure
+    env.pageserver.stop()
+
+    # duplicate the tenant in remote storage
+    src_timelines_dir: Path = remote_storage.tenant_path(template_tenant) / "timelines"
+    assert src_timelines_dir.is_dir(), f"{src_timelines_dir} is not a directory"
+    tenants = [template_tenant]
+    for i in range(0, 200):
+        new_tenant = TenantId.generate()
+        tenants.append(new_tenant)
+        log.info("Duplicating tenant #%s: %s", i, new_tenant)
+
+        dst_timelines_dir: Path = remote_storage.tenant_path(new_tenant) / "timelines"
+        dst_timelines_dir.parent.mkdir(parents=False, exist_ok=False)
+        dst_timelines_dir.mkdir(parents=False, exist_ok=False)
+
+        for tl in src_timelines_dir.iterdir():
+            src_tl_dir = src_timelines_dir / tl.name
+            assert src_tl_dir.is_dir(), f"{src_tl_dir} is not a directory"
+            dst_tl_dir = dst_timelines_dir / tl.name
+            dst_tl_dir.mkdir(parents=False, exist_ok=False)
+            for file in tl.iterdir():
+                shutil.copy2(file, dst_tl_dir)
+                if "__" in file.name:
+                    cmd: List[str] = [
+                        str(
+                            env.neon_binpath / "pagectl"
+                        ),  # TODO: abstract this like the other binaries
+                        "layer",
+                        "rewrite-summary",
+                        str(dst_tl_dir / file.name),
+                        "--new-tenant-id",
+                        str(new_tenant),
+                    ]
+                    subprocess.run(cmd, check=True)
+                else:
+                    # index_part etc need no patching
+                    pass
+
+    env.pageserver.start()
+    assert ps_http.tenant_list() == []
+    for tenant in tenants:
+        ps_http.tenant_attach(
+            tenant, config=tenant_config_mgmt_api, generation=template_tenant_gen + 1
+        )
+    for tenant in tenants:
+        wait_until_tenant_active(ps_http, tenant)
+
+    # ensure all layers are resident for predictiable performance
+    # TODO: ensure all kinds of eviction are disabled (per-tenant, disk-usage-based)
+    for tenant in tenants:
+        ps_http.download_all_layers(tenant, template_timeline)
+
+    # run the benchmark with one client per timeline, each doing 10k requests to random keys.
+    cmd = [
+        str(env.neon_binpath / "pagebench"),
+        "get-page-latest-lsn",
+        "--mgmt-api-endpoint",
+        ps_http.base_url,
+        "--page-service-connstring",
+        env.pageserver.connstr(password=None),
+        "--runtime",
+        "10s",
+        *[f"{tenant}/{template_timeline}" for tenant in tenants],
+    ]
+    log.info(f"command: {' '.join(cmd)}")
+    basepath = pg_bin.run_capture(cmd)
+    results_path = Path(basepath + ".stdout")
+    log.info(f"Benchmark results at: {results_path}")
+
+    with open(results_path, "r") as f:
+        results = json.load(f)
+
+    log.info(f"Results:\n{json.dumps(results, sort_keys=True, indent=2)}")