Refactor SchedulerState to use per-handle waiter queues and a single FIFO queue

[coderabbitai]

Summary

The SchedulerState in src/client/pool/scheduler.rs currently uses three separate collections keyed by handle ID:

• waiters: HashMap<u64, WaiterSender<…>>
• fifo_waiters: VecDeque<u64>
• round_robin_handles: VecDeque<u64>

This design requires keeping all three structures consistent and involves repeated retain scans in the round-robin and deregister paths, which adds maintenance burden and cross-collection consistency concerns.

Proposed improvement

Restructure SchedulerState to use per-handle waiter queues for round-robin and a single anonymous FIFO queue for FIFO policy. This would:

• Eliminate the global waiters map and the fifo_waiters ID queue.
• Remove retain scans from take_next_round_robin_waiter and deregister_handle.
• Simplify the take_next_* selection logic by popping directly from per-handle queues.
• Eliminate cross-collection consistency concerns entirely.

All existing behaviour would be preserved: FIFO policy retains strict arrival-order servicing, and round-robin policy continues to rotate fairly among registered handles.

Context

This was flagged as an architectural suggestion during the review of PR #498 (#498). The current implementation is correct and all tests pass; this is a follow-up refactoring opportunity to reduce complexity and improve maintainability.

Requested by @leynos.

[traycerai]

Warning

Insufficient credits. Please purchase more credits or upgrade your plan to continue using Traycer.

[leynos]

@CodeRabbit plan

[coderabbitai]

Coding Plan

Summary

• Introduce handle_waiters: HashMap<u64, VecDeque<WaiterSender>> for per-handle round-robin queues, enabling direct pops without retain scans during selection.
• Replace fifo_waiters with fifo_queue: VecDeque<(u64, WaiterSender)>, tagging each entry with its handle ID so deregistered handles can be skipped at pop time rather than scanned at deregister time.
• Pass PoolFairnessPolicy into enqueue_waiter so each policy populates only its relevant structure.
• Preserve all behavioural contracts—strict FIFO ordering and fair round-robin alternation—verified by existing tests in src/client/tests/pool_handle.rs.

Implementation Steps

Restructure SchedulerState and Update All Dependent Methods

This phase replaces the three-collection design with two policy-specific structures: a per-handle queue map for round-robin and a handle-tagged queue for FIFO. All methods on SchedulerState and the relevant call sites in PoolScheduler are updated accordingly.

Task 1: Replace SchedulerState Fields

Refactor the struct definition to use the new data model.

• Remove waiters: HashMap<u64, WaiterSender<S, P, C>> entirely
• Replace fifo_waiters: VecDeque<u64> with fifo_queue: VecDeque<(u64, WaiterSender<S, P, C>)>, storing handle ID alongside each sender for deregistration support
• Rename round_robin_handles to round_robin_order for clarity (retains VecDeque<u64> type)
• Add handle_waiters: HashMap<u64, VecDeque<WaiterSender<S, P, C>>> for per-handle queues used by round-robin policy
• Update SchedulerState::new() to initialise all four fields as empty collections

Task 2: Update Registration and Deregistration Methods

Adapt register_handle and deregister_handle to operate on the new structures.

• register_handle: insert an empty VecDeque into handle_waiters for the new ID; push the ID onto round_robin_order
• deregister_handle: remove the entry from handle_waiters (this drops all pending round-robin waiters for that handle); retain-remove the ID from round_robin_order
• Note: the retain on round_robin_order remains—this is acceptable as it operates on a small, bounded collection of registered handles and is not on the hot selection path

Task 3: Update Waiter Enqueueing

Modify enqueue_waiter to dispatch based on policy and update its call sites.

• Add a policy: PoolFairnessPolicy parameter to SchedulerState::enqueue_waiter
• For PoolFairnessPolicy::Fifo: push (handle_id, sender) onto fifo_queue
• For PoolFairnessPolicy::RoundRobin: push sender onto the VecDeque in handle_waiters for the given handle ID; if the handle is not registered, drop the sender silently (matching current behaviour)
• Update PoolScheduler::acquire_for_handle in src/client/pool/scheduler.rs to pass self.fairness_policy when invoking state.enqueue_waiter

Task 4: Update Waiter Selection Methods

Rewrite take_next_fifo_waiter, take_next_round_robin_waiter, and has_waiters to operate on the new structures.

• take_next_fifo_waiter: pop (handle_id, sender) from the front of fifo_queue; if the handle ID is no longer present in handle_waiters (deregistered), discard the sender and continue popping until a valid entry is found or the queue is exhausted
• take_next_round_robin_waiter: iterate through round_robin_order with pop-front/push-back rotation; for each handle ID, attempt to pop a sender from its queue in handle_waiters; return the first sender found; this eliminates the previous retain scan on fifo_waiters
• has_waiters: return true if fifo_queue is non-empty OR any queue in handle_waiters is non-empty; a concise implementation is !self.fifo_queue.is_empty() || self.handle_waiters.values().any(|q| !q.is_empty())
• take_next_waiter remains unchanged in structure (dispatches to the appropriate helper based on policy)

🤖 Prompt for AI agents

Refactor `SchedulerState` in `src/client/pool/scheduler.rs` to replace the
three-collection design with two policy-specific structures. The goal is to
eliminate the global `waiters` map, introduce per-handle round-robin queues, and
tag FIFO entries with handle IDs.

**Task 1.1 — Replace SchedulerState Fields**
- Remove the `waiters: HashMap<u64, WaiterSender<S, P, C>>` field entirely
- Replace `fifo_waiters: VecDeque<u64>` with `fifo_queue: VecDeque<(u64,
WaiterSender<S, P, C>)>`
- Rename `round_robin_handles` to `round_robin_order` (keep the `VecDeque<u64>`
type)
- Add `handle_waiters: HashMap<u64, VecDeque<WaiterSender<S, P, C>>>`
- Update `SchedulerState::new()` to initialise all four fields as empty
collections

