Bounding concurrency in a tree of tasks without deadlock

"""Run the lending model on a deep tree and a wide tree; check the bounds."""
import asyncio
import sys
import time

from lending import Pool, Run

live: set = set()                      # admitted, not finished
peak = {"in_flight": 0, "executing": 0, "leaves": 0}
leaves_running = 0


def check(pool: Pool, cap: int) -> None:
    executing = sum(1 for r in live if not r.lent)   # holds a token not lent out
    # A token granted to a waiter that has not resumed yet is in use but not
    # in `live`, so the count can be lower than the tokens taken, never higher.
    assert executing <= cap - pool.free, (executing, cap, pool.free)
    peak["in_flight"] = max(peak["in_flight"], len(live))
    peak["executing"] = max(peak["executing"], executing)


async def node(pool, cap, parent, depth, fan_out):
    global leaves_running
    run = Run(pool, parent)
    await run.admit()                  # a mount waits here
    live.add(run)
    check(pool, cap)
    try:
        if depth == 0:
            leaves_running += 1
            peak["leaves"] = max(peak["leaves"], leaves_running)
            await asyncio.sleep(0.001)
            leaves_running -= 1
        else:
            await asyncio.gather(*(node(pool, cap, run, depth - 1, fan_out)
                                   for _ in range(fan_out)))
    finally:
        live.discard(run)
        run.finish()
        check(pool, cap)


async def main(cap, depth, fan_out):
    pool = Pool(cap)
    for k in peak:
        peak[k] = 0
    t = time.perf_counter()
    await asyncio.wait_for(node(pool, cap, None, depth, fan_out), timeout=60)
    nodes = sum(fan_out ** d for d in range(depth + 1))
    print(f"cap {cap:>2}, depth {depth:>2}, fan-out {fan_out:>4}, {nodes:>5} tasks: "
          f"peak executing {peak['executing']}, peak in flight {peak['in_flight']} "
          f"(bound {cap * (depth + 1)}), leaves at once {peak['leaves']}, "
          f"{time.perf_counter() - t:.2f} s")
    assert peak["executing"] <= cap and peak["in_flight"] <= cap * (depth + 1)
    assert pool.free == cap and not live      # every token came back


if __name__ == "__main__":
    for cap, depth, fan_out in [(2, 3, 2), (1, 20, 1), (1, 10, 2), (2, 10, 2), (1, 1, 1000),
                                (32, 1, 3000), (4, 3, 10)]:
        asyncio.run(main(cap, depth, fan_out))

"""A parent that ends while its token is lent (CocoIndex: live components)."""
import asyncio

from lending import Pool, Run


async def main():
    pool = Pool(1)
    parent = Run(pool)
    await parent.admit()                       # takes the only pool token
    child = Run(pool, parent)
    await child.admit()                        # runs on the parent's token
    parent.finish()                            # parent ends; child still runs
    print("after parent ends:", pool.free, "free")   # 0: the loan still holds it
    assert pool.free == 0
    other = Run(pool)
    waiting = asyncio.create_task(other.admit())
    await asyncio.sleep(0)
    print("other run waiting:", not waiting.done())
    child.finish()                             # the loan returns the parent's token
    await waiting
    print("other run admitted:", other.token)
    other.finish()
    print("at the end:", pool.free, "free")
    assert pool.free == 1


asyncio.run(main())

"""One workload, four strategies: a walk mounts FILES file tasks; each file
reads its contents, then mounts CHUNKS chunk tasks that do the slow step."""
import asyncio
import time

from lending import Pool, Run

CAP, FILES, CHUNKS, SIZE = 4, 2000, 2, 100_000
stats = {}


def enter(kind):
    stats[kind] = stats.get(kind, 0) + 1
    stats["peak_" + kind] = max(stats.get("peak_" + kind, 0), stats[kind])


def leave(kind):
    stats[kind] -= 1


