Fix/reload jobs block event loop

python/apsis/jobs.py

@@ -4,16 +4,15 @@
 import os
 from pathlib import Path
 import random
+import re
 import string
 import yaml
 from ruamel.yaml import YAML
-from ruamel.yaml.constructor import DuplicateKeyError
 
 from .actions import Action
 from .actions.schedule import successor_from_jso
 from .cond import Condition
 from .exc import JobError, JobsDirErrors, SchemaError
-from .lib import itr
 from .lib.json import to_array, to_narray, check_schema
 from .lib.py import tupleize, format_ctor
 from .program import Program, NoOpProgram
@@ -155,9 +154,87 @@ def dump_yaml(file, job):
     YAML().dump(job_to_jso(job), file)
 
 
+class DuplicateKeyError(Exception):
+    """A YAML mapping contains a duplicate key."""
+
+
+class _DupCheckSafeLoader(yaml.CSafeLoader):
+    """
+    Fast libyaml loader that matches ruamel: it rejects duplicate keys and
+    resolves scalars per the YAML 1.2 core schema, whereas PyYAML defaults to
+    YAML 1.1 (e.g. `12:00:00` -> int, `NO` -> bool).
+
+    Constructors and YAML 1.2 resolvers are wired up in `_build_yaml_loader`.
+    """
+
+    def construct_mapping(self, node, deep=False):
+        # Detect duplicates among the explicit keys before expanding `<<`
+        # merges, so an explicit key that overrides a merged one isn't itself
+        # flagged as a duplicate.
+        seen = set()
+        for key_node, value_node in node.value:
+            if key_node.tag == "tag:yaml.org,2002:merge":
+                continue
+            key = self.construct_object(key_node, deep=deep)
+            if key in seen:
+                raise DuplicateKeyError(
+                    f'found duplicate key "{key}" with value "{value_node.value}"'
+                )
+            seen.add(key)
+        self.flatten_mapping(node)
+        return {
+            self.construct_object(k, deep=deep): self.construct_object(v, deep=deep)
+            for k, v in node.value
+        }
+
+    def construct_int(self, node):
+        # YAML 1.2 treats a leading zero as decimal and `0x`/`0o` as hex/octal;
+        # PyYAML's inherited constructor would read a leading zero as octal.
+        value = self.construct_scalar(node)
+        sign = -1 if value.startswith("-") else 1
+        digits = value[1:] if value[0] in "+-" else value
+        base = 16 if digits[:2] in ("0x", "0X") else 8 if digits[:2] in ("0o", "0O") else 10
+        return sign * int(digits, base)
+
+
+def _build_yaml_loader():
+    """
+    Build the job-file YAML loader: `_DupCheckSafeLoader` with its constructors
+    registered and PyYAML's YAML 1.1 scalar resolvers replaced by the YAML 1.2
+    core schema.  Each resolver is (tag, pattern, first-characters), where the
+    last is a PyYAML lookup hint listing the chars a match may start with.
+    """
+    resolvers = (
+        ("tag:yaml.org,2002:bool", r"^(?:true|True|TRUE|false|False|FALSE)$", "tTfF"),
+        ("tag:yaml.org,2002:int", r"^[-+]?(?:[0-9]+|0o[0-7]+|0x[0-9a-fA-F]+)$", "-+0123456789"),
+        (
+            "tag:yaml.org,2002:float",
+            r"^(?:[-+]?(?:\.[0-9]+|[0-9]+(?:\.[0-9]*)?)(?:[eE][-+]?[0-9]+)?"
+            r"|[-+]?\.(?:inf|Inf|INF)|\.(?:nan|NaN|NAN))$",
+            "-+0123456789.",
+        ),
+        # "" matches an empty scalar (`key:`) so it resolves to null, not "".
+        ("tag:yaml.org,2002:null", r"^(?:~|null|Null|NULL|)$", ["~", "n", "N", ""]),
+        ("tag:yaml.org,2002:merge", r"^(?:<<)$", "<"),
+    )
+
+    loader = _DupCheckSafeLoader
+    loader.add_constructor("tag:yaml.org,2002:map", loader.construct_mapping)
+    loader.add_constructor("tag:yaml.org,2002:int", loader.construct_int)
+    loader.yaml_implicit_resolvers = {}
+    for tag, pattern, first in resolvers:
+        loader.add_implicit_resolver(tag, re.compile(pattern), list(first))
+    return loader
+
+
+DupCheckSafeLoader = _build_yaml_loader()
+
+
 def list_yaml_files(dir_path):
     dir_path = Path(dir_path)
-    for dir, _, names in os.walk(dir_path):
+    for dir, dirs, names in os.walk(dir_path):
+        # Don't go into hidden dirs (e.g. `.git`)
+        dirs[:] = [d for d in dirs if not d.startswith(".")]
         dir = Path(dir)
         paths = (dir / n for n in names if not n.startswith("."))
         paths = (p for p in paths if p.suffix == ".yaml")
