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()