# 1. Hold the slot while waiting for children.
async def hold():
    sem = asyncio.Semaphore(CAP)

    async def chunk(text):
        async with sem:
            enter("exec"); await asyncio.sleep(0.01); leave("exec")

    async def file():
        async with sem:
            enter("files"); text = bytes(SIZE)
            await asyncio.gather(*(chunk(text) for _ in range(CHUNKS)))
            leave("files")

    async with sem:
        await asyncio.gather(*(file() for _ in range(FILES)))


# 2. Release the slot before waiting (CocoIndex before 1.0.26: at the first child mount).
async def release():
    sem = asyncio.Semaphore(CAP)

    async def chunk(text):
        async with sem:
            enter("exec"); await asyncio.sleep(0.01); leave("exec")

    async def file():
        await sem.acquire()
        enter("files"); text = bytes(SIZE)
        sem.release()
        await asyncio.gather(*(chunk(text) for _ in range(CHUNKS)))
        leave("files")

    await sem.acquire()
    sem.release()
    await asyncio.gather(*(file() for _ in range(FILES)))


# 3. One semaphore per level of the tree.
async def per_level():
    level = [asyncio.Semaphore(CAP) for _ in range(3)]

    async def chunk(text):
        async with level[2]:
            enter("exec"); await asyncio.sleep(0.01); leave("exec")

    async def file():
        async with level[1]:
            enter("files"); text = bytes(SIZE)
            await asyncio.gather(*(chunk(text) for _ in range(CHUNKS)))
            leave("files")

    async with level[0]:
        await asyncio.gather(*(file() for _ in range(FILES)))


# 4. Lend the parent's slot to one child at a time (lending.py).
async def lend():
    pool = Pool(CAP)

    async def task(parent, body):
        run = Run(pool, parent)
        await run.admit()
        try:
            await body(run)
        finally:
            run.finish()

    async def chunk(run, text):
        enter("exec"); await asyncio.sleep(0.01); leave("exec")

    async def file(run):
        enter("files"); text = bytes(SIZE)
        await asyncio.gather(*(task(run, lambda r: chunk(r, text)) for _ in range(CHUNKS)))
        leave("files")

    async def walk(run):
        await asyncio.gather(*(task(run, file) for _ in range(FILES)))

    await task(None, walk)


async def main():
    for strategy in (hold, release, per_level, lend):
        stats.clear()
        t = time.perf_counter()
        try:
            await asyncio.wait_for(strategy(), timeout=30)
            result = f"{time.perf_counter() - t:5.2f} s"
        except TimeoutError:
            result = "stuck (30 s timeout)"
        files = stats.get("peak_files", 0)
        print(f"{strategy.__name__:>9}: {result:>19}, peak files in flight {files:>4} "
              f"({files * SIZE / 1e6:5.1f} MB), peak chunks running {stats.get('peak_exec', 0)}")


asyncio.run(main())

"""The incident's shape on the real engine. A walk mounts FILES file
components. Each file reads 100 KB, mounts a quick first child, then waits on
an application gate (an asyncio.Semaphore of GATE) around the slow step: its
CHUNKS chunk components. Prints the peak number of file components in flight.

usage: python engine_bench.py <max_inflight_components> <db_dir>
"""
import asyncio
import os
import pathlib
import re
import resource
import subprocess
import sys
import threading
import time

import cocoindex as coco

FILES, CHUNKS, SIZE, CHUNK_SECONDS, GATE = 2000, 2, 100_000, 0.02, 4
gate: asyncio.Semaphore
lock = threading.Lock()
count = {"files": 0, "peak_files": 0, "chunks": 0, "peak_chunks": 0}


def enter(kind):
    with lock:
        count[kind] += 1
        count["peak_" + kind] = max(count["peak_" + kind], count[kind])


def leave(kind):
    with lock:
        count[kind] -= 1


