pw_rpc/py/pw_rpc/callback_client/call.py - pigweed/pigweed - Git at Google

 # Copyright 2021 The Pigweed Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
 # use this file except in compliance with the License. You may obtain a copy of
 # the License at
 #
 #     https://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations under
 # the License.
 """Classes for handling ongoing RPC calls."""

 import enum
 import logging
 import math
 import queue
 from typing import (Any, Callable, Iterable, Iterator, NamedTuple, Union,
                     Optional, Sequence, TypeVar)

 from pw_protobuf_compiler.python_protos import proto_repr
 from pw_status import Status
 from google.protobuf.message import Message

 from pw_rpc.callback_client.errors import RpcTimeout, RpcError
 from pw_rpc.client import PendingRpc, PendingRpcs
 from pw_rpc.descriptors import Method

 _LOG = logging.getLogger(__package__)


 class UseDefault(enum.Enum):
     """Marker for args that should use a default value, when None is valid."""
     VALUE = 0


 CallType = TypeVar('CallType', 'UnaryCall', 'ServerStreamingCall',
                    'ClientStreamingCall', 'BidirectionalStreamingCall')

 OnNextCallback = Callable[[CallType, Any], Any]
 OnCompletedCallback = Callable[[CallType, Any], Any]
 OnErrorCallback = Callable[[CallType, Any], Any]

 OptionalTimeout = Union[UseDefault, float, None]


 class UnaryResponse(NamedTuple):
     """Result from a unary or client streaming RPC: status and response."""
     status: Status
     response: Any

     def __repr__(self) -> str:
         return f'({self.status}, {proto_repr(self.response)})'


 class StreamResponse(NamedTuple):
     """Results from a server or bidirectional streaming RPC."""
     status: Status
     responses: Sequence[Any]

     def __repr__(self) -> str:
         return (f'({self.status}, '
                 f'[{", ".join(proto_repr(r) for r in self.responses)}])')


 class Call:
     """Represents an in-progress or completed RPC call."""
     def __init__(self, rpcs: PendingRpcs, rpc: PendingRpc,
                  default_timeout_s: Optional[float],
                  on_next: Optional[OnNextCallback],
                  on_completed: Optional[OnCompletedCallback],
                  on_error: Optional[OnErrorCallback]) -> None:
         self._rpcs = rpcs
         self._rpc = rpc
         self.default_timeout_s = default_timeout_s

         self.status: Optional[Status] = None
         self.error: Optional[Status] = None
         self._callback_exception: Optional[Exception] = None
         self._responses: list = []
         self._response_queue: queue.SimpleQueue = queue.SimpleQueue()

         self.on_next = on_next or Call._default_response
         self.on_completed = on_completed or Call._default_completion
         self.on_error = on_error or Call._default_error

     def _invoke(self, request: Optional[Message], ignore_errors: bool) -> None:
         """Calls the RPC. This must be called immediately after __init__."""
         previous = self._rpcs.send_request(self._rpc,
                                            request,
                                            self,
                                            ignore_errors=ignore_errors,
                                            override_pending=True)

         # TODO(hepler): Remove the cancel_duplicate_calls option.
         if (self._rpcs.cancel_duplicate_calls and  # type: ignore[attr-defined]
                 previous is not None and not previous.completed()):
             previous._handle_error(Status.CANCELLED)  # pylint: disable=protected-access

     def _default_response(self, response: Message) -> None:
         _LOG.debug('%s received response: %s', self._rpc, response)

     def _default_completion(self, status: Status) -> None:
         _LOG.info('%s completed: %s', self._rpc, status)

     def _default_error(self, error: Status) -> None:
         _LOG.warning('%s terminated due to an error: %s', self._rpc, error)

     @property
     def method(self) -> Method:
         return self._rpc.method

     def completed(self) -> bool:
         """True if the RPC call has completed, successfully or from an error."""
         return self.status is not None or self.error is not None

     def _send_client_stream(self, request_proto: Optional[Message],
                             request_fields: dict) -> None:
         """Sends a client to the server in the client stream.

         Sending a client stream packet on a closed RPC raises an exception.
         """
         self._check_errors()

         if self.status is not None:
             raise RpcError(self._rpc, Status.FAILED_PRECONDITION)

         self._rpcs.send_client_stream(
             self._rpc, self.method.get_request(request_proto, request_fields))

     def _finish_client_stream(self, requests: Iterable[Message]) -> None:
         for request in requests:
             self._send_client_stream(request, {})

         if not self.completed():
             self._rpcs.send_client_stream_end(self._rpc)

     def _unary_wait(self, timeout_s: OptionalTimeout) -> UnaryResponse:
         """Waits until the RPC has completed."""
         for _ in self._get_responses(timeout_s=timeout_s):
             pass

         assert self.status is not None and self._responses
         return UnaryResponse(self.status, self._responses[-1])

     def _stream_wait(self, timeout_s: OptionalTimeout) -> StreamResponse:
         """Waits until the RPC has completed."""
         for _ in self._get_responses(timeout_s=timeout_s):
             pass

         assert self.status is not None
         return StreamResponse(self.status, self._responses)

     def _get_responses(self,
                        *,
                        count: int = None,
                        timeout_s: OptionalTimeout) -> Iterator:
         """Returns an iterator of stream responses.

         Args:
           count: Responses to read before returning; None reads all
           timeout_s: max time in seconds to wait between responses; 0 doesn't
               block, None blocks indefinitely
         """
         self._check_errors()

         if self.completed() and self._response_queue.empty():
             return

         if timeout_s is UseDefault.VALUE:
             timeout_s = self.default_timeout_s

         remaining = math.inf if count is None else count

         try:
             while remaining:
                 response = self._response_queue.get(True, timeout_s)

                 self._check_errors()

                 if response is None:
                     return

                 yield response
                 remaining -= 1
         except queue.Empty:
             raise RpcTimeout(self._rpc, timeout_s)

     def cancel(self) -> bool:
         """Cancels the RPC; returns whether the RPC was active."""
         if self.completed():
             return False

         self.error = Status.CANCELLED
         return self._rpcs.send_cancel(self._rpc)

     def _check_errors(self) -> None:
         if self._callback_exception:
             raise self._callback_exception

         if self.error:
             raise RpcError(self._rpc, self.error)

     def _handle_response(self, response: Any) -> None:
         # TODO(frolv): These lists could grow very large for persistent
         # streaming RPCs such as logs. The size should be limited.
         self._responses.append(response)
         self._response_queue.put(response)

         self._invoke_callback('on_next', response)

     def _handle_completion(self, status: Status) -> None:
         self.status = status
         self._response_queue.put(None)

         self._invoke_callback('on_completed', status)

     def _handle_error(self, error: Status) -> None:
         self.error = error
         self._response_queue.put(None)

         self._invoke_callback('on_error', error)

     def _invoke_callback(self, callback_name: str, arg: Any) -> None:
         """Invokes a user-provided callback function for an RPC event."""

         # Catch and log any exceptions from the user-provided callback so that
         # exceptions don't terminate the thread handling RPC packets.
         callback: Callable[[Call, Any], None] = getattr(self, callback_name)

         try:
             callback(self, arg)
         except Exception as callback_exception:  # pylint: disable=broad-except
             msg = (f'The {callback_name} callback ({callback}) for '
                    f'{self._rpc} raised an exception')
             _LOG.exception(msg)

             self._callback_exception = RuntimeError(msg)
             self._callback_exception.__cause__ = callback_exception

     def __enter__(self) -> 'Call':
         return self

     def __exit__(self, exc_type, exc_value, traceback) -> None:
         self.cancel()

     def __repr__(self) -> str:
         return f'{type(self).__name__}({self.method})'


 class UnaryCall(Call):
     """Tracks the state of a unary RPC call."""
     @property
     def response(self) -> Any:
         return self._responses[-1] if self._responses else None

     def wait(self,
              timeout_s: OptionalTimeout = UseDefault.VALUE) -> UnaryResponse:
         return self._unary_wait(timeout_s)


 class ServerStreamingCall(Call):
     """Tracks the state of a server streaming RPC call."""
     @property
     def responses(self) -> Sequence:
         return self._responses

     def wait(self,
              timeout_s: OptionalTimeout = UseDefault.VALUE) -> StreamResponse:
         return self._stream_wait(timeout_s)

     def get_responses(
             self,
             *,
             count: int = None,
             timeout_s: OptionalTimeout = UseDefault.VALUE) -> Iterator:
         return self._get_responses(count=count, timeout_s=timeout_s)

     def __iter__(self) -> Iterator:
         return self.get_responses()


 class ClientStreamingCall(Call):
     """Tracks the state of a client streaming RPC call."""
     @property
     def response(self) -> Any:
         return self._responses[-1] if self._responses else None

     # TODO(hepler): Use / to mark the first arg as positional-only
     #     when when Python 3.7 support is no longer required.
     def send(self,
              _rpc_request_proto: Message = None,
              **request_fields) -> None:
         """Sends client stream request to the server."""
         self._send_client_stream(_rpc_request_proto, request_fields)

     def finish_and_wait(
             self,
             requests: Iterable[Message] = (),
             *,
             timeout_s: OptionalTimeout = UseDefault.VALUE) -> UnaryResponse:
         """Ends the client stream and waits for the RPC to complete."""
         self._finish_client_stream(requests)
         return self._unary_wait(timeout_s)


 class BidirectionalStreamingCall(Call):
     """Tracks the state of a bidirectional streaming RPC call."""
     @property
     def responses(self) -> Sequence:
         return self._responses

     # TODO(hepler): Use / to mark the first arg as positional-only
     #     when when Python 3.7 support is no longer required.
     def send(self,
              _rpc_request_proto: Message = None,
              **request_fields) -> None:
         """Sends a message to the server in the client stream."""
         self._send_client_stream(_rpc_request_proto, request_fields)

     def finish_and_wait(
             self,
             requests: Iterable[Message] = (),
             *,
             timeout_s: OptionalTimeout = UseDefault.VALUE) -> StreamResponse:
         """Ends the client stream and waits for the RPC to complete."""
         self._finish_client_stream(requests)
         return self._stream_wait(timeout_s)

     def get_responses(
             self,
             *,
             count: int = None,
             timeout_s: OptionalTimeout = UseDefault.VALUE) -> Iterator:
         return self._get_responses(count=count, timeout_s=timeout_s)

     def __iter__(self) -> Iterator:
         return self.get_responses()
	# Copyright 2021 The Pigweed Authors
	#
	# Licensed under the Apache License, Version 2.0 (the "License"); you may not
	# use this file except in compliance with the License. You may obtain a copy of
	# the License at
	#
	# https://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	# License for the specific language governing permissions and limitations under
	# the License.
	"""Classes for handling ongoing RPC calls."""

	import enum
	import logging
	import math
	import queue
	from typing import (Any, Callable, Iterable, Iterator, NamedTuple, Union,
	Optional, Sequence, TypeVar)

	from pw_protobuf_compiler.python_protos import proto_repr
	from pw_status import Status
	from google.protobuf.message import Message

	from pw_rpc.callback_client.errors import RpcTimeout, RpcError
	from pw_rpc.client import PendingRpc, PendingRpcs
	from pw_rpc.descriptors import Method

	_LOG = logging.getLogger(__package__)


	class UseDefault(enum.Enum):
	"""Marker for args that should use a default value, when None is valid."""
	VALUE = 0


	CallType = TypeVar('CallType', 'UnaryCall', 'ServerStreamingCall',
	'ClientStreamingCall', 'BidirectionalStreamingCall')

	OnNextCallback = Callable[[CallType, Any], Any]
	OnCompletedCallback = Callable[[CallType, Any], Any]
	OnErrorCallback = Callable[[CallType, Any], Any]

	OptionalTimeout = Union[UseDefault, float, None]


	class UnaryResponse(NamedTuple):
	"""Result from a unary or client streaming RPC: status and response."""
	status: Status
	response: Any

	def __repr__(self) -> str:
	return f'({self.status}, {proto_repr(self.response)})'


	class StreamResponse(NamedTuple):
	"""Results from a server or bidirectional streaming RPC."""
	status: Status
	responses: Sequence[Any]

	def __repr__(self) -> str:
	return (f'({self.status}, '
	f'[{", ".join(proto_repr(r) for r in self.responses)}])')


	class Call:
	"""Represents an in-progress or completed RPC call."""
	def __init__(self, rpcs: PendingRpcs, rpc: PendingRpc,
	default_timeout_s: Optional[float],
	on_next: Optional[OnNextCallback],
	on_completed: Optional[OnCompletedCallback],
	on_error: Optional[OnErrorCallback]) -> None:
	self._rpcs = rpcs
	self._rpc = rpc
	self.default_timeout_s = default_timeout_s

	self.status: Optional[Status] = None
	self.error: Optional[Status] = None
	self._callback_exception: Optional[Exception] = None
	self._responses: list = []
	self._response_queue: queue.SimpleQueue = queue.SimpleQueue()

	self.on_next = on_next or Call._default_response
	self.on_completed = on_completed or Call._default_completion
	self.on_error = on_error or Call._default_error

	def _invoke(self, request: Optional[Message], ignore_errors: bool) -> None:
	"""Calls the RPC. This must be called immediately after __init__."""
	previous = self._rpcs.send_request(self._rpc,
	request,
	self,
	ignore_errors=ignore_errors,
	override_pending=True)

	# TODO(hepler): Remove the cancel_duplicate_calls option.
	if (self._rpcs.cancel_duplicate_calls and # type: ignore[attr-defined]
	previous is not None and not previous.completed()):
	previous._handle_error(Status.CANCELLED) # pylint: disable=protected-access

	def _default_response(self, response: Message) -> None:
	_LOG.debug('%s received response: %s', self._rpc, response)

	def _default_completion(self, status: Status) -> None:
	_LOG.info('%s completed: %s', self._rpc, status)

	def _default_error(self, error: Status) -> None:
	_LOG.warning('%s terminated due to an error: %s', self._rpc, error)

	@property
	def method(self) -> Method:
	return self._rpc.method

	def completed(self) -> bool:
	"""True if the RPC call has completed, successfully or from an error."""
	return self.status is not None or self.error is not None

	def _send_client_stream(self, request_proto: Optional[Message],
	request_fields: dict) -> None:
	"""Sends a client to the server in the client stream.

	Sending a client stream packet on a closed RPC raises an exception.
	"""
	self._check_errors()

	if self.status is not None:
	raise RpcError(self._rpc, Status.FAILED_PRECONDITION)

	self._rpcs.send_client_stream(
	self._rpc, self.method.get_request(request_proto, request_fields))

	def _finish_client_stream(self, requests: Iterable[Message]) -> None:
	for request in requests:
	self._send_client_stream(request, {})

	if not self.completed():
	self._rpcs.send_client_stream_end(self._rpc)

	def _unary_wait(self, timeout_s: OptionalTimeout) -> UnaryResponse:
	"""Waits until the RPC has completed."""
	for _ in self._get_responses(timeout_s=timeout_s):
	pass

	assert self.status is not None and self._responses
	return UnaryResponse(self.status, self._responses[-1])

	def _stream_wait(self, timeout_s: OptionalTimeout) -> StreamResponse:
	"""Waits until the RPC has completed."""
	for _ in self._get_responses(timeout_s=timeout_s):
	pass

	assert self.status is not None
	return StreamResponse(self.status, self._responses)

	def _get_responses(self,
	*,
	count: int = None,
	timeout_s: OptionalTimeout) -> Iterator:
	"""Returns an iterator of stream responses.

	Args:
	count: Responses to read before returning; None reads all
	timeout_s: max time in seconds to wait between responses; 0 doesn't
	block, None blocks indefinitely
	"""
	self._check_errors()

	if self.completed() and self._response_queue.empty():
	return

	if timeout_s is UseDefault.VALUE:
	timeout_s = self.default_timeout_s

	remaining = math.inf if count is None else count

	try:
	while remaining:
	response = self._response_queue.get(True, timeout_s)

	self._check_errors()

	if response is None:
	return

	yield response
	remaining -= 1
	except queue.Empty:
	raise RpcTimeout(self._rpc, timeout_s)

	def cancel(self) -> bool:
	"""Cancels the RPC; returns whether the RPC was active."""
	if self.completed():
	return False

	self.error = Status.CANCELLED
	return self._rpcs.send_cancel(self._rpc)

	def _check_errors(self) -> None:
	if self._callback_exception:
	raise self._callback_exception

	if self.error:
	raise RpcError(self._rpc, self.error)

	def _handle_response(self, response: Any) -> None:
	# TODO(frolv): These lists could grow very large for persistent
	# streaming RPCs such as logs. The size should be limited.
	self._responses.append(response)
	self._response_queue.put(response)

	self._invoke_callback('on_next', response)

	def _handle_completion(self, status: Status) -> None:
	self.status = status
	self._response_queue.put(None)

	self._invoke_callback('on_completed', status)

	def _handle_error(self, error: Status) -> None:
	self.error = error
	self._response_queue.put(None)

	self._invoke_callback('on_error', error)

	def _invoke_callback(self, callback_name: str, arg: Any) -> None:
	"""Invokes a user-provided callback function for an RPC event."""

	# Catch and log any exceptions from the user-provided callback so that
	# exceptions don't terminate the thread handling RPC packets.
	callback: Callable[[Call, Any], None] = getattr(self, callback_name)

	try:
	callback(self, arg)
	except Exception as callback_exception: # pylint: disable=broad-except
	msg = (f'The {callback_name} callback ({callback}) for '
	f'{self._rpc} raised an exception')
	_LOG.exception(msg)

	self._callback_exception = RuntimeError(msg)
	self._callback_exception.__cause__ = callback_exception

	def __enter__(self) -> 'Call':
	return self

	def __exit__(self, exc_type, exc_value, traceback) -> None:
	self.cancel()

	def __repr__(self) -> str:
	return f'{type(self).__name__}({self.method})'


	class UnaryCall(Call):
	"""Tracks the state of a unary RPC call."""
	@property
	def response(self) -> Any:
	return self._responses[-1] if self._responses else None

	def wait(self,
	timeout_s: OptionalTimeout = UseDefault.VALUE) -> UnaryResponse:
	return self._unary_wait(timeout_s)


	class ServerStreamingCall(Call):
	"""Tracks the state of a server streaming RPC call."""
	@property
	def responses(self) -> Sequence:
	return self._responses

	def wait(self,
	timeout_s: OptionalTimeout = UseDefault.VALUE) -> StreamResponse:
	return self._stream_wait(timeout_s)

	def get_responses(
	self,
	*,
	count: int = None,
	timeout_s: OptionalTimeout = UseDefault.VALUE) -> Iterator:
	return self._get_responses(count=count, timeout_s=timeout_s)

	def __iter__(self) -> Iterator:
	return self.get_responses()


	class ClientStreamingCall(Call):
	"""Tracks the state of a client streaming RPC call."""
	@property
	def response(self) -> Any:
	return self._responses[-1] if self._responses else None

	# TODO(hepler): Use / to mark the first arg as positional-only
	# when when Python 3.7 support is no longer required.
	def send(self,
	_rpc_request_proto: Message = None,
	**request_fields) -> None:
	"""Sends client stream request to the server."""
	self._send_client_stream(_rpc_request_proto, request_fields)

	def finish_and_wait(
	self,
	requests: Iterable[Message] = (),
	*,
	timeout_s: OptionalTimeout = UseDefault.VALUE) -> UnaryResponse:
	"""Ends the client stream and waits for the RPC to complete."""
	self._finish_client_stream(requests)
	return self._unary_wait(timeout_s)


	class BidirectionalStreamingCall(Call):
	"""Tracks the state of a bidirectional streaming RPC call."""
	@property
	def responses(self) -> Sequence:
	return self._responses

	# TODO(hepler): Use / to mark the first arg as positional-only
	# when when Python 3.7 support is no longer required.
	def send(self,
	_rpc_request_proto: Message = None,
	**request_fields) -> None:
	"""Sends a message to the server in the client stream."""
	self._send_client_stream(_rpc_request_proto, request_fields)

	def finish_and_wait(
	self,
	requests: Iterable[Message] = (),
	*,
	timeout_s: OptionalTimeout = UseDefault.VALUE) -> StreamResponse:
	"""Ends the client stream and waits for the RPC to complete."""
	self._finish_client_stream(requests)
	return self._stream_wait(timeout_s)

	def get_responses(
	self,
	*,
	count: int = None,
	timeout_s: OptionalTimeout = UseDefault.VALUE) -> Iterator:
	return self._get_responses(count=count, timeout_s=timeout_s)

	def __iter__(self) -> Iterator:
	return self.get_responses()