From a2495502df6133bdd939be6ef83ce417249e7ec3 Mon Sep 17 00:00:00 2001 From: Eric Gribkoff Date: Wed, 20 Feb 2019 15:54:11 -0800 Subject: [PATCH] add enter_graceful_shutdown() to health service --- .../grpc_health/v1/health.py | 17 +++++++++++++ .../health_check/_health_servicer_test.py | 24 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/src/python/grpcio_health_checking/grpc_health/v1/health.py b/src/python/grpcio_health_checking/grpc_health/v1/health.py index b08297a5d71..04ea3b4ecfb 100644 --- a/src/python/grpcio_health_checking/grpc_health/v1/health.py +++ b/src/python/grpcio_health_checking/grpc_health/v1/health.py @@ -82,6 +82,7 @@ class HealthServicer(_health_pb2_grpc.HealthServicer): self._send_response_callbacks = {} self.Watch.__func__.experimental_non_blocking = experimental_non_blocking self.Watch.__func__.experimental_thread_pool = experimental_thread_pool + self._gracefully_shutting_down = False def _on_close_callback(self, send_response_callback, service): @@ -135,9 +136,25 @@ class HealthServicer(_health_pb2_grpc.HealthServicer): the service """ with self._lock: + if self._gracefully_shutting_down: + return 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 enter_graceful_shutdown(self): + """Permanently sets the status of all services to NOT_SERVING. + + This should be invoked when the server is entering a graceful shutdown + period. After this method is invoked, future attempts to set the status + of a service will be ignored. + """ + with self._lock: + if self._gracefully_shutting_down: + return + for service in self._server_status: + self.set(service, _health_pb2.HealthCheckResponse.NOT_SERVING) # pylint: disable=no-member + self._gracefully_shutting_down = True diff --git a/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py b/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py index f92596f5c9b..62eef8e2291 100644 --- a/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py +++ b/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py @@ -203,6 +203,30 @@ class BaseWatchTests(object): 'watch set should be empty') self.assertTrue(response_queue.empty()) + def test_graceful_shutdown(self): + request = health_pb2.HealthCheckRequest(service='') + response_queue = queue.Queue() + rendezvous = self._stub.Watch(request) + thread = threading.Thread( + target=_consume_responses, args=(rendezvous, response_queue)) + thread.start() + + response = response_queue.get(timeout=test_constants.SHORT_TIMEOUT) + self.assertEqual(health_pb2.HealthCheckResponse.SERVING, + response.status) + + self._servicer.enter_graceful_shutdown() + response = response_queue.get(timeout=test_constants.SHORT_TIMEOUT) + self.assertEqual(health_pb2.HealthCheckResponse.NOT_SERVING, + response.status) + + # This should be a no-op. + self._servicer.set('', health_pb2.HealthCheckResponse.SERVING) + + rendezvous.cancel() + thread.join() + self.assertTrue(response_queue.empty()) + class HealthServicerTest(BaseWatchTests.WatchTests):