[perf, fix] Reuse a long-lived ZMQ context instead of creating one per call

tests/test_zmq_shared_context.py

@@ -0,0 +1,338 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Regression tests for the shared long-lived ZMQ context in with_zmq_socket.
+
+with_zmq_socket used to create and term() a context per RPC call, which churned libzmq
+signaler file descriptors and crashed the process under concurrency (Bad file descriptor
+-> SIGABRT). It now reuses the owner's context and only creates the socket per call, so
+these tests assert every concurrent call sees the SAME context, alive until close().
+"""
+
+import asyncio
+from threading import Thread
+from unittest.mock import patch
+
+import pytest
+import zmq
+
+import transfer_queue.utils.zmq_utils as zmq_utils
+from transfer_queue.client import AsyncTransferQueueClient, TransferQueueClient
+from transfer_queue.metadata import BatchMeta
+from transfer_queue.storage.managers.base import StorageManager
+from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager
+from transfer_queue.utils.enum_utils import Role
+from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo
+
+
+class _EchoController:
+    """Minimal in-process ROUTER controller that answers GET_META requests."""
+
+    def __init__(self, controller_id="controller_0"):
+        self.controller_id = controller_id
+        self.context = zmq.Context()
+        self.request_socket = self.context.socket(zmq.ROUTER)
+        self.request_port = self.request_socket.bind_to_random_port("tcp://127.0.0.1")
+        self.zmq_server_info = ZMQServerInfo(
+            role=Role.CONTROLLER,
+            id=controller_id,
+            ip="127.0.0.1",
+            ports={"request_handle_socket": self.request_port},
+        )
+        self.running = True
+        self.request_thread = Thread(target=self._handle_requests, daemon=True)
+        self.request_thread.start()
+
+    def _handle_requests(self):
+        poller = zmq.Poller()
+        poller.register(self.request_socket, zmq.POLLIN)
+        while self.running:
+            try:
+                socks = dict(poller.poll(100))
+                if self.request_socket not in socks:
+                    continue
+                messages = self.request_socket.recv_multipart(copy=False)
+                identity = messages.pop(0)
+                request_msg = ZMQMessage.deserialize(messages)
+
+                batch_size = request_msg.body.get("batch_size", 1)
+                data_fields = request_msg.body.get("data_fields", [])
+                field_schema = {
+                    name: {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": False}
+                    for name in data_fields
+                }
+                metadata = BatchMeta(
+                    global_indexes=list(range(batch_size)),
+                    partition_ids=["0"] * batch_size,
+                    field_schema=field_schema,
+                )
+                response_msg = ZMQMessage.create(
+                    request_type=ZMQRequestType.GET_META_RESPONSE,
+                    sender_id=self.controller_id,
+                    receiver_id=request_msg.sender_id,
+                    body={"metadata": metadata},
+                )
+                self.request_socket.send_multipart([identity, *response_msg.serialize()])
+            except zmq.Again:
+                continue
+            except Exception as e:  # pragma: no cover - surfaced via test failure
+                print(f"_EchoController ERROR: {e}")
+
+    def stop(self):
+        self.running = False
+        self.request_thread.join(timeout=2.0)
+        self.request_socket.close(linger=0)
+        self.context.term()
+
+
+@pytest.fixture
+def echo_controller():
+    controller = _EchoController()
+    yield controller
+    controller.stop()
+
+
+@pytest.mark.asyncio
+async def test_shared_context_reused_across_concurrent_calls(echo_controller, monkeypatch):
+    """Many concurrent decorated RPCs must all reuse the client's single context.
+
+    This is the core regression guard: pre-fix, each call created and term()ed its own
+    context, which is what corrupted libzmq's signaler FDs under concurrency.
+    """
+    client = AsyncTransferQueueClient(
+        client_id="client_shared_ctx",
+        controller_info=echo_controller.zmq_server_info,
+    )
+
+    # Record the context object handed to every socket creation inside the decorator.
+    seen_contexts = []
+    original_create = zmq_utils.create_zmq_socket
+
+    def _spy_create(ctx, *args, **kwargs):
+        seen_contexts.append(ctx)
+        return original_create(ctx, *args, **kwargs)
+
+    # The decorator resolves create_zmq_socket from the zmq_utils module globals.
+    monkeypatch.setattr(zmq_utils, "create_zmq_socket", _spy_create)
+
+    num_calls = 200
+    coros = [
+        client.async_get_meta(data_fields=["tokens", "labels"], batch_size=2, partition_id="0")
+        for _ in range(num_calls)
+    ]
+    # wait_for guards against the hang failure mode.
+    results = await asyncio.wait_for(asyncio.gather(*coros), timeout=60)
+
+    assert len(results) == num_calls
+    assert all(isinstance(meta, BatchMeta) for meta in results)
+
+    # Every call must have used the SAME context, and it must be the client's context.
+    assert len(seen_contexts) == num_calls
+    assert all(ctx is client.zmq_context for ctx in seen_contexts)
+    # The shared context must NOT have been terminated by any call.
+    assert not client.zmq_context.closed
+
+    client.close()
+
+
+def test_client_context_has_fixed_io_thread_pool(echo_controller):
+    client = AsyncTransferQueueClient(
+        client_id="client_fixed_context_pool",
+        controller_info=echo_controller.zmq_server_info,
+        zmq_io_threads=4,
+    )
+
+    assert client.zmq_context.get(zmq.IO_THREADS) == 4
+
+    client.close()
+
+
+def test_client_rejects_invalid_context_pool_size(echo_controller):
+    with pytest.raises(ValueError, match="at least 1"):
+        AsyncTransferQueueClient(
+            client_id="client_invalid_context_pool",
+            controller_info=echo_controller.zmq_server_info,
+            zmq_io_threads=0,
+        )
+
+
+def test_simple_storage_borrows_client_context(echo_controller):
+    client = AsyncTransferQueueClient(
+        client_id="client_simple_storage_context",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    config = {"zmq_info": {}}
+
+    with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager:
+        client.initialize_storage_manager("SimpleStorage", config)
+
+    create_manager.assert_called_once_with(
+        "SimpleStorage",
+        controller_info=echo_controller.zmq_server_info,
+        config=config,
+        zmq_context=client.zmq_context,
+    )
+
+    client.close()
+
+
+def test_simple_storage_does_not_destroy_borrowed_context(echo_controller):
+    client = AsyncTransferQueueClient(
+        client_id="client_borrowed_context_lifecycle",
+        controller_info=echo_controller.zmq_server_info,
+    )
+
+    with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"):
+        manager = AsyncSimpleStorageManager(
+            echo_controller.zmq_server_info,
+            {"zmq_info": {"storage_0": echo_controller.zmq_server_info}},
+            zmq_context=client.zmq_context,
+        )
+
+    assert manager.zmq_context is client.zmq_context
+    assert not manager._owns_zmq_context
+
+    manager.close()
+    assert not client.zmq_context.closed
+
+    client.close()
+    assert client.zmq_context.closed
+
+
+@pytest.mark.asyncio
+async def test_close_destroys_context(echo_controller):
+    """close() must terminate the shared context exactly once (no leak, no hang)."""
+    client = AsyncTransferQueueClient(
+        client_id="client_close_ctx",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    assert not client.zmq_context.closed
+
+    # A normal call before shutdown leaves the context alive.
+    await asyncio.wait_for(
+        client.async_get_meta(data_fields=["tokens"], batch_size=1, partition_id="0"),
+        timeout=30,
+    )
+    assert not client.zmq_context.closed
+
+    client.close()
+    assert client.zmq_context.closed
+
+
+def test_non_numeric_max_sockets_env_var_names_the_variable(echo_controller):
+    """A typo'd value must say which env var is wrong, not raise a bare int() error."""
+    with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "not-a-number"):
+        with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer"):
+            AsyncTransferQueueClient(
+                client_id="client_max_sockets_garbage",
+                controller_info=echo_controller.zmq_server_info,
+            )
+
+
+def test_close_skips_destroy_while_loop_thread_alive(echo_controller):
+    """destroy() is not thread-safe, so a stuck loop thread must veto it.
+
+    TransferQueueClient.close() joins its loop thread with a timeout that only warns on
+    expiry. Falling through to destroy() with that thread still holding sockets is the
+    documented hazard, so the context is leaked instead.
+    """
+    client = TransferQueueClient(
+        client_id="client_stuck_thread",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    context = client.zmq_context
+
+    # Simulate a loop thread that outlived its join timeout.
+    with patch.object(client._thread, "is_alive", return_value=True):
+        assert client._can_destroy_zmq_context() is False
+        client.close()
+
+    assert not context.closed, "context must be leaked, not destroyed unsafely"
+
+    # With the thread genuinely gone, the veto lifts.
+    assert client._can_destroy_zmq_context() is True
+    context.destroy(linger=0)
+
+
+def _make_borrowing_manager(zmq_context):
+    """A minimal manager that borrows a caller's context, like SimpleStorage does."""
+
+    class Borrower(StorageManager):
+        def _connect_to_controller(self):
+            pass
+
+        async def put_data(self, *args, **kwargs):
+            return None
+
+        async def get_data(self, *args, **kwargs):
+            return None
+
+        async def clear_data(self, *args, **kwargs):
+            return None
+
+    return Borrower(None, {}, zmq_context=zmq_context)
+
+
+def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller):
+    """A borrowing manager's stuck notify thread must veto the owner's destroy().
+
+    The manager has no destroy() of its own to skip, so staying silent would let the client
+    destroy a context whose sockets that thread may still hold. It must be asked first.
+    """
+    client = AsyncTransferQueueClient(
+        client_id="client_notify_thread_veto",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    client.storage_manager = _make_borrowing_manager(client.zmq_context)
+    context = client.zmq_context
+
+    with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True):
+        assert client._can_destroy_zmq_context() is False
+        client.close()
+        assert not context.closed, "context must be leaked while the notify thread lives"
+
+    # Veto lifts once the thread is genuinely gone.
+    assert client._can_destroy_zmq_context() is True
+    context.destroy(linger=0)
+
+
+def test_healthy_notify_thread_does_not_block_destroy(echo_controller):
+    """The veto must not over-trigger: a clean manager shutdown still destroys."""
+    client = AsyncTransferQueueClient(
+        client_id="client_notify_thread_clean",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    client.storage_manager = _make_borrowing_manager(client.zmq_context)
+
+    client.close()
+    assert client.zmq_context.closed
+
+
+def test_manager_with_own_context_does_not_veto(echo_controller):
+    """A manager holding its own context has no say in the client's teardown."""
+    client = AsyncTransferQueueClient(
+        client_id="client_independent_manager",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    client.storage_manager = _make_borrowing_manager(None)  # creates its own context
+    assert client.storage_manager.zmq_context is not client.zmq_context
+
+    # Even a stuck notify thread on an unrelated context must not block the client.
+    with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True):
+        assert client._can_destroy_zmq_context() is True
+
+    client.storage_manager.zmq_context.destroy(linger=0)
+    client.close()
+    assert client.zmq_context.closed