async def test_grpc_connection_pool_real_sending_timeout(): server1_ready_event = multiprocessing.Event() def listen(port, event: multiprocessing.Event): class DummyServer: async def process_control(self, request, *args): returned_msg = ControlRequest(command='DEACTIVATE') await asyncio.sleep(0.1) return returned_msg async def start_grpc_server(): grpc_server = grpc.aio.server(options=[ ('grpc.max_send_request_length', -1), ('grpc.max_receive_message_length', -1), ]) jina_pb2_grpc.add_JinaControlRequestRPCServicer_to_server( DummyServer(), grpc_server) service_names = ( jina_pb2.DESCRIPTOR.services_by_name['JinaControlRequestRPC']. full_name, reflection.SERVICE_NAME, ) reflection.enable_server_reflection(service_names, grpc_server) grpc_server.add_insecure_port(f'localhost:{port}') await grpc_server.start() event.set() await grpc_server.wait_for_termination() asyncio.run(start_grpc_server()) port1 = random_port() server_process1 = Process( target=listen, args=( port1, server1_ready_event, ), ) server_process1.start() time.sleep(0.1) server1_ready_event.wait() pool = GrpcConnectionPool() pool.add_connection(deployment='encoder', head=False, address=f'localhost:{port1}') sent_msg = ControlRequest(command='STATUS') results_call_1 = pool.send_request(request=sent_msg, deployment='encoder', head=False, timeout=1.0) assert len(results_call_1) == 1 response1, meta = await results_call_1[0] assert response1.command == 'DEACTIVATE' results_call_2 = pool.send_request(request=sent_msg, deployment='encoder', head=False, timeout=0.05) assert len(results_call_2) == 1 with pytest.raises(AioRpcError): await results_call_2[0] await pool.close() server_process1.kill() server_process1.join()
async def test_connection_pool(mocker, monkeypatch): close_mock_object, create_mock = await _mock_grpc(mocker, monkeypatch) pool = GrpcConnectionPool() send_mock = mocker.Mock() pool._send_requests = lambda messages, connection, endpoint: mock_send( send_mock) pool.add_connection(deployment='encoder', head=False, address='1.1.1.1:53') pool.add_connection(deployment='encoder', head=False, address='1.1.1.2:53') results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='encoder', head=False) assert len(results) == 1 assert send_mock.call_count == 1 assert create_mock.call_count == 2 results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='encoder', head=False) assert len(results) == 1 assert send_mock.call_count == 2 assert create_mock.call_count == 2 # indexer was not added yet, so there isnt anything being sent results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='indexer', head=False) assert len(results) == 0 assert send_mock.call_count == 2 assert create_mock.call_count == 2 # add indexer now so it can be send pool.add_connection(deployment='indexer', head=False, address='2.1.1.1:53') results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='indexer', head=False) assert len(results) == 1 assert send_mock.call_count == 3 assert create_mock.call_count == 3 # polling only applies to shards, there are no shards here, so it only sends one message pool.add_connection(deployment='encoder', head=False, address='1.1.1.3:53') results = pool.send_request( request=ControlRequest(command='STATUS'), deployment='encoder', head=False, polling_type=PollingType.ALL, ) assert len(results) == 1 assert send_mock.call_count == 4 assert create_mock.call_count == 4 # polling only applies to shards, so we add a shard now and expect 2 messages being sent pool.add_connection(deployment='encoder', head=False, address='1.1.1.3:53', shard_id=1) # adding the same connection again is a noop pool.add_connection(deployment='encoder', head=False, address='1.1.1.3:53', shard_id=1) results = pool.send_request( request=ControlRequest(command='STATUS'), deployment='encoder', head=False, polling_type=PollingType.ALL, ) assert len(results) == 2 assert send_mock.call_count == 6 assert create_mock.call_count == 5 # sending to one specific shard should only send one message results = pool.send_request( request=ControlRequest(command='STATUS'), deployment='encoder', head=False, polling_type=PollingType.ANY, shard_id=1, ) assert len(results) == 1 assert send_mock.call_count == 7 # doing the same with polling ALL ignores the shard id results = pool.send_request( request=ControlRequest(command='STATUS'), deployment='encoder', head=False, polling_type=PollingType.ALL, shard_id=1, ) assert len(results) == 2 assert send_mock.call_count == 9 # removing a replica for shard 0 works and does not prevent messages to be sent to the shard assert await pool.remove_connection(deployment='encoder', head=False, address='1.1.1.2:53', shard_id=0) assert close_mock_object.call_count == 1 results = pool.send_request( request=ControlRequest(command='STATUS'), deployment='encoder', head=False, polling_type=PollingType.ANY, shard_id=0, ) assert len(results) == 1 assert send_mock.call_count == 10 # encoder pod has no head registered yet so sending to the head will not work results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='encoder', head=True) assert len(results) == 0 assert send_mock.call_count == 10 # after registering a head for encoder, sending to head should work pool.add_connection(deployment='encoder', head=True, address='1.1.1.10:53') results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='encoder', head=True) assert len(results) == 1 assert send_mock.call_count == 11 # after remove the head again, sending will not work assert await pool.remove_connection(deployment='encoder', head=True, address='1.1.1.10:53') assert close_mock_object.call_count == 2 results = pool.send_request(request=ControlRequest(command='STATUS'), deployment='encoder', head=True) assert len(results) == 0 assert send_mock.call_count == 11 # check that remove/add order is handled well pool.add_connection(deployment='encoder', head=False, address='1.1.1.4:53') assert await pool.remove_connection(deployment='encoder', head=False, address='1.1.1.1:53') assert await pool.remove_connection(deployment='encoder', head=False, address='1.1.1.4:53') assert close_mock_object.call_count == 4 assert not (await pool.remove_connection( deployment='encoder', head=False, address='1.1.1.2:53')) await pool.close()
async def test_grpc_connection_pool_real_sending(): server1_ready_event = multiprocessing.Event() server2_ready_event = multiprocessing.Event() def listen(port, event: multiprocessing.Event): class DummyServer: async def process_control(self, request, *args): returned_msg = ControlRequest(command='DEACTIVATE') return returned_msg async def start_grpc_server(): grpc_server = grpc.aio.server( options=[ ('grpc.max_send_request_length', -1), ('grpc.max_receive_message_length', -1), ] ) jina_pb2_grpc.add_JinaControlRequestRPCServicer_to_server( DummyServer(), grpc_server ) grpc_server.add_insecure_port(f'localhost:{port}') await grpc_server.start() event.set() await grpc_server.wait_for_termination() asyncio.run(start_grpc_server()) port1 = random_port() server_process1 = Process( target=listen, args=( port1, server1_ready_event, ), ) server_process1.start() port2 = random_port() server_process2 = Process( target=listen, args=( port2, server2_ready_event, ), ) server_process2.start() time.sleep(0.1) server1_ready_event.wait() server2_ready_event.wait() pool = GrpcConnectionPool() pool.add_connection(deployment='encoder', head=False, address=f'localhost:{port1}') pool.add_connection(deployment='encoder', head=False, address=f'localhost:{port2}') sent_msg = ControlRequest(command='STATUS') results_call_1 = pool.send_request( request=sent_msg, deployment='encoder', head=False ) results_call_2 = pool.send_request( request=sent_msg, deployment='encoder', head=False ) assert len(results_call_1) == 1 assert len(results_call_2) == 1 response1, meta = await results_call_1[0] assert response1.command == 'DEACTIVATE' response2, meta = await results_call_2[0] assert response2.command == 'DEACTIVATE' await pool.close() server_process1.kill() server_process2.kill() server_process1.join() server_process2.join()