From e3e4b14eb546b0a1e553bf494caaaaafe14a6cb4 Mon Sep 17 00:00:00 2001 From: Wasim Sandhu Date: Wed, 1 Jul 2026 16:59:42 -0700 Subject: [PATCH] Watcher: Move manual-mode uploads to a dedicated worker thread Manual-mode upload-queue polling ran on the heartbeat tick, so a slow or large upload delayed heartbeats and the dashboard flagged busy watchers as offline (prompting operator restarts). On stop, a long upload could also outlive the heartbeat join and write to the state DB after close, raising "Cannot operate on a closed database". Introduce `UploadQueueWorker`, a long-lived thread that owns the poll loop and shares a stop event with `Uploader` so shutdown can interrupt the retry backoff and abort between queued files. `stop_runtime` now stops the worker before closing the state DB, skipping the close if it can't stop in time, restoring the writer-threads-joined-before-close invariant. Bumps the watcher version to 0.4.0. Co-authored-by: Cursor --- developer-docs/reference/watcher.md | 9 +- uv.lock | 2 +- watcher/pyproject.toml | 2 +- watcher/src/data_hub_watcher/constants.py | 6 + watcher/src/data_hub_watcher/runtime.py | 51 +++++--- watcher/src/data_hub_watcher/uploader.py | 108 +++++++++++++++-- watcher/tests/test_runtime.py | 120 ++++++++++++++----- watcher/tests/test_uploader.py | 135 +++++++++++++++++++++- 8 files changed, 374 insertions(+), 59 deletions(-) diff --git a/developer-docs/reference/watcher.md b/developer-docs/reference/watcher.md index 13d9e03b..bf86aa62 100644 --- a/developer-docs/reference/watcher.md +++ b/developer-docs/reference/watcher.md @@ -67,8 +67,9 @@ While running: - **File monitor** watches the directory for new/modified files using `watchdog` and waits for each file to stabilize (size + mtime unchanged for the configured stability period). Files that keep changing for longer than 5 minutes are abandoned and surface as a `stability_timeout` error event. - **Run detector** groups stable files into runs by applying the configured regex to each file's relative path. The first file for a run triggers `POST /instruments/:id/runs`; subsequent files for the same run incrementally `PATCH` only the new entries onto the manifest. Files inside the watch tree that don't match the pattern emit a `pattern_mismatch` event (throttled to one per parent directory) so misconfigured patterns surface in the dashboard. -- **Uploader** requests a presigned S3 URL from the API and uploads each file via HTTP PUT (auto mode), or polls the server's upload queue (manual mode). The watcher does not need AWS credentials. Each upload retries up to 3 times with exponential backoff (1, 2, 4 s) and is recorded locally with its SHA-256 so retries and restarts don't re-upload the same bytes. In manual mode, queue-poll failures are throttled (1st failure, then every 10th) to keep a sustained outage visible without flooding the events stream. -- **Heartbeat loop** sends periodic heartbeats (every 60 seconds) to the API. The payload includes the watcher version, instrument ID, watch directory, upload mode, per-interval activity counters, and process uptime; a final `status="stopped"` heartbeat is sent on graceful shutdown. In manual mode, the tick also polls the upload queue. +- **Uploader** requests a presigned S3 URL from the API and uploads each file via HTTP PUT (auto mode), or processes the server's upload queue (manual mode). The watcher does not need AWS credentials. Each upload retries up to 3 times with exponential backoff (1, 2, 4 s) and is recorded locally with its SHA-256 so retries and restarts don't re-upload the same bytes. In manual mode, queue-poll failures are throttled (1st failure, then every 10th) to keep a sustained outage visible without flooding the events stream. +- **Upload worker** (manual mode only) polls the server's upload queue on its own long-lived thread every 60 seconds, decoupled from the heartbeat so a slow or large upload can't delay heartbeats and make a busy watcher look offline. On shutdown it is stopped and joined before the state DB is closed. Auto mode has no worker: uploads run on the monitor's stability-checker thread via the run detector's upload callback. +- **Heartbeat loop** sends periodic heartbeats (every 60 seconds) to the API. The payload includes the watcher version, instrument ID, watch directory, upload mode, per-interval activity counters, and process uptime; a final `status="stopped"` heartbeat is sent on graceful shutdown. - **Event reporter** batches and flushes lifecycle events (started, stopped, file uploaded, errors) to the API. See [Observability](#observability) for the full taxonomy. - **Auto-updater** runs from the same heartbeat tick on every platform — not only Windows services. It polls `GET /watchers/:id/update-check` roughly hourly and applies new releases when the watcher has been idle long enough not to clobber an in-flight run. The full activity-window guard, mandatory-update behavior, and rollback flow are documented in [Upgrading the watcher](../guides/upgrading-the-watcher.md); auto-update is hard-disabled in the `preview` environment. @@ -207,12 +208,12 @@ Upgrading an existing watcher is unaffected: the environment's database already ### Upload modes - **`auto`**: Files are uploaded to S3 immediately after run detection. -- **`manual`**: Runs are reported to the API without uploading. The server decides which files to upload via a queue, polled on each heartbeat tick. Useful when uploads need human approval. +- **`manual`**: Runs are reported to the API without uploading. The server decides which files to upload via a queue, polled by the upload worker thread every 60 seconds. Useful when uploads need human approval. Queued files are resolved against the current `watch_directory` (each queue entry carries a `relative_path` anchored to the root that was active when the file was detected). Two safeguards keep a stale queue entry from erroring forever (ENG-1397): - **On `watch_directory` change**: the server reverts every pending upload request for that instrument back to `detected` (clearing `upload_requested_at`) as soon as the new config is pushed, so the queue drains immediately. The reverted files remain re-requestable detections; an operator can queue them again from their new location. -- **Per-file 3-try cap (`MAX_QUEUE_FILE_ATTEMPTS`)**: a queued file that keeps failing — missing on disk or failing to upload — is retried on at most three heartbeat polls. After that the watcher cancels the request server-side (revert to `detected`) so the file leaves the queue instead of re-erroring each tick. The attempt count resets on watcher restart, so a transient outage longer than three ticks is recovered on the next start. +- **Per-file 3-try cap (`MAX_QUEUE_FILE_ATTEMPTS`)**: a queued file that keeps failing — missing on disk or failing to upload — is retried on at most three upload-queue polls. After that the watcher cancels the request server-side (revert to `detected`) so the file leaves the queue instead of re-erroring each poll. The attempt count resets on watcher restart, so a transient outage longer than three polls is recovered on the next start. ## Local state diff --git a/uv.lock b/uv.lock index a532adb1..b0e94c01 100644 --- a/uv.lock +++ b/uv.lock @@ -416,7 +416,7 @@ requires-dist = [ [[package]] name = "data-hub-watcher" -version = "0.3.0" +version = "0.4.0" source = { editable = "watcher" } dependencies = [ { name = "click" }, diff --git a/watcher/pyproject.toml b/watcher/pyproject.toml index 242d2a00..7348515c 100644 --- a/watcher/pyproject.toml +++ b/watcher/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "data-hub-watcher" -version = "0.3.0" +version = "0.4.0" description = "File-watcher agent for lab instrument PCs that ingests data into Data Hub." readme = "README.md" requires-python = ">=3.12" diff --git a/watcher/src/data_hub_watcher/constants.py b/watcher/src/data_hub_watcher/constants.py index 9c0a4abe..4c5cef79 100644 --- a/watcher/src/data_hub_watcher/constants.py +++ b/watcher/src/data_hub_watcher/constants.py @@ -71,6 +71,12 @@ def _resolve_watcher_log_dir() -> Path: SUPPORTED_ENVIRONMENTS: tuple[str, ...] = ("staging", "production", "preview") HEARTBEAT_INTERVAL_SECONDS = 60 +# Kept separate from ``HEARTBEAT_INTERVAL_SECONDS`` so the poll cadence can +# diverge now that uploads no longer ride the heartbeat tick. +UPLOAD_POLL_INTERVAL_SECONDS = 60 +# Bounded so a service stop doesn't hang on a large in-flight PUT; past this, +# shutdown stops waiting for the worker and leaves the state DB open. +UPLOAD_WORKER_STOP_TIMEOUT_SECONDS = 30 DEFAULT_STABILITY_PERIOD_SECONDS = 5 # Built-in presets for the ``init`` / ``config edit`` wizard. diff --git a/watcher/src/data_hub_watcher/runtime.py b/watcher/src/data_hub_watcher/runtime.py index 16e78cf2..3d41d206 100644 --- a/watcher/src/data_hub_watcher/runtime.py +++ b/watcher/src/data_hub_watcher/runtime.py @@ -22,6 +22,7 @@ DEFAULT_CONFIG_DIR, HEARTBEAT_INTERVAL_SECONDS, PRUNE_DAYS, + UPLOAD_WORKER_STOP_TIMEOUT_SECONDS, ) from data_hub_watcher.events import EventReporter, EventType, WatcherEvent from data_hub_watcher.heartbeat import HeartbeatLoop, WatcherCounters @@ -35,7 +36,7 @@ clear_upgrade_result, read_upgrade_result, ) -from data_hub_watcher.uploader import Uploader +from data_hub_watcher.uploader import Uploader, UploadQueueWorker logger = logging.getLogger(__name__) @@ -74,6 +75,10 @@ class WatcherRuntime: # in their stop wait so a shutdown can be triggered from any thread. shutdown_event: threading.Event = field(default_factory=threading.Event) upgrade_restart_event: threading.Event = field(default_factory=threading.Event) + # Manual mode only: the thread that polls the upload queue off the + # heartbeat. ``None`` in auto mode, where uploads run on the monitor's + # stability-checker thread via the run detector's callback. + upload_worker: UploadQueueWorker | None = None @dataclass(frozen=True) @@ -251,6 +256,11 @@ def _request_upgrade_restart(target_version: str) -> None: is_auto = inst.upload_mode == "auto" + # Shared with the manual-mode ``UploadQueueWorker`` so a shutdown can + # interrupt an in-flight upload's backoff and abort between queued files; + # unused (but harmless) in auto mode. + upload_stop_event = threading.Event() + uploader = Uploader( client=client, state_db=state_db, @@ -262,8 +272,13 @@ def _request_upgrade_restart(target_version: str) -> None: # Per-instrument knob, defaulted on the model so older configs # transparently inherit the new parallel-upload behaviour. upload_parallelism=inst.upload_parallelism, + stop_event=upload_stop_event, ) + # Manual mode polls the server queue on its own thread; auto mode uploads + # via the run detector's callback on the stability-checker thread instead. + upload_worker = None if is_auto else UploadQueueWorker(uploader, stop_event=upload_stop_event) + detector = RunDetector( pattern=inst.run_detection.pattern, instrument_id=inst.id, @@ -287,19 +302,10 @@ def _request_upgrade_restart(target_version: str) -> None: seed_baseline=seed_baseline, ) - # The heartbeat's `on_tick` hook is now multi-purpose: - # 1. In manual mode, poll the server's upload queue (uploads - # naturally inherit the heartbeat cadence). - # 2. Always: feed the in-process auto-updater so it can count - # idle ticks and run a server update-check roughly hourly. - # Each side wraps its own try/except so a failure on one side - # never blocks the other. + # Feeds the in-process auto-updater on every tick. Manual-mode upload + # polling used to run here too but moved to ``UploadQueueWorker`` so a + # slow upload can't delay a heartbeat. def _on_tick() -> None: - if not is_auto: - try: - uploader.poll_upload_queue() - except Exception: - logger.exception("Upload queue poll failed") try: updater.on_tick() except Exception: @@ -329,6 +335,7 @@ def _on_tick() -> None: config_dir=effective_config_dir, shutdown_event=shutdown_event, upgrade_restart_event=upgrade_restart_event, + upload_worker=upload_worker, ) @@ -480,6 +487,10 @@ def start_runtime(rt: WatcherRuntime, *, started_message: str) -> None: rt.detector.hydrate_from_state_db() rt.heartbeat.start() + # Start manual-mode upload polling before the (potentially long) initial + # scan so queued uploads keep draining while the scan walks the backlog. + if rt.upload_worker is not None: + rt.upload_worker.start() rt.monitor.start() @@ -551,6 +562,17 @@ def sync_config_to_api( def stop_runtime(rt: WatcherRuntime, *, stopped_message: str) -> None: """Shut everything down in reverse order and flush pending events.""" rt.monitor.stop() + # ``StateDB.close`` assumes writer threads have joined; a large upload + # outliving the join once wrote post-close. Stop the worker first, and if + # it won't stop in time, skip the close rather than race it (the OS reaps). + upload_worker_stopped = True + if rt.upload_worker is not None: + upload_worker_stopped = rt.upload_worker.stop(timeout=UPLOAD_WORKER_STOP_TIMEOUT_SECONDS) + if not upload_worker_stopped: + logger.warning( + "Upload worker still running at shutdown; leaving the state DB " + "open so the in-flight upload can finish without a closed-DB error" + ) rt.reporter.queue_event( WatcherEvent( event_type=EventType.WATCHER_STOPPED, @@ -559,4 +581,5 @@ def stop_runtime(rt: WatcherRuntime, *, stopped_message: str) -> None: ) rt.heartbeat.stop() rt.reporter.flush() - rt.state_db.close() + if upload_worker_stopped: + rt.state_db.close() diff --git a/watcher/src/data_hub_watcher/uploader.py b/watcher/src/data_hub_watcher/uploader.py index a593e87f..9ed5525b 100644 --- a/watcher/src/data_hub_watcher/uploader.py +++ b/watcher/src/data_hub_watcher/uploader.py @@ -22,6 +22,7 @@ from data_hub_watcher.api_client import ApiError, DataHubClient from data_hub_watcher.constants import ( MAX_QUEUE_FILE_ATTEMPTS, + UPLOAD_POLL_INTERVAL_SECONDS, UPLOAD_RETRY_BASE_DELAY, UPLOAD_RETRY_MAX, ) @@ -105,6 +106,7 @@ def __init__( watcher_id: str, watch_directory: Path, upload_parallelism: int = 1, + stop_event: threading.Event | None = None, ) -> None: self._client = client self._state_db = state_db @@ -116,18 +118,22 @@ def __init__( if upload_parallelism < 1: raise ValueError(f"upload_parallelism must be >= 1, got {upload_parallelism}") self._parallelism = upload_parallelism + # When set (by the owning ``UploadQueueWorker``), a shutdown can + # interrupt the retry backoff and abort between queued files; ``None`` + # for the one-shot ``upload`` CLI path, which never races a teardown. + self._stop_event = stop_event # Track consecutive upload-queue poll failures so the watcher # surfaces a ``kind=upload_queue_poll_failed`` event on the # 1st failure and every 10th repeat. The unthrottled case - # would emit one event per heartbeat tick during an outage, - # crowding out other signals on the dashboard. - # Mutated only from the heartbeat thread (manual mode), so + # would emit one event per poll during an outage, crowding out + # other signals on the dashboard. + # Mutated only from the upload worker thread (manual mode), so # not under any explicit lock. self._consecutive_queue_poll_failures = 0 # Per-file upload-queue attempt bookkeeping, keyed by server file id. # Bounds retries before giving up (see ``_process_queued_file``) and # doubles as the emit-once throttle for the missing-file error. Pruned - # each poll, so a re-requested id starts fresh. Heartbeat thread only. + # each poll, so a re-requested id starts fresh. Upload worker thread only. self._queue_attempts: dict[int, _QueueAttempt] = {} # Single ``requests.Session`` shared across every S3 PUT # (parallel or serial). Keeps TLS connections alive between @@ -230,10 +236,15 @@ def upload_files(self, run_id: str, files: list[FileInfo]) -> int: # Manual-mode: poll the server queue # ------------------------------------------------------------------ + def _stop_requested(self) -> bool: + return self._stop_event is not None and self._stop_event.is_set() + def poll_upload_queue(self) -> None: """Fetch the upload queue and process each file. - Intended to be called on heartbeat ticks in manual mode. + Driven by the manual-mode ``UploadQueueWorker`` loop on its own + thread, decoupled from the heartbeat so a slow upload can't starve + the liveness signal. """ try: queue = self._client.get_upload_queue(self._watcher_id) @@ -278,13 +289,19 @@ def poll_upload_queue(self) -> None: logger.info("Upload queue has %d file(s)", len(queue.files)) for qf in queue.files: + # Bail between files on shutdown so the worker's ``stop()`` can + # join promptly instead of draining the whole queue; the + # remaining files are picked up on the next start's poll. + if self._stop_requested(): + logger.info("Stop requested; deferring %d queued file(s)", len(queue.files)) + break self._process_queued_file(qf) def _process_queued_file(self, qf: UploadQueueFile) -> None: - """Attempt one queued file, bounding retries across heartbeat polls. + """Attempt one queued file, bounding retries across polls. - Manual-mode polling runs every heartbeat, so a file that can't be - uploaded -- missing on disk after a watch-directory change, or a + Manual-mode polling repeats on the worker's cadence, so a file that + can't be uploaded -- missing on disk after a watch-directory change, or a persistent upload error -- would otherwise re-error forever. We cap attempts at ``MAX_QUEUE_FILE_ATTEMPTS`` and then cancel the request server-side (revert to ``detected``) so it leaves the queue. The @@ -509,6 +526,7 @@ def _upload_single(self, path: Path, run_id: str) -> bool: return True last_exc: Exception | None = None + put_ok = False # Exponential backoff: 1s, 2s, 4s. Retries protect against transient # network errors common on lab-PC networks. @@ -516,6 +534,7 @@ def _upload_single(self, path: Path, run_id: str) -> bool: for attempt in range(UPLOAD_RETRY_MAX): try: self._put_to_presigned_url(presigned.upload_url, path, content_type) + put_ok = True break except Exception as exc: last_exc = exc @@ -528,8 +547,21 @@ def _upload_single(self, path: Path, run_id: str) -> bool: exc, delay, ) - time.sleep(delay) - else: + # Wait on the stop event when present so a shutdown cuts the + # backoff short; falls back to a plain sleep for the one-shot + # ``upload`` path that has no worker/event. + if self._stop_event is not None: + self._stop_event.wait(delay) + else: + time.sleep(delay) + # Abandon the remaining retries on shutdown. Returning False (not a + # hard failure event) leaves the request pending so the next start + # re-uploads it, rather than recording a spurious upload error. + if self._stop_requested(): + logger.info("Stop requested mid-upload; deferring %s", path.name) + return False + + if not put_ok: logger.error("Upload failed after %d attempts: %s", UPLOAD_RETRY_MAX, path.name) self._reporter.queue_event( WatcherEvent( @@ -583,3 +615,59 @@ def _upload_single(self, path: Path, run_id: str) -> bool: ) logger.info("Uploaded %s → s3://%s/%s", path.name, s3_bucket, s3_key) return True + + +class UploadQueueWorker: + """Polls the manual-mode upload queue on a dedicated long-lived thread. + + Uploads used to run on the heartbeat tick, so a slow or large transfer + delayed heartbeats and the dashboard flagged a busy watcher as offline. + Owning the poll loop here keeps the heartbeat free, and the shared + ``stop_event`` lets a shutdown interrupt an upload so ``stop()`` can join + before ``StateDB.close`` runs (which assumes writer threads have joined). + """ + + def __init__( + self, + uploader: Uploader, + *, + stop_event: threading.Event, + interval_seconds: int = UPLOAD_POLL_INTERVAL_SECONDS, + ) -> None: + self._uploader = uploader + self._stop_event = stop_event + self._interval = interval_seconds + self._thread: threading.Thread | None = None + + def start(self) -> None: + self._stop_event.clear() + self._thread = threading.Thread(target=self._run, daemon=True, name="upload-worker") + self._thread.start() + + def stop(self, timeout: float | None = None) -> bool: + """Signal the loop and join it. Returns whether the thread exited. + + A ``False`` return means an upload is still in flight past *timeout* + (a single S3 PUT can run up to its request timeout); the caller uses + this to avoid closing the state DB out from under the live upload. + """ + self._stop_event.set() + if self._thread is None: + return True + self._thread.join(timeout=timeout) + return not self._thread.is_alive() + + def _run(self) -> None: + # Wait first (parity with the previous heartbeat-driven cadence), + # then poll each interval until stop. + while not self._stop_event.wait(timeout=self._interval): + self._poll_once() + + def _poll_once(self) -> None: + # ``poll_upload_queue`` already handles and reports poll failures; this + # guard only keeps an unexpected error from killing the thread and + # silently ending all future polls. + try: + self._uploader.poll_upload_queue() + except Exception: + logger.exception("Upload queue poll failed") diff --git a/watcher/tests/test_runtime.py b/watcher/tests/test_runtime.py index 3b77dc35..fb3380a9 100644 --- a/watcher/tests/test_runtime.py +++ b/watcher/tests/test_runtime.py @@ -6,11 +6,13 @@ forgot to wire `on_tick` on the `HeartbeatLoop`. These tests lock in the wiring contract per `upload_mode` so any future drift fails loudly: -* auto mode -> `detector._upload_cb` is `uploader.upload_files` - and `heartbeat._on_tick` ticks the auto-updater only -* manual mode -> `detector._upload_cb` is `None` - and `heartbeat._on_tick` polls `uploader.poll_upload_queue` - *and* ticks the auto-updater +* auto mode -> `detector._upload_cb` is `uploader.upload_files`, + `heartbeat._on_tick` ticks the auto-updater only, and + `rt.upload_worker` is `None` +* manual mode -> `detector._upload_cb` is `None`, `heartbeat._on_tick` + ticks the auto-updater only (uploads now run on the + dedicated `UploadQueueWorker` thread, not the heartbeat), + and `rt.upload_worker` is set """ from __future__ import annotations @@ -29,10 +31,12 @@ from data_hub_watcher.run_detector import RunDetector from data_hub_watcher.runtime import ( ShutdownReason, + WatcherRuntime, _summarize_worker_failure, build_runtime, classify_shutdown, start_runtime, + stop_runtime, ) from data_hub_watcher.state import StateDB from data_hub_watcher.updater import Updater, write_upgrade_marker @@ -116,9 +120,20 @@ def test_heartbeat_on_tick_only_drives_updater(self, tmp_path: Path, db_path: Pa finally: rt.state_db.close() + def test_auto_mode_has_no_upload_worker(self, tmp_path: Path, db_path: Path) -> None: + cfg = _make_config(tmp_path, upload_mode="auto") + rt = build_runtime(client=MagicMock(), cfg=cfg, db_path=db_path) + + try: + # Auto-mode uploads run on the monitor's stability-checker thread + # via the detector callback, so there is no upload-queue worker. + assert rt.upload_worker is None + finally: + rt.state_db.close() + class TestBuildRuntimeManualMode: - """Manual mode: heartbeat polls the upload queue, detector does not upload.""" + """Manual mode: a dedicated worker polls the upload queue off the heartbeat.""" def test_detector_upload_callback_is_none(self, tmp_path: Path, db_path: Path) -> None: cfg = _make_config(tmp_path, upload_mode="manual") @@ -131,44 +146,37 @@ def test_detector_upload_callback_is_none(self, tmp_path: Path, db_path: Path) - finally: rt.state_db.close() - def test_heartbeat_on_tick_polls_upload_queue(self, tmp_path: Path, db_path: Path) -> None: + def test_heartbeat_on_tick_does_not_poll_upload_queue( + self, tmp_path: Path, db_path: Path + ) -> None: cfg = _make_config(tmp_path, upload_mode="manual") rt = build_runtime(client=MagicMock(), cfg=cfg, db_path=db_path) try: - # The heartbeat must call `uploader.poll_upload_queue` on every - # tick — this is the bug the runtime extraction was fixing. + # Uploads moved off the heartbeat thread onto the worker, so the + # tick must only feed the updater — a slow upload can no longer + # delay a heartbeat and make a busy watcher look offline. assert rt.heartbeat._on_tick is not None - rt.uploader.poll_upload_queue = MagicMock() # type: ignore[method-assign] rt.updater.on_tick = MagicMock(return_value=None) # type: ignore[method-assign] rt.heartbeat._on_tick() - rt.uploader.poll_upload_queue.assert_called_once_with() - # The same hook must also feed the auto-updater so its idle - # counter advances regardless of upload_mode. + rt.uploader.poll_upload_queue.assert_not_called() rt.updater.on_tick.assert_called_once_with() finally: rt.state_db.close() - def test_on_tick_swallows_poll_exceptions(self, tmp_path: Path, db_path: Path) -> None: - """Polling errors must not propagate out of the heartbeat tick, - otherwise one transient server blip kills the heartbeat thread - and the watcher goes silent until restart.""" + def test_manual_mode_builds_upload_worker_sharing_stop_event( + self, tmp_path: Path, db_path: Path + ) -> None: cfg = _make_config(tmp_path, upload_mode="manual") rt = build_runtime(client=MagicMock(), cfg=cfg, db_path=db_path) try: - rt.uploader.poll_upload_queue = MagicMock( # type: ignore[method-assign] - side_effect=RuntimeError("boom") - ) - rt.updater.on_tick = MagicMock(return_value=None) # type: ignore[method-assign] - assert rt.heartbeat._on_tick is not None - rt.heartbeat._on_tick() - rt.uploader.poll_upload_queue.assert_called_once_with() - # Updater must still tick even when the upload-queue poll - # blew up, otherwise a permanently-failing manual-mode - # poll would also disable auto-updates. - rt.updater.on_tick.assert_called_once_with() + # The worker must wrap this runtime's uploader and share its stop + # event so a shutdown can interrupt an in-flight upload. + assert rt.upload_worker is not None + assert rt.upload_worker._uploader is rt.uploader + assert rt.upload_worker._stop_event is rt.uploader._stop_event finally: rt.state_db.close() @@ -189,6 +197,62 @@ def test_on_tick_swallows_updater_exceptions(self, tmp_path: Path, db_path: Path rt.state_db.close() +class TestStopRuntimeUploadWorkerOrdering: + """`stop_runtime` must stop the upload worker before closing the state DB. + + Regression guard for the prod "Cannot operate on a closed database" race: + a still-running upload wrote to the DB after `close()` because teardown + didn't wait for the upload thread. + """ + + @staticmethod + def _runtime_with_mocks(state_db: MagicMock, upload_worker: Any) -> WatcherRuntime: + return WatcherRuntime( + state_db=state_db, + counters=MagicMock(), + reporter=MagicMock(), + uploader=MagicMock(), + detector=MagicMock(), + monitor=MagicMock(), + heartbeat=MagicMock(), + updater=MagicMock(), + config_dir=Path("/tmp"), + upload_worker=upload_worker, + ) + + def test_closes_db_when_worker_stops_cleanly(self) -> None: + state_db = MagicMock() + worker = MagicMock() + worker.stop.return_value = True + rt = self._runtime_with_mocks(state_db, worker) + + stop_runtime(rt, stopped_message="Watcher stopped") + + worker.stop.assert_called_once() + state_db.close.assert_called_once_with() + + def test_skips_close_when_worker_still_running(self) -> None: + # A large PUT outliving the join must not have the DB yanked out from + # under it; teardown leaves the connection open and lets the OS reap it. + state_db = MagicMock() + worker = MagicMock() + worker.stop.return_value = False + rt = self._runtime_with_mocks(state_db, worker) + + stop_runtime(rt, stopped_message="Watcher stopped") + + worker.stop.assert_called_once() + state_db.close.assert_not_called() + + def test_closes_db_in_auto_mode_without_worker(self) -> None: + state_db = MagicMock() + rt = self._runtime_with_mocks(state_db, None) + + stop_runtime(rt, stopped_message="Watcher stopped") + + state_db.close.assert_called_once_with() + + class TestBuildRuntimeSharedDependencies: """Cross-object wiring invariants that apply to both upload modes.""" diff --git a/watcher/tests/test_uploader.py b/watcher/tests/test_uploader.py index 6d83dc8c..43f1d08a 100644 --- a/watcher/tests/test_uploader.py +++ b/watcher/tests/test_uploader.py @@ -21,7 +21,7 @@ UploadQueueResponse, ) from data_hub_watcher.state import StateDB -from data_hub_watcher.uploader import Uploader +from data_hub_watcher.uploader import Uploader, UploadQueueWorker @pytest.fixture() @@ -742,3 +742,136 @@ def test_cancel_failure_is_retried_next_poll( # The next poll retries the cancel rather than re-erroring on upload. uploader.poll_upload_queue() assert mock_client.cancel_upload_request.call_count == 2 + + +class TestUploaderStopEvent: + """A shutdown must interrupt uploads without recording a spurious failure.""" + + def _uploader_with_stop( + self, + mock_client: MagicMock, + state_db: StateDB, + tmp_path: Path, + stop_event: threading.Event, + ) -> Uploader: + return Uploader( + client=mock_client, + state_db=state_db, + event_reporter=MagicMock(spec=EventReporter), + counters=WatcherCounters(), + instrument_id="test-instrument", + watcher_id="watcher-123", + watch_directory=tmp_path, + stop_event=stop_event, + ) + + def test_stop_during_backoff_defers_without_failure( + self, + mock_client: MagicMock, + state_db: StateDB, + tmp_file: Path, + tmp_path: Path, + ) -> None: + stop = threading.Event() + stop.set() # already stopping when the first attempt fails + up = self._uploader_with_stop(mock_client, state_db, tmp_path, stop) + mock_client.request_upload_url.return_value = PresignedUploadResponse( + upload_url="https://s3.example.com/presigned", + s3_bucket="test-bucket", + s3_key="k", + file_id=42, + expires_in=3600, + already_uploaded=False, + ) + + with patch.object( + Uploader, "_put_to_presigned_url", side_effect=ConnectionError("network down") + ) as mock_put: + result = up._upload_single(tmp_file, "RUN-001") + + # Deferred, not failed: one attempt, no success PATCH, no error event + # or counter bump, so the request stays pending for the next start. + assert result is False + assert mock_put.call_count == 1 + mock_client.mark_file_uploaded.assert_not_called() + assert up._counters.errors == 0 + cast(MagicMock, up._reporter).queue_event.assert_not_called() + + def test_poll_bails_between_files_when_stopping( + self, + mock_client: MagicMock, + state_db: StateDB, + tmp_path: Path, + ) -> None: + stop = threading.Event() + stop.set() + up = self._uploader_with_stop(mock_client, state_db, tmp_path, stop) + mock_client.get_upload_queue.return_value = UploadQueueResponse( + files=[ + UploadQueueFile( + id=1, + instrument_id="test-instrument", + run_id="R1", + filename="a.csv", + relative_path="a.csv", + ), + UploadQueueFile( + id=2, + instrument_id="test-instrument", + run_id="R1", + filename="b.csv", + relative_path="b.csv", + ), + ] + ) + + with patch.object(Uploader, "_process_queued_file") as mock_process: + up.poll_upload_queue() + + mock_process.assert_not_called() + + +class TestUploadQueueWorker: + """The worker owns the poll loop and must stop cleanly for shutdown.""" + + def test_poll_once_swallows_exceptions(self) -> None: + uploader = MagicMock() + uploader.poll_upload_queue.side_effect = RuntimeError("boom") + worker = UploadQueueWorker(uploader, stop_event=threading.Event()) + + # A poll blowup must not escape and kill the thread. + worker._poll_once() + + uploader.poll_upload_queue.assert_called_once_with() + + def test_stop_joins_idle_worker(self) -> None: + uploader = MagicMock() + worker = UploadQueueWorker(uploader, stop_event=threading.Event(), interval_seconds=60) + worker.start() + + # The loop waits on the stop event, so setting it returns the join + # immediately rather than after the 60s interval. + assert worker.stop(timeout=5) is True + + def test_stop_reports_false_when_upload_in_flight(self) -> None: + # A poll stuck mid-upload past the join timeout must report unfinished + # so teardown skips closing the state DB out from under it. + in_poll = threading.Event() + release = threading.Event() + + def blocking_poll() -> None: + in_poll.set() + release.wait(timeout=5) + + uploader = MagicMock() + uploader.poll_upload_queue.side_effect = blocking_poll + # interval=0 so the loop enters the (blocking) poll right away. + worker = UploadQueueWorker(uploader, stop_event=threading.Event(), interval_seconds=0) + worker.start() + assert in_poll.wait(timeout=5) + + try: + assert worker.stop(timeout=0.1) is False + finally: + # Let the blocked poll finish so the daemon thread exits cleanly. + release.set()