[perf, fix] Pool and reuse ZMQ request sockets

tests/test_client.py

@@ -68,6 +68,7 @@ def __init__(self, controller_id="controller_0"):
         )
 
         self.running = True
+        self.caller_identities: set[bytes] = set()
         self.request_thread = Thread(target=self._handle_requests, daemon=True)
         self.request_thread.start()
 
@@ -85,6 +86,9 @@ def _handle_requests(self):
                 if self.request_socket in socks:
                     messages = self.request_socket.recv_multipart(copy=False)
                     identity = messages.pop(0)
+                    # Each pooled socket dials with its own identity, so this records how
+                    # many distinct connections the client opened.
+                    self.caller_identities.add(bytes(identity))
                     serialized_msg = messages
                     request_msg = ZMQMessage.deserialize(serialized_msg)
 
@@ -1330,6 +1334,26 @@ def test_kv_retrieve_keys_type_validation(self, client_setup):
             )
 
 
+# =====================================================
+# Controller RPC Socket Pool Tests
+# =====================================================
+
+
+def test_controller_rpc_reuses_one_socket(client_setup):
+    """Repeated calls must share a single connection to the controller.
+
+    This is the client-side payoff of the socket pool: TransferQueueClient runs every call
+    on the one event loop it builds in __init__, so all of them land in the same pool
+    bucket. Counted from the controller, which sees one ZMQ identity per socket dialled.
+    """
+    client, mock_controller, _ = client_setup
+
+    for _ in range(5):
+        assert client.get_partition_list() == ["partition_0", "partition_1", "test_partition"]
+
+    assert len(mock_controller.caller_identities) == 1, "each request opened its own socket"
+
+
 # =====================================================
 # Checkpoint Interface Tests
 # =====================================================

tests/test_controller.py

@@ -21,7 +21,7 @@
 import torch
 import zmq
 
-from transfer_queue.controller import TransferQueueController
+from transfer_queue.controller import TQ_CONTROLLER_ZMQ_MAX_SOCKETS, TransferQueueController
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket
 
 # Set up logging
@@ -1465,6 +1465,27 @@ def test_controller_checkpoint_load_nonexistent_file(self, ray_setup, tmp_path):
         print("✓ load_checkpoint raises on missing file")
 
 
