From 838a6d304d44aec22ed6e693db6ce7b32bb7c202 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Fri, 5 Jan 2024 13:41:12 +0000
Subject: [PATCH] test_snapshot_dir fixture with marker file for finished
 snapshot

---
 test_runner/fixtures/neon_fixtures.py      | 42 ++++++++++++++++++++--
 test_runner/performance/test_pageserver.py | 22 +++++++-----
 2 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index 82ec504706..04c114d49d 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -3268,8 +3268,8 @@ def get_test_output_dir(request: FixtureRequest, top_output_dir: Path) -> Path:
 def get_test_overlay_dir(request: FixtureRequest, top_output_dir: Path) -> Path:
     return _get_test_dir(request, top_output_dir, "overlay-")
 
-def get_test_snapshot_dir(request: FixtureRequest, top_output_dir: Path) -> Path:
-    return _get_test_dir(request, top_output_dir, "snapshot")
+def get_test_snapshot_dir_path(request: FixtureRequest, top_output_dir: Path) -> Path:
+    return _get_test_dir(request, top_output_dir, "snapshot-")
 
 def get_test_repo_dir(request: FixtureRequest, top_output_dir: Path) -> Path:
     return get_test_output_dir(request, top_output_dir) / "repo"
@@ -3314,6 +3314,44 @@ def test_output_dir(request: FixtureRequest, top_output_dir: Path, test_overlay_
 
     allure_attach_from_dir(test_dir)
 
+class SnapshotDir:
+    _path: Path
+
+    def __init__(self, path: Path):
+        self._path = path
+        if not self._path.exists():
+            self._path.mkdir()
+        else:
+            assert self._path.is_dir()
+
+    @property
+    def path(self) -> Path:
+        return self._path / "snapshot"
+
+    @property
+    def _marker_file_path(self) -> Path:
+        return self._path / "initialized.marker"
+
+    def is_initialized(self):
+        # TODO: in the future, take a `tag` as argument and store it in the marker in set_initialized.
+        # Then, in this function, compare marker file contents with the tag to invalidate the snapshot if the tag changed.
+        return self._marker_file_path.exists()
+
+    def set_initialized(self):
+        self._marker_file_path.write_text("")
+
+
+@pytest.fixture(scope="function", autouse=True)
+def test_snapshot_dir(request: FixtureRequest, top_output_dir: Path, test_overlay_dir: Path) -> SnapshotDir:
+    """Create the working directory for an individual test."""
+
+    _ = test_overlay_dir # overlay mounts use snapshot dir as the lowerdir => request it so it can unmount stale overlay mounts first
+
+    # one directory per test
+    snapshot_dir = get_test_snapshot_dir_path(request, top_output_dir)
+    log.info(f"test_snapshot_dir is {snapshot_dir}")
+    return SnapshotDir(snapshot_dir)
+
 @pytest.fixture(scope="function", autouse=True)
 def test_overlay_dir(request: FixtureRequest, top_output_dir: Path) -> Optional[Path]:
     """Create the overlay state directory for an individual test."""
diff --git a/test_runner/performance/test_pageserver.py b/test_runner/performance/test_pageserver.py
index ef12eb6799..67d646c6f0 100644
--- a/test_runner/performance/test_pageserver.py
+++ b/test_runner/performance/test_pageserver.py
@@ -8,7 +8,7 @@ from typing import List, Tuple
 import pytest
 from fixtures.benchmark_fixture import NeonBenchmarker
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder, PgBin, last_flush_lsn_upload
+from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder, PgBin, SnapshotDir, 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, TimelineId
@@ -17,7 +17,9 @@ from fixtures.types import TenantId, TimelineId
 @pytest.fixture(scope="function")
 @pytest.mark.timeout(1000)
 def snapshotting_env(
-    neon_env_builder: NeonEnvBuilder, pg_bin: PgBin, test_output_dir: Path,
+    neon_env_builder: NeonEnvBuilder,
+    pg_bin: PgBin,
+    test_snapshot_dir: SnapshotDir,
 ) -> Tuple[NeonEnv, TimelineId, List[TenantId]]:
     """
     The fixture prepares environment or restores it from a snapshot.
@@ -28,7 +30,6 @@ def snapshotting_env(
     - if the fixture is executed on CI (it has CI=true in the environment), the snapshot is not saved
     """
 
-    snapshot_dir = test_output_dir.parent / f"snapshot-{test_output_dir.name}"
     save_snapshot = os.getenv("CI", "false") != "true"
 
     neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS)
@@ -44,10 +45,14 @@ def snapshotting_env(
         "image_creation_threshold": 3,
     }
 
-    if snapshot_dir.exists():
-        env = neon_env_builder.from_repo_dir(snapshot_dir)
+
+    if test_snapshot_dir.is_initialized():
+        save_snapshot = False
+        env = neon_env_builder.from_repo_dir(test_snapshot_dir.path)
         ps_http = env.pageserver.http_client()
-        tenants = list({TenantId(t.name) for t in (snapshot_dir.glob("pageserver_*/tenants/*"))})
+        tenants = list(
+            {TenantId(t.name) for t in (test_snapshot_dir.path.glob("pageserver_*/tenants/*"))}
+        )
         template_timeline = env.initial_timeline
 
         neon_env_builder.start()
@@ -128,8 +133,9 @@ def snapshotting_env(
         ps_http.download_all_layers(tenant, template_timeline)
 
     # take snapshot after download all layers so tenant dir restoration is fast
-    if save_snapshot and not snapshot_dir.exists():
-        shutil.copytree(env.repo_dir, snapshot_dir)
+    if save_snapshot:
+        shutil.copytree(env.repo_dir, test_snapshot_dir.path)
+        test_snapshot_dir.set_initialized()
 
     return env, template_timeline, tenants