modify simple storage

Author: ji-huazhong

transfer_queue/storage/simple_storage.py

@@ -16,7 +16,7 @@
 import os
 import time
 import weakref
-from threading import Event, Thread
+from threading import Event
 from typing import Any
 from uuid import uuid4
 
@@ -31,9 +31,8 @@
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
+    ZMQServerTransport,
     create_zmq_socket,
-    format_zmq_address,
-    get_free_port,
     get_node_ip_address,
 )
 
@@ -159,84 +158,60 @@ def __init__(self, storage_unit_size: int):
         # Shutdown event for graceful termination
         self._shutdown_event = Event()
 
-        # Placeholder for zmq_context, proxy_thread and worker_threads
-        self.zmq_context: zmq.Context | None = None
-        self.put_get_socket: zmq.Socket | None = None
-        self.proxy_thread: Thread | None = None
-        self.worker_thread: Thread | None = None
+        # Placeholder for worker_socket (DEALER, not managed by transport)
+        self.worker_socket: zmq.Socket | None = None
 
-        self._init_zmq_socket()
+        self._init_zmq_transport()
         self._start_process_put_get()
 
         # Register finalizer for graceful cleanup when garbage collected
         self._finalizer = weakref.finalize(
             self,
             self._shutdown_resources,
             self._shutdown_event,
-            self.worker_thread,
-            self.proxy_thread,
-            self.zmq_context,
-            self.put_get_socket,
+            self._transport,
+            self.worker_socket,
         )
 
-    def _init_zmq_socket(self) -> None:
+    def _init_zmq_transport(self) -> None:
         """
-        Initialize ZMQ socket connections between storage unit and controller/clients:
+        Initialize ZMQ transport layer and socket connections:
         - put_get_socket (ROUTER): Handle put/get requests from clients.
         - worker_socket (DEALER): Backend socket for worker communication.
         """
-        self.zmq_context = zmq.Context()
-        self._node_ip = get_node_ip_address()
-
-        # Frontend: ROUTER for receiving client requests
-        self.put_get_socket = create_zmq_socket(self.zmq_context, zmq.ROUTER, self._node_ip)
-
-        while True:
-            try:
-                self._put_get_socket_port = get_free_port(ip=self._node_ip)
-                self.put_get_socket.bind(format_zmq_address(self._node_ip, self._put_get_socket_port))
-                break
-            except zmq.ZMQError:
-                logger.warning(f"[{self.storage_unit_id}]: Try to bind ZMQ sockets failed, retrying...")
-                continue
+        self._transport = ZMQServerTransport(node_ip=get_node_ip_address())
+        self._transport.create_router_socket("put_get_socket")
 
         # Backend: DEALER for worker communication (connected via zmq.proxy)
-        self.worker_socket = create_zmq_socket(self.zmq_context, zmq.DEALER, self._node_ip)
+        self.worker_socket = create_zmq_socket(self._transport.zmq_ctx, zmq.DEALER, self._transport.node_ip)
         self.worker_socket.bind(self._inproc_addr)
 
-        self.zmq_server_info = ZMQServerInfo(
+        self.zmq_server_info = self._transport.build_server_info(
             role=Role.STORAGE,
-            id=str(self.storage_unit_id),
-            ip=self._node_ip,
-            ports={"put_get_socket": self._put_get_socket_port},
+            server_id=str(self.storage_unit_id),
         )
 
     def _start_process_put_get(self) -> None:
         """Start worker threads and ZMQ proxy for handling requests."""
-
         # Start worker thread
-        self.worker_thread = Thread(
+        self._transport.start_daemon_thread(
             target=self._worker_routine,
             name=f"StorageUnitWorkerThread-{self.storage_unit_id}",
-            daemon=True,
         )
-        self.worker_thread.start()
-
         time.sleep(0.5)  # make sure worker thread is ready before zmq.proxy forwarding messages
 
         # Start proxy thread (ROUTER <-> DEALER)
-        self.proxy_thread = Thread(
+        self._transport.start_daemon_thread(
             target=self._proxy_routine,
             name=f"StorageUnitProxyThread-{self.storage_unit_id}",
-            daemon=True,
         )
-        self.proxy_thread.start()
 
     def _proxy_routine(self) -> None:
         """ZMQ proxy for message forwarding between frontend ROUTER and backend DEALER."""
+        put_get_socket = self._transport.get_socket("put_get_socket")
         logger.info(f"[{self.storage_unit_id}]: start ZMQ proxy...")
         try:
-            zmq.proxy(self.put_get_socket, self.worker_socket)
+            zmq.proxy(put_get_socket, self.worker_socket)
         except zmq.ContextTerminated:
             logger.info(f"[{self.storage_unit_id}]: ZMQ Proxy stopped gracefully (Context Terminated)")
         except Exception as e:
@@ -248,7 +223,7 @@ def _proxy_routine(self) -> None:
     def _worker_routine(self) -> None:
         """Worker thread for processing requests."""
 
-        worker_socket = create_zmq_socket(self.zmq_context, zmq.DEALER, self._node_ip)
+        worker_socket = create_zmq_socket(self._transport.zmq_ctx, zmq.DEALER, self._transport.node_ip)
         worker_socket.connect(self._inproc_addr)
 
         poller = zmq.Poller()
@@ -481,30 +456,34 @@ def _handle_clear(self, data_parts: ZMQMessage) -> ZMQMessage:
     @staticmethod
     def _shutdown_resources(
         shutdown_event: Event,
-        worker_thread: Thread | None,
-        proxy_thread: Thread | None,
-        zmq_context: zmq.Context | None,
-        put_get_socket: zmq.Socket | None,
+        transport: ZMQServerTransport | None,
+        worker_socket: zmq.Socket | None,
     ) -> None:
         """Clean up resources on garbage collection."""
         logger.info("Shutting down SimpleStorageUnit resources...")
 
         # Signal all threads to stop
         shutdown_event.set()
 
-        # Terminate put_get_socket
-        if put_get_socket:
-            put_get_socket.close(linger=0)
+        # Close worker socket
+        if worker_socket:
+            worker_socket.close(linger=0)
+
+        # Close put_get_socket
+        if transport:
+            put_get_socket = transport.get_socket("put_get_socket")
+            if put_get_socket:
+                put_get_socket.close(linger=0)
 
         # Terminate ZMQ context to unblock proxy and workers
-        if zmq_context:
-            zmq_context.term()
+        if transport and transport.zmq_ctx:
+            transport.zmq_ctx.term()
 
         # Wait for threads to finish (with timeout)
-        if worker_thread and worker_thread.is_alive():
-            worker_thread.join(timeout=5)
-        if proxy_thread and proxy_thread.is_alive():
-            proxy_thread.join(timeout=5)
+        if transport:
+            for t in transport.threads:
+                if t.is_alive():
+                    t.join(timeout=5)
 
         logger.info("SimpleStorageUnit resources shutdown complete.")