diff --git a/sdk/README.md b/sdk/README.md index 8d62a83..75f0502 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -11,7 +11,9 @@ API. This crate provides: the tonic gRPC server and client. - `server` - a function-spec-compliant runtime: mTLS or insecure serving, gRPC server reflection, the gRPC health service, graceful shutdown, and - the standard CLI arguments. + 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. - `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 82e6d7c..75f1eb8 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}; +pub use server::{Args, serve, serve_service}; /// Errors returned by the SDK. #[derive(Debug, thiserror::Error)] diff --git a/sdk/src/server.rs b/sdk/src/server.rs index 52f83ee..3a99ef2 100644 --- a/sdk/src/server.rs +++ b/sdk/src/server.rs @@ -3,6 +3,10 @@ use std::net::SocketAddr; use std::path::{Path, PathBuf}; +use tonic::body::Body; +use tonic::codegen::Service; +use tonic::codegen::http; +use tonic::server::NamedService; use tonic::transport::{Certificate, Identity, Server, ServerTlsConfig}; use crate::Error; @@ -45,6 +49,39 @@ pub struct Args { /// reflection is enabled for both the v1 and v1alpha reflection APIs, and /// the gRPC health service reports the function as serving. 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 +} + +/// 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. +/// +/// 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(); @@ -56,11 +93,6 @@ pub async fn serve(function: F, args: &Args) -> Result builder = builder.tls_config(tls_config(dir)?)?; } - 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); - } - let reflection_v1 = tonic_reflection::server::Builder::configure() .register_encoded_file_descriptor_set(crate::proto::FILE_DESCRIPTOR_SET) .build_v1()?; @@ -69,14 +101,12 @@ pub async fn serve(function: F, args: &Args) -> Result .build_v1alpha()?; let (health_reporter, health_service) = tonic_health::server::health_reporter(); - health_reporter - .set_serving::>() - .await; + health_reporter.set_serving::().await; - tracing::info!(%address, insecure = args.insecure, "serving FunctionRunnerService"); + tracing::info!(%address, insecure = args.insecure, service = S::NAME, "serving"); builder - .add_service(function_service) + .add_service(service) .add_service(health_service) .add_service(reflection_v1) .add_service(reflection_v1alpha) diff --git a/sdk/tests/serve.rs b/sdk/tests/serve.rs index 45df9e1..8c656b2 100644 --- a/sdk/tests/serve.rs +++ b/sdk/tests/serve.rs @@ -1,7 +1,8 @@ use function_sdk_rust::proto::v1::function_runner_service_client::FunctionRunnerServiceClient; 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}; +use function_sdk_rust::{Args, response, serve, serve_service}; use tonic::{Request, Response, Status}; struct EchoFunction; @@ -69,6 +70,96 @@ async fn serve_handles_a_run_function_request() { ); } +/// A hand-rolled tonic service - what serve_service exists for. It +/// delegates to the generated server, but as far as serve_service can tell +/// it is an arbitrary Service + NamedService implementation. +#[derive(Clone)] +struct CustomService(FunctionRunnerServiceServer); + +impl tonic::server::NamedService for CustomService { + // NAME both routes requests and keys the health service, so a service + // that stands in for the FunctionRunnerService keeps its name. + const NAME: &'static str = + as tonic::server::NamedService>::NAME; +} + +impl tonic::codegen::Service> for CustomService { + type Response = tonic::codegen::http::Response; + type Error = std::convert::Infallible; + type Future = as tonic::codegen::Service< + tonic::codegen::http::Request, + >>::Future; + + fn poll_ready( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + as tonic::codegen::Service< + tonic::codegen::http::Request, + >>::poll_ready(&mut self.0, cx) + } + + fn call(&mut self, req: tonic::codegen::http::Request) -> Self::Future { + as tonic::codegen::Service< + tonic::codegen::http::Request, + >>::call(&mut self.0, req) + } +} + +#[tokio::test] +async fn serve_service_serves_a_custom_service() { + let port = free_port(); + tokio::spawn(async move { + 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 = CustomService(FunctionRunnerServiceServer::new(EchoFunction)); + serve_service(service, &args) + .await + .expect("serve must start"); + }); + + let mut client = connect_with_retry(port).await; + + let rsp = client + .run_function(RunFunctionRequest { + meta: Some(RequestMeta { + tag: "custom".to_string(), + capabilities: vec![], + }), + ..Default::default() + }) + .await + .expect("RunFunction must succeed") + .into_inner(); + + assert_eq!(rsp.meta.unwrap().tag, "custom"); + assert_eq!(rsp.results[0].message, "echo"); + + // Health is keyed by S::NAME. + 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 = tonic_health::pb::health_client::HealthClient::new(channel); + let status = health + .check(tonic_health::pb::HealthCheckRequest { + service: "apiextensions.fn.proto.v1.FunctionRunnerService".to_string(), + }) + .await + .expect("health check must succeed") + .into_inner(); + assert_eq!( + status.status, + 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")