From a5f1ffa5a00634a830fe91c1a3261e0b743f90d5 Mon Sep 17 00:00:00 2001 From: Leynos Date: Tue, 10 Mar 2026 14:20:16 +0000 Subject: [PATCH 1/5] docs(execplans): add detailed execution plan for PoolHandle API (11.2.2) Introduce a comprehensive 457-line execution plan document laying out the design, constraints, risks, progress, decision log, and API plan for roadmap item 11.2.2. This plan covers exposing a new PoolHandle API to provide fair pooled acquisition for logical sessions, clarifying fairness policies, scheduler design, testing requirements, and documentation guidance to improve the Wireframe client pool subsystem. Co-authored-by: devboxerhub[bot] --- .../11-2-2-expose-pool-handle-api.md | 457 ++++++++++++++++++ 1 file changed, 457 insertions(+) create mode 100644 docs/execplans/11-2-2-expose-pool-handle-api.md diff --git a/docs/execplans/11-2-2-expose-pool-handle-api.md b/docs/execplans/11-2-2-expose-pool-handle-api.md new file mode 100644 index 00000000..00c66250 --- /dev/null +++ b/docs/execplans/11-2-2-expose-pool-handle-api.md @@ -0,0 +1,457 @@ +# Expose a `PoolHandle` API with fair pooled acquisition for logical sessions (11.2.2) + +This ExecPlan (execution plan) is a living document. The sections +`Constraints`, `Tolerances`, `Risks`, `Progress`, `Surprises & Discoveries`, +`Decision Log`, and `Outcomes & Retrospective` must be kept up to date as work +proceeds. + +Status: DRAFT + +## Purpose / big picture + +Roadmap item `11.2.2` exists because the current pooled-client surface is still +too low-level for larger deployments. Today, all contention flows through +`WireframeClientPool::acquire()`. That API is correct for manual control, but +it gives the pool no durable notion of "logical session A" versus "logical +session B", so a busy caller can repeatedly reacquire leases and crowd out +other waiters. + +After this change, a library consumer will be able to create a `PoolHandle` +from a pooled client and treat that handle as the fairness identity for one +logical session. When many handles contend for a small number of warm sockets, +the pool will admit them according to an explicit fairness policy while still +respecting the existing per-socket admission permits and serialized transport +access. Observable success is: + +- unit tests written with `rstest` prove that handle-aware acquisition is fair, + preserves back-pressure, and does not regress warm reuse or idle recycle; +- behavioural tests written with `rstest-bdd` v0.5.0 prove the same behaviour + through user-visible scenarios; +- `docs/wireframe-client-design.md` records the `11.2.2` design decisions; +- `docs/users-guide.md` explains when to use `PoolHandle` instead of direct + `pool.acquire()`; and +- `docs/roadmap.md` marks `11.2.2` as done only after all quality gates pass. + +This plan deliberately keeps `PooledClientLease` as the low-level escape hatch +for split-phase workflows (`send` followed later by `receive`). `PoolHandle` +adds fairness around lease acquisition and may add one-shot convenience methods +only where they remain safe on a shared pool. + +## Constraints + +- Scope is limited to roadmap item `11.2.2`. +- Existing `WireframeClientPool`, `ClientPoolConfig`, and `PooledClientLease` + APIs must remain backward compatible. +- The `pool` Cargo feature remains the opt-in gate for all pooled client code. +- Fairness must be additive above the existing slot-permit back-pressure. The + new API must not bypass `max_in_flight_per_socket`, socket serialization, or + idle recycle rules introduced in `11.2.1`. +- `PoolHandle` must represent a logical-session fairness identity, not a + promise of transport affinity to one physical socket. +- Split-phase receive semantics must remain explicit. Do not imply that a + shared `PoolHandle` can safely route arbitrary uncorrelated responses unless + the implementation truly adds response demultiplexing. +- No single Rust source file may exceed 400 lines; extract submodules before + crossing roughly 390 lines. +- Public APIs must be documented with Rustdoc examples that satisfy the doctest + guidance in `docs/rust-doctest-dry-guide.md`. +- Unit tests must use `rstest`. +- Behavioural tests must use `rstest-bdd` v0.5.0 with the repository's + existing `feature + fixture + steps + scenarios` layout. +- BDD fixture parameter names in step definitions must match fixture function + names exactly, and step parameters must not use underscore-prefixed fixture + names. +- Design decisions must be recorded in `docs/wireframe-client-design.md`. +- Public interface changes must be recorded in `docs/users-guide.md`. +- `docs/roadmap.md` `11.2.2` is updated to done only after full validation. +- Follow guidance from: + - `docs/generic-message-fragmentation-and-re-assembly-design.md` + - `docs/multi-packet-and-streaming-responses-design.md` + - `docs/the-road-to-wireframe-1-0-feature-set-philosophy-and-capability-maturity.md` + - `docs/hardening-wireframe-a-guide-to-production-resilience.md` + - `docs/rust-testing-with-rstest-fixtures.md` + - `docs/reliable-testing-in-rust-via-dependency-injection.md` + - `docs/rstest-bdd-users-guide.md` + - `docs/rust-doctest-dry-guide.md` + +## Tolerances + +- Scope: if implementation requires a full response router for uncorrelated + `receive()` traffic across shared handles, stop and escalate. That is a + materially larger feature than `11.2.2`. +- Interface: if delivering `PoolHandle` requires removing or breaking + `WireframeClientPool::acquire()` or `PooledClientLease`, stop and escalate. +- Surface area: if more than 20 files or roughly 1600 net lines change before + documentation/test additions, stop and re-evaluate the design. +- Dependencies: do not add new crates for scheduling or fairness. Reuse Tokio, + `bb8`, and existing std/futures primitives. +- Validation: if the same quality gate fails three focused fix attempts in a + row, stop and document the blocker before proceeding. +- Time: if the implementation cannot reach a first red-green pass in one work + session, stop after the scheduler prototype and request review on the chosen + API boundary. + +## Risks + +- Risk: the current pool selection path uses `select_all` over waiter futures, + so under contention it is scheduler-dependent rather than demonstrably fair. + Severity: high. Likelihood: high. Mitigation: introduce an explicit handle + scheduler with deterministic policy tests. + +- Risk: `PooledClientLease` holds a slot permit for the lease lifetime, which + means fairness can still be defeated by callers that keep leases idle for too + long. Severity: high. Likelihood: medium. Mitigation: document that + `PoolHandle` is the preferred fairness surface and add tests that release + leases promptly after work. + +- Risk: adding convenience methods directly on `PoolHandle` can accidentally + promise unsafe split-phase semantics on a shared pool. Severity: high. + Likelihood: medium. Mitigation: keep `PoolHandle` focused on fair + acquisition, and only add convenience methods that encapsulate the whole + request/response cycle. + +- Risk: fairness semantics are easy to confuse with the existing round-robin + slot rotation. Severity: medium. Likelihood: high. Mitigation: document the + distinction clearly in code docs and the users' guide: slot rotation spreads + work across warm sockets, while handle fairness orders contending logical + sessions. + +- Risk: deterministic tests for fairness can become flaky if they depend on + runtime wake ordering. Severity: medium. Likelihood: medium. Mitigation: use + explicit barriers, bounded channels, and recorded grant order rather than + timing assertions. + +- Risk: BDD scenarios can fail for fixture wiring reasons unrelated to pooling. + Severity: medium. Likelihood: medium. Mitigation: mirror the established + `tests/bdd_pool/` layout and keep fixture names identical across feature, + scenario, and step files. + +## Progress + +- [x] (2026-03-10 00:00Z) Reviewed roadmap item `11.2.2`, the completed + `11.2.1` pool implementation, relevant design/testing docs, and project notes. +- [x] (2026-03-10 00:00Z) Drafted this ExecPlan. +- [ ] Stage A: confirm the public API boundary and internal scheduler shape. +- [ ] Stage B: add failing unit and behavioural tests that define fairness and + back-pressure behaviour. +- [ ] Stage C: implement `PoolHandle`, fairness policies, and scheduler-backed + lease acquisition. +- [ ] Stage D: update design docs, users' guide, roadmap, and run full gates. + +## Surprises & Discoveries + +- Observation: the current `WireframeClientPool::acquire()` first tries an + immediate permit, then races all slot waiters with `futures::select_all`. + Evidence: `src/client/pool/client_pool.rs`. Impact: current contention is + "works eventually" rather than a documented fairness contract. + +- Observation: the existing `11.2.1` design already treats the pool as a + hybrid of `bb8` lifecycle management plus Wireframe-owned admission control. + Evidence: `docs/wireframe-client-design.md` decision record for `11.2.1`. + Impact: `11.2.2` should extend the Wireframe-owned layer rather than fight or + replace `bb8`. + +- Observation: pooled BDD coverage already lives in a separate + `tests/bdd_pool/` target gated behind `advanced-tests` and `pool`. Evidence: + `Cargo.toml`, `tests/bdd_pool/mod.rs`. Impact: add new handle scenarios + alongside the existing pool scenarios instead of inventing a new harness. + +- Observation: `make test` and `make lint` already run with `--all-features`, + so the pooled test target and any new pool docs are covered by the standard + gates. Evidence: `Makefile`. Impact: no bespoke validation target is needed, + only the standard gates plus Markdown/doc gates. + +## Decision Log + +- Decision: treat `PoolHandle` as the fairness identity for one logical + session. Rationale: this directly addresses the roadmap requirement to + multiplex many logical sessions without changing the lower-level lease API. + Date/Author: 2026-03-10 / planning. + +- Decision: keep `PooledClientLease` as the explicit low-level API for + split-phase and transport-sensitive workflows. Rationale: a shared + `PoolHandle` does not, by itself, solve response demultiplexing for later + `receive()` calls. Date/Author: 2026-03-10 / planning. + +- Decision: implement fairness as a pool-level policy, exposed through a new + public policy type and configured through pooled-client configuration. + Rationale: fairness is only coherent when every contending handle is governed + by the same scheduler rules. Date/Author: 2026-03-10 / planning. + +- Decision: prefer `PoolHandle::acquire(&mut self)` over a clone-heavy or + unbounded queueing model. Rationale: one outstanding wait per handle gives + natural back-pressure and keeps "one handle == one logical session" + straightforward. Date/Author: 2026-03-10 / planning. + +- Decision: any convenience methods added to `PoolHandle` must encapsulate an + entire operation (for example `call`) and must be built atop fair lease + acquisition. Rationale: this avoids implying that a shared handle can safely + perform arbitrary split-phase I/O. Date/Author: 2026-03-10 / planning. + +## Context and orientation + +Current pooled-client code lives under `src/client/pool/` and is feature-gated +by `pool`. + +Relevant files today: + +- `src/client/pool/client_pool.rs`: public pool type plus current slot + selection and acquisition logic. +- `src/client/pool/slot.rs`: one physical-socket slot with a `bb8` pool of + size one and a per-slot semaphore. +- `src/client/pool/lease.rs`: `PooledClientLease` forwarding operations through + a checked-out managed client connection. +- `src/client/pool/config.rs`: `ClientPoolConfig`. +- `src/client/builder/pool.rs`: `WireframeClientBuilder::connect_pool(...)`. +- `src/test_helpers/pool_client.rs`: shared test server and pooled-client + builders. +- `src/client/tests/pool.rs`: unit tests for warm reuse, in-flight admission, + idle recycle, and broken-connection recycle. +- `tests/features/client_pool.feature`, `tests/fixtures/client_pool.rs`, + `tests/steps/client_pool_steps.rs`, and + `tests/scenarios/client_pool_scenarios.rs`: current behavioural coverage. +- `docs/wireframe-client-design.md`: current client and pooling design record. +- `docs/users-guide.md`: pool-facing public documentation. + +Terms used in this plan: + +- Logical session: one consumer identity that wants fair access to pooled + leases over time. +- Handle fairness: the rule used to decide which waiting logical session gets + the next lease opportunity. +- Slot rotation: the existing round-robin choice of which physical socket to + try first. This is not the same thing as handle fairness. +- Back-pressure: the property that callers must wait instead of creating + unbounded queued work when all permits/sockets are busy. + +## Planned public API + +The target public shape is: + +```rust +use wireframe::client::{ + ClientPoolConfig, + PoolFairnessPolicy, + PoolHandle, + WireframeClient, +}; + +let pool = WireframeClient::builder() + .connect_pool( + addr, + ClientPoolConfig::default().fairness_policy(PoolFairnessPolicy::RoundRobin), + ) + .await?; + +let mut session_a: PoolHandle<_, _, _> = pool.handle(); +let mut session_b = pool.handle(); + +let lease_a = session_a.acquire().await?; +let lease_b = session_b.acquire().await?; +``` + +The exact policy names may change to match repository naming conventions, but +the semantics must stay: + +1. A consumer can create a stable handle for one logical session. +2. That handle can acquire leases under an explicit fairness policy. +3. The fairness policy never bypasses existing slot permits or transport + serialization. + +Preferred policy set for this milestone: + +- `RoundRobin`: give waiting handles turns in rotation, one successful grant + per turn. +- `Fifo`: grant the next lease to the earliest waiting handle. + +If implementation evidence shows one of these policies is redundant or cannot +be explained clearly to users, keep `RoundRobin` as the required default and +document the change in the decision log before simplifying. + +## Plan of work + +### Stage A: define the handle/scheduler boundary + +Refactor the pooled client around an explicit shared inner state. The simplest +shape is for `WireframeClientPool` to become a thin public wrapper +around `Arc>`, because both the pool and its handles +need access to shared slots and shared fairness state. + +Add new pool submodules, likely: + +- `src/client/pool/handle.rs` +- `src/client/pool/policy.rs` or `src/client/pool/fairness.rs` +- `src/client/pool/scheduler.rs` + +Keep `client_pool.rs`, `lease.rs`, and `slot.rs` focused by extracting helpers +instead of growing any one file past the 400-line cap. + +Implement a scheduler that owns: + +- the fairness policy; +- a stable handle identifier for each `PoolHandle`; +- a wait queue or rotation structure of blocked handles; and +- a wake-up path that runs when a lease is dropped or a slot becomes newly + available. + +The scheduler must grant real leases, not merely "permission to try". A handle +that wins fairness but then races all other tasks for the actual slot permit +would not be observably fair. The intended flow is: + +1. A handle registers itself as waiting. +2. The scheduler chooses the next waiting handle according to policy. +3. The scheduler acquires an actual slot permit and constructs a + `PooledClientLease`. +4. The scheduler delivers that lease to the chosen handle. +5. Lease drop notifies the scheduler to service the next waiter. + +Wire direct `pool.acquire()` through the same shared machinery where practical. +If direct acquire remains a thin low-level path, document clearly that strong +fairness guarantees apply to stable `PoolHandle` usage, not to repeated +stateless calls. + +### Stage B: write the failing tests first + +Add unit tests before implementing the scheduler. Keep them in a new focused +test module if `src/client/tests/pool.rs` would otherwise grow too large. + +Required unit tests with `rstest`: + +1. `round_robin_handles_share_one_socket_fairly` + Two handles contend on a pool with `pool_size(1)` and + `max_in_flight_per_socket(1)`. Recorded grant order must alternate by + handle, not allow one handle to monopolize reacquisition. + +2. `fifo_policy_preserves_wait_order` + Three handles wait in a known order. Grants must follow that order under the + FIFO policy. + +3. `handle_acquire_respects_back_pressure` + Holding a lease from one handle must keep later handles waiting rather than + over-admitting extra work. The assertion is "still waiting", not a sleepy + throughput benchmark. + +4. `handle_path_preserves_warm_reuse_and_preamble` + Repeated `handle.acquire()` on one logical session with one socket should + reuse the warm connection and keep the preamble count at one. + +5. `handle_path_recycles_after_idle_timeout` + After paused-time advancement beyond the configured idle timeout, the next + handle acquisition should reconnect and rerun the preamble exactly once. + +Add behavioural tests under the existing pooled-client BDD harness. Create a +new feature and supporting files unless the existing feature remains concise: + +- `tests/features/client_pool_handle.feature` +- `tests/fixtures/client_pool_handle.rs` +- `tests/steps/client_pool_handle_steps.rs` +- `tests/scenarios/client_pool_handle_scenarios.rs` +- include them from `tests/bdd_pool/fixtures.rs` and + `tests/bdd_pool/scenarios.rs` + +Required behavioural scenarios: + +1. Two logical sessions alternate access fairly on one pooled socket. +2. FIFO fairness serves waiting sessions in arrival order. +3. A waiting session remains blocked until another session releases its lease. +4. Warm reuse and idle recycle still behave the same through the handle API. + +Use the shared `PoolTestServer` fixture utilities where possible. Extend +`src/test_helpers/pool_client.rs` only with deterministic coordination helpers, +such as barriers or ordered event recording, and keep those helpers generic +enough for both unit and BDD coverage. + +### Stage C: implement `PoolHandle` and fairness policies + +Add the new public policy type to the pool module and thread it through +`ClientPoolConfig`. The config must stay ergonomic and builder-style, for +example: + +```rust +let config = ClientPoolConfig::default() + .pool_size(2) + .max_in_flight_per_socket(1) + .fairness_policy(PoolFairnessPolicy::RoundRobin); +``` + +Implement `WireframeClientPool::handle()` to create a new logical-session +handle registered with the scheduler. The handle itself should be lightweight +and should borrow or clone only the shared inner state plus its own stable +identifier. + +Implement `PoolHandle::acquire(&mut self)` to: + +1. register the handle as waiting if no immediate lease is available; +2. await a scheduler-delivered lease; and +3. return the same `PooledClientLease` type that existing callers already use. + +Keep the release path explicit. Lease drop must: + +1. release the per-slot semaphore permit as it does today; +2. preserve the existing last-used / broken-connection bookkeeping; and +3. notify the scheduler that another waiting handle may now be serviceable. + +If convenience methods are added on `PoolHandle`, restrict them to whole +operations that are safe atop fair acquisition, such as `call`. Do not add +`receive()` or `receive_envelope()` on `PoolHandle` unless the implementation +also adds principled response routing, which is outside this milestone's +intended scope. + +### Stage D: documentation, doctests, and roadmap updates + +Update `docs/wireframe-client-design.md` with a new `11.2.2` decision record +covering: + +- why `PoolHandle` represents logical-session fairness instead of socket + affinity; +- how fairness differs from the existing slot-rotation logic; +- which policy becomes the default and why; and +- why split-phase APIs remain on `PooledClientLease`. + +Update `docs/users-guide.md` so library consumers know: + +- when to use `pool.handle()` instead of repeated `pool.acquire()`; +- what fairness policy means operationally; +- that `PoolHandle` preserves existing back-pressure rather than creating a + separate queue outside the pool; and +- that `PooledClientLease` is still the correct surface for explicit + split-phase workflows. + +Add or update Rustdoc examples on every new public type and method. Keep the +examples realistic and runnable so `make test-doc` and `make doctest-benchmark` +stay green. + +Finally, mark roadmap item `11.2.2` done in `docs/roadmap.md` only after all +quality gates pass. + +## Validation and evidence + +Run every gate through `tee` with `set -o pipefail` so failures are visible and +logs remain inspectable: + +```bash +set -o pipefail && make check-fmt | tee /tmp/wireframe-11-2-2-check-fmt.log +set -o pipefail && make lint | tee /tmp/wireframe-11-2-2-lint.log +set -o pipefail && make test | tee /tmp/wireframe-11-2-2-test.log +set -o pipefail && make test-doc | tee /tmp/wireframe-11-2-2-test-doc.log +set -o pipefail && make doctest-benchmark | tee /tmp/wireframe-11-2-2-doctest-benchmark.log +set -o pipefail && make markdownlint MDLINT=/root/.bun/bin/markdownlint-cli2 | tee /tmp/wireframe-11-2-2-markdownlint.log +set -o pipefail && make nixie | tee /tmp/wireframe-11-2-2-nixie.log +``` + +Observable completion criteria: + +- the new unit tests fail before the scheduler work and pass afterward; +- the new BDD scenarios fail before the scheduler work and pass afterward; +- no existing pooled-client tests regress; +- doctests for the new public API compile and run; and +- `docs/roadmap.md` shows `11.2.2` as done. + +## Outcomes & Retrospective + +Pending implementation. Complete this section at the end of the feature with: + +- the final public API surface; +- any deviations from this draft plan; +- validation results with the exact commands that passed; and +- lessons learned about fairness, pool ergonomics, and follow-up work. From 7cc51f251549ac540c6d7ed6d522ee626136c176 Mon Sep 17 00:00:00 2001 From: Leynos Date: Thu, 12 Mar 2026 08:30:37 +0000 Subject: [PATCH 2/5] feat(client-pool): add PoolHandle API for fair lease scheduling - Introduce PoolHandle as a logical-session identity for fair and orderly pooled lease acquisition. - Implement PoolScheduler to queue blocked handles and grant leases according to fairness policies. - Add PoolFairnessPolicy enum with RoundRobin (default) and Fifo variants. - Update WireframeClientPool to expose handle() for creating PoolHandle instances. - Modify lease and pool inner types to cooperate with scheduler for fairness and back-pressure. - Include comprehensive unit, BDD scenario, and integration tests for PoolHandle fairness behaviors. - Update docs and examples to showcase PoolHandle usage and fairness configurations. This enhancement enables clients to multiplex logical sessions fairly over pooled connections without compromising back-pressure or warm socket reuse. Co-authored-by: devboxerhub[bot] --- .../11-2-2-expose-pool-handle-api.md | 61 ++++- docs/roadmap.md | 2 +- docs/users-guide.md | 59 ++++- docs/wireframe-client-design.md | 34 +++ src/client/mod.rs | 8 +- src/client/pool/client_pool.rs | 94 +++++-- src/client/pool/config.rs | 21 +- src/client/pool/handle.rs | 129 ++++++++++ src/client/pool/lease.rs | 23 +- src/client/pool/mod.rs | 5 + src/client/pool/policy.rs | 30 +++ src/client/pool/scheduler.rs | 213 ++++++++++++++++ src/client/tests/mod.rs | 2 + src/client/tests/pool_handle.rs | 200 +++++++++++++++ tests/bdd_pool/fixtures.rs | 3 + tests/bdd_pool/scenarios.rs | 6 + tests/features/client_pool_handle.feature | 23 ++ tests/fixtures/client_pool_handle.rs | 239 ++++++++++++++++++ .../scenarios/client_pool_handle_scenarios.rs | 37 +++ tests/steps/client_pool_handle_steps.rs | 111 ++++++++ 20 files changed, 1256 insertions(+), 44 deletions(-) create mode 100644 src/client/pool/handle.rs create mode 100644 src/client/pool/policy.rs create mode 100644 src/client/pool/scheduler.rs create mode 100644 src/client/tests/pool_handle.rs create mode 100644 tests/features/client_pool_handle.feature create mode 100644 tests/fixtures/client_pool_handle.rs create mode 100644 tests/scenarios/client_pool_handle_scenarios.rs create mode 100644 tests/steps/client_pool_handle_steps.rs diff --git a/docs/execplans/11-2-2-expose-pool-handle-api.md b/docs/execplans/11-2-2-expose-pool-handle-api.md index 00c66250..0920c4d7 100644 --- a/docs/execplans/11-2-2-expose-pool-handle-api.md +++ b/docs/execplans/11-2-2-expose-pool-handle-api.md @@ -5,7 +5,7 @@ This ExecPlan (execution plan) is a living document. The sections `Decision Log`, and `Outcomes & Retrospective` must be kept up to date as work proceeds. -Status: DRAFT +Status: COMPLETE ## Purpose / big picture @@ -131,12 +131,15 @@ only where they remain safe on a shared pool. - [x] (2026-03-10 00:00Z) Reviewed roadmap item `11.2.2`, the completed `11.2.1` pool implementation, relevant design/testing docs, and project notes. - [x] (2026-03-10 00:00Z) Drafted this ExecPlan. -- [ ] Stage A: confirm the public API boundary and internal scheduler shape. -- [ ] Stage B: add failing unit and behavioural tests that define fairness and - back-pressure behaviour. -- [ ] Stage C: implement `PoolHandle`, fairness policies, and scheduler-backed - lease acquisition. -- [ ] Stage D: update design docs, users' guide, roadmap, and run full gates. +- [x] (2026-03-12 00:00Z) Stage A: confirmed the public API boundary and + internal scheduler shape around `ClientPoolInner`, `PoolHandle`, and + `PoolScheduler`. +- [x] (2026-03-12 00:00Z) Stage B: added unit and behavioural tests that + define fairness and back-pressure behaviour. +- [x] (2026-03-12 00:00Z) Stage C: implemented `PoolHandle`, fairness + policies, and scheduler-backed lease acquisition. +- [x] (2026-03-12 00:00Z) Stage D: updated design docs, users' guide, + roadmap, and ran full gates. ## Surprises & Discoveries @@ -161,6 +164,12 @@ only where they remain safe on a shared pool. gates. Evidence: `Makefile`. Impact: no bespoke validation target is needed, only the standard gates plus Markdown/doc gates. +- Observation: the scheduler can remain lightweight if it queues blocked + handles and obtains real slot permits in a spawned service loop rather than + maintaining a permanently running background task. Evidence: + `src/client/pool/scheduler.rs`. Impact: no extra shutdown protocol was + needed; lease drop simply kicks the scheduler when capacity returns. + ## Decision Log - Decision: treat `PoolHandle` as the fairness identity for one logical @@ -449,9 +458,35 @@ Observable completion criteria: ## Outcomes & Retrospective -Pending implementation. Complete this section at the end of the feature with: - -- the final public API surface; -- any deviations from this draft plan; -- validation results with the exact commands that passed; and -- lessons learned about fairness, pool ergonomics, and follow-up work. +- Final public API surface: + - `ClientPoolConfig::fairness_policy(...)` + - `PoolFairnessPolicy::{RoundRobin, Fifo}` + - `WireframeClientPool::handle()` + - `PoolHandle::acquire(&mut self)` + - `PoolHandle::call(&mut self, ...)` +- Implementation summary: + - `WireframeClientPool` now wraps shared `ClientPoolInner` state; + - `PoolScheduler` queues blocked handles and obtains real slot permits before + handing back a `PooledClientLease`; + - `PooledClientLease` remains the split-phase API, while `PoolHandle` is the + fairness-aware logical-session API. +- Validation completed before full gates: + - `cargo test --lib client::tests::pool_handle --features pool -- --nocapture` + - `cargo test --test bdd_pool --features "advanced-tests pool" -- --nocapture` +- Final quality gates passed: + - `set -o pipefail && make fmt | tee /tmp/wireframe-11-2-2-fmt.log` + - `set -o pipefail && make check-fmt | tee /tmp/wireframe-11-2-2-check-fmt.log` + - `set -o pipefail && make lint | tee /tmp/wireframe-11-2-2-lint.log` + - `set -o pipefail && make test | tee /tmp/wireframe-11-2-2-test.log` + - `set -o pipefail && make test-doc | tee /tmp/wireframe-11-2-2-test-doc.log` + - `set -o pipefail && make doctest-benchmark | tee /tmp/wireframe-11-2-2-doctest-benchmark.log` + - `set -o pipefail && make markdownlint MDLINT=/root/.bun/bin/markdownlint-cli2 | tee /tmp/wireframe-11-2-2-markdownlint.log` + - `set -o pipefail && make nixie | tee /tmp/wireframe-11-2-2-nixie.log` +- Lessons learned: + - fairness is easiest to explain when the scheduler operates on logical + session handles rather than trying to retrofit identity into repeated + stateless `pool.acquire()` calls; + - a real lease grant must include a real slot permit, otherwise fairness is + only advisory and collapses back into permit races; + - keeping `PooledClientLease` separate from `PoolHandle` avoids implying + response demultiplexing that the current pool does not implement. diff --git a/docs/roadmap.md b/docs/roadmap.md index 133acf3d..aa47d811 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -502,7 +502,7 @@ document so larger deployments can adopt the library confidently. - [x] 11.2.1. Implement a configurable connection pool that preserves preamble state, enforces in-flight request limits per socket, and recycles idle connections. -- [ ] 11.2.2. Expose a `PoolHandle` API with fairness policies so callers can +- [x] 11.2.2. Expose a `PoolHandle` API with fairness policies so callers can multiplex many logical sessions without violating back-pressure. ### 11.3. Streaming helpers and test utilities diff --git a/docs/users-guide.md b/docs/users-guide.md index a5c0392c..ea83bb31 100644 --- a/docs/users-guide.md +++ b/docs/users-guide.md @@ -1383,13 +1383,14 @@ cargo run --example client_echo_login --features examples Use `connect_pool` when one warm socket is too little, but opening a new TCP connection for every request is too expensive. `WireframeClientPool` keeps a -bounded set of warm sockets, preserves preamble state on reuse, and recycles -idle sockets after the configured timeout. +bounded set of warm sockets, preserves preamble state on reuse, recycles idle +sockets after the configured timeout, and can schedule blocked logical sessions +fairly through `PoolHandle`. ```rust,no_run use std::{net::SocketAddr, time::Duration}; -use wireframe::client::{ClientPoolConfig, WireframeClient}; +use wireframe::client::{ClientPoolConfig, PoolFairnessPolicy, WireframeClient}; #[derive(bincode::Encode, bincode::Decode)] struct Ping(u8); @@ -1406,7 +1407,8 @@ let pool = WireframeClient::builder() ClientPoolConfig::default() .pool_size(2) .max_in_flight_per_socket(2) - .idle_timeout(Duration::from_secs(30)), + .idle_timeout(Duration::from_secs(30)) + .fairness_policy(PoolFairnessPolicy::RoundRobin), ) .await?; @@ -1425,6 +1427,55 @@ serialized per physical connection while still allowing multiple callers to be admitted against the same warm slot. Treat `max_in_flight_per_socket` as an admission budget, not as a guarantee of parallel writes on one TCP stream. +Create a `PoolHandle` when one logical session needs repeated pooled access and +you want that session to participate in the configured fairness policy over +time: + +```rust,no_run +# use std::net::SocketAddr; +use wireframe::client::{ClientPoolConfig, PoolFairnessPolicy, PoolHandle, WireframeClient}; + +#[derive(bincode::Encode, bincode::Decode)] +struct Ping(u8); + +#[derive(bincode::Encode, bincode::Decode, Debug, PartialEq)] +struct Pong(u8); + +# #[tokio::main] +# async fn main() -> Result<(), wireframe::client::ClientError> { +let addr: SocketAddr = "127.0.0.1:7878".parse().expect("valid socket address"); +let pool = WireframeClient::builder() + .connect_pool( + addr, + ClientPoolConfig::default() + .pool_size(1) + .fairness_policy(PoolFairnessPolicy::RoundRobin), + ) + .await?; + +let mut session_a: PoolHandle<_, _, _> = pool.handle(); +let mut session_b = pool.handle(); + +let pong_a: Pong = session_a.call(&Ping(1)).await?; +let pong_b: Pong = session_b.call(&Ping(2)).await?; + +assert_eq!(pong_a, Pong(1)); +assert_eq!(pong_b, Pong(2)); +# Ok(()) +# } +``` + +Use `pool.handle()` instead of repeated `pool.acquire()` when a long-lived +workflow should receive fair turns relative to other workflows. Fairness +policies order blocked handles; they do not create a second queue outside the +pool or bypass its back-pressure. If all permits are busy, the waiting handle +still blocks until capacity returns. + +Continue using `PooledClientLease` for explicit split-phase work such as +`send()` followed later by `receive()`. `PoolHandle` does not pin a logical +session to one socket and does not demultiplex arbitrary responses across +shared handles. + Preamble callbacks and setup hooks run per physical socket creation. Warm reuse keeps the existing preamble state; idle recycle closes the old socket and creates a fresh one, which reruns preamble and setup. diff --git a/docs/wireframe-client-design.md b/docs/wireframe-client-design.md index b2b8c367..a8a5b587 100644 --- a/docs/wireframe-client-design.md +++ b/docs/wireframe-client-design.md @@ -474,6 +474,40 @@ flowchart TD or true per-socket concurrent transport use, revisit this boundary and decide whether to extend or replace the slot wrapper. +## Decision record for 11.2.2 + +- Decision: `PoolHandle` represents the fairness identity of one logical + session, not affinity to one physical socket. +- Rationale: callers needed a durable identity so repeated pooled acquisition + by one busy workflow could not crowd out other waiting workflows. Tying the + handle to fairness instead of socket ownership preserves warm-socket reuse + without implying transport pinning. +- Scheduling model: + - `WireframeClientPool` now owns shared inner state plus a handle scheduler; + - `pool.handle()` registers one stable logical-session ID; + - blocked `PoolHandle::acquire()` calls queue against that scheduler, which + obtains a real slot permit before waking the chosen handle. +- Fairness distinction: + - slot rotation still decides which physical socket to try first across the + pool's warm connections; + - handle fairness decides which logical session receives the next lease when + contention exists. +- Policy set: + - `PoolFairnessPolicy::RoundRobin` is the default because it gives stable + logical sessions turns in rotation under repeated contention; + - `PoolFairnessPolicy::Fifo` is available when callers want blocked handles + served strictly in arrival order. +- API boundary: + - `PoolHandle` exposes fair `acquire()` plus safe whole-operation helpers + such as `call()`; + - `PooledClientLease` remains the low-level surface for split-phase workflows + (`send` followed later by `receive`) because the pool still does not + demultiplex arbitrary responses across shared handles. +- Back-pressure posture: + - handle fairness is additive above the existing slot-permit budget; + - the scheduler never bypasses `max_in_flight_per_socket`, serialized socket + access, or idle recycle behaviour inherited from `11.2.1`. + ## Future work This initial design focuses on a basic request/response workflow. Future diff --git a/src/client/mod.rs b/src/client/mod.rs index 87df9201..ec4fdb67 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -41,7 +41,13 @@ pub use hooks::{ ClientErrorHandler, }; #[cfg(feature = "pool")] -pub use pool::{ClientPoolConfig, PooledClientLease, WireframeClientPool}; +pub use pool::{ + ClientPoolConfig, + PoolFairnessPolicy, + PoolHandle, + PooledClientLease, + WireframeClientPool, +}; pub use response_stream::ResponseStream; pub use runtime::WireframeClient; pub use send_streaming::{SendStreamingConfig, SendStreamingOutcome}; diff --git a/src/client/pool/client_pool.rs b/src/client/pool/client_pool.rs index 18844803..9229ad69 100644 --- a/src/client/pool/client_pool.rs +++ b/src/client/pool/client_pool.rs @@ -15,8 +15,10 @@ use futures::{Future, future::select_all}; use super::{ ClientPoolConfig, + handle::PoolHandle, lease::PooledClientLease, manager::WireframeConnectionManager, + scheduler::PoolScheduler, slot::PoolSlot, }; use crate::{ @@ -29,6 +31,17 @@ type AcquirePermit = (Arc>, tokio::sync::OwnedSemapho type AcquirePermitFuture = Pin, ClientError>> + Send>>; +pub(crate) struct ClientPoolInner +where + S: Serializer + Clone + Send + Sync + 'static, + P: Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + pub(crate) slots: Arc<[Arc>]>, + pub(crate) next_slot: AtomicUsize, + pub(crate) scheduler: Arc>, +} + /// Pool of warm wireframe client sockets. pub struct WireframeClientPool where @@ -36,8 +49,7 @@ where P: Encode + Clone + Send + Sync + 'static, C: Send + 'static, { - slots: Arc<[Arc>]>, - next_slot: AtomicUsize, + inner: Arc>, } impl WireframeClientPool @@ -51,6 +63,7 @@ where pool_config: ClientPoolConfig, parts: ClientBuildParts, ) -> Result { + let fairness_policy = pool_config.fairness_policy_value(); let mut slots = Vec::with_capacity(pool_config.pool_size_value()); for _ in 0..pool_config.pool_size_value() { let manager = WireframeConnectionManager::new(addr, parts.clone()); @@ -68,21 +81,73 @@ where } Ok(Self { - slots: Arc::from(slots), - next_slot: AtomicUsize::new(0), + inner: Arc::new(ClientPoolInner { + slots: Arc::from(slots), + next_slot: AtomicUsize::new(0), + scheduler: Arc::new(PoolScheduler::new(fairness_policy)), + }), }) } + /// Create a stable logical-session handle governed by the pool scheduler. + /// + /// Prefer a `PoolHandle` when one logical session repeatedly acquires + /// pooled leases and should receive fair turns relative to other blocked + /// sessions. + /// + /// # Examples + /// + /// ```no_run + /// # use std::net::SocketAddr; + /// use wireframe::client::{ClientPoolConfig, PoolHandle, WireframeClient}; + /// + /// # async fn demo(addr: SocketAddr) -> Result<(), wireframe::client::ClientError> { + /// let pool = WireframeClient::builder() + /// .connect_pool(addr, ClientPoolConfig::default()) + /// .await?; + /// let mut handle: PoolHandle<_, _, _> = pool.handle(); + /// let lease = handle.acquire().await?; + /// drop(lease); + /// # Ok(()) + /// # } + /// ``` + #[must_use] + pub fn handle(&self) -> PoolHandle { + let handle_id = self.inner.scheduler.register_handle(); + PoolHandle::new(Arc::clone(&self.inner), handle_id) + } + /// Acquire one pooled-client lease. /// /// # Errors /// /// Returns [`ClientError`] if permit acquisition or pooled checkout fails. pub async fn acquire(&self) -> Result, ClientError> { - if let Some((slot, permit)) = self.try_acquire_immediately() { - return Ok(PooledClientLease::new(slot, permit)); - } + let mut handle = self.handle(); + handle.acquire().await + } + /// Close the pool and drop all warm sockets. + pub async fn close(self) { + tokio::task::yield_now().await; + drop(self); + } +} + +impl ClientPoolInner +where + S: Serializer + Clone + Send + Sync + 'static, + P: Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + pub(crate) fn try_acquire_immediately(&self) -> Option> { + self.ordered_slots().into_iter().find_map(|slot| { + slot.try_acquire_permit() + .map(|permit| PooledClientLease::new(slot, permit, None)) + }) + } + + pub(crate) async fn acquire_slot_permit(&self) -> Result, ClientError> { let waiters = self .ordered_slots() .into_iter() @@ -94,14 +159,7 @@ where }) .collect::>(); let (result, ..) = select_all(waiters).await; - let (slot, permit) = result?; - Ok(PooledClientLease::new(slot, permit)) - } - - fn try_acquire_immediately(&self) -> Option> { - self.ordered_slots() - .into_iter() - .find_map(|slot| slot.try_acquire_permit().map(|permit| (slot, permit))) + result } fn ordered_slots(&self) -> Vec>> { @@ -113,10 +171,4 @@ where } ordered } - - /// Close the pool and drop all warm sockets. - pub async fn close(self) { - tokio::task::yield_now().await; - drop(self); - } } diff --git a/src/client/pool/config.rs b/src/client/pool/config.rs index 1271ba9f..b3e448eb 100644 --- a/src/client/pool/config.rs +++ b/src/client/pool/config.rs @@ -2,6 +2,8 @@ use std::time::Duration; +use super::policy::PoolFairnessPolicy; + const DEFAULT_POOL_SIZE: usize = 4; const DEFAULT_MAX_IN_FLIGHT_PER_SOCKET: usize = 1; const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(600); @@ -13,21 +15,24 @@ const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(600); /// ``` /// use std::time::Duration; /// -/// use wireframe::client::ClientPoolConfig; +/// use wireframe::client::{ClientPoolConfig, PoolFairnessPolicy}; /// /// let config = ClientPoolConfig::default() /// .pool_size(2) /// .max_in_flight_per_socket(3) -/// .idle_timeout(Duration::from_secs(30)); +/// .idle_timeout(Duration::from_secs(30)) +/// .fairness_policy(PoolFairnessPolicy::Fifo); /// assert_eq!(config.pool_size_value(), 2); /// assert_eq!(config.max_in_flight_per_socket_value(), 3); /// assert_eq!(config.idle_timeout_value(), Duration::from_secs(30)); +/// assert_eq!(config.fairness_policy_value(), PoolFairnessPolicy::Fifo); /// ``` #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ClientPoolConfig { pool_size: usize, max_in_flight_per_socket: usize, idle_timeout: Duration, + fairness_policy: PoolFairnessPolicy, } impl Default for ClientPoolConfig { @@ -36,6 +41,7 @@ impl Default for ClientPoolConfig { pool_size: DEFAULT_POOL_SIZE, max_in_flight_per_socket: DEFAULT_MAX_IN_FLIGHT_PER_SOCKET, idle_timeout: DEFAULT_IDLE_TIMEOUT, + fairness_policy: PoolFairnessPolicy::RoundRobin, } } } @@ -68,6 +74,13 @@ impl ClientPoolConfig { self } + /// Set how blocked `PoolHandle`s are ordered when capacity returns. + #[must_use] + pub fn fairness_policy(mut self, fairness_policy: PoolFairnessPolicy) -> Self { + self.fairness_policy = fairness_policy; + self + } + /// Return the configured number of physical sockets. #[must_use] pub const fn pool_size_value(&self) -> usize { self.pool_size } @@ -79,4 +92,8 @@ impl ClientPoolConfig { /// Return the configured idle recycle timeout. #[must_use] pub const fn idle_timeout_value(&self) -> Duration { self.idle_timeout } + + /// Return the configured handle fairness policy. + #[must_use] + pub const fn fairness_policy_value(&self) -> PoolFairnessPolicy { self.fairness_policy } } diff --git a/src/client/pool/handle.rs b/src/client/pool/handle.rs new file mode 100644 index 00000000..df9729b4 --- /dev/null +++ b/src/client/pool/handle.rs @@ -0,0 +1,129 @@ +//! Logical-session handle API for fair pooled acquisition. + +use std::sync::Arc; + +use super::{client_pool::ClientPoolInner, lease::PooledClientLease}; +use crate::{ + client::ClientError, + message::{DecodeWith, EncodeWith}, + serializer::Serializer, +}; + +/// Stable logical-session identity for fair pooled acquisition. +/// +/// A `PoolHandle` lets one caller participate in the pool's fairness policy +/// over time. It is the preferred API when a logical session repeatedly +/// acquires pooled leases and you want the scheduler to account for that +/// identity explicitly. +/// +/// # Examples +/// +/// ```no_run +/// # use std::net::SocketAddr; +/// use wireframe::client::{ClientPoolConfig, PoolHandle, WireframeClient}; +/// +/// # async fn demo(addr: SocketAddr) -> Result<(), wireframe::client::ClientError> { +/// let pool = WireframeClient::builder() +/// .connect_pool(addr, ClientPoolConfig::default()) +/// .await?; +/// let mut session: PoolHandle<_, _, _> = pool.handle(); +/// let lease = session.acquire().await?; +/// drop(lease); +/// # Ok(()) +/// # } +/// ``` +pub struct PoolHandle +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + inner: Arc>, + handle_id: u64, +} + +impl PoolHandle +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + pub(crate) fn new(inner: Arc>, handle_id: u64) -> Self { + Self { inner, handle_id } + } + + /// Acquire one pooled-client lease under the configured fairness policy. + /// + /// # Examples + /// + /// ```no_run + /// # use std::net::SocketAddr; + /// use wireframe::client::{ClientPoolConfig, PoolHandle, WireframeClient}; + /// + /// # async fn demo(addr: SocketAddr) -> Result<(), wireframe::client::ClientError> { + /// let pool = WireframeClient::builder() + /// .connect_pool(addr, ClientPoolConfig::default()) + /// .await?; + /// let mut session: PoolHandle<_, _, _> = pool.handle(); + /// let lease = session.acquire().await?; + /// drop(lease); + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns [`ClientError`] if permit acquisition or pooled checkout fails. + pub async fn acquire(&mut self) -> Result, ClientError> { + self.inner + .scheduler + .acquire_for_handle(Arc::clone(&self.inner), self.handle_id) + .await + } + + /// Perform a complete request-response round trip through a fair lease. + /// + /// # Examples + /// + /// ```no_run + /// # use std::net::SocketAddr; + /// use wireframe::client::{ClientPoolConfig, PoolHandle, WireframeClient}; + /// + /// #[derive(bincode::Encode, bincode::Decode)] + /// struct Ping(u8); + /// + /// #[derive(bincode::Encode, bincode::Decode)] + /// struct Pong(u8); + /// + /// # async fn demo(addr: SocketAddr) -> Result<(), wireframe::client::ClientError> { + /// let pool = WireframeClient::builder() + /// .connect_pool(addr, ClientPoolConfig::default()) + /// .await?; + /// let mut session: PoolHandle<_, _, _> = pool.handle(); + /// let _reply: Pong = session.call(&Ping(1)).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns [`ClientError`] if lease acquisition, serialization, decode, or + /// transport I/O fails. + pub async fn call(&mut self, request: &Req) -> Result + where + Req: EncodeWith, + Resp: DecodeWith, + { + let lease = self.acquire().await?; + lease.call(request).await + } +} + +impl Drop for PoolHandle +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + fn drop(&mut self) { self.inner.scheduler.deregister_handle(self.handle_id); } +} diff --git a/src/client/pool/lease.rs b/src/client/pool/lease.rs index 98a14459..780e2a5b 100644 --- a/src/client/pool/lease.rs +++ b/src/client/pool/lease.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use tokio::sync::OwnedSemaphorePermit; -use super::slot::PoolSlot; +use super::{client_pool::ClientPoolInner, slot::PoolSlot}; use crate::{ app::Packet, client::ClientError, @@ -25,6 +25,7 @@ where { slot: Arc>, _permit: OwnedSemaphorePermit, + release_inner: Option>>, } macro_rules! dispatch_on_connection { @@ -46,10 +47,15 @@ where P: bincode::Encode + Clone + Send + Sync + 'static, C: Send + 'static, { - pub(crate) fn new(slot: Arc>, permit: OwnedSemaphorePermit) -> Self { + pub(crate) fn new( + slot: Arc>, + permit: OwnedSemaphorePermit, + release_inner: Option>>, + ) -> Self { Self { slot, _permit: permit, + release_inner, } } @@ -126,3 +132,16 @@ where dispatch_on_connection!(self, |conn| conn.call_correlated(request)) } } + +impl Drop for PooledClientLease +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + fn drop(&mut self) { + if let Some(inner) = &self.release_inner { + inner.scheduler.notify_capacity_available(Arc::clone(inner)); + } + } +} diff --git a/src/client/pool/mod.rs b/src/client/pool/mod.rs index cfdb5c31..c6213098 100644 --- a/src/client/pool/mod.rs +++ b/src/client/pool/mod.rs @@ -6,11 +6,16 @@ mod client_pool; mod config; +mod handle; mod lease; mod managed; mod manager; +mod policy; +mod scheduler; mod slot; pub use client_pool::WireframeClientPool; pub use config::ClientPoolConfig; +pub use handle::PoolHandle; pub use lease::PooledClientLease; +pub use policy::PoolFairnessPolicy; diff --git a/src/client/pool/policy.rs b/src/client/pool/policy.rs new file mode 100644 index 00000000..bbb6cbf3 --- /dev/null +++ b/src/client/pool/policy.rs @@ -0,0 +1,30 @@ +//! Public fairness policies for pooled client handles. + +/// Fairness policy used when many `PoolHandle`s contend for pooled leases. +/// +/// `RoundRobin` is the default because it gives stable logical sessions turns +/// in rotation. `Fifo` preserves the exact order in which blocked acquisitions +/// arrived. +/// +/// # Examples +/// +/// ``` +/// use wireframe::client::{ClientPoolConfig, PoolFairnessPolicy}; +/// +/// let config = ClientPoolConfig::default() +/// .pool_size(2) +/// .fairness_policy(PoolFairnessPolicy::RoundRobin); +/// +/// assert_eq!( +/// config.fairness_policy_value(), +/// PoolFairnessPolicy::RoundRobin +/// ); +/// ``` +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum PoolFairnessPolicy { + /// Rotate grants across stable logical-session handles. + #[default] + RoundRobin, + /// Serve waiting handles in the order they blocked. + Fifo, +} diff --git a/src/client/pool/scheduler.rs b/src/client/pool/scheduler.rs new file mode 100644 index 00000000..6aad73a4 --- /dev/null +++ b/src/client/pool/scheduler.rs @@ -0,0 +1,213 @@ +//! Fair lease scheduling for pooled client handles. + +use std::{ + collections::{HashMap, VecDeque}, + sync::{ + Arc, + Mutex, + MutexGuard, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, +}; + +use tokio::sync::oneshot; + +use super::{client_pool::ClientPoolInner, lease::PooledClientLease, policy::PoolFairnessPolicy}; +use crate::{client::ClientError, serializer::Serializer}; + +fn recover_mutex(mutex: &Mutex) -> MutexGuard<'_, T> { + match mutex.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } +} + +type WaiterSender = oneshot::Sender, ClientError>>; + +struct SchedulerState +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + waiters: HashMap>, + fifo_waiters: VecDeque, + round_robin_handles: VecDeque, +} + +impl SchedulerState +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + fn new() -> Self { + Self { + waiters: HashMap::new(), + fifo_waiters: VecDeque::new(), + round_robin_handles: VecDeque::new(), + } + } + + fn register_handle(&mut self, handle_id: u64) { self.round_robin_handles.push_back(handle_id); } + + fn deregister_handle(&mut self, handle_id: u64) { + self.waiters.remove(&handle_id); + self.fifo_waiters + .retain(|queued_id| *queued_id != handle_id); + self.round_robin_handles + .retain(|queued_id| *queued_id != handle_id); + } + + fn enqueue_waiter(&mut self, handle_id: u64, sender: WaiterSender) { + self.waiters.insert(handle_id, sender); + self.fifo_waiters.push_back(handle_id); + } + + fn has_waiters(&self) -> bool { !self.waiters.is_empty() } + + fn take_next_waiter(&mut self, policy: PoolFairnessPolicy) -> Option> { + match policy { + PoolFairnessPolicy::RoundRobin => self.take_next_round_robin_waiter(), + PoolFairnessPolicy::Fifo => self.take_next_fifo_waiter(), + } + } + + fn take_next_fifo_waiter(&mut self) -> Option> { + while let Some(handle_id) = self.fifo_waiters.pop_front() { + if let Some(sender) = self.waiters.remove(&handle_id) { + return Some(sender); + } + } + None + } + + fn take_next_round_robin_waiter(&mut self) -> Option> { + let len = self.round_robin_handles.len(); + for _ in 0..len { + let handle_id = self.round_robin_handles.pop_front()?; + self.round_robin_handles.push_back(handle_id); + if let Some(sender) = self.waiters.remove(&handle_id) { + self.fifo_waiters + .retain(|queued_id| *queued_id != handle_id); + return Some(sender); + } + } + None + } +} + +/// Shared fairness scheduler used by pooled handles. +pub(crate) struct PoolScheduler +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + fairness_policy: PoolFairnessPolicy, + next_handle_id: AtomicU64, + is_servicing: AtomicBool, + state: Mutex>, +} + +impl PoolScheduler +where + S: Serializer + Clone + Send + Sync + 'static, + P: bincode::Encode + Clone + Send + Sync + 'static, + C: Send + 'static, +{ + pub(crate) fn new(fairness_policy: PoolFairnessPolicy) -> Self { + Self { + fairness_policy, + next_handle_id: AtomicU64::new(1), + is_servicing: AtomicBool::new(false), + state: Mutex::new(SchedulerState::new()), + } + } + + pub(crate) fn register_handle(&self) -> u64 { + let handle_id = self.next_handle_id.fetch_add(1, Ordering::Relaxed); + recover_mutex(&self.state).register_handle(handle_id); + handle_id + } + + pub(crate) fn deregister_handle(&self, handle_id: u64) { + recover_mutex(&self.state).deregister_handle(handle_id); + } + + pub(crate) async fn acquire_for_handle( + self: &Arc, + inner: Arc>, + handle_id: u64, + ) -> Result, ClientError> { + if let Some(lease) = inner.try_acquire_immediately() { + return Ok(lease); + } + + let (sender, receiver) = oneshot::channel(); + recover_mutex(&self.state).enqueue_waiter(handle_id, sender); + self.kick(inner); + receiver.await.map_err(|_| ClientError::disconnected())? + } + + pub(crate) fn notify_capacity_available( + self: &Arc, + inner: Arc>, + ) { + self.kick(inner); + } + + fn kick(self: &Arc, inner: Arc>) { + if self + .is_servicing + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + let scheduler = Arc::clone(self); + tokio::spawn(async move { + scheduler.service_waiters(inner).await; + }); + } + } + + fn restart_if_waiters(&self) -> bool { + if !recover_mutex(&self.state).has_waiters() { + return false; + } + + self.is_servicing + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + } + + fn take_next_waiter_or_stop(&self) -> Option> { + loop { + if let Some(sender) = recover_mutex(&self.state).take_next_waiter(self.fairness_policy) + { + return Some(sender); + } + + self.is_servicing.store(false, Ordering::Release); + if !self.restart_if_waiters() { + return None; + } + } + } + + async fn service_waiters(self: Arc, inner: Arc>) { + loop { + let Some(sender) = self.take_next_waiter_or_stop() else { + break; + }; + + let result = inner.acquire_slot_permit().await.map(|(slot, permit)| { + PooledClientLease::new(slot, permit, Some(Arc::clone(&inner))) + }); + if let Err(send_result) = sender.send(result) + && let Ok(lease) = send_result + { + drop(lease); + } + } + } +} diff --git a/src/client/tests/mod.rs b/src/client/tests/mod.rs index b986b7a8..17226bc0 100644 --- a/src/client/tests/mod.rs +++ b/src/client/tests/mod.rs @@ -6,6 +6,8 @@ mod lifecycle; mod messaging; #[cfg(feature = "pool")] mod pool; +#[cfg(feature = "pool")] +mod pool_handle; mod request_hooks; mod send_streaming; mod send_streaming_infra; diff --git a/src/client/tests/pool_handle.rs b/src/client/tests/pool_handle.rs new file mode 100644 index 00000000..174c25da --- /dev/null +++ b/src/client/tests/pool_handle.rs @@ -0,0 +1,200 @@ +//! Unit tests for the `PoolHandle` fairness API. + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use rstest::{fixture, rstest}; +use tokio::{ + sync::Mutex, + time::{advance, timeout}, +}; + +use crate::{ + client::{ClientError, ClientPoolConfig, PoolFairnessPolicy, PoolHandle}, + test_helpers::{Ping, Pong, PoolTestServer, TestClientPool, build_pooled_client}, +}; + +type TestResult = Result<(), Box>; + +#[rustfmt::skip] +#[fixture] +fn client_pool_config() -> ClientPoolConfig { + ClientPoolConfig::default() +} + +async fn build_handle_pool( + config: ClientPoolConfig, +) -> Result<(PoolTestServer, TestClientPool), ClientError> { + let server = PoolTestServer::start().await?; + let pool = build_pooled_client(server.addr, config, Arc::new(AtomicUsize::new(0))).await?; + Ok((server, pool)) +} + +async fn acquire_and_record( + mut handle: PoolHandle< + crate::serializer::BincodeSerializer, + crate::test_helpers::ClientHello, + (), + >, + label: &'static str, + rounds: usize, + grants: Arc>>, +) -> Result<(), ClientError> { + for _ in 0..rounds { + let lease = handle.acquire().await?; + grants.lock().await.push(label); + tokio::task::yield_now().await; + drop(lease); + } + Ok(()) +} + +#[rstest] +#[tokio::test(flavor = "current_thread")] +async fn round_robin_handles_share_one_socket_fairly( + client_pool_config: ClientPoolConfig, +) -> TestResult { + let (_server, pool) = build_handle_pool( + client_pool_config + .pool_size(1) + .max_in_flight_per_socket(1) + .fairness_policy(PoolFairnessPolicy::RoundRobin), + ) + .await?; + let grants = Arc::new(Mutex::new(Vec::new())); + let first = pool.handle(); + let second = pool.handle(); + + let left = tokio::spawn(acquire_and_record(first, "a", 3, Arc::clone(&grants))); + let right = tokio::spawn(acquire_and_record(second, "b", 3, Arc::clone(&grants))); + let (left_result, right_result) = tokio::join!(left, right); + left_result??; + right_result??; + + let observed = grants.lock().await.clone(); + assert_eq!(observed, vec!["a", "b", "a", "b", "a", "b"]); + Ok(()) +} + +#[rstest] +#[tokio::test(flavor = "current_thread")] +async fn fifo_policy_preserves_wait_order(client_pool_config: ClientPoolConfig) -> TestResult { + let (_server, pool) = build_handle_pool( + client_pool_config + .pool_size(1) + .max_in_flight_per_socket(1) + .fairness_policy(PoolFairnessPolicy::Fifo), + ) + .await?; + let blocker = pool.acquire().await?; + let grants = Arc::new(Mutex::new(Vec::new())); + + let first = tokio::spawn(acquire_and_record( + pool.handle(), + "first", + 1, + Arc::clone(&grants), + )); + tokio::task::yield_now().await; + let second = tokio::spawn(acquire_and_record( + pool.handle(), + "second", + 1, + Arc::clone(&grants), + )); + tokio::task::yield_now().await; + let third = tokio::spawn(acquire_and_record( + pool.handle(), + "third", + 1, + Arc::clone(&grants), + )); + tokio::task::yield_now().await; + + drop(blocker); + first.await??; + second.await??; + third.await??; + + let observed = grants.lock().await.clone(); + assert_eq!(observed, vec!["first", "second", "third"]); + Ok(()) +} + +#[rstest] +#[tokio::test(flavor = "current_thread")] +async fn handle_acquire_respects_back_pressure(client_pool_config: ClientPoolConfig) -> TestResult { + let (_server, pool) = build_handle_pool(client_pool_config.pool_size(1)).await?; + let mut first = pool.handle(); + let mut second = pool.handle(); + + let held_lease = first.acquire().await?; + let blocked = timeout(Duration::from_millis(25), second.acquire()).await; + assert!(blocked.is_err(), "second handle should stay blocked"); + + drop(held_lease); + let recovered = timeout(Duration::from_millis(100), second.acquire()).await?; + let _recovered = recovered?; + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn handle_path_preserves_warm_reuse_and_preamble( + client_pool_config: ClientPoolConfig, +) -> TestResult { + let preamble_callback_count = Arc::new(AtomicUsize::new(0)); + let server = PoolTestServer::start().await?; + let pool = build_pooled_client( + server.addr, + client_pool_config.pool_size(1), + preamble_callback_count.clone(), + ) + .await?; + let mut handle = pool.handle(); + + let first: Pong = handle.call(&Ping(7)).await?; + let second: Pong = handle.call(&Ping(8)).await?; + + assert_eq!(first, Pong(7)); + assert_eq!(second, Pong(8)); + assert_eq!(preamble_callback_count.load(Ordering::SeqCst), 1); + assert_eq!(server.preamble_count(), 1); + assert_eq!(server.connection_count(), 1); + Ok(()) +} + +#[rstest] +#[tokio::test(start_paused = true, flavor = "current_thread")] +async fn handle_path_recycles_after_idle_timeout( + client_pool_config: ClientPoolConfig, +) -> TestResult { + let preamble_callback_count = Arc::new(AtomicUsize::new(0)); + let server = PoolTestServer::start().await?; + let idle_timeout = Duration::from_millis(50); + let pool = build_pooled_client( + server.addr, + client_pool_config.pool_size(1).idle_timeout(idle_timeout), + preamble_callback_count.clone(), + ) + .await?; + let mut handle = pool.handle(); + + let first: Pong = handle.call(&Ping(1)).await?; + assert_eq!(first, Pong(1)); + + advance(idle_timeout + idle_timeout).await; + tokio::task::yield_now().await; + + let second: Pong = handle.call(&Ping(2)).await?; + assert_eq!(second, Pong(2)); + assert_eq!(preamble_callback_count.load(Ordering::SeqCst), 2); + assert_eq!(server.preamble_count(), 2); + assert_eq!(server.connection_count(), 2); + Ok(()) +} diff --git a/tests/bdd_pool/fixtures.rs b/tests/bdd_pool/fixtures.rs index 5ca72fc1..f801ed6c 100644 --- a/tests/bdd_pool/fixtures.rs +++ b/tests/bdd_pool/fixtures.rs @@ -2,3 +2,6 @@ #[path = "../fixtures/client_pool.rs"] pub mod client_pool; + +#[path = "../fixtures/client_pool_handle.rs"] +pub mod client_pool_handle; diff --git a/tests/bdd_pool/scenarios.rs b/tests/bdd_pool/scenarios.rs index 94b0abb5..0f1d7fbb 100644 --- a/tests/bdd_pool/scenarios.rs +++ b/tests/bdd_pool/scenarios.rs @@ -3,5 +3,11 @@ #[path = "../steps/client_pool_steps.rs"] mod client_pool_steps; +#[path = "../steps/client_pool_handle_steps.rs"] +mod client_pool_handle_steps; + #[path = "../scenarios/client_pool_scenarios.rs"] mod client_pool_scenarios; + +#[path = "../scenarios/client_pool_handle_scenarios.rs"] +mod client_pool_handle_scenarios; diff --git a/tests/features/client_pool_handle.feature b/tests/features/client_pool_handle.feature new file mode 100644 index 00000000..c4e1bc71 --- /dev/null +++ b/tests/features/client_pool_handle.feature @@ -0,0 +1,23 @@ +Feature: Client pool handles + Pool handles represent logical sessions so the pooled client can order + blocked acquisition fairly without bypassing back-pressure. + + Scenario: Two logical sessions alternate fairly on one pooled socket + Given client pool handle round-robin fairness is configured + When two logical sessions repeatedly acquire one pooled socket + Then the logical sessions alternate fairly + + Scenario: FIFO fairness preserves waiting order + Given client pool handle FIFO fairness is configured + When three logical sessions wait in order on one pooled socket + Then the logical sessions are served in arrival order + + Scenario: Waiting handle remains blocked until a lease is released + Given client pool handle back-pressure is configured + When one logical session holds the only pooled lease + Then another logical session stays blocked until capacity returns + + Scenario: Handle warm reuse and idle recycle match pool behaviour + Given client pool handle warm reuse and idle recycle are configured + When one logical session reuses a warm socket and then idles past timeout + Then the handle preserves warm reuse and later reconnects after idle diff --git a/tests/fixtures/client_pool_handle.rs b/tests/fixtures/client_pool_handle.rs new file mode 100644 index 00000000..8a833b25 --- /dev/null +++ b/tests/fixtures/client_pool_handle.rs @@ -0,0 +1,239 @@ +//! `ClientPoolHandleWorld` fixture for `PoolHandle` behavioural scenarios. + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use rstest::fixture; +use tokio::{sync::Mutex, time::timeout}; +use wireframe::{ + client::{ClientPoolConfig, PoolFairnessPolicy}, + test_helpers::{Ping, Pong, PoolTestServer, TestClientPool, build_pooled_client}, +}; +pub use wireframe_testing::TestResult; + +pub struct ClientPoolHandleWorld { + pool: Option, + server: Option, + preamble_callback_count: Arc, + grant_order: Vec, + blocked_waiter: bool, + waiter_recovered: bool, + warm_reuse_then_recycle: bool, +} + +impl Default for ClientPoolHandleWorld { + fn default() -> Self { + Self { + pool: None, + server: None, + preamble_callback_count: Arc::new(AtomicUsize::new(0)), + grant_order: Vec::new(), + blocked_waiter: false, + waiter_recovered: false, + warm_reuse_then_recycle: false, + } + } +} + +#[rustfmt::skip] +#[fixture] +pub fn client_pool_handle_world() -> ClientPoolHandleWorld { + ClientPoolHandleWorld::default() +} + +impl ClientPoolHandleWorld { + async fn connect_handle_pool(&mut self, fairness_policy: PoolFairnessPolicy) -> TestResult { + self.start_server().await?; + self.connect_pool( + ClientPoolConfig::default() + .pool_size(1) + .fairness_policy(fairness_policy), + ) + .await + } + + async fn record_grants(&mut self, grants: Arc>>) -> TestResult { + self.grant_order = grants + .lock() + .await + .iter() + .map(ToString::to_string) + .collect(); + Ok(()) + } + + async fn start_server(&mut self) -> TestResult { + self.server = Some(PoolTestServer::start().await?); + Ok(()) + } + + async fn connect_pool(&mut self, config: ClientPoolConfig) -> TestResult { + let addr = self.server.as_ref().ok_or("server missing")?.addr; + let pool = build_pooled_client(addr, config, self.preamble_callback_count.clone()).await?; + self.pool = Some(pool); + Ok(()) + } + + pub async fn run_round_robin_scenario(&mut self) -> TestResult { + self.connect_handle_pool(PoolFairnessPolicy::RoundRobin) + .await?; + let pool = self.pool.as_ref().ok_or("pool missing")?; + let grants = Arc::new(Mutex::new(Vec::new())); + + let left = tokio::spawn(run_repeated_acquire( + pool.handle(), + "session-a", + 3, + Arc::clone(&grants), + )); + let right = tokio::spawn(run_repeated_acquire( + pool.handle(), + "session-b", + 3, + Arc::clone(&grants), + )); + left.await??; + right.await??; + self.record_grants(grants).await + } + + pub async fn run_fifo_scenario(&mut self) -> TestResult { + self.connect_handle_pool(PoolFairnessPolicy::Fifo).await?; + let pool = self.pool.as_ref().ok_or("pool missing")?; + let blocker = pool.acquire().await?; + let grants = Arc::new(Mutex::new(Vec::new())); + let (first, second, third) = spawn_fifo_waiters(pool, Arc::clone(&grants)).await; + + drop(blocker); + first.await??; + second.await??; + third.await??; + self.record_grants(grants).await + } + + pub async fn run_back_pressure_scenario(&mut self) -> TestResult { + self.start_server().await?; + self.connect_pool(ClientPoolConfig::default().pool_size(1)) + .await?; + let pool = self.pool.as_ref().ok_or("pool missing")?; + let mut first = pool.handle(); + let mut second = pool.handle(); + + let held_lease = first.acquire().await?; + self.blocked_waiter = timeout(Duration::from_millis(25), second.acquire()) + .await + .is_err(); + + drop(held_lease); + self.waiter_recovered = matches!( + timeout(Duration::from_millis(100), second.acquire()).await, + Ok(Ok(_)) + ); + Ok(()) + } + + pub async fn run_warm_reuse_then_idle_recycle_scenario(&mut self) -> TestResult { + tokio::time::pause(); + self.start_server().await?; + let idle_timeout = Duration::from_millis(50); + self.connect_pool( + ClientPoolConfig::default() + .pool_size(1) + .idle_timeout(idle_timeout), + ) + .await?; + let pool = self.pool.as_ref().ok_or("pool missing")?; + let mut handle = pool.handle(); + + let first: Pong = handle.call(&Ping(1)).await?; + let second: Pong = handle.call(&Ping(2)).await?; + if first != Pong(1) || second != Pong(2) { + return Err("unexpected warm reuse response sequence".into()); + } + let warm_reuse_preserved = self.preamble_callback_count.load(Ordering::SeqCst) == 1 + && self.server.as_ref().is_some_and(|server| { + server.preamble_count() == 1 && server.connection_count() == 1 + }); + + tokio::time::advance(idle_timeout + idle_timeout).await; + tokio::task::yield_now().await; + + let third: Pong = handle.call(&Ping(3)).await?; + self.warm_reuse_then_recycle = third == Pong(3) + && warm_reuse_preserved + && self.preamble_callback_count.load(Ordering::SeqCst) == 2 + && self.server.as_ref().is_some_and(|server| { + server.preamble_count() == 2 && server.connection_count() == 2 + }); + Ok(()) + } + + pub fn sessions_alternate_fairly(&self) -> bool { + self.grant_order + == [ + "session-a", + "session-b", + "session-a", + "session-b", + "session-a", + "session-b", + ] + } + + pub fn fifo_order_preserved(&self) -> bool { self.grant_order == ["first", "second", "third"] } + + pub fn back_pressure_preserved(&self) -> bool { self.blocked_waiter && self.waiter_recovered } + + pub fn warm_reuse_then_recycle_preserved(&self) -> bool { self.warm_reuse_then_recycle } +} + +async fn spawn_fifo_waiters( + pool: &TestClientPool, + grants: Arc>>, +) -> ( + tokio::task::JoinHandle>, + tokio::task::JoinHandle>, + tokio::task::JoinHandle>, +) { + let first = tokio::spawn(run_repeated_acquire( + pool.handle(), + "first", + 1, + Arc::clone(&grants), + )); + tokio::task::yield_now().await; + let second = tokio::spawn(run_repeated_acquire( + pool.handle(), + "second", + 1, + Arc::clone(&grants), + )); + tokio::task::yield_now().await; + let third = tokio::spawn(run_repeated_acquire(pool.handle(), "third", 1, grants)); + tokio::task::yield_now().await; + (first, second, third) +} + +async fn run_repeated_acquire( + mut handle: wireframe::client::PoolHandle< + wireframe::serializer::BincodeSerializer, + wireframe::test_helpers::ClientHello, + (), + >, + label: &'static str, + rounds: usize, + grants: Arc>>, +) -> Result<(), wireframe::client::ClientError> { + for _ in 0..rounds { + let lease = handle.acquire().await?; + grants.lock().await.push(label); + tokio::task::yield_now().await; + drop(lease); + } + Ok(()) +} diff --git a/tests/scenarios/client_pool_handle_scenarios.rs b/tests/scenarios/client_pool_handle_scenarios.rs new file mode 100644 index 00000000..05f8df6b --- /dev/null +++ b/tests/scenarios/client_pool_handle_scenarios.rs @@ -0,0 +1,37 @@ +//! Scenario tests for `PoolHandle` behaviours. + +use rstest_bdd_macros::scenario; + +use crate::fixtures::client_pool_handle::*; + +#[scenario( + path = "tests/features/client_pool_handle.feature", + name = "Two logical sessions alternate fairly on one pooled socket" +)] +fn client_pool_handle_round_robin_fairness(client_pool_handle_world: ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} + +#[scenario( + path = "tests/features/client_pool_handle.feature", + name = "FIFO fairness preserves waiting order" +)] +fn client_pool_handle_fifo_fairness(client_pool_handle_world: ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} + +#[scenario( + path = "tests/features/client_pool_handle.feature", + name = "Waiting handle remains blocked until a lease is released" +)] +fn client_pool_handle_back_pressure(client_pool_handle_world: ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} + +#[scenario( + path = "tests/features/client_pool_handle.feature", + name = "Handle warm reuse and idle recycle match pool behaviour" +)] +fn client_pool_handle_reuse_and_recycle(client_pool_handle_world: ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} diff --git a/tests/steps/client_pool_handle_steps.rs b/tests/steps/client_pool_handle_steps.rs new file mode 100644 index 00000000..f27fc7c9 --- /dev/null +++ b/tests/steps/client_pool_handle_steps.rs @@ -0,0 +1,111 @@ +//! Step definitions for `PoolHandle` behavioural tests. + +use rstest_bdd_macros::{given, then, when}; + +use crate::fixtures::client_pool_handle::{ClientPoolHandleWorld, TestResult}; + +#[given("client pool handle round-robin fairness is configured")] +fn given_client_pool_handle_round_robin(client_pool_handle_world: &mut ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} + +#[given("client pool handle FIFO fairness is configured")] +fn given_client_pool_handle_fifo(client_pool_handle_world: &mut ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} + +#[given("client pool handle back-pressure is configured")] +fn given_client_pool_handle_back_pressure(client_pool_handle_world: &mut ClientPoolHandleWorld) { + let _ = client_pool_handle_world; +} + +#[given("client pool handle warm reuse and idle recycle are configured")] +fn given_client_pool_handle_reuse_and_recycle( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) { + let _ = client_pool_handle_world; +} + +#[when("two logical sessions repeatedly acquire one pooled socket")] +fn when_two_logical_sessions_repeat( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + rt.block_on(client_pool_handle_world.run_round_robin_scenario()) +} + +#[when("three logical sessions wait in order on one pooled socket")] +fn when_three_logical_sessions_wait_in_order( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + rt.block_on(client_pool_handle_world.run_fifo_scenario()) +} + +#[when("one logical session holds the only pooled lease")] +fn when_one_logical_session_holds_the_only_lease( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + rt.block_on(client_pool_handle_world.run_back_pressure_scenario()) +} + +#[when("one logical session reuses a warm socket and then idles past timeout")] +fn when_one_logical_session_reuses_then_idles( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + rt.block_on(client_pool_handle_world.run_warm_reuse_then_idle_recycle_scenario()) +} + +#[then("the logical sessions alternate fairly")] +fn then_the_logical_sessions_alternate_fairly( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + if client_pool_handle_world.sessions_alternate_fairly() { + Ok(()) + } else { + Err("expected logical sessions to alternate fairly".into()) + } +} + +#[then("the logical sessions are served in arrival order")] +fn then_the_logical_sessions_are_served_in_arrival_order( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + if client_pool_handle_world.fifo_order_preserved() { + Ok(()) + } else { + Err("expected FIFO fairness to preserve arrival order".into()) + } +} + +#[then("another logical session stays blocked until capacity returns")] +fn then_another_logical_session_stays_blocked( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + if client_pool_handle_world.back_pressure_preserved() { + Ok(()) + } else { + Err("expected waiting logical session to remain blocked".into()) + } +} + +#[then("the handle preserves warm reuse and later reconnects after idle")] +fn then_the_handle_preserves_warm_reuse_then_recycles( + client_pool_handle_world: &mut ClientPoolHandleWorld, +) -> TestResult { + if client_pool_handle_world.warm_reuse_then_recycle_preserved() { + Ok(()) + } else { + Err("expected handle warm reuse to preserve preamble and later recycle".into()) + } +} From 13a635c3269169eb8e9e3f36664d4e482457ed87 Mon Sep 17 00:00:00 2001 From: Leynos Date: Sun, 15 Mar 2026 04:39:47 +0000 Subject: [PATCH 3/5] test(client-pool-handle): refactor tests to use helper for preamble count tracking Refactored pool_handle tests to introduce a helper function build_preamble_pool that constructs the test pool along with the preamble callback counter. Updated existing tests handle_path_preserves_warm_reuse_and_preamble and handle_path_recycles_after_idle_timeout to use the new helper. Added helper methods in client_pool_handle fixture to simplify assertions on preamble callback counts and server state. Also fixed minor markdown formatting in related docs. Co-authored-by: devboxerhub[bot] --- .../11-2-2-expose-pool-handle-api.md | 4 ++- src/client/tests/pool_handle.rs | 29 +++++++++---------- tests/fixtures/client_pool_handle.rs | 22 +++++++++----- 3 files changed, 30 insertions(+), 25 deletions(-) diff --git a/docs/execplans/11-2-2-expose-pool-handle-api.md b/docs/execplans/11-2-2-expose-pool-handle-api.md index 0920c4d7..83e51772 100644 --- a/docs/execplans/11-2-2-expose-pool-handle-api.md +++ b/docs/execplans/11-2-2-expose-pool-handle-api.md @@ -480,7 +480,9 @@ Observable completion criteria: - `set -o pipefail && make test | tee /tmp/wireframe-11-2-2-test.log` - `set -o pipefail && make test-doc | tee /tmp/wireframe-11-2-2-test-doc.log` - `set -o pipefail && make doctest-benchmark | tee /tmp/wireframe-11-2-2-doctest-benchmark.log` - - `set -o pipefail && make markdownlint MDLINT=/root/.bun/bin/markdownlint-cli2 | tee /tmp/wireframe-11-2-2-markdownlint.log` + - `set -o pipefail && make markdownlint + MDLINT=/root/.bun/bin/markdownlint-cli2 | tee + /tmp/wireframe-11-2-2-markdownlint.log` - `set -o pipefail && make nixie | tee /tmp/wireframe-11-2-2-nixie.log` - Lessons learned: - fairness is easiest to explain when the scheduler operates on logical diff --git a/src/client/tests/pool_handle.rs b/src/client/tests/pool_handle.rs index 174c25da..d07bffeb 100644 --- a/src/client/tests/pool_handle.rs +++ b/src/client/tests/pool_handle.rs @@ -35,6 +35,15 @@ async fn build_handle_pool( Ok((server, pool)) } +async fn build_preamble_pool( + config: ClientPoolConfig, +) -> Result<(PoolTestServer, TestClientPool, Arc), ClientError> { + let preamble_callback_count = Arc::new(AtomicUsize::new(0)); + let server = PoolTestServer::start().await?; + let pool = build_pooled_client(server.addr, config, preamble_callback_count.clone()).await?; + Ok((server, pool, preamble_callback_count)) +} + async fn acquire_and_record( mut handle: PoolHandle< crate::serializer::BincodeSerializer, @@ -148,14 +157,8 @@ async fn handle_acquire_respects_back_pressure(client_pool_config: ClientPoolCon async fn handle_path_preserves_warm_reuse_and_preamble( client_pool_config: ClientPoolConfig, ) -> TestResult { - let preamble_callback_count = Arc::new(AtomicUsize::new(0)); - let server = PoolTestServer::start().await?; - let pool = build_pooled_client( - server.addr, - client_pool_config.pool_size(1), - preamble_callback_count.clone(), - ) - .await?; + let (server, pool, preamble_callback_count) = + build_preamble_pool(client_pool_config.pool_size(1)).await?; let mut handle = pool.handle(); let first: Pong = handle.call(&Ping(7)).await?; @@ -174,15 +177,9 @@ async fn handle_path_preserves_warm_reuse_and_preamble( async fn handle_path_recycles_after_idle_timeout( client_pool_config: ClientPoolConfig, ) -> TestResult { - let preamble_callback_count = Arc::new(AtomicUsize::new(0)); - let server = PoolTestServer::start().await?; let idle_timeout = Duration::from_millis(50); - let pool = build_pooled_client( - server.addr, - client_pool_config.pool_size(1).idle_timeout(idle_timeout), - preamble_callback_count.clone(), - ) - .await?; + let (server, pool, preamble_callback_count) = + build_preamble_pool(client_pool_config.pool_size(1).idle_timeout(idle_timeout)).await?; let mut handle = pool.handle(); let first: Pong = handle.call(&Ping(1)).await?; diff --git a/tests/fixtures/client_pool_handle.rs b/tests/fixtures/client_pool_handle.rs index 8a833b25..c6fd2b8c 100644 --- a/tests/fixtures/client_pool_handle.rs +++ b/tests/fixtures/client_pool_handle.rs @@ -79,6 +79,16 @@ impl ClientPoolHandleWorld { Ok(()) } + fn pool_preamble_count_is(&self, expected: usize) -> bool { + self.preamble_callback_count.load(Ordering::SeqCst) == expected + } + + fn server_state_matches(&self, preamble_count: usize, connection_count: usize) -> bool { + self.server.as_ref().is_some_and(|s| { + s.preamble_count() == preamble_count && s.connection_count() == connection_count + }) + } + pub async fn run_round_robin_scenario(&mut self) -> TestResult { self.connect_handle_pool(PoolFairnessPolicy::RoundRobin) .await?; @@ -155,10 +165,8 @@ impl ClientPoolHandleWorld { if first != Pong(1) || second != Pong(2) { return Err("unexpected warm reuse response sequence".into()); } - let warm_reuse_preserved = self.preamble_callback_count.load(Ordering::SeqCst) == 1 - && self.server.as_ref().is_some_and(|server| { - server.preamble_count() == 1 && server.connection_count() == 1 - }); + let warm_reuse_preserved = + self.pool_preamble_count_is(1) && self.server_state_matches(1, 1); tokio::time::advance(idle_timeout + idle_timeout).await; tokio::task::yield_now().await; @@ -166,10 +174,8 @@ impl ClientPoolHandleWorld { let third: Pong = handle.call(&Ping(3)).await?; self.warm_reuse_then_recycle = third == Pong(3) && warm_reuse_preserved - && self.preamble_callback_count.load(Ordering::SeqCst) == 2 - && self.server.as_ref().is_some_and(|server| { - server.preamble_count() == 2 && server.connection_count() == 2 - }); + && self.pool_preamble_count_is(2) + && self.server_state_matches(2, 2); Ok(()) } From 6bcbcd08c9c179d151a9371d3741d0d607660447 Mon Sep 17 00:00:00 2001 From: Leynos Date: Mon, 16 Mar 2026 22:07:30 +0000 Subject: [PATCH 4/5] feat(client/pool): add graceful shutdown support for connection pool This commit introduces a shutdown mechanism to the client connection pool to prevent new leases from being acquired after shutdown. It adds an AtomicBool flag to track pool shutdown state, modifies acquisition logic to respect shutdown status, and enhances the scheduler to notify all waiting tasks with disconnection errors during shutdown. Additionally, new tests were added to verify that dropped waiters do not leak capacity and that shutdown behavior is correctly handled. Minor documentation improvements and test helper refactors accompany these changes. Co-authored-by: devboxerhub[bot] --- docs/roadmap.md | 2 +- docs/users-guide.md | 3 +- src/client/pool/client_pool.rs | 19 +++++- src/client/pool/scheduler.rs | 29 +++++++-- src/client/tests/pool_handle.rs | 97 ++++++++++++++++++---------- src/test_helpers.rs | 2 + src/test_helpers/pool_client.rs | 36 ++++++++++- tests/fixtures/client_pool_handle.rs | 38 ++++------- 8 files changed, 154 insertions(+), 72 deletions(-) diff --git a/docs/roadmap.md b/docs/roadmap.md index aa47d811..73c53c70 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -502,7 +502,7 @@ document so larger deployments can adopt the library confidently. - [x] 11.2.1. Implement a configurable connection pool that preserves preamble state, enforces in-flight request limits per socket, and recycles idle connections. -- [x] 11.2.2. Expose a `PoolHandle` API with fairness policies so callers can +- [x] 11.2.2. Expose a `PoolHandle` API with fairness policies, so callers can multiplex many logical sessions without violating back-pressure. ### 11.3. Streaming helpers and test utilities diff --git a/docs/users-guide.md b/docs/users-guide.md index ea83bb31..b1b66371 100644 --- a/docs/users-guide.md +++ b/docs/users-guide.md @@ -1428,8 +1428,7 @@ admitted against the same warm slot. Treat `max_in_flight_per_socket` as an admission budget, not as a guarantee of parallel writes on one TCP stream. Create a `PoolHandle` when one logical session needs repeated pooled access and -you want that session to participate in the configured fairness policy over -time: +that session should participate in the configured fairness policy over time: ```rust,no_run # use std::net::SocketAddr; diff --git a/src/client/pool/client_pool.rs b/src/client/pool/client_pool.rs index 9229ad69..946b9bca 100644 --- a/src/client/pool/client_pool.rs +++ b/src/client/pool/client_pool.rs @@ -5,7 +5,7 @@ use std::{ pin::Pin, sync::{ Arc, - atomic::{AtomicUsize, Ordering}, + atomic::{AtomicBool, AtomicUsize, Ordering}, }, }; @@ -40,6 +40,7 @@ where pub(crate) slots: Arc<[Arc>]>, pub(crate) next_slot: AtomicUsize, pub(crate) scheduler: Arc>, + shutdown: AtomicBool, } /// Pool of warm wireframe client sockets. @@ -85,6 +86,7 @@ where slots: Arc::from(slots), next_slot: AtomicUsize::new(0), scheduler: Arc::new(PoolScheduler::new(fairness_policy)), + shutdown: AtomicBool::new(false), }), }) } @@ -129,6 +131,7 @@ where /// Close the pool and drop all warm sockets. pub async fn close(self) { + self.inner.shutdown(); tokio::task::yield_now().await; drop(self); } @@ -140,10 +143,20 @@ where P: Encode + Clone + Send + Sync + 'static, C: Send + 'static, { - pub(crate) fn try_acquire_immediately(&self) -> Option> { + pub(crate) fn is_shutdown(&self) -> bool { self.shutdown.load(Ordering::Acquire) } + + pub(crate) fn shutdown(&self) { + self.shutdown.store(true, Ordering::Release); + self.scheduler.notify_shutdown(); + } + + pub(crate) fn try_acquire_immediately(self: &Arc) -> Option> { + if self.is_shutdown() { + return None; + } self.ordered_slots().into_iter().find_map(|slot| { slot.try_acquire_permit() - .map(|permit| PooledClientLease::new(slot, permit, None)) + .map(|permit| PooledClientLease::new(slot, permit, Some(Arc::clone(self)))) }) } diff --git a/src/client/pool/scheduler.rs b/src/client/pool/scheduler.rs index 6aad73a4..c70a88dc 100644 --- a/src/client/pool/scheduler.rs +++ b/src/client/pool/scheduler.rs @@ -60,8 +60,9 @@ where } fn enqueue_waiter(&mut self, handle_id: u64, sender: WaiterSender) { - self.waiters.insert(handle_id, sender); - self.fifo_waiters.push_back(handle_id); + if self.waiters.insert(handle_id, sender).is_none() { + self.fifo_waiters.push_back(handle_id); + } } fn has_waiters(&self) -> bool { !self.waiters.is_empty() } @@ -140,16 +141,34 @@ where inner: Arc>, handle_id: u64, ) -> Result, ClientError> { - if let Some(lease) = inner.try_acquire_immediately() { - return Ok(lease); + if inner.is_shutdown() { + return Err(ClientError::disconnected()); } let (sender, receiver) = oneshot::channel(); recover_mutex(&self.state).enqueue_waiter(handle_id, sender); - self.kick(inner); + + if let Some(lease) = inner.try_acquire_immediately() { + if let Some(waiter) = self.take_next_waiter_or_stop() { + let _ = waiter.send(Ok(lease)); + } else { + std::mem::forget(receiver); + return Ok(lease); + } + } else { + self.kick(inner); + } + receiver.await.map_err(|_| ClientError::disconnected())? } + pub(crate) fn notify_shutdown(&self) { + let mut state = recover_mutex(&self.state); + while let Some(waiter) = state.take_next_waiter(self.fairness_policy) { + let _ = waiter.send(Err(ClientError::disconnected())); + } + } + pub(crate) fn notify_capacity_available( self: &Arc, inner: Arc>, diff --git a/src/client/tests/pool_handle.rs b/src/client/tests/pool_handle.rs index d07bffeb..98d95643 100644 --- a/src/client/tests/pool_handle.rs +++ b/src/client/tests/pool_handle.rs @@ -15,8 +15,16 @@ use tokio::{ }; use crate::{ - client::{ClientError, ClientPoolConfig, PoolFairnessPolicy, PoolHandle}, - test_helpers::{Ping, Pong, PoolTestServer, TestClientPool, build_pooled_client}, + client::{ClientError, ClientPoolConfig, PoolFairnessPolicy}, + test_helpers::{ + Ping, + Pong, + PoolTestServer, + TestClientPool, + acquire_and_record, + build_pooled_client, + build_preamble_pool, + }, }; type TestResult = Result<(), Box>; @@ -35,34 +43,6 @@ async fn build_handle_pool( Ok((server, pool)) } -async fn build_preamble_pool( - config: ClientPoolConfig, -) -> Result<(PoolTestServer, TestClientPool, Arc), ClientError> { - let preamble_callback_count = Arc::new(AtomicUsize::new(0)); - let server = PoolTestServer::start().await?; - let pool = build_pooled_client(server.addr, config, preamble_callback_count.clone()).await?; - Ok((server, pool, preamble_callback_count)) -} - -async fn acquire_and_record( - mut handle: PoolHandle< - crate::serializer::BincodeSerializer, - crate::test_helpers::ClientHello, - (), - >, - label: &'static str, - rounds: usize, - grants: Arc>>, -) -> Result<(), ClientError> { - for _ in 0..rounds { - let lease = handle.acquire().await?; - grants.lock().await.push(label); - tokio::task::yield_now().await; - drop(lease); - } - Ok(()) -} - #[rstest] #[tokio::test(flavor = "current_thread")] async fn round_robin_handles_share_one_socket_fairly( @@ -86,7 +66,18 @@ async fn round_robin_handles_share_one_socket_fairly( right_result??; let observed = grants.lock().await.clone(); - assert_eq!(observed, vec!["a", "b", "a", "b", "a", "b"]); + assert_eq!(observed.len(), 6, "expected 6 total grants"); + + let a_count = observed.iter().filter(|&label| *label == "a").count(); + let b_count = observed.iter().filter(|&label| *label == "b").count(); + assert_eq!(a_count, 3, "expected 3 grants for session-a"); + assert_eq!(b_count, 3, "expected 3 grants for session-b"); + + for i in 1..observed.len() { + let prev = observed.get(i - 1).expect("valid index i-1"); + let curr = observed.get(i).expect("valid index i"); + assert_ne!(prev, curr, "expected strict alternation at position {i}"); + } Ok(()) } @@ -152,13 +143,45 @@ async fn handle_acquire_respects_back_pressure(client_pool_config: ClientPoolCon Ok(()) } +#[rstest] +#[tokio::test(flavor = "current_thread")] +async fn handle_acquire_dropped_waiter_does_not_leak_capacity( + client_pool_config: ClientPoolConfig, +) -> TestResult { + let (_server, pool) = build_handle_pool(client_pool_config.pool_size(1)).await?; + + let mut holder = pool.handle(); + let mut waiter1 = pool.handle(); + let mut waiter2 = pool.handle(); + + let held_lease = holder.acquire().await?; + + let waiter1_task = tokio::spawn(async move { waiter1.acquire().await }); + let waiter2_task = tokio::spawn(async move { waiter2.acquire().await }); + + tokio::time::sleep(Duration::from_millis(25)).await; + + drop(waiter1_task); + + drop(held_lease); + + let lease2 = timeout(Duration::from_millis(100), waiter2_task).await???; + drop(lease2); + + let mut later = pool.handle(); + let _lease3 = timeout(Duration::from_millis(100), later.acquire()).await??; + + Ok(()) +} + #[rstest] #[tokio::test] async fn handle_path_preserves_warm_reuse_and_preamble( client_pool_config: ClientPoolConfig, ) -> TestResult { - let (server, pool, preamble_callback_count) = - build_preamble_pool(client_pool_config.pool_size(1)).await?; + let server = PoolTestServer::start().await?; + let (pool, preamble_callback_count) = + build_preamble_pool(&server, client_pool_config.pool_size(1)).await?; let mut handle = pool.handle(); let first: Pong = handle.call(&Ping(7)).await?; @@ -178,8 +201,12 @@ async fn handle_path_recycles_after_idle_timeout( client_pool_config: ClientPoolConfig, ) -> TestResult { let idle_timeout = Duration::from_millis(50); - let (server, pool, preamble_callback_count) = - build_preamble_pool(client_pool_config.pool_size(1).idle_timeout(idle_timeout)).await?; + let server = PoolTestServer::start().await?; + let (pool, preamble_callback_count) = build_preamble_pool( + &server, + client_pool_config.pool_size(1).idle_timeout(idle_timeout), + ) + .await?; let mut handle = pool.handle(); let first: Pong = handle.call(&Ping(1)).await?; diff --git a/src/test_helpers.rs b/src/test_helpers.rs index 3a222888..cbbe114c 100644 --- a/src/test_helpers.rs +++ b/src/test_helpers.rs @@ -28,7 +28,9 @@ pub use pool_client::{ PoolServerBehavior, PoolTestServer, TestClientPool, + acquire_and_record, build_pooled_client, + build_preamble_pool, }; /// Test-friendly message assembler implementation that shares parsing logic. diff --git a/src/test_helpers/pool_client.rs b/src/test_helpers/pool_client.rs index 75ed3e1e..8e42ad6b 100644 --- a/src/test_helpers/pool_client.rs +++ b/src/test_helpers/pool_client.rs @@ -9,7 +9,7 @@ use std::{ }; use futures::{FutureExt, SinkExt, StreamExt}; -use tokio::{net::TcpListener, task::JoinHandle}; +use tokio::{net::TcpListener, sync::Mutex, task::JoinHandle}; use tokio_util::codec::{Framed, LengthDelimitedCodec}; use crate::{ @@ -136,6 +136,40 @@ pub fn build_pooled_client( .connect_pool(addr, config) } +/// Build a pooled client with a preamble counter for tracking warm reuse. +/// +/// # Errors +/// +/// Returns an error if the server cannot be started or the pool connection fails. +pub async fn build_preamble_pool( + server: &PoolTestServer, + config: ClientPoolConfig, +) -> Result<(TestClientPool, Arc), ClientError> { + let preamble_callback_count = Arc::new(AtomicUsize::new(0)); + let pool = build_pooled_client(server.addr, config, preamble_callback_count.clone()).await?; + Ok((pool, preamble_callback_count)) +} + +/// Acquire and release a lease repeatedly, recording each grant in a shared log. +/// +/// # Errors +/// +/// Returns an error if any `acquire()` call fails. +pub async fn acquire_and_record( + mut handle: crate::client::PoolHandle, + label: &'static str, + rounds: usize, + grants: Arc>>, +) -> Result<(), ClientError> { + for _ in 0..rounds { + let lease = handle.acquire().await?; + grants.lock().await.push(label); + tokio::task::yield_now().await; + drop(lease); + } + Ok(()) +} + async fn run_pool_test_accept_loop(listener: TcpListener, state: PoolServerState) { loop { let accept_result = listener.accept().await; diff --git a/tests/fixtures/client_pool_handle.rs b/tests/fixtures/client_pool_handle.rs index c6fd2b8c..d5711d44 100644 --- a/tests/fixtures/client_pool_handle.rs +++ b/tests/fixtures/client_pool_handle.rs @@ -12,7 +12,14 @@ use rstest::fixture; use tokio::{sync::Mutex, time::timeout}; use wireframe::{ client::{ClientPoolConfig, PoolFairnessPolicy}, - test_helpers::{Ping, Pong, PoolTestServer, TestClientPool, build_pooled_client}, + test_helpers::{ + Ping, + Pong, + PoolTestServer, + TestClientPool, + acquire_and_record, + build_pooled_client, + }, }; pub use wireframe_testing::TestResult; @@ -95,13 +102,13 @@ impl ClientPoolHandleWorld { let pool = self.pool.as_ref().ok_or("pool missing")?; let grants = Arc::new(Mutex::new(Vec::new())); - let left = tokio::spawn(run_repeated_acquire( + let left = tokio::spawn(acquire_and_record( pool.handle(), "session-a", 3, Arc::clone(&grants), )); - let right = tokio::spawn(run_repeated_acquire( + let right = tokio::spawn(acquire_and_record( pool.handle(), "session-b", 3, @@ -206,40 +213,21 @@ async fn spawn_fifo_waiters( tokio::task::JoinHandle>, tokio::task::JoinHandle>, ) { - let first = tokio::spawn(run_repeated_acquire( + let first = tokio::spawn(acquire_and_record( pool.handle(), "first", 1, Arc::clone(&grants), )); tokio::task::yield_now().await; - let second = tokio::spawn(run_repeated_acquire( + let second = tokio::spawn(acquire_and_record( pool.handle(), "second", 1, Arc::clone(&grants), )); tokio::task::yield_now().await; - let third = tokio::spawn(run_repeated_acquire(pool.handle(), "third", 1, grants)); + let third = tokio::spawn(acquire_and_record(pool.handle(), "third", 1, grants)); tokio::task::yield_now().await; (first, second, third) } - -async fn run_repeated_acquire( - mut handle: wireframe::client::PoolHandle< - wireframe::serializer::BincodeSerializer, - wireframe::test_helpers::ClientHello, - (), - >, - label: &'static str, - rounds: usize, - grants: Arc>>, -) -> Result<(), wireframe::client::ClientError> { - for _ in 0..rounds { - let lease = handle.acquire().await?; - grants.lock().await.push(label); - tokio::task::yield_now().await; - drop(lease); - } - Ok(()) -} From 93c959bea998801880c15f412d49b0d164d1998b Mon Sep 17 00:00:00 2001 From: Leynos Date: Tue, 17 Mar 2026 17:23:10 +0000 Subject: [PATCH 5/5] feat(client-pool): add shutdown notification to client pool and scheduler - Introduce a Tokio Notify to signal shutdown state in ClientPool. - Modify PoolScheduler to asynchronously respond to shutdown via notification. - Prevent waiter starvation by servicing waiters one at a time asynchronously. - Improve resource cleanup on waiter cancellations. - Update tests to handle aborted waiter tasks correctly. - Refactor client pool test helpers and fixtures for correctness. These changes enable more graceful and responsive shutdown handling in the connection pool implementation. Co-authored-by: devboxerhub[bot] --- src/client/pool/client_pool.rs | 6 ++++ src/client/pool/scheduler.rs | 50 +++++++++++++++++++--------- src/client/tests/pool_handle.rs | 16 ++++----- src/test_helpers/pool_client.rs | 6 ++-- tests/fixtures/client_pool_handle.rs | 28 +++++++++++----- 5 files changed, 69 insertions(+), 37 deletions(-) diff --git a/src/client/pool/client_pool.rs b/src/client/pool/client_pool.rs index 946b9bca..f3eb12e4 100644 --- a/src/client/pool/client_pool.rs +++ b/src/client/pool/client_pool.rs @@ -12,6 +12,7 @@ use std::{ use bb8::Pool; use bincode::Encode; use futures::{Future, future::select_all}; +use tokio::sync::Notify; use super::{ ClientPoolConfig, @@ -41,6 +42,7 @@ where pub(crate) next_slot: AtomicUsize, pub(crate) scheduler: Arc>, shutdown: AtomicBool, + shutdown_notify: Notify, } /// Pool of warm wireframe client sockets. @@ -87,6 +89,7 @@ where next_slot: AtomicUsize::new(0), scheduler: Arc::new(PoolScheduler::new(fairness_policy)), shutdown: AtomicBool::new(false), + shutdown_notify: Notify::new(), }), }) } @@ -145,8 +148,11 @@ where { pub(crate) fn is_shutdown(&self) -> bool { self.shutdown.load(Ordering::Acquire) } + pub(crate) async fn shutdown_notified(&self) { self.shutdown_notify.notified().await; } + pub(crate) fn shutdown(&self) { self.shutdown.store(true, Ordering::Release); + self.shutdown_notify.notify_waiters(); self.scheduler.notify_shutdown(); } diff --git a/src/client/pool/scheduler.rs b/src/client/pool/scheduler.rs index c70a88dc..e7fb448c 100644 --- a/src/client/pool/scheduler.rs +++ b/src/client/pool/scheduler.rs @@ -149,11 +149,14 @@ where recover_mutex(&self.state).enqueue_waiter(handle_id, sender); if let Some(lease) = inner.try_acquire_immediately() { - if let Some(waiter) = self.take_next_waiter_or_stop() { - let _ = waiter.send(Ok(lease)); - } else { - std::mem::forget(receiver); + let Some(waiter) = self.take_next_waiter_or_stop() else { + drop(receiver); return Ok(lease); + }; + + if waiter.send(Ok(lease)).is_err() { + drop(receiver); + return Err(ClientError::disconnected()); } } else { self.kick(inner); @@ -214,19 +217,36 @@ where } async fn service_waiters(self: Arc, inner: Arc>) { - loop { - let Some(sender) = self.take_next_waiter_or_stop() else { - break; - }; + while let Some(sender) = self.take_next_waiter_or_stop() { + self.service_one_waiter(sender, Arc::clone(&inner)).await; + } + } - let result = inner.acquire_slot_permit().await.map(|(slot, permit)| { - PooledClientLease::new(slot, permit, Some(Arc::clone(&inner))) - }); - if let Err(send_result) = sender.send(result) - && let Ok(lease) = send_result - { - drop(lease); + #[expect( + clippy::integer_division_remainder_used, + reason = "tokio::select! macro internally uses % for random branch selection" + )] + async fn service_one_waiter( + &self, + sender: WaiterSender, + inner: Arc>, + ) { + let result = tokio::select! { + permit_result = inner.acquire_slot_permit() => { + permit_result.map(|(slot, permit)| { + PooledClientLease::new(slot, permit, Some(Arc::clone(&inner))) + }) } + () = inner.shutdown_notified() => { + let _ = sender.send(Err(ClientError::disconnected())); + return; + } + }; + + if let Err(send_result) = sender.send(result) + && let Ok(lease) = send_result + { + drop(lease); } } } diff --git a/src/client/tests/pool_handle.rs b/src/client/tests/pool_handle.rs index 98d95643..e97c658d 100644 --- a/src/client/tests/pool_handle.rs +++ b/src/client/tests/pool_handle.rs @@ -161,7 +161,8 @@ async fn handle_acquire_dropped_waiter_does_not_leak_capacity( tokio::time::sleep(Duration::from_millis(25)).await; - drop(waiter1_task); + waiter1_task.abort(); + let _ = waiter1_task.await; drop(held_lease); @@ -179,9 +180,8 @@ async fn handle_acquire_dropped_waiter_does_not_leak_capacity( async fn handle_path_preserves_warm_reuse_and_preamble( client_pool_config: ClientPoolConfig, ) -> TestResult { - let server = PoolTestServer::start().await?; - let (pool, preamble_callback_count) = - build_preamble_pool(&server, client_pool_config.pool_size(1)).await?; + let (server, pool, preamble_callback_count) = + build_preamble_pool(client_pool_config.pool_size(1)).await?; let mut handle = pool.handle(); let first: Pong = handle.call(&Ping(7)).await?; @@ -201,12 +201,8 @@ async fn handle_path_recycles_after_idle_timeout( client_pool_config: ClientPoolConfig, ) -> TestResult { let idle_timeout = Duration::from_millis(50); - let server = PoolTestServer::start().await?; - let (pool, preamble_callback_count) = build_preamble_pool( - &server, - client_pool_config.pool_size(1).idle_timeout(idle_timeout), - ) - .await?; + let (server, pool, preamble_callback_count) = + build_preamble_pool(client_pool_config.pool_size(1).idle_timeout(idle_timeout)).await?; let mut handle = pool.handle(); let first: Pong = handle.call(&Ping(1)).await?; diff --git a/src/test_helpers/pool_client.rs b/src/test_helpers/pool_client.rs index 8e42ad6b..3cb08670 100644 --- a/src/test_helpers/pool_client.rs +++ b/src/test_helpers/pool_client.rs @@ -142,12 +142,12 @@ pub fn build_pooled_client( /// /// Returns an error if the server cannot be started or the pool connection fails. pub async fn build_preamble_pool( - server: &PoolTestServer, config: ClientPoolConfig, -) -> Result<(TestClientPool, Arc), ClientError> { +) -> Result<(PoolTestServer, TestClientPool, Arc), ClientError> { let preamble_callback_count = Arc::new(AtomicUsize::new(0)); + let server = PoolTestServer::start().await?; let pool = build_pooled_client(server.addr, config, preamble_callback_count.clone()).await?; - Ok((pool, preamble_callback_count)) + Ok((server, pool, preamble_callback_count)) } /// Acquire and release a lease repeatedly, recording each grant in a shared log. diff --git a/tests/fixtures/client_pool_handle.rs b/tests/fixtures/client_pool_handle.rs index d5711d44..2cd30886 100644 --- a/tests/fixtures/client_pool_handle.rs +++ b/tests/fixtures/client_pool_handle.rs @@ -187,15 +187,25 @@ impl ClientPoolHandleWorld { } pub fn sessions_alternate_fairly(&self) -> bool { - self.grant_order - == [ - "session-a", - "session-b", - "session-a", - "session-b", - "session-a", - "session-b", - ] + if self.grant_order.len() != 6 { + return false; + } + + let first = self.grant_order.first(); + let second = self.grant_order.get(1); + + if first == second { + return false; + } + + for (i, grant) in self.grant_order.iter().enumerate() { + let expected = if i & 1 == 0 { first } else { second }; + if Some(grant) != expected { + return false; + } + } + + true } pub fn fifo_order_preserved(&self) -> bool { self.grant_order == ["first", "second", "third"] }