**Task 1.2 — Update Registration and Deregistration Methods**
- In `register_handle`: insert an empty `VecDeque` into `handle_waiters` for the
new handle ID; push the ID onto `round_robin_order`
- In `deregister_handle`: remove the entry from `handle_waiters` (dropping all
pending round-robin waiters); retain-remove the ID from `round_robin_order`
- The `retain` on `round_robin_order` is acceptable — it is not on the hot
selection path

**Task 1.3 — Update Waiter Enqueueing**
- Add a `policy: PoolFairnessPolicy` parameter to
`SchedulerState::enqueue_waiter`
- For `PoolFairnessPolicy::Fifo`: push `(handle_id, sender)` onto `fifo_queue`
- For `PoolFairnessPolicy::RoundRobin`: push `sender` onto the `VecDeque` in
`handle_waiters` for the given handle ID; if the handle is not registered, drop
the sender silently
- In `PoolScheduler::acquire_for_handle`, pass `self.fairness_policy` when
invoking `state.enqueue_waiter`

**Task 1.4 — Update Waiter Selection Methods**
- `take_next_fifo_waiter`: pop `(handle_id, sender)` from the front of
`fifo_queue`; if the handle ID is absent from `handle_waiters` (i.e.
deregistered), discard the sender and continue popping until a valid entry is
found or the queue is exhausted
- `take_next_round_robin_waiter`: iterate through `round_robin_order` using
pop-front/push-back rotation; for each handle ID, attempt to pop a sender from
its queue in `handle_waiters`; return the first sender found
- `has_waiters`: implement as `!self.fifo_queue.is_empty() ||
self.handle_waiters.values().any(|q| !q.is_empty())`
- Leave `take_next_waiter` structurally unchanged — it dispatches to the
appropriate helper based on policy

Research

The scheduler resides entirely in src/client/pool/scheduler.rs, comprising a private SchedulerState<S, P, C> struct and a pub(crate) PoolScheduler<S, P, C> wrapper. The current SchedulerState holds three collections—waiters (HashMap), fifo_waiters (VecDeque of IDs), and round_robin_handles (VecDeque of IDs)—which must be kept consistent across operations. PoolScheduler wraps this state in a Mutex and exposes registration, enqueueing, and selection methods. ClientPoolInner (in client_pool.rs) holds an Arc<PoolScheduler> and calls notify_shutdown and notify_capacity_available; PoolHandle (in handle.rs) calls acquire_for_handle and deregister_handle. Unit tests in src/client/tests/pool_handle.rs enforce strict alternation for round-robin and exact arrival-order for FIFO policies.

---

🚀 Next Steps

🤖 All AI agent prompts combined

Refactor `SchedulerState` in `src/client/pool/scheduler.rs` to replace the
three-collection design with two policy-specific structures. The goal is to
eliminate the global `waiters` map, introduce per-handle round-robin queues, and
tag FIFO entries with handle IDs.

**Task 1.1 — Replace SchedulerState Fields**
- Remove the `waiters: HashMap<u64, WaiterSender<S, P, C>>` field entirely
- Replace `fifo_waiters: VecDeque<u64>` with `fifo_queue: VecDeque<(u64,
WaiterSender<S, P, C>)>`
- Rename `round_robin_handles` to `round_robin_order` (keep the `VecDeque<u64>`
type)
- Add `handle_waiters: HashMap<u64, VecDeque<WaiterSender<S, P, C>>>`
- Update `SchedulerState::new()` to initialise all four fields as empty
collections

**Task 1.2 — Update Registration and Deregistration Methods**
- In `register_handle`: insert an empty `VecDeque` into `handle_waiters` for the
new handle ID; push the ID onto `round_robin_order`
- In `deregister_handle`: remove the entry from `handle_waiters` (dropping all
pending round-robin waiters); retain-remove the ID from `round_robin_order`
- The `retain` on `round_robin_order` is acceptable — it is not on the hot
selection path

**Task 1.3 — Update Waiter Enqueueing**
- Add a `policy: PoolFairnessPolicy` parameter to
`SchedulerState::enqueue_waiter`
- For `PoolFairnessPolicy::Fifo`: push `(handle_id, sender)` onto `fifo_queue`
- For `PoolFairnessPolicy::RoundRobin`: push `sender` onto the `VecDeque` in
`handle_waiters` for the given handle ID; if the handle is not registered, drop
the sender silently
- In `PoolScheduler::acquire_for_handle`, pass `self.fairness_policy` when
invoking `state.enqueue_waiter`

**Task 1.4 — Update Waiter Selection Methods**
- `take_next_fifo_waiter`: pop `(handle_id, sender)` from the front of
`fifo_queue`; if the handle ID is absent from `handle_waiters` (i.e.
deregistered), discard the sender and continue popping until a valid entry is
found or the queue is exhausted
- `take_next_round_robin_waiter`: iterate through `round_robin_order` using
pop-front/push-back rotation; for each handle ID, attempt to pop a sender from
its queue in `handle_waiters`; return the first sender found
- `has_waiters`: implement as `!self.fifo_queue.is_empty() ||
self.handle_waiters.values().any(|q| !q.is_empty())`
- Leave `take_next_waiter` structurally unchanged — it dispatches to the
appropriate helper based on policy

💡 Iterate on the plan with: @coderabbitai <feedback>

Example Feedback
- `@coderabbitai` You can skip phase 3. Add a simple unit test case for phase 2.
- `@coderabbitai` For design choice 1 go ahead with option 3 and replan.

---

💬 Have feedback or questions? Drop into our discord!