Skip to content

Commit 7d9dadd

Browse files
committed
remove unnecessary poller
Signed-off-by: ji-huazhong <hzji210@gmail.com>
1 parent 9ff85a9 commit 7d9dadd

2 files changed

Lines changed: 31 additions & 42 deletions

File tree

transfer_queue/controller.py

Lines changed: 30 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1040,51 +1040,44 @@ def _start_daemon_threads(self):
10401040
)
10411041

10421042
def _wait_connection(self):
1043-
"""Wait for storage instances to complete handshake with retransmission support."""
1043+
"""Wait for storage instances to complete handshake."""
10441044
handshake_socket = self._transport.get_socket("handshake_socket")
1045-
poller = zmq.Poller()
1046-
poller.register(handshake_socket, zmq.POLLIN)
10471045

10481046
logger.debug(f"Controller {self.controller_id} started waiting for storage connections...")
10491047

10501048
while True:
1051-
socks = dict(poller.poll(1000))
1049+
try:
1050+
messages = handshake_socket.recv_multipart(copy=False)
1051+
identity = messages.pop(0)
1052+
serialized_msg = messages
1053+
request_msg = ZMQMessage.deserialize(serialized_msg)
10521054

1053-
if handshake_socket in socks:
1054-
try:
1055-
messages = handshake_socket.recv_multipart(copy=False)
1056-
identity = messages.pop(0)
1057-
serialized_msg = messages
1058-
request_msg = ZMQMessage.deserialize(serialized_msg)
1055+
if request_msg.request_type == ZMQRequestType.HANDSHAKE:
1056+
storage_manager_id = request_msg.sender_id
10591057

1060-
if request_msg.request_type == ZMQRequestType.HANDSHAKE:
1061-
storage_manager_id = request_msg.sender_id
1062-
1063-
# Always send ACK for HANDSHAKE
1064-
response_msg = ZMQMessage.create(
1065-
request_type=ZMQRequestType.HANDSHAKE_ACK,
1066-
sender_id=self.controller_id,
1067-
body={},
1068-
).serialize()
1069-
handshake_socket.send_multipart([identity, *response_msg])
1070-
1071-
# Track new connections
1072-
if storage_manager_id not in self._connected_storage_managers:
1073-
self._connected_storage_managers.add(storage_manager_id)
1074-
storage_manager_type = request_msg.body.get("storage_manager_type", "Unknown")
1075-
logger.debug(
1076-
f"[{self.controller_id}]: received handshake from "
1077-
f"storage manager {storage_manager_id} (type: {storage_manager_type}). "
1078-
f"Total connected: {len(self._connected_storage_managers)}"
1079-
)
1080-
else:
1081-
logger.debug(
1082-
f"[{self.controller_id}]: received duplicate handshake from "
1083-
f"storage manager {storage_manager_id}. Resending ACK."
1084-
)
1058+
response_msg = ZMQMessage.create(
1059+
request_type=ZMQRequestType.HANDSHAKE_ACK,
1060+
sender_id=self.controller_id,
1061+
body={},
1062+
).serialize()
1063+
handshake_socket.send_multipart([identity, *response_msg])
1064+
1065+
if storage_manager_id not in self._connected_storage_managers:
1066+
self._connected_storage_managers.add(storage_manager_id)
1067+
storage_manager_type = request_msg.body.get("storage_manager_type", "Unknown")
1068+
logger.debug(
1069+
f"[{self.controller_id}]: received handshake from "
1070+
f"storage manager {storage_manager_id} (type: {storage_manager_type}). "
1071+
f"Total connected: {len(self._connected_storage_managers)}"
1072+
)
1073+
else:
1074+
logger.debug(
1075+
f"[{self.controller_id}]: received duplicate handshake from "
1076+
f"storage manager {storage_manager_id}. Resending ACK."
1077+
)
10851078

1086-
except Exception as e:
1087-
logger.error(f"[{self.controller_id}]: error processing handshake: {e}")
1079+
except Exception as e:
1080+
logger.error(f"[{self.controller_id}]: error processing handshake: {e}")
10881081

10891082
def _process_request(self):
10901083
"""Main request processing loop - adapted for partition-based operations."""
@@ -1701,7 +1694,6 @@ def get_metadata(
17011694

17021695
elif mode == "force_fetch":
17031696
batch_global_indexes = self.index_manager.get_indexes_for_partition(partition_id)
1704-
consumed_indexes = []
17051697

17061698
# Package into metadata
17071699
metadata = self.generate_batch_meta(partition_id, batch_global_indexes, data_fields, mode)

transfer_queue/utils/zmq_utils.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -188,10 +188,7 @@ def deserialize(cls, frames: list) -> "ZMQMessage":
188188

189189

190190
class ZMQServerTransport:
191-
"""Unified ZMQ transport abstraction for Controller / StorageUnit.
192-
Encapsulates socket creation, binding, inproc endpoint, daemon thread,
193-
ZMQ proxy lifecycle, and unified resource cleanup.
194-
"""
191+
"""Unified management of ZMQ Router Sockets, port binding, daemon threads, and message I/O."""
195192

196193
def __init__(self, node_ip: str, ctx: zmq.Context | None = None):
197194
self.node_ip = node_ip

0 commit comments

Comments
 (0)