feat(sdk): hand health reporter to the caller

sdk/README.md

@@ -13,7 +13,8 @@ API. This crate provides:
   gRPC server reflection, the gRPC health service, graceful shutdown, and
   the standard CLI arguments. `serve` runs a typed `FunctionRunnerService`;
   `serve_service` runs any tonic service (a custom codec, an instrumented
-  wrapper) over the same transport.
+  wrapper) over the same transport; `serve_service_with_health` also hands
+  back the health reporter, so readiness can flip after start-up work.
 - `request`, `response`, `resource` - helpers for typed function input,
   context keys, required resources and schemas, credentials, capabilities,
   results, conditions, readiness, and desired resource updates from any

sdk/src/lib.rs

@@ -22,7 +22,7 @@ pub mod resource;
 pub mod response;
 pub mod server;
 
-pub use server::{Args, serve, serve_service};
+pub use server::{Args, HealthReporter, serve, serve_service, serve_service_with_health};
 
 /// Errors returned by the SDK.
 #[derive(Debug, thiserror::Error)]

sdk/src/server.rs

@@ -1,5 +1,6 @@
 //! A spec-compliant gRPC server runtime for composition functions.
 
+use std::future::Future;
 use std::net::SocketAddr;
 use std::path::{Path, PathBuf};
 
@@ -8,6 +9,7 @@ use tonic::codegen::Service;
 use tonic::codegen::http;
 use tonic::server::NamedService;
 use tonic::transport::{Certificate, Identity, Server, ServerTlsConfig};
+pub use tonic_health::server::HealthReporter;
 
 use crate::Error;
 use crate::proto::v1::function_runner_service_server::{
@@ -70,6 +72,61 @@ pub async fn serve<F: FunctionRunnerService>(function: F, args: &Args) -> Result
 /// decoding, so `--max-recv-message-size` is its business to enforce;
 /// [`Args::max_recv_message_size`] carries what the caller asked for.
 pub async fn serve_service<S>(service: S, args: &Args) -> Result<(), Error>
+where
+    S: Service<
+            http::Request<Body>,
+            Response = http::Response<Body>,
+            Error = std::convert::Infallible,
+        > + NamedService
+        + Clone
+        + Send
+        + Sync
+        + 'static,
+    S::Future: Send + 'static,
+{
+    let (health_reporter, server) = serve_service_with_health(service, args).await?;
+    health_reporter.set_serving::<S>().await;
+    server.await
+}
+
+/// Like [`serve_service`], but hands back the gRPC health reporter next to
+/// the server future instead of reporting the service as serving up front.
+///
+/// The service starts as NOT SERVING; flip it with
+/// [`HealthReporter::set_serving`] once the function is ready - typically
+/// from a start-up task (warming caches, loading modules) running while the
+/// server already answers. Requests are served regardless of health status:
+/// health is what probes read, not a gate. Nothing listens until the
+/// returned future is awaited.
+///
+/// ```no_run
+/// # async fn example<S>(service: S, args: &function_sdk_rust::Args)
+/// # -> Result<(), function_sdk_rust::Error>
+/// # where S: tonic::codegen::Service<
+/// #         tonic::codegen::http::Request<tonic::body::Body>,
+/// #         Response = tonic::codegen::http::Response<tonic::body::Body>,
+/// #         Error = std::convert::Infallible,
+/// #     > + tonic::server::NamedService + Clone + Send + Sync + 'static,
+/// #     S::Future: Send + 'static,
+/// # {
+/// let (health, server) = function_sdk_rust::serve_service_with_health(service, args).await?;
+/// tokio::spawn(async move {
+///     // ... warm caches ...
+///     health.set_serving::<S>().await;
+/// });
+/// server.await
+/// # }
+/// ```
+pub async fn serve_service_with_health<S>(
+    service: S,
+    args: &Args,
+) -> Result<
+    (
+        HealthReporter,
+        impl Future<Output = Result<(), Error>> + use<S>,
+    ),
+    Error,
+>
 where
     S: Service<
             http::Request<Body>,
@@ -101,19 +158,21 @@ where
         .build_v1alpha()?;
 
     let (health_reporter, health_service) = tonic_health::server::health_reporter();
-    health_reporter.set_serving::<S>().await;
-
-    tracing::info!(%address, insecure = args.insecure, service = S::NAME, "serving");
-
-    builder
-        .add_service(service)
-        .add_service(health_service)
-        .add_service(reflection_v1)
-        .add_service(reflection_v1alpha)
-        .serve_with_shutdown(address, shutdown_signal())
-        .await?;
-
-    Ok(())
+    health_reporter.set_not_serving::<S>().await;
+
+    let insecure = args.insecure;
+    let server = async move {
+        tracing::info!(%address, insecure, service = S::NAME, "serving");
+        builder
+            .add_service(service)
+            .add_service(health_service)
+            .add_service(reflection_v1)
+            .add_service(reflection_v1alpha)
+            .serve_with_shutdown(address, shutdown_signal())
+            .await?;
+        Ok(())
+    };
+    Ok((health_reporter, server))
 }
 
 fn tls_config(dir: &Path) -> Result<ServerTlsConfig, Error> {

sdk/tests/serve.rs

@@ -2,7 +2,7 @@ use function_sdk_rust::proto::v1::function_runner_service_client::FunctionRunner
 use function_sdk_rust::proto::v1::function_runner_service_server::FunctionRunnerService;
 use function_sdk_rust::proto::v1::function_runner_service_server::FunctionRunnerServiceServer;
 use function_sdk_rust::proto::v1::{RequestMeta, RunFunctionRequest, RunFunctionResponse};
-use function_sdk_rust::{Args, response, serve, serve_service};
+use function_sdk_rust::{Args, response, serve, serve_service, serve_service_with_health};
 use tonic::{Request, Response, Status};
 
 struct EchoFunction;
@@ -160,6 +160,72 @@ async fn serve_service_serves_a_custom_service() {
     );
 }
 
