feat(router): Hive Console Usage Reporting

Author: ardatan

Cargo.lock

@@ -26,8 +26,10 @@ lto = true
 codegen-units = 1
 
 [workspace.dependencies]
-graphql-tools = "0.4.0"
-graphql-parser = "0.4.1"
+graphql-parser = { version = "0.5.0", package = "graphql-parser-hive-fork" }
+graphql-tools = { version = "0.4.0", features = [
+    "graphql_parser_fork",
+], default-features = false }
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.142"
 sonic-rs = "0.5.3"

Cargo.toml

@@ -19,6 +19,7 @@ path = "src/main.rs"
 hive-router-query-planner = { path = "../../lib/query-planner", version = "2.0.0" }
 hive-router-plan-executor = { path = "../../lib/executor", version = "4.0.0" }
 hive-router-config = { path = "../../lib/router-config", version = "0.0.7" }
+hive-console-sdk = { path = "../../../graphql-hive/packages/libraries/sdk-rs" }
 
 tokio = { workspace = true }
 futures = { workspace = true }

bin/router/Cargo.toml

@@ -8,7 +8,7 @@ mod schema_state;
 mod shared_state;
 mod supergraph;
 
-use std::sync::Arc;
+use std::{sync::Arc};
 
 use crate::{
     background_tasks::BackgroundTasksManager,
@@ -110,12 +110,17 @@ pub async fn configure_app_from_config(
         Some(jwt_config) => Some(JwtAuthRuntime::init(bg_tasks_manager, jwt_config).await?),
         None => None,
     };
+    let usage_agent = pipeline::usage::from_config(&router_config).map(Arc::new);
+    
+    if let Some(usage_agent) = usage_agent.clone() {
+        bg_tasks_manager.register_task(usage_agent);
+    }
 
     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)?);
+    let shared_state = Arc::new(RouterSharedState::new(router_config_arc, jwt_runtime, usage_agent)?);
 
     Ok((shared_state, schema_state_arc))
 }

bin/router/src/lib.rs

@@ -1,4 +1,4 @@
-use std::{borrow::Cow, sync::Arc};
+use std::{borrow::Cow, sync::Arc, time::Instant};
 
 use hive_router_plan_executor::execution::plan::PlanExecutionOutput;
 use hive_router_query_planner::utils::cancellation::CancellationToken;
@@ -40,6 +40,7 @@ pub mod normalize;
 pub mod parser;
 pub mod progressive_override;
 pub mod query_plan;
+pub mod usage;
 pub mod validation;
 
 static GRAPHIQL_HTML: &str = include_str!("../../static/graphiql.html");
@@ -100,6 +101,7 @@ pub async fn execute_pipeline(
     shared_state: &Arc<RouterSharedState>,
     schema_state: &Arc<SchemaState>,
 ) -> Result<PlanExecutionOutput, PipelineError> {
+    let start = Instant::now();
     perform_csrf_prevention(req, &shared_state.router_config.csrf)?;
 
     let execution_request = get_execution_request(req, body_bytes).await?;
@@ -117,6 +119,7 @@ pub async fn execute_pipeline(
     )
     .await?;
     let query = Cow::Owned(execution_request.query.clone());
+    let query_clone: Cow<'_, str> = Cow::Owned(execution_request.query.clone());
     let variable_payload =
         coerce_request_variables(req, supergraph, execution_request, &normalize_payload)?;
 
@@ -144,5 +147,22 @@ pub async fn execute_pipeline(
     )
     .await?;
 
+    let operation_name = normalize_payload.operation_for_plan.name.as_ref();
+
+    shared_state.usage_agent.as_ref().and_then(|usage_agent| {
+        shared_state.router_config.usage.as_ref().map(|usage_config| {
+            usage::send_usage_report(
+                supergraph.schema.clone(),
+                start,
+                req,
+                operation_name,
+                &query_clone,
+                usage_agent,
+                usage_config,
+                &execution_result,
+            )
+        })
+    });
+
     Ok(execution_result)
 }

bin/router/src/pipeline/mod.rs

@@ -0,0 +1,91 @@
+use std::{sync::Arc, time::{Duration, Instant, SystemTime, UNIX_EPOCH}};
+
+use async_trait::async_trait;
+use graphql_parser::schema::Document;
+use hive_console_sdk::agent::{ExecutionReport, UsageAgent};
+use hive_router_config::{HiveRouterConfig, usage::UsageConfig};
+use hive_router_plan_executor::execution::plan::PlanExecutionOutput;
+use ntex::web::HttpRequest;
+use rand::Rng;
+use tokio_util::sync::CancellationToken;
+
+use crate::background_tasks::BackgroundTask;
+
+pub fn from_config(
+    router_config: &HiveRouterConfig,
+) -> Option<UsageAgent> {
+    router_config.usage.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.buffer_size,
+            usage_config.connect_timeout,
+            usage_config.request_timeout,
+            usage_config.accept_invalid_certs,
+            flush_interval,
+            "hive-router".to_string(),
+        )
+    })
+}
+
+pub fn send_usage_report(
+    schema: Arc<Document<'static, String>>,
+    start: Instant,
+    req: &HttpRequest,
+    operation_name: Option<&String>,
+    operation_body: &str,
+    usage_agent: &UsageAgent,
+    usage_config: &UsageConfig,
+    execution_result: &PlanExecutionOutput,
+) {
+    let mut rng = rand::rng();
+    let sampled = rng.random::<f64>() < usage_config.sample_rate;
+    if !sampled {
+        return;
+    }
+    if operation_name.is_some_and(|op_name| usage_config.exclude.contains(op_name)) {
+        return;
+    }
+    let client_name = get_header_value(req, &usage_config.client_name_header);
+    let client_version = get_header_value(req, &usage_config.client_version_header);
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_secs()
+        * 1000;
+    let duration = start.elapsed();
+    let execution_report = ExecutionReport {
+        schema,
+        client_name,
+        client_version,
+        timestamp,
+        duration,
+        ok: execution_result.error_count == 0,
+        errors: execution_result.error_count,
+        operation_body: operation_body.to_owned(),
+        operation_name: operation_name.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));
+}
+
+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())
+}
+
+#[async_trait]
+impl BackgroundTask for UsageAgent {
+    fn id(&self) -> &str {
+        "usage_report_flush_interval"
+    }
+
+    async fn run(&self, token: CancellationToken) {
+        self.start_flush_interval(Some(token)).await
+    }
+}

bin/router/src/pipeline/usage.rs

@@ -1,6 +1,7 @@
 use arc_swap::{ArcSwap, Guard};
 use async_trait::async_trait;
-use graphql_tools::validation::utils::ValidationError;
+use graphql_parser::schema::Document;
+use graphql_tools::{validation::utils::ValidationError};
 use hive_router_config::{supergraph::SupergraphSource, HiveRouterConfig};
 use hive_router_plan_executor::{
     executors::error::SubgraphExecutorError,
@@ -21,7 +22,7 @@ use tracing::{debug, error, trace};
 
 use crate::{
     background_tasks::{BackgroundTask, BackgroundTasksManager},
-    pipeline::normalize::GraphQLNormalizationPayload,
+    pipeline::{normalize::GraphQLNormalizationPayload},
     supergraph::{
         base::{LoadSupergraphError, ReloadSupergraphResult, SupergraphLoader},
         resolve_from_config,
@@ -39,6 +40,7 @@ pub struct SupergraphData {
     pub metadata: SchemaMetadata,
     pub planner: Planner,
     pub subgraph_executor_map: SubgraphExecutorMap,
+    pub schema: Arc<Document<'static, String>>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -124,6 +126,7 @@ impl SchemaState {
         )?;
 
         Ok(SupergraphData {
+            schema: Arc::new(parsed_supergraph_sdl),
             metadata,
             planner,
             subgraph_executor_map,