@@ -221,13 +298,13 @@ def get_jobs(self, *, ad_hoc=None):
         return jobs
 
 
-async def load_jobs_dir(path, yaml_loader=None):
+async def load_jobs_dir(path, yaml_loader=DupCheckSafeLoader):
     """
     Attempts to loads jobs from a jobs dir.
 
     :param yaml_loader:
-      An optional PyYAML loader class (e.g. ``yaml.CSafeLoader``) to use
-      instead of the default ruamel YAML loader.
+      The PyYAML loader class used to parse each job file.  Defaults to
+      `DupCheckSafeLoader` (libyaml-backed, fast, rejects duplicate keys).
     :return:
       The successfully loaded `JobsDir`.
     :raise NotADirectoryError:
@@ -252,32 +329,27 @@ async def load_job(path, job_id):
                 content = await file.read()
 
             def _parse():
-                if yaml_loader is not None:
-                    job_jso = yaml.load(content, Loader=yaml_loader)
-                else:
-                    job_jso = YAML().load(content)
+                job_jso = yaml.load(content, Loader=yaml_loader)
                 return Job.from_jso(job_jso, job_id)
 
             job = await asyncio.to_thread(_parse)
             return job_id, job, None
-        except DuplicateKeyError as exc:
-            err_msg = exc.problem if exc.problem else str(exc)
-            schema_err = SchemaError(err_msg)
+        except (DuplicateKeyError, yaml.YAMLError) as exc:
+            schema_err = SchemaError(str(exc))
             schema_err.job_id = job_id
             return job_id, None, schema_err
         except SchemaError as exc:
             log.debug(f"error: {path}: {exc}", exc_info=True)
             exc.job_id = job_id
             return job_id, None, exc
 
-    load_coros = [load_job(path, job_id) for path, job_id in list_yaml_files(jobs_path)]
-    for chunk in itr.chunks(load_coros, 100):
-        results = await asyncio.gather(*chunk)
-        for job_id, job, exc in results:
-            if job is not None:
-                jobs[job_id] = job
-            if exc is not None:
-                errors.append(exc)
+    # Load one file at a time. `load_job` parses in a thread, which yields to the loop, so a reload doesn't block it.
+    for path, job_id in list_yaml_files(jobs_path):
+        _, job, exc = await load_job(path, job_id)
+        if job is not None:
+            jobs[job_id] = job
+        if exc is not None:
+            errors.append(exc)
 
     jobs_dir = JobsDir(jobs_path, jobs)
 

test/unit/test_load_jobs_yield.py

@@ -0,0 +1,171 @@
+"""Tests that `load_jobs_dir` doesn't block the event loop while loading."""
+
+import asyncio
+import os
+import time
+
+import pytest
+import yaml
+
+import apsis.exc
+import apsis.jobs
+from apsis.jobs import DupCheckSafeLoader, DuplicateKeyError
+
+# -------------------------------------------------------------------------------
+
+
+def _write_job(path):
+    path.write_text("params: []\nprogram:\n  type: no-op\n")
+
+
+def _make_jobs_tree(root, n_dirs, per_dir=1):
+    for d in range(n_dirs):
+        sub = root / f"d{d}"
+        sub.mkdir()
+        for f in range(per_dir):
+            _write_job(sub / f"job{f}.yaml")
+
+
+@pytest.mark.asyncio
+async def test_load_jobs_dir_interleaves_slow_walk(tmp_path, monkeypatch):
+    """A slow walk shouldn't block the loop for its whole duration."""
+    # Lots of dirs so the cooperative stall (~1 dir) is well under the threshold.
+    n_dirs = 20
+    _make_jobs_tree(tmp_path, n_dirs)
+
+    delay = 0.02
+    real_walk = os.walk
+
+    def slow_walk(*args, **kwargs):
+        # Simulate slow NFS metadata reads, one blocking step per directory.
+        for entry in real_walk(*args, **kwargs):
+            time.sleep(delay)
+            yield entry
+
+    monkeypatch.setattr(apsis.jobs.os, "walk", slow_walk)
+
+    gaps = []
+    stop = False
+
+    # Ticker records the gap between its wakeups; a blocked loop -> big gap.
+    async def ticker():
+        last = time.perf_counter()
+        while not stop:
+            await asyncio.sleep(0)
+            now = time.perf_counter()
+            gaps.append(now - last)
+            last = now
+
+    task = asyncio.create_task(ticker())
+    await asyncio.sleep(0.01)  # let the ticker establish a baseline
+    jobs_dir = await apsis.jobs.load_jobs_dir(tmp_path)
+    stop = True
+    await task
+
+    assert len(list(jobs_dir.get_jobs())) == n_dirs
+    # Eager loading would block for the whole walk; cooperative bounds the stall
+    # to a file-bearing dir plus any yaml-less dirs preceding it (here, 1 each).
+    total_walk = delay * (n_dirs + 1)
+    assert max(gaps) < total_walk / 2, (
+        f"event loop blocked {max(gaps) * 1000:.0f}ms; whole walk is {total_walk * 1000:.0f}ms"
+    )
+
+
+@pytest.mark.asyncio
+async def test_load_jobs_dir_loads_all_jobs(tmp_path):
+    """The refactored loop still loads every job, with correct job ids."""
+    _make_jobs_tree(tmp_path, n_dirs=3, per_dir=4)
+    jobs_dir = await apsis.jobs.load_jobs_dir(tmp_path)
+    job_ids = {j.job_id for j in jobs_dir.get_jobs()}
+    assert job_ids == {f"d{d}/job{f}" for d in range(3) for f in range(4)}
+
+
+@pytest.mark.asyncio
+async def test_load_jobs_dir_skips_hidden_dirs(tmp_path):
+    """Hidden dirs (e.g. `.git`) are not descended into."""
+    _write_job(tmp_path / "real.yaml")
+    git = tmp_path / ".git" / "objects"
+    git.mkdir(parents=True)
+    _write_job(git / "nope.yaml")
+    jobs_dir = await apsis.jobs.load_jobs_dir(tmp_path)
+    assert {j.job_id for j in jobs_dir.get_jobs()} == {"real"}
+
+
+@pytest.mark.asyncio
+async def test_load_jobs_dir_reports_bad_yaml(tmp_path):
+    """A malformed job still surfaces as an error through the sequential loop."""
+    _write_job(tmp_path / "good.yaml")
+    # Valid YAML, but missing the required `program` -> collected SchemaError.
+    (tmp_path / "bad.yaml").write_text("params: [x]\n")
+    # Malformed YAML syntax -> collected, not raised out of the whole reload.
+    (tmp_path / "broken.yaml").write_text("params: [x\nprogram\n")
+    with pytest.raises(apsis.exc.JobsDirErrors) as exc_info:
+        await apsis.jobs.load_jobs_dir(tmp_path)
+    assert {e.job_id for e in exc_info.value.errors} == {"bad", "broken"}
+
+
+# -------------------------------------------------------------------------------
+# DupCheckSafeLoader
+
+
+def test_dup_check_loader_rejects_duplicate_keys():
+    with pytest.raises(DuplicateKeyError):
+        yaml.load("command: one\ncommand: two\n", Loader=DupCheckSafeLoader)
+
+
+def test_dup_check_loader_parses_normal_mapping():
+    assert yaml.load("a: 1\nb: two\n", Loader=DupCheckSafeLoader) == {"a": 1, "b": "two"}
+
+
+def test_dup_check_loader_empty_scalar_is_null():
+    assert yaml.load("a:\n", Loader=DupCheckSafeLoader) == {"a": None}
+
+
+@pytest.mark.parametrize(
+    "text, expected",
+    [
+        # YAML 1.1 would parse these as an int (43200) and a bool; the YAML 1.2
+        # core schema (like ruamel) keeps them as strings.
+        ("12:00:00", "12:00:00"),
+        ("NO", "NO"),
+        ("no", "no"),
+        ("on", "on"),
+        ("off", "off"),
+        ("yes", "yes"),
+        # ...while these still resolve as scalars.
+        ("true", True),
+        ("false", False),
+        ("null", None),
+        ("42", 42),
+        ("1.5", 1.5),
+        # Leading-zero ints are decimal (YAML 1.2), not YAML 1.1 octal; `09`
+        # isn't even valid octal and would raise under the inherited constructor.
+        ("010", 10),
+        ("0123", 123),
+        ("09", 9),
+        ("0x1A", 26),
+        ("0o17", 15),
+    ],
+)
+def test_dup_check_loader_uses_yaml_1_2_scalars(text, expected):
+    result = yaml.load(f"x: {text}\n", Loader=DupCheckSafeLoader)["x"]
+    assert result == expected
+    assert type(result) is type(expected)
+
+
+def test_dup_check_loader_honors_merge_keys():
+    doc = "base: &b {a: 1, b: 2}\nchild:\n  <<: *b\n  c: 3\n"
+    assert yaml.load(doc, Loader=DupCheckSafeLoader)["child"] == {"a": 1, "b": 2, "c": 3}
+
+
+def test_dup_check_loader_merge_override_is_not_a_duplicate():
+    # An explicit key overriding one from `<<` must win, not raise.
+    doc = "base: &b {a: 1, p: 9}\nchild:\n  <<: *b\n  p: 3\n"
+    assert yaml.load(doc, Loader=DupCheckSafeLoader)["child"] == {"a": 1, "p": 3}
+
+
+def test_dup_check_loader_still_rejects_real_duplicate_with_merge():
+    # A genuine duplicate among explicit keys still raises, merge present or not.
+    doc = "base: &b {a: 1}\nchild:\n  <<: *b\n  p: 1\n  p: 2\n"
+    with pytest.raises(DuplicateKeyError):
+        yaml.load(doc, Loader=DupCheckSafeLoader)