def test_send_sync_rpc_response(self): expected = SyncRPCResponse(reqId=123, respBody=self._expected_resp) self._sync_rpc_client._response_queue.put(expected) res = self._sync_rpc_client.send_sync_rpc_response() actual = next(res) self.assertEqual(expected, actual) expected = SyncRPCResponse(heartBeat=True) actual = next(res) self.assertEqual(expected, actual)
async def _await_gateway_response(self, client, stream_id, body, req_id, response_queue, conn_closed_table): await client.send_data(stream_id, body, end_stream=True) resp_headers = await client.recv_response(stream_id) status = self._get_resp_status(resp_headers) curr_payload = await self._read_stream(client, stream_id, req_id, response_queue, conn_closed_table) next_payload = await self._read_stream(client, stream_id, req_id, response_queue, conn_closed_table) while True: trailers = await client.recv_trailers(stream_id) \ if not next_payload else [] headers = self._get_resp_headers(resp_headers, trailers) res = GatewayResponse(status=status, headers=headers, payload=curr_payload) response_queue.put( SyncRPCResponse(heartBeat=False, reqId=req_id, respBody=res)) if not next_payload: break curr_payload = next_payload next_payload = await self._read_stream(client, stream_id, req_id, response_queue, conn_closed_table)
def test_send_request_done(self): """ Test if send_request_done puts the right SyncRPCResponse in response_queue based on the result of future. Returns: None """ self.assertTrue(self._sync_rpc_client._response_queue.empty()) req_id = 356 future = MockFuture(is_error=False, expected_result=self._expected_resp, expected_err_msg="") # future has result, send_request_done should enqueue a SyncRPCResponse # with the req_id, and the expected_resp set in MockFuture self._sync_rpc_client.send_request_done(req_id, future) self.assertFalse(self._sync_rpc_client._response_queue.empty()) res = self._sync_rpc_client._response_queue.get(block=False) self.assertEqual( res, SyncRPCResponse(reqId=req_id, respBody=self._expected_resp, heartBeat=False)) self.assertTrue(self._sync_rpc_client._response_queue.empty()) req_id = 234 expected_err_resp = GatewayResponse(err=self._expected_err_msg) future = MockFuture(is_error=True, expected_result=GatewayResponse(), expected_err_msg=self._expected_err_msg) self._sync_rpc_client.send_request_done(req_id, future) res = self._sync_rpc_client._response_queue.get(block=False) self.assertEqual( res, SyncRPCResponse(reqId=req_id, respBody=expected_err_resp, heartBeat=False))
def send_request_done(self, req_id, future): """ A future that has a GatewayResponse is done. Check if a exception is raised. If so, log the error and enqueue an empty SyncRPCResponse. Else, enqueue a SyncRPCResponse that contains the GatewayResponse that became available in the future. Args: req_id: request id that's associated with the response future: A future that contains a GatewayResponse that is done. Returns: None """ err = future.exception() if err: logging.error("[SyncRPC] Forward to control proxy error: %s", err) self._response_queue.put( SyncRPCResponse(heartBeat=False, reqId=req_id, respBody=GatewayResponse(err=str(err)))) else: res = future.result() self._response_queue.put( SyncRPCResponse(heartBeat=False, reqId=req_id, respBody=res))
async def try_read_stream(): while True: try: payload = await asyncio.wait_for( client.read_stream(stream_id), timeout=10.0) if conn_closed_table.get(req_id, False): raise ConnectionAbortedError return payload except asyncio.TimeoutError: if conn_closed_table.get(req_id, False): raise ConnectionAbortedError response_queue.put( SyncRPCResponse( heartBeat=False, reqId=req_id, respBody=GatewayResponse(keepConnActive=True)))
def send_sync_rpc_response(self): """ Retrieve SyncRPCResponse from queue. If no response is available yet, block for at most response_timeout seconds, and send a heartBeat if timeout. Returns: A generator of SyncRPCResponse """ while True: try: resp = self._response_queue.get(block=True, timeout=self._response_timeout) yield resp except queue.Empty: # response_queue is empty, send heartbeat # as the function itself has no knowledge on when it's # the first time it's called # this heartbeat response could be periodically called logging.debug("[SyncRPC] Sending heartbeat") yield SyncRPCResponse(heartBeat=True)
async def send(self, gateway_request, req_id, sync_rpc_response_queue, conn_closed_table): """ Forwards the given request to the service provided in :authority and awaits a response. If a exception is raised, log the error and enqueue an empty SyncRPCResponse. Else, enqueue SyncRPCResponse(s) that contains the GatewayResponse. Args: gateway_request: gateway_request: A GatewayRequest that is defined in the sync_rpc_service.proto. It has fields gwId, authority, path, headers, and payload. req_id: request id that's associated with the response sync_rpc_response_queue: the response queue that responses will be put in conn_closed_table: table that maps req ids to if the conn is closed Returns: None. """ client = await self._get_client(gateway_request.authority) # Small hack to set PingReceived to no-op because the log gets spammed # with KeyError messages since aioh2 doesn't have a handler for # PingReceived. Remove if future versions support it. # pylint: disable=protected-access if hasattr(h2.events, "PingReceived"): # Need the hasattr here because some older versions of h2 may not # have the PingReceived event client._event_handlers[h2.events.PingReceived] = lambda _: None # pylint: enable=protected-access if req_id in self._connection_table: logging.error( "[SyncRPC] proxy_client is already handling " "request ID %s", req_id) sync_rpc_response_queue.put( SyncRPCResponse( heartBeat=False, reqId=req_id, respBody=GatewayResponse( err=str("request ID {} is already being handled". format(req_id))))) client.close_connection() return self._connection_table[req_id] = client try: await client.wait_functional() req_headers = self._get_req_headers(gateway_request.headers, gateway_request.path, gateway_request.authority) body = gateway_request.payload stream_id = await client.start_request(req_headers) await self._await_gateway_response(client, stream_id, body, req_id, sync_rpc_response_queue, conn_closed_table) except ConnectionAbortedError: logging.error("[SyncRPC] proxy_client connection " "terminated by cloud") except Exception as e: # pylint: disable=broad-except logging.error("[SyncRPC] Exception in proxy_client: %s", e) sync_rpc_response_queue.put( SyncRPCResponse(heartBeat=False, reqId=req_id, respBody=GatewayResponse(err=str(e)))) finally: del self._connection_table[req_id] client.close_connection()