feat: send cost model via state

Author: gusinacio

common/src/indexer_service/http/indexer_service.rs

@@ -33,6 +33,7 @@ use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
 use tower_http::trace::TraceLayer;
 use tracing::{info, info_span};
 
+use crate::tap::{ValueCheckReceiver, ValueCheckSender};
 use crate::{
     address::public_key,
     indexer_service::http::{
@@ -167,6 +168,8 @@ where
     pub url_namespace: &'static str,
     pub metrics_prefix: &'static str,
     pub extra_routes: Router<Arc<IndexerServiceState<I>>>,
+    pub value_check_receiver: ValueCheckReceiver,
+    pub value_check_sender: ValueCheckSender,
 }
 
 pub struct IndexerServiceState<I>
@@ -178,6 +181,7 @@ where
     pub tap_manager: Manager<IndexerTapContext>,
     pub service_impl: Arc<I>,
     pub metrics: IndexerServiceMetrics,
+    pub value_check_sender: ValueCheckSender,
 }
 
 pub struct IndexerService {}
@@ -294,6 +298,7 @@ impl IndexerService {
             allocations,
             escrow_accounts,
             domain_separator.clone(),
+            options.value_check_receiver,
         )
         .await;
 
@@ -305,6 +310,7 @@ impl IndexerService {
             tap_manager,
             service_impl: Arc::new(options.service_impl),
             metrics,
+            value_check_sender: options.value_check_sender,
         });
 
         // Rate limits by allowing bursts of 10 requests and requiring 100ms of

common/src/tap.rs

@@ -4,6 +4,7 @@
 use crate::tap::checks::allocation_eligible::AllocationEligible;
 use crate::tap::checks::deny_list_check::DenyListCheck;
 use crate::tap::checks::sender_balance_check::SenderBalanceCheck;
+use crate::tap::checks::value_check::MinimumValue;
 use crate::{escrow_accounts::EscrowAccounts, prelude::Allocation};
 use alloy_sol_types::Eip712Domain;
 use eventuals::Eventual;
@@ -17,6 +18,10 @@ use tracing::error;
 mod checks;
 mod receipt_store;
 
+pub use checks::value_check::{
+    create_value_check, CostModelSource, ValueCheckReceiver, ValueCheckSender,
+};
+
 #[derive(Clone)]
 pub struct IndexerTapContext {
     pgpool: PgPool,
@@ -35,6 +40,7 @@ impl IndexerTapContext {
         indexer_allocations: Eventual<HashMap<Address, Allocation>>,
         escrow_accounts: Eventual<EscrowAccounts>,
         domain_separator: Eip712Domain,
+        value_check_receiver: ValueCheckReceiver,
     ) -> Vec<ReceiptCheck> {
         vec![
             Arc::new(AllocationEligible::new(indexer_allocations)),
@@ -43,6 +49,7 @@ impl IndexerTapContext {
                 domain_separator.clone(),
             )),
             Arc::new(DenyListCheck::new(pgpool, escrow_accounts, domain_separator).await),
+            Arc::new(MinimumValue::new(value_check_receiver)),
         ]
     }
 

common/src/tap/checks/value_check.rs

@@ -10,7 +10,10 @@ use std::{
     time::Duration,
 };
 use thegraph::types::DeploymentId;
-use tokio::{sync::mpsc::Receiver, task::JoinHandle};
+use tokio::{
+    sync::mpsc::{Receiver, Sender},
+    task::JoinHandle,
+};
 use ttl_cache::TtlCache;
 
 use anyhow::anyhow;
@@ -28,10 +31,39 @@ pub struct MinimumValue {
     query_handle: JoinHandle<()>,
 }
 
