refactor: use just pglistener

gusinacio · gusinacio · commit 6de9f43a6a03 · 2024-10-18T20:48:53.000+02:00
Signed-off-by: Gustavo Inacio &lt;gustavo@semiotic.ai&gt;
diff --git a/common/src/indexer_service/http/indexer_service.rs b/common/src/indexer_service/http/indexer_service.rs
@@ -37,7 +37,6 @@ use tracing::{info, info_span};
 
 use crate::escrow_accounts::EscrowAccounts;
 use crate::escrow_accounts::EscrowAccountsError;
-use crate::tap::{ValueCheckReceiver, ValueCheckSender};
 use crate::{
     address::public_key,
     indexer_service::http::static_subgraph::static_subgraph_request_handler,
@@ -178,8 +177,6 @@ where
     pub release: IndexerServiceRelease,
     pub url_namespace: &'static str,
     pub extra_routes: Router<Arc<IndexerServiceState<I>>>,
-    pub value_check_receiver: ValueCheckReceiver,
-    pub value_check_sender: ValueCheckSender,
 }
 
 pub struct IndexerServiceState<I>
@@ -194,7 +191,6 @@ where
     // tap
     pub escrow_accounts: Eventual<EscrowAccounts>,
     pub domain_separator: Eip712Domain,
-    pub value_check_sender: ValueCheckSender,
 }
 
 pub struct IndexerService {}
@@ -317,7 +313,6 @@ impl IndexerService {
             domain_separator.clone(),
             timestamp_error_tolerance,
             receipt_max_value,
-            options.value_check_receiver,
         )
         .await;
 
@@ -334,7 +329,6 @@ impl IndexerService {
             service_impl: Arc::new(options.service_impl),
             escrow_accounts,
             domain_separator,
-            value_check_sender: options.value_check_sender,
         });
 
         // Rate limits by allowing bursts of 10 requests and requiring 100ms of
diff --git a/common/src/tap.rs b/common/src/tap.rs
@@ -24,9 +24,7 @@ use tracing::error;
 mod checks;
 mod receipt_store;
 
-pub use checks::value_check::{
-    create_value_check, AgoraQuery, CostModelSource, ValueCheckReceiver, ValueCheckSender,
-};
+pub use checks::value_check::{AgoraQuery, CostModelSource};
 
 #[derive(Clone)]
 pub struct IndexerTapContext {
@@ -50,7 +48,6 @@ impl IndexerTapContext {
         domain_separator: Eip712Domain,
         timestamp_error_tolerance: Duration,
         receipt_max_value: u128,
-        value_check_receiver: ValueCheckReceiver,
     ) -> Vec<ReceiptCheck> {
         vec![
             Arc::new(AllocationEligible::new(indexer_allocations)),
@@ -59,9 +56,9 @@ impl IndexerTapContext {
                 domain_separator.clone(),
             )),
             Arc::new(TimestampCheck::new(timestamp_error_tolerance)),
-            Arc::new(DenyListCheck::new(pgpool, escrow_accounts, domain_separator).await),
+            Arc::new(DenyListCheck::new(pgpool.clone(), escrow_accounts, domain_separator).await),
             Arc::new(ReceiptMaxValueCheck::new(receipt_max_value)),
-            Arc::new(MinimumValue::new(value_check_receiver)),
+            Arc::new(MinimumValue::new(pgpool).await),
         ]
     }
 
diff --git a/common/src/tap/checks/value_check.rs b/common/src/tap/checks/value_check.rs
@@ -5,18 +5,15 @@ use anyhow::anyhow;
 use bigdecimal::ToPrimitive;
 use cost_model::CostModel;
 use sqlx::{postgres::PgListener, PgPool};
-use tracing::error;
 use std::{
     cmp::min,
     collections::HashMap,
     sync::{Arc, Mutex},
     time::Duration,
 };
 use thegraph_core::DeploymentId;
