feat: Query Trace Processing

entity/src/lib.rs

@@ -11,6 +11,9 @@ pub mod content_audit;
 pub mod execution_metadata;
 pub mod key_value;
 pub mod node;
+pub mod query_response;
+pub mod query_response_node;
+pub mod query_trace;
 pub mod record;
 pub mod state_roots;
 pub mod test;

entity/src/query_response.rs

@@ -0,0 +1,63 @@
+use std::{collections::HashMap, hash::Hash};
+
+use super::query_response_node;
+use anyhow::Result;
+use enr::NodeId;
+use ethportal_api::types::query_trace::QueryResponse;
+use sea_orm::{entity::prelude::*, ActiveValue::NotSet, Set};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
+#[sea_orm(table_name = "query_response")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    pub query_trace_id: i32,
+    pub node_id: i32,
+    pub duration_ms: u32,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::query_trace::Entity",
+        from = "Column::QueryTraceId",
+        to = "super::query_trace::Column::Id",
+        on_update = "Cascade",
+        on_delete = "Cascade"
+    )]
+    QueryTrace,
+    #[sea_orm(has_many = "super::query_response_node::Entity")]
+    QueryResponseNode,
+}
+
+impl Related<super::query_trace::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::QueryTrace.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}
+
+// Create a new query response entity.
+pub async fn create(
+    query_trace_id: i32,
+    query_responses: HashMap<NodeId, QueryResponse>,
+    conn: &DatabaseConnection,
+) -> Result<()> {
+    //  let duration_ms = query_response.duration_ms as u32;
+
+    // let query_response_model = ActiveModel {
+    //     id: NotSet,
+    //     query_trace_id: Set(query_trace_id),
+    //     node_id: Set(node.id),
+    //     duration_ms: Set(duration_ms),
+    // };
+    // let query_response_model = query_response_model.insert(conn).await?;
+
+    // Lookup node IDs of all nodes that responded.
+
+    // For each node that responded, look up the node and create a new query response node entity.
+    // query_response_node::create(query_response_model.id, query_response, conn).await?;
+
+    Ok(())
+}

entity/src/query_response_node.rs

@@ -0,0 +1,64 @@
+use anyhow::Result;
+use sea_orm::{entity::prelude::*, ActiveValue::NotSet, Set};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
+#[sea_orm(table_name = "query_response_node")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    pub query_response_id: i32,
+    pub node_id: i32,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::query_response::Entity",
+        from = "Column::QueryResponseId",
+        to = "super::query_response::Column::Id",
+        on_update = "Cascade",
+        on_delete = "Cascade"
+    )]
+    QueryResponse,
+    #[sea_orm(
+        belongs_to = "super::node::Entity",
+        from = "Column::NodeId",
+        to = "super::node::Column::Id",
+        on_update = "Cascade",
+        on_delete = "Cascade"
+    )]
+    Node,
+}
+
+impl Related<super::query_response::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::QueryResponse.def()
+    }
+}
+
+impl Related<super::node::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Node.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}
+
+// Create a new query response node entity.
+pub async fn create(
+    node_id: i32,
+    query_response_ids: Vec<i32>,
+    conn: &DatabaseConnection,
+) -> Result<()> {
+    let models = query_response_ids
+        .into_iter()
+        .map(|query_response_id| ActiveModel {
+            id: NotSet,
+            query_response_id: Set(query_response_id),
+            node_id: Set(node_id),
+        })
+        .collect::<Vec<_>>();
+
+    Entity::insert_many(models).exec(conn).await?;
+    Ok(())
+}

entity/src/query_trace.rs

