Manager for dealer connections, supporting multiplexing and connection reuse
| 77 | |
| 78 | |
| 79 | class DealerConnectionManager: |
| 80 | """ |
| 81 | Manager for dealer connections, supporting multiplexing and connection reuse |
| 82 | """ |
| 83 | |
| 84 | def __init__(self, pid, max_connections=10): |
| 85 | self.pid = pid |
| 86 | self.max_connections = max(max_connections, 10) |
| 87 | self.connections = [] |
| 88 | self.connection_load = [] |
| 89 | self.connection_heap = [] |
| 90 | self.request_map = {} # request_id -> response_queue |
| 91 | self.request_num = {} # request_id -> num_choices |
| 92 | self.lock = asyncio.Lock() |
| 93 | self.connection_tasks = [] |
| 94 | self.running = False |
| 95 | |
| 96 | async def initialize(self): |
| 97 | """initialize all connections""" |
| 98 | self.running = True |
| 99 | for index in range(self.max_connections): |
| 100 | await self._add_connection(index) |
| 101 | api_server_logger.info(f"Started {self.max_connections} connections, pid {self.pid}") |
| 102 | |
| 103 | async def _add_connection(self, index): |
| 104 | """create a new connection and start listening task""" |
| 105 | try: |
| 106 | dealer = await aiozmq.create_zmq_stream( |
| 107 | zmq.DEALER, |
| 108 | connect=f"ipc:///dev/shm/router_{self.pid}.ipc", |
| 109 | ) |
| 110 | async with self.lock: |
| 111 | self.connections.append(dealer) |
| 112 | self.connection_load.append(0) |
| 113 | heapq.heappush(self.connection_heap, (0, index)) |
| 114 | |
| 115 | # start listening |
| 116 | task = asyncio.create_task(self._listen_connection(dealer, index)) |
| 117 | self.connection_tasks.append(task) |
| 118 | return True |
| 119 | except Exception as e: |
| 120 | api_server_logger.error(f"Failed to create dealer: {str(e)}") |
| 121 | return False |
| 122 | |
| 123 | async def _listen_connection(self, dealer, conn_index): |
| 124 | """ |
| 125 | listen for messages from the dealer connection |
| 126 | """ |
| 127 | while self.running: |
| 128 | try: |
| 129 | raw_data = await dealer.read() |
| 130 | response = ForkingPickler.loads(raw_data[-1]) |
| 131 | _zmq_metrics_stats = ZMQMetricsStats() |
| 132 | _zmq_metrics_stats.msg_recv_total += 1 |
| 133 | if "zmq_send_time" in response: |
| 134 | _zmq_metrics_stats.zmq_latency = time.perf_counter() - response["zmq_send_time"] |
| 135 | address = dealer.transport.getsockopt(zmq.LAST_ENDPOINT) |
| 136 | main_process_metrics.record_zmq_stats(_zmq_metrics_stats, address) |
no outgoing calls