code review fix 2

Author: NeptuneHub

app_helper.py

@@ -46,6 +46,7 @@
     load_map_projection,
     get_task_info_from_db,
     get_task_statuses,
+    main_task_start_lock,
     get_tracks_by_ids,
     save_track_analysis_and_embedding,
     # Used internally by the build_and_store_* projection orchestration below.
@@ -448,6 +449,15 @@ def cancel_job_and_children_recursive(
     REVOKED row for the requested `job_id` (so UI sees one canonical cancelled task).
     This is intentionally simple and destructive (as requested).
     """
+    # Serialize against the start paths. Cancel scans RQ, then wipes task_status;
+    # a starter that had committed its PENDING row but not yet enqueued was invisible
+    # to the scan AND had its row deleted, so it enqueued afterwards and ran after the
+    # user pressed Cancel.
+    with main_task_start_lock():
+        return _cancel_job_and_children_locked(job_id, reason)
+
+
+def _cancel_job_and_children_locked(job_id, reason):
     cancelled_count = 0
 
     # --- Scan RQ for job ids to cancel ---

app_provider_migration.py

@@ -97,14 +97,15 @@ def __getattr__(self, name):
 _COMPLETED_SESSIONS_KEPT = 10
 
 
-def _migration_job_in_flight(cur):
+def _migration_job_in_flight(cur, keys=('dry_run_task_id', 'source_refresh_task_id',
+                                        'exec_task_id')):
     # Dry-run and source-refresh write no task row and hold no lock, so their only
     # durable trace is the RQ job id parked on the session. Without this a second
     # and third session start would prune the session out from under a running dry
     # run, and the job would fail or silently update nothing.
     cur.execute(
-        "SELECT state->>'dry_run_task_id', state->>'source_refresh_task_id', "
-        "state->>'exec_task_id' FROM migration_session"
+        "SELECT " + ", ".join("state->>'%s'" % k for k in keys)
+        + " FROM migration_session"  # nosec B608 - keys are module constants
     )
     job_ids = [j for row in (cur.fetchall() or []) for j in row if j]
     if not job_ids:
@@ -1584,6 +1585,20 @@ def execute():
         return jsonify({'error': 'session not found'}), 404
     target_type, status, is_current_session = row[0], row[1], row[2]
 
+    # Refuse while a dry run or source refresh is still producing a plan: the
+    # finalized numbers this request confirmed would be applied against a mapping
+    # another worker is in the middle of rewriting.
+    with db.cursor() as planning:
+        if _migration_job_in_flight(
+            planning, keys=('dry_run_task_id', 'source_refresh_task_id')
+        ):
+            return jsonify(
+                {
+                    'error': 'A dry run is still building the plan. Wait for it to '
+                             'finish, then confirm the numbers again.'
+                }
+            ), 409
+
     # A session kept alive by the prune is still reachable by id. Executing an older
     # one would repoint the catalogue using a superseded plan.
     if not is_current_session:

restart_listener.py

@@ -23,13 +23,15 @@
 
 import logging
 import os
+import socket
 import threading
 import time
 
 from app_logging import configure_logging
 from taskqueue import new_redis_connection
 from restart_manager import (
     RESTART_CHANNEL,
+    WORKER_PRESENCE_PREFIX,
     restart_supervisor_workers,
     stop_supervisor_workers,
     start_supervisor_workers,
@@ -38,6 +40,21 @@
 logger = logging.getLogger(__name__)
 configure_logging()
 
+PRESENCE_REFRESH_SECONDS = 15
+
+PRESENCE_TTL_SECONDS = 60
+
+
+def presence_key():
+    return f"{WORKER_PRESENCE_PREFIX}{socket.gethostname()}:{os.getpid()}"
+
+
+def announce_presence(redis_conn):
+    try:
+        redis_conn.set(presence_key(), '1', ex=PRESENCE_TTL_SECONDS)
+    except Exception:
+        logger.exception('Could not refresh the worker restart-listener presence key')
+
 try:
     from plugin.manager import worker_presync
 except Exception:
@@ -76,14 +93,22 @@ def main():
             pubsub.subscribe(channel)
             logger.info('Subscribed to restart channel. Waiting for restart messages...')
 
-            for message in pubsub.listen():
+            # get_message with a timeout instead of listen(): the presence key below
+            # has to be refreshed on a schedule, and listen() blocks forever between
+            # messages. Only a WORKER-role listener registers, because only it acts
+            # on a signal - the Flask container subscribes and ignores, which is why
+            # counting PUBLISH subscribers could never prove delivery.
+            while True:
+                service_type = os.environ.get('SERVICE_TYPE', '').lower()
+                if service_type == 'worker':
+                    announce_presence(redis_conn)
+                message = pubsub.get_message(timeout=PRESENCE_REFRESH_SECONDS)
                 if not message:
                     continue
                 if message.get('type') != 'message':
                     continue
                 payload = message.get('data')
                 logger.info('Control listener received signal: %s', payload)
-                service_type = os.environ.get('SERVICE_TYPE', '').lower()
                 if service_type != 'worker':
                     logger.info('Control signal received, but SERVICE_TYPE is not worker; skipping')
                     continue

restart_manager.py

@@ -33,6 +33,8 @@
 SUPERVISOR_CONF = os.environ.get('SUPERVISOR_CONF', '/etc/supervisor/conf.d/supervisord.conf')
 logger = logging.getLogger(__name__)
 
+WORKER_PRESENCE_PREFIX = 'audiomuse:worker_restart_listener:'
+
 FLASK_SERVICE = ['flask']
 WORKER_SERVICES = ['rq-worker-default', 'rq-worker-high', 'rq-janitor']
 
@@ -44,14 +46,24 @@ def publish_control_request(action):
             socket_timeout=5,
             decode_responses=True,
         )
-        delivered = redis_conn.publish(RESTART_CHANNEL, action)
-        if not delivered:
+        # Subscriber COUNT proves nothing: the Flask container also subscribes and
+        # then ignores every signal, so publish() returns >= 1 even when the worker
+        # listener - the only one that acts - is dead. Worker listeners announce
+        # themselves with a TTL key instead, so absence is a real answer.
+        listeners = 0
+        try:
+            for _ in redis_conn.scan_iter(match=f'{WORKER_PRESENCE_PREFIX}*', count=100):
+                listeners += 1
+                break
+        except Exception:
+            logger.exception('Could not check for live worker restart listeners')
+            listeners = -1
+        redis_conn.publish(RESTART_CHANNEL, action)
+        if listeners == 0:
             logger.error(
-                'PUBLISHED %s TO REDIS BUT NO LISTENER RECEIVED IT. Every deployment '
-                'runs a restart-listener (supervisord config-restart-listener, or the '
-                'native supervisor role), so zero subscribers means it is not running '
-                'and the workers are STILL ON THE OLD CONFIGURATION. Restart the '
-                'container or the app.',
+                'PUBLISHED %s BUT NO WORKER RESTART-LISTENER IS ALIVE. The workers are '
+                'STILL ON THE OLD CONFIGURATION. Check that the worker container (or '
+                'the native worker role) is running its restart listener.',
                 action,
             )
             return False

tasks/multiserver_sync.py

@@ -257,6 +257,10 @@ def recover_abandoned_sweeps():
             "a fresh alignment of all servers was enqueued."
         )
         details = json.dumps({'message': message, 'status_message': message})
+        # One transaction: revoking the abandoned sweep and inserting its
+        # replacement used to be two autocommits, so a crash in between retired the
+        # alignment and left nothing to take its place.
+        db.autocommit = False
         cur = db.cursor()
         try:
             cur.execute(
@@ -270,15 +274,21 @@ def recover_abandoned_sweeps():
             )
             revoked_count = cur.rowcount
             if not revoked_count:
+                db.rollback()
                 return None
             new_task_id = str(uuid.uuid4())
             insert_pending_sweep_row(
                 cur, new_task_id,
                 'Server alignment queued for all servers.',
                 full_refresh=full_refresh,
             )
+            db.commit()
+        except Exception:
+            db.rollback()
+            raise
         finally:
             cur.close()
+            db.autocommit = True
         rq_queue_high.enqueue(
             'tasks.multiserver_sync.sweep_all_secondary_servers',
             kwargs={'task_id': new_task_id, 'full_refresh': full_refresh},
@@ -299,6 +309,8 @@ def recover_abandoned_sweeps():
             logger.debug("Recovery connection close failed", exc_info=True)
 
 
+_JANITOR_LOCK_KEY = 6193044728150337
+
 _ORPHAN_GRACE_SECONDS = 120
 
 _INLINE_STALE_SECONDS = 1800
@@ -366,6 +378,14 @@ def reap_orphaned_tasks():
 
     db = connect_raw()
     db.autocommit = True
+    # One janitor at a time. Two of them probing the same abandoned job both
+    # decided to requeue it, because retrying had no claim anywhere. The lock is
+    # session scoped, so closing this connection releases it.
+    with db.cursor() as claim:
+        claim.execute("SELECT pg_try_advisory_lock(%s)", (_JANITOR_LOCK_KEY,))
+        if not claim.fetchone()[0]:
+            db.close()
+            return 0
     missing = []
     terminal = {}
     restarted = []

test/unit/test_app_cancel.py

@@ -22,11 +22,19 @@
 * A failed status QUERY is not an empty answer, and leaves the task running
 """
 
+from contextlib import nullcontext
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 
+@pytest.fixture(autouse=True)
+def stub_the_start_lock(monkeypatch):
+    import app_helper
+
+    monkeypatch.setattr(app_helper, 'main_task_start_lock', nullcontext)
+
+
 class _FakeCursor:
     """Records executed SQL and answers the snapshot SELECT."""