diff --git a/Dockerfile b/Dockerfile index 3377783..59b03ff 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,4 +6,4 @@ WORKDIR /app COPY requirements.txt requirements.txt RUN pip install -r requirements.txt -CMD ["gunicorn", "-c", "gunicorn.conf.py", "app:app"] +CMD ["python", "main.py"] diff --git a/app.py b/app.py deleted file mode 100644 index 2c9e15a..0000000 --- a/app.py +++ /dev/null @@ -1,89 +0,0 @@ -import os -import logging -import queue -import secrets -from flask import Flask, request, jsonify -from dispatcher.constant import Language -from dispatcher.dispatcher import Dispatcher -from dispatcher import file_manager -from dispatcher.testdata import ( - ensure_testdata, - get_problem_meta, - get_problem_root, -) -from dispatcher.config import (SANDBOX_TOKEN, SUBMISSION_DIR) - -logging.basicConfig(filename='logs/sandbox.log') -app = Flask(__name__) -if __name__ != '__main__': - # let flask app use gunicorn's logger - gunicorn_logger = logging.getLogger('gunicorn.error') - app.logger.handlers = gunicorn_logger.handlers - app.logger.setLevel(gunicorn_logger.level) - logging.getLogger().setLevel(gunicorn_logger.level) -logger = app.logger - -# setup dispatcher -DISPATCHER_CONFIG = os.getenv( - 'DISPATCHER_CONFIG', - '.config/dispatcher.json.example', -) -DISPATCHER = Dispatcher(DISPATCHER_CONFIG) -DISPATCHER.start() - - -@app.post('/submit/') -def submit(submission_id: str): - token = request.values.get('token', '') - if not secrets.compare_digest(token, SANDBOX_TOKEN): - logger.debug(f'get invalid token: {token}') - return 'invalid token', 403 - # Ensure the testdata is up to data - problem_id = request.form.get('problem_id', type=int) - if problem_id is None: - return 'missing problen id', 400 - ensure_testdata(problem_id) - language = Language(request.form.get('language', type=int)) - try: - file_manager.extract( - root_dir=SUBMISSION_DIR, - job_id=submission_id, - meta=get_problem_meta(problem_id, language), - source=request.files['src'], - testdata=get_problem_root(problem_id), - ) - except ValueError as e: - return str(e), 400 - logger.debug(f'send submission {submission_id} to dispatcher') - try: - DISPATCHER.handle(job_id=submission_id, submission_id=submission_id) - except queue.Full: - return jsonify({ - 'status': 'err', - 'msg': 'task queue is full now.\n' - 'please wait a moment and re-send the submission.', - 'data': None, - }), 500 - return jsonify({ - 'status': 'ok', - 'msg': 'ok', - 'data': 'ok', - }) - - -@app.get('/status') -def status(): - ret = { - 'load': DISPATCHER.queue.qsize() / DISPATCHER.MAX_TASK_COUNT, - } - # if token is provided - if secrets.compare_digest(SANDBOX_TOKEN, request.args.get('token', '')): - ret.update({ - 'queueSize': DISPATCHER.queue.qsize(), - 'maxTaskCount': DISPATCHER.MAX_TASK_COUNT, - 'containerCount': DISPATCHER.container_count, - 'maxContainerCount': DISPATCHER.MAX_TASK_COUNT, - 'submissions': [*DISPATCHER.result.keys()], - 'running': DISPATCHER.do_run, - }) - return jsonify(ret), 200 diff --git a/dispatcher/config.py b/dispatcher/config.py index 7dac6ca..b9b3fa8 100644 --- a/dispatcher/config.py +++ b/dispatcher/config.py @@ -2,11 +2,13 @@ from pathlib import Path # backend config -BACKEND_API = os.getenv( - 'BACKEND_API', +BACKEND_URL = os.getenv( + 'BACKEND_URL', 'http://web:8080', ) -# sandbox token +# Shared secret that now ONLY authenticates the legacy testdata channel +# (backend /problem//meta|testdata|checksum ?token= query). Its retirement +# is coordinated with the backend keystone re-authing those endpoints. SANDBOX_TOKEN = os.getenv( 'SANDBOX_TOKEN', 'KoNoSandboxDa', diff --git a/dispatcher/dispatcher.py b/dispatcher/dispatcher.py index 78bc702..3cd1423 100644 --- a/dispatcher/dispatcher.py +++ b/dispatcher/dispatcher.py @@ -2,28 +2,59 @@ import os import threading import time -import requests import pathlib import queue -import tempfile -from datetime import datetime from executor.submission import SubmissionExecutor -from . import job, file_manager, config +from . import job, config from .exception import * from .meta import Meta from .constant import Language from .utils import logger +class JobContext: + """Everything one handle() generation owns, in one object. + + A fresh instance per handle() IS the generation token: queue entries and + worker threads carry the instance they were created under, and anything + presenting a stale instance is dropped on an identity check. Holding ALL + per-generation state here -- instead of parallel dicts keyed by job_id -- + is what makes the binding airtight: there is no per-job state left that a + stale generation could reach, or poison, through its job_id alone. + """ + + __slots__ = ('config', 'submission_id', 'cases', 'compile_result', + 'started') + + def __init__(self, config, submission_id, cases): + # the parsed Meta for this generation + self.config = config + self.submission_id = submission_id + # case_no -> result dict; None until the case resolves + self.cases = cases + # written once by this generation's compile worker + self.compile_result = None + # set by the run loop when this generation's work is first picked up; + # a drain only returns generations that never started + self.started = False + + class Dispatcher(threading.Thread): def __init__( self, dispatcher_config='.config/dispatcher.json', submission_config='.config/submission.json', + *, + on_complete=None, ): super().__init__() + # Daemon so fail-fast exit paths (heartbeat 401 -> exit without + # drain) are not blocked by this thread: the interpreter waits for + # non-daemon threads, and the run loop only ends via stop(). The + # graceful drain path still stops and joins explicitly. + self.daemon = True self.testing = False # read config d_config = {} @@ -34,19 +65,32 @@ def __init__( self.do_run = True # submission location self.SUBMISSION_DIR = config.SUBMISSION_DIR - # task queue - # type Queue[Tuple[job_id, task_no]] - self.MAX_TASK_COUNT = d_config.get('QUEUE_SIZE', 16) - self.queue = queue.Queue(self.MAX_TASK_COUNT) - # task result - # type: Dict[job_id, Tuple[submission_info, List[result]]] + # Unbounded task queue: in pull mode the poller is the only producer + # and total task volume is capped by the job-level capacity gate + # (max_concurrent_jobs) × per-job case count, so a task-level bound + # only creates spurious queue.Full failures (see handle()). + # type Queue[Union[job.Compile, job.Execute]] + self.queue = queue.Queue() + # type: Dict[job_id, JobContext] -- the CURRENT generation per job. + # Everything else a generation owns lives on its JobContext, so a + # release or re-handle atomically retires the whole generation at + # once; queue entries and workers hold the instance they were created + # under and are dropped on an identity mismatch. self.result = {} - # threading locks for each job - self.locks = {} - self.compile_locks = {} - self.compile_results = {} - # maps job_id -> submission_id, for the backend callback - self.submission_ids = {} + # invoked with (job_id, submission_id, tasks) when a job finishes; + # main.py wires this to the result queue consumed by ResultSenderThread. + self.on_complete = on_complete + # Cleared by stop_accepting()/drain_unstarted(): from then on handle() + # refuses new jobs and the run loop starts no further ones, so nothing + # begins judging after shutdown was requested. + self.accepting = True + # Guards every mutation of the job state above: publication (handle), + # release, the started-set transition, and the generation check a + # worker does before writing a result. Nothing blocking runs under it + # (no network, no executor, no joins -- logging aside), so every + # critical section is short and the drain's budgeted waits stay + # meaningful. + self.state_lock = threading.Lock() # manage containers self.MAX_CONTAINER_SIZE = d_config.get('MAX_CONTAINER_NUMBER', 8) self.container_count_lock = threading.Lock() @@ -56,8 +100,6 @@ def __init__( s_config = json.load(f) self.submission_executor_cwd = pathlib.Path( s_config['working_dir']) - self.timeout = 300 - self.created_at = {} def compile_need(self, lang: Language): return lang in {Language.C, Language.CPP} @@ -73,12 +115,6 @@ def dec_container(self): with self.container_count_lock: self.container_count -= 1 - def is_timed_out(self, job_id: str): - if not self.contains(job_id): - return False - delta = (datetime.now() - self.created_at[job_id]).seconds - return delta > self.timeout - def handle(self, job_id: str, submission_id: str): ''' handle a job, save its config and push into task queue @@ -86,6 +122,23 @@ def handle(self, job_id: str, submission_id: str): multiple jobs for the same submission are legal concurrency (e.g. rejudge racing an in-flight judge run); they are kept fully separate here, keyed by job_id. + + handle() is atomic: every validation that can fail (missing dir, not a + directory, meta parse) runs BEFORE any state assignment or queue + mutation, so on any raise no state and no queue entries exist. The + unbounded queue is what makes this hold -- put_nowait can no longer + raise queue.Full, so there is no partial-enqueue rollback to do. + + Raises DispatcherDrainingError once a drain has latched the + dispatcher (stop_accepting). Nothing is published in that case; the + poller returns the claim as an attempt-neutral 'drain' abort. + + Returns True when this job_id was already present, i.e. the previous + generation was superseded by this call (a re-claim of the same job + after a backend reclaim). The superseded generation can never produce + an outcome -- its queue entries and workers fail the identity checks + from now on -- so the poller uses the return value to release the + dead claim's tracker count. Returns False otherwise. ''' logger().info(f'receive job {job_id} (submission {submission_id}).') job_path = self.SUBMISSION_DIR / job_id @@ -98,18 +151,29 @@ def handle(self, job_id: str, submission_id: str): with (job_path / 'meta.json').open() as f: submission_config = Meta.parse_obj(json.load(f)) - # assign job context + # Publish the job state and its tasks under state_lock, so a drain + # scan holding that lock sees either all of this job or none of it. task_content = {} - self.result[job_id] = (submission_config, task_content) - self.locks[job_id] = threading.Lock() - self.compile_locks[job_id] = threading.Lock() - self.created_at[job_id] = datetime.now() - self.submission_ids[job_id] = submission_id + with self.state_lock: + # The latch check shares the publication's critical section, and + # that is what makes the drain's SINGLE scan complete: every job + # that ever published did so before the latch (so the scan sees + # it), and everything after raises here without publishing. A + # dispatch can therefore never land in a dispatcher the drain has + # already scanned, no matter how late the poller runs. + if not self.accepting: + raise DispatcherDrainingError( + f'draining; refusing job {job_id}') + superseded = job_id in self.result + if superseded: + logger().warning(f're-handling {job_id}; superseding the ' + 'previous generation') + ctx = JobContext(submission_config, submission_id, task_content) + self.result[job_id] = ctx - logger().debug(f'current jobs: {[*self.result.keys()]}') - try: + logger().debug(f'current jobs: {[*self.result.keys()]}') if self.compile_need(submission_config.language): - self.queue.put_nowait(job.Compile(job_id=job_id)) + self.queue.put_nowait(job.Compile(job_id=job_id, ctx=ctx)) for i, task in enumerate(submission_config.tasks): for j in range(task.caseCount): case_no = f'{i:02d}{j:02d}' @@ -118,26 +182,62 @@ def handle(self, job_id: str, submission_id: str): job_id=job_id, task_id=i, case_id=j, + ctx=ctx, ) self.queue.put_nowait(_job) - except queue.Full as e: - self.release(job_id) - raise e + return superseded def release(self, job_id: str): + with self.state_lock: + self._release_locked(job_id) + + def _release_locked(self, job_id: str): ''' - Release variable about job + Retire the job's current generation. Caller must hold state_lock. ''' - for v in ( - self.result, - self.compile_locks, - self.compile_results, - self.locks, - self.created_at, - self.submission_ids, - ): - if job_id in v: - del v[job_id] + self.result.pop(job_id, None) + + def stop_accepting(self): + '''Refuse new jobs and start no more queued ones; in-flight ones keep + running. + + This is the first thing a drain does -- before even stopping the + poller. The poller can complete a prep and dispatch at any point on + its way out, and until this latch is set the run loop is free to start + that job. It would then be judging work that began after shutdown, get + abandoned when the process exits, and come back through reclaim with + its claim attempt already spent, instead of being requeued + attempt-neutrally by drain_unstarted(). Latching first also makes the + window harmless rather than merely small: a dispatch that arrives + after the latch is refused by handle() outright, and one that got in + just before cannot start and is returned by the scan. + ''' + with self.state_lock: + self.accepting = False + + def drain_unstarted(self): + '''Release every job with no started work, and return them. + + Returns [(job_id, submission_id)] for the released jobs. Call after + stop_accepting(): handle() refuses new jobs from then on, so a single + scan sees every job that ever published, regardless of whether the + poller is still running. In-flight jobs are left to finish and report + normally. + + Latches on its own too, so a caller that skips stop_accepting() still + gets a consistent scan -- just a wider window before it. + ''' + with self.state_lock: + self.accepting = False + victims = [(job_id, ctx.submission_id) + for job_id, ctx in list(self.result.items()) + if not ctx.started] + for job_id, _ in victims: + self._release_locked(job_id) + return victims + + def has_jobs(self): + return bool(self.result) def run(self): self.do_run = True @@ -158,54 +258,79 @@ def run(self): # get a case _job = self.queue.get() job_id = _job.job_id - # if a job was discarded, it will not appear in the `self.result` - if not self.contains(job_id): - logger().info(f'discarded job [id={job_id}]') - continue - if self.is_timed_out(job_id): - logger().info(f'job timed out [id={job_id}]') - continue - # get task info - submission_config, _ = self.result[job_id] - if isinstance(_job, job.Compile): - threading.Thread( - target=self.compile, - args=( - job_id, - submission_config.language, - ), - ).start() - # if this job needs compile and it haven't finished - elif self.compile_need(submission_config.language) \ - and self.compile_results.get(job_id) is None: - self.queue.put(_job) - else: - task_info = submission_config.tasks[_job.task_id] - case_no = f'{_job.task_id:02d}{_job.case_id:02d}' - logger().info(f'create container [task={job_id}/{case_no}]') - logger().debug(f'task info: {task_info}') - # output path should be the container path - base_path = self.SUBMISSION_DIR / job_id / 'testcase' - out_path = str((base_path / f'{case_no}.out').absolute()) - # input path should be the host path - base_path = self.submission_executor_cwd / job_id / 'testcase' - in_path = str((base_path / f'{case_no}.in').absolute()) - # debug log - logger().debug('in path: ' + in_path) - logger().debug('out path: ' + out_path) - # assign a new executor - threading.Thread( - target=self.create_container, - args=( - job_id, - case_no, - task_info.memoryLimit, - task_info.timeLimit, - in_path, - out_path, - submission_config.language, - ), - ).start() + # Read job context and transition to started atomically w.r.t. + # drain_unstarted(): once we mark a job started here, drain leaves + # it to finish; before that, drain may abort it out from under us. + with self.state_lock: + # A queue entry belongs to the generation that enqueued it + # (handle() binds the JobContext in). An identity mismatch + # means the job was released -- or re-handled under the same + # job_id -- since the entry was queued: letting it through + # would bind the old generation's work to the NEW attempt, + # double-running cases or crashing on a task shape the new + # meta no longer has. Fetching state by job_id here would be + # exactly that bug, so the entry's own token is all we use. + ctx = _job.ctx + if self.result.get(job_id) is not ctx: + logger().info(f'discarded stale task [id={job_id}]') + continue + # Draining: only generations already in flight may keep + # consuming their remaining cases. Anything still unstarted + # belongs to drain_unstarted(), which returns it. + if not self.accepting and not ctx.started: + logger().info(f'draining, not starting [id={job_id}]') + continue + submission_config = ctx.config + if isinstance(_job, job.Compile): + ctx.started = True + threading.Thread( + target=self.compile, + args=( + job_id, + submission_config.language, + ), + kwargs={ + 'ctx': ctx + }, + ).start() + # if this job needs compile and it haven't finished + elif self.compile_need(submission_config.language) \ + and ctx.compile_result is None: + # unbounded queue: never blocks, safe under the lock + self.queue.put(_job) + else: + ctx.started = True + task_info = submission_config.tasks[_job.task_id] + case_no = f'{_job.task_id:02d}{_job.case_id:02d}' + logger().info( + f'create container [task={job_id}/{case_no}]') + logger().debug(f'task info: {task_info}') + # output path should be the container path + base_path = self.SUBMISSION_DIR / job_id / 'testcase' + out_path = str((base_path / f'{case_no}.out').absolute()) + # input path should be the host path + base_path = (self.submission_executor_cwd / job_id / + 'testcase') + in_path = str((base_path / f'{case_no}.in').absolute()) + # debug log + logger().debug('in path: ' + in_path) + logger().debug('out path: ' + out_path) + # assign a new executor + threading.Thread( + target=self.create_container, + args=( + job_id, + case_no, + task_info.memoryLimit, + task_info.timeLimit, + in_path, + out_path, + submission_config.language, + ), + kwargs={ + 'ctx': ctx + }, + ).start() def stop(self): self.do_run = False @@ -214,20 +339,24 @@ def compile( self, job_id: str, lang: Language, + *, + ctx, ): - # another thread is compiling this job, bye - if self.compile_locks[job_id].locked(): - logger().error(f'start a compile thread on locked job {job_id}') - return + '''Compile a job; ``ctx`` is the generation token run() spawned us + with, presented back before the result is written.''' # this job should not be compiled! if not self.compile_need(lang): logger().warning( f'try to compile job {job_id}' f' with language {lang}', ) return - # compile this job. don't forget to acquire the lock - with self.compile_locks[job_id]: - logger().info(f'start compiling {job_id}') + logger().info(f'start compiling {job_id}') + # A JE fallback keeps the drain invariant "in-flight jobs always + # finish": an unhandled error here (e.g. docker APIError from + # create_container) would otherwise never write a compile result, + # so the job would stay leased forever. JudgeError is already + # wrapped inside SubmissionExecutor; this catches what it does not. + try: res = SubmissionExecutor( job_id=job_id, time_limit=-1, @@ -236,8 +365,23 @@ def compile( testdata_output_path='', lang=['c11', 'cpp17'][int(lang)], ).compile() - self.compile_results[job_id] = res - logger().debug(f'finish compiling, get status {res["Status"]}') + except Exception: + logger().exception(f'compile crashed for {job_id}') + res = {'Status': 'JE'} + # Generation check and write share one critical section: a job + # released while we compiled (drain), or re-handled under the same + # job_id (backend reclaim), owns a different JobContext, and this + # stale result must not leak into it. The result lives ON the + # context, so even the dropped write could never reach another + # generation -- the check is what keeps the log honest and the + # generation's lifecycle explicit. + with self.state_lock: + if self.result.get(job_id) is not ctx: + logger().warning( + f'stale compile result for {job_id}; dropping') + return + ctx.compile_result = res + logger().debug(f'finish compiling, get status {res["Status"]}') def create_container( self, @@ -248,28 +392,54 @@ def create_container( case_in_path: str, case_out_path: str, lang: Language, + *, + ctx, ): - lang = ['c11', 'cpp17', 'python3'][int(lang)] - executor = SubmissionExecutor( - job_id, - time_limit, - mem_limit, - case_in_path, - case_out_path, - lang=lang, - case_no=case_no, - ) - res = self.extract_compile_result(job_id, lang) - # Execute if compile successfully - if res['Status'] != 'CE': - try: - self.inc_container() - res = executor.run() - finally: - self.dec_container() + # JE fallback: an unhandled error anywhere before the case result + # exists (executor construction reading its config file, docker + # APIError from run()) must still resolve the case, or the job would + # never complete and stay leased forever. + try: + lang_name = ['c11', 'cpp17', 'python3'][int(lang)] + executor = SubmissionExecutor( + job_id, + time_limit, + mem_limit, + case_in_path, + case_out_path, + lang=lang_name, + case_no=case_no, + ) + # Pass the Language enum, not lang_name: compile_need() compares + # against the enum, so the defensive CE fallback inside would + # otherwise be unreachable for C/C++. + res = self.extract_compile_result(ctx, lang) + # Execute if compile successfully + if res['Status'] != 'CE': + try: + self.inc_container() + res = executor.run() + finally: + self.dec_container() + except Exception: + logger().exception(f'judge crashed for {job_id}/{case_no}') + res = {'Status': 'JE'} logger().info(f'finish task {job_id}/{case_no}') - with self.locks[job_id]: - self.on_case_complete( + # Generation check and write share one critical section. The job may + # have been released while this case ran (a drain), or re-handled + # under the SAME job_id (an abort/reclaim requeues it and this runner + # can claim it again). Looking state up by job_id alone would then + # write this stale result into the new attempt -- completing it early + # with another run's data -- and a bare lookup after a release would + # raise outside the JE guard and kill this thread. The identity check + # on the exact JobContext we were spawned with rejects both, + # atomically with the write. + with self.state_lock: + if self.result.get(job_id) is not ctx: + logger().warning( + f'stale case result for {job_id}/{case_no}; dropping') + return + self._on_case_complete_locked( job_id=job_id, case_no=case_no, stdout=res.get('Stdout', ''), @@ -280,16 +450,15 @@ def create_container( prob_status=res['Status'], ) - def extract_compile_result(self, job_id: str, lang: Language): + def extract_compile_result(self, ctx, lang: Language): ''' - Get compile result for specific job. If the language does - not need to be compiled, return a AC result. + Get the compile result of a job's generation. If the language does + not need to be compiled, return an AC result. ''' - try: - return self.compile_results[job_id] - except KeyError: - status = 'CE' if self.compile_need(lang) else 'AC' - return {'Status': status} + if ctx.compile_result is not None: + return ctx.compile_result + status = 'CE' if self.compile_need(lang) else 'AC' + return {'Status': status} def on_case_complete( self, @@ -302,12 +471,39 @@ def on_case_complete( mem_usage: int, prob_status: str, ): + '''Test-facing wrapper. Generation-UNSAFE: it resolves the job by + job_id, so production code must go through the identity-checked path + in create_container() instead of calling this.''' + with self.state_lock: + self._on_case_complete_locked( + job_id=job_id, + case_no=case_no, + stdout=stdout, + stderr=stderr, + exit_code=exit_code, + exec_time=exec_time, + mem_usage=mem_usage, + prob_status=prob_status, + ) + + def _on_case_complete_locked( + self, + job_id: str, + case_no: str, + stdout: str, + stderr: str, + exit_code: int, + exec_time: int, + mem_usage: int, + prob_status: str, + ): + '''Record one case result. Caller must hold state_lock.''' # if id not exists if job_id not in self.result: raise JobIdNotFoundError( f'received case result for unknown job {job_id}') # update case result - _, results = self.result[job_id] + results = self.result[job_id].cases if case_no not in results: raise ValueError(f'{job_id}/{case_no} not found.') results[case_no] = { @@ -322,16 +518,23 @@ def on_case_complete( _results = [k for k, v in results.items() if not v] logger().debug(f'tasks wait for judge: {_results}') if all(results.values()): - self.on_job_complete(job_id) + self._on_job_complete_locked(job_id) def on_job_complete(self, job_id: str): + '''Test-facing wrapper. Generation-UNSAFE, like on_case_complete().''' + with self.state_lock: + self._on_job_complete_locked(job_id) + + def _on_job_complete_locked(self, job_id: str): + '''Finalize a finished job. Caller must hold state_lock.''' if not self.contains(job_id): raise JobIdNotFoundError(f'{job_id} not found!') if self.testing: logger().info( f'skip job post processing in testing [job_id={job_id}]') return True - _, results = self.result[job_id] + ctx = self.result[job_id] + results = ctx.cases # parse results submission_result = {} for no, r in results.items(): @@ -346,33 +549,18 @@ def on_job_complete(self, job_id: str): submission_result[task_no] = [*cases.values()] assert [*submission_result.keys()] == [*range(len(submission_result))] submission_result = [*submission_result.values()] - # post data - submission_id = self.submission_ids[job_id] - with tempfile.NamedTemporaryFile("w") as tmpf: - submission_data = { - 'tasks': submission_result, - 'token': config.SANDBOX_TOKEN - } - # write payload to file - json.dump(submission_data, tmpf) - tmpf.flush() - # release resources - del submission_data - self.release(job_id) - - logger().info( - f'send to BE [job_id={job_id}, submission_id={submission_id}]') - # open in binary mode as requests needs a binary stream - with open(tmpf.name, "rb") as payload: - resp = requests.put( - f'{config.BACKEND_API}/submission/{submission_id}/complete', - data=payload, - headers={'Content-Type': 'application/json'}, - ) - logger().debug(f'get BE response: [{resp.status_code}] {resp.text}', ) - # clear - if resp.ok: - file_manager.clean_data(job_id) - # copy to another place + # Hand the result off to the result queue via on_complete. Local dir + # cleanup/backup and tracker bookkeeping are now entirely the + # ResultSenderThread's job; release() only drops dispatcher state. + submission_id = ctx.submission_id + # Enqueue BEFORE release: has_jobs() must stay True until the result + # is in the queue, or a concurrent drain could observe an empty + # dispatcher, stop the sender, and lose this result. + if self.on_complete is None: + logger().error( + f'no on_complete wired; dropping result for job {job_id}') else: - file_manager.backup_data(job_id) + logger().info(f'enqueue result [job_id={job_id},' + f' submission_id={submission_id}]') + self.on_complete(job_id, submission_id, submission_result) + self._release_locked(job_id) diff --git a/dispatcher/exception.py b/dispatcher/exception.py index 1091044..e1baca1 100644 --- a/dispatcher/exception.py +++ b/dispatcher/exception.py @@ -2,3 +2,10 @@ class JobIdNotFoundError(Exception): ''' raise this error when job id not found ''' + + +class DispatcherDrainingError(Exception): + ''' + raised by handle() when the dispatcher has been latched by a drain + (stop_accepting) and refuses to take on new jobs + ''' diff --git a/dispatcher/job.py b/dispatcher/job.py index d253c6f..8f5771d 100644 --- a/dispatcher/job.py +++ b/dispatcher/job.py @@ -4,6 +4,11 @@ @dataclass class Compile: job_id: str + # The JobContext generation this entry was enqueued under. The run loop + # discards the entry when it is no longer the current generation, so a + # task queued by a released or superseded handle() can never bind itself + # to the attempt that replaced it. + ctx: object @dataclass @@ -11,3 +16,5 @@ class Execute: job_id: str task_id: int case_id: int + # See Compile.ctx. + ctx: object diff --git a/dispatcher/testdata.py b/dispatcher/testdata.py index 4b748ac..abff995 100644 --- a/dispatcher/testdata.py +++ b/dispatcher/testdata.py @@ -14,7 +14,7 @@ logger, ) from .config import ( - BACKEND_API, + BACKEND_URL, SANDBOX_TOKEN, TESTDATA_ROOT, ) @@ -22,6 +22,17 @@ META_DIR = TESTDATA_ROOT / 'meta' META_DIR.mkdir(exist_ok=True) +# Unbounded network I/O here would stall a SIGTERM drain until docker's stop +# grace period SIGKILLs us (losing in-flight results), because prep runs inside +# the poller iteration that stop() has to interrupt. Every request and every +# lock wait must be bounded so the poller's stop() takes effect within one +# prep attempt. +# redis-py raises LockError when the lock cannot be acquired within +# LOCK_BLOCKING_TIMEOUT, which the poller's prep retry already treats as a +# prep failure. +HTTP_TIMEOUT = (5, 30) # (connect, read) seconds +LOCK_BLOCKING_TIMEOUT = 30 # seconds + def calc_checksum(data: bytes) -> str: return hashlib.md5(data).hexdigest() @@ -41,10 +52,11 @@ def handle_problem_response(resp: rq.Response): def fetch_problem_meta(problem_id: int) -> str: logger().debug(f'fetch problem meta [problem_id={problem_id}]') resp = rq.get( - f'{BACKEND_API}/problem/{problem_id}/meta', + f'{BACKEND_URL}/problem/{problem_id}/meta', params={ 'token': SANDBOX_TOKEN, }, + timeout=HTTP_TIMEOUT, ) handle_problem_response(resp) content = json.dumps(resp.json()['data']) @@ -71,10 +83,11 @@ def fetch_testdata(problem_id: int): ''' logger().debug(f'fetch problem testdata [problem_id={problem_id}]') resp = rq.get( - f'{BACKEND_API}/problem/{problem_id}/testdata', + f'{BACKEND_URL}/problem/{problem_id}/testdata', params={ 'token': SANDBOX_TOKEN, }, + timeout=HTTP_TIMEOUT, ) handle_problem_response(resp) return resp.content @@ -82,10 +95,11 @@ def fetch_testdata(problem_id: int): def get_checksum(problem_id: int) -> str: resp = rq.get( - f'{BACKEND_API}/problem/{problem_id}/checksum', + f'{BACKEND_URL}/problem/{problem_id}/checksum', params={ 'token': SANDBOX_TOKEN, }, + timeout=HTTP_TIMEOUT, ) handle_problem_response(resp) return resp.json()['data'] @@ -98,7 +112,9 @@ def ensure_testdata(problem_id: int): client = get_redis_client() key = f'problem-{problem_id}-checksum' lock_key = f'{key}-lock' - with client.lock(lock_key, timeout=60): + with client.lock(lock_key, + timeout=60, + blocking_timeout=LOCK_BLOCKING_TIMEOUT): curr_checksum = client.get(key) if curr_checksum is not None: curr_checksum = curr_checksum.decode() diff --git a/dispatcher/utils.py b/dispatcher/utils.py index dea4bbc..8c4202e 100644 --- a/dispatcher/utils.py +++ b/dispatcher/utils.py @@ -1,14 +1,10 @@ import logging import os import redis -from flask import current_app def logger() -> logging.Logger: - try: - return current_app.logger - except RuntimeError: - return logging.getLogger('gunicorn.error') + return logging.getLogger('sandbox') # Fake redis server @@ -33,5 +29,12 @@ def get_redis_client(): # Create connection pool global redis_pool if redis_pool is None: - redis_pool = redis.ConnectionPool.from_url(REDIS_URL) + # Bounded socket I/O for the same reason as the testdata HTTP + # timeouts: a wedged redis must not be able to hang a SIGTERM drain + # until docker's stop grace period kills us. + redis_pool = redis.ConnectionPool.from_url( + REDIS_URL, + socket_timeout=10, + socket_connect_timeout=5, + ) return redis.Redis(connection_pool=redis_pool) diff --git a/gunicorn.conf.py b/gunicorn.conf.py deleted file mode 100644 index a18dd9d..0000000 --- a/gunicorn.conf.py +++ /dev/null @@ -1,10 +0,0 @@ -port = 1450 -bind = f'0.0.0.0:{port}' -timeout = 60 - -# loglevel = 'debug' -accesslog = 'logs/access.log' -errorlog = 'logs/error.log' - -worker_class = 'gthread' -threads = 5 diff --git a/main.py b/main.py new file mode 100644 index 0000000..cb674ea --- /dev/null +++ b/main.py @@ -0,0 +1,291 @@ +"""Pull-based runner entrypoint (spec §10). + +Wires the coordination threads together: the dispatcher runs judge work, the +poller claims jobs, the heartbeat renews leases, and the result sender reports +outcomes. There is no Flask and no inbound HTTP server anymore -- the runner +only talks *out* to the backend runner API. + +Import-safe: nothing here starts a thread, talks to the network, or configures +logging at import time, so tests can import main and exercise the helpers +directly. (Importing dispatcher.config on the way in does create the working +directories it owns -- that predates this module.) +""" + +import logging +import os +import queue +import signal +import sys +import threading +import time + +from runner.client import BackendClient, BackendAuthError +from runner.active_jobs import ActiveJobTracker +from runner.config import DRAIN_TIMEOUT_SEC +from runner.heartbeat import HeartbeatThread +from runner.poller import PollerThread +from runner.registration import register_with_backoff +from runner.result_sender import ( + ResultSenderThread, + CompleteRequest, + AbortRequest, +) + +logger = logging.getLogger('sandbox') + +# Fractions of the drain budget (see drain()). Waiting for the poller is the +# least valuable thing the drain does -- an unstarted claim costs nothing to +# hand back -- and the likeliest to get stuck, since a slow-drip response can +# keep one prep attempt alive past any per-request timeout. Reserving a slice +# for the reporting tail keeps a long-running job from leaving finished +# results with no time left to send them. +POLLER_JOIN_SHARE = 0.25 +REPORTING_RESERVE_SHARE = 0.25 + + +def setup_logging(): + """Configure root logging for the runner process. + + Logs go to both a file (persisted via the bind-mounted logs/ dir) and + stdout so ``docker logs`` shows them. Level comes from LOG_LEVEL. + """ + os.makedirs('logs', exist_ok=True) + logging.basicConfig( + level=os.getenv('LOG_LEVEL', 'INFO'), + format='%(asctime)s %(levelname)s %(name)s: %(message)s', + handlers=[ + logging.FileHandler('logs/sandbox.log'), + logging.StreamHandler(), + ], + ) + + +def register(client, shutdown): + """Register with the backend, honouring shutdown during retry backoff. + + A docker stop mid-registration must not hang, so the retry sleep waits on + ``shutdown`` and raises SystemExit the moment it is set. A rejected token + (401) is fatal per ADR-0004: exit 1 so the restart policy re-registers + with a fresh identity. Returns a RunnerIdentity, or None on 401. + """ + + def interruptible_sleep(seconds): + if shutdown.wait(seconds): + raise SystemExit(0) + + from runner.config import RUNNER_NAME + try: + return register_with_backoff( + client, + _registration_token(), + RUNNER_NAME, + sleep=interruptible_sleep, + ) + except BackendAuthError: + logger.error( + 'registration token rejected (401); exiting for restart policy ' + 'to re-register') + return None + + +def _registration_token(): + # Read at call time (not import) so tests can monkeypatch runner.config. + from runner import config + return config.RUNNER_REGISTRATION_TOKEN + + +def install_signal_handlers(shutdown): + """Route SIGTERM (docker stop) and SIGINT into the graceful drain. + + The handler only sets ``shutdown``; the drain itself runs on the main + thread once ``shutdown.wait()`` returns, so no drain work happens in + signal context. + """ + + def _on_signal(signum, _frame): + logger.info('received signal %s; shutting down', signum) + shutdown.set() + + signal.signal(signal.SIGTERM, _on_signal) + signal.signal(signal.SIGINT, _on_signal) + + +def drain(poller, + dispatcher, + sender, + heartbeat, + result_queue, + tracker, + sleep=time.sleep, + timeout_sec=DRAIN_TIMEOUT_SEC, + monotonic=time.monotonic): + """SIGTERM graceful drain (spec §10). Returns 0. + + The ordering here is the whole point of the slice: + 1. Latch the dispatcher FIRST, before touching anything else. The + invariant is "once a drain starts, no job begins judging", so the latch + has to precede every other step -- including stopping the poller. The + latch is also what makes the single scan below COMPLETE: handle() + checks it inside the same critical section that publishes a job, so + every job that ever published is visible to the scan, and any dispatch + after the latch is refused and returned by the poller itself as a + 'drain' abort. Quiescence comes from the latch, not from guessing + whether the poller might still be running -- there is deliberately no + rescan and no is_alive() heuristic here. + 2. Stop and JOIN the poller, capped at its share of the budget. The join + is an optimization, not a correctness step: a poller that exits now + makes the reporting tail below trivial, but a stuck one can no longer + publish anything the scan would miss. + 3. Scan, and abort everything unstarted with reason 'drain' (the backend + does not count those against a submission's attempts). + 4. Wait for in-flight (already-started) jobs to finish; they push + CompleteRequests. The dispatcher has to stay running for this: a job + whose first case is in flight still needs its remaining cases dispatched + or it never completes at all. This wait keeps the reporting reserve. + 5. Stop the dispatcher. Its join is bounded away from the reserve too: + the reserve exists for the sender, so no earlier step may consume it. + 6. Reporting tail: keep the sender alive until the poller thread has + exited AND every claimed job's outcome has been finalized (the tracker + is empty), or the budget runs out. A poller stuck in prep I/O can + still wake up and queue a late 'drain' abort; stopping the sender any + earlier would throw that report away, so a stuck poller holds the + sender open to the deadline. In-flight jobs the wait above gave up on + also stay in the tracker -- if one finishes during this tail, its + result still goes out instead of being lost to lease expiry. + 7. Stop the sender -- it exits only once the queue is empty, so every + queued complete and drain abort is actually reported first. (The + tracker can go empty a moment before the last abort's HTTP send + finishes -- the sender finalizes local state first -- and the join + here covers that last send within what is left of the budget.) + 8. Stop the heartbeat LAST: complete/abort need a live lease until sent, + so the heartbeat must outlive the sender. + + Every wait shares one budget (``timeout_sec``). Draining is best-effort: + prep and result I/O can be bounded per request but never in total, so + waiting forever would just hand the decision to docker's SIGKILL and throw + away the results already collected. Anything that misses the budget is left + to lease expiry and backend reclaim. + + The budget is not first-come-first-served, because the steps are not + equally worth waiting for. Reporting a finished judge run matters more than + waiting out a poller stuck in prep, so the poller gets a capped share and + the reporting tail keeps a reserve. + """ + deadline = monotonic() + timeout_sec + reserve = timeout_sec * REPORTING_RESERVE_SHARE + + def remaining(reserve=0.0): + return max(0.0, deadline - monotonic() - reserve) + + def join_within(component, name, limit=None, reserve=0.0): + budget = remaining(reserve) + if limit is not None: + budget = min(budget, limit) + component.join(timeout=budget) + if component.is_alive(): + logger.error( + 'drain: %s outlived the budget; its work is left to lease ' + 'expiry and backend reclaim', name) + + # Before anything else, including stopping the poller: see step 1. + dispatcher.stop_accepting() + poller.stop() + join_within(poller, 'poller', limit=timeout_sec * POLLER_JOIN_SHARE) + + # The single scan (step 3): complete by construction, see the docstring. + for job_id, submission_id in dispatcher.drain_unstarted(): + result_queue.put(AbortRequest(job_id, submission_id, 'drain')) + + waited = 0 + while dispatcher.has_jobs(): + if not remaining(reserve=reserve): + logger.error('drain: budget exhausted with jobs still in flight') + break + sleep(0.5) + waited += 1 + if waited % 20 == 0: + logger.info('drain: waiting for in-flight jobs to finish') + + dispatcher.stop() + join_within(dispatcher, 'dispatcher', reserve=reserve) + + # The reporting tail (step 6). + waited = 0 + while remaining() and (poller.is_alive() or len(tracker)): + sleep(0.5) + waited += 1 + if waited % 20 == 0: + logger.info( + 'drain: reporting tail waiting (poller alive: %s, ' + 'unreported jobs: %d)', poller.is_alive(), len(tracker)) + + sender.stop() + join_within(sender, 'result sender') + heartbeat.stop() + join_within(heartbeat, 'heartbeat') + return 0 + + +def main(): + from dispatcher.dispatcher import Dispatcher + from runner.config import BACKEND_URL + + shutdown = threading.Event() + fatal = threading.Event() + + install_signal_handlers(shutdown) + + if not _registration_token(): + logger.error( + 'RUNNER_REGISTRATION_TOKEN is empty; cannot register. Set it in ' + 'the runner environment (ADR-0005).') + return 1 + + client = BackendClient(BACKEND_URL) + identity = register(client, shutdown) + if identity is None: + return 1 + + result_queue = queue.Queue() + tracker = ActiveJobTracker() + + dispatcher = Dispatcher( + os.getenv('DISPATCHER_CONFIG', '.config/dispatcher.json.example'), + on_complete=lambda job_id, submission_id, tasks: result_queue.put( + CompleteRequest(job_id, submission_id, tasks)), + ) + sender = ResultSenderThread(client, identity, tracker, result_queue) + poller = PollerThread( + client, + identity, + tracker, + result_queue, + dispatch=dispatcher.handle, + ) + heartbeat = HeartbeatThread( + client, + identity, + tracker, + on_fatal=lambda: (fatal.set(), shutdown.set()), + ) + + dispatcher.start() + sender.start() + poller.start() + heartbeat.start() + logger.info('runner %s up', identity.runner_id) + + shutdown.wait() + + if fatal.is_set(): + # 401 fail-fast: no drain. Our identity has evaporated, so complete / + # abort would 401 anyway; compose restart re-registers a fresh one. + logger.error('heartbeat fail-fast; exiting without drain') + return 1 + + return drain(poller, dispatcher, sender, heartbeat, result_queue, tracker) + + +if __name__ == '__main__': + setup_logging() + sys.exit(main()) diff --git a/requirements.txt b/requirements.txt index 42fd0ab..bd96e3c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,5 @@ docker==7.1.0 requests~=2.27 -gunicorn~=20.1 -flask~=2.0 yapf~=0.32 pydantic~=1.9 redis~=4.1.4 diff --git a/runner/__init__.py b/runner/__init__.py index c324208..a66e65a 100644 --- a/runner/__init__.py +++ b/runner/__init__.py @@ -2,6 +2,6 @@ Registration, heartbeat, HTTP client, active-job tracking, job polling and result sending used by the runner process to talk to the backend runner API. -Job prep reuses the dispatcher's file/testdata helpers; still no docker and -no Flask in this package. +Job prep reuses the dispatcher's file/testdata helpers; no docker and no +Flask live in this package. """ diff --git a/runner/active_jobs.py b/runner/active_jobs.py index eea62dd..17adf0b 100644 --- a/runner/active_jobs.py +++ b/runner/active_jobs.py @@ -2,30 +2,45 @@ class ActiveJobTracker: - """Thread-safe set of the job ids this runner currently holds. - - The snapshot is sent with every heartbeat to renew the backend leases - (spec §7.2); jobs absent from the snapshot stop being renewed and become - orphans once their lease expires. + """Thread-safe multiset of the claims this runner currently holds. + + Keyed by job_id but counted per CLAIM: the backend can reclaim a job + whose lease expired and requeue it, and this same runner can claim it + again while the first claim's outcome is still being finalized. A plain + set would collapse the two -- the first outcome's remove would also drop + the still-active claim, so the heartbeat would stop renewing its lease, + the capacity gate would free a slot that is still busy, and the drain's + reporting tail would believe everything was already reported. + + snapshot() deduplicates to job_ids (the heartbeat renews one lease per + job); len() counts claims (capacity and the drain's reporting tail need + the un-collapsed number). """ def __init__(self): self._lock = threading.Lock() - self._jobs = set() + # job_id -> number of outstanding claims + self._claims = {} def add(self, job_id): with self._lock: - self._jobs.add(job_id) + self._claims[job_id] = self._claims.get(job_id, 0) + 1 def remove(self, job_id): - # Idempotent: removing an absent id is a no-op. + # Releases ONE claim; idempotent once none are left. Every claim is + # released exactly once: by the sender when its outcome is finalized, + # or by the poller when a re-claim supersedes it (see _poll_once). with self._lock: - self._jobs.discard(job_id) + count = self._claims.get(job_id, 0) + if count <= 1: + self._claims.pop(job_id, None) + else: + self._claims[job_id] = count - 1 def snapshot(self): with self._lock: - return list(self._jobs) + return list(self._claims) def __len__(self): with self._lock: - return len(self._jobs) + return sum(self._claims.values()) diff --git a/runner/config.py b/runner/config.py index 273be9b..b6785c8 100644 --- a/runner/config.py +++ b/runner/config.py @@ -1,13 +1,10 @@ import os import socket -# Backend base URL for the pull-based runner API. Coexists with the legacy -# BACKEND_API (dispatcher/config.py) until the keystone slice removes the -# old push path. -BACKEND_URL = os.getenv( - 'BACKEND_URL', - 'http://web:8080', -) +# Single source of truth for the backend base URL, shared with the legacy +# testdata channel. runner/ already imports from dispatcher/ so this reuse +# introduces no layering issue or import cycle. +from dispatcher.config import BACKEND_URL # Shared secret presented to POST /runners/register. RUNNER_REGISTRATION_TOKEN = os.getenv( 'RUNNER_REGISTRATION_TOKEN', @@ -34,3 +31,8 @@ PREP_BACKOFF_SCHEDULE = (1, 2) # Backoff (seconds) between complete/abort resends; len == max retries (spec §7). SEND_RETRY_BACKOFF_SCHEDULE = (1, 2, 4, 8, 16) +# Total budget (seconds) for the SIGTERM drain. Prep and result I/O can only +# be bounded per request, never in total, so the drain caps itself instead of +# letting docker's SIGKILL decide -- leaving headroom under the +# `docker stop --time=600` the deployment uses (spec §10). +DRAIN_TIMEOUT_SEC = 540 diff --git a/runner/poller.py b/runner/poller.py index 9171996..dbb858e 100644 --- a/runner/poller.py +++ b/runner/poller.py @@ -2,12 +2,12 @@ import logging import shutil import threading -import time import requests from dispatcher import file_manager, testdata from dispatcher.config import SUBMISSION_DIR +from dispatcher.exception import DispatcherDrainingError from dispatcher.meta import Meta from .client import BackendAPIError from .config import REQUEST_TIMEOUT, PREP_MAX_ATTEMPTS, PREP_BACKOFF_SCHEDULE @@ -46,6 +46,18 @@ class PollerThread(threading.Thread): Only polls when there is spare capacity, adds the claimed job to the tracker before prep so the heartbeat renews the lease while downloading, and retries prep locally before giving up with a prep_failed abort. + + Drain safety: a claim that is still in prep when stop() lands is returned + to the backend as a 'drain' abort (see _poll_once). A claim whose prep + happens to finish after stop() still attempts to dispatch, and the + DISPATCHER decides: handle() checks the drain latch inside the same + critical section that publishes the job, so a dispatch either lands + before the latch (and main.drain()'s scan returns it) or raises + DispatcherDrainingError without publishing anything (and _poll_once + returns the claim as a 'drain' abort). Checking the stop event here + before dispatching could not close that race -- stop() can land between + the check and the dispatch -- which is why the latch check lives with + the publication instead. """ def __init__( @@ -58,7 +70,7 @@ def __init__( *, prepare=None, poll_interval_sec=None, - sleep=time.sleep, + sleep=None, ): super().__init__(daemon=True) self._client = client @@ -70,8 +82,10 @@ def __init__( self._poll_interval_sec = (poll_interval_sec if poll_interval_sec is not None else identity.config.poll_interval_sec) - self._sleep = sleep self._stop_event = threading.Event() + # The prep backoff must be interruptible by stop(), or a drain waits + # out the full schedule before the poller can return its claim. + self._sleep = sleep if sleep is not None else self._stop_event.wait def run(self): while not self._stop_event.is_set(): @@ -89,9 +103,17 @@ def run(self): def _poll_once(self): """Claim and prep one job. Returns True when the loop should idle.""" - # Capacity gate: only GET next-job when there is room (spec §10). + # Capacity gate: only GET next-job when there is room (spec §10). This + # also bounds the task dimension: the dispatcher task queue is unbounded + # and its volume is bounded by this job-level gate × per-job case count, + # so dispatch can never fail on capacity. if len(self._tracker) >= self._identity.config.max_concurrent_jobs: return True + # Do not take on new work once draining. This narrows but cannot close + # the window -- stop() can still land while the claim HTTP call is in + # flight -- so the prep-loop check below is what guarantees correctness. + if self._stop_event.is_set(): + return True try: payload = self._client.next_job(self._identity) except (BackendAPIError, requests.RequestException) as err: @@ -106,6 +128,16 @@ def _poll_once(self): self._tracker.add(payload.job_id) for i in range(PREP_MAX_ATTEMPTS): + if self._stop_event.is_set(): + # This claim never started any work, so it has to go back as a + # 'drain' abort: drain is attempt-neutral and requeued + # immediately, whereas 'prep_failed' counts against the + # submission's attempts and would let a rolling restart push + # otherwise healthy submissions to JE (spec §7.5, §12). + logger.info('draining: returning unstarted claim %s', + payload.job_id) + self._abort(payload, 'drain') + return False try: self._prepare(payload) break @@ -116,25 +148,53 @@ def _poll_once(self): if i < len(PREP_BACKOFF_SCHEDULE): self._sleep(PREP_BACKOFF_SCHEDULE[i]) else: - logger.error('prep for %s exhausted all attempts; aborting', - payload.job_id) - self._abort(payload) + # stop() can also land during the final attempt, i.e. after the + # loop's own check has run for the last time. The claim still + # never started any work, so drain semantics win here too. + if self._stop_event.is_set(): + logger.info( + 'draining: returning unstarted claim %s after a ' + 'failed prep', payload.job_id) + self._abort(payload, 'drain') + else: + logger.error('prep for %s exhausted all attempts; aborting', + payload.job_id) + self._abort(payload, 'prep_failed') return False try: - self._dispatch(payload.job_id, payload.submission_id) + superseded = self._dispatch(payload.job_id, payload.submission_id) + except DispatcherDrainingError: + # The drain latched the dispatcher while this claim was in prep. + # Nothing was published and nothing started, so the claim goes + # back attempt-neutrally, same as a prep interrupted by stop(). + logger.info('draining: dispatcher refused %s; returning claim', + payload.job_id) + self._abort(payload, 'drain') except Exception as err: - # One shot only: a failed handle() may have partially enqueued - # task entries, and calling it again for the same job_id would - # revive them (duplicate execution). Requeue via abort instead; - # slice 4 makes handle() atomic and closes this for good. + # handle() is atomic (it validates before any state/queue + # mutation), so an exception here means the prepped job dir/meta + # is genuinely bad -- abort(prep_failed) is the right response even + # mid-drain, because the job must converge rather than be requeued + # onto another runner that would hit the same bad payload. logger.warning('dispatch for %s failed: %s', payload.job_id, err) - self._abort(payload) + self._abort(payload, 'prep_failed') + else: + if superseded: + # This claim replaced an earlier generation of the same + # job_id (the backend reclaimed its expired lease, requeued, + # and we claimed it again). The dead generation fails every + # identity check from now on, so nothing else will ever + # produce an outcome for it -- release its tracker count + # here, keeping adds and removes balanced at one per claim. + logger.info('claim of %s superseded a previous generation', + payload.job_id) + self._tracker.remove(payload.job_id) return False - def _abort(self, payload): + def _abort(self, payload, reason): self._result_queue.put( - AbortRequest(payload.job_id, payload.submission_id, 'prep_failed')) + AbortRequest(payload.job_id, payload.submission_id, reason)) # Do NOT remove from tracker; the sender does that as part of # finalizing the abort. diff --git a/tests/test_active_jobs.py b/tests/test_active_jobs.py index 6529112..7b47afa 100644 --- a/tests/test_active_jobs.py +++ b/tests/test_active_jobs.py @@ -12,13 +12,35 @@ def test_add_and_snapshot(): assert len(tracker) == 2 -def test_add_is_idempotent(): +def test_add_counts_claims_snapshot_dedupes(): + # The same job_id can be claimed twice (lease expired, backend requeued, + # this runner re-claimed). Each claim counts -- capacity and the drain's + # reporting tail need the real number -- but the heartbeat renews one + # lease per job, so the snapshot stays deduplicated. tracker = ActiveJobTracker() tracker.add('jb_1') tracker.add('jb_1') assert tracker.snapshot() == ['jb_1'] + assert len(tracker) == 2 + + +def test_interleaved_claims_of_the_same_job_id(): + # The regression the set-based tracker had: claim 1's outcome finalizing + # (remove) must not also drop claim 2 -- the heartbeat would stop + # renewing the still-active claim's lease and the drain's reporting tail + # would think everything was reported. + tracker = ActiveJobTracker() + tracker.add('jb_1') # claim 1 + tracker.add('jb_1') # claim 2, same job_id + + tracker.remove('jb_1') # claim 1's outcome finalized assert len(tracker) == 1 + assert tracker.snapshot() == ['jb_1'] # claim 2 still renewed + + tracker.remove('jb_1') # claim 2's outcome finalized + assert len(tracker) == 0 + assert tracker.snapshot() == [] def test_remove(): diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py index b41cfd5..afb5869 100644 --- a/tests/test_dispatcher.py +++ b/tests/test_dispatcher.py @@ -1,10 +1,44 @@ +import json +import threading +import time + +import pytest + +from dispatcher import job +from dispatcher.constant import Language from dispatcher.dispatcher import Dispatcher +from dispatcher.exception import DispatcherDrainingError from tests.submission_generator import SubmissionGenerator +def _write_job_dir(dispatcher, job_id, language, case_counts): + '''Lay out a minimal job dir (meta.json only) so handle() can read it.''' + job_dir = dispatcher.SUBMISSION_DIR / job_id + job_dir.mkdir(parents=True, exist_ok=True) + n = len(case_counts) + # taskScore must sum to 100 (Meta validator); spread it evenly. + scores = [100 // n] * n + scores[-1] += 100 - sum(scores) + tasks = [{ + 'taskScore': scores[i], + 'memoryLimit': 65536, + 'timeLimit': 1000, + 'caseCount': c, + } for i, c in enumerate(case_counts)] + (job_dir / 'meta.json').write_text( + json.dumps({ + 'language': int(language), + 'tasks': tasks, + })) + + def test_create_dispatcher(): docker_dispatcher = Dispatcher() assert docker_dispatcher is not None + # Daemon: fail-fast exit paths (heartbeat 401) return from main() without + # stopping the dispatcher; a non-daemon thread would keep the interpreter + # alive forever and the restart policy would never re-register. + assert docker_dispatcher.daemon def test_start_dispatcher(docker_dispatcher: Dispatcher): @@ -52,3 +86,469 @@ def test_same_submission_parallel_jobs( assert not docker_dispatcher.contains(job1) assert docker_dispatcher.contains(job2) + + +def test_handle_enqueues_all_cases_no_size_limit( + docker_dispatcher: Dispatcher): + # A C job with 20 cases (> the old QUEUE_SIZE bound of 16) must enqueue + # every case plus the compile job with no queue.Full failure. + _write_job_dir(docker_dispatcher, 'big', Language.C, [20]) + + docker_dispatcher.handle(job_id='big', submission_id='sub-big') + + queued = list(docker_dispatcher.queue.queue) + compiles = [j for j in queued if isinstance(j, job.Compile)] + executes = [j for j in queued if isinstance(j, job.Execute)] + assert len(compiles) == 1 + assert len(executes) == 20 + assert docker_dispatcher.contains('big') + + +def test_handle_missing_dir_is_atomic(docker_dispatcher: Dispatcher): + # A missing job dir raises before any state is assigned: no leftover + # entry in result and nothing enqueued. + with pytest.raises(FileNotFoundError): + docker_dispatcher.handle(job_id='ghost', submission_id='sub-ghost') + + assert not docker_dispatcher.contains('ghost') + assert docker_dispatcher.result == {} + assert docker_dispatcher.queue.empty() + + +def test_handle_refuses_jobs_once_draining(docker_dispatcher: Dispatcher): + # The latch check shares handle()'s publication critical section, which is + # what makes the drain's single scan complete: a dispatch that loses the + # race with the latch raises and publishes NOTHING, so no job can ever + # appear in the dispatcher after the scan ran. The poller turns this + # refusal into an attempt-neutral 'drain' abort. + _write_job_dir(docker_dispatcher, 'late', Language.PY, [1]) + docker_dispatcher.stop_accepting() + + with pytest.raises(DispatcherDrainingError): + docker_dispatcher.handle(job_id='late', submission_id='sub-late') + + assert not docker_dispatcher.contains('late') + assert docker_dispatcher.result == {} + assert docker_dispatcher.queue.empty() + # Nothing published means nothing for the scan to (re)turn either. + assert docker_dispatcher.drain_unstarted() == [] + + +def test_drain_unstarted_releases_only_unstarted( + docker_dispatcher: Dispatcher): + _write_job_dir(docker_dispatcher, 'j1', Language.PY, [1]) + _write_job_dir(docker_dispatcher, 'j2', Language.PY, [1]) + docker_dispatcher.handle(job_id='j1', submission_id='sub-j1') + docker_dispatcher.handle(job_id='j2', submission_id='sub-j2') + # j1 has picked up work; j2 has not. + docker_dispatcher.result['j1'].started = True + + victims = docker_dispatcher.drain_unstarted() + + assert victims == [('j2', 'sub-j2')] + assert docker_dispatcher.contains('j1') + assert not docker_dispatcher.contains('j2') + + +def test_latch_keeps_run_loop_from_starting_queued_jobs( + docker_dispatcher: Dispatcher): + # A job dispatched BEFORE the latch (so it published fine) but not yet + # started must stay unstarted once the latch is set: it belongs to the + # drain scan. Otherwise it could begin judging after shutdown, only to be + # abandoned when the process exits -- and it would be reclaimed with its + # attempts already spent, not requeued attempt-neutrally like a drain + # abort. (Jobs dispatched AFTER the latch never publish at all; that is + # test_handle_refuses_jobs_once_draining.) + _write_job_dir(docker_dispatcher, 'late', Language.PY, [1]) + docker_dispatcher.handle(job_id='late', submission_id='sub-late') + docker_dispatcher.stop_accepting() + assert docker_dispatcher.accepting is False + docker_dispatcher.start() + + # The run loop consuming the task is the signal that it has had its + # chance at this job; waiting a fixed slice instead would pass vacuously + # whenever the loop happened to still be in its idle sleep. + deadline = time.time() + 5.0 + while not docker_dispatcher.queue.empty(): + assert time.time() < deadline, 'run loop never consumed the task' + time.sleep(0.02) + assert not docker_dispatcher.result['late'].started + assert docker_dispatcher.container_count == 0 + # Still releasable as an unstarted job, which is what drain does with it. + assert docker_dispatcher.drain_unstarted() == [('late', 'sub-late')] + + +def test_latched_dispatcher_still_finishes_started_jobs( + docker_dispatcher: Dispatcher, + monkeypatch, +): + # The latch must let an in-flight job consume its REMAINING cases. If it + # blocked those too, a multi-case job would never complete: drain would + # wait out its whole budget and the results already computed would be lost. + ran = [] + first_case_running = threading.Event() + release = threading.Event() + + class StubExecutor: + + def __init__(self, job_id, *args, **kwargs): + self.case_no = kwargs.get('case_no') + + def run(self): + ran.append(self.case_no) + if len(ran) == 1: + first_case_running.set() + release.wait(5.0) + return {'Status': 'AC'} + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + StubExecutor) + # One container at a time, so the second case can only run after the first + # finishes -- i.e. after the latch is already in place. + docker_dispatcher.MAX_CONTAINER_SIZE = 1 + _write_job_dir(docker_dispatcher, 'multi', Language.PY, [2]) + docker_dispatcher.handle(job_id='multi', submission_id='sub-multi') + docker_dispatcher.start() + + assert first_case_running.wait(5.0), 'first case never started' + assert docker_dispatcher.drain_unstarted() == [] # in flight, so kept + release.set() + + deadline = time.time() + 5.0 + while len(ran) < 2: + assert time.time() < deadline, 'second case never ran after the latch' + time.sleep(0.02) + assert sorted(ran) == ['0000', '0001'] + + +def test_case_result_for_released_job_is_dropped_not_fatal( + docker_dispatcher: Dispatcher, + monkeypatch, +): + # A case can finish after its job was released (a drain). The worker holds + # the generation token it was spawned with; presenting it after the + # release must drop the result quietly -- raising here would kill the + # worker thread without resolving anything. + class StubExecutor: + + def __init__(self, job_id, *args, **kwargs): + pass + + def run(self): + return {'Status': 'AC'} + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + StubExecutor) + _write_job_dir(docker_dispatcher, 'gone', Language.PY, [1]) + docker_dispatcher.handle(job_id='gone', submission_id='sub-gone') + ctx = docker_dispatcher.result['gone'] + docker_dispatcher.release('gone') + + docker_dispatcher.create_container( + job_id='gone', + case_no='0000', + mem_limit=65536, + time_limit=1000, + case_in_path='/nonexistent/0000.in', + case_out_path='/nonexistent/0000.out', + lang=Language.PY, + ctx=ctx, + ) + + assert not docker_dispatcher.contains('gone') + + +def test_stale_case_result_cannot_poison_a_rehandled_job( + docker_dispatcher: Dispatcher, + monkeypatch, +): + # After an abort the backend requeues the job, and this same runner can + # claim it again under the SAME job_id while a worker from the first + # attempt is still running. A lookup by job_id alone would fetch the new + # attempt's state and complete it with the old run's data; the generation + # token pins the write to the exact handle() the worker was spawned under. + class StubExecutor: + + def __init__(self, job_id, *args, **kwargs): + pass + + def run(self): + return {'Status': 'AC'} + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + StubExecutor) + _write_job_dir(docker_dispatcher, 'aba', Language.PY, [1]) + docker_dispatcher.handle(job_id='aba', submission_id='sub-1') + stale_ctx = docker_dispatcher.result['aba'] + docker_dispatcher.release('aba') + docker_dispatcher.handle(job_id='aba', submission_id='sub-2') + fresh_ctx = docker_dispatcher.result['aba'] + assert fresh_ctx is not stale_ctx + + kwargs = dict( + job_id='aba', + case_no='0000', + mem_limit=65536, + time_limit=1000, + case_in_path='in.txt', + case_out_path='out.txt', + lang=Language.PY, + ) + docker_dispatcher.create_container(**kwargs, ctx=stale_ctx) + + # The old attempt's result never reached the new attempt's state. + assert docker_dispatcher.result['aba'].cases['0000'] is None + + # The current generation still writes normally (testing=True, so the job + # is not post-processed/released and the state stays inspectable). + docker_dispatcher.create_container(**kwargs, ctx=fresh_ctx) + assert docker_dispatcher.result['aba'].cases['0000']['status'] == 'AC' + + +def test_stale_compile_result_is_dropped( + docker_dispatcher: Dispatcher, + monkeypatch, +): + # Same generation contract for the compile path: a compile that finishes + # after its job was released / re-handled must not write its result into + # the new attempt (a stale CE would fail every case of a healthy rerun). + class StubExecutor: + + def __init__(self, *args, **kwargs): + pass + + def compile(self): + return {'Status': 'CE'} + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + StubExecutor) + _write_job_dir(docker_dispatcher, 'cs', Language.C, [1]) + docker_dispatcher.handle(job_id='cs', submission_id='sub-1') + stale_ctx = docker_dispatcher.result['cs'] + docker_dispatcher.release('cs') + docker_dispatcher.handle(job_id='cs', submission_id='sub-2') + + docker_dispatcher.compile('cs', Language.C, ctx=stale_ctx) + # Dropped before the write: neither generation saw the stale result. + assert stale_ctx.compile_result is None + assert docker_dispatcher.result['cs'].compile_result is None + + docker_dispatcher.compile('cs', + Language.C, + ctx=docker_dispatcher.result['cs']) + assert docker_dispatcher.result['cs'].compile_result['Status'] == 'CE' + + +def test_stale_queue_entry_is_discarded_not_bound_to_new_generation( + docker_dispatcher: Dispatcher, + monkeypatch, +): + # Queue entries carry the generation they were enqueued under. A task + # left over from a released/re-handled generation must be discarded at + # dequeue -- binding it to the CURRENT generation would double-run the + # case (or IndexError on a task shape the new meta no longer has, killing + # the run loop). The old entry is re-injected ahead of the new one, so + # the run loop meets the stale entry first. + ran = [] + + class StubExecutor: + + def __init__(self, job_id, *args, **kwargs): + pass + + def run(self): + ran.append('run') + return {'Status': 'AC'} + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + StubExecutor) + _write_job_dir(docker_dispatcher, 'q1', Language.PY, [1]) + docker_dispatcher.handle(job_id='q1', submission_id='sub-1') + stale_task = docker_dispatcher.queue.get_nowait() + assert docker_dispatcher.queue.empty() + + docker_dispatcher.release('q1') + docker_dispatcher.handle(job_id='q1', submission_id='sub-2') + fresh_task = docker_dispatcher.queue.get_nowait() + assert stale_task.ctx is not fresh_task.ctx + + # Stale entry first, then the current one. + docker_dispatcher.queue.put_nowait(stale_task) + docker_dispatcher.queue.put_nowait(fresh_task) + docker_dispatcher.start() + + deadline = time.time() + 5.0 + while docker_dispatcher.result['q1'].cases['0000'] is None: + assert time.time() < deadline, 'current generation never judged' + time.sleep(0.02) + assert ran == ['run'] # the stale entry never spawned work + assert docker_dispatcher.result['q1'].cases['0000']['status'] == 'AC' + + +def test_handle_reports_when_it_supersedes_a_generation( + docker_dispatcher: Dispatcher): + # The poller balances the claim tracker on this return value: a re-claim + # of the same job_id kills the previous generation, which will never + # produce an outcome, so its claim count has to be released somewhere. + _write_job_dir(docker_dispatcher, 'sup', Language.PY, [1]) + + assert docker_dispatcher.handle(job_id='sup', + submission_id='sub-1') is False + first_ctx = docker_dispatcher.result['sup'] + + assert docker_dispatcher.handle(job_id='sup', + submission_id='sub-2') is True + second_ctx = docker_dispatcher.result['sup'] + assert second_ctx is not first_ctx + # The new generation starts clean; nothing carries over by job_id. + assert second_ctx.compile_result is None + assert second_ctx.started is False + assert second_ctx.submission_id == 'sub-2' + + # A released job_id is a fresh handle, not a supersede. + docker_dispatcher.release('sup') + assert docker_dispatcher.handle(job_id='sup', + submission_id='sub-3') is False + + +def test_drain_unstarted_leaves_in_flight_jobs_running( + docker_dispatcher: Dispatcher): + # The latch must not strand a job that is already in flight: its remaining + # cases still need to run, or the job never completes and drain waits out + # its whole budget. + _write_job_dir(docker_dispatcher, 'inflight', Language.PY, [2]) + docker_dispatcher.handle(job_id='inflight', submission_id='sub-inflight') + docker_dispatcher.result['inflight'].started = True + + assert docker_dispatcher.drain_unstarted() == [] + assert docker_dispatcher.contains('inflight') + + +def test_on_complete_wiring_receives_result_and_releases( + docker_dispatcher: Dispatcher, + submission_generator, +): + recorded = [] + docker_dispatcher.testing = False + docker_dispatcher.on_complete = lambda *a: recorded.append(a) + + job_id, prob = next( + (i, p) for i, p in submission_generator.submission_ids.items() + if p == 'normal-submission') + docker_dispatcher.handle(job_id=job_id, submission_id=f'sub-{job_id}') + + task_content = docker_dispatcher.result[job_id].cases + for case_no in list(task_content): + docker_dispatcher.on_case_complete( + job_id=job_id, + case_no=case_no, + stdout='', + stderr='', + exit_code=0, + exec_time=1, + mem_usage=1, + prob_status='AC', + ) + + assert len(recorded) == 1 + got_job_id, got_submission_id, tasks = recorded[0] + assert got_job_id == job_id + assert got_submission_id == f'sub-{job_id}' + # normal-submission has two single-case tasks -> [[case], [case]]. + assert [len(t) for t in tasks] == [1, 1] + assert all(c['status'] == 'AC' for t in tasks for c in t) + # release() ran as part of on_job_complete. + assert not docker_dispatcher.contains(job_id) + + +def test_on_job_complete_enqueues_before_releasing( + docker_dispatcher: Dispatcher): + # has_jobs() is bool(self.result), so releasing before the result is + # handed to on_complete would open a window where a concurrent drain sees + # an empty dispatcher, stops the sender, and loses this result entirely. + observed = {} + + def on_complete(job_id, submission_id, tasks): + observed['contains'] = docker_dispatcher.contains(job_id) + observed['has_jobs'] = docker_dispatcher.has_jobs() + + docker_dispatcher.testing = False + docker_dispatcher.on_complete = on_complete + _write_job_dir(docker_dispatcher, 'oc', Language.PY, [2]) + docker_dispatcher.handle(job_id='oc', submission_id='sub-oc') + + task_content = docker_dispatcher.result['oc'].cases + for case_no in list(task_content): + docker_dispatcher.on_case_complete( + job_id='oc', + case_no=case_no, + stdout='', + stderr='', + exit_code=0, + exec_time=1, + mem_usage=1, + prob_status='AC', + ) + + assert observed == {'contains': True, 'has_jobs': True} + # ...and the release still happens, right after. + assert not docker_dispatcher.contains('oc') + + +def test_compile_je_fallback(docker_dispatcher: Dispatcher, monkeypatch): + # A crash the executor does not wrap (e.g. docker APIError) must still + # produce a JE compile result so the job can complete instead of hanging. + class BoomExecutor: + + def __init__(self, *args, **kwargs): + pass + + def compile(self): + raise RuntimeError('docker exploded') + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + BoomExecutor) + _write_job_dir(docker_dispatcher, 'cj', Language.C, [1]) + docker_dispatcher.handle(job_id='cj', submission_id='sub-cj') + + docker_dispatcher.compile('cj', + Language.C, + ctx=docker_dispatcher.result['cj']) + + assert docker_dispatcher.result['cj'].compile_result['Status'] == 'JE' + + +def test_execute_je_fallback(docker_dispatcher: Dispatcher, monkeypatch): + # Same contract as the compile path, for the execute path: an error the + # executor does not wrap must still resolve the case as JE. Left + # unresolved, the case never reports, the job never completes, and a drain + # waits out its whole budget on it before giving up. + # + # run() is the failure point here, but the guard covers the whole body -- + # executor construction and extract_compile_result included -- so a stub + # that blew up in __init__ would land on JE just the same. + class BoomExecutor: + + def __init__(self, *args, **kwargs): + pass + + def run(self): + raise RuntimeError('docker exploded') + + monkeypatch.setattr('dispatcher.dispatcher.SubmissionExecutor', + BoomExecutor) + _write_job_dir(docker_dispatcher, 'ej', Language.PY, [1]) + docker_dispatcher.handle(job_id='ej', submission_id='sub-ej') + + docker_dispatcher.create_container( + 'ej', + '0000', + 65536, + 1000, + 'in.txt', + 'out.txt', + Language.PY, + ctx=docker_dispatcher.result['ej'], + ) + + assert docker_dispatcher.result['ej'].cases['0000']['status'] == 'JE' diff --git a/tests/test_drain_integration.py b/tests/test_drain_integration.py new file mode 100644 index 0000000..9d22187 --- /dev/null +++ b/tests/test_drain_integration.py @@ -0,0 +1,575 @@ +"""End-to-end drain tests against REAL runner threads. + +tests/test_main.py pins the drain *ordering* with synchronous fakes, which +never exercises a live PollerThread, a poller blocked inside prep, or an +actual ResultSenderThread flush. These tests wire the real PollerThread / +ResultSenderThread / HeartbeatThread / Dispatcher together and only fake the +three things that would reach outside the process: the backend HTTP client, +the container executor, and the local job-dir cleanup. +""" + +import queue +import threading +import time +from types import SimpleNamespace + +import pytest + +import main +import dispatcher.dispatcher as dispatcher_mod +from dispatcher import job as job_mod +from dispatcher.constant import Language +from dispatcher.dispatcher import Dispatcher +from runner.active_jobs import ActiveJobTracker +from runner.client import JobPayload, RunnerConfig, RunnerIdentity +from runner.heartbeat import HeartbeatThread +from runner.poller import PollerThread +from runner.result_sender import CompleteRequest, ResultSenderThread +from tests.test_dispatcher import _write_job_dir + +# Same path conftest.py uses; it is absent from the repo, so the Dispatcher +# falls back to its defaults and the test sets what it needs explicitly. +TEST_CONFIG_PATH = '.config/dispatcher.test.json' +# Every blocking wait in this file is bounded, so a regression fails the test +# instead of hanging the suite. +WAIT_TIMEOUT = 10.0 + + +def wait_until(predicate, message, timeout=WAIT_TIMEOUT): + """Bounded poll-until-deadline; fails loudly instead of hanging.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.01) + raise AssertionError(f'timed out after {timeout}s waiting for {message}') + + +def release_when(predicate, event, message): + """Set ``event`` once ``predicate`` holds, from a helper thread. + + Releases have to be event-driven rather than racing a fixed delay: drain + runs on the test thread, and a slow CI box would otherwise let the release + land before the drain step it is supposed to follow, silently testing a + different interleaving. The event is set even on timeout so a wrong + predicate fails an assertion instead of hanging the drain. + """ + + def _run(): + try: + wait_until(predicate, message) + finally: + event.set() + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + return thread + + +def make_payload(job_id, submission_id): + # python3 (language 2) on purpose: the dispatcher enqueues no Compile job + # for it, so one payload == exactly one Execute task and the capacity + # bookkeeping below stays trivial to reason about. + return JobPayload( + job_id=job_id, + submission_id=submission_id, + problem_id=42, + language=int(Language.PY), + code_url='http://minio/code.zip', + checker=None, + tasks=[{ + 'taskScore': 100, + 'memoryLimit': 65536, + 'timeLimit': 1000, + 'caseCount': 1, + }], + ) + + +class RecordingClient: + """The four backend calls the runner threads make, all recorded. + + ``next_job`` walks a script; a ``threading.Event`` entry in the script is a + gate the test opens when it wants the *next* claim to happen, which is how + these tests pin down claim ordering without sleeping. Nothing here raises, + so the sender's retry path stays out of the picture. + """ + + def __init__(self, script): + self._lock = threading.Lock() + self._script = list(script) + self.completes = [] + self.aborts = [] + self.heartbeats = [] + self.gate_timed_out = False + + def _pop(self): + with self._lock: + return self._script.pop(0) if self._script else None + + def next_job(self, identity): + item = self._pop() + if isinstance(item, threading.Event): + # Generous: the dispatcher run loop sleeps a full second when its + # queue is empty, so the test needs several of those to line up + # its preconditions before opening the gate. + if not item.wait(WAIT_TIMEOUT): + self.gate_timed_out = True + return None + item = self._pop() + return item + + def complete(self, identity, job_id, tasks): + with self._lock: + self.completes.append((job_id, tasks)) + + def abort(self, identity, job_id, reason): + with self._lock: + self.aborts.append((job_id, reason)) + + def heartbeat(self, identity, active_job_ids): + with self._lock: + self.heartbeats.append(active_job_ids) + + +class ExecutorRecorder: + """Shared state for the StubExecutor installed over SubmissionExecutor.""" + + def __init__(self): + self.started = threading.Event() + self.release = threading.Event() + self.lock = threading.Lock() + self.ran = [] + + +def install_stub_executor(monkeypatch, recorder): + """Replace the container executor with one the test drives by hand.""" + + class StubExecutor: + + def __init__(self, job_id, *args, **kwargs): + self.job_id = job_id + + def run(self): + with recorder.lock: + recorder.ran.append(self.job_id) + recorder.started.set() + # Blocks until the test releases it, so "in-flight work" is a + # state the test fully controls. Bounded so a stuck test fails. + recorder.release.wait(5.0) + return { + 'Status': 'AC', + 'Stdout': '', + 'Stderr': '', + 'DockerExitCode': 0, + 'Duration': 1, + 'MemUsage': 1, + } + + def compile(self): + return {'Status': 'AC'} + + monkeypatch.setattr(dispatcher_mod, 'SubmissionExecutor', StubExecutor) + + +class Rig: + """Owns every thread/event so teardown always unblocks and joins them.""" + + def __init__(self): + self.threads = [] + self.events = [] + + def add_thread(self, thread): + self.threads.append(thread) + return thread + + def add_event(self, event): + self.events.append(event) + return event + + def shutdown(self): + # Release blocked executors / preps first: a thread parked on one of + # these would otherwise make the joins below time out. + for event in self.events: + event.set() + for thread in reversed(self.threads): + thread.stop() + for thread in reversed(self.threads): + if thread.is_alive(): + thread.join(timeout=WAIT_TIMEOUT) + + +@pytest.fixture +def rig(): + r = Rig() + try: + yield r + finally: + # Runs even when an assertion fails, so a broken drain can never hang + # the rest of the suite. + r.shutdown() + + +def build_stack(rig, tmp_path, client, prepare, wrap_dispatch=None): + """Real dispatcher + sender + poller + heartbeat, started and tracked.""" + identity = RunnerIdentity( + 'rn_x', + 'tok', + RunnerConfig( + heartbeat_interval_sec=60, + poll_interval_sec=0.01, + max_concurrent_jobs=8, + ), + ) + result_queue = queue.Queue() + tracker = ActiveJobTracker() + + d = Dispatcher( + TEST_CONFIG_PATH, + on_complete=lambda job_id, submission_id, tasks: result_queue.put( + CompleteRequest(job_id, submission_id, tasks)), + ) + d.SUBMISSION_DIR = tmp_path + # testing stays False: on_job_complete's real post-processing (the nested + # task/case result shape handed to on_complete) is exactly what we assert. + assert d.testing is False + # One container at a time. This is what makes the "claimed but never + # started" state reachable deterministically: while job A's executor is + # blocked, the run loop sees container_count >= MAX_CONTAINER_SIZE, never + # calls queue.get(), and job B stays in result/ but out of started/. + d.MAX_CONTAINER_SIZE = 1 + + sender = ResultSenderThread( + client, + identity, + tracker, + result_queue, + # Fake both so nothing touches the real submissions/ tree: the default + # cleanup resolves job dirs against the module-level config path, not + # the tmp_path we gave the dispatcher. + cleanup=lambda job_id: None, + backup=lambda job_id: None, + ) + poller = PollerThread( + client, + identity, + tracker, + result_queue, + dispatch=(wrap_dispatch(d.handle) + if wrap_dispatch is not None else d.handle), + prepare=prepare, + ) + heartbeat = HeartbeatThread(client, + identity, + tracker, + on_fatal=lambda: None) + + for thread in (d, sender, poller, heartbeat): + rig.add_thread(thread) + thread.start() + return d, sender, poller, heartbeat, result_queue, tracker + + +def test_drain_completes_in_flight_and_returns_unstarted( + tmp_path, + monkeypatch, + rig, +): + # A started job must still report complete; a claimed-but-never-started + # job must come back as abort('drain'). Both through the real sender. + recorder = ExecutorRecorder() + rig.add_event(recorder.release) + install_stub_executor(monkeypatch, recorder) + + gate = rig.add_event(threading.Event()) + # The gate sits between A and B so B is only claimed once the test says so. + client = RecordingClient( + [make_payload('A', 'sub_a'), gate, + make_payload('B', 'sub_b')]) + + # The poller starts claiming inside build_stack, so prep cannot close over + # the dispatcher object; it only needs the job-dir root, which is tmp_path. + job_root = SimpleNamespace(SUBMISSION_DIR=tmp_path) + + def prepare(payload): + _write_job_dir(job_root, payload.job_id, Language.PY, [1]) + + d, sender, poller, heartbeat, result_queue, tracker = build_stack( + rig, tmp_path, client, prepare) + + # Open the gate only once A's executor is genuinely running AND the + # container slot is taken AND the dispatcher marked A started. That exact + # ordering is what guarantees the run loop is saturated when B arrives. + assert recorder.started.wait(WAIT_TIMEOUT), 'executor for A never ran' + wait_until(lambda: d.container_count == 1, 'container slot to be taken') + wait_until(lambda: d.contains('A') and d.result['A'].started, + 'A to be marked started') + gate.set() + + wait_until(lambda: d.contains('B'), 'B to be handled by the dispatcher') + assert not d.result['B'].started, \ + 'B must stay unstarted for this test to mean anything' + + # Let A finish only once drain_unstarted() has already returned B (B is + # gone from the dispatcher). Releasing earlier would free the container + # slot, let the run loop start B, and quietly turn this into a test of a + # different interleaving. + release_when( + lambda: not d.contains('B'), + recorder.release, + 'B to be released by drain_unstarted', + ) + t0 = time.monotonic() + rc = main.drain( + poller, + d, + sender, + heartbeat, + result_queue, + tracker, + sleep=lambda s: time.sleep(0.02), + # Bound the budget: drain runs on the test thread, so a regression must + # fail the test rather than hang the whole suite. + timeout_sec=10.0, + ) + elapsed = time.monotonic() - t0 + + assert rc == 0 + assert client.aborts == [('B', 'drain')] + assert [job_id for job_id, _ in client.completes] == ['A'] + tasks = client.completes[0][1] + assert len(tasks) == 1 and len(tasks[0]) == 1 + assert tasks[0][0]['status'] == 'AC' + assert tracker.snapshot() == [] + assert not client.gate_timed_out + # The dispatcher run loop sleeps up to 1s before noticing stop(), so the + # bound is generous; the point is that drain returns on its own. + assert elapsed < 8.0, f'drain took {elapsed:.2f}s' + + +class GhostFeeder: + """Keeps the dispatcher run loop hot so it can react within a test. + + The run loop sleeps a full second whenever its queue is empty -- longer + than a whole drain takes here, so by default it never even wakes up and + any "it did not start the job" assertion passes for the wrong reason. This + keeps a backlog of tasks for jobs that do not exist (the loop discards + them), so the loop is genuinely spinning and will pounce on anything + dispatched during the drain. + + ``fed`` only grows once the loop consumes what it was given, which is how + a test can wait for the loop to actually be awake. + """ + + def __init__(self, dispatcher, done): + self._dispatcher = dispatcher + self._done = done + self.fed = 0 + + def run(self): + while not self._done.is_set(): + while self._dispatcher.queue.qsize() < 32 and \ + not self._done.is_set(): + # ctx is a throwaway object, so the identity check discards + # these instantly -- which is the point: keep the loop hot. + self._dispatcher.queue.put( + job_mod.Execute(job_id='ghost', + task_id=0, + case_id=0, + ctx=object())) + self.fed += 1 + time.sleep(0.001) + + +def keep_run_loop_busy(rig, dispatcher): + """Start a GhostFeeder and wait until the run loop is demonstrably awake.""" + feeder = GhostFeeder(dispatcher, rig.add_event(threading.Event())) + threading.Thread(target=feeder.run, daemon=True).start() + wait_until(lambda: feeder.fed > 64, 'run loop to start consuming tasks') + return feeder + + +def test_drain_latches_dispatcher_before_waiting_for_poller( + tmp_path, + monkeypatch, + rig, +): + # End-to-end shape of the ordering: the poller finishes a prep and + # dispatches while drain is joining it, with the run loop awake and hungry + # rather than parked in its idle sleep. The job must still come back as a + # drain abort and must never be judged. + # + # This exercises the path; it does not pin the interleaving. Whichever + # side of the latch the dispatch lands on, the outcome is the same by + # construction: before the latch it publishes and the scan returns it; + # after the latch handle() refuses it and the poller aborts it as drain. + # The guarantees themselves are pinned elsewhere: the required call order + # in tests/test_main.py, the refusal in test_handle_refuses_jobs_once_ + # draining, and the latch's effect on the run loop in + # tests/test_dispatcher.py. + recorder = ExecutorRecorder() + rig.add_event(recorder.release) + install_stub_executor(monkeypatch, recorder) + + client = RecordingClient([make_payload('L', 'sub_l')]) + prep_entered = threading.Event() + prep_release = rig.add_event(threading.Event()) + job_root = SimpleNamespace(SUBMISSION_DIR=tmp_path) + + def prepare(payload): + prep_entered.set() + prep_release.wait(WAIT_TIMEOUT) + _write_job_dir(job_root, payload.job_id, Language.PY, [1]) + + def linger_after_dispatch(handle): + # Widen the window between the dispatch and the poller thread actually + # ending. In production that window is one GIL slice -- narrow, but it + # is precisely the window this ordering must close, and a test that + # cannot lose the race is not testing the ordering at all. + def _dispatch(job_id, submission_id): + handle(job_id, submission_id) + time.sleep(0.3) + + return _dispatch + + d, sender, poller, heartbeat, result_queue, tracker = build_stack( + rig, tmp_path, client, prepare, wrap_dispatch=linger_after_dispatch) + keep_run_loop_busy(rig, d) + + assert prep_entered.wait(WAIT_TIMEOUT), 'prep was never entered' + + # Release prep only once drain has stopped the poller, so the dispatch + # lands squarely inside the join window. + release_when( + poller._stop_event.is_set, + prep_release, + 'poller stop to be requested', + ) + rc = main.drain( + poller, + d, + sender, + heartbeat, + result_queue, + tracker, + sleep=lambda s: time.sleep(0.02), + timeout_sec=10.0, + ) + + assert rc == 0 + assert client.aborts == [('L', 'drain')] + assert client.completes == [] + assert recorder.ran == [], 'job started judging after shutdown' + assert not d.contains('L') + + +def test_drain_returns_claim_whose_prep_succeeded_after_stop( + tmp_path, + monkeypatch, + rig, +): + # The nastier variant: stop() lands while prep is blocked, but prep then + # SUCCEEDS, so the poller tries to dispatch the job on its way out. It + # must come back as a drain abort and must never be judged. + # + # The dispatch here lands after the latch, so handle() refuses it and the + # poller itself queues the drain abort. This test pins the end-to-end + # outcome; the refusal semantics are pinned in + # test_handle_refuses_jobs_once_draining. + recorder = ExecutorRecorder() + rig.add_event(recorder.release) + install_stub_executor(monkeypatch, recorder) + + client = RecordingClient([make_payload('E', 'sub_e')]) + prep_entered = threading.Event() + prep_release = rig.add_event(threading.Event()) + job_root = SimpleNamespace(SUBMISSION_DIR=tmp_path) + + def prepare(payload): + prep_entered.set() + prep_release.wait(WAIT_TIMEOUT) + # Succeeds this time: the claim is fully prepped and will be + # dispatched even though the poller is already stopping. + _write_job_dir(job_root, payload.job_id, Language.PY, [1]) + + d, sender, poller, heartbeat, result_queue, tracker = build_stack( + rig, tmp_path, client, prepare) + + assert prep_entered.wait(WAIT_TIMEOUT), 'prep was never entered' + + release_when( + poller._stop_event.is_set, + prep_release, + 'poller stop to be requested', + ) + rc = main.drain( + poller, + d, + sender, + heartbeat, + result_queue, + tracker, + sleep=lambda s: time.sleep(0.02), + timeout_sec=10.0, + ) + + assert rc == 0 + assert client.aborts == [('E', 'drain')] + assert client.completes == [] + # Never judged: no work may begin once shutdown was requested. + assert recorder.ran == [] + assert not d.contains('E') + assert tracker.snapshot() == [] + + +def test_drain_returns_claim_blocked_in_prep(tmp_path, monkeypatch, rig): + # The P1 fix end-to-end: a claim still stuck in prep when SIGTERM lands + # comes back as 'drain' (attempt-neutral), never 'prep_failed'. + recorder = ExecutorRecorder() + rig.add_event(recorder.release) + install_stub_executor(monkeypatch, recorder) + + client = RecordingClient([make_payload('C', 'sub_c')]) + prep_entered = threading.Event() + prep_release = rig.add_event(threading.Event()) + + def prepare(payload): + prep_entered.set() + # Models the slow backend I/O the reviewer described: blocked, then + # failing, with stop() landing in between. + prep_release.wait(WAIT_TIMEOUT) + raise RuntimeError('backend gone') + + d, sender, poller, heartbeat, result_queue, tracker = build_stack( + rig, tmp_path, client, prepare) + + assert prep_entered.wait(WAIT_TIMEOUT), 'prep was never entered' + + # Unblock prep only after stop() has landed: that is the exact state the + # P1 fix is about (I/O that fails *while* draining must not be reported as + # prep_failed). Releasing before stop() would just be an ordinary retry. + release_when( + poller._stop_event.is_set, + prep_release, + 'poller stop to be requested', + ) + t0 = time.monotonic() + rc = main.drain( + poller, + d, + sender, + heartbeat, + result_queue, + tracker, + sleep=lambda s: time.sleep(0.02), + # Bound the budget: drain runs on the test thread, so a regression must + # fail the test rather than hang the whole suite. + timeout_sec=10.0, + ) + elapsed = time.monotonic() - t0 + + assert rc == 0 + assert client.aborts == [('C', 'drain')] + assert client.completes == [] + assert tracker.snapshot() == [] + assert elapsed < 8.0, f'drain took {elapsed:.2f}s' diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..b08e539 --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,572 @@ +import os +import queue +import signal +import threading + +import pytest + +import main +import runner.config +from runner.active_jobs import ActiveJobTracker +from runner.client import BackendAuthError, RunnerConfig, RunnerIdentity +from runner.result_sender import ( + AbortRequest, + CompleteRequest, + ResultSenderThread, +) + + +class ManualClock: + """A monotonic() stand-in that only moves when a fake join consumes it.""" + + def __init__(self): + self.t = 0.0 + + def __call__(self): + return self.t + + +class FakeComponent: + """Records stop()/join() calls into a shared ordered log. + + With ``clock`` set, join() consumes its whole granted timeout from that + clock -- modelling a component that blocks for everything it was given, + which is what exposes budget leaks between drain steps. + """ + + def __init__(self, log, name, alive_after_join=False, clock=None): + self._log = log + self._name = name + self._alive_after_join = alive_after_join + self._clock = clock + self.join_timeouts = [] + + def stop(self): + self._log.append(f'{self._name}.stop') + + def join(self, timeout=None): + self._log.append(f'{self._name}.join') + self.join_timeouts.append(timeout) + if self._clock is not None and timeout: + self._clock.t += timeout + + def is_alive(self): + return self._alive_after_join + + +class FakeDispatcher(FakeComponent): + + def __init__(self, + log, + victims, + has_jobs_seq, + always_has_jobs=False, + clock=None): + super().__init__(log, 'dispatcher', clock=clock) + self._victims = victims + self._has_jobs_seq = list(has_jobs_seq) + self._always_has_jobs = always_has_jobs + + def stop_accepting(self): + self._log.append('dispatcher.stop_accepting') + + def drain_unstarted(self): + self._log.append('dispatcher.drain_unstarted') + return self._victims + + def has_jobs(self): + # Logged like the other calls so tests can pin *when* the in-flight + # wait happened relative to dispatcher.stop. + self._log.append('dispatcher.has_jobs') + if self._always_has_jobs: + return True + return self._has_jobs_seq.pop(0) if self._has_jobs_seq else False + + +def test_drain_orders_shutdown_and_aborts_unstarted(): + log = [] + victims = [('jb_1', 'sub_1'), ('jb_2', 'sub_2')] + poller = FakeComponent(log, 'poller') + dispatcher = FakeDispatcher(log, victims, [True, True, False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + result_queue = queue.Queue() + + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + result_queue, + ActiveJobTracker(), + sleep=lambda s: None, + ) + + assert rc == 0 + # poller fully stopped before the dispatcher is scanned for victims. + assert log.index('poller.stop') < log.index('poller.join') + assert log.index('poller.join') < log.index('dispatcher.drain_unstarted') + # The latch comes FIRST -- before the poller is even asked to stop. Any + # later and there is a window where the poller finishes a prep, dispatches, + # and the run loop starts that job while the dispatcher still accepts work: + # judging that begins after shutdown, which the scan then leaves alone + # because it looks in-flight. + assert log[0] == 'dispatcher.stop_accepting', log + assert log.index('dispatcher.stop_accepting') < log.index('poller.stop') + # both unstarted jobs aborted with reason 'drain'. + aborts = [] + while not result_queue.empty(): + aborts.append(result_queue.get_nowait()) + assert all( + isinstance(a, AbortRequest) and a.reason == 'drain' for a in aborts) + assert {(a.job_id, a.submission_id) for a in aborts} == set(victims) + # dispatcher stops before sender, sender before heartbeat. + assert log.index('dispatcher.stop') < log.index('sender.stop') + assert log.index('sender.stop') < log.index('heartbeat.stop') + # sender.stop precedes sender.join (drains the queue before joining). + assert log.index('sender.stop') < log.index('sender.join') + # every join is capped by the shared budget, never unbounded. + for component in (poller, dispatcher, sender, heartbeat): + assert component.join_timeouts and all( + t is not None and t > 0 for t in component.join_timeouts) + + +def test_drain_stops_dispatcher_only_after_the_in_flight_wait(): + # The dispatcher has to outlive the whole has_jobs() wait. A job whose + # first case is in flight still needs the run loop alive to dispatch its + # remaining cases; stop the dispatcher before the wait and that job never + # completes, has_jobs() never clears, the drain burns its entire budget, + # and the results already computed are thrown away. + log = [] + poller = FakeComponent(log, 'poller') + # At least one True, so the wait loop genuinely iterates instead of + # falling straight through and passing for the wrong reason. + dispatcher = FakeDispatcher(log, [], [True, True, False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + ) + + assert rc == 0 + assert log.count('dispatcher.has_jobs') >= 2, log + last_wait = len(log) - 1 - log[::-1].index('dispatcher.has_jobs') + assert last_wait < log.index('dispatcher.stop'), log + + +def test_drain_budget_is_shared_not_per_component(): + # Each join must get what is LEFT of the budget, not a fresh copy of it. + # With a fresh copy per component, four slow components could stretch the + # drain to 4x the budget and blow past docker's stop grace period. + log = [] + poller = FakeComponent(log, 'poller') + dispatcher = FakeDispatcher(log, [], [False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + # Each monotonic() call advances one second, so the budget visibly drains. + ticks = iter(range(1000)) + + main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + timeout_sec=100.0, + monotonic=lambda: float(next(ticks)), + ) + + # Nobody gets a fresh copy of the budget, and each grant reflects what the + # earlier steps already consumed. + for component in (poller, dispatcher, sender, heartbeat): + assert component.join_timeouts[0] < 100.0 + # The sender and heartbeat draw from the same shrinking pool, in order. + assert sender.join_timeouts[0] > heartbeat.join_timeouts[0] + # The poller is capped at its own share of the budget. + assert poller.join_timeouts[0] <= 100.0 * main.POLLER_JOIN_SHARE + # The dispatcher's grant excludes the reporting reserve: that slice + # belongs to the sender, whatever the dispatcher does with its join. + assert dispatcher.join_timeouts[0] <= \ + 100.0 * (1 - main.REPORTING_RESERVE_SHARE) + + +def test_drain_keeps_a_reserve_for_reporting(): + # A long-running in-flight job must not be able to consume the budget the + # reporting tail needs. If the wait ran to the last second, the sender + # would be joined with nothing left and every complete already computed -- + # plus the drain aborts just queued -- would go unsent. + log = [] + poller = FakeComponent(log, 'poller') + dispatcher = FakeDispatcher(log, [], [], always_has_jobs=True) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + ticks = iter(range(1000)) + + main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + timeout_sec=100.0, + monotonic=lambda: float(next(ticks)), + ) + + assert sender.join_timeouts[0] > 0, 'sender left with no time to flush' + assert heartbeat.join_timeouts[0] > 0 + + +def test_drain_dispatcher_join_cannot_eat_the_reporting_reserve(): + # The reserve is only real if EVERY step between the in-flight wait and + # the sender honours it. A dispatcher join granted the full remainder + # would consume the reserve right after the wait loop so carefully kept + # it, and the sender would still end up with nothing. The fake joins here + # consume everything they are granted, which is exactly the case that + # exposes the leak. + log = [] + clock = ManualClock() + poller = FakeComponent(log, 'poller', clock=clock) + dispatcher = FakeDispatcher(log, [], [False], clock=clock) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + + main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + timeout_sec=100.0, + monotonic=clock, + ) + + # poller consumed its full share (25), dispatcher its full grant; the + # sender must still be handed at least the whole reserve (25). + reserve = 100.0 * main.REPORTING_RESERVE_SHARE + assert sender.join_timeouts[0] >= reserve, sender.join_timeouts + + +def test_drain_scans_exactly_once_when_the_poller_finished(): + # Quiescence comes from the dispatcher latch (handle() refuses jobs after + # it), so one scan is complete by construction and a rescan would only + # suggest otherwise. + log = [] + poller = FakeComponent(log, 'poller') + dispatcher = FakeDispatcher(log, [], [False]) + + main.drain( + poller, + dispatcher, + FakeComponent(log, 'sender'), + FakeComponent(log, 'heartbeat'), + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + ) + + assert log.count('dispatcher.drain_unstarted') == 1, log + + +def test_drain_holds_the_sender_open_while_the_poller_lives(): + # A poller that outlived its join share is still capable of waking up and + # queueing a late 'drain' abort (its dispatch would be refused by the + # latched dispatcher). The sender must therefore stay alive until the + # poller thread has actually exited or the budget runs out -- never + # stopped early on an is_alive() guess. + log = [] + sleeps = [] + poller = FakeComponent(log, 'poller', alive_after_join=True) + dispatcher = FakeDispatcher(log, [], [False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + ticks = iter(range(1000)) + + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: sleeps.append(s), + timeout_sec=100.0, + monotonic=lambda: float(next(ticks)), + ) + + assert rc == 0 + # Still exactly one scan: the latch already guarantees nothing the scan + # would miss can be published, so a stuck poller changes the WAIT, not + # the scan count. + assert log.count('dispatcher.drain_unstarted') == 1, log + # The reporting tail genuinely waited for the (never-exiting) poller + # until the budget ran out, and only then stopped the sender. + assert len(sleeps) >= 50, len(sleeps) + assert 'sender.stop' in log + + +def test_drain_reporting_tail_waits_for_unreported_outcomes(): + # tracker entries are removed by the sender only after an outcome is + # finalized, so a non-empty tracker means reports are still owed. The + # tail must keep the sender alive until they are out -- and must exit as + # soon as they are, not burn the rest of the budget. + log = [] + tracker = ActiveJobTracker() + tracker.add('jb_x') + sleeps = [] + + def sleep(seconds): + # Models the sender finalizing the last outcome a moment later. + sleeps.append(seconds) + if len(sleeps) == 3: + tracker.remove('jb_x') + + poller = FakeComponent(log, 'poller') + dispatcher = FakeDispatcher(log, [], [False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + ticks = iter(range(1000)) + + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + tracker, + sleep=sleep, + timeout_sec=100.0, + monotonic=lambda: float(next(ticks)), + ) + + assert rc == 0 + # It waited (3 sleeps until the tracker emptied), then moved on promptly + # instead of spinning to the deadline (~90+ sleeps). + assert 3 <= len(sleeps) <= 20, len(sleeps) + assert 'sender.stop' in log + + +def test_drain_reporting_tail_survives_same_job_id_reclaim(): + # Two claims of the SAME job_id can be outstanding at once: the first + # claim's lease expired, the backend requeued the job, and this runner + # claimed it again. The first claim's outcome finalizing must not + # convince the tail that the second was reported too -- that was the + # failure mode of a set-based tracker, where one remove dropped both. + log = [] + tracker = ActiveJobTracker() + tracker.add('jb_x') # claim 1 + tracker.add('jb_x') # claim 2, same job_id + sleeps = [] + + def sleep(seconds): + sleeps.append(seconds) + if len(sleeps) == 1: + tracker.remove('jb_x') # claim 1's outcome finalized + if len(sleeps) == 4: + tracker.remove('jb_x') # claim 2's outcome finalized + + poller = FakeComponent(log, 'poller') + dispatcher = FakeDispatcher(log, [], [False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + ticks = iter(range(1000)) + + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + tracker, + sleep=sleep, + timeout_sec=100.0, + monotonic=lambda: float(next(ticks)), + ) + + assert rc == 0 + # With a set, the first remove would empty the tracker and the tail + # would stop after a single sleep. It must keep waiting for claim 2, + # then move on promptly once that one is reported as well. + assert 4 <= len(sleeps) <= 20, len(sleeps) + assert 'sender.stop' in log + + +def test_drain_gives_up_on_in_flight_jobs_when_budget_runs_out(): + # Waiting forever would hand the decision to docker's SIGKILL and lose the + # results already queued; the budget makes the drain finish on its own and + # leave the stragglers to lease expiry. + log = [] + dispatcher = FakeDispatcher(log, [], [], always_has_jobs=True) + poller = FakeComponent(log, 'poller') + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + timeout_sec=0.05, + ) + + assert rc == 0 + # It still shut everything down in order rather than bailing out early. + assert log.index('dispatcher.stop') < log.index('sender.stop') + assert log.index('sender.stop') < log.index('heartbeat.stop') + assert 'heartbeat.join' in log + + +def test_drain_logs_components_that_outlive_the_budget(caplog): + log = [] + poller = FakeComponent(log, 'poller', alive_after_join=True) + dispatcher = FakeDispatcher(log, [], [False]) + sender = FakeComponent(log, 'sender') + heartbeat = FakeComponent(log, 'heartbeat') + + with caplog.at_level('ERROR'): + rc = main.drain( + poller, + dispatcher, + sender, + heartbeat, + queue.Queue(), + ActiveJobTracker(), + sleep=lambda s: None, + timeout_sec=0.05, + ) + + assert rc == 0 + assert 'poller outlived the budget' in caplog.text + + +class RecordingResultClient: + """Records what the sender actually reported to the backend.""" + + def __init__(self): + self.completes = [] + self.aborts = [] + + def complete(self, identity, job_id, tasks): + self.completes.append(job_id) + + def abort(self, identity, job_id, reason): + self.aborts.append((job_id, reason)) + + +def test_result_sender_flushes_the_queue_after_stop(): + # Drain step 6 hands the sender everything it collected -- the completes + # from the in-flight jobs and the drain aborts from the scan -- and then + # stops it. If stop() discarded what was still queued, the drain would + # silently lose results it had already computed. So this uses the REAL + # sender: the fakes in the drain-ordering tests cannot see this. + client = RecordingResultClient() + result_queue = queue.Queue() + tracker = ActiveJobTracker() + identity = RunnerIdentity( + 'rn_x', + 'tok', + RunnerConfig( + heartbeat_interval_sec=60, + poll_interval_sec=1, + max_concurrent_jobs=8, + ), + ) + for i in range(3): + tracker.add(f'jb_{i}') + result_queue.put(CompleteRequest(f'jb_{i}', f'sub_{i}', [])) + tracker.add('jb_late') + result_queue.put(AbortRequest('jb_late', 'sub_late', 'drain')) + + sender = ResultSenderThread( + client, + identity, + tracker, + result_queue, + cleanup=lambda job_id: None, + backup=lambda job_id: None, + queue_poll_sec=0.01, + ) + # Everything is queued before the thread ever runs, so "stop() was already + # requested when the work arrived" is guaranteed rather than raced. + sender.stop() + sender.start() + try: + sender.join(timeout=10.0) + assert not sender.is_alive(), 'sender never exited after stop()' + finally: + # A failed assertion above must not leave the thread behind. + sender.stop() + + assert client.completes == ['jb_0', 'jb_1', 'jb_2'] + assert client.aborts == [('jb_late', 'drain')] + assert tracker.snapshot() == [] + + +def test_missing_token_returns_1_without_registering(monkeypatch): + monkeypatch.setattr(runner.config, 'RUNNER_REGISTRATION_TOKEN', '') + + def _boom(*args, **kwargs): + raise AssertionError('must not register without a token') + + monkeypatch.setattr(main, 'register_with_backoff', _boom) + + assert main.main() == 1 + + +def test_registration_401_returns_1(monkeypatch): + monkeypatch.setattr(runner.config, 'RUNNER_REGISTRATION_TOKEN', 'tok') + + def _reject(*args, **kwargs): + raise BackendAuthError('nope', status_code=401) + + monkeypatch.setattr(main, 'register_with_backoff', _reject) + + assert main.main() == 1 + + +class RejectingClient: + + def register(self, registration_token, name): + import requests + raise requests.ConnectionError('backend down') + + +def test_interruptible_registration_sleep_raises_on_shutdown(monkeypatch): + # A docker stop during a registration retry must abort immediately rather + # than sleep through the backoff. + monkeypatch.setattr(runner.config, 'RUNNER_REGISTRATION_TOKEN', 'tok') + shutdown = threading.Event() + shutdown.set() + + with pytest.raises(SystemExit): + main.register(RejectingClient(), shutdown) + + +def test_install_signal_handlers_sets_shutdown_on_sigterm(): + # docker stop sends a real SIGTERM; exercise that path rather than calling + # the handler directly. + shutdown = threading.Event() + previous_term = signal.getsignal(signal.SIGTERM) + previous_int = signal.getsignal(signal.SIGINT) + try: + main.install_signal_handlers(shutdown) + os.kill(os.getpid(), signal.SIGTERM) + assert shutdown.wait(1.0) + finally: + signal.signal(signal.SIGTERM, previous_term) + signal.signal(signal.SIGINT, previous_int) diff --git a/tests/test_runner_poller.py b/tests/test_runner_poller.py index 67c58c2..9bc34cd 100644 --- a/tests/test_runner_poller.py +++ b/tests/test_runner_poller.py @@ -1,10 +1,12 @@ import io import queue +import threading import time import pytest import requests +from dispatcher.exception import DispatcherDrainingError from runner.client import ( BackendAPIError, BackendAuthError, @@ -14,6 +16,7 @@ ) from runner.active_jobs import ActiveJobTracker from runner.result_sender import AbortRequest +from runner.config import PREP_MAX_ATTEMPTS from runner.poller import PollerThread, prepare_job import runner.poller as poller_mod @@ -57,16 +60,20 @@ def next_job(self, identity): class DispatchRecorder: - def __init__(self, fail_times=0, exc=None): + def __init__(self, fail_times=0, exc=None, superseded=False): self.calls = [] self._fail_times = fail_times self._exc = exc or RuntimeError('dispatch failed') + # Mirrors Dispatcher.handle()'s return value: True when this dispatch + # replaced a previous generation of the same job_id. + self._superseded = superseded def __call__(self, job_id, submission_id): self.calls.append((job_id, submission_id)) if self._fail_times > 0: self._fail_times -= 1 raise self._exc + return self._superseded class PrepareRecorder: @@ -303,6 +310,194 @@ def test_run_survives_unexpected_error_from_poll(): assert dispatch.calls == [('jb_1', 'sub_1')] +# --- drain behaviour --- + + +def test_stop_during_prep_returns_claim_as_drain(): + # stop() can land while the claim HTTP call is in flight, so the prep loop + # is the last place that can still return the claim. It must go back as + # 'drain' (attempt-neutral), never 'prep_failed'. + tracker = ActiveJobTracker() + dispatch = DispatchRecorder() + holder = {} + + class StopThenFail: + + def __init__(self): + self.calls = [] + + def __call__(self, payload): + self.calls.append(payload) + holder['poller']._stop_event.set() + raise RuntimeError('prep failed') + + prepare = StopThenFail() + poller, _, tracker, q, _, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + holder['poller'] = poller + + assert poller._poll_once() is False + assert len(prepare.calls) == 1 # no further attempt after stop() + assert dispatch.calls == [] + item = q.get_nowait() + assert isinstance(item, AbortRequest) + assert item.job_id == 'jb_1' + assert item.submission_id == 'sub_1' + assert item.reason == 'drain' + assert q.empty() + # Still in tracker: the sender removes it as it finalizes the abort. + assert tracker.snapshot() == ['jb_1'] + + +def test_stop_during_final_prep_attempt_returns_claim_as_drain(): + # stop() landing during the LAST attempt lands after the loop's own check + # has run for the last time; the exhausted path has to re-check, or the + # claim burns an attempt exactly when a rolling restart is under way. + tracker = ActiveJobTracker() + dispatch = DispatchRecorder() + holder = {} + + class FailAndStopOnLastAttempt: + + def __init__(self): + self.calls = [] + + def __call__(self, payload): + self.calls.append(payload) + if len(self.calls) == PREP_MAX_ATTEMPTS: + holder['poller']._stop_event.set() + raise RuntimeError('prep failed') + + prepare = FailAndStopOnLastAttempt() + poller, _, tracker, q, _, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + holder['poller'] = poller + + assert poller._poll_once() is False + assert len(prepare.calls) == PREP_MAX_ATTEMPTS + assert dispatch.calls == [] + item = q.get_nowait() + assert isinstance(item, AbortRequest) + assert item.reason == 'drain' + assert q.empty() + + +def test_dispatch_refused_by_draining_dispatcher_aborts_as_drain(): + # A prep that finishes after the drain latched the dispatcher still tries + # to dispatch; handle() refuses it without publishing anything. That claim + # never started any work, so it must go back as 'drain' (attempt-neutral), + # never 'prep_failed' -- otherwise every rolling restart would burn an + # attempt for the claim that lost this race. + tracker = ActiveJobTracker() + prepare = PrepareRecorder(tracker) + dispatch = DispatchRecorder(fail_times=1, + exc=DispatcherDrainingError('draining')) + poller, _, tracker, q, sleep, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + + result = poller._poll_once() + + assert result is False + assert len(dispatch.calls) == 1 # refused once, never retried + assert sleep.slept == [] + item = q.get_nowait() + assert isinstance(item, AbortRequest) + assert item.job_id == 'jb_1' + assert item.submission_id == 'sub_1' + assert item.reason == 'drain' + assert q.empty() + # Still in tracker: the sender removes it as it finalizes the abort. + assert tracker.snapshot() == ['jb_1'] + + +def test_supersede_dispatch_releases_the_dead_claims_count(): + # A dispatch that superseded a previous generation (handle() returned + # True) means the old claim of this job_id can never produce an outcome: + # nothing else will ever remove its tracker count, so the poller must. + # Without this, the dead count would hold a capacity slot forever and + # keep the drain's reporting tail waiting for a report that cannot come. + tracker = ActiveJobTracker() + tracker.add('jb_1') # the dead claim, left over from an expired lease + prepare = PrepareRecorder(tracker) + dispatch = DispatchRecorder(superseded=True) + poller, _, tracker, q, _, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + + assert poller._poll_once() is False + + assert dispatch.calls == [('jb_1', 'sub_1')] + assert q.empty() + # One count for the new claim; the dead claim's count was released. + assert len(tracker) == 1 + assert tracker.snapshot() == ['jb_1'] + + +def test_stop_before_claim_skips_next_job(): + poller, client, tracker, q, _, dispatch = build_poller([make_payload()]) + poller.stop() + + assert poller._poll_once() is True + assert client.call_count == 0 + assert len(tracker) == 0 + assert q.empty() + + +def test_prep_backoff_is_interruptible(): + # The production default sleep is stop-aware, so a drain must not have to + # wait out PREP_BACKOFF_SCHEDULE (1 + 2 seconds). No sleep= injection here + # on purpose: the default is the thing under test. + tracker = ActiveJobTracker() + prepared = threading.Event() + + def always_failing_prepare(payload): + prepared.set() + raise RuntimeError('prep failed') + + q = queue.Queue() + poller = PollerThread( + ScriptedClient([make_payload()] * 10), + make_identity(poll_interval=0.01), + tracker, + q, + DispatchRecorder(), + prepare=always_failing_prepare, + poll_interval_sec=0.01, + ) + + poller.start() + assert prepared.wait(1.0) + t0 = time.time() + poller.stop() + poller.join(timeout=1.0) + elapsed = time.time() - t0 + + assert not poller.is_alive() + assert elapsed < 1.0 + while not q.empty(): + item = q.get_nowait() + assert isinstance(item, AbortRequest) + assert item.reason == 'drain' + + +def test_default_sleep_is_stop_aware(): + poller = PollerThread( + ScriptedClient([]), + make_identity(), + ActiveJobTracker(), + queue.Queue(), + DispatchRecorder(), + ) + assert poller._sleep == poller._stop_event.wait + + # --- prepare_job unit tests --- diff --git a/tests/test_sigterm_process.py b/tests/test_sigterm_process.py new file mode 100644 index 0000000..961d602 --- /dev/null +++ b/tests/test_sigterm_process.py @@ -0,0 +1,366 @@ +"""SIGTERM drain proven against a REAL process (issue #70 done-criterion). + +Every other drain test runs in-process: tests/test_main.py fakes the threads, +tests/test_drain_integration.py wires the real threads but still calls +``main.drain()`` by hand, and the signal test only checks that the handler sets +an event. None of them prove that ``docker stop`` on a live runner actually +completes the drain and exits. This one does: it launches main.py as a +subprocess against a stub backend, sends it a real SIGTERM at a point where the +child is demonstrably inside prep (it has opened a connection to our stub +redis), and asserts the backend saw abort(reason='drain') and the process +exited 0. + +NOTE: this is the slow file in the suite (a few seconds -- it spawns a process +and waits for real HTTP round trips). It needs no docker daemon and no extra +dependencies. +""" + +import json +import os +import signal +import socket +import socketserver +import subprocess +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +RUNNER_ID = 'rn_test' +JOB_ID = 'job_sigterm' +SUBMISSION_ID = 'sub_sigterm' + +# Bounded everywhere: a regression must fail this test, never hang the suite. +WAIT_TIMEOUT = 20.0 +# Generous headroom over the ~1s the drain actually needs, but far below any +# hang. The drain's own budget (DRAIN_TIMEOUT_SEC) is 540s, so a real deadlock +# would blow past this instead of quietly passing. +EXIT_TIMEOUT = 30.0 + +TASKS = [{ + 'taskScore': 100, + 'memoryLimit': 65536, + 'timeLimit': 1000, + 'caseCount': 1, +}] +# language 2 is python3: the dispatcher enqueues no Compile job for it, so this +# would be exactly one Execute task if it ever got dispatched (it must not). +JOB_PAYLOAD = { + 'job_id': JOB_ID, + 'submission_id': SUBMISSION_ID, + 'problem_id': 42, + 'language': 2, + 'code_url': 'http://127.0.0.1:1/code.zip', + 'checker': None, + 'tasks': TASKS, +} + + +class StubBackend: + """The five runner-API endpoints main.py calls, recorded thread-safely. + + ``next_job_served`` is the test's trigger: SIGTERM has to land while the + claim is still in prep, and the only observable edge for that is the + backend handing out the job. + """ + + def __init__(self): + self.lock = threading.Lock() + self.requests = [] + self.job_handed_out = False + self.next_job_served = threading.Event() + + def record(self, path, body): + with self.lock: + self.requests.append((path, body)) + + def paths(self): + with self.lock: + return [path for path, _ in self.requests] + + def find(self, suffix): + with self.lock: + return [(p, b) for p, b in self.requests if p.endswith(suffix)] + + def take_job(self): + with self.lock: + if self.job_handed_out: + return None + self.job_handed_out = True + return JOB_PAYLOAD + + +class QuietHTTPServer(ThreadingHTTPServer): + """ThreadingHTTPServer without HTTPServer's reverse-DNS bind. + + ``HTTPServer.server_bind`` resolves the bound address with + ``socket.getfqdn``, which blocks for ~35s on a macOS box with mDNS in the + resolver chain -- all of it before the test body even starts. The FQDN is + only used to fill in ``server_name``, which nothing here reads. + """ + daemon_threads = True + + def server_bind(self): + socketserver.TCPServer.server_bind(self) + self.server_name, self.server_port = self.server_address[:2] + + +def make_handler(state): + + class Handler(BaseHTTPRequestHandler): + protocol_version = 'HTTP/1.1' + + def log_message(self, *args): + pass # keep pytest output readable + + def _body(self): + length = int(self.headers.get('Content-Length') or 0) + if not length: + return None + raw = self.rfile.read(length) + try: + return json.loads(raw) + except ValueError: + return raw.decode('utf-8', 'replace') + + def _respond(self, status, payload=None): + self.send_response(status) + if payload is None: + self.send_header('Content-Length', '0') + self.end_headers() + return + encoded = json.dumps(payload).encode() + self.send_header('Content-Type', 'application/json') + self.send_header('Content-Length', str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def do_POST(self): + state.record(self.path, self._body()) + if self.path == '/runners/register': + self._respond( + 201, { + 'runner_id': RUNNER_ID, + 'token': 'tok', + 'config': { + 'heartbeat_interval_sec': 60, + 'poll_interval_sec': 0.05, + 'max_concurrent_jobs': 8, + }, + }) + elif self.path == f'/runners/{RUNNER_ID}/heartbeat': + self._respond(204) + else: + self._respond(404) + + def do_GET(self): + state.record(self.path, None) + if self.path == f'/runners/{RUNNER_ID}/next-job': + payload = state.take_job() + if payload is None: + self._respond(204) + else: + self._respond(200, payload) + state.next_job_served.set() + else: + self._respond(404) + + def do_PUT(self): + state.record(self.path, self._body()) + if self.path.endswith('/complete'): + self._respond(204) + elif self.path.endswith('/abort'): + self._respond(202) + else: + self._respond(404) + + return Handler + + +class StubRedis: + """A socket that accepts connections and then says nothing. + + This is the test's synchronisation point, and it has to be this rather + than "the backend served next-job": that only proves the response was + written, not that the child received it, let alone that it entered + ``_prepare``. SIGTERM arriving a moment too early would be handled by the + prep loop's very first stop check -- still an abort(drain), so the test + would pass while proving something weaker than it claims. + + ``prepare_job`` starts with ``ensure_testdata()``, which talks to redis, so + a connection here means the child is demonstrably inside prep. Holding the + connection open (never answering) keeps it parked there until the test has + sent the signal and drops it. + """ + + def __init__(self): + self._sock = socket.socket() + self._sock.bind(('127.0.0.1', 0)) + self._sock.listen(8) + self.port = self._sock.getsockname()[1] + self.connected = threading.Event() + self._lock = threading.Lock() + self._conns = [] + self._closed = False + + def serve(self): + while True: + try: + conn, _ = self._sock.accept() + except OSError: + return # listening socket closed during teardown + with self._lock: + if self._closed: + conn.close() + return + self._conns.append(conn) + self.connected.set() + + def drop_connections(self): + """Fail the child's in-flight redis call now. + + Without this it would sit there until its socket timeout, turning a + 1.6s test into an 11s one. + """ + with self._lock: + conns, self._conns = self._conns, [] + for conn in conns: + conn.close() + + def close(self): + with self._lock: + self._closed = True + self.drop_connections() + self._sock.close() + + +@pytest.fixture +def stub_redis(): + stub = StubRedis() + threading.Thread(target=stub.serve, daemon=True).start() + try: + yield stub + finally: + stub.close() + + +@pytest.fixture +def stub_backend(): + state = StubBackend() + server = QuietHTTPServer(('127.0.0.1', 0), make_handler(state)) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + state.url = f'http://127.0.0.1:{server.server_address[1]}' + try: + yield state + finally: + server.shutdown() + server.server_close() + thread.join(timeout=WAIT_TIMEOUT) + + +def test_real_sigterm_drains_and_exits(tmp_path, stub_backend, stub_redis): + # docker stop on a runner whose claim is still in prep: the claim must come + # back as an attempt-neutral 'drain' abort (never 'prep_failed', which + # would push otherwise healthy submissions towards JE on a rolling + # restart), nothing may be reported complete, and the process must exit 0 + # on its own rather than wait for SIGKILL. + env = { + **os.environ, + 'BACKEND_URL': stub_backend.url, + 'RUNNER_REGISTRATION_TOKEN': 'reg-token', + 'RUNNER_NAME': 'sigterm-test-runner', + # Answers the TCP handshake and nothing else, so prep parks here -- + # see StubRedis. + 'REDIS_URL': f'redis://127.0.0.1:{stub_redis.port}/0', + # Keep every directory the runner creates or writes inside tmp_path so + # the test never touches the working tree. + 'SUBMISSION_DIR': str(tmp_path / 'submissions'), + 'SUBMISSION_BACKUP_DIR': str(tmp_path / 'submissions.bk'), + 'TESTDATA_ROOT': str(tmp_path / 'testdata'), + } + out_path = tmp_path / 'runner.out' + err_path = tmp_path / 'runner.err' + + def output(): + return (f'\n--- runner stdout ---\n{out_path.read_text()}' + f'\n--- runner stderr ---\n{err_path.read_text()}') + + with out_path.open('wb') as out, err_path.open('wb') as err: + proc = subprocess.Popen( + [sys.executable, 'main.py'], + cwd=str(REPO_ROOT), + env=env, + stdout=out, + stderr=err, + ) + try: + wait_for( + proc, + lambda: '/runners/register' in stub_backend.paths(), + 'the runner to register', + output, + ) + wait_for( + proc, + stub_backend.next_job_served.is_set, + 'the runner to claim a job', + output, + ) + # The claim is not enough: wait until the child is provably inside + # prep (it has opened a redis connection, which only ensure_testdata + # does) before signalling. + wait_for( + proc, + stub_redis.connected.is_set, + 'the runner to reach redis from inside prep', + output, + ) + + # The job is now claimed, prep is running, and nothing has been + # dispatched -- the exact state whose drain behaviour is at stake. + proc.send_signal(signal.SIGTERM) + # Let the parked redis call fail immediately so prep unwinds now + # rather than after its socket timeout. + stub_redis.drop_connections() + + try: + returncode = proc.wait(timeout=EXIT_TIMEOUT) + except subprocess.TimeoutExpired: + pytest.fail(f'runner did not exit within {EXIT_TIMEOUT}s of ' + f'SIGTERM{output()}') + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout=WAIT_TIMEOUT) + + assert returncode == 0, f'runner exited {returncode}{output()}' + + aborts = stub_backend.find('/abort') + assert len(aborts) == 1, f'expected exactly one abort, got {aborts}' + path, body = aborts[0] + assert path == f'/runners/{RUNNER_ID}/jobs/{JOB_ID}/abort' + # The substantive claim: interrupted-by-shutdown, not failed. + assert body == {'reason': 'drain'} + assert stub_backend.find('/complete') == [] + + +def wait_for(proc, predicate, message, output, timeout=WAIT_TIMEOUT): + """Bounded poll that also fails fast (with output) if the runner died. + + Without the proc.poll() check, a registration failure or an import error + would show up as an opaque timeout instead of the actual traceback. + """ + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + if proc.poll() is not None: + pytest.fail(f'runner exited early ({proc.returncode}) while ' + f'waiting for {message}{output()}') + time.sleep(0.02) + pytest.fail(f'timed out after {timeout}s waiting for {message}{output()}') diff --git a/tests/test_testdata_timeouts.py b/tests/test_testdata_timeouts.py new file mode 100644 index 0000000..80489e4 --- /dev/null +++ b/tests/test_testdata_timeouts.py @@ -0,0 +1,86 @@ +"""Every prep-path network call must be bounded (drain safety). + +Unbounded HTTP or redis-lock waits inside ensure_testdata() would keep a +poller iteration alive past stop(), so a SIGTERM drain would hang until +docker SIGKILLs the runner and in-flight results are lost. +""" + +import pytest + +from dispatcher import testdata + + +class StubResponse: + + def __init__(self, payload=None, content=b''): + self.status_code = 200 + self.ok = True + self.content = content + self._payload = payload + + def json(self): + return self._payload + + +class GetRecorder: + + def __init__(self, response): + self.calls = [] + self._response = response + + def __call__(self, url, **kwargs): + self.calls.append((url, kwargs)) + return self._response + + +def test_fetch_problem_meta_is_bounded(tmp_path, monkeypatch): + recorder = GetRecorder(StubResponse(payload={'data': {'tasks': []}})) + monkeypatch.setattr(testdata.rq, 'get', recorder) + monkeypatch.setattr(testdata, 'META_DIR', tmp_path) + + testdata.fetch_problem_meta(42) + + assert recorder.calls[0][1]['timeout'] == testdata.HTTP_TIMEOUT + + +def test_fetch_testdata_is_bounded(monkeypatch): + recorder = GetRecorder(StubResponse(content=b'zip')) + monkeypatch.setattr(testdata.rq, 'get', recorder) + + assert testdata.fetch_testdata(42) == b'zip' + assert recorder.calls[0][1]['timeout'] == testdata.HTTP_TIMEOUT + + +def test_get_checksum_is_bounded(monkeypatch): + recorder = GetRecorder(StubResponse(payload={'data': 'abc'})) + monkeypatch.setattr(testdata.rq, 'get', recorder) + + assert testdata.get_checksum(42) == 'abc' + assert recorder.calls[0][1]['timeout'] == testdata.HTTP_TIMEOUT + + +class LockSentinel(Exception): + pass + + +class FakeRedis: + """Records lock() kwargs, then short-circuits the rest of ensure_testdata.""" + + def __init__(self): + self.lock_kwargs = None + + def lock(self, key, **kwargs): + self.lock_kwargs = kwargs + raise LockSentinel() + + +def test_ensure_testdata_lock_has_blocking_timeout(monkeypatch): + fake = FakeRedis() + monkeypatch.setattr(testdata, 'get_redis_client', lambda: fake) + + with pytest.raises(LockSentinel): + testdata.ensure_testdata(42) + + assert fake.lock_kwargs['timeout'] == 60 + assert fake.lock_kwargs[ + 'blocking_timeout'] == testdata.LOCK_BLOCKING_TIMEOUT