control.py

#!/usr/bin/env python3

import json
import logging
import os
import subprocess
import threading
import time

from flask import Flask, jsonify

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("concurrency")

POLL_INTERVAL = int(os.environ.get("POLL_INTERVAL", "30"))
NAMESPACE = os.environ.get("NAMESPACE", "")
MACHINE_LABELS = ("builder", "tester")
PORT = int(os.environ.get("PORT", "8080"))

state = {
    "busy_machines": [],
    "pending": [],
    "running": [],
    "last_reconcile": None,
}

app = Flask(__name__)


def oc(*args):
    cmd = ["oc"]
    if NAMESPACE:
        cmd += ["-n", NAMESPACE]
    cmd += list(args)
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        log.error("oc %s failed: %s", " ".join(args), result.stderr.strip())
        return None
    return result.stdout


def get_pipelineruns():
    out = oc("get", "pipelinerun", "-o", "json")
    if not out:
        return []
    return json.loads(out).get("items", [])


def is_pending(pr):
    return pr.get("spec", {}).get("status", "") == "PipelineRunPending"


def is_running(pr):
    spec_status = pr.get("spec", {}).get("status", "")
    if spec_status == "PipelineRunPending":
        return False
    conditions = pr.get("status", {}).get("conditions", [])
    for c in conditions:
        if c.get("type") == "Succeeded":
            status = c.get("status", "")
            reason = c.get("reason", "")
            if status == "Unknown" and reason in ("Running", "Started", "Pending"):
                return True
    return False


def get_machine_labels(pr):
    labels = pr.get("metadata", {}).get("labels", {})
    machines = set()
    for key in MACHINE_LABELS:
        val = labels.get(key, "")
        if val:
            machines.add(val)
    return machines


def get_busy_machines(pipelineruns):
    busy = set()
    for pr in pipelineruns:
        if is_running(pr):
            busy.update(get_machine_labels(pr))
    return busy


def start_pipelinerun(name):
    patch = json.dumps([{"op": "remove", "path": "/spec/status"}])
    out = oc(
        "patch", "pipelinerun", name,
        "--type=json",
        f"-p={patch}",
    )
    if out is not None:
        log.info("Started PipelineRun %s", name)
        return True
    log.error("Failed to start PipelineRun %s", name)
    return False


def pr_name(pr):
    return pr["metadata"]["name"]


def reconcile():
    pipelineruns = get_pipelineruns()
    if not pipelineruns:
        return

    busy = get_busy_machines(pipelineruns)
    pending = [pr for pr in pipelineruns if is_pending(pr)]
    running = [pr for pr in pipelineruns if is_running(pr)]

    state["busy_machines"] = sorted(busy)
    state["pending"] = [pr_name(pr) for pr in pending]
    state["running"] = [pr_name(pr) for pr in running]
    state["last_reconcile"] = time.strftime("%Y-%m-%d %H:%M:%S")

    if busy:
        log.info("Busy machines: %s", ", ".join(sorted(busy)))

    if not pending:
        log.debug("No pending PipelineRuns")
        return

    log.info("Pending PipelineRuns: %d", len(pending))

    for pr in pending:
        name = pr_name(pr)
        required = get_machine_labels(pr)
        if not required:
            log.info("%s has no machine labels, starting it", name)
            start_pipelinerun(name)
            continue

        conflict = required & busy
        if conflict:
            log.info(
                "%s waiting on busy machine(s): %s",
                name, ", ".join(sorted(conflict)),
            )
            continue

        start_pipelinerun(name)
        busy.update(required)


def controller_loop():
    log.info(
        "Concurrency controller starting (poll=%ds, namespace=%s)",
        POLL_INTERVAL, NAMESPACE or "<current>",
    )
    while True:
        try:
            reconcile()
        except Exception:
            log.exception("Error during reconcile")
        time.sleep(POLL_INTERVAL)


@app.route("/healthz")
def healthz():
    return jsonify({"status": "ok"})


@app.route("/status")
def status():
    return jsonify(state)


if __name__ == "__main__":
    t = threading.Thread(target=controller_loop, daemon=True)
    t.start()
    app.run(host="0.0.0.0", port=PORT)