242 lines
9.6 KiB
Python
242 lines
9.6 KiB
Python
"""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()
|