From 1318cd9b456a9563570b46f8c96059f2830be985 Mon Sep 17 00:00:00 2001
From: Till Rohrmann <till@restate.dev>
Date: Fri, 13 Dec 2024 11:11:26 +0100
Subject: [PATCH] Extend NetworkServerBuilder to register axum routes

---
 crates/core/src/network/server_builder.rs          | 9 +++++++--
 crates/metadata-store/src/local/tests.rs           | 2 +-
 crates/node/src/network_server/service.rs          | 5 ++++-
 tools/restatectl/src/environment/metadata_store.rs | 2 +-
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/crates/core/src/network/server_builder.rs b/crates/core/src/network/server_builder.rs
index f9405a8fe..a9ced76d6 100644
--- a/crates/core/src/network/server_builder.rs
+++ b/crates/core/src/network/server_builder.rs
@@ -30,6 +30,7 @@ use super::net_util::run_hyper_server;
 pub struct NetworkServerBuilder {
     grpc_descriptors: Vec<&'static [u8]>,
     grpc_routes: Option<Routes>,
+    axum_router: Option<axum::routing::Router>,
 }
 
 impl NetworkServerBuilder {
@@ -56,10 +57,13 @@ impl NetworkServerBuilder {
         self
     }
 
+    pub fn register_axum_routes(&mut self, routes: impl Into<axum::routing::Router>) {
+        self.axum_router = Some(self.axum_router.take().unwrap_or_default().merge(routes));
+    }
+
     pub async fn run(
         self,
         node_rpc_health: HealthStatus<NodeRpcStatus>,
-        axum_router: Option<axum::routing::Router>,
         bind_address: &BindAddress,
     ) -> Result<(), anyhow::Error> {
         node_rpc_health.update(NodeRpcStatus::StartingUp);
@@ -68,7 +72,8 @@ impl NetworkServerBuilder {
             .include_headers(true)
             .level(tracing::Level::ERROR);
 
-        let axum_router = axum_router
+        let axum_router = self
+            .axum_router
             .unwrap_or_default()
             .layer(TraceLayer::new_for_http().make_span_with(span_factory.clone()))
             .fallback(handler_404);
diff --git a/crates/metadata-store/src/local/tests.rs b/crates/metadata-store/src/local/tests.rs
index 26b70b805..987d77215 100644
--- a/crates/metadata-store/src/local/tests.rs
+++ b/crates/metadata-store/src/local/tests.rs
@@ -322,7 +322,7 @@ async fn start_metadata_store(
         let rpc_server_health_status = rpc_server_health_status.clone();
         async move {
             server_builder
-                .run(rpc_server_health_status, None, &bind_address)
+                .run(rpc_server_health_status, &bind_address)
                 .await
         }
     })?;
diff --git a/crates/node/src/network_server/service.rs b/crates/node/src/network_server/service.rs
index fd45720b8..3081e928f 100644
--- a/crates/node/src/network_server/service.rs
+++ b/crates/node/src/network_server/service.rs
@@ -92,7 +92,10 @@ impl NetworkServer {
             )
             .with_state(shared_state);
 
+        server_builder.register_axum_routes(axum_router);
+
         let node_health = health.node_status();
+        let node_rpc_health = health.node_rpc_status();
 
         server_builder.register_grpc_service(
             NodeCtlSvcServer::new(NodeCtlSvcHandler::new(
@@ -125,7 +128,7 @@ impl NetworkServer {
         );
 
         server_builder
-            .run(node_health, axum_router, &options.bind_address.unwrap())
+            .run(node_rpc_health, &options.bind_address.unwrap())
             .await?;
 
         Ok(())
diff --git a/tools/restatectl/src/environment/metadata_store.rs b/tools/restatectl/src/environment/metadata_store.rs
index eef88c27a..5abdb204e 100644
--- a/tools/restatectl/src/environment/metadata_store.rs
+++ b/tools/restatectl/src/environment/metadata_store.rs
@@ -48,7 +48,7 @@ pub async fn start_metadata_store(
         let rpc_server_health_status = rpc_server_health_status.clone();
         async move {
             server_builder
-                .run(rpc_server_health_status, None, &bind_address)
+                .run(rpc_server_health_status, &bind_address)
                 .await
         }
     })?;