diff --git a/sdk/README.md b/sdk/README.md index 75f0502..4995ad9 100644 --- a/sdk/README.md +++ b/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 diff --git a/sdk/src/lib.rs b/sdk/src/lib.rs index 75f1eb8..550fd60 100644 --- a/sdk/src/lib.rs +++ b/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)] diff --git a/sdk/src/server.rs b/sdk/src/server.rs index 3a99ef2..702b650 100644 --- a/sdk/src/server.rs +++ b/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(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(service: S, args: &Args) -> Result<(), Error> +where + S: Service< + http::Request, + Response = http::Response, + 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::().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(service: S, args: &function_sdk_rust::Args) +/// # -> Result<(), function_sdk_rust::Error> +/// # where S: tonic::codegen::Service< +/// # tonic::codegen::http::Request, +/// # Response = tonic::codegen::http::Response, +/// # 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::().await; +/// }); +/// server.await +/// # } +/// ``` +pub async fn serve_service_with_health( + service: S, + args: &Args, +) -> Result< + ( + HealthReporter, + impl Future> + use, + ), + Error, +> where S: Service< http::Request, @@ -101,19 +158,21 @@ where .build_v1alpha()?; let (health_reporter, health_service) = tonic_health::server::health_reporter(); - health_reporter.set_serving::().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::().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 { diff --git a/sdk/tests/serve.rs b/sdk/tests/serve.rs index 8c656b2..52efafa 100644 --- a/sdk/tests/serve.rs +++ b/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| { + 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::>() + .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")