-use tokio::{
-    sync::mpsc::{Receiver, Sender},
-    task::JoinHandle,
-};
+use tokio::task::JoinHandle;
+use tracing::error;
 use ttl_cache::TtlCache;
 
 use tap_core::receipt::{
@@ -30,60 +27,25 @@ pub struct MinimumValue {
     model_handle: JoinHandle<()>,
 }
 
-#[derive(Clone)]
-pub struct ValueCheckSender {
-    pub tx_cost_model: Sender<CostModelSource>,
-}
-
-pub struct ValueCheckReceiver {
-    rx_cost_model: Receiver<CostModelSource>,
-}
+impl MinimumValue {
+    pub async fn new(pgpool: PgPool) -> Self {
+        let cost_model_cache = Arc::new(Mutex::new(HashMap::<DeploymentId, CostModelCache>::new()));
 
-pub fn create_value_check(size: usize) -> (ValueCheckSender, ValueCheckReceiver) {
-    let (tx_cost_model, rx_cost_model) = tokio::sync::mpsc::channel(size);
+        let mut pglistener = PgListener::connect_with(&pgpool.clone()).await.unwrap();
+        pglistener.listen("cost_models_update_notify").await.expect(
+            "should be able to subscribe to Postgres Notify events on the channel \
+                'cost_models_update_notify'",
+        );
 
-    (
-        ValueCheckSender { tx_cost_model },
-        ValueCheckReceiver { rx_cost_model },
-    )
-}
+        // TODO start watcher
+        let cancel_token = tokio_util::sync::CancellationToken::new();
 
-impl MinimumValue {
-    pub fn new(ValueCheckReceiver { mut rx_cost_model }: ValueCheckReceiver) -> Self {
-        let cost_model_cache = Arc::new(Mutex::new(HashMap::<DeploymentId, CostModelCache>::new()));
-        let cache = cost_model_cache.clone();
-        let model_handle = tokio::spawn(async move {
-            loop {
-                let model = rx_cost_model.recv().await;
-                match model {
-                    Some(value) => {
-                        let deployment_id = value.deployment_id;
-
-                        if let Some(query) = cache.lock().unwrap().get_mut(&deployment_id) {
-                            let _ = query.insert_model(value).inspect_err(|err| {
-                                tracing::error!(
-                                    "Error while compiling cost model for deployment id {}. Error: {}",
-                                    deployment_id, err
-                                )
-                            });
-                        } else {
-                            match CostModelCache::new(value) {
-                                Ok(value) => {
-                                    cache.lock().unwrap().insert(deployment_id, value);
-                                }
-                                Err(err) => {
-                                    tracing::error!(
-                                        "Error while compiling cost model for deployment id {}. Error: {}",
-                                        deployment_id, err
-                                    )
-                                }
-                            }
-                        }
-                    }
-                    None => break,
-                }
-            }
-        });
+        let model_handle = tokio::spawn(Self::cost_models_watcher(
+            pgpool.clone(),
+            pglistener,
+            cost_model_cache.clone(),
+            cancel_token.clone(),
+        ));
 
         Self {
             cost_model_cache,
@@ -92,17 +54,11 @@ impl MinimumValue {
     }
 
     async fn cost_models_watcher(
-        pgpool: PgPool,
+        _pgpool: PgPool,
         mut pglistener: PgListener,
-        denylist: Arc<Mutex<HashMap<DeploymentId, CostModelCache>>>,
+        cost_model_cache: Arc<Mutex<HashMap<DeploymentId, CostModelCache>>>,
         cancel_token: tokio_util::sync::CancellationToken,
     ) {
-        #[derive(serde::Deserialize)]
-        struct DenylistNotification {
-            tg_op: String,
-            deployment: DeploymentId,
-        }
-
         loop {
             tokio::select! {
                 _ = cancel_token.cancelled() => {
@@ -112,39 +68,58 @@ impl MinimumValue {
                 pg_notification = pglistener.recv() => {
                     let pg_notification = pg_notification.expect(
                     "should be able to receive Postgres Notify events on the channel \
-                    'scalar_tap_deny_notification'",
+                    'cost_models_update_notify'",
                     );
 
-                    let denylist_notification: DenylistNotification =
+                    let cost_model_notification: CostModelNotification =
                         serde_json::from_str(pg_notification.payload()).expect(
                             "should be able to deserialize the Postgres Notify event payload as a \
-                            DenylistNotification",
+                            CostModelNotification",
                         );
 
-                    match denylist_notification.tg_op.as_str() {
+                    let deployment_id = cost_model_notification.deployment;
+
+                    match cost_model_notification.tg_op.as_str() {
                         "INSERT" => {
-                            denylist
-                                .write()
-                                .unwrap()
-                                .insert(denylist_notification.sender_address);
+                            let cost_model_source: CostModelSource = cost_model_notification.into();
+                            let mut cost_model_cache = cost_model_cache
+                                .lock()
+                                .unwrap();
+
+                            match cost_model_cache.get_mut(&deployment_id) {
+                                Some(cache) => {
+                                    let _ = cache.insert_model(cost_model_source);
+                                },
+                                None => {
+                                     if let Ok(cache) = CostModelCache::new(cost_model_source).inspect_err(|err| {
+                                        tracing::error!(
+                                            "Error while compiling cost model for deployment id {}. Error: {}",
+                                            deployment_id, err
+                                        )
+                                    }) {
+                                        cost_model_cache.insert(deployment_id, cache);
+                                    }
+                                },
+                            }
                         }
                         "DELETE" => {
-                            denylist
-                                .write()
+                            cost_model_cache
+                                .lock()
                                 .unwrap()
-                                .remove(&denylist_notification.sender_address);
+                                .remove(&cost_model_notification.deployment);
                         }
-                        // UPDATE and TRUNCATE are not expected to happen. Reload the entire denylist.
+                        // UPDATE and TRUNCATE are not expected to happen. Reload the entire cost
+                        // model cache.
                         _ => {
                             error!(
-                                "Received an unexpected denylist table notification: {}. Reloading entire \
-                                denylist.",
-                                denylist_notification.tg_op
+                                "Received an unexpected cost model table notification: {}. Reloading entire \
+                                cost model.",
+                                cost_model_notification.tg_op
                             );
 
-                            Self::sender_denylist_reload(pgpool.clone(), denylist.clone())
-                                .await
-                                .expect("should be able to reload the sender denylist")
+                            // Self::sender_denylist_reload(pgpool.clone(), denylist.clone())
+                            //     .await
+                            //     .expect("should be able to reload cost models")
                         }
                     }
                 }
@@ -229,6 +204,24 @@ pub struct CostModelSource {
     pub variables: String,
 }
 
+#[derive(serde::Deserialize)]
+struct CostModelNotification {
+    tg_op: String,
+    deployment: DeploymentId,
+    model: String,
+    variables: String,
+}
+
+impl From<CostModelNotification> for CostModelSource {
+    fn from(value: CostModelNotification) -> Self {
+        CostModelSource {
+            deployment_id: value.deployment,
+            model: value.model,
+            variables: value.variables,
+        }
+    }
+}
+
 pub struct CostModelCache {
     models: TtlCache<CostModelSource, CostModel>,
     latest_model: CostModel,
diff --git a/service/src/routes/cost.rs b/service/src/routes/cost.rs
@@ -7,12 +7,12 @@ 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 lazy_static::lazy_static;
 use prometheus::{
     register_counter, register_counter_vec, register_histogram, register_histogram_vec, Counter,
     CounterVec, Histogram, HistogramVec,
 };
-use indexer_common::tap::CostModelSource;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use thegraph_core::DeploymentId;
@@ -140,18 +140,9 @@ impl Query {
     ) -> Result<Vec<GraphQlCostModel>, anyhow::Error> {
         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())
     }
 
@@ -160,19 +151,9 @@ impl Query {
         ctx: &Context<'_>,
         deployment_id: DeploymentId,
     ) -> Result<Option<GraphQlCostModel>, anyhow::Error> {
-
-        let state = &ctx.data_unchecked::<Arc<SubgraphServiceState>>();
-        let cost_model_sender = &state.value_check_sender;
         let pool = &ctx.data_unchecked::<Arc<SubgraphServiceState>>().database;
         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))
     }
 }
diff --git a/service/src/service.rs b/service/src/service.rs
@@ -7,9 +7,8 @@ use std::time::Duration;
 use super::{config::Config, error::SubgraphServiceError, routes};
 use anyhow::anyhow;
 use axum::{async_trait, routing::post, Json, Router};
-use indexer_common::{
-    indexer_service::http::{AttestationOutput, IndexerServiceImpl, IndexerServiceResponse},
-    tap::{create_value_check, ValueCheckSender},
+use indexer_common::indexer_service::http::{
+    AttestationOutput, IndexerServiceImpl, IndexerServiceResponse,
 };
 use indexer_config::Config as MainConfig;
 use reqwest::Url;
@@ -68,7 +67,6 @@ 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 {
@@ -148,9 +146,6 @@ pub async fn run() -> anyhow::Result<()> {
     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 subgraph service configuration goes into the so-called
     // "state", which will be passed to any request handler, middleware etc.
     // that is involved in serving requests
@@ -177,7 +172,6 @@ pub async fn run() -> anyhow::Result<()> {
             .expect("config must have `common.graph_node.query_url` set")
             .query_base_url
             .clone(),
-        value_check_sender: value_check_sender.clone(),
     });
 
     IndexerService::run(IndexerServiceOptions {
@@ -189,8 +183,6 @@ pub async fn run() -> anyhow::Result<()> {
             .route("/cost", post(routes::cost::cost))
             .route("/status", post(routes::status))
             .with_state(state),
-        value_check_receiver,
-        value_check_sender,
     })
     .await
 }
