|
20 | 20 | import json |
21 | 21 | import logging |
22 | 22 | import os |
| 23 | +import shutil |
23 | 24 | import subprocess |
24 | 25 | import sys |
25 | | -import tempfile |
26 | 26 | from enum import Enum |
27 | 27 | from typing import TypedDict, Tuple, List, Dict |
28 | 28 |
|
|
53 | 53 | init_logger() |
54 | 54 |
|
55 | 55 |
|
| 56 | +class EphemeralOrDebugWorkdir: |
| 57 | + """Context manager similar to tempfile.TemporaryDirectory with debug + auto-clean. |
| 58 | +
|
| 59 | + Behavior: |
| 60 | + - If DEBUG_WORKDIR env var is set (non-empty): that directory is created (if needed), |
| 61 | + returned, and never deleted nor cleaned. |
| 62 | + - Else: creates a temporary directory under the provided dir (or WORKDIR_ROOT env var), |
| 63 | + removes sibling directories older than a TTL (default 3600s / override via WORKDIR_MAX_AGE_SECONDS), |
| 64 | + and deletes the created directory on exit. |
| 65 | +
|
| 66 | + Only directories whose names start with the fixed CLEANUP_PREFIX are considered for cleanup |
| 67 | + to avoid deleting unrelated folders that might exist under the same root. |
| 68 | +
|
| 69 | + The final on-disk directory name always starts with the hardcoded prefix 'pmtiles_'. |
| 70 | + The caller-supplied prefix (if any) is appended verbatim after that. |
| 71 | + """ |
| 72 | + |
| 73 | + CLEANUP_PREFIX = "pmtiles_" |
| 74 | + |
| 75 | + def __init__( |
| 76 | + self, |
| 77 | + dir: str | None = None, |
| 78 | + prefix: str | None = None, |
| 79 | + logger: logging.Logger | None = None, |
| 80 | + ): |
| 81 | + import tempfile |
| 82 | + |
| 83 | + self._debug_dir = os.getenv("DEBUG_WORKDIR") or None |
| 84 | + self._root = dir or os.getenv("WORKDIR_ROOT", "/tmp/in-memory") |
| 85 | + self._logger = logger or get_logger("Workdir") |
| 86 | + self._temp: tempfile.TemporaryDirectory[str] | None = None |
| 87 | + self.name: str |
| 88 | + |
| 89 | + os.makedirs(self._root, exist_ok=True) |
| 90 | + |
| 91 | + self._ttl_seconds = int(os.getenv("WORKDIR_MAX_AGE_SECONDS", "3600")) |
| 92 | + |
| 93 | + if self._debug_dir: |
| 94 | + os.makedirs(self._debug_dir, exist_ok=True) |
| 95 | + self.name = self._debug_dir |
| 96 | + return |
| 97 | + |
| 98 | + self._cleanup_old() |
| 99 | + |
| 100 | + # Simple prefix: fixed manager prefix + raw user prefix (if any) |
| 101 | + combined_prefix = self.CLEANUP_PREFIX + (prefix or "") |
| 102 | + |
| 103 | + self._temp = tempfile.TemporaryDirectory(dir=self._root, prefix=combined_prefix) |
| 104 | + self.name = self._temp.name |
| 105 | + |
| 106 | + def _cleanup_old(self): |
| 107 | + """ |
| 108 | + Delete stale work directories created by this manager (names starting with CLEANUP_PREFIX) |
| 109 | + whose modification time is older than the configured TTL. |
| 110 | + """ |
| 111 | + import time |
| 112 | + |
| 113 | + # If in debug mode, dont cleanup anything |
| 114 | + if self._debug_dir: |
| 115 | + return |
| 116 | + |
| 117 | + now = time.time() |
| 118 | + deleted_count = 0 |
| 119 | + try: |
| 120 | + entries = list(os.scandir(self._root)) |
| 121 | + except OSError as e: |
| 122 | + self._logger.warning("Could not scan workdir root %s: %s", self._root, e) |
| 123 | + return |
| 124 | + |
| 125 | + for entry in entries: |
| 126 | + try: |
| 127 | + if not entry.is_dir(follow_symlinks=False): |
| 128 | + continue |
| 129 | + if not entry.name.startswith(self.CLEANUP_PREFIX): |
| 130 | + continue |
| 131 | + try: |
| 132 | + age = now - entry.stat(follow_symlinks=False).st_mtime |
| 133 | + except OSError: |
| 134 | + continue |
| 135 | + if age > self._ttl_seconds: |
| 136 | + shutil.rmtree(entry.path, ignore_errors=True) |
| 137 | + deleted_count += 1 |
| 138 | + self._logger.debug( |
| 139 | + "Removed expired workdir: %s age=%.0fs", entry.path, age |
| 140 | + ) |
| 141 | + except OSError as e: |
| 142 | + self._logger.warning("Failed to remove %s: %s", entry.path, e) |
| 143 | + |
| 144 | + if deleted_count: |
| 145 | + self._logger.info( |
| 146 | + "Cleanup removed %d expired workdirs from %s", deleted_count, self._root |
| 147 | + ) |
| 148 | + |
| 149 | + def __enter__(self) -> str: # Return path like TemporaryDirectory |
| 150 | + return self.name |
| 151 | + |
| 152 | + def __exit__(self, exc_type, exc, tb): |
| 153 | + if self._temp: |
| 154 | + self._temp.cleanup() |
| 155 | + return False # do not suppress exceptions |
| 156 | + |
| 157 | + |
56 | 158 | class RouteCoordinates(TypedDict): |
57 | 159 | shape_id: str |
58 | 160 | trip_ids: List[str] |
@@ -93,18 +195,11 @@ def build_pmtiles_handler(request: flask.Request) -> dict: |
93 | 195 | } |
94 | 196 |
|
95 | 197 | try: |
96 | | - # Create a temporary folder to work in. It will be deleted when exiting the block. |
97 | | - with tempfile.TemporaryDirectory(prefix="build_pmtiles_") as temp_dir: |
98 | | - # If DEBUG_WORKDIR is set, use it as the work directory so it survives at the end and can be examined. |
99 | | - # In that case temp_dir will not be used but still deleted at the end of the block. |
100 | | - |
101 | | - debug_workdir = os.getenv("DEBUG_WORKDIR") |
102 | | - if debug_workdir: |
103 | | - os.makedirs(debug_workdir, exist_ok=True) |
104 | | - workdir = debug_workdir |
105 | | - else: |
106 | | - workdir = temp_dir |
107 | | - |
| 198 | + workdir_root = os.getenv("WORKDIR_ROOT", "/tmp/in-memory") |
| 199 | + # Use combined context manager that also cleans old directories |
| 200 | + with EphemeralOrDebugWorkdir( |
| 201 | + dir=workdir_root, prefix=f"{dataset_stable_id}_" |
| 202 | + ) as workdir: |
108 | 203 | result = { |
109 | 204 | "params": { |
110 | 205 | "feed_stable_id": feed_stable_id, |
|
0 commit comments