More

Author: ardatan

bin/router/src/lib.rs

@@ -19,7 +19,7 @@ use crate::{
     },
     jwt::JwtAuthRuntime,
     logger::configure_logging,
-    pipeline::graphql_request_handler,
+    pipeline::{graphql_request_handler, usage_reporting::create_hive_user_agent},
 };
 
 pub use crate::{schema_state::SchemaState, shared_state::RouterSharedState};
@@ -110,7 +110,10 @@ pub async fn configure_app_from_config(
         true => Some(JwtAuthRuntime::init(bg_tasks_manager, &router_config.jwt).await?),
         false => None,
     };
-    let usage_agent = pipeline::usage_reporting::from_config(&router_config).map(Arc::new);
+    let usage_agent = router_config
+        .usage_reporting
+        .as_ref()
+        .map(|usage_config| Arc::new(create_hive_user_agent(usage_config)));
 
     if let Some(usage_agent) = &usage_agent {
         bg_tasks_manager.register_task(usage_agent.clone());

bin/router/src/pipeline/mod.rs

@@ -192,23 +192,26 @@ pub async fn execute_pipeline(
     )
     .await?;
 
-    shared_state.usage_agent.as_ref().and_then(|usage_agent| {
-        shared_state
-            .router_config
-            .usage_reporting
-            .as_ref()
-            .map(|usage_config| {
-                usage_reporting::send_usage_report(
-                    supergraph.schema.clone(),
-                    start,
-                    req,
-                    &client_request_details,
-                    usage_agent,
-                    usage_config,
-                    &execution_result,
-                )
-            })
-    });
+    shared_state
+        .hive_usage_agent
+        .as_ref()
+        .and_then(|usage_agent| {
+            shared_state
+                .router_config
+                .usage_reporting
+                .as_ref()
+                .map(|usage_config| {
+                    usage_reporting::collect_usage_report(
+                        supergraph.supergraph_schema.clone(),
+                        start.elapsed(),
+                        req,
+                        &client_request_details,
+                        usage_agent,
+                        usage_config,
+                        &execution_result,
+                    )
+                })
+        });
 
     Ok(execution_result)
 }

bin/router/src/pipeline/usage_reporting.rs

