feat(sdk): serve any named gRPC service

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

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

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<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
+}
+
+/// 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<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();
@@ -56,11 +93,6 @@ pub async fn serve<F: FunctionRunnerService>(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<F: FunctionRunnerService>(function: F, args: &Args) -> Result
         .build_v1alpha()?;
 
     let (health_reporter, health_service) = tonic_health::server::health_reporter();
-    health_reporter
-        .set_serving::<FunctionRunnerServiceServer<F>>()
-        .await;
+    health_reporter.set_serving::<S>().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)

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<EchoFunction>);
+
+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 =
+        <FunctionRunnerServiceServer<EchoFunction> as tonic::server::NamedService>::NAME;
+}
+
+impl tonic::codegen::Service<tonic::codegen::http::Request<tonic::body::Body>> for CustomService {
+    type Response = tonic::codegen::http::Response<tonic::body::Body>;
+    type Error = std::convert::Infallible;
+    type Future = <FunctionRunnerServiceServer<EchoFunction> as tonic::codegen::Service<
+        tonic::codegen::http::Request<tonic::body::Body>,
+    >>::Future;
+
+    fn poll_ready(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        <FunctionRunnerServiceServer<EchoFunction> as tonic::codegen::Service<
+            tonic::codegen::http::Request<tonic::body::Body>,
+        >>::poll_ready(&mut self.0, cx)
+    }
+
+    fn call(&mut self, req: tonic::codegen::http::Request<tonic::body::Body>) -> Self::Future {
+        <FunctionRunnerServiceServer<EchoFunction> as tonic::codegen::Service<
+            tonic::codegen::http::Request<tonic::body::Body>,
+        >>::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")