The First Brick: Building a Priority Queue for Deterministic Partition Scheduling

Message Overview

The subject message ([msg 2867]) is deceptively brief:

[assistant] First, add the BTreeMap import and the PriorityWorkQueue struct: [edit] /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs Edit applied successfully.

This single sentence and tool call represent the opening move in a carefully planned refactor of the cuzk proving engine's scheduling architecture. The message applies the first edit of a multi-step transformation that replaces FIFO channels with priority queues, solving a critical performance bug where all proving pipelines stalled together instead of completing sequentially.

The Problem That Demanded This Edit

To understand why this message was written, one must understand the scheduling pathology it was designed to fix. The cuzk proving daemon processes proof requests in a pipelined fashion: CPU-based synthesis transforms proof data into a form suitable for GPU proving, then hands the synthesized result to GPU workers. Between these stages, work items flow through channels.

The original design used tokio::sync::mpsc::channel — a FIFO (first-in, first-out) queue — for both the synthesis work channel and the GPU work channel. Under light load, this worked fine. But under realistic load with multiple jobs competing for resources, a devastating interaction emerged.

When multiple proof pipelines (jobs) were active simultaneously, all their partitions were dispatched as independent tokio tasks racing on a shared Notify-based budget acquire mechanism. This caused a thundering herd problem: every time budget freed up, all waiting tasks woke up, raced to acquire the slot, and the losers went back to sleep. The result was not just wasteful wakeups but, more critically, random partition selection across pipelines. The GPU channel, being FIFO, simply consumed whatever synthesized partition happened to finish first — which was essentially random because all partitions from all jobs started at roughly the same time with enough synthesis workers to handle them concurrently.

This randomness had a catastrophic consequence: instead of completing one pipeline at a time (which would free its memory and GPU resources), all pipelines progressed in lockstep, each partially done. Memory pressure remained high across all pipelines simultaneously, and none could finish quickly. The system was thrashing.

The Design Decision

The assistant's reasoning in the preceding message ([msg 2855]) reveals a thorough design exploration. The core insight was that both the synthesis queue and the GPU queue needed to be priority-based, not FIFO. The priority was defined as (job_seq ASC, partition_idx ASC) — meaning the lowest-numbered partition from the oldest (earliest-dispatched) job always gets processed first.

The choice of BTreeMap as the underlying data structure was deliberate. A BTreeMap maintains its entries in sorted order by key, which means the highest-priority item (smallest key) is always at the front. This is exactly the semantics needed: workers can pop the first entry and know they are processing the most important work. Alternatives like BinaryHeap were considered but BTreeMap offers O(log n) insertion and removal with the natural ordering semantics required.

The synchronization design was equally careful. The assistant considered several approaches for coordinating multiple worker tasks competing for queue entries:

Naive Mutex + Notify: Check the queue, if empty, wait for notification. But this has a race condition — a notification could arrive between unlocking the mutex and calling notified().await.
Pinned Notified with enable(): Create the Notified future before checking the queue to register as a waiter early. But enable() doesn't exist on Notified in the way expected.
Permit-based approach: Rely on Notify's internal permit mechanism. When notify_one() is called and no waiter is present, a permit is stored. The next notified().await consumes this permit immediately. This handles the race between try_pop() returning None and actually waiting. The assistant settled on approach 3, recognizing that notify_one() stores a permit when no one is waiting, and that permit is consumed by the next notified().await. This elegantly handles the race condition: if a push happens between the queue check and the await, the stored permit ensures the worker wakes up and retries. A further subtlety was the shutdown signaling. Rather than overloading the Notify for shutdown (which would require waking all waiters simultaneously), the assistant leveraged the existing tokio::sync::watch channel that was already used elsewhere in the engine. Workers would select! between queue.pop() and shutdown_rx.changed(), keeping the concerns separate.

Input Knowledge Required

Understanding this message requires familiarity with several layers of the system:

