update

Signed-off-by: ji-huazhong <hzji210@gmail.com>

Author: ji-huazhong

transfer_queue/client.py

@@ -36,7 +36,7 @@
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
-    dynamic_zmq_socket,
+    with_zmq_socket,
 )
 
 logger = logging.getLogger(__name__)
@@ -51,10 +51,10 @@
 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
 
 # Pre-bound decorator for controller socket operations.
-_controller_socket = dynamic_zmq_socket(
+_controller_socket = with_zmq_socket(
     "request_handle_socket",
-    owner_id_attr="client_id",
-    server_attr="_controller",
+    get_identity=lambda self: self.client_id,
+    get_peer=lambda self, target: self._controller,
 )
 
 

transfer_queue/storage/managers/simple_backend_manager.py

@@ -34,7 +34,7 @@
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
-    dynamic_zmq_socket,
+    with_zmq_socket,
 )
 
 logger = logging.getLogger(__name__)
@@ -49,11 +49,11 @@
 TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200))  # seconds
 
 # Pre-bound decorator for storage-unit socket operations.
-_storage_unit_socket = dynamic_zmq_socket(
+_storage_unit_socket = with_zmq_socket(
     "put_get_socket",
-    owner_id_attr="storage_manager_id",
-    server_attr="storage_unit_infos",
-    target_kwarg="target_storage_unit",
+    get_identity=lambda self: self.storage_manager_id,
+    get_peer=lambda self, target: self.storage_unit_infos[target],
+    resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"),
     timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT,
 )
 

transfer_queue/utils/zmq_utils.py

@@ -304,12 +304,12 @@ def create_zmq_socket(
     return socket
 
 
-def dynamic_zmq_socket(
+def with_zmq_socket(
     socket_name: str,
     *,
-    owner_id_attr: str,
-    server_attr: str,
-    target_kwarg: Optional[str] = None,
+    get_identity: Callable[[Any], str],
+    get_peer: Callable[[Any, Optional[str]], ZMQServerInfo],
+    resolve_target: Optional[Callable[[tuple, dict], Optional[str]]] = None,
     timeout: Optional[int] = None,
 ):
     """Create a reusable async decorator for request sockets.
@@ -320,50 +320,34 @@ def dynamic_zmq_socket(
 
     Args:
         socket_name: Socket port key in ``ZMQServerInfo.ports``.
-        owner_id_attr: Attribute name on ``self`` used in identity/log prefix
-            (e.g., ``client_id`` or ``storage_manager_id``).
-        server_attr: Attribute name on ``self`` that stores server info.
-            - ``ZMQServerInfo`` for single-target calls.
-            - ``Mapping[str, ZMQServerInfo]`` for multi-target calls.
-        target_kwarg: Optional kwarg name that provides target server id when
-            ``server_attr`` is a mapping.
+        get_identity: Callable that extracts owner identity from ``self``.
+            Example: ``lambda self: self.client_id``
+        get_peer: Callable that returns ``ZMQServerInfo`` for the target.
+            For single-target scenarios, ignore the target parameter.
+            Example: ``lambda self, target: self.server_info``
+            Example: ``lambda self, target: self.storage_unit_infos[target]``
+        resolve_target: Optional callable that extracts target identifier from
+            function arguments. Receives (args, kwargs) and returns target name.
+            Example: ``lambda args, kwargs: kwargs.get("target_storage_unit")``
         timeout: Optional timeout (seconds) for both send/recv operations.
     """
 
     def decorator(func: Callable):
         @wraps(func)
         async def wrapper(self, *args, **kwargs):
-            owner_id = getattr(self, owner_id_attr, None)
+            owner_id = get_identity(self)
             if owner_id is None:
-                raise RuntimeError(f"Missing owner id attribute: {owner_id_attr}")
-
-            server_obj = getattr(self, server_attr, None)
-            if server_obj is None:
-                raise RuntimeError(f"Missing server registry attribute: {server_attr}")
+                raise RuntimeError("get_identity returned None")
 
             target_name: Optional[str] = None
-            if target_kwarg is not None:
-                target_name = kwargs.get(target_kwarg)
-                if target_name is None:
-                    for arg in args:
-                        if isinstance(arg, str):
-                            target_name = arg
-                            break
-
-            if isinstance(server_obj, ZMQServerInfo):
-                if target_name is not None and target_name != server_obj.id:
-                    raise RuntimeError(
-                        f"Target mismatch: target '{target_name}' does not match registered server '{server_obj.id}'"
-                    )
-                server_info = server_obj
-            elif isinstance(server_obj, Mapping):
-                if target_name is None:
-                    raise RuntimeError(f"Missing target server identifier via '{target_kwarg}'")
-                server_info = server_obj.get(target_name)
-                if server_info is None:
-                    raise RuntimeError(f"Server '{target_name}' not found in registered servers")
-            else:
-                raise RuntimeError(f"Unsupported server registry type for '{server_attr}': {type(server_obj).__name__}")
+            if resolve_target is not None:
+                target_name = resolve_target(args, kwargs)
+
+            server_info = get_peer(self, target_name)
+            if server_info is None:
+                raise RuntimeError(
+                    f"get_peer returned None for target '{target_name}'"
+                )
 
             port = server_info.ports.get(socket_name)
             if port is None: