diff --git a/api/changelog.d/attack-paths-stale-cleanup.fixed.md b/api/changelog.d/attack-paths-stale-cleanup.fixed.md new file mode 100644 index 0000000000..758ef9d630 --- /dev/null +++ b/api/changelog.d/attack-paths-stale-cleanup.fixed.md @@ -0,0 +1 @@ +`attack-paths-cleanup-stale-scans` now retries worker pings and checks recent scan activity before failing scans and removing temporary databases diff --git a/api/src/backend/api/attack_paths/views_helpers.py b/api/src/backend/api/attack_paths/views_helpers.py index d1b351f454..1c1d3e7888 100644 --- a/api/src/backend/api/attack_paths/views_helpers.py +++ b/api/src/backend/api/attack_paths/views_helpers.py @@ -152,10 +152,10 @@ def execute_custom_query( scan: AttackPathsScan, ) -> dict[str, Any]: # Defense-in-depth for custom queries: - # 1. `neo4j.READ_ACCESS` — prevents mutations at the driver level - # 2. `inject_provider_label()` — regex-based label injection scopes node patterns - # 3. `_serialize_graph()` — post-query filter drops nodes without the provider label - # 4. `USING QUERY:TIMEOUTMILLISECONDS` on Neptune — server-side runaway cutoff + # 1. `neo4j.READ_ACCESS` - prevents mutations at the driver level + # 2. `inject_provider_label()` - regex-based label injection scopes node patterns + # 3. `_serialize_graph()` - post-query filter drops nodes without the provider label + # 4. `USING QUERY:TIMEOUTMILLISECONDS` on Neptune - server-side runaway cutoff # # Layer 2 is best-effort (regex can't fully parse Cypher); # layer 3 is the safety net that guarantees provider isolation. diff --git a/api/src/backend/config/django/base.py b/api/src/backend/config/django/base.py index 23af0cc6d3..04251b4479 100644 --- a/api/src/backend/config/django/base.py +++ b/api/src/backend/config/django/base.py @@ -308,6 +308,9 @@ CSRF_COOKIE_SECURE = True SESSION_COOKIE_SECURE = True # Attack Paths +ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES = env.int( + "ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES", 30 +) ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES = env.int( "ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES", 2880 ) # 48h diff --git a/api/src/backend/tasks/jobs/attack_paths/cleanup.py b/api/src/backend/tasks/jobs/attack_paths/cleanup.py index 83192f18d0..11868b3fe0 100644 --- a/api/src/backend/tasks/jobs/attack_paths/cleanup.py +++ b/api/src/backend/tasks/jobs/attack_paths/cleanup.py @@ -1,40 +1,50 @@ from datetime import UTC, datetime, timedelta +from functools import partial from api.attack_paths import database as graph_database from api.db_router import MainRouter from api.db_utils import rls_transaction from api.models import AttackPathsScan, StateChoices -from celery import states +from celery import current_app, states from celery.utils.log import get_task_logger -from config.django.base import ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES +from config.django.base import ( + ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES, + ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES, +) +from django.db import DatabaseError +from django.db.transaction import on_commit from tasks.jobs.attack_paths.db_utils import ( mark_scan_finished, recover_graph_data_ready, ) -from tasks.jobs.orphan_recovery import is_worker_alive as _is_worker_alive from tasks.jobs.orphan_recovery import revoke_task as _revoke_task logger = get_task_logger(__name__) +WORKER_PING_BASE_TIMEOUT_SECONDS = 5 +WORKER_PING_MAX_ATTEMPTS = 3 + def cleanup_stale_attack_paths_scans() -> dict: """ Mark stale `AttackPathsScan` rows as `FAILED`. Covers two stuck-state scenarios: - 1. `EXECUTING` scans whose workers are dead, or that have exceeded the - stale threshold while alive. - 2. `SCHEDULED` scans that never made it to a worker — parent scan + 1. `EXECUTING` scans whose workers are unresponsive and whose rows have + stopped receiving progress updates, or that exceeded the stale threshold. + 2. `SCHEDULED` scans that never made it to a worker - parent scan crashed before dispatch, broker lost the message, etc. Detected by age plus the parent `Scan` no longer being in flight. """ - threshold = timedelta(minutes=ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES) now = datetime.now(tz=UTC) - cutoff = now - threshold + stale_cutoff = now - timedelta(minutes=ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES) + inactivity_cutoff = now - timedelta( + minutes=ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES + ) cleaned_up: list[str] = [] - cleaned_up.extend(_cleanup_stale_executing_scans(cutoff)) - cleaned_up.extend(_cleanup_stale_scheduled_scans(cutoff)) + cleaned_up.extend(_cleanup_stale_executing_scans(stale_cutoff, inactivity_cutoff)) + cleaned_up.extend(_cleanup_stale_scheduled_scans(stale_cutoff)) logger.info( f"Stale `AttackPathsScan` cleanup: {len(cleaned_up)} scan(s) cleaned up" @@ -42,13 +52,57 @@ def cleanup_stale_attack_paths_scans() -> dict: return {"cleaned_up_count": len(cleaned_up), "scan_ids": cleaned_up} -def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]: +def _ping_workers(workers: set[str]) -> tuple[set[str], set[str] | None]: + """Ping worker destinations in parallel and retry only missing workers. + + The second tuple item is `None` when the final ping attempt raises. In that + case the pending workers have unknown liveness and their scans must be kept. + """ + pending = set(workers) + responsive: set[str] = set() + + for attempt in range(WORKER_PING_MAX_ATTEMPTS): + if not pending: + return responsive, set() + + timeout = WORKER_PING_BASE_TIMEOUT_SECONDS * 2**attempt + try: + response = current_app.control.inspect( + destination=sorted(pending), timeout=timeout + ).ping() + except Exception: + attempts_remaining = WORKER_PING_MAX_ATTEMPTS - attempt - 1 + if attempts_remaining: + logger.warning( + f"Attack Paths worker ping attempt {attempt + 1} failed; " + f"retrying pending workers with {attempts_remaining} " + "attempt(s) remaining", + exc_info=True, + ) + continue + + logger.exception( + "Attack Paths worker ping attempts exhausted; preserving scans " + "for workers with unknown liveness" + ) + return responsive, None + + responded = pending.intersection((response or {}).keys()) + responsive.update(responded) + pending.difference_update(responded) + + return responsive, pending + + +def _cleanup_stale_executing_scans( + stale_cutoff: datetime, inactivity_cutoff: datetime +) -> list[str]: """ Two-pass detection for `EXECUTING` scans: - 1. If `TaskResult.worker` exists, ping the worker. - - Dead worker: cleanup immediately (any age). - - Alive + past threshold: revoke the task, then cleanup. - - Alive + within threshold: skip. + 1. Ping all recorded workers in parallel with bounded retries. + - Responsive + past stale threshold: cleanup. + - Unresponsive + past inactivity threshold: cleanup. + - Unknown after a final ping exception: preserve. 2. If no worker field: fall back to time-based heuristic only. """ executing_scans = list( @@ -57,14 +111,13 @@ def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]: .select_related("task__task_runner_task") ) - # Cache worker liveness so each worker is pinged at most once workers = { tr.worker for scan in executing_scans if (tr := getattr(scan.task, "task_runner_task", None) if scan.task else None) and tr.worker } - worker_alive = {w: _is_worker_alive(w) for w in workers} + responsive_workers, unresponsive_workers = _ping_workers(workers) cleaned_up: list[str] = [] @@ -75,27 +128,50 @@ def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]: worker = task_result.worker if task_result else None if worker: - alive = worker_alive.get(worker, True) - - if alive: - if scan.started_at and scan.started_at >= cutoff: + if worker in responsive_workers: + if scan.started_at is None or scan.started_at >= stale_cutoff: continue - # Alive but stale — revoke before cleanup - _revoke_task(task_result) - reason = "Scan exceeded stale threshold — cleaned up by periodic task" + reason = "Scan exceeded stale threshold - cleaned up by periodic task" + recheck_activity_cutoff = None + elif unresponsive_workers is None or worker not in unresponsive_workers: + logger.info( + f"Preserving scan {scan.id}: worker {worker} liveness is " + f"unknown (progress={scan.progress}, updated_at={scan.updated_at})" + ) + continue else: - reason = "Worker dead — cleaned up by periodic task" + if scan.updated_at >= inactivity_cutoff: + logger.info( + f"Preserving scan {scan.id}: worker {worker} is unresponsive " + f"but activity is recent (progress={scan.progress}, " + f"updated_at={scan.updated_at})" + ) + continue + + reason = ( + "Worker unresponsive and scan inactive for " + f"{ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES} minutes - " + "cleaned up by periodic task" + ) + recheck_activity_cutoff = inactivity_cutoff else: # No worker recorded, time-based heuristic only - if scan.started_at and scan.started_at >= cutoff: + if scan.started_at is None or scan.started_at >= stale_cutoff: continue reason = ( - "No worker recorded, scan exceeded stale threshold — " + "No worker recorded, scan exceeded stale threshold - " "cleaned up by periodic task" ) + recheck_activity_cutoff = None - if _cleanup_scan(scan, task_result, reason): + if _cleanup_scan( + scan, + task_result, + reason, + revoke=worker is not None, + inactivity_cutoff=recheck_activity_cutoff, + ): cleaned_up.append(str(scan.id)) return cleaned_up @@ -112,10 +188,9 @@ def _cleanup_stale_scheduled_scans(cutoff: datetime) -> list[str]: avoids cleaning up rows whose parent Prowler scan is legitimately still running. - For each match: revoke the queued task (best-effort; harmless if already - consumed), atomically flip to `FAILED`, and mark the `TaskResult`. The - temp Neo4j database is never created while `SCHEDULED`, so no drop is - needed. + For each match: lock and recheck the row, mark the scan and `TaskResult` as + failed, then revoke the queued task after the transaction commits. The temp + Neo4j database is never created while `SCHEDULED`, so no drop is needed. """ scheduled_scans = list( AttackPathsScan.all_objects.using(MainRouter.admin_db) @@ -141,42 +216,54 @@ def _cleanup_stale_scheduled_scans(cutoff: datetime) -> list[str]: task_result = ( getattr(scan.task, "task_runner_task", None) if scan.task else None ) - if task_result: - _revoke_task(task_result, terminate=False) - - reason = "Scan never started — cleaned up by periodic task" + reason = "Scan never started - cleaned up by periodic task" if _cleanup_scheduled_scan(scan, task_result, reason): cleaned_up.append(str(scan.id)) return cleaned_up -def _cleanup_scan(scan, task_result, reason: str) -> bool: +def _cleanup_scan( + scan, + task_result, + reason: str, + *, + revoke: bool = False, + inactivity_cutoff: datetime | None = None, +) -> bool: """ Clean up a single stale `AttackPathsScan`: - drop temp DB, mark `FAILED`, update `TaskResult`, recover `graph_data_ready`. + lock and recheck, mark `FAILED`, revoke after commit, drop the temp DB, and + recover graph readiness. Returns `True` if the scan was actually cleaned up, `False` if skipped. """ scan_id_str = str(scan.id) - # Drop temp Neo4j database + try: + fresh_scan = _finalize_failed_scan( + scan, + StateChoices.EXECUTING, + reason, + task_result=task_result, + revoke=revoke, + inactivity_cutoff=inactivity_cutoff, + ) + except DatabaseError: + logger.exception( + f"Failed to mark stale Attack Paths scan {scan_id_str} as failed" + ) + return False + + if fresh_scan is None: + return False + tmp_db_name = graph_database.get_database_name(scan.id, temporary=True) try: graph_database.drop_database(tmp_db_name) except Exception: logger.exception(f"Failed to drop temp database {tmp_db_name}") - fresh_scan = _finalize_failed_scan(scan, StateChoices.EXECUTING, reason) - if fresh_scan is None: - return False - - # Mark `TaskResult` as `FAILURE` (not RLS-protected, outside lock) - if task_result: - task_result.status = states.FAILURE - task_result.date_done = datetime.now(tz=UTC) - task_result.save(update_fields=["status", "date_done"]) - recover_graph_data_ready(fresh_scan) logger.info(f"Cleaned up stale scan {scan_id_str}: {reason}") @@ -187,31 +274,49 @@ def _cleanup_scheduled_scan(scan, task_result, reason: str) -> bool: """ Clean up a `SCHEDULED` scan that never reached a worker. - Skips the temp Neo4j drop — the database is only created once the worker + Skips the temp Neo4j drop - the database is only created once the worker enters `EXECUTING`, so dropping it here just produces noisy log output. Returns `True` if the scan was actually cleaned up, `False` if skipped. """ scan_id_str = str(scan.id) - fresh_scan = _finalize_failed_scan(scan, StateChoices.SCHEDULED, reason) - if fresh_scan is None: + try: + fresh_scan = _finalize_failed_scan( + scan, + StateChoices.SCHEDULED, + reason, + task_result=task_result, + revoke=task_result is not None, + terminate=False, + ) + except DatabaseError: + logger.exception( + f"Failed to mark scheduled Attack Paths scan {scan_id_str} as failed" + ) return False - if task_result: - task_result.status = states.FAILURE - task_result.date_done = datetime.now(tz=UTC) - task_result.save(update_fields=["status", "date_done"]) + if fresh_scan is None: + return False logger.info(f"Cleaned up scheduled scan {scan_id_str}: {reason}") return True -def _finalize_failed_scan(scan, expected_state: str, reason: str): +def _finalize_failed_scan( + scan, + expected_state: str, + reason: str, + *, + task_result=None, + revoke: bool = False, + terminate: bool = True, + inactivity_cutoff: datetime | None = None, +): """ - Atomically lock the row, verify it's still in `expected_state`, and - mark it `FAILED`. Returns the locked row on success, `None` if the - row is gone or has already moved on. + Atomically lock the row, verify it's still eligible, and mark it `FAILED`. + If requested, register revocation after commit. Returns the locked row on + success, `None` if the row is gone or has already moved on. """ scan_id_str = str(scan.id) with rls_transaction(str(scan.tenant_id)): @@ -225,6 +330,23 @@ def _finalize_failed_scan(scan, expected_state: str, reason: str): logger.info(f"Scan {scan_id_str} is now {fresh_scan.state}, skipping") return None + if inactivity_cutoff is not None and fresh_scan.updated_at >= inactivity_cutoff: + logger.info( + f"Scan {scan_id_str} received activity during worker checks, skipping" + ) + return None + mark_scan_finished(fresh_scan, StateChoices.FAILED, {"global_error": reason}) + if task_result: + task_result.status = states.FAILURE + task_result.date_done = datetime.now(tz=UTC) + task_result.save(update_fields=["status", "date_done"]) + + if revoke and task_result: + on_commit( + partial(_revoke_task, task_result, terminate=terminate), + using=fresh_scan._state.db, + ) + return fresh_scan diff --git a/api/src/backend/tasks/jobs/attack_paths/db_utils.py b/api/src/backend/tasks/jobs/attack_paths/db_utils.py index c444a62602..327d4463e2 100644 --- a/api/src/backend/tasks/jobs/attack_paths/db_utils.py +++ b/api/src/backend/tasks/jobs/attack_paths/db_utils.py @@ -126,14 +126,17 @@ def starting_attack_paths_scan( if locked.state != StateChoices.SCHEDULED: return False + now = datetime.now(tz=UTC) locked.state = StateChoices.EXECUTING - locked.started_at = datetime.now(tz=UTC) + locked.started_at = now + locked.updated_at = now locked.update_tag = cartography_config.update_tag - locked.save(update_fields=["state", "started_at", "update_tag"]) + locked.save(update_fields=["state", "started_at", "updated_at", "update_tag"]) # Keep the in-memory object the caller is holding in sync. attack_paths_scan.state = locked.state attack_paths_scan.started_at = locked.started_at + attack_paths_scan.updated_at = locked.updated_at attack_paths_scan.update_tag = locked.update_tag return True @@ -181,7 +184,8 @@ def update_attack_paths_scan_progress( ) -> None: with rls_transaction(attack_paths_scan.tenant_id): attack_paths_scan.progress = progress - attack_paths_scan.save(update_fields=["progress"]) + attack_paths_scan.updated_at = datetime.now(tz=UTC) + attack_paths_scan.save(update_fields=["progress", "updated_at"]) def set_graph_data_ready( diff --git a/api/src/backend/tasks/jobs/orphan_recovery.py b/api/src/backend/tasks/jobs/orphan_recovery.py index 7211f1a1d5..05a2083dbe 100644 --- a/api/src/backend/tasks/jobs/orphan_recovery.py +++ b/api/src/backend/tasks/jobs/orphan_recovery.py @@ -172,11 +172,9 @@ def reconcile_orphans( window_hours: int = 6, dry_run: bool = False, ) -> dict: - """Run the full orphan sweep under a single-flight advisory lock. + """Run the orphan task sweep under a single-flight advisory lock. - Recovers any orphaned in-flight task and delegates attack-paths scans that - never reached a worker to their existing stale-cleanup. Returns a summary; - a no-op (lock not won) is reported too. + Returns a recovery summary. A no-op is reported when the lock is not acquired. """ with advisory_lock() as acquired: if not acquired: @@ -200,11 +198,6 @@ def reconcile_orphans( logger.info("Orphan task recovery disabled by feature flag") result = {"recovered": [], "failed": [], "skipped": [], "enabled": False} - if not dry_run: - from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans - - result["attack_paths"] = cleanup_stale_attack_paths_scans() - return {"acquired": True, **result} diff --git a/api/src/backend/tasks/tests/test_attack_paths_scan.py b/api/src/backend/tasks/tests/test_attack_paths_scan.py index db4dc4b0c6..3be885a7fc 100644 --- a/api/src/backend/tasks/tests/test_attack_paths_scan.py +++ b/api/src/backend/tasks/tests/test_attack_paths_scan.py @@ -17,7 +17,7 @@ from api.models import ( StatusChoices, Task, ) -from django.db import DEFAULT_DB_ALIAS +from django.db import DEFAULT_DB_ALIAS, DatabaseError from django_celery_results.models import TaskResult from prowler.lib.check.models import Severity from tasks.jobs.attack_paths import findings as findings_module @@ -2686,10 +2686,194 @@ class TestAttackPathsDbUtilsGraphDataReady: assert ap_scan_b.graph_data_ready is True +class TestAttackPathsWorkerPing: + @patch("tasks.jobs.attack_paths.cleanup.current_app") + def test_pings_workers_in_parallel_and_retries_only_missing(self, mock_app): + from tasks.jobs.attack_paths.cleanup import _ping_workers + + first_ping = MagicMock(return_value={"worker-a@host": {"ok": "pong"}}) + second_ping = MagicMock(return_value={"worker-b@host": {"ok": "pong"}}) + third_ping = MagicMock(return_value={"worker-c@host": {"ok": "pong"}}) + mock_app.control.inspect.side_effect = [ + MagicMock(ping=first_ping), + MagicMock(ping=second_ping), + MagicMock(ping=third_ping), + ] + + responsive, unresponsive = _ping_workers( + {"worker-c@host", "worker-a@host", "worker-b@host"} + ) + + assert responsive == { + "worker-a@host", + "worker-b@host", + "worker-c@host", + } + assert unresponsive == set() + assert mock_app.control.inspect.call_args_list == [ + call( + destination=["worker-a@host", "worker-b@host", "worker-c@host"], + timeout=5, + ), + call(destination=["worker-b@host", "worker-c@host"], timeout=10), + call(destination=["worker-c@host"], timeout=20), + ] + + @patch("tasks.jobs.attack_paths.cleanup.logger") + @patch("tasks.jobs.attack_paths.cleanup.current_app") + def test_retries_intermediate_ping_exceptions(self, mock_app, mock_logger): + from tasks.jobs.attack_paths.cleanup import _ping_workers + + mock_app.control.inspect.side_effect = [ + MagicMock(ping=MagicMock(side_effect=ConnectionError("first"))), + MagicMock(ping=MagicMock(side_effect=ConnectionError("second"))), + MagicMock(ping=MagicMock(return_value={})), + ] + + responsive, unresponsive = _ping_workers({"worker@host"}) + + assert responsive == set() + assert unresponsive == {"worker@host"} + assert mock_logger.warning.call_count == 2 + assert all( + warning.kwargs["exc_info"] is True + for warning in mock_logger.warning.call_args_list + ) + mock_logger.exception.assert_not_called() + + @patch("tasks.jobs.attack_paths.cleanup.logger") + @patch("tasks.jobs.attack_paths.cleanup.current_app") + def test_final_ping_exception_leaves_pending_workers_unknown( + self, mock_app, mock_logger + ): + from tasks.jobs.attack_paths.cleanup import _ping_workers + + mock_app.control.inspect.side_effect = [ + MagicMock(ping=MagicMock(return_value={"worker-a@host": {"ok": "pong"}})), + MagicMock(ping=MagicMock(return_value={})), + MagicMock(ping=MagicMock(side_effect=ConnectionError("final"))), + ] + + responsive, unresponsive = _ping_workers({"worker-a@host", "worker-b@host"}) + + assert responsive == {"worker-a@host"} + assert unresponsive is None + mock_logger.exception.assert_called_once() + + @patch("tasks.jobs.attack_paths.cleanup.logger") + @patch("tasks.jobs.attack_paths.cleanup.current_app") + def test_worker_can_respond_after_an_intermediate_exception( + self, mock_app, mock_logger + ): + from tasks.jobs.attack_paths.cleanup import _ping_workers + + mock_app.control.inspect.side_effect = [ + MagicMock(ping=MagicMock(side_effect=ConnectionError("first"))), + MagicMock(ping=MagicMock(side_effect=ConnectionError("second"))), + MagicMock(ping=MagicMock(return_value={"worker@host": {"ok": "pong"}})), + ] + + responsive, unresponsive = _ping_workers({"worker@host"}) + + assert responsive == {"worker@host"} + assert unresponsive == set() + assert mock_logger.warning.call_count == 2 + mock_logger.exception.assert_not_called() + + +class TestAttackPathsCleanupTask: + @patch( + "tasks.tasks.cleanup_stale_attack_paths_scans", + return_value={"cleaned_up_count": 1, "scan_ids": ["scan-id"]}, + ) + def test_hourly_task_invokes_attack_paths_cleanup(self, mock_cleanup): + from tasks.tasks import cleanup_stale_attack_paths_scans_task + + result = cleanup_stale_attack_paths_scans_task.run() + + assert result == {"cleaned_up_count": 1, "scan_ids": ["scan-id"]} + mock_cleanup.assert_called_once_with() + + +@pytest.mark.django_db +class TestAttackPathsDbUtilsActivity: + @patch( + "tasks.jobs.attack_paths.db_utils.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + def test_starting_scan_refreshes_updated_at( + self, tenants_fixture, aws_provider, scans_fixture + ): + from tasks.jobs.attack_paths.db_utils import starting_attack_paths_scan + + old_updated_at = datetime.now(tz=UTC) - timedelta(hours=1) + attack_paths_scan = AttackPathsScan.objects.create( + tenant_id=tenants_fixture[0].id, + provider=aws_provider, + scan=scans_fixture[0], + state=StateChoices.SCHEDULED, + ) + AttackPathsScan.objects.filter(id=attack_paths_scan.id).update( + updated_at=old_updated_at + ) + attack_paths_scan.refresh_from_db() + + started = starting_attack_paths_scan( + attack_paths_scan, SimpleNamespace(update_tag=123) + ) + + assert attack_paths_scan.updated_at > old_updated_at + attack_paths_scan.refresh_from_db() + assert started is True + assert attack_paths_scan.updated_at > old_updated_at + + @patch( + "tasks.jobs.attack_paths.db_utils.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + def test_progress_update_refreshes_updated_at( + self, tenants_fixture, aws_provider, scans_fixture + ): + from tasks.jobs.attack_paths.db_utils import update_attack_paths_scan_progress + + old_updated_at = datetime.now(tz=UTC) - timedelta(hours=1) + attack_paths_scan = AttackPathsScan.objects.create( + tenant_id=tenants_fixture[0].id, + provider=aws_provider, + scan=scans_fixture[0], + state=StateChoices.EXECUTING, + ) + AttackPathsScan.objects.filter(id=attack_paths_scan.id).update( + updated_at=old_updated_at + ) + attack_paths_scan.refresh_from_db() + + update_attack_paths_scan_progress(attack_paths_scan, 42) + + assert attack_paths_scan.updated_at > old_updated_at + attack_paths_scan.refresh_from_db() + assert attack_paths_scan.progress == 42 + assert attack_paths_scan.updated_at > old_updated_at + + @pytest.mark.django_db class TestCleanupStaleAttackPathsScans: + @pytest.fixture(autouse=True) + def execute_on_commit_callbacks(self): + with patch( + "tasks.jobs.attack_paths.cleanup.on_commit", + side_effect=lambda callback, **kwargs: callback(), + ): + yield + def _create_executing_scan( - self, tenant, provider, scan=None, started_at=None, worker=None + self, + tenant, + provider, + scan=None, + started_at=None, + updated_at=None, + worker=None, ): """Helper to create an EXECUTING AttackPathsScan with optional Task+TaskResult.""" ap_scan = AttackPathsScan.objects.create( @@ -2716,18 +2900,66 @@ class TestCleanupStaleAttackPathsScans: ap_scan.task = task ap_scan.save(update_fields=["task_id"]) + if updated_at is not None: + AttackPathsScan.objects.filter(id=ap_scan.id).update(updated_at=updated_at) + ap_scan.updated_at = updated_at + return ap_scan, task_result + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + def test_defers_revoke_until_scan_failure_is_persisted( + self, + mock_revoke, + tenants_fixture, + aws_provider, + ): + from tasks.jobs.attack_paths.cleanup import _finalize_failed_scan + + ap_scan, task_result = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + worker="unresponsive-worker@host", + ) + + with patch("tasks.jobs.attack_paths.cleanup.on_commit") as mock_on_commit: + finalized_scan = _finalize_failed_scan( + ap_scan, + StateChoices.EXECUTING, + "Cleanup reason", + task_result=task_result, + revoke=True, + ) + + assert finalized_scan is not None + ap_scan.refresh_from_db() + task_result.refresh_from_db() + assert ap_scan.state == StateChoices.FAILED + assert task_result.status == "FAILURE" + mock_revoke.assert_not_called() + mock_on_commit.assert_called_once() + assert mock_on_commit.call_args.kwargs == {"using": DEFAULT_DB_ALIAS} + + callback = mock_on_commit.call_args.args[0] + callback() + + mock_revoke.assert_called_once_with(task_result, terminate=True) + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") @patch( "tasks.jobs.attack_paths.cleanup.rls_transaction", new=lambda *args, **kwargs: nullcontext(), ) - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False) - def test_cleans_up_scan_with_dead_worker( + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_cleans_up_inactive_scan_with_unresponsive_worker( self, - mock_alive, + mock_ping, + mock_revoke, mock_drop_db, mock_recover, tenants_fixture, @@ -2735,19 +2967,39 @@ class TestCleanupStaleAttackPathsScans: scans_fixture, ): from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + from tasks.jobs.attack_paths.db_utils import mark_scan_finished tenant = tenants_fixture[0] provider = aws_provider - # Recent scan — should still be cleaned up because worker is dead + updated_at = datetime.now(tz=UTC) - timedelta(minutes=31) ap_scan, task_result = self._create_executing_scan( - tenant, provider, worker="dead-worker@host" + tenant, + provider, + updated_at=updated_at, + worker="unresponsive-worker@host", ) + mock_ping.return_value = (set(), {"unresponsive-worker@host"}) - result = cleanup_stale_attack_paths_scans() + with patch( + "tasks.jobs.attack_paths.cleanup.mark_scan_finished", + wraps=mark_scan_finished, + ) as mock_mark_failed: + call_order = MagicMock() + call_order.attach_mock(mock_revoke, "revoke") + call_order.attach_mock(mock_mark_failed, "mark_failed") + call_order.attach_mock(mock_drop_db, "drop_database") + + result = cleanup_stale_attack_paths_scans() assert result["cleaned_up_count"] == 1 assert str(ap_scan.id) in result["scan_ids"] + assert [entry[0] for entry in call_order.mock_calls] == [ + "mark_failed", + "revoke", + "drop_database", + ] + mock_revoke.assert_called_once_with(task_result, terminate=True) mock_drop_db.assert_called_once() mock_recover.assert_called_once() @@ -2756,7 +3008,10 @@ class TestCleanupStaleAttackPathsScans: assert ap_scan.progress == 100 assert ap_scan.completed_at is not None assert ap_scan.ingestion_exceptions == { - "global_error": "Worker dead — cleaned up by periodic task" + "global_error": ( + "Worker unresponsive and scan inactive for 30 minutes - " + "cleaned up by periodic task" + ) } task_result.refresh_from_db() @@ -2770,10 +3025,10 @@ class TestCleanupStaleAttackPathsScans: new=lambda *args, **kwargs: nullcontext(), ) @patch("tasks.jobs.attack_paths.cleanup._revoke_task") - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=True) - def test_revokes_and_cleans_scan_exceeding_threshold_on_live_worker( + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_revokes_and_cleans_scan_exceeding_threshold_on_responsive_worker( self, - mock_alive, + mock_ping, mock_revoke, mock_drop_db, mock_recover, @@ -2790,11 +3045,12 @@ class TestCleanupStaleAttackPathsScans: ap_scan, task_result = self._create_executing_scan( tenant, provider, started_at=old_start, worker="live-worker@host" ) + mock_ping.return_value = ({"live-worker@host"}, set()) result = cleanup_stale_attack_paths_scans() assert result["cleaned_up_count"] == 1 - mock_revoke.assert_called_once_with(task_result) + mock_revoke.assert_called_once_with(task_result, terminate=True) mock_recover.assert_called_once() ap_scan.refresh_from_db() @@ -2806,10 +3062,10 @@ class TestCleanupStaleAttackPathsScans: "tasks.jobs.attack_paths.cleanup.rls_transaction", new=lambda *args, **kwargs: nullcontext(), ) - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=True) - def test_ignores_recent_executing_scans_on_live_worker( + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_ignores_recent_executing_scans_on_responsive_worker( self, - mock_alive, + mock_ping, mock_drop_db, mock_recover, tenants_fixture, @@ -2821,8 +3077,8 @@ class TestCleanupStaleAttackPathsScans: tenant = tenants_fixture[0] provider = aws_provider - # Recent scan on live worker — should be skipped self._create_executing_scan(tenant, provider, worker="live-worker@host") + mock_ping.return_value = ({"live-worker@host"}, set()) result = cleanup_stale_attack_paths_scans() @@ -2830,6 +3086,130 @@ class TestCleanupStaleAttackPathsScans: mock_drop_db.assert_not_called() mock_recover.assert_not_called() + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_preserves_recent_scan_on_unresponsive_worker( + self, + mock_ping, + mock_revoke, + mock_drop_db, + mock_recover, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + self._create_executing_scan( + tenants_fixture[0], + aws_provider, + updated_at=datetime.now(tz=UTC) - timedelta(minutes=29), + worker="unresponsive-worker@host", + ) + mock_ping.return_value = (set(), {"unresponsive-worker@host"}) + + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == 0 + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() + + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers", return_value=(set(), None)) + def test_final_ping_exception_preserves_pending_worker_scan( + self, + mock_ping, + mock_revoke, + mock_drop_db, + mock_recover, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + ap_scan, _ = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + updated_at=datetime.now(tz=UTC) - timedelta(hours=1), + worker="unknown-worker@host", + ) + + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == 0 + mock_ping.assert_called_once_with({"unknown-worker@host"}) + ap_scan.refresh_from_db() + assert ap_scan.state == StateChoices.EXECUTING + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() + + @pytest.mark.parametrize( + ("inactive_seconds", "should_clean"), + [(29 * 60 + 59, False), (30 * 60, False), (30 * 60 + 1, True)], + ) + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_inactivity_boundary_is_strict( + self, + mock_ping, + mock_revoke, + mock_drop_db, + mock_recover, + inactive_seconds, + should_clean, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + now = datetime.now(tz=UTC) + ap_scan, task_result = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + updated_at=now - timedelta(seconds=inactive_seconds), + worker="unresponsive-worker@host", + ) + mock_ping.return_value = (set(), {"unresponsive-worker@host"}) + + with patch("tasks.jobs.attack_paths.cleanup.datetime") as mock_datetime: + mock_datetime.now.return_value = now + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == int(should_clean) + ap_scan.refresh_from_db() + expected_state = StateChoices.FAILED if should_clean else StateChoices.EXECUTING + assert ap_scan.state == expected_state + if should_clean: + mock_revoke.assert_called_once_with(task_result, terminate=True) + mock_drop_db.assert_called_once() + mock_recover.assert_called_once() + else: + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") @patch( @@ -2868,16 +3248,20 @@ class TestCleanupStaleAttackPathsScans: @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") @patch( "tasks.jobs.attack_paths.cleanup.graph_database.drop_database", - side_effect=Exception("Neo4j unreachable"), + side_effect=[Exception("Neo4j unreachable"), None], ) @patch( "tasks.jobs.attack_paths.cleanup.rls_transaction", new=lambda *args, **kwargs: nullcontext(), ) - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False) - def test_handles_drop_database_failure_gracefully( + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + @patch("tasks.jobs.attack_paths.cleanup.logger") + def test_neo4j_failure_leaves_scan_failed_and_continues( self, - mock_alive, + mock_logger, + mock_ping, + mock_revoke, mock_drop_db, mock_recover, tenants_fixture, @@ -2889,12 +3273,75 @@ class TestCleanupStaleAttackPathsScans: tenant = tenants_fixture[0] provider = aws_provider - self._create_executing_scan(tenant, provider, worker="dead-worker@host") + updated_at = datetime.now(tz=UTC) - timedelta(minutes=31) + ap_scan_1, _ = self._create_executing_scan( + tenant, + provider, + updated_at=updated_at, + worker="unresponsive-worker-1@host", + ) + ap_scan_2, _ = self._create_executing_scan( + tenant, + provider, + updated_at=updated_at, + worker="unresponsive-worker-2@host", + ) + mock_ping.return_value = ( + set(), + {"unresponsive-worker-1@host", "unresponsive-worker-2@host"}, + ) result = cleanup_stale_attack_paths_scans() - assert result["cleaned_up_count"] == 1 - mock_drop_db.assert_called_once() + assert result["cleaned_up_count"] == 2 + assert mock_revoke.call_count == 2 + assert mock_drop_db.call_count == 2 + mock_logger.exception.assert_called_once() + ap_scan_1.refresh_from_db() + ap_scan_2.refresh_from_db() + assert ap_scan_1.state == StateChoices.FAILED + assert ap_scan_2.state == StateChoices.FAILED + + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch( + "tasks.jobs.attack_paths.cleanup.mark_scan_finished", + side_effect=DatabaseError("PostgreSQL unavailable"), + ) + @patch("tasks.jobs.attack_paths.cleanup.logger") + def test_postgresql_failure_prevents_revoke_and_neo4j_deletion( + self, + mock_logger, + mock_mark_failed, + mock_ping, + mock_revoke, + mock_drop_db, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + self._create_executing_scan( + tenants_fixture[0], + aws_provider, + updated_at=datetime.now(tz=UTC) - timedelta(minutes=31), + worker="unresponsive-worker@host", + ) + mock_ping.return_value = (set(), {"unresponsive-worker@host"}) + + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == 0 + mock_mark_failed.assert_called_once() + mock_logger.exception.assert_called_once() + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") @@ -2902,10 +3349,12 @@ class TestCleanupStaleAttackPathsScans: "tasks.jobs.attack_paths.cleanup.rls_transaction", new=lambda *args, **kwargs: nullcontext(), ) - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") def test_cross_tenant_cleanup( self, - mock_alive, + mock_ping, + mock_revoke, mock_drop_db, mock_recover, tenants_fixture, @@ -2924,16 +3373,28 @@ class TestCleanupStaleAttackPathsScans: tenant_id=tenant2.id, ) + updated_at = datetime.now(tz=UTC) - timedelta(minutes=31) ap_scan1, _ = self._create_executing_scan( - tenant1, provider1, worker="dead-worker-1@host" + tenant1, + provider1, + updated_at=updated_at, + worker="unresponsive-worker-1@host", ) ap_scan2, _ = self._create_executing_scan( - tenant2, provider2, worker="dead-worker-2@host" + tenant2, + provider2, + updated_at=updated_at, + worker="unresponsive-worker-2@host", + ) + mock_ping.return_value = ( + set(), + {"unresponsive-worker-1@host", "unresponsive-worker-2@host"}, ) result = cleanup_stale_attack_paths_scans() assert result["cleaned_up_count"] == 2 + assert mock_revoke.call_count == 2 assert mock_recover.call_count == 2 ap_scan1.refresh_from_db() @@ -2947,10 +3408,12 @@ class TestCleanupStaleAttackPathsScans: "tasks.jobs.attack_paths.cleanup.rls_transaction", new=lambda *args, **kwargs: nullcontext(), ) - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") def test_recovers_graph_data_ready_for_stale_scan( self, - mock_alive, + mock_ping, + mock_revoke, mock_drop_db, mock_recover, tenants_fixture, @@ -2963,11 +3426,16 @@ class TestCleanupStaleAttackPathsScans: provider = aws_provider ap_scan, _ = self._create_executing_scan( - tenant, provider, worker="dead-worker@host" + tenant, + provider, + updated_at=datetime.now(tz=UTC) - timedelta(minutes=31), + worker="unresponsive-worker@host", ) + mock_ping.return_value = (set(), {"unresponsive-worker@host"}) cleanup_stale_attack_paths_scans() + mock_revoke.assert_called_once() mock_recover.assert_called_once() recovered_scan = mock_recover.call_args[0][0] assert recovered_scan.id == ap_scan.id @@ -3013,10 +3481,57 @@ class TestCleanupStaleAttackPathsScans: "tasks.jobs.attack_paths.cleanup.rls_transaction", new=lambda *args, **kwargs: nullcontext(), ) - @patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False) - def test_shared_worker_is_pinged_only_once( + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_preserves_scans_without_a_started_at_timestamp( self, - mock_alive, + mock_ping, + mock_revoke, + mock_drop_db, + mock_recover, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + responsive_scan, _ = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + worker="responsive-worker@host", + ) + AttackPathsScan.objects.filter(id=responsive_scan.id).update(started_at=None) + no_worker_scan = AttackPathsScan.objects.create( + tenant_id=tenants_fixture[0].id, + provider=aws_provider, + state=StateChoices.EXECUTING, + started_at=None, + ) + mock_ping.return_value = ({"responsive-worker@host"}, set()) + + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == 0 + responsive_scan.refresh_from_db() + no_worker_scan.refresh_from_db() + assert responsive_scan.state == StateChoices.EXECUTING + assert no_worker_scan.state == StateChoices.EXECUTING + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() + + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + @patch("tasks.jobs.attack_paths.cleanup._ping_workers") + def test_shared_worker_is_collected_only_once( + self, + mock_ping, + mock_revoke, mock_drop_db, mock_recover, tenants_fixture, @@ -3028,15 +3543,114 @@ class TestCleanupStaleAttackPathsScans: tenant = tenants_fixture[0] provider = aws_provider - # Two scans on the same dead worker - self._create_executing_scan(tenant, provider, worker="shared-worker@host") - self._create_executing_scan(tenant, provider, worker="shared-worker@host") + updated_at = datetime.now(tz=UTC) - timedelta(minutes=31) + self._create_executing_scan( + tenant, + provider, + updated_at=updated_at, + worker="shared-worker@host", + ) + self._create_executing_scan( + tenant, + provider, + updated_at=updated_at, + worker="shared-worker@host", + ) + mock_ping.return_value = (set(), {"shared-worker@host"}) result = cleanup_stale_attack_paths_scans() assert result["cleaned_up_count"] == 2 - # Worker should be pinged exactly once — cache prevents second ping - mock_alive.assert_called_once_with("shared-worker@host") + assert mock_revoke.call_count == 2 + mock_ping.assert_called_once_with({"shared-worker@host"}) + + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + def test_locked_recheck_preserves_scan_with_new_activity( + self, + mock_revoke, + mock_drop_db, + mock_recover, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + ap_scan, _ = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + updated_at=datetime.now(tz=UTC) - timedelta(minutes=31), + worker="unresponsive-worker@host", + ) + + def record_activity(_workers): + AttackPathsScan.objects.filter(id=ap_scan.id).update( + updated_at=datetime.now(tz=UTC) + ) + return set(), {"unresponsive-worker@host"} + + with patch( + "tasks.jobs.attack_paths.cleanup._ping_workers", + side_effect=record_activity, + ): + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == 0 + ap_scan.refresh_from_db() + assert ap_scan.state == StateChoices.EXECUTING + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() + + @patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready") + @patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database") + @patch( + "tasks.jobs.attack_paths.cleanup.rls_transaction", + new=lambda *args, **kwargs: nullcontext(), + ) + @patch("tasks.jobs.attack_paths.cleanup._revoke_task") + def test_locked_recheck_preserves_scan_that_changed_state( + self, + mock_revoke, + mock_drop_db, + mock_recover, + tenants_fixture, + aws_provider, + scans_fixture, + ): + from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans + + ap_scan, _ = self._create_executing_scan( + tenants_fixture[0], + aws_provider, + updated_at=datetime.now(tz=UTC) - timedelta(minutes=31), + worker="unresponsive-worker@host", + ) + + def complete_scan(_workers): + AttackPathsScan.objects.filter(id=ap_scan.id).update( + state=StateChoices.COMPLETED + ) + return set(), {"unresponsive-worker@host"} + + with patch( + "tasks.jobs.attack_paths.cleanup._ping_workers", + side_effect=complete_scan, + ): + result = cleanup_stale_attack_paths_scans() + + assert result["cleaned_up_count"] == 0 + ap_scan.refresh_from_db() + assert ap_scan.state == StateChoices.COMPLETED + mock_revoke.assert_not_called() + mock_drop_db.assert_not_called() + mock_recover.assert_not_called() # `SCHEDULED` state cleanup def _create_scheduled_scan( @@ -3123,7 +3737,7 @@ class TestCleanupStaleAttackPathsScans: assert ap_scan.progress == 100 assert ap_scan.completed_at is not None assert ap_scan.ingestion_exceptions == { - "global_error": "Scan never started — cleaned up by periodic task" + "global_error": "Scan never started - cleaned up by periodic task" } # SCHEDULED revoke must NOT terminate a running worker diff --git a/api/src/backend/tasks/tests/test_orphan_recovery.py b/api/src/backend/tasks/tests/test_orphan_recovery.py index b78aca4e63..074e311b70 100644 --- a/api/src/backend/tasks/tests/test_orphan_recovery.py +++ b/api/src/backend/tasks/tests/test_orphan_recovery.py @@ -7,6 +7,7 @@ from celery import states from django.test import override_settings from django_celery_results.models import TaskResult from tasks.jobs.orphan_recovery import ( + _SKIP_RECOVERY, _decode_celery_field, _reconcile_task_results, _recovery_attempt_count, @@ -14,6 +15,7 @@ from tasks.jobs.orphan_recovery import ( is_worker_alive, reconcile_orphans, reenqueueable_tasks, + revoke_task, ) @@ -180,10 +182,18 @@ class TestReconcileTaskResults: assert tr.task_id in result["failed"] mock_count.assert_not_called() - def test_scan_task_is_skipped_entirely(self, tenants_fixture): + @pytest.mark.parametrize( + "task_name", + [ + "scan-perform", + "attack-paths-scan-perform", + "attack-paths-cleanup-stale-scans", + ], + ) + def test_scan_task_is_skipped_entirely(self, tenants_fixture, task_name): """Scan tasks are excluded from recovery: the watchdog never touches them.""" tr = _orphan_result( - name="scan-perform", + name=task_name, kwargs={ "tenant_id": str(tenants_fixture[0].id), "scan_id": str(uuid4()), @@ -339,6 +349,15 @@ class TestOrphanRecoveryHelpers: ): assert is_worker_alive("w@h") is False + def test_revoke_task_terminates_with_sigterm_by_default(self): + task_result = MagicMock(task_id="task-id") + with patch( + "tasks.jobs.orphan_recovery.current_app.control.revoke" + ) as mock_revoke: + revoke_task(task_result) + + mock_revoke.assert_called_once_with("task-id", terminate=True, signal="SIGTERM") + def test_recovery_attempt_count_increments(self): # Unique signature so the Valkey counter starts fresh for this test. kwargs_repr = repr({"probe": str(uuid4())}) @@ -350,6 +369,12 @@ class TestOrphanRecoveryHelpers: class TestRecoveryFeatureFlags: + def test_attack_paths_tasks_are_excluded_from_generic_recovery(self): + assert { + "attack-paths-scan-perform", + "attack-paths-cleanup-stale-scans", + } <= _SKIP_RECOVERY + def test_all_groups_enabled_by_default(self): tasks = reenqueueable_tasks() assert "scan-summary" in tasks @@ -374,33 +399,23 @@ class TestRecoveryFeatureFlags: class TestRecoveryMasterFlag: @override_settings(TASK_RECOVERY_ENABLED=False) def test_master_flag_disables_task_recovery(self): - with ( - patch( - "tasks.jobs.orphan_recovery._reconcile_task_results" - ) as mock_reconcile, - patch( - "tasks.jobs.attack_paths.cleanup.cleanup_stale_attack_paths_scans", - return_value={}, - ), - ): + with patch( + "tasks.jobs.orphan_recovery._reconcile_task_results" + ) as mock_reconcile: result = reconcile_orphans(grace_minutes=2, max_attempts=3, dry_run=False) mock_reconcile.assert_not_called() assert result["acquired"] is True assert result["enabled"] is False + assert "attack_paths" not in result @override_settings(TASK_RECOVERY_ENABLED=True) def test_master_flag_enabled_runs_task_recovery(self): - with ( - patch( - "tasks.jobs.orphan_recovery._reconcile_task_results", - return_value={"recovered": [], "failed": [], "skipped": []}, - ) as mock_reconcile, - patch( - "tasks.jobs.attack_paths.cleanup.cleanup_stale_attack_paths_scans", - return_value={}, - ), - ): - reconcile_orphans(grace_minutes=2, max_attempts=3, dry_run=False) + with patch( + "tasks.jobs.orphan_recovery._reconcile_task_results", + return_value={"recovered": [], "failed": [], "skipped": []}, + ) as mock_reconcile: + result = reconcile_orphans(grace_minutes=2, max_attempts=3, dry_run=False) mock_reconcile.assert_called_once() + assert "attack_paths" not in result