feat(mcp): queue long-running calculations

2026-10-29 18:25:15 +08:00
parent 2865811f1b
commit 63af998045
32 changed files with 1588 additions and 72 deletions
@@ -0,0 +1,241 @@
+"""Worker boundaries use real managed subprocesses, without external databases."""
+import asyncio
+from functools import wraps
+from types import SimpleNamespace
+from contextlib import asynccontextmanager
+import os
+from pathlib import Path
+import signal
+import sys
+
+import pytest
+
+from app.services.calculation_worker import CalculationWorker, WorkerSettings, stop_process_group
+
+
+@pytest.fixture(autouse=True)
+def no_external_cleanup(monkeypatch):
+    async def cleanup(*args):
+        pass
+    monkeypatch.setattr('app.services.calculation_worker.cleanup_job_resources', cleanup)
+
+
+def run_async(function):
+    @wraps(function)
+    def wrapper(*args, **kwargs):
+        return asyncio.run(function(*args, **kwargs))
+    return wrapper
+
+
+def test_worker_settings_require_explicit_valid_scope(monkeypatch):
+    monkeypatch.setenv('CALCULATION_WORKER_PROJECT_IDS', '6b39b128-2ba8-454f-9bbe-e9553902ab53')
+    settings = WorkerSettings.from_environment()
+    assert settings.max_concurrency == 2
+    monkeypatch.setenv('CALCULATION_WORKER_HEARTBEAT_SECONDS', '30')
+    with pytest.raises(ValueError):
+        WorkerSettings.from_environment()
+
+
+@run_async
+async def test_process_group_cleanup_kills_native_style_descendant(tmp_path):
+    pid_file = tmp_path / 'pid'
+    code = ('import subprocess,time,pathlib,sys; '
+            'p=subprocess.Popen([sys.executable,"-c","import time; time.sleep(60)"]); '
+            'pathlib.Path(sys.argv[1]).write_text(str(p.pid)); time.sleep(60)')
+    process = await asyncio.create_subprocess_exec(sys.executable, '-c', code,
+        str(pid_file), start_new_session=True)
+    try:
+        for _ in range(100):
+            if pid_file.exists():
+                break
+            await asyncio.sleep(.02)
+        child_pid = int(pid_file.read_text())
+        await stop_process_group(process)
+        assert process.returncode is not None
+        # Linux may retain a dead grandchild as a zombie until PID 1 reaps it.
+        stat_path = Path(f'/proc/{child_pid}/stat')
+        for _ in range(100):
+            if not stat_path.exists() or stat_path.read_text().split()[2] == 'Z':
+                break
+            await asyncio.sleep(.01)
+        assert not stat_path.exists() or stat_path.read_text().split()[2] == 'Z'
+    finally:
+        await stop_process_group(process)
+
+
+@run_async
+@pytest.mark.parametrize('lease_status, expected', [('cancel_requested', 'cancelled'), (None, None)])
+async def test_cancel_and_lost_lease_stop_process_before_persisting(monkeypatch, lease_status, expected):
+    processes = []
+    original = asyncio.create_subprocess_exec
+    async def spawn(*args, **kwargs):
+        process = await original(sys.executable, '-c', 'import time; time.sleep(60)',
+                                 start_new_session=True)
+        processes.append(process)
+        return process
+    monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
+    finishes = []
+    class Repository:
+        async def heartbeat(self, *args):
+            return lease_status
+        async def finish(self, row, owner, status, **kwargs):
+            assert processes[0].returncode is not None
+            finishes.append(status)
+            return True
+    @asynccontextmanager
+    async def repository(project):
+        yield Repository()
+    worker = CalculationWorker(WorkerSettings(('project',), heartbeat_seconds=.05), repository=repository)
+    await worker.run_job({'project_id': 'project', 'job_id': 'job', 'timeout_seconds': 60})
+    assert finishes == ([expected] if expected else [])
+    assert processes[0].returncode is not None
+
+
+@run_async
+async def test_timeout_stops_process_before_recording(monkeypatch):
+    original = asyncio.create_subprocess_exec
+    processes = []
+    async def spawn(*args, **kwargs):
+        p = await original(sys.executable, '-c', 'import time; time.sleep(60)', start_new_session=True)
+        processes.append(p)
+        return p
+    monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
+    statuses = []
+    class Repository:
+        async def heartbeat(self, *args):
+            return 'running'
+        async def finish(self, row, owner, status, **kwargs):
+            assert processes[0].returncode is not None
+            statuses.append(status)
+            return True
+    @asynccontextmanager
+    async def repository(project):
+        yield Repository()
+    worker = CalculationWorker(WorkerSettings(('project',), heartbeat_seconds=.02), repository=repository)
+    await worker.run_job({'project_id': 'project', 'timeout_seconds': 0})
+    assert statuses == ['timed_out']
+
+
+@run_async
+async def test_duplicate_supervisor_never_claims(monkeypatch):
+    from app.infra.db.metadb import database
+    class Lock:
+        async def scalar(self, statement):
+            assert 'pg_try_advisory_lock' in str(statement)
+            return False
+        async def commit(self):
+            pass
+    @asynccontextmanager
+    async def connection():
+        yield Lock()
+    monkeypatch.setattr(database, 'engine', SimpleNamespace(connect=connection))
+    @asynccontextmanager
+    async def repository(project):
+        pytest.fail('A duplicate supervisor must never open the queue')
+        yield
+    with pytest.raises(RuntimeError, match='already active'):
+        await CalculationWorker(WorkerSettings(('project',)), repository=repository).run()
+
+
+@run_async
+async def test_repository_terminal_updates_are_fenced():
+    from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
+    statements = []
+    class Cursor:
+        async def fetchone(self):
+            return None
+    class Connection:
+        async def execute(self, sql, parameters):
+            statements.append((sql, parameters))
+            return Cursor()
+    repository = CalculationJobRepository(Connection(), 'project')
+    row = {'job_id': 'job', 'lease_generation': 7}
+    assert await repository.finish(row, 'owner', 'succeeded', result={'original': True}) is False
+    sql, parameters = statements[-1]
+    assert 'lease_generation=%s' in sql and 'lease_owner=%s' in sql
+    assert 'lease_expires_at>clock_timestamp()' in sql
+    assert "status!='cancel_requested'" in sql
+    assert 7 in parameters
+
+
+@run_async
+async def test_guard_parent_death_pipe_kills_entire_group(tmp_path):
+    pid_file = tmp_path / 'guard-child'
+    child = 'import os,time,pathlib,sys; pathlib.Path(sys.argv[1]).write_text(str(os.getpid())); time.sleep(60)'
+    launcher = '''import asyncio,sys
+from app.services.calculation_worker import guard_child
+original = asyncio.create_subprocess_exec
+async def spawn(*args, **kwargs):
+    return await original(sys.executable, '-c', sys.argv[2], sys.argv[1])
+asyncio.create_subprocess_exec = spawn
+asyncio.run(guard_child('unused','unused'))
+'''
+    guard = await asyncio.create_subprocess_exec(sys.executable, '-c', launcher,
+        str(pid_file), child, stdin=asyncio.subprocess.PIPE, start_new_session=True)
+    try:
+        for _ in range(100):
+            if pid_file.exists():
+                break
+            await asyncio.sleep(.02)
+        child_pid = int(pid_file.read_text())
+        guard.stdin.close()
+        await asyncio.wait_for(guard.wait(), 3)
+        assert guard.returncode == -signal.SIGKILL
+        stat = Path(f'/proc/{child_pid}/stat')
+        for _ in range(100):
+            if not stat.exists() or stat.read_text().split()[2] == 'Z':
+                break
+            await asyncio.sleep(.01)
+        assert not stat.exists() or stat.read_text().split()[2] == 'Z'
+    finally:
+        await stop_process_group(guard)
+
+
+@run_async
+async def test_cleanup_failure_is_not_reported_cancelled(monkeypatch, tmp_path):
+    monkeypatch.setenv('CALCULATION_WORKER_STATE_DIR', str(tmp_path))
+    async def cleanup(row, directory):
+        (Path(directory) / 'cleanup-failed').touch()
+        raise RuntimeError('external resource unavailable')
+    monkeypatch.setattr('app.services.calculation_worker.cleanup_job_resources', cleanup)
+    original = asyncio.create_subprocess_exec
+    async def spawn(*args, **kwargs):
+        return await original(sys.executable, '-c', 'import time; time.sleep(60)', start_new_session=True)
+    monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
+    results = []
+    class Repository:
+        async def heartbeat(self, *args):
+            return 'cancel_requested'
+        async def finish(self, row, owner, status, **kwargs):
+            results.append((status, kwargs['error']['code']))
+            return True
+    @asynccontextmanager
+    async def repository(project):
+        yield Repository()
+    worker = CalculationWorker(WorkerSettings(('project',)), repository=repository)
+    await worker.run_job({'project_id': 'project', 'timeout_seconds': 60})
+    assert results == [('failed', 'CLEANUP_FAILED')]
+    assert len(list(tmp_path.glob('*/cleanup-failed'))) == 1
+
+
+@run_async
+async def test_startup_recovery_checks_persisted_generation(monkeypatch, tmp_path):
+    import json
+    from app.services.calculation_worker import recover_job_resources
+    project = '6b39b128-2ba8-454f-9bbe-e9553902ab53'
+    job = '2d3a6231-a8ee-4acb-bb13-58328d29b88f'
+    directory = tmp_path / (job + '-owned')
+    directory.mkdir()
+    (directory / 'scope.json').write_text(json.dumps({
+        'project_id': project, 'job_id': job, 'lease_generation': 1}))
+    monkeypatch.setenv('CALCULATION_WORKER_STATE_DIR', str(tmp_path))
+    class Repository:
+        async def get(self, job_id):
+            return {'lease_generation': 2}
+    @asynccontextmanager
+    async def repository(project):
+        yield Repository()
+    monkeypatch.setattr('app.services.calculation_worker.queue_repository', repository)
+    with pytest.raises(RuntimeError, match='ownership'):
+        await recover_job_resources((project,))
+    assert directory.exists()