@@ -1,41 +1,40 @@
 use std::{
     sync::Arc,
-    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
+    time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use async_trait::async_trait;
 use graphql_parser::schema::Document;
 use hive_console_sdk::agent::{ExecutionReport, UsageAgent};
-use hive_router_config::{usage_reporting::UsageReportingConfig, HiveRouterConfig};
+use hive_router_config::usage_reporting::UsageReportingConfig;
 use hive_router_plan_executor::execution::{
     client_request_details::ClientRequestDetails, plan::PlanExecutionOutput,
 };
 use ntex::web::HttpRequest;
 use rand::Rng;
 use tokio_util::sync::CancellationToken;
 
-use crate::background_tasks::BackgroundTask;
+use crate::{background_tasks::BackgroundTask, consts::ROUTER_VERSION};
 
-pub fn from_config(router_config: &HiveRouterConfig) -> Option<UsageAgent> {
-    router_config.usage_reporting.as_ref().map(|usage_config| {
-        let flush_interval = Duration::from_secs(usage_config.flush_interval);
-        hive_console_sdk::agent::UsageAgent::new(
-            usage_config.token.clone(),
-            usage_config.endpoint.clone(),
-            usage_config.target_id.clone(),
-            usage_config.buffer_size,
-            usage_config.connect_timeout,
-            usage_config.request_timeout,
-            usage_config.accept_invalid_certs,
-            flush_interval,
-            "hive-router".to_string(),
-        )
-    })
+pub fn create_hive_user_agent(usage_config: &UsageReportingConfig) -> UsageAgent {
+    let user_agent = format!("hive-router/{}", ROUTER_VERSION);
+    hive_console_sdk::agent::UsageAgent::new(
+        usage_config.access_token.clone(),
+        usage_config.endpoint.clone(),
+        usage_config.target_id.clone(),
+        usage_config.buffer_size,
+        usage_config.connect_timeout,
+        usage_config.request_timeout,
+        usage_config.accept_invalid_certs,
+        usage_config.flush_interval,
+        user_agent,
+    )
 }
 
-pub fn send_usage_report(
+#[inline]
+pub fn collect_usage_report(
     schema: Arc<Document<'static, String>>,
-    start: Instant,
+    duration: Duration,
     req: &HttpRequest,
     client_request_details: &ClientRequestDetails,
     usage_agent: &UsageAgent,
@@ -59,13 +58,11 @@ pub fn send_usage_report(
     let timestamp = SystemTime::now()
         .duration_since(UNIX_EPOCH)
         .unwrap()
-        .as_secs()
-        * 1000;
-    let duration = start.elapsed();
+        .as_millis() as u64;
     let execution_report = ExecutionReport {
         schema,
-        client_name,
-        client_version,
+        client_name: client_name.map(|s| s.to_owned()),
+        client_version: client_version.map(|s| s.to_owned()),
         timestamp,
         duration,
         ok: execution_result.error_count == 0,
@@ -77,22 +74,20 @@ pub fn send_usage_report(
             .map(|op_name| op_name.to_owned()),
         persisted_document_hash: None,
     };
-    usage_agent
-        .add_report(execution_report)
-        .unwrap_or_else(|err| tracing::error!("Failed to send usage report: {}", err));
+
+    if let Err(err) = usage_agent.add_report(execution_report) {
+        tracing::error!("Failed to send usage report: {}", err);
+    }
 }
 
-fn get_header_value(req: &HttpRequest, header_name: &str) -> Option<String> {
-    req.headers()
-        .get(header_name)
-        .and_then(|v| v.to_str().ok())
-        .map(|s| s.to_string())
+fn get_header_value<'req>(req: &'req HttpRequest, header_name: &str) -> Option<&'req str> {
+    req.headers().get(header_name).and_then(|v| v.to_str().ok())
 }
 
 #[async_trait]
 impl BackgroundTask for UsageAgent {
     fn id(&self) -> &str {
-        "usage_report_flush_interval"
+        "hive_console_usage_report_task"
     }
 
     async fn run(&self, token: CancellationToken) {

bin/router/src/schema_state.rs

@@ -40,7 +40,7 @@ pub struct SupergraphData {
     pub metadata: SchemaMetadata,
     pub planner: Planner,
     pub subgraph_executor_map: SubgraphExecutorMap,
-    pub schema: Arc<Document<'static, String>>,
+    pub supergraph_schema: Arc<Document<'static, String>>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -126,7 +126,7 @@ impl SchemaState {
         )?;
 
         Ok(SupergraphData {
-            schema: Arc::new(parsed_supergraph_sdl),
+            supergraph_schema: Arc::new(parsed_supergraph_sdl),
             metadata,
             planner,
             subgraph_executor_map,

bin/router/src/shared_state.rs

@@ -19,7 +19,7 @@ pub struct RouterSharedState {
     pub override_labels_evaluator: OverrideLabelsEvaluator,
     pub cors_runtime: Option<Cors>,
     pub jwt_auth_runtime: Option<JwtAuthRuntime>,
-    pub usage_agent: Option<Arc<UsageAgent>>,
+    pub hive_usage_agent: Option<Arc<UsageAgent>>,
 }
 
 impl RouterSharedState {
@@ -39,7 +39,7 @@ impl RouterSharedState {
             )
             .map_err(Box::new)?,
             jwt_auth_runtime,
-            usage_agent,
+            hive_usage_agent: usage_agent,
         })
     }
 }
@@ -52,6 +52,6 @@ pub enum SharedStateError {
     CORSConfig(#[from] Box<CORSConfigError>),
     #[error("invalid override labels config: {0}")]
     OverrideLabelsCompile(#[from] Box<OverrideLabelsCompileError>),
-    #[error("error creating usage agent: {0}")]
+    #[error("error creating hive usage agent: {0}")]
     UsageAgent(#[from] Box<hive_console_sdk::agent::AgentError>),
 }

lib/router-config/src/usage_reporting.rs

@@ -1,11 +1,13 @@
+use std::time::Duration;
+
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct UsageReportingConfig {
     /// Your [Registry Access Token](https://the-guild.dev/graphql/hive/docs/management/targets#registry-access-tokens) with write permission.
-    pub token: String,
+    pub access_token: String,
     /// A target ID, this can either be a slug following the format “$organizationSlug/$projectSlug/$targetSlug” (e.g “the-guild/graphql-hive/staging”) or an UUID (e.g. “a0f4c605-6541-4350-8cfe-b31f21a4bf80”). To be used when the token is configured with an organization access token.
     pub target_id: Option<String>,
     /// For self-hosting, you can override `/usage` endpoint (defaults to `https://app.graphql-hive.com/usage`).
@@ -45,8 +47,13 @@ pub struct UsageReportingConfig {
     pub request_timeout: u64,
     /// Frequency of flushing the buffer to the server
     /// Default: 5 seconds
-    #[serde(default = "default_flush_interval")]
-    pub flush_interval: u64,
+    #[serde(
+        default = "default_flush_interval",
+        deserialize_with = "humantime_serde::deserialize",
+        serialize_with = "humantime_serde::serialize"
+    )]
+    #[schemars(with = "String")]
+    pub flush_interval: Duration,
 }
 
 fn default_endpoint() -> String {
@@ -81,6 +88,6 @@ fn default_connect_timeout() -> u64 {
     5
 }
 
-fn default_flush_interval() -> u64 {
-    5
+fn default_flush_interval() -> Duration {
+    Duration::from_secs(5)
 }