The cuzk proving pipeline: How synthesis (CPU work) feeds into GPU proving, and the role of channels as the glue between stages.
The budget-based memory manager: A 400 GiB budget gates how many partitions can be in-flight simultaneously. The scheduling fix interacts with this budget because workers claim items from the priority queue before acquiring budget.
tokio async primitives: Specifically Mutex, Notify, watch::Receiver, and the subtleties of Notify's permit-based wakeup mechanism. The assistant's reasoning shows deep understanding of tokio's internal behavior.
BTreeMap ordering semantics: The key type (u64, u32) sorts first by job_seq (the 64-bit job sequence number) and then by partition_idx (the 32-bit partition index within that job). This ensures all partitions of older jobs are processed before any partition of newer jobs.
The AtomicU64 job counter: A monotonically increasing sequence number assigned to each job at dispatch time, ensuring total ordering across all jobs regardless of when they arrive.

Output Knowledge Created

This message created the foundational data structure for the entire priority scheduling rewrite. The PriorityWorkQueue<T> struct and the BTreeMap import are the first building blocks. Subsequent messages in the same round ([msg 2868], [msg 2869]) add job_seq fields to PartitionWorkItem and SynthesizedJob, threading the priority key through the data flow.

The struct itself encapsulates:

A Mutex<BTreeMap<(u64, u32), T>> for storage
A tokio::sync::Notify for coordination
A closed flag for graceful shutdown
push() method: inserts an item with its priority key and notifies one waiter
pop() method: returns the highest-priority item, blocking if empty, using the permit-based race avoidance pattern
close() method: signals shutdown to all waiters This structure would be used to replace both the synthesis work channel and the GPU work channel, ensuring deterministic, priority-ordered processing across all pipelines.

Assumptions and Potential Mistakes

The assistant made several assumptions worth examining:

The priority queue would fix the thundering herd problem: The assumption was that replacing FIFO channels with priority queues would cause pipelines to complete sequentially rather than in lockstep. This is correct for the ordering of work, but the thundering herd on budget acquire is a separate concern — the priority queue ensures the right work is done first, but workers still race on budget. The assistant acknowledged this tradeoff, noting that workers claim an item from the queue before acquiring budget, which means a worker holding a lower-priority item could block higher-priority work from being seen. This was deemed acceptable because items are pushed in dispatch order.
BTreeMap is the right data structure: A BTreeMap provides ordered iteration and efficient insertion/removal. However, for a priority queue, a BinaryHeap would also work and might have better cache behavior. The choice of BTreeMap likely reflects the need to potentially iterate over all items (e.g., for debugging or status reporting), which BTreeMap supports naturally while BinaryHeap does not.
The Notify permit model is sufficient: The assistant correctly identified that notify_one() stores a permit when no waiter is present. However, there's a subtlety: only one permit is stored even if multiple pushes happen. The chain-reaction pattern (pop notifies next waiter if items remain) handles this, but it means the queue relies on workers cooperatively passing the wakeup torch. If a worker crashes or stalls between popping and notifying, other workers could starve.
The job_seq counter is monotonic and non-wrapping: Using an AtomicU64 that monotonically increases assumes the system won't run long enough to wrap around. For a proving daemon that might run for months, this is a safe assumption — even at 1,000 jobs per second, wrapping would take over 500 million years.

The Broader Context

This message sits at a transition point in the conversation. The preceding segments (16–20) had implemented a unified budget-based memory manager, a JSON status API, and a live monitoring UI. These were deployed and tested on a remote machine. But during testing, the fundamental scheduling flaw was discovered: all pipelines stalled together. The priority queue fix was designed in message 2855 and begins implementation here.

The subsequent messages in the same round ([msg 2868], [msg 2869]) continue the implementation by adding job_seq to the work item structs. Later messages will replace the channel initialization, update the worker loops, and thread the priority counter through the dispatch chain.

What makes this message significant is not its size — it's a single edit — but its role as the keystone of a architectural change. The PriorityWorkQueue struct is the foundation upon which deterministic scheduling is built. Without it, the GPU workers would continue consuming synthesized partitions in random order, and all pipelines would continue to thrash. With it, the proving engine gains predictable, FIFO-per-pipeline behavior that matches user expectations: submit jobs in order, get proofs back in order, with maximum throughput.