From fec7a6bfb5fdce691d1008e3cf60578aba503a42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonasz=20=C5=81asut-Balcerzak?= Date: Wed, 26 Aug 2026 00:07:16 +0200 Subject: [PATCH] feat!(sdk): replace serve_service with a builder serve_customized(args) starts a ServerBuilder composing the server from parts, in the spirit of function-sdk-go's options: function() adds the generated typed server, service() any tonic service (a custom codec, an instrumented wrapper, services beside the function), file_descriptor_set() their reflection descriptors, and health_reporter() hands readiness to the caller (services otherwise report serving when serving starts). serve delegates to it; the serve_service function from #1 is removed before any release ships it. Co-Authored-By: Claude Fable 5 --- sdk/README.md | 6 +- sdk/src/lib.rs | 2 +- sdk/src/server.rs | 218 +++++++++++++++++++++++++++++++++------------ sdk/tests/serve.rs | 77 +++++++++++++++- 4 files changed, 241 insertions(+), 62 deletions(-) diff --git a/sdk/README.md b/sdk/README.md index 75f0502..131e137 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -12,8 +12,10 @@ API. This crate provides: - `server` - a function-spec-compliant runtime: mTLS or insecure serving, 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. + `serve_customized` builds the server from parts - the typed function or + any tonic service (a custom codec, an instrumented wrapper), further + services, reflection descriptors, and a health reporter the caller can + own to flip readiness 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..9078e4b 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, ServerBuilder, serve, serve_customized}; /// Errors returned by the SDK. #[derive(Debug, thiserror::Error)] diff --git a/sdk/src/server.rs b/sdk/src/server.rs index 3a99ef2..af466a2 100644 --- a/sdk/src/server.rs +++ b/sdk/src/server.rs @@ -7,7 +7,9 @@ use tonic::body::Body; use tonic::codegen::Service; use tonic::codegen::http; use tonic::server::NamedService; +use tonic::service::Routes; use tonic::transport::{Certificate, Identity, Server, ServerTlsConfig}; +pub use tonic_health::server::HealthReporter; use crate::Error; use crate::proto::v1::function_runner_service_server::{ @@ -48,72 +50,176 @@ pub struct Args { /// to authenticate Crossplane) unless `--insecure` is set. gRPC server /// reflection is enabled for both the v1 and v1alpha reflection APIs, and /// the gRPC health service reports the function as serving. +/// +/// This is the default server; [`serve_customized`] builds one from parts. pub async fn serve(function: F, args: &Args) -> Result<(), Error> { - let mut function_service = FunctionRunnerServiceServer::new(function); - if let Some(size) = args.max_recv_message_size { - function_service = function_service.max_decoding_message_size(size); - } - serve_service(function_service, args).await + serve_customized(args).function(function).serve().await } -/// Like [`serve`], but serves any gRPC service in place of the generated -/// FunctionRunnerService server, with the same spec-compliant transport: -/// mTLS from `--tls-certs-dir` unless `--insecure`, v1 and v1alpha gRPC -/// server reflection, the gRPC health service reporting the service (by its -/// [`NamedService::NAME`]) as serving, and graceful shutdown on SIGTERM or -/// SIGINT. +/// Starts building a customized function server: the same spec-compliant +/// transport as [`serve`] (mTLS from `--tls-certs-dir` unless `--insecure`, +/// v1 and v1alpha gRPC server reflection, the gRPC health service, graceful +/// shutdown on SIGTERM or SIGINT), composed from the components you add. /// -/// Use it when the generated server's prost codec is not enough - a runtime -/// that must forward request bytes verbatim (prost drops fields newer than -/// the generated types, so a transparent proxy needs its own codec), or a -/// service wrapped for instrumentation. The service owns its own message -/// 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 address: SocketAddr = args.address.parse()?; - - let mut builder = Server::builder(); - if !args.insecure { - let dir = args - .tls_certs_dir - .as_deref() - .ok_or(Error::MissingTlsCertsDir)?; - builder = builder.tls_config(tls_config(dir)?)?; +/// Add the function itself with [`ServerBuilder::function`] (the typed +/// service [`serve`] runs) or [`ServerBuilder::service`] (any tonic +/// service - a custom codec, an instrumented wrapper), further gRPC +/// services with more [`ServerBuilder::service`] calls, and reflection +/// descriptors for them with [`ServerBuilder::file_descriptor_set`]. Take +/// the health reporter with [`ServerBuilder::health_reporter`] to own +/// readiness - for example, to flip to serving only after start-up work. +/// Finish with [`ServerBuilder::serve`]. +/// +/// ```no_run +/// # use function_sdk_rust::proto::v1::function_runner_service_server::{ +/// # FunctionRunnerService, FunctionRunnerServiceServer, +/// # }; +/// # async fn example( +/// # function: F, +/// # args: &function_sdk_rust::Args, +/// # ) -> Result<(), function_sdk_rust::Error> { +/// let mut builder = function_sdk_rust::serve_customized(args) +/// .function(function); +/// let health = builder.health_reporter(); +/// health.set_not_serving::>().await; +/// tokio::spawn({ +/// let health = health.clone(); +/// async move { +/// // ... warm caches ... +/// health.set_serving::>().await; +/// } +/// }); +/// builder.serve().await +/// # } +/// ``` +pub fn serve_customized(args: &Args) -> ServerBuilder { + let (health_reporter, health_service) = tonic_health::server::health_reporter(); + ServerBuilder { + address: args.address.clone(), + insecure: args.insecure, + tls_certs_dir: args.tls_certs_dir.clone(), + max_recv_message_size: args.max_recv_message_size, + routes: Routes::default().add_service(health_service), + service_names: Vec::new(), + descriptor_sets: vec![crate::proto::FILE_DESCRIPTOR_SET], + health_reporter, + health_managed: true, } +} - let reflection_v1 = tonic_reflection::server::Builder::configure() - .register_encoded_file_descriptor_set(crate::proto::FILE_DESCRIPTOR_SET) - .build_v1()?; - let reflection_v1alpha = tonic_reflection::server::Builder::configure() - .register_encoded_file_descriptor_set(crate::proto::FILE_DESCRIPTOR_SET) - .build_v1alpha()?; +/// A function server under construction; see [`serve_customized`]. +pub struct ServerBuilder { + address: String, + insecure: bool, + tls_certs_dir: Option, + max_recv_message_size: Option, + routes: Routes, + service_names: Vec<&'static str>, + descriptor_sets: Vec<&'static [u8]>, + health_reporter: HealthReporter, + health_managed: bool, +} - let (health_reporter, health_service) = tonic_health::server::health_reporter(); - health_reporter.set_serving::().await; +impl ServerBuilder { + /// Adds the function as the generated typed FunctionRunnerService + /// server, honoring `--max-recv-message-size` - the service [`serve`] + /// runs. + pub fn function(self, function: F) -> Self { + let mut service = FunctionRunnerServiceServer::new(function); + if let Some(size) = self.max_recv_message_size { + service = service.max_decoding_message_size(size); + } + self.service(service) + } + + /// Adds any gRPC service. Use it for the function itself when the + /// generated server's prost codec is not enough (prost drops fields + /// newer than the generated types, so a transparent proxy needs its own + /// codec - which then owns its message-size limits; + /// [`Args::max_recv_message_size`] carries what the caller asked for), + /// or for further services beside the function. Each added service is + /// registered with the health service under its + /// [`NamedService::NAME`]. + pub fn service(mut self, service: S) -> Self + where + S: Service< + http::Request, + Response = http::Response, + Error = std::convert::Infallible, + > + NamedService + + Clone + + Send + + Sync + + 'static, + S::Future: Send + 'static, + { + self.routes = self.routes.add_service(service); + self.service_names.push(S::NAME); + self + } - tracing::info!(%address, insecure = args.insecure, service = S::NAME, "serving"); + /// Registers another encoded file descriptor set with gRPC server + /// reflection, next to this crate's own - for services added beside the + /// function. + pub fn file_descriptor_set(mut self, encoded: &'static [u8]) -> Self { + self.descriptor_sets.push(encoded); + self + } - builder - .add_service(service) - .add_service(health_service) - .add_service(reflection_v1) - .add_service(reflection_v1alpha) - .serve_with_shutdown(address, shutdown_signal()) - .await?; + /// Hands out the gRPC health reporter and makes health the caller's: + /// [`serve`](ServerBuilder::serve) then reports nothing by itself, so + /// set each service's initial status (typically + /// [`HealthReporter::set_not_serving`] before start-up work, flipped + /// with [`HealthReporter::set_serving`] when ready). Without this call, + /// every added service is reported as serving when serving starts. + pub fn health_reporter(&mut self) -> HealthReporter { + self.health_managed = false; + self.health_reporter.clone() + } - Ok(()) + /// Starts the server and serves until SIGTERM or SIGINT, then shuts + /// down gracefully. Nothing listens before this. + pub async fn serve(self) -> Result<(), Error> { + let address: SocketAddr = self.address.parse()?; + + let mut builder = Server::builder(); + if !self.insecure { + let dir = self + .tls_certs_dir + .as_deref() + .ok_or(Error::MissingTlsCertsDir)?; + builder = builder.tls_config(tls_config(dir)?)?; + } + + let reflection = || { + let mut builder = tonic_reflection::server::Builder::configure(); + for encoded in &self.descriptor_sets { + builder = builder.register_encoded_file_descriptor_set(encoded); + } + builder + }; + let reflection_v1 = reflection().build_v1()?; + let reflection_v1alpha = reflection().build_v1alpha()?; + + if self.health_managed { + for name in &self.service_names { + self.health_reporter + .set_service_status(name, tonic_health::ServingStatus::Serving) + .await; + } + } + + tracing::info!(%address, insecure = self.insecure, services = ?self.service_names, "serving"); + + builder + .add_routes(self.routes) + .add_service(reflection_v1) + .add_service(reflection_v1alpha) + .serve_with_shutdown(address, shutdown_signal()) + .await?; + + Ok(()) + } } fn tls_config(dir: &Path) -> Result { diff --git a/sdk/tests/serve.rs b/sdk/tests/serve.rs index 8c656b2..a2b37cd 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_customized}; use tonic::{Request, Response, Status}; struct EchoFunction; @@ -107,7 +107,7 @@ impl tonic::codegen::Service> f } #[tokio::test] -async fn serve_service_serves_a_custom_service() { +async fn serve_customized_serves_a_custom_service() { let port = free_port(); tokio::spawn(async move { let args = Args { @@ -118,7 +118,9 @@ async fn serve_service_serves_a_custom_service() { max_recv_message_size: None, }; let service = CustomService(FunctionRunnerServiceServer::new(EchoFunction)); - serve_service(service, &args) + serve_customized(&args) + .service(service) + .serve() .await .expect("serve must start"); }); @@ -160,6 +162,75 @@ async fn serve_service_serves_a_custom_service() { ); } +#[tokio::test] +async fn serve_customized_hands_health_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 mut builder = serve_customized(&args).function(EchoFunction); + let health = builder.health_reporter(); + health + .set_not_serving::>() + .await; + tokio::spawn(async move { + builder.serve().await.expect("serve must start"); + }); + + 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")