@@ -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 )
0 commit comments