match patterns

Author: ardatan

bin/router/src/lib.rs

@@ -19,7 +19,7 @@ use crate::{
     },
     jwt::JwtAuthRuntime,
     logger::configure_logging,
-    pipeline::{graphql_request_handler, usage_reporting::create_hive_user_agent},
+    pipeline::{graphql_request_handler, usage_reporting::init_hive_user_agent},
 };
 
 pub use crate::{schema_state::SchemaState, shared_state::RouterSharedState};
@@ -111,26 +111,22 @@ pub async fn configure_app_from_config(
         false => None,
     };
 
-    let usage_agent = if router_config.usage_reporting.enabled {
-        Some(Arc::new(create_hive_user_agent(
+    let hive_usage_agent = match router_config.usage_reporting.enabled {
+        true => Some(init_hive_user_agent(
+            bg_tasks_manager,
             &router_config.usage_reporting,
-        )))
-    } else {
-        None
+        )),
+        false => None,
     };
 
-    if let Some(usage_agent) = &usage_agent {
-        bg_tasks_manager.register_task(usage_agent.clone());
-    }
-
     let router_config_arc = Arc::new(router_config);
     let schema_state =
         SchemaState::new_from_config(bg_tasks_manager, router_config_arc.clone()).await?;
     let schema_state_arc = Arc::new(schema_state);
     let shared_state = Arc::new(RouterSharedState::new(
         router_config_arc,
         jwt_runtime,
-        usage_agent,
+        hive_usage_agent,
     )?);
 
     Ok((shared_state, schema_state_arc))

bin/router/src/pipeline/mod.rs

@@ -193,13 +193,13 @@ pub async fn execute_pipeline(
     .await?;
 
     if shared_state.router_config.usage_reporting.enabled {
-        if let Some(usage_agent) = &shared_state.hive_usage_agent {
+        if let Some(hive_usage_agent) = &shared_state.hive_usage_agent {
             usage_reporting::collect_usage_report(
                 supergraph.supergraph_schema.clone(),
                 start.elapsed(),
                 req,
                 &client_request_details,
-                usage_agent,
+                hive_usage_agent,
                 &shared_state.router_config.usage_reporting,
                 &execution_result,
             );

bin/router/src/pipeline/usage_reporting.rs

@@ -14,11 +14,17 @@ use ntex::web::HttpRequest;
 use rand::Rng;
 use tokio_util::sync::CancellationToken;
 
-use crate::{background_tasks::BackgroundTask, consts::ROUTER_VERSION};
+use crate::{
+    background_tasks::{BackgroundTask, BackgroundTasksManager},
+    consts::ROUTER_VERSION,
+};
 
-pub fn create_hive_user_agent(usage_config: &UsageReportingConfig) -> UsageAgent {
+pub fn init_hive_user_agent(
+    bg_tasks_manager: &mut BackgroundTasksManager,
+    usage_config: &UsageReportingConfig,
+) -> Arc<UsageAgent> {
     let user_agent = format!("hive-router/{}", ROUTER_VERSION);
-    hive_console_sdk::agent::UsageAgent::new(
+    let hive_user_agent = hive_console_sdk::agent::UsageAgent::new(
         usage_config.access_token.clone(),
         usage_config.endpoint.clone(),
         usage_config.target_id.clone(),
@@ -28,7 +34,10 @@ pub fn create_hive_user_agent(usage_config: &UsageReportingConfig) -> UsageAgent
         usage_config.accept_invalid_certs,
         usage_config.flush_interval,
         user_agent,
-    )
+    );
+    let hive_user_agent_arc = Arc::new(hive_user_agent);
+    bg_tasks_manager.register_task(hive_user_agent_arc.clone());
+    hive_user_agent_arc
 }
 
 #[inline]
@@ -37,7 +46,7 @@ pub fn collect_usage_report(
     duration: Duration,
     req: &HttpRequest,
     client_request_details: &ClientRequestDetails,
-    usage_agent: &UsageAgent,
+    hive_usage_agent: &UsageAgent,
     usage_config: &UsageReportingConfig,
     execution_result: &PlanExecutionOutput,
 ) {
@@ -75,7 +84,7 @@ pub fn collect_usage_report(
         persisted_document_hash: None,
     };
 
-    if let Err(err) = usage_agent.add_report(execution_report) {
+    if let Err(err) = hive_usage_agent.add_report(execution_report) {
         tracing::error!("Failed to send usage report: {}", err);
     }
 }

bin/router/src/shared_state.rs

@@ -26,7 +26,7 @@ impl RouterSharedState {
     pub fn new(
         router_config: Arc<HiveRouterConfig>,
         jwt_auth_runtime: Option<JwtAuthRuntime>,
-        usage_agent: Option<Arc<UsageAgent>>,
+        hive_usage_agent: Option<Arc<UsageAgent>>,
     ) -> Result<Self, SharedStateError> {
         Ok(Self {
             validation_plan: graphql_tools::validation::rules::default_rules_validation_plan(),
@@ -39,7 +39,7 @@ impl RouterSharedState {
             )
             .map_err(Box::new)?,
             jwt_auth_runtime,
-            hive_usage_agent: usage_agent,
+            hive_usage_agent,
         })
     }
 }