feat!(sdk): replace serve_service with a builder

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

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)]

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<F: FunctionRunnerService>(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<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 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<F: FunctionRunnerService>(
+/// #     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::<FunctionRunnerServiceServer<F>>().await;
+/// tokio::spawn({
+///     let health = health.clone();
+///     async move {
+///         // ... warm caches ...
+///         health.set_serving::<FunctionRunnerServiceServer<F>>().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<PathBuf>,
+    max_recv_message_size: Option<usize>,
+    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::<S>().await;
+impl ServerBuilder {
+    /// Adds the function as the generated typed FunctionRunnerService
+    /// server, honoring `--max-recv-message-size` - the service [`serve`]
+    /// runs.
+    pub fn function<F: FunctionRunnerService>(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<S>(mut self, service: S) -> Self
+    where
+        S: Service<
+                http::Request<Body>,
+                Response = http::Response<Body>,
+                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<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_customized};
 use tonic::{Request, Response, Status};
 
 struct EchoFunction;
@@ -107,7 +107,7 @@ impl tonic::codegen::Service<tonic::codegen::http::Request<tonic::body::Body>> 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::<FunctionRunnerServiceServer<EchoFunction>>()
+        .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<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")