From e1d1926824afbd5452f330b0dd855a1abdcb1893 Mon Sep 17 00:00:00 2001
From: Jack Kleeman <jackkleeman@gmail.com>
Date: Thu, 19 Dec 2024 13:49:28 +0000
Subject: [PATCH] Install prometheus recorder earlier so kafka metrics work

---
 crates/node/src/lib.rs                    | 10 +++++++++-
 crates/node/src/network_server/mod.rs     |  1 +
 crates/node/src/network_server/service.rs | 11 ++++++-----
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/crates/node/src/lib.rs b/crates/node/src/lib.rs
index 3b6cd49f8f..d83e51ca3b 100644
--- a/crates/node/src/lib.rs
+++ b/crates/node/src/lib.rs
@@ -280,6 +280,14 @@ impl Node {
 
         let config = self.updateable_config.pinned();
 
+        let prometheus_handle = if config.common.disable_prometheus {
+            None
+        } else {
+            Some(network_server::install_global_prometheus_recorder(
+                &config.common,
+            ))
+        };
+
         if let Some(metadata_store) = self.metadata_store_role {
             tc.spawn(
                 TaskKind::MetadataStore,
@@ -431,7 +439,7 @@ impl Node {
             TaskKind::RpcServer,
             "node-rpc-server",
             None,
-            self.server.run(config.common.clone()),
+            self.server.run(config.common.clone(), prometheus_handle),
         )?;
 
         Ok(())
diff --git a/crates/node/src/network_server/mod.rs b/crates/node/src/network_server/mod.rs
index a79efc4f58..6a836725f4 100644
--- a/crates/node/src/network_server/mod.rs
+++ b/crates/node/src/network_server/mod.rs
@@ -15,4 +15,5 @@ mod prometheus_helpers;
 mod service;
 mod state;
 
+pub(crate) use metrics::install_global_prometheus_recorder;
 pub use service::{AdminDependencies, NetworkServer, WorkerDependencies};
diff --git a/crates/node/src/network_server/service.rs b/crates/node/src/network_server/service.rs
index 6817391eee..65df99cc21 100644
--- a/crates/node/src/network_server/service.rs
+++ b/crates/node/src/network_server/service.rs
@@ -34,7 +34,6 @@ use restate_types::config::CommonOptions;
 use crate::network_server::handler;
 use crate::network_server::handler::cluster_ctrl::ClusterCtrlSvcHandler;
 use crate::network_server::handler::node::NodeSvcHandler;
-use crate::network_server::metrics::install_global_prometheus_recorder;
 use crate::network_server::multiplex::MultiplexService;
 use crate::network_server::state::NodeCtrlHandlerStateBuilder;
 
@@ -57,15 +56,17 @@ impl NetworkServer {
         }
     }
 
-    pub async fn run(self, options: CommonOptions) -> Result<(), anyhow::Error> {
+    pub async fn run(
+        self,
+        options: CommonOptions,
+        prometheus_handle: Option<metrics_exporter_prometheus::PrometheusHandle>,
+    ) -> Result<(), anyhow::Error> {
         let tc = task_center();
         // Configure Metric Exporter
         let mut state_builder = NodeCtrlHandlerStateBuilder::default();
         state_builder.task_center(tc.clone());
 
-        if !options.disable_prometheus {
-            let prometheus_handle = install_global_prometheus_recorder(&options);
-
+        if let Some(prometheus_handle) = prometheus_handle {
             tc.spawn_child(
                 TaskKind::SystemService,
                 "prometheus-metrics-upkeep",