[perf, fix] Pool and reuse ZMQ request sockets - #167
Merged
Conversation
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
OutstanderWang
force-pushed
the
feat_socket_pool
branch
from
September 7, 2026 08:13
b9b4f2b to
6707683
Compare
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
0oshowero0
reviewed
Sep 7, 2026
| self.zmq_context, | ||
| client_id, | ||
| "request_handle_socket", | ||
| maxsize=TQ_CLIENT_ZMQ_POOL_SIZE, |
Collaborator
There was a problem hiding this comment.
TQ_CLIENT_ZMQ_POOL_SIZE only affects controller RPC pool. The name is misleading
Contributor
Author
There was a problem hiding this comment.
Thank you for pointing it out. I have renamed it as TQ_CONTROLLER_RPC_POOL_SIZE
0oshowero0
reviewed
Sep 7, 2026
| return owner.is_closed() | ||
|
|
||
|
|
||
| class ZMQSocketPool: |
Collaborator
There was a problem hiding this comment.
I believe we need to figure out which cases can be benefit from the socket pool. Now the pool is managed by (event_loop, address). So in different asyncio.run, they cannot benefit from the refactor. Besides, the notify_pool uses its own loop, makes it more difficult to figure out which cases this PR can improve.
Contributor
Author
There was a problem hiding this comment.
Right on the keying, and right that a fresh asyncio.run() per call gets no reuse — test_finished_loop_releases_its_sockets pins exactly that.
That's by design rather than a limitation to work around: each role owns its own pool, so sockets are never shared across roles. notify_pool having its own loop is what makes it eligible at all — a ZMQ socket can't safely cross loops, so it could never have shared with another scenario regardless.
Every supported entry point sits on a long-lived owner, so all four pools reuse in practice:
| Pool | Owner | Reuse |
|---|---|---|
controller_rpc_pool (client) |
one loop built in TransferQueueClient.__init__, calls via run_coroutine_threadsafe |
full |
storage_rpc_pool (storage manager) |
the caller's loop | full |
notify_pool (storage manager) |
dedicated loop + thread, created once | full |
| metrics pool (collector) | synchronous; one daemon thread, keyed by thread | full |
AsyncTransferQueueClient isn't exported, so reaching it directly is already off the supported path.
Where you were dead right: the docstrings showed the wrong thing — every example wrapped a single call in asyncio.run(), the put/get walkthrough three times in a row. Fixed in 29e40a2, with the reuse condition now stated on the class instead of left to _lease_owner.
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
0oshowero0
reviewed
Sep 8, 2026
Comment on lines
+381
to
+384
| # A subclass may reject its config before calling super().__init__(), as the KV | ||
| # managers do, leaving nothing here allocated. __del__ calls close() anyway, so | ||
| # return rather than burying the constructor's error under an AttributeError. | ||
| # notify_pool is the last thing __init__ sets before its first fallible step. |
Collaborator
There was a problem hiding this comment.
We can remove some over complicated comments
0oshowero0
reviewed
Sep 8, 2026
| ... partition_id="train_0" | ||
| ... )) | ||
| >>> print(f"Global index: {global_index}, Consumption status: {consumption_status}") | ||
| >>> async def main(): |
Collaborator
There was a problem hiding this comment.
We don't need to add async def main(): here.
0oshowero0
reviewed
Sep 8, 2026
|
|
||
| @pytest.mark.asyncio | ||
| @pytest.mark.parametrize("failure", ["exception", "cancellation"]) | ||
| async def test_poisoned_lease_is_discarded(peer, failure): |
Collaborator
There was a problem hiding this comment.
【AI Comment】
poisoned() never sends a request, so peer.callers == 1 holds whether the socket is discarded or incorrectly returned to the pool: the peer only observes the subsequent successful request.
I verified that both parameterized cases still pass after modifying lease() to return the socket to the pool even on exception or cancellation.
Please send the first request and wait until the peer receives it before triggering the failure, while delaying its reply. Then send a distinct second request and assert that it uses a new connection identity and receives its own response. This would verify the discard behavior and protect against a late reply being mistaken for the next request’s response.
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
lease() treats a clean exit as proof the socket completed a send/recv and hands it to _release(), which parked it without checking whether it was still open. StorageManager._notify_and_wait exits exactly that way: on a missing ACK it closes the socket, so a reply still in flight cannot be read by the next lessee as its own, then swallows the error because a slow controller must not fail the put that triggered the notification. The closed socket went back into the pool and held a slot until the next lease popped and discarded it, costing a reconnect after every failed notification. _release now refuses a closed socket, which is the invariant the pool wanted all along: nothing closed is ever parked. Fixing it here rather than at the call site keeps the notify path able to distinguish a slow ACK from the configuration and lifecycle faults lease() itself raises, which it deliberately does not swallow. The regression test was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
storage_rpc_pool is assigned after super().__init__(), which runs the controller handshake and raises TimeoutError once TQ_STORAGE_HANDSHAKE_MAX_RETRIES is exhausted -- an ordinary outcome when the controller is slow or unreachable at startup. close() dereferenced the pool unconditionally, so on that path it raised AttributeError before reaching super().close(), leaving the notify pool and the context with its native I/O threads behind. __del__ swallows the error into a log line, so the leak surfaced only as a confusing message. Guarding the attribute matches the base class, whose close() already reaches for every teardown attribute through hasattr/getattr precisely because __del__ must cope with a partially built object. This override was the only one that did not. The regression test drives a handshake failure and asserts the context is destroyed; it was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…it__() KVStorageManager, RayStorageManager and YuanrongStorageManager all validate their config and raise before delegating to super().__init__(), so none of the attributes close() reaches for exist on the resulting object. __del__ calls close() anyway, and the very first line dereferenced controller_handshake_socket unguarded. The AttributeError was caught and logged, which buried the real ValueError under "Exception during __del__: object has no attribute controller_handshake_socket". Nothing leaks on this path, since super().__init__() is what allocates the context, the notify pool and the notify thread; the cost is a misleading error that sends readers after the wrong fault. The existing hasattr guards further down show close() was already meant to cope with a partially built object, just not with one where the base constructor never ran at all. Guarding on notify_pool covers the whole method in one check: it is the last attribute __init__ assigns before its first fallible step, so if it is present every attribute close() touches is too, and the existing guards already handle a handshake that failed after that point. Predates the socket pool work. The regression test was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The pool carried a per-peer map of current addresses so that a peer moving to a new port under a stable id had the sockets at its old address retired. Only the metrics collector opted in, via register_storage_units(), and that runs exactly once: interface.init() calls it during first-time initialization, returns early through _init_from_existing() on every later call, and tq.close() ray.kill()s the controller actor rather than reconfiguring it. Nothing else remaps an endpoint at runtime -- client._controller, manager.controller_info and manager.storage_unit_infos are each assigned once at construction. Reuse was always keyed by address, never by peer id, so a moved peer is still dialed correctly without any of this; test_reregistered_peer_is_not_served_a_ stale_socket keeps that guarantee and passes unchanged. What the tracking added on top was closing the sockets left behind at the abandoned address, and its absence is now recorded where _idle is declared so that anyone adding runtime remapping knows to retire those buckets. Removing it takes with it the follow_endpoint_changes flag, the _endpoints map, _mark_current, _superseded, and the peer_id argument that only those two needed from _take and _release. Four tests covering the deleted state go too; the remaining eleven cover behaviour that is still reachable. Net 199 lines lighter, 60 of them in the pool itself. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The exporter minted its own zmq.Context() the first time it queried a storage unit, and nothing ever terminated it: the exporter is created once per Ray actor by start_metrics(), which is idempotent, and the actor is torn down with ray.kill(), so there is no shutdown path to close it from. Adding one would have meant a graceful actor-exit protocol -- a stop signal for the collect loop, a shutdown for the Prometheus HTTP server, and an RPC in interface.close() ahead of ray.kill() -- to release a context in a process that is about to exit anyway. The controller already holds a long-lived synchronous context, which is exactly what these queries need, so the exporter now borrows it and the second context disappears along with the question of who closes it. This is the arrangement ZMQSocketPool already documents and the one the client uses when it lends its context to the SimpleStorage manager. The storage-role exporter passes nothing and needs nothing: only the controller role starts the collection loop, so only it ever builds a pool. Asking for a pool without a context now raises instead of quietly creating one, which keeps the lending explicit. Metrics sockets, one per storage unit, now count against the controller context's socket budget. That is libzmq's default 1023 and the context holds two ROUTERs today, so the headroom covers roughly a thousand storage units; the socket count itself is unchanged, only which context they belong to. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The cap bounds only how many idle sockets are kept per (owner, endpoint) bucket, not concurrency: a burst beyond it is still served, but the excess is closed on return and pays a fresh handshake next time. At 8 that made reuse fall off well below the concurrency these paths actually reach. Both defaults move together so all four pools follow: the client's controller RPC pool reads TQ_CLIENT_ZMQ_POOL_SIZE, while the storage RPC, notify, and metrics pools take the constructor default. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
At this point explicitly_requested is true exactly when max_sockets is not None, so the two guards accept the same inputs. Testing the value itself lets a type checker see that the comparison below operates on an int. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
TQ_CLIENT_ZMQ_POOL_SIZE reads as a client-wide budget but is applied at a single call site, the controller RPC pool. The storage RPC, notify, and metrics pools take ZMQSocketPool's own default, so the CLIENT prefix promises a reach the knob does not have. Rename it to TQ_CONTROLLER_RPC_POOL_SIZE, which names exactly what it sets. The knob is new in this branch and unreferenced outside it, so no configuration in the wild has to change. Document the two properties the name cannot carry: the cap counts per (owner, address) rather than per pool, so a pool dialling N peers may park N*maxsize idle sockets, and the pools that take the default are the ones whose concurrency does not warrant a knob -- notify serializes onto its own loop and metrics collects sequentially, each holding one socket at a time. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…r call Every example on AsyncTransferQueueClient wrapped a single call in asyncio.run(), and the put/get walkthrough did so three times in a row. Pooled sockets are keyed by their owning event loop, so that pattern retires its socket at the end of each call and pays a fresh connect handshake on the next one -- the examples demonstrated precisely the shape that cannot reuse a connection. Await the calls inside one async def instead, which is both how the supported entry points drive the client and the shape the pool is built for. TransferQueueClient keeps a loop of its own for callers that have none, so state the reuse condition on the class itself rather than leaving it to be inferred from the pool internals. No behavior change; docstrings only. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The pool tests read pool._idle through a helper and asserted on _ctx, _socket_name and _timeout, so they described the implementation rather than the behavior and broke on any refactor that kept the contract intact. Every pooled socket dials with its own ZMQ identity, so the peer can count distinct callers: one identity across many requests means the connection was reused, a new one means the old socket was discarded. That is the same property observed from the outside, and it needs no access to pool state. Drop the cases reachable only by construction -- identity collisions between same-pid pools, connect failure caught by patching create_zmq_socket, an unknown socket name, close() idempotency, and a size check already covered on the client -- and fold the cancelled and raising leases into one parametrized test. Eight tests remain, covering reuse, the timed-out socket that must not answer the next request, poisoned leases, loop keying, peer re-registration, synchronous callers, and bursts above maxsize. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Reuse was only covered against ZMQSocketPool directly, which is internal and not something a user reaches. Each pool belongs to a role, so the reuse each role depends on now lives with that role. test_client.py drives it through the public client: repeated get_partition_list() calls must reach the controller over one connection, which is the user-visible payoff of pooling and something the file asserted nowhere before. MockController records caller identities to make that observable; it answers requests exactly as before. test_metrics.py covers the collector, whose pool keys by thread because it runs with no event loop at all, against a ROUTER that counts callers. Both were checked by breaking _take() so no socket is ever handed back: each then fails with one identity per request, so neither passes vacuously. What stays in test_zmq_socket_pool.py needs control over reply timing and loop lifetime that a caller-level test cannot reach -- the late reply that must not answer the next request, poisoned leases, loop keying, peer re-registration, and bursts past maxsize. Its docstring now says where each role's own reuse is covered, so the split is not mistaken for an omission. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Review found test_poisoned_lease_is_discarded passing vacuously. The lease sent nothing before failing, so the peer never saw that socket and the caller count held whether the socket was discarded or wrongly parked. Confirmed by releasing the socket on the exception path instead of closing it: both parametrized cases still passed. Send the request and wait until the peer has it before failing, with the reply delayed so it is genuinely outstanding, then assert the next request both uses a different connection and receives its own reply. Sabotaging the discard now fails both cases on the reply belonging to the abandoned request, which is the misattribution the discard exists to prevent. Drop the async def main() wrapper from examples that make a single call; it only earns its keep where consecutive calls show reuse under one loop. Bare await matches the async examples already in interface.py. Shorten the close() guard comment to the reason it exists. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Five blocks ran past the four-line ceiling and three more repeated what the code or a nearby docstring already says. Comments only in code, none removed outright. Most of the excess was duplication: lease exclusivity and the maxsize semantics are documented on ZMQSocketPool, so restating them at each call site left three places to keep in sync. Two blocks explained a state no supported path reaches -- a bucket orphaned by endpoint remapping, and pool isolation that the field names already carry -- which the simplicity gate asks us not to describe. What is left is the part a reader cannot infer: why the pre-context checks come first, why reuse makes the socket budget track concurrency rather than request count, why a late ACK forces a close, and why owner_id alone cannot make a ZMQ identity unique. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The knob only reached the client's controller RPC pool, so the storage RPC, notify, and metrics pools were fixed at the constructor default with no way to tune them. Review asked for one name covering all four; TQ_SOCKET_POOL_SIZE replaces TQ_CONTROLLER_RPC_POOL_SIZE, keeping the same default of 64. It lives beside ZMQSocketPool because that is what the cap belongs to, and is resolved in __init__ rather than bound as a signature default: a default read at import freezes whatever the environment held when the module first loaded, which is the failure the per-call socket timeout already had to fix. An explicit maxsize still wins, so a caller can opt out. The cap remains per (owner, address), so one value means different totals for a pool dialling one peer and a pool dialling N. That is unchanged behaviour, and the docstring says so. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…ontext The cap is per (owner, address), so it multiplies by peer count while the context's socket ceiling does not. At 64 a storage manager addressing two thousand units could park 128k idle sockets against a budget of 8192, and the lease that hits the limit fails with EMFILE rather than degrading. Drop the default to 8. One socket per peer already removes the repeated handshake, which is what pooling was for; a larger cap only helps when a single peer sees concurrent requests, and that is the case worth opting into explicitly rather than paying for by default. Warn at construction when TQ_SOCKET_POOL_SIZE times the registered unit count exceeds the context's ZMQ_MAX_SOCKETS, naming both knobs. The storage manager is the first point that knows how many peers it will dial, and the pool cannot tell on its own because it connects lazily. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…unit Pooling these queries made the collector hold a socket per unit for as long as the process ran. Sockets are bucketed per address and only swept when their owner ends, and the collector is a permanent daemon thread, so nothing ever retired them: a few thousand units meant a few thousand resident sockets. They land on the controller's context, which never raises MAX_SOCKETS from libzmq's default of 1023. Past that a lease fails with EMFILE, so the units past the limit lose their metrics every cycle, and the controller itself cannot open a socket either -- a monitoring change taking out the control plane. TQ_SOCKET_POOL_SIZE does not bound this. It caps sockets per (owner, address) and every bucket here holds exactly one; the count comes from the number of addresses. Close the socket after each query, as the notify path already does. Collection walks every unit once per cycle, so a kept socket is reused only a cycle later while occupying budget for the whole walk. Resident sockets go from one per unit to zero, and the per-cycle handshake is immaterial against a 10s interval. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
maxsize bounded only what the pool parked, not what it opened: a burst of N concurrent requests opened N sockets and closed all but maxsize on return. Peak socket count therefore tracked real concurrency, which is what makes the budget unpredictable -- maxsize times peer count was a number nothing enforced, and past ZMQ_MAX_SOCKETS a lease fails with EMFILE. Async callers now use alease, which takes a permit per (owner, address) before leasing. Concurrency past the cap waits for a socket to come back rather than opening another, so sockets in flight follow configuration. It is also faster, since waiting costs less than the handshake it replaces: locally, 300-way concurrency ran 48ms against 71ms, with peak sockets 8 against 300. Permits are keyed like the buckets, so a fan-out across peers is not serialized by a cap meant to bound one peer's concurrency. A permit lives until its body ends, so leasing inside a lease can wait on one the task -- or a sibling mid-cycle -- already holds, and that hangs with no timeout. Any nesting now raises, including across pools: separate semaphores do not break a cycle, only the ordering does. An RPC's reply is what frees its permit, so a second RPC belongs after the first returns; notify already works that way, running once the puts it reports have completed. The synchronous lessee keeps lease(): the metrics collector issues one request at a time, so it never queues and cannot await. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…fault The controller never set MAX_SOCKETS, so its context kept libzmq's default of 1023. Its own two ROUTERs are nowhere near that, but the metrics exporter borrows this context to query storage units, and a controller that cannot open a socket stops answering requests at all -- a monitoring cost taking out the control plane. Set it to 4096, clamped to this build's ZMQ_SOCKET_LIMIT, before the first socket is opened: libzmq applies the ceiling at socket creation. Metrics now closes each query's socket, so today's peak is one; the headroom is for a collector that fans out instead of walking units in sequence. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
With alease the cap also bounds sockets in flight to one peer, so it multiplies by peer count against a fixed context budget. At 8 the worst case for two thousand storage units was 16000 sockets against a budget of 8192 -- reachable, not theoretical: a single peer under 50-way concurrency does saturate its cap. Four keeps that product inside the budget (8004) while still absorbing the overlap a peer actually sees, since one put or get sends a single merged request per unit and concurrency there comes only from operations overlapping. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
One block ran past the four-line ceiling and seven more said something the code or a nearby docstring already says. Comments only. Two were duplication introduced by earlier edits: the permit keying rationale now sits on the field declaration rather than repeated at its accessor, and the per-role isolation argument belongs to ZMQSocketPool's docstring rather than to each decorator that picks a pool. One narrated what the code plainly does -- naming a daemon thread that the thread's own name and target already give. What is left is the part a reader cannot infer: why the pre-context checks come first, why ZMQ_SOCKET_LIMIT forces the last check below allocation, why a late ACK forces a close, why the pool cap is resolved at construction rather than at import, and why owner_id alone cannot make a ZMQ identity unique. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Rebasing onto main replayed this branch over the identity filter added in Ascend#168, and the pool's changes to the same region dropped the prefixes it filters on. Restore them: the storage proxy rejects any identity lacking one, so without these constants the filter has nothing to match and the module does not import. The pooled sockets already satisfy the filter -- the storage manager pool passes storage_manager_id and the metrics pool now derives its owner id from METRICS_COLLECTOR_IDENTITY_PREFIX rather than repeating the literal, so a change to the prefix reaches both ends. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
OutstanderWang
force-pushed
the
feat_socket_pool
branch
from
September 10, 2026 06:55
dceba42 to
56f2e7d
Compare
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
0oshowero0
approved these changes
Sep 10, 2026
adoda
pushed a commit
to adoda/TransferQueue
that referenced
this pull request
Sep 10, 2026
Keep storage RPC pooling from Ascend#167 while preserving fresh-connection retry: a failed lease is discarded, and the diagnostic probe uses its own short-timeout pool. Compact the retry tests without dropping behavior coverage. Signed-off-by: jathonzhang <jathonzhang@tencent.com>
0oshowero0
pushed a commit
that referenced
this pull request
Sep 11, 2026
## Summary A storage-unit request can be lost between the manager and the unit without either end raising an error. #171 recovers such a request by retrying it; this PR closes one of the paths it can be lost on, and adds the instrumentation needed to tell the loss apart from a slow unit. It also corrects two defects in the diagnosis that #171 introduced — both are on `main` today. ## What's here ### 1. Raise the client-facing ROUTER's accept queue ZMQ defaults `BACKLOG` to 100. That was simply the unset default: a full accept queue is drained **without an RST** (`tcp_abort_on_overflow=0`), so the client stays in `ESTABLISHED` waiting for a reply nobody will ever send, and the server never learns the request existed. Neither end raises. Default is now 4096, tunable via `TQ_STORAGE_ZMQ_BACKLOG` (set it to 100 to restore ZMQ's default for an A/B run). This only raises a queue ceiling — it cannot reject a connection that previously succeeded. ### 2. Make the loss measurable rather than inferred - **Arrival counters** (`requests_arrived`, `arrivals_by_op` in `get_metrics`). The per-op counters advance inside `monitor.measure()`, so they only move once a request *completes* — a request that arrived and then stalled is indistinguishable from one that never arrived. Counting at decode time separates the two, and a shortfall against the caller's send count localizes the loss to one side of the connection. - **Accept-queue probe** (`transfer_queue/utils/accept_probe.py`). Samples `Recv-Q` / `sk_drops` / `ListenOverflows` from the kernel's own view. **Opt-in and off by default** (`TQ_ACCEPT_PROBE_INTERVAL=0`): it shells out to `ss` on a timer, and depth must be sampled sub-second because the queue drains in milliseconds — which is why inspecting a unit *after* it hung always read zero. ### 3. Report an accept-queue drop once per drop `sk_drops` is a cumulative kernel counter, but the initial check compared it against the probe's *first* sample. Once a socket had ever dropped a connection the condition stayed true for the life of the process — at a 0.1 s interval that is ten errors per second per storage unit. The flood also destroyed the signal: repeating the same cumulative total says nothing about *when* drops happened. Now compares against the previous sample and reports both deltas. ### 4. Base the timeout diagnosis on evidence (fixes two defects on `main`) `_diagnose_storage_unit` reported `verdict=request_lost_in_flight` **unconditionally** whenever the post-failure probe answered. But a unit that resumed inside the 10 s diagnostic window answers that probe too, having merely finished the original request late — so the verdict asserted *where* the request was lost on evidence that only shows the unit is serving *now*. It also printed `ops=` from `op_stats`, which `_handle_get_metrics` populates only when Prometheus is enabled, so the line read `ops={}` as if the unit had served nothing. The verdict now follows the unit's own arrival counter for the failed operation: | Probe answered, and… | Verdict | |---|---| | no arrival recorded for that op | `request_lost_in_flight` | | an arrival recorded for that op | `arrived_but_unfinished` | | counters unavailable (e.g. older unit) | `unit_serving_again` | `op_stats` is now reported as `unavailable(prometheus_disabled)` rather than as an empty dict. **The counters could not actually be correlated as intended.** The worker keyed `arrivals_by_op` by `str(operation)`, which for a `(str, Enum)` member renders as `"ZMQRequestType.GET_DATA"`, while `op_stats` is keyed `"GET_DATA"` — the two dicts shared no keys at all. Fixed by keying on `operation.name`. Note `.value` would *also* miss: the enum values are short wire tokens (`GET_DATA = "GET"`). ### 5. Stop the accept probe on shutdown `AcceptQueueProbe.start()` was called but `stop()` had no caller. The finalizer tore down only the ZMQ resources, so a unit finalized without its process exiting left the daemon thread spawning `ss` on a timer indefinitely, and the window summary was never logged. Now stopped in `_shutdown_resources`, before the ZMQ teardown. ## Behavior and compatibility - **No public API change.** `transfer_queue/utils/zmq_utils.py` is untouched. - **One default changes:** ROUTER `BACKLOG` 100 → 4096. - **The probe is off by default** and imported lazily, so a run that does not enable it never loads the module or spawns the thread. - **Steady-state log volume is unchanged**; the near-full-queue warning fires once per probe, not per sample. ## New configuration | Variable | Default | Purpose | |---|---|---| | `TQ_STORAGE_ZMQ_BACKLOG` | `4096` | Accept-queue depth for the client-facing ROUTER | | `TQ_ACCEPT_PROBE_INTERVAL` | `0` (off) | Accept-queue sampling period, seconds; sub-second when enabled | ## Note on the rebase This branch previously carried #171's commits as its base and included a per-call socket-timeout change. Since then #171 merged and #167 replaced per-call sockets with `ZMQSocketPool`, which resolves timeouts from the pool instead of the decorator. That change therefore had nowhere left to live and has been dropped — `main` already routes the two timeouts through separate `storage_rpc_pool` / `storage_probe_pool` instances, which satisfies the same intent. Its test file went with it; pool timeout behaviour is covered by `tests/test_zmq_socket_pool.py::test_timed_out_socket_is_not_reused`. The branch is now four commits on top of `ba8880c`, all mine. ## Testing - `tests/test_accept_probe.py` — utilization and delta arithmetic, peak tracking across a drained burst, drop reported once per drop, each new drop still reported, warn-once behaviour, zero-backlog not dividing by zero, and the probe being stopped on shutdown. - `tests/test_storage_request_retry.py` — the diagnosis table above, plus the empty-`op_stats` case and the arrival-key spelling. - Full suite: **644 passed, 10 skipped**, against a **623**-passing baseline measured on `ba8880c` itself — the +21 are the new tests, no regressions. - `check_license.py`, `check_docstrings.py`, `ruff check` and `ruff format --check` all clean. `tests/test_yuanrong_storage_client_e2e.py` is excluded from those runs: it errors identically on the unmodified base (optional Yuanrong backend not installed locally). --------- Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Every RPC built a socket from scratch: create DEALER → connect → send → recv → close. That pays a TCP and ZMTP handshake per request, and makes the peer's ROUTER accrete a fresh identity on every call.
Approach
ZMQSocketPoollends connected DEALER sockets and takes them back after a clean send/recv. The decoratorwith_zmq_socketnow leases from a pool instead of building a socket per call.Four properties carry the correctness argument:
A lease is exclusive. Responses do not echo their request's
request_id(seeZMQMessage.create), so a reply is matched to its request only by arrival order. Two concurrent users of one socket would read each other's replies.A failed lease discards its socket. Any exception — including timeout and
CancelledError, whichasyncio.gatherraises into siblings on the first failure — closes the socket rather than returning it. The request may already be on the wire, so its reply could still arrive and the next lessee would read it as its own. The notify path closes its socket explicitly for the same reason, without raising, and the pool refuses to park an already-closed socket.Sockets are keyed by lease owner, then by address. The owner is the running event loop, or the thread outside one. pyzmq silently rebinds an async socket to whatever loop it next sees, and one bound to a closed loop is the
Bad file descriptor/ SIGABRT failure that per-call context churn used to cause; a finished owner's sockets are swept on the next access. Keying by address rather than peer id means a peer restarted under the same id at a new address is never handed a socket still connected to the old one.One pool per role. Each role owns the pools for the requests it makes, and sockets are never shared across roles. The pool borrows its context and never terminates it.
Concurrency waits rather than opening more. An async lease takes a permit per (owner, address) first, so sockets in flight follow
TQ_SOCKET_POOL_SIZEinstead of peak concurrency — the context's budget becomes a property of configuration. Waiting also costs less than the handshake it replaces: locally, 300-way concurrency to one peer ran 48ms against 71ms, with peak sockets 4 against 300.One pool per role
The split is deliberate rather than incidental: a pool is scoped to one role and one kind of request, so every socket it holds is interchangeable but for the address it is connected to.
CLIENTcontroller_rpc_poolrequest_handle_socketTransferQueueClientbuilds in__init__, or the caller's for the async clientSTORAGEstorage_rpc_poolput_get_socketSTORAGEnotify_poolrequest_handle_socketCONTROLLERput_get_socketSharing one pool across these would reuse nothing in the first place: they differ by owning loop, socket name, or timeout, and a socket can only be reused when all three match. A ZMQ socket also cannot safely cross event loops, so
notify_poolrunning on its own loop could never have shared with another scenario regardless of how the pools were arranged.Keeping them separate buys isolation on top of that. One role's timeouts and cancellations only ever poison its own sockets, a role's pool is closed with the component that owns it, and the eviction of a finished loop's sockets cannot reach into another role's buckets. Reuse itself does not depend on the split — every entry point above sits on a long-lived owner — but attributing a failure does.
Sizing
TQ_SOCKET_POOL_SIZE(default 4, must be >= 1) is the cap for every pool: idle sockets per (owner, address) bucket, and with an async lease the sockets in flight to that address. Concurrency past it waits for one to come back.The cap multiplies by peer count while a context's ceiling does not, which is what decides the default. A storage manager addressing two thousand units can hold
4 x 2000 = 8004sockets against the client context's 8192 — inside the budget, where 8 would have put it at 16000. That worst case is reachable rather than theoretical: a single peer under 50-way concurrency does saturate its cap. Raise it if one peer routinely sees more overlap than that, and lower it if descriptors are tighter than latency; a singleputorgetsends one merged request per unit, so per-peer concurrency comes only from operations overlapping.AsyncSimpleStorageManagerwarns at construction when the product exceeds its context'sZMQ_MAX_SOCKETS, naming both knobs. Past that ceiling a lease fails outright withEMFILE, so it is worth saying before the first request rather than during one.Ceilings, for reference: the client context allows 8192 (
TQ_CLIENT_ZMQ_MAX_SOCKETS), the controller's 4096 (TQ_CONTROLLER_ZMQ_MAX_SOCKETS), and a storage unit keeps libzmq's 1023 — it binds three sockets regardless of fleet size. Raising any of them needs file descriptors to match (ulimit -n).Fixes found along the way
AsyncTransferQueueClient.__init__leaked the context and its native I/O threads when it raised after allocating it (an out-of-rangemax_sockets). Everything checkable without a live context is now checked first, and the context is destroyed if the remaining check fails.StorageManager.close()raisedAttributeErrorwhen a subclass rejected its config before callingsuper().__init__(), as the KV managers do — burying the constructor's real error under a teardown failure. Pre-existing onmain.AsyncSimpleStorageManager.close()had the same problem forstorage_rpc_poolwhen the base constructor raised, e.g. on a controller handshake timeout._notify_and_waitdecremented its ACK budget by the poll interval, so unrelated traffic on the socket could stretch the wait. It now uses one deadline for the whole wait.TQMetricsExporterlazily minted its ownzmq.Contextwith nothing to close it. It now borrows the controller's.Behaviour and API changes
with_zmq_socket(...)dropssocket_name,get_identity,get_contextandtimeoutin favour ofget_pool; socket name and timeout live on the pool.TQMetricsExporter(role=..., zmq_context=...). Without a context, querying storage units raisesRuntimeErrorinstead of silently building a second context. The controller always passes one.TQ_SOCKET_POOL_SIZE(default 4, must be >= 1) sizes every role's pool. Reuse cannot be disabled: below 1 nothing is ever parked, so every request would pay a fresh connect while still looking pooled.ZMQSocketPool.aleaseis the async entry point and enforces the cap by waiting;leasestays for the synchronous collector, which issues one request at a time and cannot await. Nesting one lease inside another raises, including across pools — a permit lives until its body ends, so nesting can wait on one the task or a sibling already holds, and separate semaphores do not break that cycle. An RPC's reply is what frees its permit, so a second RPC belongs after the first returns.MAX_SOCKETS(4096, clamped toZMQ_SOCKET_LIMIT) instead of keeping libzmq's default 1023. Its own two ROUTERs are nowhere near that, but the metrics exporter borrows this context, and a controller that cannot open a socket stops answering requests at all.StorageManagergainsnotify_pool;AsyncSimpleStorageManagergainsstorage_rpc_pool. Both are closed before the context they live on.Tests
Reuse is asserted from the peer's side rather than from pool state. Every socket dials with its own ZMQ identity, so one identity across many requests means the connection was reused and a new one means the old socket was discarded — the same property observed from outside, with no access to the pool's internals.
Each role's own reuse is covered where that role is tested, since a pool belongs to a role.
test_client.pydrives it through the public client: repeatedget_partition_list()calls must reach the controller over one connection, which is the user-visible payoff of pooling.test_metrics.pycovers the collector, whose pool keys by thread because it runs with no event loop at all. Both were checked by breaking the pool's take path so no socket is ever handed back — each then fails with one identity per request, so neither passes vacuously.What remains in
tests/test_zmq_socket_pool.pyneeds control over reply timing and loop lifetime that a caller-level test cannot reach: the timed-out socket that must not answer the next request; poisoned leases from both an exception and a cancellation; no reuse across event loops; no stale socket after a peer re-registers at a new address; and a burst above the cap.The cap and its permits are covered there too: 40-way concurrency against a cap of 4 must peak at 4 sockets and drop nothing; a busy peer must not stall requests to another, since permits are keyed like the buckets; nesting raises, both on one pool and across two; and consecutive leases are not mistaken for nesting.
The rest cover context sharing and pool wiring (
test_zmq_shared_context.py), the exporter borrowing its owner's context (test_metrics.py), and quiet teardown after a rejected config (test_kv_storage_manager.py).Verification
python -m compileall -q transfer_queue tutorial tests— clean.ruff check,ruff format --checkandmypy— clean.python -m pytest -q:mainThe 8 errors are the same on both:
tests/test_yuanrong_storage_client_e2e.pycannot import the optional Yuanrong backend in this environment.