+class TestControllerSocketBudget:
+    """The controller's context must hold more sockets than libzmq's default 1023.
+
+    The metrics exporter borrows this context to query storage units, so the budget has to
+    cover the fleet rather than just the controller's own two ROUTERs. Past the ceiling a
+    socket cannot be opened at all, which would take out the control plane too.
+    """
+
+    def test_context_ceiling_is_raised_above_the_libzmq_default(self, ray_setup):
+        controller = TransferQueueController.remote()
+        try:
+            # Read it from inside the actor: the ceiling is only meaningful on the context
+            # the controller actually built.
+            budget = ray.get(controller.__ray_call__.remote(lambda self: self.zmq_context.get(zmq.MAX_SOCKETS)))
+        finally:
+            ray.kill(controller)
+
+        assert budget > 1023, "the controller kept libzmq's default socket ceiling"
+        assert budget == TQ_CONTROLLER_ZMQ_MAX_SOCKETS
+
+
 class TestTransferQueueControllerBadRequests:
     """The request loop must survive requests it cannot decode, does not handle, or fails on.
 

tests/test_kv_storage_manager.py

@@ -23,6 +23,28 @@
 from transfer_queue.storage.managers.base import KVStorageManager
 
 
+def test_close_is_quiet_when_config_is_rejected_before_super_init():
+    """A config rejected before super().__init__() leaves nothing for close() to release.
+
+    Every KV manager validates its config first, so none of the base attributes close()
+    reaches for exist yet. __del__ calls close() regardless, and an AttributeError there
+    buries the real ValueError under a spurious teardown error.
+    """
+    built = []
+
+    class Probe(KVStorageManager):
+        def __init__(self, *args, **kwargs):
+            built.append(self)
+            super().__init__(*args, **kwargs)
+
+    with pytest.raises(ValueError, match="Missing client_name"):
+        Probe(None, {})
+
+    manager = built[0]
+    assert not hasattr(manager, "notify_pool"), "the test no longer exercises the pre-super path"
+    manager.close()
+
+
 def get_meta(data, global_indexes=None):
     if not global_indexes:
         global_indexes = list(range(data.batch_size[0]))

tests/test_metrics.py

@@ -16,12 +16,17 @@
 """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics)."""
 
 import time
-from unittest.mock import MagicMock
+from threading import Thread
+from unittest.mock import MagicMock, patch
 
 import pytest
 
 try:
+    import zmq
+
     from transfer_queue.metrics import TQMetricsExporter
+    from transfer_queue.utils.enum_utils import Role
+    from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo
 
     _HAS_DEPS = True
 except (ImportError, OSError):
@@ -207,6 +212,77 @@ def test_multiple_ops_tracked_independently(self):
 # ---------------------------------------------------------------------------
 
 
+class TestStorageQuerySocketPool:
+    def test_pool_borrows_the_owner_context(self):
+        """The exporter must query storage units over the context it was handed.
+
+        The controller already holds a long-lived synchronous context. A second one would
+        add another native I/O thread and leave a context nobody closes, since the exporter
+        lives for the whole life of its Ray actor.
+        """
+        ctx = zmq.Context()
+        try:
+            exporter = TQMetricsExporter(zmq_context=ctx)
+            with patch("zmq.Context") as minted:
+                exporter._get_socket_pool()
+            minted.assert_not_called()
+        finally:
+            ctx.destroy(linger=0)
+
+    def test_missing_context_is_reported(self):
+        """Without a context there is nothing to query over, so say so rather than mint one."""
+        exporter = TQMetricsExporter()
+        with pytest.raises(RuntimeError, match="without a ZMQ context"):
+            exporter._get_socket_pool()
+
+    def test_collector_parks_no_socket_between_queries(self):
+        """A queried unit must leave nothing in the pool.
+
+        Collection walks every unit once per cycle, so a parked socket is reused only a
+        cycle later while holding a slot in the controller context's budget for the whole
+        walk. At a few thousand units that budget is what runs out first.
+        """
+        identities: set[bytes] = set()
+        ctx_peer = zmq.Context()
+        router = ctx_peer.socket(zmq.ROUTER)
+        port = router.bind_to_random_port("tcp://127.0.0.1")
+        running = True
+
+        def serve():
+            poller = zmq.Poller()
+            poller.register(router, zmq.POLLIN)
+            while running:
+                if not dict(poller.poll(50)):
+                    continue
+                identity, _ = router.recv_multipart()
+                identities.add(bytes(identity))
+                response = ZMQMessage.create(
+                    request_type=ZMQRequestType.METRICS_RESPONSE,
+                    sender_id="storage_0",
+                    body={},
+                )
+                router.send_multipart([identity, *response.serialize()])
+
+        server = Thread(target=serve, daemon=True)
+        server.start()
+
+        su_info = ZMQServerInfo(role=Role.STORAGE, id="storage_0", ip="127.0.0.1", ports={"put_get_socket": port})
+        ctx = zmq.Context()
+        try:
+            exporter = TQMetricsExporter(zmq_context=ctx)
+            for _ in range(3):
+                assert exporter._query_storage_unit(su_info, "storage_0") == {}
+            # A parked socket would be reused, so a fresh identity per query is the
+            # externally visible proof that nothing was kept.
+            assert len(identities) == 3, "a queried unit left its socket in the pool"
+        finally:
+            running = False
+            server.join(timeout=2.0)
+            ctx.destroy(linger=0)
+            router.close(linger=0)
+            ctx_peer.term()
+
+
 class TestStorageMetricsCollection:
     def test_collect_with_no_storage_units(self):
         """No storage units registered — collect should be a no-op."""