@coco.fn
async def chunk(i: int, j: int) -> None:
    enter("chunks")
    await asyncio.sleep(CHUNK_SECONDS)   # the slow downstream step
    leave("chunks")


@coco.fn
async def metadata(i: int) -> None:
    pass


@coco.fn
async def file(i: int) -> None:
    enter("files")
    text = os.urandom(SIZE)              # the file's contents (random, so memory counts them)
    await coco.use_mount(coco.component_subpath("meta"), metadata, i)
    async with gate:                     # the application's own limit
        await asyncio.gather(*(
            coco.use_mount(coco.component_subpath(f"c{j}"), chunk, i, j)
            for j in range(CHUNKS)
        ))
    del text
    leave("files")


async def walk() -> None:
    global gate
    gate = asyncio.Semaphore(GATE)       # created on the engine's loop
    for i in range(FILES):
        await coco.mount(coco.component_subpath(str(i)), file, i)


def peak_memory_mb() -> float:
    """Peak physical footprint. On macOS ru_maxrss undercounts (memory
    compression), so read phys_footprint_peak from `footprint`."""
    if sys.platform == "darwin":
        out = subprocess.run(["footprint", str(os.getpid())], capture_output=True, text=True).stdout
        m = re.search(r"phys_footprint_peak: ([\d.]+) (KB|MB|GB)", out)
        if m:
            return float(m.group(1)) * {"KB": 1e-3, "MB": 1, "GB": 1e3}[m.group(2)]
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e3  # KB on Linux


if __name__ == "__main__":
    cap, db = int(sys.argv[1]), pathlib.Path(sys.argv[2])
    env = coco.Environment(coco.Settings.from_env(db_path=db))
    app = coco.App(coco.AppConfig(name="bench", environment=env,
                                  max_inflight_components=cap), walk)
    t = time.perf_counter()
    app.update_blocking()
    rss_mb = peak_memory_mb()
    print(f"cocoindex {coco.__version__ if hasattr(coco, '__version__') else '?'}"
          f" cap {cap}: peak {count['peak_files']} file components in flight, "
          f"peak {count['peak_chunks']} chunks running, "
          f"{time.perf_counter() - t:.1f} s, peak memory {rss_mb:.0f} MB")

package main

import (
	"fmt"
	"time"

	"golang.org/x/sync/errgroup"
)

// walk visits a binary tree of the given depth; every node waits for its children.
func walk(g *errgroup.Group, depth int) error {
	if depth == 0 {
		time.Sleep(time.Millisecond)
		return nil
	}
	done := make(chan error, 2)
	for i := 0; i < 2; i++ {
		g.Go(func() error { err := walk(g, depth-1); done <- err; return err })
	}
	<-done
	return <-done
}

func main() {
	var g errgroup.Group
	g.SetLimit(2)
	finished := make(chan error, 1)
	go func() { finished <- walk(&g, 3) }()
	select {
	case <-finished:
		fmt.Println("finished")
	case <-time.After(2 * time.Second):
		fmt.Println("stuck after 2 s: both goroutines the limit admits are blocked in g.Go")
	}
}

package main

import (
	"fmt"
	"time"

	"golang.org/x/sync/errgroup"
)

// walk visits a binary tree of the given depth; every node waits for its children.
func walk(g *errgroup.Group, depth int) error {
	if depth == 0 {
		time.Sleep(time.Millisecond)
		return nil
	}
	done := make(chan error, 2)
	for i := 0; i < 2; i++ {
		f := func() error { err := walk(g, depth-1); done <- err; return err }
		if !g.TryGo(f) {
			f() // no slot free: run the child inline
		}
	}
	<-done
	return <-done
}

func main() {
	var g errgroup.Group
	g.SetLimit(2)
	finished := make(chan error, 1)
	go func() { finished <- walk(&g, 3) }()
	select {
	case <-finished:
		fmt.Println("finished")
	case <-time.After(2 * time.Second):
		fmt.Println("stuck after 2 s: both goroutines the limit admits are blocked in g.Go")
	}
}

