1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141 |
- # Copyright 2016 gRPC 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
- #
- # http://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.
- """Service-side implementation of gRPC Python."""
- from __future__ import annotations
- import collections
- from concurrent import futures
- import enum
- import logging
- import threading
- import time
- import traceback
- from typing import (Any, Callable, Iterable, Iterator, List, Mapping, Optional,
- Sequence, Set, Tuple, Union)
- import grpc # pytype: disable=pyi-error
- from grpc import _common # pytype: disable=pyi-error
- from grpc import _compression # pytype: disable=pyi-error
- from grpc import _interceptor # pytype: disable=pyi-error
- from grpc._cython import cygrpc
- from grpc._typing import ArityAgnosticMethodHandler
- from grpc._typing import ChannelArgumentType
- from grpc._typing import DeserializingFunction
- from grpc._typing import MetadataType
- from grpc._typing import NullaryCallbackType
- from grpc._typing import ResponseType
- from grpc._typing import SerializingFunction
- from grpc._typing import ServerCallbackTag
- from grpc._typing import ServerTagCallbackType
- _LOGGER = logging.getLogger(__name__)
- _SHUTDOWN_TAG = 'shutdown'
- _REQUEST_CALL_TAG = 'request_call'
- _RECEIVE_CLOSE_ON_SERVER_TOKEN = 'receive_close_on_server'
- _SEND_INITIAL_METADATA_TOKEN = 'send_initial_metadata'
- _RECEIVE_MESSAGE_TOKEN = 'receive_message'
- _SEND_MESSAGE_TOKEN = 'send_message'
- _SEND_INITIAL_METADATA_AND_SEND_MESSAGE_TOKEN = (
- 'send_initial_metadata * send_message')
- _SEND_STATUS_FROM_SERVER_TOKEN = 'send_status_from_server'
- _SEND_INITIAL_METADATA_AND_SEND_STATUS_FROM_SERVER_TOKEN = (
- 'send_initial_metadata * send_status_from_server')
- _OPEN = 'open'
- _CLOSED = 'closed'
- _CANCELLED = 'cancelled'
- _EMPTY_FLAGS = 0
- _DEALLOCATED_SERVER_CHECK_PERIOD_S = 1.0
- _INF_TIMEOUT = 1e9
- def _serialized_request(request_event: cygrpc.BaseEvent) -> bytes:
- return request_event.batch_operations[0].message()
- def _application_code(code: grpc.StatusCode) -> cygrpc.StatusCode:
- cygrpc_code = _common.STATUS_CODE_TO_CYGRPC_STATUS_CODE.get(code)
- return cygrpc.StatusCode.unknown if cygrpc_code is None else cygrpc_code
- def _completion_code(state: _RPCState) -> cygrpc.StatusCode:
- if state.code is None:
- return cygrpc.StatusCode.ok
- else:
- return _application_code(state.code)
- def _abortion_code(state: _RPCState,
- code: cygrpc.StatusCode) -> cygrpc.StatusCode:
- if state.code is None:
- return code
- else:
- return _application_code(state.code)
- def _details(state: _RPCState) -> bytes:
- return b'' if state.details is None else state.details
- class _HandlerCallDetails(
- collections.namedtuple('_HandlerCallDetails', (
- 'method',
- 'invocation_metadata',
- )), grpc.HandlerCallDetails):
- pass
- class _RPCState(object):
- condition: threading.Condition
- due = Set[str]
- request: Any
- client: str
- initial_metadata_allowed: bool
- compression_algorithm: Optional[grpc.Compression]
- disable_next_compression: bool
- trailing_metadata: Optional[MetadataType]
- code: Optional[grpc.StatusCode]
- details: Optional[bytes]
- statused: bool
- rpc_errors: List[Exception]
- callbacks: Optional[List[NullaryCallbackType]]
- aborted: bool
- def __init__(self):
- self.condition = threading.Condition()
- self.due = set()
- self.request = None
- self.client = _OPEN
- self.initial_metadata_allowed = True
- self.compression_algorithm = None
- self.disable_next_compression = False
- self.trailing_metadata = None
- self.code = None
- self.details = None
- self.statused = False
- self.rpc_errors = []
- self.callbacks = []
- self.aborted = False
- def _raise_rpc_error(state: _RPCState) -> None:
- rpc_error = grpc.RpcError()
- state.rpc_errors.append(rpc_error)
- raise rpc_error
- def _possibly_finish_call(state: _RPCState,
- token: str) -> ServerTagCallbackType:
- state.due.remove(token)
- if not _is_rpc_state_active(state) and not state.due:
- callbacks = state.callbacks
- state.callbacks = None
- return state, callbacks
- else:
- return None, ()
- def _send_status_from_server(state: _RPCState, token: str) -> ServerCallbackTag:
- def send_status_from_server(unused_send_status_from_server_event):
- with state.condition:
- return _possibly_finish_call(state, token)
- return send_status_from_server
- def _get_initial_metadata(
- state: _RPCState,
- metadata: Optional[MetadataType]) -> Optional[MetadataType]:
- with state.condition:
- if state.compression_algorithm:
- compression_metadata = (
- _compression.compression_algorithm_to_metadata(
- state.compression_algorithm),)
- if metadata is None:
- return compression_metadata
- else:
- return compression_metadata + tuple(metadata)
- else:
- return metadata
- def _get_initial_metadata_operation(
- state: _RPCState, metadata: Optional[MetadataType]) -> cygrpc.Operation:
- operation = cygrpc.SendInitialMetadataOperation(
- _get_initial_metadata(state, metadata), _EMPTY_FLAGS)
- return operation
- def _abort(state: _RPCState, call: cygrpc.Call, code: cygrpc.StatusCode,
- details: bytes) -> None:
- if state.client is not _CANCELLED:
- effective_code = _abortion_code(state, code)
- effective_details = details if state.details is None else state.details
- if state.initial_metadata_allowed:
- operations = (
- _get_initial_metadata_operation(state, None),
- cygrpc.SendStatusFromServerOperation(state.trailing_metadata,
- effective_code,
- effective_details,
- _EMPTY_FLAGS),
- )
- token = _SEND_INITIAL_METADATA_AND_SEND_STATUS_FROM_SERVER_TOKEN
- else:
- operations = (cygrpc.SendStatusFromServerOperation(
- state.trailing_metadata, effective_code, effective_details,
- _EMPTY_FLAGS),)
- token = _SEND_STATUS_FROM_SERVER_TOKEN
- call.start_server_batch(operations,
- _send_status_from_server(state, token))
- state.statused = True
- state.due.add(token)
- def _receive_close_on_server(state: _RPCState) -> ServerCallbackTag:
- def receive_close_on_server(receive_close_on_server_event):
- with state.condition:
- if receive_close_on_server_event.batch_operations[0].cancelled():
- state.client = _CANCELLED
- elif state.client is _OPEN:
- state.client = _CLOSED
- state.condition.notify_all()
- return _possibly_finish_call(state, _RECEIVE_CLOSE_ON_SERVER_TOKEN)
- return receive_close_on_server
- def _receive_message(
- state: _RPCState, call: cygrpc.Call,
- request_deserializer: Optional[DeserializingFunction]
- ) -> ServerCallbackTag:
- def receive_message(receive_message_event):
- serialized_request = _serialized_request(receive_message_event)
- if serialized_request is None:
- with state.condition:
- if state.client is _OPEN:
- state.client = _CLOSED
- state.condition.notify_all()
- return _possibly_finish_call(state, _RECEIVE_MESSAGE_TOKEN)
- else:
- request = _common.deserialize(serialized_request,
- request_deserializer)
- with state.condition:
- if request is None:
- _abort(state, call, cygrpc.StatusCode.internal,
- b'Exception deserializing request!')
- else:
- state.request = request
- state.condition.notify_all()
- return _possibly_finish_call(state, _RECEIVE_MESSAGE_TOKEN)
- return receive_message
- def _send_initial_metadata(state: _RPCState) -> ServerCallbackTag:
- def send_initial_metadata(unused_send_initial_metadata_event):
- with state.condition:
- return _possibly_finish_call(state, _SEND_INITIAL_METADATA_TOKEN)
- return send_initial_metadata
- def _send_message(state: _RPCState, token: str) -> ServerCallbackTag:
- def send_message(unused_send_message_event):
- with state.condition:
- state.condition.notify_all()
- return _possibly_finish_call(state, token)
- return send_message
- class _Context(grpc.ServicerContext):
- _rpc_event: cygrpc.BaseEvent
- _state: _RPCState
- request_deserializer: Optional[DeserializingFunction]
- def __init__(self, rpc_event: cygrpc.BaseEvent, state: _RPCState,
- request_deserializer: Optional[DeserializingFunction]):
- self._rpc_event = rpc_event
- self._state = state
- self._request_deserializer = request_deserializer
- def is_active(self) -> bool:
- with self._state.condition:
- return _is_rpc_state_active(self._state)
- def time_remaining(self) -> float:
- return max(self._rpc_event.call_details.deadline - time.time(), 0)
- def cancel(self) -> None:
- self._rpc_event.call.cancel()
- def add_callback(self, callback: NullaryCallbackType) -> bool:
- with self._state.condition:
- if self._state.callbacks is None:
- return False
- else:
- self._state.callbacks.append(callback)
- return True
- def disable_next_message_compression(self) -> None:
- with self._state.condition:
- self._state.disable_next_compression = True
- def invocation_metadata(self) -> Optional[MetadataType]:
- return self._rpc_event.invocation_metadata
- def peer(self) -> str:
- return _common.decode(self._rpc_event.call.peer())
- def peer_identities(self) -> Optional[Sequence[bytes]]:
- return cygrpc.peer_identities(self._rpc_event.call)
- def peer_identity_key(self) -> Optional[str]:
- id_key = cygrpc.peer_identity_key(self._rpc_event.call)
- return id_key if id_key is None else _common.decode(id_key)
- def auth_context(self) -> Mapping[str, Sequence[bytes]]:
- auth_context = cygrpc.auth_context(self._rpc_event.call)
- auth_context_dict = {} if auth_context is None else auth_context
- return {
- _common.decode(key): value
- for key, value in auth_context_dict.items()
- }
- def set_compression(self, compression: grpc.Compression) -> None:
- with self._state.condition:
- self._state.compression_algorithm = compression
- def send_initial_metadata(self, initial_metadata: MetadataType) -> None:
- with self._state.condition:
- if self._state.client is _CANCELLED:
- _raise_rpc_error(self._state)
- else:
- if self._state.initial_metadata_allowed:
- operation = _get_initial_metadata_operation(
- self._state, initial_metadata)
- self._rpc_event.call.start_server_batch(
- (operation,), _send_initial_metadata(self._state))
- self._state.initial_metadata_allowed = False
- self._state.due.add(_SEND_INITIAL_METADATA_TOKEN)
- else:
- raise ValueError('Initial metadata no longer allowed!')
- def set_trailing_metadata(self, trailing_metadata: MetadataType) -> None:
- with self._state.condition:
- self._state.trailing_metadata = trailing_metadata
- def trailing_metadata(self) -> Optional[MetadataType]:
- return self._state.trailing_metadata
- def abort(self, code: grpc.StatusCode, details: str) -> None:
- # treat OK like other invalid arguments: fail the RPC
- if code == grpc.StatusCode.OK:
- _LOGGER.error(
- 'abort() called with StatusCode.OK; returning UNKNOWN')
- code = grpc.StatusCode.UNKNOWN
- details = ''
- with self._state.condition:
- self._state.code = code
- self._state.details = _common.encode(details)
- self._state.aborted = True
- raise Exception()
- def abort_with_status(self, status: grpc.Status) -> None:
- self._state.trailing_metadata = status.trailing_metadata
- self.abort(status.code, status.details)
- def set_code(self, code: grpc.StatusCode) -> None:
- with self._state.condition:
- self._state.code = code
- def code(self) -> grpc.StatusCode:
- return self._state.code
- def set_details(self, details: str) -> None:
- with self._state.condition:
- self._state.details = _common.encode(details)
- def details(self) -> bytes:
- return self._state.details
- def _finalize_state(self) -> None:
- pass
- class _RequestIterator(object):
- _state: _RPCState
- _call: cygrpc.Call
- _request_deserializer: Optional[DeserializingFunction]
- def __init__(self, state: _RPCState, call: cygrpc.Call,
- request_deserializer: Optional[DeserializingFunction]):
- self._state = state
- self._call = call
- self._request_deserializer = request_deserializer
- def _raise_or_start_receive_message(self) -> None:
- if self._state.client is _CANCELLED:
- _raise_rpc_error(self._state)
- elif not _is_rpc_state_active(self._state):
- raise StopIteration()
- else:
- self._call.start_server_batch(
- (cygrpc.ReceiveMessageOperation(_EMPTY_FLAGS),),
- _receive_message(self._state, self._call,
- self._request_deserializer))
- self._state.due.add(_RECEIVE_MESSAGE_TOKEN)
- def _look_for_request(self) -> Any:
- if self._state.client is _CANCELLED:
- _raise_rpc_error(self._state)
- elif (self._state.request is None and
- _RECEIVE_MESSAGE_TOKEN not in self._state.due):
- raise StopIteration()
- else:
- request = self._state.request
- self._state.request = None
- return request
- raise AssertionError() # should never run
- def _next(self) -> Any:
- with self._state.condition:
- self._raise_or_start_receive_message()
- while True:
- self._state.condition.wait()
- request = self._look_for_request()
- if request is not None:
- return request
- def __iter__(self) -> _RequestIterator:
- return self
- def __next__(self) -> Any:
- return self._next()
- def next(self) -> Any:
- return self._next()
- def _unary_request(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- request_deserializer: Optional[DeserializingFunction]
- ) -> Callable[[], Any]:
- def unary_request():
- with state.condition:
- if not _is_rpc_state_active(state):
- return None
- else:
- rpc_event.call.start_server_batch(
- (cygrpc.ReceiveMessageOperation(_EMPTY_FLAGS),),
- _receive_message(state, rpc_event.call,
- request_deserializer))
- state.due.add(_RECEIVE_MESSAGE_TOKEN)
- while True:
- state.condition.wait()
- if state.request is None:
- if state.client is _CLOSED:
- details = '"{}" requires exactly one request message.'.format(
- rpc_event.call_details.method)
- _abort(state, rpc_event.call,
- cygrpc.StatusCode.unimplemented,
- _common.encode(details))
- return None
- elif state.client is _CANCELLED:
- return None
- else:
- request = state.request
- state.request = None
- return request
- return unary_request
- def _call_behavior(
- rpc_event: cygrpc.BaseEvent,
- state: _RPCState,
- behavior: ArityAgnosticMethodHandler,
- argument: Any,
- request_deserializer: Optional[DeserializingFunction],
- send_response_callback: Optional[Callable[[ResponseType], None]] = None
- ) -> Tuple[Union[ResponseType, Iterator[ResponseType]], bool]:
- from grpc import _create_servicer_context # pytype: disable=pyi-error
- with _create_servicer_context(rpc_event, state,
- request_deserializer) as context:
- try:
- response_or_iterator = None
- if send_response_callback is not None:
- response_or_iterator = behavior(argument, context,
- send_response_callback)
- else:
- response_or_iterator = behavior(argument, context)
- return response_or_iterator, True
- except Exception as exception: # pylint: disable=broad-except
- with state.condition:
- if state.aborted:
- _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
- b'RPC Aborted')
- elif exception not in state.rpc_errors:
- try:
- details = 'Exception calling application: {}'.format(
- exception)
- except Exception: # pylint: disable=broad-except
- details = 'Calling application raised unprintable Exception!'
- traceback.print_exc()
- _LOGGER.exception(details)
- _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
- _common.encode(details))
- return None, False
- def _take_response_from_response_iterator(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- response_iterator: Iterator[ResponseType]) -> Tuple[ResponseType, bool]:
- try:
- return next(response_iterator), True
- except StopIteration:
- return None, True
- except Exception as exception: # pylint: disable=broad-except
- with state.condition:
- if state.aborted:
- _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
- b'RPC Aborted')
- elif exception not in state.rpc_errors:
- details = 'Exception iterating responses: {}'.format(exception)
- _LOGGER.exception(details)
- _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
- _common.encode(details))
- return None, False
- def _serialize_response(
- rpc_event: cygrpc.BaseEvent, state: _RPCState, response: Any,
- response_serializer: Optional[SerializingFunction]) -> Optional[bytes]:
- serialized_response = _common.serialize(response, response_serializer)
- if serialized_response is None:
- with state.condition:
- _abort(state, rpc_event.call, cygrpc.StatusCode.internal,
- b'Failed to serialize response!')
- return None
- else:
- return serialized_response
- def _get_send_message_op_flags_from_state(
- state: _RPCState) -> Union[int, cygrpc.WriteFlag]:
- if state.disable_next_compression:
- return cygrpc.WriteFlag.no_compress
- else:
- return _EMPTY_FLAGS
- def _reset_per_message_state(state: _RPCState) -> None:
- with state.condition:
- state.disable_next_compression = False
- def _send_response(rpc_event: cygrpc.BaseEvent, state: _RPCState,
- serialized_response: bytes) -> bool:
- with state.condition:
- if not _is_rpc_state_active(state):
- return False
- else:
- if state.initial_metadata_allowed:
- operations = (
- _get_initial_metadata_operation(state, None),
- cygrpc.SendMessageOperation(
- serialized_response,
- _get_send_message_op_flags_from_state(state)),
- )
- state.initial_metadata_allowed = False
- token = _SEND_INITIAL_METADATA_AND_SEND_MESSAGE_TOKEN
- else:
- operations = (cygrpc.SendMessageOperation(
- serialized_response,
- _get_send_message_op_flags_from_state(state)),)
- token = _SEND_MESSAGE_TOKEN
- rpc_event.call.start_server_batch(operations,
- _send_message(state, token))
- state.due.add(token)
- _reset_per_message_state(state)
- while True:
- state.condition.wait()
- if token not in state.due:
- return _is_rpc_state_active(state)
- def _status(rpc_event: cygrpc.BaseEvent, state: _RPCState,
- serialized_response: Optional[bytes]) -> None:
- with state.condition:
- if state.client is not _CANCELLED:
- code = _completion_code(state)
- details = _details(state)
- operations = [
- cygrpc.SendStatusFromServerOperation(state.trailing_metadata,
- code, details,
- _EMPTY_FLAGS),
- ]
- if state.initial_metadata_allowed:
- operations.append(_get_initial_metadata_operation(state, None))
- if serialized_response is not None:
- operations.append(
- cygrpc.SendMessageOperation(
- serialized_response,
- _get_send_message_op_flags_from_state(state)))
- rpc_event.call.start_server_batch(
- operations,
- _send_status_from_server(state, _SEND_STATUS_FROM_SERVER_TOKEN))
- state.statused = True
- _reset_per_message_state(state)
- state.due.add(_SEND_STATUS_FROM_SERVER_TOKEN)
- def _unary_response_in_pool(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- behavior: ArityAgnosticMethodHandler, argument_thunk: Callable[[], Any],
- request_deserializer: Optional[SerializingFunction],
- response_serializer: Optional[SerializingFunction]) -> None:
- cygrpc.install_context_from_request_call_event(rpc_event)
- try:
- argument = argument_thunk()
- if argument is not None:
- response, proceed = _call_behavior(rpc_event, state, behavior,
- argument, request_deserializer)
- if proceed:
- serialized_response = _serialize_response(
- rpc_event, state, response, response_serializer)
- if serialized_response is not None:
- _status(rpc_event, state, serialized_response)
- except Exception: # pylint: disable=broad-except
- traceback.print_exc()
- finally:
- cygrpc.uninstall_context()
- def _stream_response_in_pool(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- behavior: ArityAgnosticMethodHandler, argument_thunk: Callable[[], Any],
- request_deserializer: Optional[DeserializingFunction],
- response_serializer: Optional[SerializingFunction]) -> None:
- cygrpc.install_context_from_request_call_event(rpc_event)
- def send_response(response: Any) -> None:
- if response is None:
- _status(rpc_event, state, None)
- else:
- serialized_response = _serialize_response(rpc_event, state,
- response,
- response_serializer)
- if serialized_response is not None:
- _send_response(rpc_event, state, serialized_response)
- try:
- argument = argument_thunk()
- if argument is not None:
- if hasattr(behavior, 'experimental_non_blocking'
- ) and behavior.experimental_non_blocking:
- _call_behavior(rpc_event,
- state,
- behavior,
- argument,
- request_deserializer,
- send_response_callback=send_response)
- else:
- response_iterator, proceed = _call_behavior(
- rpc_event, state, behavior, argument, request_deserializer)
- if proceed:
- _send_message_callback_to_blocking_iterator_adapter(
- rpc_event, state, send_response, response_iterator)
- except Exception: # pylint: disable=broad-except
- traceback.print_exc()
- finally:
- cygrpc.uninstall_context()
- def _is_rpc_state_active(state: _RPCState) -> bool:
- return state.client is not _CANCELLED and not state.statused
- def _send_message_callback_to_blocking_iterator_adapter(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- send_response_callback: Callable[[ResponseType], None],
- response_iterator: Iterator[ResponseType]) -> None:
- while True:
- response, proceed = _take_response_from_response_iterator(
- rpc_event, state, response_iterator)
- if proceed:
- send_response_callback(response)
- if not _is_rpc_state_active(state):
- break
- else:
- break
- def _select_thread_pool_for_behavior(
- behavior: ArityAgnosticMethodHandler,
- default_thread_pool: futures.ThreadPoolExecutor
- ) -> futures.ThreadPoolExecutor:
- if hasattr(behavior, 'experimental_thread_pool') and isinstance(
- behavior.experimental_thread_pool, futures.ThreadPoolExecutor):
- return behavior.experimental_thread_pool
- else:
- return default_thread_pool
- def _handle_unary_unary(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- method_handler: grpc.RpcMethodHandler,
- default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
- unary_request = _unary_request(rpc_event, state,
- method_handler.request_deserializer)
- thread_pool = _select_thread_pool_for_behavior(method_handler.unary_unary,
- default_thread_pool)
- return thread_pool.submit(_unary_response_in_pool, rpc_event, state,
- method_handler.unary_unary, unary_request,
- method_handler.request_deserializer,
- method_handler.response_serializer)
- def _handle_unary_stream(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- method_handler: grpc.RpcMethodHandler,
- default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
- unary_request = _unary_request(rpc_event, state,
- method_handler.request_deserializer)
- thread_pool = _select_thread_pool_for_behavior(method_handler.unary_stream,
- default_thread_pool)
- return thread_pool.submit(_stream_response_in_pool, rpc_event, state,
- method_handler.unary_stream, unary_request,
- method_handler.request_deserializer,
- method_handler.response_serializer)
- def _handle_stream_unary(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- method_handler: grpc.RpcMethodHandler,
- default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
- request_iterator = _RequestIterator(state, rpc_event.call,
- method_handler.request_deserializer)
- thread_pool = _select_thread_pool_for_behavior(method_handler.stream_unary,
- default_thread_pool)
- return thread_pool.submit(_unary_response_in_pool, rpc_event, state,
- method_handler.stream_unary,
- lambda: request_iterator,
- method_handler.request_deserializer,
- method_handler.response_serializer)
- def _handle_stream_stream(
- rpc_event: cygrpc.BaseEvent, state: _RPCState,
- method_handler: grpc.RpcMethodHandler,
- default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
- request_iterator = _RequestIterator(state, rpc_event.call,
- method_handler.request_deserializer)
- thread_pool = _select_thread_pool_for_behavior(method_handler.stream_stream,
- default_thread_pool)
- return thread_pool.submit(_stream_response_in_pool, rpc_event, state,
- method_handler.stream_stream,
- lambda: request_iterator,
- method_handler.request_deserializer,
- method_handler.response_serializer)
- def _find_method_handler(
- rpc_event: cygrpc.BaseEvent, generic_handlers: List[grpc.GenericRpcHandler],
- interceptor_pipeline: Optional[_interceptor._ServicePipeline]
- ) -> Optional[grpc.RpcMethodHandler]:
- def query_handlers(
- handler_call_details: _HandlerCallDetails
- ) -> Optional[grpc.RpcMethodHandler]:
- for generic_handler in generic_handlers:
- method_handler = generic_handler.service(handler_call_details)
- if method_handler is not None:
- return method_handler
- return None
- handler_call_details = _HandlerCallDetails(
- _common.decode(rpc_event.call_details.method),
- rpc_event.invocation_metadata)
- if interceptor_pipeline is not None:
- return interceptor_pipeline.execute(query_handlers,
- handler_call_details)
- else:
- return query_handlers(handler_call_details)
- def _reject_rpc(rpc_event: cygrpc.BaseEvent, status: cygrpc.StatusCode,
- details: bytes) -> _RPCState:
- rpc_state = _RPCState()
- operations = (
- _get_initial_metadata_operation(rpc_state, None),
- cygrpc.ReceiveCloseOnServerOperation(_EMPTY_FLAGS),
- cygrpc.SendStatusFromServerOperation(None, status, details,
- _EMPTY_FLAGS),
- )
- rpc_event.call.start_server_batch(operations, lambda ignored_event: (
- rpc_state,
- (),
- ))
- return rpc_state
- def _handle_with_method_handler(
- rpc_event: cygrpc.BaseEvent, method_handler: grpc.RpcMethodHandler,
- thread_pool: futures.ThreadPoolExecutor
- ) -> Tuple[_RPCState, futures.Future]:
- state = _RPCState()
- with state.condition:
- rpc_event.call.start_server_batch(
- (cygrpc.ReceiveCloseOnServerOperation(_EMPTY_FLAGS),),
- _receive_close_on_server(state))
- state.due.add(_RECEIVE_CLOSE_ON_SERVER_TOKEN)
- if method_handler.request_streaming:
- if method_handler.response_streaming:
- return state, _handle_stream_stream(rpc_event, state,
- method_handler, thread_pool)
- else:
- return state, _handle_stream_unary(rpc_event, state,
- method_handler, thread_pool)
- else:
- if method_handler.response_streaming:
- return state, _handle_unary_stream(rpc_event, state,
- method_handler, thread_pool)
- else:
- return state, _handle_unary_unary(rpc_event, state,
- method_handler, thread_pool)
- def _handle_call(
- rpc_event: cygrpc.BaseEvent, generic_handlers: List[grpc.GenericRpcHandler],
- interceptor_pipeline: Optional[_interceptor._ServicePipeline],
- thread_pool: futures.ThreadPoolExecutor, concurrency_exceeded: bool
- ) -> Tuple[Optional[_RPCState], Optional[futures.Future]]:
- if not rpc_event.success:
- return None, None
- if rpc_event.call_details.method is not None:
- try:
- method_handler = _find_method_handler(rpc_event, generic_handlers,
- interceptor_pipeline)
- except Exception as exception: # pylint: disable=broad-except
- details = 'Exception servicing handler: {}'.format(exception)
- _LOGGER.exception(details)
- return _reject_rpc(rpc_event, cygrpc.StatusCode.unknown,
- b'Error in service handler!'), None
- if method_handler is None:
- return _reject_rpc(rpc_event, cygrpc.StatusCode.unimplemented,
- b'Method not found!'), None
- elif concurrency_exceeded:
- return _reject_rpc(rpc_event, cygrpc.StatusCode.resource_exhausted,
- b'Concurrent RPC limit exceeded!'), None
- else:
- return _handle_with_method_handler(rpc_event, method_handler,
- thread_pool)
- else:
- return None, None
- @enum.unique
- class _ServerStage(enum.Enum):
- STOPPED = 'stopped'
- STARTED = 'started'
- GRACE = 'grace'
- class _ServerState(object):
- lock: threading.RLock
- completion_queue: cygrpc.CompletionQueue
- server: cygrpc.Server
- generic_handlers: List[grpc.GenericRpcHandler]
- interceptor_pipeline: Optional[_interceptor._ServicePipeline]
- thread_pool: futures.ThreadPoolExecutor
- stage: _ServerStage
- termination_event: threading.Event
- shutdown_events: List[threading.Event]
- maximum_concurrent_rpcs: Optional[int]
- active_rpc_count: int
- rpc_states: Set[_RPCState]
- due: Set[str]
- server_deallocated: bool
- # pylint: disable=too-many-arguments
- def __init__(self, completion_queue: cygrpc.CompletionQueue,
- server: cygrpc.Server,
- generic_handlers: Sequence[grpc.GenericRpcHandler],
- interceptor_pipeline: Optional[_interceptor._ServicePipeline],
- thread_pool: futures.ThreadPoolExecutor,
- maximum_concurrent_rpcs: Optional[int]):
- self.lock = threading.RLock()
- self.completion_queue = completion_queue
- self.server = server
- self.generic_handlers = list(generic_handlers)
- self.interceptor_pipeline = interceptor_pipeline
- self.thread_pool = thread_pool
- self.stage = _ServerStage.STOPPED
- self.termination_event = threading.Event()
- self.shutdown_events = [self.termination_event]
- self.maximum_concurrent_rpcs = maximum_concurrent_rpcs
- self.active_rpc_count = 0
- # TODO(https://github.com/grpc/grpc/issues/6597): eliminate these fields.
- self.rpc_states = set()
- self.due = set()
- # A "volatile" flag to interrupt the daemon serving thread
- self.server_deallocated = False
- def _add_generic_handlers(
- state: _ServerState,
- generic_handlers: Iterable[grpc.GenericRpcHandler]) -> None:
- with state.lock:
- state.generic_handlers.extend(generic_handlers)
- def _add_insecure_port(state: _ServerState, address: bytes) -> int:
- with state.lock:
- return state.server.add_http2_port(address)
- def _add_secure_port(state: _ServerState, address: bytes,
- server_credentials: grpc.ServerCredentials) -> int:
- with state.lock:
- return state.server.add_http2_port(address,
- server_credentials._credentials)
- def _request_call(state: _ServerState) -> None:
- state.server.request_call(state.completion_queue, state.completion_queue,
- _REQUEST_CALL_TAG)
- state.due.add(_REQUEST_CALL_TAG)
- # TODO(https://github.com/grpc/grpc/issues/6597): delete this function.
- def _stop_serving(state: _ServerState) -> bool:
- if not state.rpc_states and not state.due:
- state.server.destroy()
- for shutdown_event in state.shutdown_events:
- shutdown_event.set()
- state.stage = _ServerStage.STOPPED
- return True
- else:
- return False
- def _on_call_completed(state: _ServerState) -> None:
- with state.lock:
- state.active_rpc_count -= 1
- def _process_event_and_continue(state: _ServerState,
- event: cygrpc.BaseEvent) -> bool:
- should_continue = True
- if event.tag is _SHUTDOWN_TAG:
- with state.lock:
- state.due.remove(_SHUTDOWN_TAG)
- if _stop_serving(state):
- should_continue = False
- elif event.tag is _REQUEST_CALL_TAG:
- with state.lock:
- state.due.remove(_REQUEST_CALL_TAG)
- concurrency_exceeded = (
- state.maximum_concurrent_rpcs is not None and
- state.active_rpc_count >= state.maximum_concurrent_rpcs)
- rpc_state, rpc_future = _handle_call(event, state.generic_handlers,
- state.interceptor_pipeline,
- state.thread_pool,
- concurrency_exceeded)
- if rpc_state is not None:
- state.rpc_states.add(rpc_state)
- if rpc_future is not None:
- state.active_rpc_count += 1
- rpc_future.add_done_callback(
- lambda unused_future: _on_call_completed(state))
- if state.stage is _ServerStage.STARTED:
- _request_call(state)
- elif _stop_serving(state):
- should_continue = False
- else:
- rpc_state, callbacks = event.tag(event)
- for callback in callbacks:
- try:
- callback()
- except Exception: # pylint: disable=broad-except
- _LOGGER.exception('Exception calling callback!')
- if rpc_state is not None:
- with state.lock:
- state.rpc_states.remove(rpc_state)
- if _stop_serving(state):
- should_continue = False
- return should_continue
- def _serve(state: _ServerState) -> None:
- while True:
- timeout = time.time() + _DEALLOCATED_SERVER_CHECK_PERIOD_S
- event = state.completion_queue.poll(timeout)
- if state.server_deallocated:
- _begin_shutdown_once(state)
- if event.completion_type != cygrpc.CompletionType.queue_timeout:
- if not _process_event_and_continue(state, event):
- return
- # We want to force the deletion of the previous event
- # ~before~ we poll again; if the event has a reference
- # to a shutdown Call object, this can induce spinlock.
- event = None
- def _begin_shutdown_once(state: _ServerState) -> None:
- with state.lock:
- if state.stage is _ServerStage.STARTED:
- state.server.shutdown(state.completion_queue, _SHUTDOWN_TAG)
- state.stage = _ServerStage.GRACE
- state.due.add(_SHUTDOWN_TAG)
- def _stop(state: _ServerState, grace: Optional[float]) -> threading.Event:
- with state.lock:
- if state.stage is _ServerStage.STOPPED:
- shutdown_event = threading.Event()
- shutdown_event.set()
- return shutdown_event
- else:
- _begin_shutdown_once(state)
- shutdown_event = threading.Event()
- state.shutdown_events.append(shutdown_event)
- if grace is None:
- state.server.cancel_all_calls()
- else:
- def cancel_all_calls_after_grace():
- shutdown_event.wait(timeout=grace)
- with state.lock:
- state.server.cancel_all_calls()
- thread = threading.Thread(target=cancel_all_calls_after_grace)
- thread.start()
- return shutdown_event
- shutdown_event.wait()
- return shutdown_event
- def _start(state: _ServerState) -> None:
- with state.lock:
- if state.stage is not _ServerStage.STOPPED:
- raise ValueError('Cannot start already-started server!')
- state.server.start()
- state.stage = _ServerStage.STARTED
- _request_call(state)
- thread = threading.Thread(target=_serve, args=(state,))
- thread.daemon = True
- thread.start()
- def _validate_generic_rpc_handlers(
- generic_rpc_handlers: Iterable[grpc.GenericRpcHandler]) -> None:
- for generic_rpc_handler in generic_rpc_handlers:
- service_attribute = getattr(generic_rpc_handler, 'service', None)
- if service_attribute is None:
- raise AttributeError(
- '"{}" must conform to grpc.GenericRpcHandler type but does '
- 'not have "service" method!'.format(generic_rpc_handler))
- def _augment_options(
- base_options: Sequence[ChannelArgumentType],
- compression: Optional[grpc.Compression]
- ) -> Sequence[ChannelArgumentType]:
- compression_option = _compression.create_channel_option(compression)
- return tuple(base_options) + compression_option
- class _Server(grpc.Server):
- _state: _ServerState
- # pylint: disable=too-many-arguments
- def __init__(self, thread_pool: futures.ThreadPoolExecutor,
- generic_handlers: Sequence[grpc.GenericRpcHandler],
- interceptors: Sequence[grpc.ServerInterceptor],
- options: Sequence[ChannelArgumentType],
- maximum_concurrent_rpcs: Optional[int],
- compression: Optional[grpc.Compression], xds: bool):
- completion_queue = cygrpc.CompletionQueue()
- server = cygrpc.Server(_augment_options(options, compression), xds)
- server.register_completion_queue(completion_queue)
- self._state = _ServerState(completion_queue, server, generic_handlers,
- _interceptor.service_pipeline(interceptors),
- thread_pool, maximum_concurrent_rpcs)
- def add_generic_rpc_handlers(
- self,
- generic_rpc_handlers: Iterable[grpc.GenericRpcHandler]) -> None:
- _validate_generic_rpc_handlers(generic_rpc_handlers)
- _add_generic_handlers(self._state, generic_rpc_handlers)
- def add_insecure_port(self, address: str) -> int:
- return _common.validate_port_binding_result(
- address, _add_insecure_port(self._state, _common.encode(address)))
- def add_secure_port(self, address: str,
- server_credentials: grpc.ServerCredentials) -> int:
- return _common.validate_port_binding_result(
- address,
- _add_secure_port(self._state, _common.encode(address),
- server_credentials))
- def start(self) -> None:
- _start(self._state)
- def wait_for_termination(self, timeout: Optional[float] = None) -> bool:
- # NOTE(https://bugs.python.org/issue35935)
- # Remove this workaround once threading.Event.wait() is working with
- # CTRL+C across platforms.
- return _common.wait(self._state.termination_event.wait,
- self._state.termination_event.is_set,
- timeout=timeout)
- def stop(self, grace: Optional[float]) -> threading.Event:
- return _stop(self._state, grace)
- def __del__(self):
- if hasattr(self, '_state'):
- # We can not grab a lock in __del__(), so set a flag to signal the
- # serving daemon thread (if it exists) to initiate shutdown.
- self._state.server_deallocated = True
- def create_server(thread_pool: futures.ThreadPoolExecutor,
- generic_rpc_handlers: Sequence[grpc.GenericRpcHandler],
- interceptors: Sequence[grpc.ServerInterceptor],
- options: Sequence[ChannelArgumentType],
- maximum_concurrent_rpcs: Optional[int],
- compression: Optional[grpc.Compression],
- xds: bool) -> _Server:
- _validate_generic_rpc_handlers(generic_rpc_handlers)
- return _Server(thread_pool, generic_rpc_handlers, interceptors, options,
- maximum_concurrent_rpcs, compression, xds)
|