cms-SiteForge/app/worker.py

'''Sole build executor: lease-based queue, isolated process, cancellation, timeout, retries.'''
import json
import os
import signal
import subprocess
import sys
import time
import uuid
from . import db, enterprise, builder, jobs, worker_health

running=True
WORKER_ID=uuid.uuid4().hex[:12]

def heartbeat():
    """Container health heartbeat, independent of the HTTP server.

    This single-worker probe is deliberately not used as a multi-worker
    membership system; multi-worker orchestration needs separate leases.
    """
    path=worker_health.heartbeat_path()
    path.parent.mkdir(parents=True,exist_ok=True)
    temp=path.with_suffix(f'.{WORKER_ID}.tmp')
    temp.write_text(WORKER_ID,encoding='ascii')
    os.replace(temp,path)


def stop(*_):
    global running
    running=False

def claim():
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        query='''SELECT id FROM build_jobs AS j WHERE j.state IN ('queued','retrying')
            AND (j.available_at IS NULL OR j.available_at <= ?)
            AND NOT EXISTS (SELECT 1 FROM build_jobs AS other WHERE other.site_id=j.site_id AND other.state='running')
            ORDER BY j.created_at,j.id LIMIT 1'''
        if db.POSTGRES:query+=' FOR UPDATE SKIP LOCKED'
        item=c.execute(query,(db.now(),)).fetchone()
        if not item:return None
        c.execute('''UPDATE build_jobs SET state='running',attempts=attempts+1,
          started_at=?,worker_id=?,lease_until=?,heartbeat_at=? WHERE id=?''',
          (db.now(),WORKER_ID,jobs.iso_after(45),db.now(),item['id']))
        return dict(c.execute('SELECT * FROM build_jobs WHERE id=?',(item['id'],)).fetchone())

def complete(job, value=None, error=None, retryable=False, cancelled=False, timedout=False):
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        current=c.execute('SELECT * FROM build_jobs WHERE id=?',(job['id'],)).fetchone()
        if not current or current['state']!='running' or current['worker_id']!=WORKER_ID:
            return
        # A networked release may already have changed the remote pointer.
        # Cancellation/timeout/crash is UNKNOWN remote outcome, not safe retry.
        external=current['kind'] in ('remote_deploy','remote_rollback')
        if cancelled or current['cancel_requested']:
            if external:
                state='failed'
                error='Release interrupted; reconcile target before manual retry'
            else:state='cancelled'
        elif error is None:
            state='done'
        elif retryable and not external and current['attempts']<current['max_attempts']:
            state='retrying'
        else:state='failed'
        delay=min(60,2**current['attempts']) if state=='retrying' else 0
        c.execute('''UPDATE build_jobs SET state=?,result_json=?,error=?,finished_at=?,
                 available_at=?,worker_id=NULL,lease_until=NULL WHERE id=?''',
                 (state,json.dumps(value or {},ensure_ascii=False),error or '',
                  db.now() if state not in ('retrying',) else None,
                  jobs.iso_after(delay) if delay else db.now(),job['id']))

def execute_subprocess(job):
    env=os.environ.copy()
    env['SITEFORGE_DATA_DIR']=str(db.DATA)
    env['SITEFORGE_DB_FILE']=str(db.DB_FILE)
    output=db.DATA/'job_results'/f"{job['id']}.json"
    output.parent.mkdir(parents=True,exist_ok=True)
    output.unlink(missing_ok=True)
    proc=subprocess.Popen([sys.executable,'-m','app.job_runner',job['id']],
        stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,env=env)
    began=time.monotonic()
    interval=float(os.getenv('SITEFORGE_WORKER_POLL_SECONDS','0.5'))
    while proc.poll() is None:
        heartbeat()
        row=db.one('SELECT cancel_requested FROM build_jobs WHERE id=?',(job['id'],))
        if not running or (row and row['cancel_requested']):
            proc.terminate()
            try:proc.wait(timeout=3)
            except subprocess.TimeoutExpired:proc.kill();proc.wait()
            return None,'Task cancelled',False,True
        if time.monotonic()-began>job['timeout_seconds']:
            proc.terminate()
            try:proc.wait(timeout=3)
            except subprocess.TimeoutExpired:proc.kill();proc.wait()
            return None,'Build timed out',True,False
        db.execute('UPDATE build_jobs SET lease_until=?,heartbeat_at=? WHERE id=? AND state=? AND worker_id=?',
            (jobs.iso_after(45),db.now(),job['id'],'running',WORKER_ID))
        time.sleep(interval)
    if output.exists():
        try:
            data=json.loads(output.read_text(encoding='utf-8'))
            if data.get('ok'):return data['result'],None,False,False
            return None,data.get('error','Build failed'),bool(data.get('retryable',False)),False
        except Exception:return None,'Invalid job result',True,False
    return None,f'Build process exited with code {proc.returncode}',True,False

def run_once(executor=None):
    job=claim()
    if job is None:return None
    try:
        if executor is None:
            value,error,retryable,cancelled=execute_subprocess(job)
        else:
            # In-process executor is explicitly for deterministic unit tests only.
            try:value=executor(job['site_id'],json.loads(job['payload_json']));error=None;retryable=False;cancelled=False
            except Exception as exc:value=None;error=str(exc);retryable=isinstance(exc,(TimeoutError,ConnectionError));cancelled=False
        complete(job,value,error,retryable,cancelled)
    except Exception as exc:
        complete(job,error=str(exc),retryable=True)
    return job['id']

def main():
    db.init();enterprise.setup()
    signal.signal(signal.SIGTERM,stop);signal.signal(signal.SIGINT,stop)
    while running:
        heartbeat()
        jobs.recover_expired()
        if not run_once():
            from . import redis_signal
            redis_signal.wait(float(os.getenv('SITEFORGE_JOB_POLL_SECONDS','1.5')))

if __name__=='__main__':main()