+#[derive(Clone)]
+pub struct ValueCheckSender {
+    pub tx_cost_model: Sender<CostModelSource>,
+    pub tx_query: Sender<AgoraQuery>,
+}
+
+pub struct ValueCheckReceiver {
+    rx_cost_model: Receiver<CostModelSource>,
+    rx_query: Receiver<AgoraQuery>,
+}
+
+pub fn create_value_check(size: usize) -> (ValueCheckSender, ValueCheckReceiver) {
+    let (tx_cost_model, rx_cost_model) = tokio::sync::mpsc::channel(size);
+    let (tx_query, rx_query) = tokio::sync::mpsc::channel(size);
+
+    (
+        ValueCheckSender {
+            tx_query,
+            tx_cost_model,
+        },
+        ValueCheckReceiver {
+            rx_cost_model,
+            rx_query,
+        },
+    )
+}
+
 impl MinimumValue {
     pub fn new(
-        mut rx_cost_model: Receiver<CostModelSource>,
-        mut rx_query: Receiver<AgoraQuery>,
+        ValueCheckReceiver {
+            mut rx_query,
+            mut rx_cost_model,
+        }: ValueCheckReceiver,
     ) -> Self {
         let cost_model_cache = Arc::new(Mutex::new(HashMap::<DeploymentId, CostModelCache>::new()));
         let query_ids = Arc::new(Mutex::new(HashMap::new()));
@@ -174,9 +206,9 @@ pub struct AgoraQuery {
 
 #[derive(Clone, Eq, Hash, PartialEq)]
 pub struct CostModelSource {
-    deployment_id: DeploymentId,
-    model: String,
-    variables: String,
+    pub deployment_id: DeploymentId,
+    pub model: String,
+    pub variables: String,
 }
 
 pub struct CostModelCache {

service/src/main.rs

@@ -12,9 +12,12 @@ use axum::{
     Json, Router,
 };
 use clap::Parser;
-use indexer_common::indexer_service::http::{
-    IndexerService, IndexerServiceImpl, IndexerServiceOptions, IndexerServiceRelease,
-    IndexerServiceResponse,
+use indexer_common::{
+    indexer_service::http::{
+        IndexerService, IndexerServiceImpl, IndexerServiceOptions, IndexerServiceRelease,
+        IndexerServiceResponse,
+    },
+    tap::{create_value_check, ValueCheckSender},
 };
 use reqwest::{StatusCode, Url};
 use serde_json::{json, Value};
@@ -104,6 +107,7 @@ pub struct SubgraphServiceState {
     pub graph_node_client: reqwest::Client,
     pub graph_node_status_url: String,
     pub graph_node_query_base_url: String,
+    pub value_check_sender: ValueCheckSender,
 }
 
 struct SubgraphService {
@@ -186,6 +190,9 @@ async fn main() -> Result<(), Error> {
     build_info::build_info!(fn build_info);
     let release = IndexerServiceRelease::from(build_info());
 
+    // arbitrary value
+    let (value_check_sender, value_check_receiver) = create_value_check(10);
+
     // Some of the subgrpah service configuration goes into the so-called
     // "state", which will be passed to any request handler, middleware etc.
     // that is involved in serving requests
@@ -212,6 +219,7 @@ async fn main() -> Result<(), Error> {
             .expect("config must have `common.graph_node.query_url` set")
             .query_base_url
             .clone(),
+        value_check_sender: value_check_sender.clone(),
     });
 
     IndexerService::run(IndexerServiceOptions {
@@ -224,6 +232,8 @@ async fn main() -> Result<(), Error> {
             .route("/cost", post(routes::cost::cost))
             .route("/status", post(routes::status))
             .with_state(state),
+        value_check_receiver,
+        value_check_sender,
     })
     .await
 }

service/src/routes/cost.rs

@@ -7,6 +7,7 @@ use std::sync::Arc;
 use async_graphql::{Context, EmptyMutation, EmptySubscription, Object, Schema, SimpleObject};
 use async_graphql_axum::{GraphQLRequest, GraphQLResponse};
 use axum::extract::State;
+use indexer_common::tap::CostModelSource;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use thegraph::types::DeploymentId;
@@ -21,6 +22,16 @@ pub struct GraphQlCostModel {
     pub variables: Option<Value>,
 }
 
+impl From<CostModel> for CostModelSource {
+    fn from(value: CostModel) -> Self {
+        Self {
+            deployment_id: value.deployment,
+            model: value.model.unwrap_or_default(),
+            variables: value.variables.unwrap_or_default().to_string(),
+        }
+    }
+}
+
 impl From<CostModel> for GraphQlCostModel {
     fn from(model: CostModel) -> Self {
         Self {
@@ -45,8 +56,20 @@ impl Query {
             .into_iter()
             .map(|s| DeploymentId::from_str(&s))
             .collect::<Result<Vec<DeploymentId>, _>>()?;
-        let pool = &ctx.data_unchecked::<Arc<SubgraphServiceState>>().database;
+        let state = &ctx.data_unchecked::<Arc<SubgraphServiceState>>();
+
+        let cost_model_sender = &state.value_check_sender;
+
+        let pool = &state.database;
         let cost_models = database::cost_models(pool, &deployment_ids).await?;
+
+        for model in &cost_models {
+            let _ = cost_model_sender
+                .tx_cost_model
+                .send(CostModelSource::from(model.clone()))
+                .await;
+        }
+
         Ok(cost_models.into_iter().map(|m| m.into()).collect())
     }
 
@@ -56,10 +79,20 @@ impl Query {
         deployment: String,
     ) -> Result<Option<GraphQlCostModel>, anyhow::Error> {
         let deployment_id = DeploymentId::from_str(&deployment)?;
+
+        let state = &ctx.data_unchecked::<Arc<SubgraphServiceState>>();
+        let cost_model_sender = &state.value_check_sender;
         let pool = &ctx.data_unchecked::<Arc<SubgraphServiceState>>().database;
-        database::cost_model(pool, &deployment_id)
-            .await
-            .map(|model_opt| model_opt.map(GraphQlCostModel::from))
+        let model = database::cost_model(pool, &deployment_id).await?;
+
+        if let Some(model) = &model {
+            let _ = cost_model_sender
+                .tx_cost_model
+                .send(CostModelSource::from(model.clone()))
+                .await;
+        }
+
+        Ok(model.map(GraphQlCostModel::from))
     }
 }