feat: Use the in-memory file system and clean it at startup (#1381)

jcpitre · web-flow · commit 2d818ca6beb9 · 2025-10-02T15:02:51.000-04:00
diff --git a/functions-python/pmtiles_builder/src/csv_cache.py b/functions-python/pmtiles_builder/src/csv_cache.py
@@ -16,6 +16,8 @@
 import csv
 import logging
 import os
+import subprocess
+from pathlib import Path
 from typing import TypedDict, List, Dict
 
 from gtfs import stop_txt_is_lat_log_required
@@ -39,6 +41,40 @@ class ShapeTrips(TypedDict):
     trip_ids: List[str]
 
 
+def get_volume_size(mountpoint: str):
+    """
+    Returns the total size of the specified filesystem mount point in a human-readable format.
+
+    This function uses the `df` command-line utility to determine the size of the filesystem
+    mounted at the path specified by `mountpoint`. If the mount point does not exist, the function
+    prints an error message to the standard error and returns "N/A".
+
+    Parameters:
+    mountpoint: str
+        The filesystem mount point path to check.
+
+    Returns:
+    str
+        The total size of the specified filesystem mount point in human-readable format. If the
+        mount point is not found, returns "N/A".
+    """
+    mp = Path(mountpoint)
+    if not mp.exists():
+        logging.warning("Mountpoint not found: %s", mountpoint)
+        return "N/A"
+    cmd = [
+        "bash",
+        "-c",
+        "df -h \"$1\" | awk 'NR==2 {print $2}'",
+        "_",  # $0 placeholder (ignored)
+        str(mp),  # $1
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+    size = result.stdout.strip()
+
+    return size
+
+
 class CsvCache:
     """
     CsvCache provides cached access to GTFS CSV files in a specified working directory.
@@ -68,6 +104,7 @@ def __init__(
         self.trips_no_shapes_per_route: Dict[str, List[str]] = {}
 
         self.logger.info("Using work directory: %s", self.workdir)
+        self.logger.info("Size of workdir: %s", get_volume_size(self.workdir))
 
     def debug_log_size(self, label: str, obj: object) -> None:
         """Log the deep size of an object in bytes when DEBUG is enabled."""
diff --git a/functions-python/pmtiles_builder/src/main.py b/functions-python/pmtiles_builder/src/main.py
@@ -20,9 +20,9 @@
 import json
 import logging
 import os
+import shutil
 import subprocess
 import sys
-import tempfile
 from enum import Enum
 from typing import TypedDict, Tuple, List, Dict
 
@@ -53,6 +53,108 @@
 init_logger()
 
 
+class EphemeralOrDebugWorkdir:
+    """Context manager similar to tempfile.TemporaryDirectory with debug + auto-clean.
+
+    Behavior:
+      - If DEBUG_WORKDIR env var is set (non-empty): that directory is created (if needed),
+        returned, and never deleted nor cleaned.
+      - Else: creates a temporary directory under the provided dir (or WORKDIR_ROOT env var),
+        removes sibling directories older than a TTL (default 3600s / override via WORKDIR_MAX_AGE_SECONDS),
+        and deletes the created directory on exit.
+
+    Only directories whose names start with the fixed CLEANUP_PREFIX are considered for cleanup
+    to avoid deleting unrelated folders that might exist under the same root.
+
+    The final on-disk directory name always starts with the hardcoded prefix 'pmtiles_'.
+    The caller-supplied prefix (if any) is appended verbatim after that.
+    """
+
+    CLEANUP_PREFIX = "pmtiles_"
+
+    def __init__(
+        self,
+        dir: str | None = None,
+        prefix: str | None = None,
+        logger: logging.Logger | None = None,
+    ):
+        import tempfile
+
+        self._debug_dir = os.getenv("DEBUG_WORKDIR") or None
+        self._root = dir or os.getenv("WORKDIR_ROOT", "/tmp/in-memory")
+        self._logger = logger or get_logger("Workdir")
+        self._temp: tempfile.TemporaryDirectory[str] | None = None
+        self.name: str
+
+        os.makedirs(self._root, exist_ok=True)
+
+        self._ttl_seconds = int(os.getenv("WORKDIR_MAX_AGE_SECONDS", "3600"))
+
+        if self._debug_dir:
+            os.makedirs(self._debug_dir, exist_ok=True)
+            self.name = self._debug_dir
+            return
+
+        self._cleanup_old()
+
+        # Simple prefix: fixed manager prefix + raw user prefix (if any)
+        combined_prefix = self.CLEANUP_PREFIX + (prefix or "")
+
+        self._temp = tempfile.TemporaryDirectory(dir=self._root, prefix=combined_prefix)
+        self.name = self._temp.name
+
+    def _cleanup_old(self):
+        """
+        Delete stale work directories created by this manager (names starting with CLEANUP_PREFIX)
+        whose modification time is older than the configured TTL.
+        """
+        import time
+
+        # If in debug mode, dont cleanup anything
+        if self._debug_dir:
+            return
+
+        now = time.time()
+        deleted_count = 0
+        try:
+            entries = list(os.scandir(self._root))
+        except OSError as e:
+            self._logger.warning("Could not scan workdir root %s: %s", self._root, e)
+            return
+
+        for entry in entries:
+            try:
+                if not entry.is_dir(follow_symlinks=False):
+                    continue
+                if not entry.name.startswith(self.CLEANUP_PREFIX):
+                    continue
+                try:
+                    age = now - entry.stat(follow_symlinks=False).st_mtime
+                except OSError:
+                    continue
+                if age > self._ttl_seconds:
+                    shutil.rmtree(entry.path, ignore_errors=True)
+                    deleted_count += 1
+                    self._logger.debug(
+                        "Removed expired workdir: %s age=%.0fs", entry.path, age
+                    )
+            except OSError as e:
+                self._logger.warning("Failed to remove %s: %s", entry.path, e)
+
+        if deleted_count:
+            self._logger.info(
+                "Cleanup removed %d expired workdirs from %s", deleted_count, self._root
+            )
+
+    def __enter__(self) -> str:  # Return path like TemporaryDirectory
+        return self.name
+
+    def __exit__(self, exc_type, exc, tb):
+        if self._temp:
+            self._temp.cleanup()
+        return False  # do not suppress exceptions
+
+
 class RouteCoordinates(TypedDict):
     shape_id: str
     trip_ids: List[str]
@@ -93,18 +195,11 @@ def build_pmtiles_handler(request: flask.Request) -> dict:
         }
 
     try:
-        # Create a temporary folder to work in. It will be deleted when exiting the block.
-        with tempfile.TemporaryDirectory(prefix="build_pmtiles_") as temp_dir:
-            # If DEBUG_WORKDIR is set, use it as the work directory so it survives at the end and can be examined.
-            # In that case temp_dir will not be used but still deleted at the end of the block.
-
-            debug_workdir = os.getenv("DEBUG_WORKDIR")
-            if debug_workdir:
-                os.makedirs(debug_workdir, exist_ok=True)
-                workdir = debug_workdir
-            else:
-                workdir = temp_dir
-
+        workdir_root = os.getenv("WORKDIR_ROOT", "/tmp/in-memory")
+        # Use combined context manager that also cleans old directories
+        with EphemeralOrDebugWorkdir(
+            dir=workdir_root, prefix=f"{dataset_stable_id}_"
+        ) as workdir:
             result = {
                 "params": {
                     "feed_stable_id": feed_stable_id,
diff --git a/functions-python/pmtiles_builder/tests/test_build_pmtiles.py b/functions-python/pmtiles_builder/tests/test_build_pmtiles.py
@@ -480,25 +480,26 @@ def test_build_pmtiles_handler_feed_not_prefix(self, mock_builder):
         self.assertIn("is not a prefix of dataset_stable_id", result["error"])
 
     @patch("main.PmtilesBuilder")
-    def test_build_pmtiles_handler_success(self, mock_builder):
+    def test_build_pmtiles_handler_failure(self, mock_builder):
         # Set up environment and request
-        debug_dir = tempfile.mkdtemp()
-        os.environ["DEBUG_WORKDIR"] = debug_dir
-        payload = {
-            "feed_stable_id": self.feed_stable_id,
-            "dataset_stable_id": self.dataset_stable_id,
-        }
-        request = MagicMock()
-        request.get_json.return_value = payload
-
-        # Simulate FAILURE status
-        instance = mock_builder.return_value
-        instance.build_pmtiles.return_value = (
-            instance.OperationStatus.FAILURE,
-            "fail msg",
-        )
-        result = build_pmtiles_handler(request)
-        self.assertIn("Successfully", result["message"])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            os.environ["WORKDIR_ROOT"] = temp_dir
+            payload = {
+                "feed_stable_id": self.feed_stable_id,
+                "dataset_stable_id": self.dataset_stable_id,
+            }
+            request = MagicMock()
+            request.get_json.return_value = payload
+
+            # Simulate FAILURE status
+            instance = mock_builder.return_value
+            instance.build_pmtiles.return_value = (
+                instance.OperationStatus.FAILURE,
+                "fail msg",
+            )
+            result = build_pmtiles_handler(request)
+            self.assertIn("Successfully", result["message"])
+            self.assertEqual(os.listdir(temp_dir), [], "Expected empty workdir root")
 
 
 class TestPmtilesBuilderUpload(unittest.TestCase):
diff --git a/functions-python/pmtiles_builder/tests/test_workdir.py b/functions-python/pmtiles_builder/tests/test_workdir.py
@@ -0,0 +1,66 @@
+import os
+import tempfile
+import time
+from pathlib import Path
+
+from src.main import EphemeralOrDebugWorkdir
+
+
+def _backdate(path: Path, seconds_ago: int):
+    t = time.time() - seconds_ago
+    os.utime(path, (t, t))
+
+
+def test_old_prefixed_dir_deleted_and_root_auto_removed(monkeypatch):
+    monkeypatch.delenv("DEBUG_WORKDIR", raising=False)
+    ttl = 30
+    monkeypatch.setenv("WORKDIR_MAX_AGE_SECONDS", str(ttl))
+
+    with tempfile.TemporaryDirectory() as root:
+        root_path = Path(root)
+
+        stale_dir = root_path / "pmtiles_stale123"
+        stale_dir.mkdir()
+        _backdate(stale_dir, ttl + 300)
+
+        keep_other = root_path / "unrelated_dir"
+        keep_other.mkdir()
+
+        fresh_prefixed = root_path / "pmtiles_fresh"
+        fresh_prefixed.mkdir()
+        _backdate(fresh_prefixed, 1)
+
+        with EphemeralOrDebugWorkdir(dir=root) as new_dir:
+            assert not stale_dir.exists()
+            assert keep_other.exists()
+            assert fresh_prefixed.exists()
+            assert new_dir.startswith(root)
+            assert Path(new_dir).exists()
+            assert Path(new_dir).name.startswith("pmtiles_")
+
+        # Temp workdir created by EphemeralOrDebugWorkdir should be gone
+        assert not Path(new_dir).exists()
+
+        # Root still exists inside the TemporaryDirectory context
+        assert root_path.exists()
+
+    # After exiting TemporaryDirectory, root is removed (Not strictly a test of the code, but check anyway)
+    assert not root_path.exists()
+
+
+def test_fresh_prefixed_dir_retained(monkeypatch, tmp_path):
+    monkeypatch.delenv("DEBUG_WORKDIR", raising=False)
+    monkeypatch.setenv("WORKDIR_MAX_AGE_SECONDS", "1000")
+
+    root = tmp_path / "workroot2"
+    root.mkdir()
+
+    recent_dir = root / "pmtiles_recent"
+    recent_dir.mkdir()
+    _backdate(recent_dir, 10)  # Younger than TTL
+
+    with EphemeralOrDebugWorkdir(dir=str(root)):
+        assert recent_dir.exists()
+
+    # Recent directory still present
+    assert recent_dir.exists()