@@ -0,0 +1,143 @@
+use std::time::UNIX_EPOCH;
+
+use super::query_response;
+use anyhow::Result;
+use chrono::{DateTime, TimeZone, Utc};
+use ethportal_api::types::query_trace::QueryTrace;
+use sea_orm::{entity::prelude::*, ActiveValue::NotSet, Set};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
+#[sea_orm(table_name = "query_trace")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    pub timestamp: DateTime<Utc>,
+    pub origin_record_id: i32,
+    pub content_successfully_received_from: Option<i32>,
+    pub target_content: i32,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(has_many = "super::query_response::Entity")]
+    QueryResponse,
+}
+
+impl Related<super::query_response::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::QueryResponse.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}
+
+pub async fn create(query_trace: QueryTrace, conn: &DatabaseConnection) -> Result<Model> {
+    // Get query timestamp.
+    let timestamp = query_trace.started_at_ms;
+    let timestamp: DateTime<Utc> = {
+        let duration = timestamp.duration_since(UNIX_EPOCH).unwrap();
+        Utc.timestamp_opt(duration.as_secs() as i64, duration.as_nanos() as u32)
+            .single()
+            .unwrap()
+    };
+
+    // Get origin record.
+    let origin_node_id = query_trace.origin;
+    let origin_record = match super::record::Entity::find()
+        .filter(super::record::Column::NodeId.eq(origin_node_id.raw().to_vec()))
+        .one(conn)
+        .await?
+    {
+        Some(record) => record,
+        None => {
+            return Err(anyhow::anyhow!(
+                "No record found for node ID {}",
+                origin_node_id
+            ));
+        }
+    };
+
+    // Get the most recent record for the node that responded (if any).
+    let received_from_node_id = query_trace.received_from;
+    let received_from_record = if let Some(node_id) = received_from_node_id {
+        super::record::Entity::find()
+            .filter(super::record::Column::NodeId.eq(node_id.raw().to_vec()))
+            .one(conn)
+            .await?
+    } else {
+        None
+    };
+
+    // Get the target content's db ID.
+    let content_id_bytes = query_trace.target_id.to_vec();
+    let target_content = match super::content::Entity::find()
+        .filter(super::content::Column::ContentId.eq(content_id_bytes.clone()))
+        .one(conn)
+        .await?
+    {
+        Some(content) => content.id,
+        None => {
+            return Err(anyhow::anyhow!(
+                "No content found for content ID {:?}",
+                content_id_bytes
+            ));
+        }
+    };
+
+    // Save query_trace model into the database.
+    let query_trace_model = ActiveModel {
+        id: NotSet,
+        timestamp: Set(timestamp),
+        origin_record_id: Set(origin_record.id),
+        content_successfully_received_from: Set(received_from_record.map(|r| r.id)),
+        target_content: Set(target_content),
+    };
+
+    let query_trace_model = match query_trace_model.insert(conn).await {
+        Ok(query_trace) => query_trace,
+        Err(err) => {
+            return Err(anyhow::anyhow!("Failed to save query trace: {:?}", err));
+        }
+    };
+
+    // Create all of the QueryResponse entries for this QueryTrace.
+    query_response::create(query_trace_model.id, query_trace.responses, conn).await?;
+
+    // Create QueryResponse objects for each response record, tied to the QueryTrace.
+    Ok(query_trace_model)
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use crate::test::setup_database;
+
+    #[tokio::test]
+    async fn test_create_query_trace() -> Result<(), DbErr> {
+        let conn = setup_database().await?;
+
+        // 1.) create a set of ENRs/Node IDs to use as origin and response records.
+
+        // 2.) create a query trace with the origin and response records and a content ID.
+        // 3.) create a new query trace in the DB.
+        // 4.) read it and verify that it was created succesfully.
+
+        // let query_trace = QueryTrace {
+        //     started_at_ms: Utc::now(),
+        //     origin: 1,
+        //     received_from: Some(2),
+        //     target_id: [1, 2, 3, 4, 5],
+        // };
+
+        // let query_trace = create(query_trace, &conn.0).await.unwrap();
+
+        // assert_eq!(query_trace.id, 1);
+        // assert_eq!(query_trace.timestamp, Utc::now());
+        // assert_eq!(query_trace.origin_record_id, 1);
+        // assert_eq!(query_trace.content_successfully_received_from, Some(2));
+        // assert_eq!(query_trace.target_content, 1);
+
+        Ok(())
+    }
+}

entity/src/test.rs

@@ -24,7 +24,7 @@ use pgtemp::PgTempDB;
 
 #[allow(dead_code)]
 // Temporary Postgres db will be deleted once PgTempDB goes out of scope, so keep it in scope.
