[perf, fix] Pool and reuse ZMQ request sockets - #167
Conversation
…uest Every request path built a fresh DEALER socket -- create, connect, send, recv, close -- paying a TCP + ZMTP handshake per call, minting a new ROUTER peer identity each time, and calling psutil.virtual_memory() per socket. Sockets are now leased from a long-lived pool: 0.43ms -> 0.19ms per round trip (2.2x), and in-flight socket count tracks concurrency rather than request count, which relieves pressure on TQ_CLIENT_ZMQ_MAX_SOCKETS during large fan-outs. A lease is exclusive. Responses do not echo their request's request_id, 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 socket therefore returns to the pool only after a clean send+recv, and is closed on timeout, error, or cancellation -- a request already on the wire may still get a reply, and the next lessee would read it as its own. asyncio.gather cancels siblings on the first failure, so cancellation is a live path, not a theoretical one. Buckets are keyed by event loop (or thread, outside one), since pyzmq rebinds a socket to whatever loop it next sees and one bound to a closed loop is the documented "Bad file descriptor / SIGABRT" hazard. Finished owners are evicted on the next lease: a pooled async socket references its own loop, so these cannot be reclaimed by garbage collection. The keying also lets the synchronous metrics collector share one pool implementation with the async callers. This fixes a live bug in metrics collection. _query_storage_unit caught zmq.error.Again before the branch that evicted the socket, and Again is an Exception subclass, so after any timeout the next cycle read the previous cycle's reply and attributed it to the wrong cycle. The pool's poison rule removes the bug along with the bookkeeping. Pool and context ownership are now enforced to arrive together: a pool built over a borrowed context would mint sockets the lender may destroy underneath it. Verified: 582 passed, 10 skipped. The 8 remaining errors are pre-existing and require the optional Yuanrong datasystem SDK. The three poison-rule tests fail against a pool without the rule, confirming they are real guards. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Requiring zmq_context and zmq_socket_pool to arrive together broke a previously supported signature. Two reachable paths raised at construction: a caller or third-party manager passing only zmq_context, and any registered manager whose __init__ predates zmq_socket_pool -- StorageManagerFactory filters the keyword it cannot accept and forwards the context alone. A lone context now builds a pool over it instead. Pool and context ownership are tracked separately, so such a manager closes only its own pooled sockets and never the lender's context. Passing both still shares the lender's pool rather than duplicating it. Teardown moves out of the _owns_zmq_context branch accordingly, or a manager-owned pool over a borrowed context would never have been closed. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The pool keyed reuse on (peer id, socket name, timeout), which assumed an id permanently identifies one address. A storage unit or controller restarted or re-registered keeps its id but gets a fresh port, so the pool kept leasing a socket still connected to the address the peer no longer answers on -- reachable via register_storage_units(), which updates ZMQServerInfo entries in place. The key now carries the formatted endpoint. Sockets for an address a peer has moved away from are dropped on the next miss, so past addresses cannot accumulate buckets that each keep libzmq reconnecting to a dead endpoint. maxsize was also unvalidated: TQ_CLIENT_ZMQ_POOL_SIZE=-1 or 0 built a client successfully but made the bucket bound permanently false, silently disabling reuse while still looking pooled. Both the pool and the client now reject sizes below 1, and the client names the environment variable. The key becomes a NamedTuple so address-vs-id and the eviction scan read clearly at the point they matter. Regression tests cover a peer moving to a new address, moving back, and not accumulating stale buckets, plus both rejection paths. Each was confirmed to fail against the specific defect it guards. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Endpoint migration cleanup only scanned idle buckets, so a lease already in flight to the old address escaped it. When that request finished, _release() recreated the bucket and parked a socket wired to the obsolete endpoint indefinitely. The pool now records the address most recently leased per (peer, socket name). Both paths consult it: _take() drops superseded buckets as before, and _release() declines to park a socket whose address has since been retired. A single predicate covers the race in either direction, and the map is bounded by peer count and cleared on close(). An in-flight request still completes normally -- only its socket is discarded rather than reused. Also move the TQ_CLIENT_ZMQ_POOL_SIZE check ahead of context construction. It ran after the context was built and configured but before the finalizer was armed, so an invalid value raised while leaving the context and its native I/O threads open. All cheap validation now happens before anything is allocated. Both regression tests were confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…ng the context _endpoints is pool-wide but the migration sweep only walked the current owner's buckets. The pool is deliberately shared between a client's RPC loop and a storage manager's notify loop, so several owners can hold idle sockets for one peer: when a migration was observed by one, the others kept sockets open against the retired address until they happened to lease again, or forever once they went quiet. The sweep now runs wherever the current endpoint is recorded and covers every owner's buckets, closing those sockets without ever handing one across owners. It is also skipped entirely when the address has not changed, which is the common path. Client construction still leaked the context on two max-sockets paths: TQ_CLIENT_ZMQ_MAX_SOCKETS parsing and the caller-supplied lower-bound check both ran after zmq.asyncio.Context() but before the finalizer was armed. Both are checkable without a context and now run before allocation. Only the upper bound needs ZMQ_SOCKET_LIMIT, so that check stays after construction and destroys the context if it fails. Regression tests cover three concurrently live owners across a migration and every invalid max-sockets input; both were confirmed to fail against the pre-fix code. The context-leak assertion is now a shared helper, since three tests need it. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
One pool served controller RPC, storage RPC and notify, but it never actually shared a socket between them: buckets are keyed by owning loop, and the three differ by loop, socket name, or timeout, so every scenario already had its own sockets. The sharing was cost without benefit -- it required lending a pool alongside a context, ownership rules for the borrowed case, and a key wide enough to separate scenarios inside one pool. Each scenario now builds its own pool: controller_rpc_pool on the client, storage_rpc_pool on the SimpleStorage manager, notify_pool on the base manager, and one in the metrics exporter. They still share the client's context, so the socket budget stays client-wide. Because a pool now serves exactly one scenario, socket_name and timeout move to its constructor, and the key collapses from a four-field _PoolKey to the address. The zmq_socket_pool= parameter disappears from five constructors -- including the three KV managers, which only ever accepted it to keep the factory from warning -- along with the pool-vs-context ownership rules and their divergence guard. Endpoint-migration tracking becomes opt-in via follow_endpoint_changes, enabled only for metrics. Verified this is the sole endpoint mapping mutated at runtime (register_storage_units uses dict.update); client._controller, manager.controller_info and manager.storage_unit_infos are each assigned once at construction, so the other three pools no longer pay for a case they cannot hit. Net 68 lines lighter, with the same 2x speedup on the RPC round trip. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
The endpoint-migration sweep filtered every idle bucket through _superseded(peer_id, addr), which reduces to "addr is not this peer's current address". Applied to the buckets of *other* peers that predicate is always true, so one peer moving closed the pooled sockets of every other peer -- and a peer seen for the first time did the same, since an unrecorded id also compares unequal. The pool that opts into this is the metrics collector's, which queries the whole storage fleet through a single pool, so a single unit re-registering dropped every unit's socket. It self-heals on the next collection cycle, so the cost is a reconnect storm rather than lost traffic, but the sweep was never meant to reach beyond the peer that moved. The sweep now pops the peer's recorded previous address and nothing else, and returns early when there is no previous address to retire. _superseded is left alone: _release calls it with the peer id of the lease being returned, where the predicate is the intended one. The regression test uses three peers, the shape metrics actually runs in, and was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
b9b4f2b to
6707683
Compare
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
| self.zmq_context, | ||
| client_id, | ||
| "request_handle_socket", | ||
| maxsize=TQ_CLIENT_ZMQ_POOL_SIZE, |
There was a problem hiding this comment.
TQ_CLIENT_ZMQ_POOL_SIZE only affects controller RPC pool. The name is misleading
There was a problem hiding this comment.
Thank you for pointing it out. I have renamed it as TQ_CONTROLLER_RPC_POOL_SIZE
| return owner.is_closed() | ||
|
|
||
|
|
||
| class ZMQSocketPool: |
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.
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.
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>
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. 👍 |
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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
…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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
…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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
…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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
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>
CLA Signature PassOutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
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.