"""Record the interactive race in fig/race-data.json.

The same tree (a walk mounts 6 files; each file reads for 1 tick, then waits on
2 chunks of 2 ticks) runs under a cap of 2 with three strategies: hold the slot
(asyncio.Semaphore), release it at the first child (asyncio.Semaphore), and lend it
(../_code/lending.py). A virtual clock makes every run deterministic: sleeps
advance time instantly. A sampler lets every hand-off due at an instant settle,
then records each node's state, which nodes hold a pool token or a loan, every
half tick. The static strips in the post (_diagrams-src/dark/gen.py) call
record() on a 3-file tree. Needs Python 3.12 or later: asyncio.Semaphore wakes
its waiters in a different order on older versions, which changes the frames.

usage: uv run --python 3.12 python race.py   (from any directory)
"""
import asyncio
import json
import pathlib
import selectors
import sys

if sys.version_info < (3, 12):
    sys.exit("race.py needs Python 3.12 or later (asyncio.Semaphore wake-up order)")

HERE = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "_code"))
from lending import Pool, Run  # noqa: E402

CAP, FILES, CHUNKS, READ, CHUNK = 2, 6, 2, 1.0, 2.0
STEP, LIMIT = 0.5, 24.0
EPS = 1e-6                               # sample just after a tick, so events due at the tick count


class _InstantSelector(selectors.SelectSelector):
    def __init__(self, loop_ref):
        super().__init__()
        self.loop_ref = loop_ref

    def select(self, timeout=None):
        if timeout:
            self.loop_ref[0].now += timeout   # jump to the next timer
        return super().select(0)


class VirtualLoop(asyncio.SelectorEventLoop):
    def __init__(self):
        ref = [None]
        super().__init__(_InstantSelector(ref))
        ref[0] = self
        self.now = 0.0

    def time(self):
        return self.now


def names_of(files, chunks=CHUNKS):
    return ["walk"] + [f"f{i}" for i in range(1, files + 1)] + [
        f"f{i}c{j}" for i in range(1, files + 1) for j in range(1, chunks + 1)]


def record(strategy, files=FILES, cap=CAP):
    """Run one strategy and return its frames.

    Each frame: t, s (name -> none, wait, run, hold, lent, mem, done), k (name ->
    the pool token, "s1" or "s2", that the node holds or has lent out; a borrower
    shows its lender's token) and l (names running on a loan; lend only)."""
    names = names_of(files)
    # Long enough for any cap to finish (one task at a time); hold never finishes.
    limit = 6.0 if strategy == "hold" else max(LIMIT, files * (READ + CHUNKS * CHUNK) + 2)
    state = {n: "none" for n in names}
    frames, finished = [], {"t": None}
    owners = {}                          # name -> token id, for nodes holding a pool token

    def tokens(pool_holders, chain_root):
        """Assign stable ids to pool tokens: a new holder gets the lowest free id."""
        for n in list(owners):
            if n not in pool_holders:
                del owners[n]
        for n in pool_holders:
            if n not in owners:
                owners[n] = next(i for i in range(1, cap + 1) if i not in owners.values())
        return {n: f"s{owners[chain_root(n)]}" for n in names if chain_root(n) in owners}

    async def sampler(snapshot, fingerprint):
        loop = asyncio.get_running_loop()
        while True:
            # Let every hand-off due at this instant settle: yield until nothing changes.
            prev = None
            for _ in range(200):
                await asyncio.sleep(0)
                cur = fingerprint()
                if cur == prev:
                    break
                prev = cur
            frames.append({"t": round(loop.time(), 2), **snapshot()})
            if finished["t"] is not None or loop.time() >= limit:
                return
            await asyncio.sleep(STEP)

    async def sampler_after_tick(snapshot, fingerprint):
        await asyncio.sleep(EPS)
        await sampler(snapshot, fingerprint)

    sem = None                           # created on the virtual loop, in main()

    if strategy in ("hold", "release"):
        def fingerprint():
            return (tuple(state.values()), sem._value)

        def snapshot():
            holders = [n for n in names if state[n] in ("run", "hold")]
            return {"s": dict(state), "k": tokens(holders, lambda n: n)}

        async def chunk(name):
            state[name] = "wait"
            async with sem:
                state[name] = "run"
                await asyncio.sleep(CHUNK)
            state[name] = "done"

        async def file(i):
            name = f"f{i}"
            state[name] = "wait"
            await sem.acquire()
            state[name] = "run"
            await asyncio.sleep(READ)                  # read the file
            if strategy == "release":
                sem.release()
                state[name] = "mem"                    # waits without a slot
            else:
                state[name] = "hold"
            await asyncio.gather(*(chunk(f"{name}c{j}") for j in range(1, CHUNKS + 1)))
            if strategy == "hold":
                sem.release()
            state[name] = "done"

        async def walk():
            state["walk"] = "wait"
            await sem.acquire()
            if strategy == "release":
                sem.release()
                state["walk"] = "mem"
            else:
                state["walk"] = "hold"
            await asyncio.gather(*(file(i) for i in range(1, files + 1)))
            if strategy == "hold":
                sem.release()
            state["walk"] = "done"
    else:
        pool = Pool(cap)
        runs = {}

        def holding(n):
            r = runs.get(n)
            return r is not None and r.token is not None and not r.finished

        def chain_root(n):               # the run that took this node's token from the pool
            r = runs.get(n)
            while r is not None and r.token == "loan":
                r = r.parent
            return r.name if r is not None else n

        def fingerprint():
            return (tuple(state.values()), pool.free, len(pool.waiters),
                    tuple((n, r.token, r.lent, r.finished) for n, r in runs.items()))

        def snapshot():
            for n, r in runs.items():
                if state[n] in ("hold", "lent", "run") and holding(n):
                    if r.lent:
                        state[n] = "lent"
                    elif state[n] == "lent":
                        state[n] = "hold"
            pool_holders = [n for n in names if holding(n) and runs[n].token == "pool"]
            k = tokens(pool_holders, lambda n: chain_root(n) if holding(n) else n)
            loans = [n for n in names if holding(n) and runs[n].token == "loan"]
            return {"s": dict(state), "k": k, "l": loans}

        async def task(name, parent, body):
            r = Run(pool, parent)
            r.name = name
            runs[name] = r
            state[name] = "wait"
            await r.admit()
            try:
                await body(r)
            finally:
                r.finish()
                state[name] = "done"

        async def chunk_body(r, name):
            state[name] = "run"
            await asyncio.sleep(CHUNK)

        async def file_body(r, name):
            state[name] = "run"
            await asyncio.sleep(READ)
            state[name] = "hold"
            await asyncio.gather(*(task(f"{name}c{j}", r, lambda rr, n=f"{name}c{j}": chunk_body(rr, n))
                                   for j in range(1, CHUNKS + 1)))

        async def walk_body(r):
            state["walk"] = "hold"
            await asyncio.gather(*(task(f"f{i}", r, lambda rr, n=f"f{i}": file_body(rr, n))
                                   for i in range(1, files + 1)))

        async def walk():
            await task("walk", None, walk_body)

    async def main():
        nonlocal sem
        if strategy in ("hold", "release"):
            sem = asyncio.Semaphore(cap)
        s = asyncio.create_task(sampler_after_tick(snapshot, fingerprint))
        w = asyncio.create_task(walk())
        done, _ = await asyncio.wait({w}, timeout=limit)
        if done:
            w.result()                   # a crash raises here instead of reading as "done"
            finished["t"] = asyncio.get_running_loop().time()
        await s
        w.cancel()
        await asyncio.gather(w, return_exceptions=True)
        return finished["t"]

    loop = VirtualLoop()
    try:
        end = loop.run_until_complete(main())
    finally:
        loop.close()
    files_peak = max(sum(1 for i in range(1, files + 1) if f["s"][f"f{i}"] in ("run", "hold", "lent", "mem"))
                     for f in frames)
    return {"frames": frames, "end": end, "stuck": end is None, "files_peak": files_peak}


if __name__ == "__main__":
    out = {"cap": CAP, "files": FILES, "chunks": CHUNKS, "names": names_of(FILES),
           "runs": {k: record(k) for k in ("hold", "release", "lend")}}
    (HERE.parent / "fig" / "race-data.json").write_text(json.dumps(out, separators=(",", ":")))
    for k, r in out["runs"].items():
        print(f"{k:>7}: {'stuck' if r['stuck'] else 'done at t=' + str(r['end'])}, "
              f"{len(r['frames'])} frames, peak files in flight {r['files_peak']}")

# Machine: Apple M5 Pro, 18 cores, 48 GB; macOS 27.0.1; Python 3.12.13 (uv); Go 1.27.3; golang.org/x/sync v0.24.0; 2026-11-05

$ python deadlock.py
stuck after 2 s: 2 tasks hold both slots, 0 of 8 leaves done

$ python release_on_wait.py
cap 4: peak 2000 files in flight, 200 MB of contents held

$ python tree.py
cap  2, depth  3, fan-out    2,    15 tasks: peak executing 2, peak in flight 7 (bound 8), leaves at once 2, 0.01 s
cap  1, depth 20, fan-out    1,    21 tasks: peak executing 1, peak in flight 21 (bound 21), leaves at once 1, 0.00 s
cap  1, depth 10, fan-out    2,  2047 tasks: peak executing 1, peak in flight 11 (bound 11), leaves at once 1, 2.12 s
cap  2, depth 10, fan-out    2,  2047 tasks: peak executing 2, peak in flight 21 (bound 22), leaves at once 2, 2.20 s
cap  1, depth  1, fan-out 1000,  1001 tasks: peak executing 1, peak in flight 2 (bound 2), leaves at once 1, 3.81 s
cap 32, depth  1, fan-out 3000,  3001 tasks: peak executing 32, peak in flight 33 (bound 64), leaves at once 32, 0.25 s
cap  4, depth  3, fan-out   10,  1111 tasks: peak executing 4, peak in flight 13 (bound 16), leaves at once 4, 0.52 s

$ python parked.py
after parent ends: 0 free
other run waiting: True
other run admitted: pool
at the end: 1 free

$ go run ./go-errgroup
stuck after 2 s: both goroutines the limit admits are blocked in g.Go
$ go run ./go-trygo
finished

$ engine_bench.py (cocoindex from PyPI, one fresh LMDB dir per run; memory = macOS phys_footprint_peak)
cocoindex 1.0.25 cap 16: peak 1131 file components in flight, peak 8 chunks running, 153.5 s, peak memory 168 MB
cocoindex 1.1.1 cap 16: peak 16 file components in flight, peak 8 chunks running, 200.8 s, peak memory 33 MB
cocoindex 1.0.25 cap 16: peak 1125 file components in flight, peak 8 chunks running, 130.8 s, peak memory 167 MB
cocoindex 1.1.1 cap 16: peak 16 file components in flight, peak 8 chunks running, 114.2 s, peak memory 33 MB
cocoindex 1.1.1 cap 256: peak 255 file components in flight, peak 8 chunks running, 142.2 s, peak memory 68 MB

$ python compare.py
     hold: stuck (30 s timeout), peak files in flight    3 (  0.3 MB), peak chunks running 0
  release:             12.26 s, peak files in flight 2000 (200.0 MB), peak chunks running 4
per_level:             13.13 s, peak files in flight    4 (  0.4 MB), peak chunks running 4
     lend:             13.86 s, peak files in flight    4 (  0.4 MB), peak chunks running 4