-async fn setup_database() -> Result<(DbConn, PgTempDB), DbErr> {
+pub async fn setup_database() -> Result<(DbConn, PgTempDB), DbErr> {
     let db: PgTempDB = PgTempDB::async_new().await;
     let conn: DbConn = Database::connect(db.connection_uri()).await?;
     Migrator::up(&conn, None).await.unwrap();

glados-web/Cargo.toml

@@ -18,6 +18,7 @@ clap = { version = "4.0.26", features = ["derive"] }
 entity = { path = "../entity" }
 env_logger = "0.9.3"
 ethportal-api = { git = "https://github.com/ethereum/trin" }
+discv5 = "0.4.1"
 glados-core = { path = "../glados-core" }
 migration = { path = "../migration" }
 itertools = "0.10.5"

glados-web/src/routes.rs

@@ -1,18 +1,22 @@
-use alloy_primitives::U256;
+use alloy_primitives::{B256, U256};
 use axum::{
     extract::{Extension, Path, Query as HttpQuery},
     http::StatusCode,
     response::IntoResponse,
     Json,
 };
 use chrono::{DateTime, Utc};
+use discv5::enr::NodeId;
 use entity::{audit_stats, census, census_node, client_info};
 use entity::{
     content,
     content_audit::{self, AuditResult},
     execution_metadata, key_value, node, record,
 };
-use ethportal_api::types::distance::{Distance, Metric, XorMetric};
+use ethportal_api::types::{
+    distance::{Distance, Metric, XorMetric},
+    query_trace::QueryTrace,
+};
 use ethportal_api::utils::bytes::{hex_decode, hex_encode};
 use ethportal_api::{jsonrpsee::core::__reexports::serde_json, BeaconContentKey, StateContentKey};
 use ethportal_api::{HistoryContentKey, OverlayContentKey};
@@ -28,6 +32,7 @@ use sea_orm::{
     FromQueryResult, LoaderTrait, ModelTrait, QueryFilter, QueryOrder, QuerySelect,
 };
 use serde::Serialize;
+use serde_json::json;
 use std::collections::{HashMap, HashSet};
 use std::fmt::Formatter;
 use std::sync::Arc;
@@ -844,6 +849,36 @@ pub async fn contentaudit_detail(
         .unwrap()
         .expect("No audit found");
 
+    let trace = &audit.trace;
+    // Deserialize trace into QueryTrace object
+    let trace: QueryTrace = serde_json::from_value(json!(trace)).unwrap();
+    // Build a HashMap of (node ID, distance) for each node in the trace
+    let node_distances: HashMap<NodeId, B256> = trace
+        .metadata
+        .iter()
+        .map(|(node_id, node_info)| {
+            let node_id = node_id.clone();
+            let distance = node_info.distance;
+            (node_id, distance)
+        })
+        .collect();
+    // Do a query to get the most recent data_radius_high for each node, passing in all node IDs.
+    let node_ids: Vec<NodeId> = node_distances.keys().cloned().collect();
+    let node_data_radius_high: HashMap<NodeId, u8> = node::Entity::find()
+        .filter(node::Column::NodeId.eq_any(node_ids))
+        .order_by_desc(census_node::Column::SurveyedAt)
+        .group_by(census_node::Column::NodeId)
+        .select(census_node::Column::DataRadiusHigh)
+        .all(&state.database_connection)
+        .await
+        .unwrap()
+        .into_iter()
+        .map(|(node_id, data_radius_high)| (node_id, data_radius_high))
+        .collect();
+    // Create a map of (node ID, bool) for each node.
+
+    // Iterate through NodeInfo, setting within_radius boolean
+
     let content = audit
         .find_related(content::Entity)
         .one(&state.database_connection)

migration/src/lib.rs

@@ -12,6 +12,7 @@ mod m20231107_004843_create_audit_stats;
 mod m20240213_190221_add_fourfours_stats;
 mod m20240322_205213_add_content_audit_index;
 mod m20240515_064320_state_roots;
+mod m20240626_205432_add_content_id_index;
 
 pub struct Migrator;
 
@@ -31,6 +32,7 @@ impl MigratorTrait for Migrator {
             Box::new(m20240213_190221_add_fourfours_stats::Migration),
             Box::new(m20240322_205213_add_content_audit_index::Migration),
             Box::new(m20240515_064320_state_roots::Migration),
+            Box::new(m20240626_205432_add_content_id_index::Migration),
         ]
     }
 }