+#[tokio::test]
+async fn serve_service_with_health_hands_readiness_to_the_caller() {
+    let port = free_port();
+    let args = Args {
+        debug: false,
+        address: format!("127.0.0.1:{port}"),
+        tls_certs_dir: None,
+        insecure: true,
+        max_recv_message_size: None,
+    };
+    let service = FunctionRunnerServiceServer::new(EchoFunction);
+    let (health, server) = serve_service_with_health(service, &args)
+        .await
+        .expect("server must build");
+    tokio::spawn(server);
+
+    let mut client = connect_with_retry(port).await;
+
+    let channel = tonic::transport::Channel::from_shared(format!("http://127.0.0.1:{port}"))
+        .expect("valid URI")
+        .connect()
+        .await
+        .expect("health channel must connect");
+    let mut health_client = tonic_health::pb::health_client::HealthClient::new(channel);
+    let check =
+        |c: &mut tonic_health::pb::health_client::HealthClient<tonic::transport::Channel>| {
+            let mut c = c.clone();
+            async move {
+                c.check(tonic_health::pb::HealthCheckRequest {
+                    service: "apiextensions.fn.proto.v1.FunctionRunnerService".to_string(),
+                })
+                .await
+                .expect("health check must succeed")
+                .into_inner()
+                .status
+            }
+        };
+
+    // Not serving until the caller says so - but requests are answered:
+    // health is what probes read, not a gate.
+    assert_eq!(
+        check(&mut health_client).await,
+        tonic_health::pb::health_check_response::ServingStatus::NotServing as i32
+    );
+    let rsp = client
+        .run_function(RunFunctionRequest {
+            meta: Some(RequestMeta {
+                tag: "early".to_string(),
+                capabilities: vec![],
+            }),
+            ..Default::default()
+        })
+        .await
+        .expect("RunFunction must succeed while not serving")
+        .into_inner();
+    assert_eq!(rsp.results[0].message, "echo");
+
+    health
+        .set_serving::<FunctionRunnerServiceServer<EchoFunction>>()
+        .await;
+    assert_eq!(
+        check(&mut health_client).await,
+        tonic_health::pb::health_check_response::ServingStatus::Serving as i32
+    );
+}
+
 fn free_port() -> u16 {
     std::net::TcpListener::bind("127.0.0.1:0")
         .expect("bind to an ephemeral port")