feat grpc: add graceful shutdown headers to trailing metadata if graceful shutdown started after initial metadata sending

commit_hash:9c719dac7c45a0c977b598a5c05049d66f734aca

Author: lexeyo

grpc/functional_tests/middleware_server/tests/test_graceful_shutdown_headers.py

@@ -1,3 +1,4 @@
+import asyncio
 from signal import SIGTERM
 
 import pytest
@@ -20,8 +21,7 @@ async def test_graceful_shutdown_headers(service_daemon_instance, grpc_client, s
     ]
 
     for headers_enabled, headers in params:
-        dynamic_config.set_values({'GRACEFUL_SHUTDOWN_HEADERS': {'enabled': headers_enabled, 'headers': headers}})
-        await service_client.update_server_state()
+        await update_graceful_shutdown_headers(service_client, dynamic_config, headers_enabled, headers)
 
         request = greeter_protos.GreetingRequest(name='Python')
         call = grpc_client.SayHello(request)
@@ -33,8 +33,7 @@ async def test_graceful_shutdown_headers(service_daemon_instance, grpc_client, s
     service_daemon_instance.process.send_signal(SIGTERM)
 
     for headers_enabled, headers in params:
-        dynamic_config.set_values({'GRACEFUL_SHUTDOWN_HEADERS': {'enabled': headers_enabled, 'headers': headers}})
-        await service_client.update_server_state()
+        await update_graceful_shutdown_headers(service_client, dynamic_config, headers_enabled, headers)
 
         request = greeter_protos.GreetingRequest(name='Python')
         call = grpc_client.SayHello(request)
@@ -52,6 +51,83 @@ async def test_graceful_shutdown_headers(service_daemon_instance, grpc_client, s
     service_daemon_instance.process.wait()
 
 
+@pytest.mark.uservice_oneshot
+async def test_graceful_shutdown_headers_streams(service_daemon_instance, grpc_client, service_client, dynamic_config):
+    headers = {'x-envoy-immediate-health-check-fail': ['true']}
+    await update_graceful_shutdown_headers(service_client, dynamic_config, True, headers)
+
+    start = f'Python{MD_ONE_REQ}{MD_TWO_REQ}'
+    end = f'{MD_TWO_RES}{MD_ONE_RES}'
+
+    service_daemon_instance.process.send_signal(SIGTERM)
+
+    call = grpc_client.SayHelloStreams(
+        _prepare_requests(['Python', '!', '!', '!'], 0.1),
+        wait_for_ready=True,
+    )
+    async for response in call:
+        assert response.greeting == f'Hello, {start}{end}'
+        start += f'!{MD_ONE_REQ}{MD_TWO_REQ}'
+
+    check_present(await call.initial_metadata(), headers)
+    check_not_present(await call.trailing_metadata(), headers)
+
+    # After a couple more seconds, the service will start shutting down.
+    service_daemon_instance.process.wait()
+
+
+@pytest.mark.uservice_oneshot
+async def test_late_graceful_shutdown_headers_streams(
+    service_daemon_instance, grpc_client, service_client, dynamic_config
+):
+    headers = {'x-envoy-immediate-health-check-fail': ['true']}
+    await update_graceful_shutdown_headers(service_client, dynamic_config, True, headers)
+
+    start = f'Python{MD_ONE_REQ}{MD_TWO_REQ}'
+    end = f'{MD_TWO_RES}{MD_ONE_RES}'
+
+    call = grpc_client.SayHelloStreams(
+        _prepare_late_requests(service_daemon_instance, ['Python', '!', '!'], 0.2),
+        wait_for_ready=True,
+    )
+    async for response in call:
+        assert response.greeting == f'Hello, {start}{end}'
+        start += f'!{MD_ONE_REQ}{MD_TWO_REQ}'
+
+    check_not_present(await call.initial_metadata(), headers)
+    check_present(await call.trailing_metadata(), headers)
+
+    # After a couple more seconds, the service will start shutting down.
+    service_daemon_instance.process.wait()
+
+
+async def _prepare_requests(names, sleep=1):
+    reqs = []
+    for name in names:
+        reqs.append(greeter_protos.GreetingRequest(name=name))
+    for req in reqs:
+        await asyncio.sleep(sleep)
+        yield req
+
+
+async def _prepare_late_requests(service_daemon_instance, names, sleep=1):
+    reqs = []
+    for name in names:
+        reqs.append(greeter_protos.GreetingRequest(name=name))
+    i = 0
+    for req in reqs:
+        if i == 1:
+            service_daemon_instance.process.send_signal(SIGTERM)
+        i += 1
+        await asyncio.sleep(sleep)
+        yield req
+
+
+async def update_graceful_shutdown_headers(service_client, dynamic_config, headers_enabled, headers):
+    dynamic_config.set_values({'GRACEFUL_SHUTDOWN_HEADERS': {'enabled': headers_enabled, 'headers': headers}})
+    await service_client.update_server_state()
+
+
 def check_present(metadata, headers: dict[str, list[str]]):
     metadata_dict = to_dict(metadata)
     for k, v in headers.items():

grpc/src/ugrpc/server/middlewares/graceful_shutdown_headers/middleware.cpp

@@ -7,6 +7,35 @@ USERVER_NAMESPACE_BEGIN
 
 namespace ugrpc::server::middlewares::graceful_shutdown_headers {
 
+namespace {
+
+struct Identity {};
+
+inline const utils::AnyStorageDataTag<ugrpc::server::StorageContext, Identity> kGracefulShutdownHeadersSet;
+
+using AddMetadataMethod = decltype(&grpc::ServerContext::AddInitialMetadata);
+
+bool SetHeadersIfNeeded(
+    const dynamic_config::Snapshot& config,
+    MiddlewareCallContext& context,
+    AddMetadataMethod add_metadata_method
+) {
+    const auto& graceful_shutdown_headers = config[::dynamic_config::GRACEFUL_SHUTDOWN_HEADERS];
+    if (!graceful_shutdown_headers.enabled) {
+        return false;
+    }
+    auto& server_context = context.GetServerContext();
+    for (const auto& [name, values] : graceful_shutdown_headers.headers.extra) {
+        auto header_name = ugrpc::impl::ToGrpcString(name);
+        for (const auto& value : values) {
+            (server_context.*add_metadata_method)(header_name, ugrpc::impl::ToGrpcString(value));
+        }
+    }
+    return true;
+}
+
+}  // namespace
+
 Middleware::Middleware(const components::State& state, dynamic_config::Source source)
     : state_(state),
       source_(std::move(source))
@@ -16,18 +45,20 @@ void Middleware::OnCallStart(MiddlewareCallContext& context) const {
     if (state_.GetServiceLifetimeStage() != components::ServiceLifetimeStage::kGracefulShutdown) {
         return;
     }
-    const auto& config = context.GetInitialDynamicConfig();
-    const auto& graceful_shutdown_headers = config[::dynamic_config::GRACEFUL_SHUTDOWN_HEADERS];
-    if (!graceful_shutdown_headers.enabled) {
+    if (SetHeadersIfNeeded(context.GetInitialDynamicConfig(), context, &grpc::ServerContext::AddInitialMetadata)) {
+        context.GetStorageContext().Set(kGracefulShutdownHeadersSet, {});
+    }
+}
+
+void Middleware::PreSendStatus(MiddlewareCallContext& context, grpc::Status&) const {
+    if (state_.GetServiceLifetimeStage() != components::ServiceLifetimeStage::kGracefulShutdown) {
         return;
     }
-    auto& server_context = context.GetServerContext();
-    for (const auto& [name, values] : graceful_shutdown_headers.headers.extra) {
-        auto header_name = ugrpc::impl::ToGrpcString(name);
-        for (const auto& value : values) {
-            server_context.AddInitialMetadata(header_name, ugrpc::impl::ToGrpcString(value));
-        }
+    if (context.GetStorageContext().GetOptional(kGracefulShutdownHeadersSet) != nullptr) {
+        return;
     }
+
+    SetHeadersIfNeeded(source_.GetSnapshot(), context, &grpc::ServerContext::AddTrailingMetadata);
 }
 
 }  // namespace ugrpc::server::middlewares::graceful_shutdown_headers

grpc/src/ugrpc/server/middlewares/graceful_shutdown_headers/middleware.hpp

@@ -13,6 +13,7 @@ class Middleware final : public MiddlewareBase {
     Middleware(const components::State& state, dynamic_config::Source source);
 
     void OnCallStart(MiddlewareCallContext& context) const override;
+    void PreSendStatus(MiddlewareCallContext& context, grpc::Status&) const override;
 
 private:
     const components::State state_;