Extend NetworkServerBuilder to register axum routes

tillrohrmann · Dec 16, 2024 · 1318cd9 · 1318cd9
1 parent fe189e6
commit 1318cd9
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
         }
     })?;