_server.py 44 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141
  1. # Copyright 2016 gRPC authors.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Service-side implementation of gRPC Python."""
  15. from __future__ import annotations
  16. import collections
  17. from concurrent import futures
  18. import enum
  19. import logging
  20. import threading
  21. import time
  22. import traceback
  23. from typing import (Any, Callable, Iterable, Iterator, List, Mapping, Optional,
  24. Sequence, Set, Tuple, Union)
  25. import grpc # pytype: disable=pyi-error
  26. from grpc import _common # pytype: disable=pyi-error
  27. from grpc import _compression # pytype: disable=pyi-error
  28. from grpc import _interceptor # pytype: disable=pyi-error
  29. from grpc._cython import cygrpc
  30. from grpc._typing import ArityAgnosticMethodHandler
  31. from grpc._typing import ChannelArgumentType
  32. from grpc._typing import DeserializingFunction
  33. from grpc._typing import MetadataType
  34. from grpc._typing import NullaryCallbackType
  35. from grpc._typing import ResponseType
  36. from grpc._typing import SerializingFunction
  37. from grpc._typing import ServerCallbackTag
  38. from grpc._typing import ServerTagCallbackType
  39. _LOGGER = logging.getLogger(__name__)
  40. _SHUTDOWN_TAG = 'shutdown'
  41. _REQUEST_CALL_TAG = 'request_call'
  42. _RECEIVE_CLOSE_ON_SERVER_TOKEN = 'receive_close_on_server'
  43. _SEND_INITIAL_METADATA_TOKEN = 'send_initial_metadata'
  44. _RECEIVE_MESSAGE_TOKEN = 'receive_message'
  45. _SEND_MESSAGE_TOKEN = 'send_message'
  46. _SEND_INITIAL_METADATA_AND_SEND_MESSAGE_TOKEN = (
  47. 'send_initial_metadata * send_message')
  48. _SEND_STATUS_FROM_SERVER_TOKEN = 'send_status_from_server'
  49. _SEND_INITIAL_METADATA_AND_SEND_STATUS_FROM_SERVER_TOKEN = (
  50. 'send_initial_metadata * send_status_from_server')
  51. _OPEN = 'open'
  52. _CLOSED = 'closed'
  53. _CANCELLED = 'cancelled'
  54. _EMPTY_FLAGS = 0
  55. _DEALLOCATED_SERVER_CHECK_PERIOD_S = 1.0
  56. _INF_TIMEOUT = 1e9
  57. def _serialized_request(request_event: cygrpc.BaseEvent) -> bytes:
  58. return request_event.batch_operations[0].message()
  59. def _application_code(code: grpc.StatusCode) -> cygrpc.StatusCode:
  60. cygrpc_code = _common.STATUS_CODE_TO_CYGRPC_STATUS_CODE.get(code)
  61. return cygrpc.StatusCode.unknown if cygrpc_code is None else cygrpc_code
  62. def _completion_code(state: _RPCState) -> cygrpc.StatusCode:
  63. if state.code is None:
  64. return cygrpc.StatusCode.ok
  65. else:
  66. return _application_code(state.code)
  67. def _abortion_code(state: _RPCState,
  68. code: cygrpc.StatusCode) -> cygrpc.StatusCode:
  69. if state.code is None:
  70. return code
  71. else:
  72. return _application_code(state.code)
  73. def _details(state: _RPCState) -> bytes:
  74. return b'' if state.details is None else state.details
  75. class _HandlerCallDetails(
  76. collections.namedtuple('_HandlerCallDetails', (
  77. 'method',
  78. 'invocation_metadata',
  79. )), grpc.HandlerCallDetails):
  80. pass
  81. class _RPCState(object):
  82. condition: threading.Condition
  83. due = Set[str]
  84. request: Any
  85. client: str
  86. initial_metadata_allowed: bool
  87. compression_algorithm: Optional[grpc.Compression]
  88. disable_next_compression: bool
  89. trailing_metadata: Optional[MetadataType]
  90. code: Optional[grpc.StatusCode]
  91. details: Optional[bytes]
  92. statused: bool
  93. rpc_errors: List[Exception]
  94. callbacks: Optional[List[NullaryCallbackType]]
  95. aborted: bool
  96. def __init__(self):
  97. self.condition = threading.Condition()
  98. self.due = set()
  99. self.request = None
  100. self.client = _OPEN
  101. self.initial_metadata_allowed = True
  102. self.compression_algorithm = None
  103. self.disable_next_compression = False
  104. self.trailing_metadata = None
  105. self.code = None
  106. self.details = None
  107. self.statused = False
  108. self.rpc_errors = []
  109. self.callbacks = []
  110. self.aborted = False
  111. def _raise_rpc_error(state: _RPCState) -> None:
  112. rpc_error = grpc.RpcError()
  113. state.rpc_errors.append(rpc_error)
  114. raise rpc_error
  115. def _possibly_finish_call(state: _RPCState,
  116. token: str) -> ServerTagCallbackType:
  117. state.due.remove(token)
  118. if not _is_rpc_state_active(state) and not state.due:
  119. callbacks = state.callbacks
  120. state.callbacks = None
  121. return state, callbacks
  122. else:
  123. return None, ()
  124. def _send_status_from_server(state: _RPCState, token: str) -> ServerCallbackTag:
  125. def send_status_from_server(unused_send_status_from_server_event):
  126. with state.condition:
  127. return _possibly_finish_call(state, token)
  128. return send_status_from_server
  129. def _get_initial_metadata(
  130. state: _RPCState,
  131. metadata: Optional[MetadataType]) -> Optional[MetadataType]:
  132. with state.condition:
  133. if state.compression_algorithm:
  134. compression_metadata = (
  135. _compression.compression_algorithm_to_metadata(
  136. state.compression_algorithm),)
  137. if metadata is None:
  138. return compression_metadata
  139. else:
  140. return compression_metadata + tuple(metadata)
  141. else:
  142. return metadata
  143. def _get_initial_metadata_operation(
  144. state: _RPCState, metadata: Optional[MetadataType]) -> cygrpc.Operation:
  145. operation = cygrpc.SendInitialMetadataOperation(
  146. _get_initial_metadata(state, metadata), _EMPTY_FLAGS)
  147. return operation
  148. def _abort(state: _RPCState, call: cygrpc.Call, code: cygrpc.StatusCode,
  149. details: bytes) -> None:
  150. if state.client is not _CANCELLED:
  151. effective_code = _abortion_code(state, code)
  152. effective_details = details if state.details is None else state.details
  153. if state.initial_metadata_allowed:
  154. operations = (
  155. _get_initial_metadata_operation(state, None),
  156. cygrpc.SendStatusFromServerOperation(state.trailing_metadata,
  157. effective_code,
  158. effective_details,
  159. _EMPTY_FLAGS),
  160. )
  161. token = _SEND_INITIAL_METADATA_AND_SEND_STATUS_FROM_SERVER_TOKEN
  162. else:
  163. operations = (cygrpc.SendStatusFromServerOperation(
  164. state.trailing_metadata, effective_code, effective_details,
  165. _EMPTY_FLAGS),)
  166. token = _SEND_STATUS_FROM_SERVER_TOKEN
  167. call.start_server_batch(operations,
  168. _send_status_from_server(state, token))
  169. state.statused = True
  170. state.due.add(token)
  171. def _receive_close_on_server(state: _RPCState) -> ServerCallbackTag:
  172. def receive_close_on_server(receive_close_on_server_event):
  173. with state.condition:
  174. if receive_close_on_server_event.batch_operations[0].cancelled():
  175. state.client = _CANCELLED
  176. elif state.client is _OPEN:
  177. state.client = _CLOSED
  178. state.condition.notify_all()
  179. return _possibly_finish_call(state, _RECEIVE_CLOSE_ON_SERVER_TOKEN)
  180. return receive_close_on_server
  181. def _receive_message(
  182. state: _RPCState, call: cygrpc.Call,
  183. request_deserializer: Optional[DeserializingFunction]
  184. ) -> ServerCallbackTag:
  185. def receive_message(receive_message_event):
  186. serialized_request = _serialized_request(receive_message_event)
  187. if serialized_request is None:
  188. with state.condition:
  189. if state.client is _OPEN:
  190. state.client = _CLOSED
  191. state.condition.notify_all()
  192. return _possibly_finish_call(state, _RECEIVE_MESSAGE_TOKEN)
  193. else:
  194. request = _common.deserialize(serialized_request,
  195. request_deserializer)
  196. with state.condition:
  197. if request is None:
  198. _abort(state, call, cygrpc.StatusCode.internal,
  199. b'Exception deserializing request!')
  200. else:
  201. state.request = request
  202. state.condition.notify_all()
  203. return _possibly_finish_call(state, _RECEIVE_MESSAGE_TOKEN)
  204. return receive_message
  205. def _send_initial_metadata(state: _RPCState) -> ServerCallbackTag:
  206. def send_initial_metadata(unused_send_initial_metadata_event):
  207. with state.condition:
  208. return _possibly_finish_call(state, _SEND_INITIAL_METADATA_TOKEN)
  209. return send_initial_metadata
  210. def _send_message(state: _RPCState, token: str) -> ServerCallbackTag:
  211. def send_message(unused_send_message_event):
  212. with state.condition:
  213. state.condition.notify_all()
  214. return _possibly_finish_call(state, token)
  215. return send_message
  216. class _Context(grpc.ServicerContext):
  217. _rpc_event: cygrpc.BaseEvent
  218. _state: _RPCState
  219. request_deserializer: Optional[DeserializingFunction]
  220. def __init__(self, rpc_event: cygrpc.BaseEvent, state: _RPCState,
  221. request_deserializer: Optional[DeserializingFunction]):
  222. self._rpc_event = rpc_event
  223. self._state = state
  224. self._request_deserializer = request_deserializer
  225. def is_active(self) -> bool:
  226. with self._state.condition:
  227. return _is_rpc_state_active(self._state)
  228. def time_remaining(self) -> float:
  229. return max(self._rpc_event.call_details.deadline - time.time(), 0)
  230. def cancel(self) -> None:
  231. self._rpc_event.call.cancel()
  232. def add_callback(self, callback: NullaryCallbackType) -> bool:
  233. with self._state.condition:
  234. if self._state.callbacks is None:
  235. return False
  236. else:
  237. self._state.callbacks.append(callback)
  238. return True
  239. def disable_next_message_compression(self) -> None:
  240. with self._state.condition:
  241. self._state.disable_next_compression = True
  242. def invocation_metadata(self) -> Optional[MetadataType]:
  243. return self._rpc_event.invocation_metadata
  244. def peer(self) -> str:
  245. return _common.decode(self._rpc_event.call.peer())
  246. def peer_identities(self) -> Optional[Sequence[bytes]]:
  247. return cygrpc.peer_identities(self._rpc_event.call)
  248. def peer_identity_key(self) -> Optional[str]:
  249. id_key = cygrpc.peer_identity_key(self._rpc_event.call)
  250. return id_key if id_key is None else _common.decode(id_key)
  251. def auth_context(self) -> Mapping[str, Sequence[bytes]]:
  252. auth_context = cygrpc.auth_context(self._rpc_event.call)
  253. auth_context_dict = {} if auth_context is None else auth_context
  254. return {
  255. _common.decode(key): value
  256. for key, value in auth_context_dict.items()
  257. }
  258. def set_compression(self, compression: grpc.Compression) -> None:
  259. with self._state.condition:
  260. self._state.compression_algorithm = compression
  261. def send_initial_metadata(self, initial_metadata: MetadataType) -> None:
  262. with self._state.condition:
  263. if self._state.client is _CANCELLED:
  264. _raise_rpc_error(self._state)
  265. else:
  266. if self._state.initial_metadata_allowed:
  267. operation = _get_initial_metadata_operation(
  268. self._state, initial_metadata)
  269. self._rpc_event.call.start_server_batch(
  270. (operation,), _send_initial_metadata(self._state))
  271. self._state.initial_metadata_allowed = False
  272. self._state.due.add(_SEND_INITIAL_METADATA_TOKEN)
  273. else:
  274. raise ValueError('Initial metadata no longer allowed!')
  275. def set_trailing_metadata(self, trailing_metadata: MetadataType) -> None:
  276. with self._state.condition:
  277. self._state.trailing_metadata = trailing_metadata
  278. def trailing_metadata(self) -> Optional[MetadataType]:
  279. return self._state.trailing_metadata
  280. def abort(self, code: grpc.StatusCode, details: str) -> None:
  281. # treat OK like other invalid arguments: fail the RPC
  282. if code == grpc.StatusCode.OK:
  283. _LOGGER.error(
  284. 'abort() called with StatusCode.OK; returning UNKNOWN')
  285. code = grpc.StatusCode.UNKNOWN
  286. details = ''
  287. with self._state.condition:
  288. self._state.code = code
  289. self._state.details = _common.encode(details)
  290. self._state.aborted = True
  291. raise Exception()
  292. def abort_with_status(self, status: grpc.Status) -> None:
  293. self._state.trailing_metadata = status.trailing_metadata
  294. self.abort(status.code, status.details)
  295. def set_code(self, code: grpc.StatusCode) -> None:
  296. with self._state.condition:
  297. self._state.code = code
  298. def code(self) -> grpc.StatusCode:
  299. return self._state.code
  300. def set_details(self, details: str) -> None:
  301. with self._state.condition:
  302. self._state.details = _common.encode(details)
  303. def details(self) -> bytes:
  304. return self._state.details
  305. def _finalize_state(self) -> None:
  306. pass
  307. class _RequestIterator(object):
  308. _state: _RPCState
  309. _call: cygrpc.Call
  310. _request_deserializer: Optional[DeserializingFunction]
  311. def __init__(self, state: _RPCState, call: cygrpc.Call,
  312. request_deserializer: Optional[DeserializingFunction]):
  313. self._state = state
  314. self._call = call
  315. self._request_deserializer = request_deserializer
  316. def _raise_or_start_receive_message(self) -> None:
  317. if self._state.client is _CANCELLED:
  318. _raise_rpc_error(self._state)
  319. elif not _is_rpc_state_active(self._state):
  320. raise StopIteration()
  321. else:
  322. self._call.start_server_batch(
  323. (cygrpc.ReceiveMessageOperation(_EMPTY_FLAGS),),
  324. _receive_message(self._state, self._call,
  325. self._request_deserializer))
  326. self._state.due.add(_RECEIVE_MESSAGE_TOKEN)
  327. def _look_for_request(self) -> Any:
  328. if self._state.client is _CANCELLED:
  329. _raise_rpc_error(self._state)
  330. elif (self._state.request is None and
  331. _RECEIVE_MESSAGE_TOKEN not in self._state.due):
  332. raise StopIteration()
  333. else:
  334. request = self._state.request
  335. self._state.request = None
  336. return request
  337. raise AssertionError() # should never run
  338. def _next(self) -> Any:
  339. with self._state.condition:
  340. self._raise_or_start_receive_message()
  341. while True:
  342. self._state.condition.wait()
  343. request = self._look_for_request()
  344. if request is not None:
  345. return request
  346. def __iter__(self) -> _RequestIterator:
  347. return self
  348. def __next__(self) -> Any:
  349. return self._next()
  350. def next(self) -> Any:
  351. return self._next()
  352. def _unary_request(
  353. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  354. request_deserializer: Optional[DeserializingFunction]
  355. ) -> Callable[[], Any]:
  356. def unary_request():
  357. with state.condition:
  358. if not _is_rpc_state_active(state):
  359. return None
  360. else:
  361. rpc_event.call.start_server_batch(
  362. (cygrpc.ReceiveMessageOperation(_EMPTY_FLAGS),),
  363. _receive_message(state, rpc_event.call,
  364. request_deserializer))
  365. state.due.add(_RECEIVE_MESSAGE_TOKEN)
  366. while True:
  367. state.condition.wait()
  368. if state.request is None:
  369. if state.client is _CLOSED:
  370. details = '"{}" requires exactly one request message.'.format(
  371. rpc_event.call_details.method)
  372. _abort(state, rpc_event.call,
  373. cygrpc.StatusCode.unimplemented,
  374. _common.encode(details))
  375. return None
  376. elif state.client is _CANCELLED:
  377. return None
  378. else:
  379. request = state.request
  380. state.request = None
  381. return request
  382. return unary_request
  383. def _call_behavior(
  384. rpc_event: cygrpc.BaseEvent,
  385. state: _RPCState,
  386. behavior: ArityAgnosticMethodHandler,
  387. argument: Any,
  388. request_deserializer: Optional[DeserializingFunction],
  389. send_response_callback: Optional[Callable[[ResponseType], None]] = None
  390. ) -> Tuple[Union[ResponseType, Iterator[ResponseType]], bool]:
  391. from grpc import _create_servicer_context # pytype: disable=pyi-error
  392. with _create_servicer_context(rpc_event, state,
  393. request_deserializer) as context:
  394. try:
  395. response_or_iterator = None
  396. if send_response_callback is not None:
  397. response_or_iterator = behavior(argument, context,
  398. send_response_callback)
  399. else:
  400. response_or_iterator = behavior(argument, context)
  401. return response_or_iterator, True
  402. except Exception as exception: # pylint: disable=broad-except
  403. with state.condition:
  404. if state.aborted:
  405. _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
  406. b'RPC Aborted')
  407. elif exception not in state.rpc_errors:
  408. try:
  409. details = 'Exception calling application: {}'.format(
  410. exception)
  411. except Exception: # pylint: disable=broad-except
  412. details = 'Calling application raised unprintable Exception!'
  413. traceback.print_exc()
  414. _LOGGER.exception(details)
  415. _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
  416. _common.encode(details))
  417. return None, False
  418. def _take_response_from_response_iterator(
  419. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  420. response_iterator: Iterator[ResponseType]) -> Tuple[ResponseType, bool]:
  421. try:
  422. return next(response_iterator), True
  423. except StopIteration:
  424. return None, True
  425. except Exception as exception: # pylint: disable=broad-except
  426. with state.condition:
  427. if state.aborted:
  428. _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
  429. b'RPC Aborted')
  430. elif exception not in state.rpc_errors:
  431. details = 'Exception iterating responses: {}'.format(exception)
  432. _LOGGER.exception(details)
  433. _abort(state, rpc_event.call, cygrpc.StatusCode.unknown,
  434. _common.encode(details))
  435. return None, False
  436. def _serialize_response(
  437. rpc_event: cygrpc.BaseEvent, state: _RPCState, response: Any,
  438. response_serializer: Optional[SerializingFunction]) -> Optional[bytes]:
  439. serialized_response = _common.serialize(response, response_serializer)
  440. if serialized_response is None:
  441. with state.condition:
  442. _abort(state, rpc_event.call, cygrpc.StatusCode.internal,
  443. b'Failed to serialize response!')
  444. return None
  445. else:
  446. return serialized_response
  447. def _get_send_message_op_flags_from_state(
  448. state: _RPCState) -> Union[int, cygrpc.WriteFlag]:
  449. if state.disable_next_compression:
  450. return cygrpc.WriteFlag.no_compress
  451. else:
  452. return _EMPTY_FLAGS
  453. def _reset_per_message_state(state: _RPCState) -> None:
  454. with state.condition:
  455. state.disable_next_compression = False
  456. def _send_response(rpc_event: cygrpc.BaseEvent, state: _RPCState,
  457. serialized_response: bytes) -> bool:
  458. with state.condition:
  459. if not _is_rpc_state_active(state):
  460. return False
  461. else:
  462. if state.initial_metadata_allowed:
  463. operations = (
  464. _get_initial_metadata_operation(state, None),
  465. cygrpc.SendMessageOperation(
  466. serialized_response,
  467. _get_send_message_op_flags_from_state(state)),
  468. )
  469. state.initial_metadata_allowed = False
  470. token = _SEND_INITIAL_METADATA_AND_SEND_MESSAGE_TOKEN
  471. else:
  472. operations = (cygrpc.SendMessageOperation(
  473. serialized_response,
  474. _get_send_message_op_flags_from_state(state)),)
  475. token = _SEND_MESSAGE_TOKEN
  476. rpc_event.call.start_server_batch(operations,
  477. _send_message(state, token))
  478. state.due.add(token)
  479. _reset_per_message_state(state)
  480. while True:
  481. state.condition.wait()
  482. if token not in state.due:
  483. return _is_rpc_state_active(state)
  484. def _status(rpc_event: cygrpc.BaseEvent, state: _RPCState,
  485. serialized_response: Optional[bytes]) -> None:
  486. with state.condition:
  487. if state.client is not _CANCELLED:
  488. code = _completion_code(state)
  489. details = _details(state)
  490. operations = [
  491. cygrpc.SendStatusFromServerOperation(state.trailing_metadata,
  492. code, details,
  493. _EMPTY_FLAGS),
  494. ]
  495. if state.initial_metadata_allowed:
  496. operations.append(_get_initial_metadata_operation(state, None))
  497. if serialized_response is not None:
  498. operations.append(
  499. cygrpc.SendMessageOperation(
  500. serialized_response,
  501. _get_send_message_op_flags_from_state(state)))
  502. rpc_event.call.start_server_batch(
  503. operations,
  504. _send_status_from_server(state, _SEND_STATUS_FROM_SERVER_TOKEN))
  505. state.statused = True
  506. _reset_per_message_state(state)
  507. state.due.add(_SEND_STATUS_FROM_SERVER_TOKEN)
  508. def _unary_response_in_pool(
  509. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  510. behavior: ArityAgnosticMethodHandler, argument_thunk: Callable[[], Any],
  511. request_deserializer: Optional[SerializingFunction],
  512. response_serializer: Optional[SerializingFunction]) -> None:
  513. cygrpc.install_context_from_request_call_event(rpc_event)
  514. try:
  515. argument = argument_thunk()
  516. if argument is not None:
  517. response, proceed = _call_behavior(rpc_event, state, behavior,
  518. argument, request_deserializer)
  519. if proceed:
  520. serialized_response = _serialize_response(
  521. rpc_event, state, response, response_serializer)
  522. if serialized_response is not None:
  523. _status(rpc_event, state, serialized_response)
  524. except Exception: # pylint: disable=broad-except
  525. traceback.print_exc()
  526. finally:
  527. cygrpc.uninstall_context()
  528. def _stream_response_in_pool(
  529. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  530. behavior: ArityAgnosticMethodHandler, argument_thunk: Callable[[], Any],
  531. request_deserializer: Optional[DeserializingFunction],
  532. response_serializer: Optional[SerializingFunction]) -> None:
  533. cygrpc.install_context_from_request_call_event(rpc_event)
  534. def send_response(response: Any) -> None:
  535. if response is None:
  536. _status(rpc_event, state, None)
  537. else:
  538. serialized_response = _serialize_response(rpc_event, state,
  539. response,
  540. response_serializer)
  541. if serialized_response is not None:
  542. _send_response(rpc_event, state, serialized_response)
  543. try:
  544. argument = argument_thunk()
  545. if argument is not None:
  546. if hasattr(behavior, 'experimental_non_blocking'
  547. ) and behavior.experimental_non_blocking:
  548. _call_behavior(rpc_event,
  549. state,
  550. behavior,
  551. argument,
  552. request_deserializer,
  553. send_response_callback=send_response)
  554. else:
  555. response_iterator, proceed = _call_behavior(
  556. rpc_event, state, behavior, argument, request_deserializer)
  557. if proceed:
  558. _send_message_callback_to_blocking_iterator_adapter(
  559. rpc_event, state, send_response, response_iterator)
  560. except Exception: # pylint: disable=broad-except
  561. traceback.print_exc()
  562. finally:
  563. cygrpc.uninstall_context()
  564. def _is_rpc_state_active(state: _RPCState) -> bool:
  565. return state.client is not _CANCELLED and not state.statused
  566. def _send_message_callback_to_blocking_iterator_adapter(
  567. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  568. send_response_callback: Callable[[ResponseType], None],
  569. response_iterator: Iterator[ResponseType]) -> None:
  570. while True:
  571. response, proceed = _take_response_from_response_iterator(
  572. rpc_event, state, response_iterator)
  573. if proceed:
  574. send_response_callback(response)
  575. if not _is_rpc_state_active(state):
  576. break
  577. else:
  578. break
  579. def _select_thread_pool_for_behavior(
  580. behavior: ArityAgnosticMethodHandler,
  581. default_thread_pool: futures.ThreadPoolExecutor
  582. ) -> futures.ThreadPoolExecutor:
  583. if hasattr(behavior, 'experimental_thread_pool') and isinstance(
  584. behavior.experimental_thread_pool, futures.ThreadPoolExecutor):
  585. return behavior.experimental_thread_pool
  586. else:
  587. return default_thread_pool
  588. def _handle_unary_unary(
  589. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  590. method_handler: grpc.RpcMethodHandler,
  591. default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
  592. unary_request = _unary_request(rpc_event, state,
  593. method_handler.request_deserializer)
  594. thread_pool = _select_thread_pool_for_behavior(method_handler.unary_unary,
  595. default_thread_pool)
  596. return thread_pool.submit(_unary_response_in_pool, rpc_event, state,
  597. method_handler.unary_unary, unary_request,
  598. method_handler.request_deserializer,
  599. method_handler.response_serializer)
  600. def _handle_unary_stream(
  601. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  602. method_handler: grpc.RpcMethodHandler,
  603. default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
  604. unary_request = _unary_request(rpc_event, state,
  605. method_handler.request_deserializer)
  606. thread_pool = _select_thread_pool_for_behavior(method_handler.unary_stream,
  607. default_thread_pool)
  608. return thread_pool.submit(_stream_response_in_pool, rpc_event, state,
  609. method_handler.unary_stream, unary_request,
  610. method_handler.request_deserializer,
  611. method_handler.response_serializer)
  612. def _handle_stream_unary(
  613. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  614. method_handler: grpc.RpcMethodHandler,
  615. default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
  616. request_iterator = _RequestIterator(state, rpc_event.call,
  617. method_handler.request_deserializer)
  618. thread_pool = _select_thread_pool_for_behavior(method_handler.stream_unary,
  619. default_thread_pool)
  620. return thread_pool.submit(_unary_response_in_pool, rpc_event, state,
  621. method_handler.stream_unary,
  622. lambda: request_iterator,
  623. method_handler.request_deserializer,
  624. method_handler.response_serializer)
  625. def _handle_stream_stream(
  626. rpc_event: cygrpc.BaseEvent, state: _RPCState,
  627. method_handler: grpc.RpcMethodHandler,
  628. default_thread_pool: futures.ThreadPoolExecutor) -> futures.Future:
  629. request_iterator = _RequestIterator(state, rpc_event.call,
  630. method_handler.request_deserializer)
  631. thread_pool = _select_thread_pool_for_behavior(method_handler.stream_stream,
  632. default_thread_pool)
  633. return thread_pool.submit(_stream_response_in_pool, rpc_event, state,
  634. method_handler.stream_stream,
  635. lambda: request_iterator,
  636. method_handler.request_deserializer,
  637. method_handler.response_serializer)
  638. def _find_method_handler(
  639. rpc_event: cygrpc.BaseEvent, generic_handlers: List[grpc.GenericRpcHandler],
  640. interceptor_pipeline: Optional[_interceptor._ServicePipeline]
  641. ) -> Optional[grpc.RpcMethodHandler]:
  642. def query_handlers(
  643. handler_call_details: _HandlerCallDetails
  644. ) -> Optional[grpc.RpcMethodHandler]:
  645. for generic_handler in generic_handlers:
  646. method_handler = generic_handler.service(handler_call_details)
  647. if method_handler is not None:
  648. return method_handler
  649. return None
  650. handler_call_details = _HandlerCallDetails(
  651. _common.decode(rpc_event.call_details.method),
  652. rpc_event.invocation_metadata)
  653. if interceptor_pipeline is not None:
  654. return interceptor_pipeline.execute(query_handlers,
  655. handler_call_details)
  656. else:
  657. return query_handlers(handler_call_details)
  658. def _reject_rpc(rpc_event: cygrpc.BaseEvent, status: cygrpc.StatusCode,
  659. details: bytes) -> _RPCState:
  660. rpc_state = _RPCState()
  661. operations = (
  662. _get_initial_metadata_operation(rpc_state, None),
  663. cygrpc.ReceiveCloseOnServerOperation(_EMPTY_FLAGS),
  664. cygrpc.SendStatusFromServerOperation(None, status, details,
  665. _EMPTY_FLAGS),
  666. )
  667. rpc_event.call.start_server_batch(operations, lambda ignored_event: (
  668. rpc_state,
  669. (),
  670. ))
  671. return rpc_state
  672. def _handle_with_method_handler(
  673. rpc_event: cygrpc.BaseEvent, method_handler: grpc.RpcMethodHandler,
  674. thread_pool: futures.ThreadPoolExecutor
  675. ) -> Tuple[_RPCState, futures.Future]:
  676. state = _RPCState()
  677. with state.condition:
  678. rpc_event.call.start_server_batch(
  679. (cygrpc.ReceiveCloseOnServerOperation(_EMPTY_FLAGS),),
  680. _receive_close_on_server(state))
  681. state.due.add(_RECEIVE_CLOSE_ON_SERVER_TOKEN)
  682. if method_handler.request_streaming:
  683. if method_handler.response_streaming:
  684. return state, _handle_stream_stream(rpc_event, state,
  685. method_handler, thread_pool)
  686. else:
  687. return state, _handle_stream_unary(rpc_event, state,
  688. method_handler, thread_pool)
  689. else:
  690. if method_handler.response_streaming:
  691. return state, _handle_unary_stream(rpc_event, state,
  692. method_handler, thread_pool)
  693. else:
  694. return state, _handle_unary_unary(rpc_event, state,
  695. method_handler, thread_pool)
  696. def _handle_call(
  697. rpc_event: cygrpc.BaseEvent, generic_handlers: List[grpc.GenericRpcHandler],
  698. interceptor_pipeline: Optional[_interceptor._ServicePipeline],
  699. thread_pool: futures.ThreadPoolExecutor, concurrency_exceeded: bool
  700. ) -> Tuple[Optional[_RPCState], Optional[futures.Future]]:
  701. if not rpc_event.success:
  702. return None, None
  703. if rpc_event.call_details.method is not None:
  704. try:
  705. method_handler = _find_method_handler(rpc_event, generic_handlers,
  706. interceptor_pipeline)
  707. except Exception as exception: # pylint: disable=broad-except
  708. details = 'Exception servicing handler: {}'.format(exception)
  709. _LOGGER.exception(details)
  710. return _reject_rpc(rpc_event, cygrpc.StatusCode.unknown,
  711. b'Error in service handler!'), None
  712. if method_handler is None:
  713. return _reject_rpc(rpc_event, cygrpc.StatusCode.unimplemented,
  714. b'Method not found!'), None
  715. elif concurrency_exceeded:
  716. return _reject_rpc(rpc_event, cygrpc.StatusCode.resource_exhausted,
  717. b'Concurrent RPC limit exceeded!'), None
  718. else:
  719. return _handle_with_method_handler(rpc_event, method_handler,
  720. thread_pool)
  721. else:
  722. return None, None
  723. @enum.unique
  724. class _ServerStage(enum.Enum):
  725. STOPPED = 'stopped'
  726. STARTED = 'started'
  727. GRACE = 'grace'
  728. class _ServerState(object):
  729. lock: threading.RLock
  730. completion_queue: cygrpc.CompletionQueue
  731. server: cygrpc.Server
  732. generic_handlers: List[grpc.GenericRpcHandler]
  733. interceptor_pipeline: Optional[_interceptor._ServicePipeline]
  734. thread_pool: futures.ThreadPoolExecutor
  735. stage: _ServerStage
  736. termination_event: threading.Event
  737. shutdown_events: List[threading.Event]
  738. maximum_concurrent_rpcs: Optional[int]
  739. active_rpc_count: int
  740. rpc_states: Set[_RPCState]
  741. due: Set[str]
  742. server_deallocated: bool
  743. # pylint: disable=too-many-arguments
  744. def __init__(self, completion_queue: cygrpc.CompletionQueue,
  745. server: cygrpc.Server,
  746. generic_handlers: Sequence[grpc.GenericRpcHandler],
  747. interceptor_pipeline: Optional[_interceptor._ServicePipeline],
  748. thread_pool: futures.ThreadPoolExecutor,
  749. maximum_concurrent_rpcs: Optional[int]):
  750. self.lock = threading.RLock()
  751. self.completion_queue = completion_queue
  752. self.server = server
  753. self.generic_handlers = list(generic_handlers)
  754. self.interceptor_pipeline = interceptor_pipeline
  755. self.thread_pool = thread_pool
  756. self.stage = _ServerStage.STOPPED
  757. self.termination_event = threading.Event()
  758. self.shutdown_events = [self.termination_event]
  759. self.maximum_concurrent_rpcs = maximum_concurrent_rpcs
  760. self.active_rpc_count = 0
  761. # TODO(https://github.com/grpc/grpc/issues/6597): eliminate these fields.
  762. self.rpc_states = set()
  763. self.due = set()
  764. # A "volatile" flag to interrupt the daemon serving thread
  765. self.server_deallocated = False
  766. def _add_generic_handlers(
  767. state: _ServerState,
  768. generic_handlers: Iterable[grpc.GenericRpcHandler]) -> None:
  769. with state.lock:
  770. state.generic_handlers.extend(generic_handlers)
  771. def _add_insecure_port(state: _ServerState, address: bytes) -> int:
  772. with state.lock:
  773. return state.server.add_http2_port(address)
  774. def _add_secure_port(state: _ServerState, address: bytes,
  775. server_credentials: grpc.ServerCredentials) -> int:
  776. with state.lock:
  777. return state.server.add_http2_port(address,
  778. server_credentials._credentials)
  779. def _request_call(state: _ServerState) -> None:
  780. state.server.request_call(state.completion_queue, state.completion_queue,
  781. _REQUEST_CALL_TAG)
  782. state.due.add(_REQUEST_CALL_TAG)
  783. # TODO(https://github.com/grpc/grpc/issues/6597): delete this function.
  784. def _stop_serving(state: _ServerState) -> bool:
  785. if not state.rpc_states and not state.due:
  786. state.server.destroy()
  787. for shutdown_event in state.shutdown_events:
  788. shutdown_event.set()
  789. state.stage = _ServerStage.STOPPED
  790. return True
  791. else:
  792. return False
  793. def _on_call_completed(state: _ServerState) -> None:
  794. with state.lock:
  795. state.active_rpc_count -= 1
  796. def _process_event_and_continue(state: _ServerState,
  797. event: cygrpc.BaseEvent) -> bool:
  798. should_continue = True
  799. if event.tag is _SHUTDOWN_TAG:
  800. with state.lock:
  801. state.due.remove(_SHUTDOWN_TAG)
  802. if _stop_serving(state):
  803. should_continue = False
  804. elif event.tag is _REQUEST_CALL_TAG:
  805. with state.lock:
  806. state.due.remove(_REQUEST_CALL_TAG)
  807. concurrency_exceeded = (
  808. state.maximum_concurrent_rpcs is not None and
  809. state.active_rpc_count >= state.maximum_concurrent_rpcs)
  810. rpc_state, rpc_future = _handle_call(event, state.generic_handlers,
  811. state.interceptor_pipeline,
  812. state.thread_pool,
  813. concurrency_exceeded)
  814. if rpc_state is not None:
  815. state.rpc_states.add(rpc_state)
  816. if rpc_future is not None:
  817. state.active_rpc_count += 1
  818. rpc_future.add_done_callback(
  819. lambda unused_future: _on_call_completed(state))
  820. if state.stage is _ServerStage.STARTED:
  821. _request_call(state)
  822. elif _stop_serving(state):
  823. should_continue = False
  824. else:
  825. rpc_state, callbacks = event.tag(event)
  826. for callback in callbacks:
  827. try:
  828. callback()
  829. except Exception: # pylint: disable=broad-except
  830. _LOGGER.exception('Exception calling callback!')
  831. if rpc_state is not None:
  832. with state.lock:
  833. state.rpc_states.remove(rpc_state)
  834. if _stop_serving(state):
  835. should_continue = False
  836. return should_continue
  837. def _serve(state: _ServerState) -> None:
  838. while True:
  839. timeout = time.time() + _DEALLOCATED_SERVER_CHECK_PERIOD_S
  840. event = state.completion_queue.poll(timeout)
  841. if state.server_deallocated:
  842. _begin_shutdown_once(state)
  843. if event.completion_type != cygrpc.CompletionType.queue_timeout:
  844. if not _process_event_and_continue(state, event):
  845. return
  846. # We want to force the deletion of the previous event
  847. # ~before~ we poll again; if the event has a reference
  848. # to a shutdown Call object, this can induce spinlock.
  849. event = None
  850. def _begin_shutdown_once(state: _ServerState) -> None:
  851. with state.lock:
  852. if state.stage is _ServerStage.STARTED:
  853. state.server.shutdown(state.completion_queue, _SHUTDOWN_TAG)
  854. state.stage = _ServerStage.GRACE
  855. state.due.add(_SHUTDOWN_TAG)
  856. def _stop(state: _ServerState, grace: Optional[float]) -> threading.Event:
  857. with state.lock:
  858. if state.stage is _ServerStage.STOPPED:
  859. shutdown_event = threading.Event()
  860. shutdown_event.set()
  861. return shutdown_event
  862. else:
  863. _begin_shutdown_once(state)
  864. shutdown_event = threading.Event()
  865. state.shutdown_events.append(shutdown_event)
  866. if grace is None:
  867. state.server.cancel_all_calls()
  868. else:
  869. def cancel_all_calls_after_grace():
  870. shutdown_event.wait(timeout=grace)
  871. with state.lock:
  872. state.server.cancel_all_calls()
  873. thread = threading.Thread(target=cancel_all_calls_after_grace)
  874. thread.start()
  875. return shutdown_event
  876. shutdown_event.wait()
  877. return shutdown_event
  878. def _start(state: _ServerState) -> None:
  879. with state.lock:
  880. if state.stage is not _ServerStage.STOPPED:
  881. raise ValueError('Cannot start already-started server!')
  882. state.server.start()
  883. state.stage = _ServerStage.STARTED
  884. _request_call(state)
  885. thread = threading.Thread(target=_serve, args=(state,))
  886. thread.daemon = True
  887. thread.start()
  888. def _validate_generic_rpc_handlers(
  889. generic_rpc_handlers: Iterable[grpc.GenericRpcHandler]) -> None:
  890. for generic_rpc_handler in generic_rpc_handlers:
  891. service_attribute = getattr(generic_rpc_handler, 'service', None)
  892. if service_attribute is None:
  893. raise AttributeError(
  894. '"{}" must conform to grpc.GenericRpcHandler type but does '
  895. 'not have "service" method!'.format(generic_rpc_handler))
  896. def _augment_options(
  897. base_options: Sequence[ChannelArgumentType],
  898. compression: Optional[grpc.Compression]
  899. ) -> Sequence[ChannelArgumentType]:
  900. compression_option = _compression.create_channel_option(compression)
  901. return tuple(base_options) + compression_option
  902. class _Server(grpc.Server):
  903. _state: _ServerState
  904. # pylint: disable=too-many-arguments
  905. def __init__(self, thread_pool: futures.ThreadPoolExecutor,
  906. generic_handlers: Sequence[grpc.GenericRpcHandler],
  907. interceptors: Sequence[grpc.ServerInterceptor],
  908. options: Sequence[ChannelArgumentType],
  909. maximum_concurrent_rpcs: Optional[int],
  910. compression: Optional[grpc.Compression], xds: bool):
  911. completion_queue = cygrpc.CompletionQueue()
  912. server = cygrpc.Server(_augment_options(options, compression), xds)
  913. server.register_completion_queue(completion_queue)
  914. self._state = _ServerState(completion_queue, server, generic_handlers,
  915. _interceptor.service_pipeline(interceptors),
  916. thread_pool, maximum_concurrent_rpcs)
  917. def add_generic_rpc_handlers(
  918. self,
  919. generic_rpc_handlers: Iterable[grpc.GenericRpcHandler]) -> None:
  920. _validate_generic_rpc_handlers(generic_rpc_handlers)
  921. _add_generic_handlers(self._state, generic_rpc_handlers)
  922. def add_insecure_port(self, address: str) -> int:
  923. return _common.validate_port_binding_result(
  924. address, _add_insecure_port(self._state, _common.encode(address)))
  925. def add_secure_port(self, address: str,
  926. server_credentials: grpc.ServerCredentials) -> int:
  927. return _common.validate_port_binding_result(
  928. address,
  929. _add_secure_port(self._state, _common.encode(address),
  930. server_credentials))
  931. def start(self) -> None:
  932. _start(self._state)
  933. def wait_for_termination(self, timeout: Optional[float] = None) -> bool:
  934. # NOTE(https://bugs.python.org/issue35935)
  935. # Remove this workaround once threading.Event.wait() is working with
  936. # CTRL+C across platforms.
  937. return _common.wait(self._state.termination_event.wait,
  938. self._state.termination_event.is_set,
  939. timeout=timeout)
  940. def stop(self, grace: Optional[float]) -> threading.Event:
  941. return _stop(self._state, grace)
  942. def __del__(self):
  943. if hasattr(self, '_state'):
  944. # We can not grab a lock in __del__(), so set a flag to signal the
  945. # serving daemon thread (if it exists) to initiate shutdown.
  946. self._state.server_deallocated = True
  947. def create_server(thread_pool: futures.ThreadPoolExecutor,
  948. generic_rpc_handlers: Sequence[grpc.GenericRpcHandler],
  949. interceptors: Sequence[grpc.ServerInterceptor],
  950. options: Sequence[ChannelArgumentType],
  951. maximum_concurrent_rpcs: Optional[int],
  952. compression: Optional[grpc.Compression],
  953. xds: bool) -> _Server:
  954. _validate_generic_rpc_handlers(generic_rpc_handlers)
  955. return _Server(thread_pool, generic_rpc_handlers, interceptors, options,
  956. maximum_concurrent_rpcs, compression, xds)