def handle_request(self, request_stream, response_stream):
try:
method_path = request_stream.headers.get(":path")
entrypoint = self.entrypoints[method_path]
except KeyError:
raise GrpcError(
status=StatusCode.UNIMPLEMENTED, details="Method not found!"
)
encoding = request_stream.headers.get("grpc-encoding", "identity")
if encoding not in SUPPORTED_ENCODINGS:
raise GrpcError(
status=StatusCode.UNIMPLEMENTED,
details="Algorithm not supported: {}".format(encoding),
)
timeout = request_stream.headers.get("grpc-timeout")
if timeout:
timeout = unbucket_timeout(timeout)
self.container.spawn_managed_thread(
partial(self.timeout, request_stream, response_stream, timeout)
)
self.container.spawn_managed_thread(
partial(entrypoint.handle_request, request_stream, response_stream)
)
def handle_request(self, request_stream, response_stream):
request = request_stream.consume(self.input_type)
if self.cardinality in (Cardinality.UNARY_STREAM, Cardinality.UNARY_UNARY):
request = next(request)
context = GrpcContext(request_stream, response_stream)
args = (request, context)
kwargs = {}
context_data = context_data_from_metadata(context.invocation_metadata())
handle_result = partial(self.handle_result, response_stream)
try:
self.container.spawn_worker(
self,
args,
kwargs,
context_data=context_data,
handle_result=handle_result,
)
except ContainerBeingKilled:
raise GrpcError(
status=StatusCode.UNAVAILABLE, details="Server shutting down"
)
def test_close_with_error(self):
stream = StreamBase(1)
error = GrpcError("boom", "details", "error string")
stream.close(error)
assert stream.closed
assert stream.queue.get() == error
def cleanup_on_exit(self):
error = None
try:
yield
except Exception:
log.info(
"ConnectionManager shutting down with error. Traceback:",
exc_info=True,
)
error = GrpcError.from_exception(sys.exc_info(),
code=StatusCode.UNAVAILABLE)
finally:
for send_stream in self.send_streams.values():
if send_stream.closed:
continue # stream.close() is idemponent but this prevents the log
log.info(f"Terminating send stream {send_stream}"
f"{f' with error {error}' if error else ''}.")
send_stream.close(error)
for receive_stream in self.receive_streams.values():
if receive_stream.closed:
continue # stream.close() is idemponent but this prevents the log
log.info(f"Terminating receive stream {receive_stream}"
f"{f' with error {error}' if error else ''}.")
receive_stream.close(error)
self.sock.close()
self.stopped.set()
def test_consume_grpc_error(self):
stream = ReceiveStream(1)
error = GrpcError("boom", "details", "message")
stream.queue.put(error)
message_type = Mock()
with pytest.raises(GrpcError):
next(stream.consume(message_type))
def test_error_on_queue(self, generate_messages):
stream = SendStream(1)
stream.populate(generate_messages(count=2, length=20))
error = GrpcError("boom", "details", "error string")
stream.close(error)
with pytest.raises(GrpcError):
stream.flush_queue_to_buffer()
def timeout(self, send_stream, response_stream, deadline):
start = time.time()
while True:
elapsed = time.time() - start
if elapsed > deadline:
error = GrpcError(status=StatusCode.DEADLINE_EXCEEDED,
details="Deadline Exceeded")
response_stream.close(error)
send_stream.close()
time.sleep(0.001)
def send_stream(self, result):
try:
for item in result:
self.send(item)
except grpc.RpcError as exc:
state = exc._state
error = GrpcError(state.code, state.details,
status_from_metadata(state.trailing_metadata))
self.send(error)
self.send(self.ENDSTREAM, close=True)
def unary_grpc_error(self, request, context):
maybe_echo_metadata(context)
maybe_sleep(request)
code = StatusCode.UNAUTHENTICATED
message = "Not allowed!"
raise GrpcError(code=code,
message=message,
status=make_status(code, message))
def send_stream(self, result):
try:
for item in result:
self.send(item)
except grpc.RpcError as exc:
state = exc._state
error = GrpcError(state.code, state.details,
state.debug_error_string)
self.send(error)
self.send(self.ENDSTREAM, close=True)
def test_error_on_queue(self, generate_messages):
stream = SendStream(1)
error = GrpcError("boom", "details")
messages = itertools.chain(generate_messages(count=2, length=20),
[error])
stream.populate(messages)
with pytest.raises(GrpcError):
stream.flush_queue_to_buffer()
def test_stream_closed_with_error(self):
stream = SendStream(1)
error = GrpcError("boom", "details", "error string")
stream.close(error)
max_bytes = 10
chunk_size = 5
with pytest.raises(GrpcError):
next(stream.read(max_bytes, chunk_size))
def handle_result(self, response_stream, worker_ctx, result, exc_info):
if self.cardinality in (Cardinality.STREAM_UNARY, Cardinality.UNARY_UNARY):
result = (result,)
if exc_info is None:
try:
response_stream.populate(result)
except Exception as exception:
message = "Exception iterating responses: {}".format(exception)
error = GrpcError(status=StatusCode.UNKNOWN, details=message)
response_stream.close(error)
else:
error = GrpcError(
status=StatusCode.UNKNOWN,
details="Exception calling application: {}".format(exc_info[1]),
)
response_stream.close(error)
return result, exc_info
def handler(exc_info, code=None, message=None):
exc_type, exc, tb = exc_info
code = code or StatusCode.PERMISSION_DENIED
message = "Not allowed!"
status = Status(
code=STATUS_CODE_ENUM_TO_INT_MAP[code],
message=message,
details=[], # don't include traceback
)
return GrpcError(code=code, message=message, status=status)
def timeout(self, request_stream, response_stream, deadline):
start = time.time()
while True:
if request_stream.closed and response_stream.closed:
break
elapsed = time.time() - start
if elapsed > deadline:
request_stream.close()
error = GrpcError(code=StatusCode.DEADLINE_EXCEEDED,
message="Deadline Exceeded")
response_stream.close(error)
break
time.sleep(0.001)
def handle_result(self, response_stream, worker_ctx, result, exc_info):
if self.cardinality in (Cardinality.STREAM_UNARY,
Cardinality.UNARY_UNARY):
result = (result, )
if exc_info is None:
try:
response_stream.populate(result)
except Exception as exception:
message = "Exception iterating responses: {}".format(exception)
error = GrpcError.from_exception(sys.exc_info(),
message=message)
response_stream.close(error)
else:
message = "Exception calling application: {}".format(exc_info[1])
error = GrpcError.from_exception(exc_info, message=message)
response_stream.close(error)
return result, exc_info
def timeout(self, request_stream, response_stream, deadline):
start = time.time()
while True:
elapsed = time.time() - start
if elapsed > deadline:
request_stream.close()
# XXX does server actually need to do this according to the spec?
# perhaps we could just close the stream.
error = GrpcError(
status=StatusCode.DEADLINE_EXCEEDED, details="Deadline Exceeded"
)
response_stream.close(error)
break
time.sleep(0.001)
def send_data(self, stream_id):
try:
super().send_data(stream_id)
except UnsupportedEncoding:
response_stream = self.receive_streams[stream_id]
request_stream = self.send_streams[stream_id]
error = GrpcError(
status=StatusCode.UNIMPLEMENTED,
details="Algorithm not supported: {}".format(
request_stream.encoding),
)
response_stream.close(error)
request_stream.close()
def stream_grpc_error(self, request, context):
metadata = extract_metadata(context)
maybe_echo_metadata(context)
message = request.value * (request.multiplier or 1)
for i in range(request.response_count):
maybe_sleep(request)
# raise on the last message
if i == request.response_count - 1:
code = StatusCode.RESOURCE_EXHAUSTED
message = "Out of tokens!"
raise GrpcError(
code=code,
message=message,
status=make_status(code, message),
)
yield ExampleReply(message=message, seqno=i + 1, metadata=metadata)
def trailers_received(self, event):
""" Called when trailers are received on a stream.
If the trailers contain an error, we should raise it here.
"""
super().trailers_received(event)
stream_id = event.stream_id
response_stream = self.receive_streams.get(stream_id)
if response_stream is None:
self.conn.reset_stream(stream_id,
error_code=ErrorCodes.PROTOCOL_ERROR)
return
trailers = response_stream.trailers
if int(trailers.get("grpc-status", 0)) > 0:
error = GrpcError.from_headers(trailers)
response_stream.close(error)
def execute(command, stub):
method = getattr(stub, command.method_name)
request = command.get_request()
compression = command.kwargs.pop("compression", None)
if compression:
command.kwargs["metadata"] = list(command.kwargs.get(
"metadata", [])) + [("grpc-internal-encoding-request", compression)
]
response_metadata = {}
try:
if command.cardinality in (Cardinality.STREAM_UNARY,
Cardinality.UNARY_UNARY):
response_future = method.future(request, **command.kwargs)
response_metadata["code"] = response_future.code()
response_metadata["details"] = response_future.details()
response_metadata["initial_metadata"] = list(
map(tuple, response_future.initial_metadata()))
response_metadata["trailing_metadata"] = list(
map(tuple, response_future.trailing_metadata()))
response = response_future.result()
else:
# .future() interface for RPCs with STREAM responses not supported
response = method(request, **command.kwargs)
except grpc.RpcError as exc:
state = exc._state
response_metadata["code"] = state.code
response_metadata["details"] = state.details
response = GrpcError(state.code, state.details,
status_from_metadata(state.trailing_metadata))
command.send_response(response)
command.send_metadata(response_metadata)