def Check(self, request, context):
with self._lock:
status = self._server_status.get(request.service)
if status is None:
context.set_code(grpc.StatusCode.NOT_FOUND)
return _health_pb2.HealthCheckResponse()
else:
return _health_pb2.HealthCheckResponse(status=status)
async def Watch(self, request: _health_pb2.HealthCheckRequest,
context) -> None:
condition = self._server_watchers[request.service]
last_status = None
try:
async with condition:
while True:
status = self._server_status.get(
request.service,
_health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
# NOTE(lidiz) If the observed status is the same, it means
# there are missing intermediate statuses. It's considered
# acceptable since peer only interested in eventual status.
if status != last_status:
# Responds with current health state
await context.write(
_health_pb2.HealthCheckResponse(status=status))
# Records the last sent status
last_status = status
# Polling on health state changes
await condition.wait()
finally:
if request.service in self._server_watchers:
del self._server_watchers[request.service]
async def Check(self, request: _health_pb2.HealthCheckRequest,
context) -> None:
status = self._server_status.get(request.service)
if status is None:
await context.abort(grpc.StatusCode.NOT_FOUND)
else:
return _health_pb2.HealthCheckResponse(status=status)
def set(self, service, status):
"""Sets the status of a service.
Args:
service: string, the name of the service.
status: HealthCheckResponse.status enum value indicating the status of
the service
"""
with self._lock:
if self._gracefully_shutting_down:
return
else:
self._server_status[service] = status
if service in self._send_response_callbacks:
for send_response_callback in self._send_response_callbacks[
service]:
send_response_callback(
_health_pb2.HealthCheckResponse(status=status))
def Watch(self, request, context, send_response_callback=None):
blocking_watcher = None
if send_response_callback is None:
# The server does not support the experimental_non_blocking
# parameter. For backwards compatibility, return a blocking response
# generator.
blocking_watcher = _Watcher()
send_response_callback = _watcher_to_send_response_callback_adapter(
blocking_watcher)
service = request.service
with self._lock:
status = self._server_status.get(service)
if status is None:
status = _health_pb2.HealthCheckResponse.SERVICE_UNKNOWN # pylint: disable=no-member
send_response_callback(
_health_pb2.HealthCheckResponse(status=status))
if service not in self._send_response_callbacks:
self._send_response_callbacks[service] = set()
self._send_response_callbacks[service].add(send_response_callback)
context.add_callback(
self._on_close_callback(send_response_callback, service))
return blocking_watcher