Skip to content

[perf, fix] Pool and reuse ZMQ request sockets - #167

Open
OutstanderWang wants to merge 27 commits into
Ascend:mainfrom
OutstanderWang:feat_socket_pool
Open

[perf, fix] Pool and reuse ZMQ request sockets#167
OutstanderWang wants to merge 27 commits into
Ascend:mainfrom
OutstanderWang:feat_socket_pool

Conversation

@OutstanderWang

@OutstanderWang OutstanderWang commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

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

ZMQSocketPool lends connected DEALER sockets and takes them back after a clean send/recv. The decorator with_zmq_socket now 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 (see ZMQMessage.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, which asyncio.gather raises 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_SIZE instead 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.

Role Pool Dials Owner
CLIENT controller_rpc_pool controller request_handle_socket the loop TransferQueueClient builds in __init__, or the caller's for the async client
STORAGE storage_rpc_pool storage units' put_get_socket the caller's loop
STORAGE notify_pool controller request_handle_socket its own dedicated loop and thread
CONTROLLER metrics pool storage units' put_get_socket the collector thread, with no loop at all

Sharing 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_pool running 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 = 8004 sockets 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 single put or get sends one merged request per unit, so per-peer concurrency comes only from operations overlapping.

AsyncSimpleStorageManager warns at construction when the product exceeds its context's ZMQ_MAX_SOCKETS, naming both knobs. Past that ceiling a lease fails outright with EMFILE, 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-range max_sockets). Everything checkable without a live context is now checked first, and the context is destroyed if the remaining check fails.
  • StorageManager.close() raised AttributeError when a subclass rejected its config before calling super().__init__(), as the KV managers do — burying the constructor's real error under a teardown failure. Pre-existing on main.
  • AsyncSimpleStorageManager.close() had the same problem for storage_rpc_pool when the base constructor raised, e.g. on a controller handshake timeout.
  • _notify_and_wait decremented 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.
  • TQMetricsExporter lazily minted its own zmq.Context with nothing to close it. It now borrows the controller's.
  • Pooling those queries then made the collector hold a socket per storage unit for the life of the process: buckets are per address and only swept when their owner ends, and the collector is a permanent daemon thread. Each query now closes its socket, as the notify path does, taking resident sockets from one per unit to zero. Collection walks every unit once per cycle, so a kept socket would be reused only a cycle later while occupying budget for the whole walk.

Behaviour and API changes

  • with_zmq_socket(...) drops socket_name, get_identity, get_context and timeout in favour of get_pool; socket name and timeout live on the pool.
  • TQMetricsExporter(role=..., zmq_context=...). Without a context, querying storage units raises RuntimeError instead of silently building a second context. The controller always passes one.
  • New 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.alease is the async entry point and enforces the cap by waiting; lease stays 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.
  • The controller's context now sets MAX_SOCKETS (4096, clamped to ZMQ_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.
  • StorageManager gains notify_pool; AsyncSimpleStorageManager gains storage_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.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. test_metrics.py covers 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.py needs 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 --check and mypy — clean.

python -m pytest -q:

passed skipped errors
this branch 607 10 8
main 584 10 8

The 8 errors are the same on both: tests/test_yuanrong_storage_client_e2e.py cannot import the optional Yuanrong backend in this environment.

…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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Comment thread transfer_queue/client.py Outdated
self.zmq_context,
client_id,
"request_handle_socket",
maxsize=TQ_CLIENT_ZMQ_POOL_SIZE,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TQ_CLIENT_ZMQ_POOL_SIZE only affects controller RPC pool. The name is misleading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for pointing it out. I have renamed it as TQ_CONTROLLER_RPC_POOL_SIZE

return owner.is_closed()


class ZMQSocketPool:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@OutstanderWang OutstanderWang Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, 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>
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants