refactor: consolidate cost model handling

src/indexers.rs

@@ -1,6 +1,5 @@
 pub use urls::*;
 
-pub mod cost_models;
 pub mod indexing_progress;
 pub mod public_poi;
 mod urls;

src/indexers/urls.rs

@@ -14,18 +14,6 @@ pub fn status_url<U: Borrow<Url>>(url: U) -> StatusUrl {
     StatusUrl(url)
 }
 
-/// Builds the URL to the cost model endpoint of the indexer.
-///
-/// # Panics
-/// The function panics if the URL cannot be built.
-pub fn cost_url<U: Borrow<Url>>(url: U) -> CostUrl {
-    let url = url
-        .borrow()
-        .join("cost/")
-        .expect("failed to build indexer cost URL");
-    CostUrl(url)
-}
-
 /// Newtype wrapper around `Url` to provide type safety.
 macro_rules! url_new_type {
     ($name:ident) => {
@@ -69,13 +57,12 @@ macro_rules! url_new_type {
 }
 
 url_new_type!(StatusUrl);
-url_new_type!(CostUrl);
 
 #[cfg(test)]
 mod tests {
     use url::Url;
 
-    use super::{cost_url, status_url};
+    use super::status_url;
 
     /// Ensure the different URL builder functions accept owned and borrowed URL parameters.
     #[test]
@@ -85,9 +72,5 @@ mod tests {
         // Status URL
         let _ = status_url(&url);
         let _ = status_url(url.clone());
-
-        // Cost URL
-        let _ = cost_url(&url);
-        let _ = cost_url(url.clone());
     }
 }

src/network.rs

@@ -2,9 +2,9 @@ pub use errors::{DeploymentError, SubgraphError};
 pub use internal::{Indexing, IndexingId};
 pub use service::{NetworkService, ResolvedSubgraphInfo};
 
+pub mod cost_model;
 mod errors;
 pub mod host_filter;
-pub mod indexer_indexing_cost_model_resolver;
 pub mod indexer_indexing_poi_blocklist;
 pub mod indexer_indexing_poi_resolver;
 pub mod indexer_indexing_progress_resolver;

src/network/cost_model.rs

@@ -0,0 +1,121 @@
+use std::collections::HashMap;
+
+use anyhow::anyhow;
+use thegraph_core::DeploymentId;
+use thegraph_graphql_http::{
+    graphql::{Document, IntoDocument, IntoDocumentWithVariables},
+    http_client::ReqwestExt,
+};
+use url::Url;
+
+pub struct CostModelResolver {
+    http: reqwest::Client,
+    cache: parking_lot::Mutex<HashMap<DeploymentId, u128>>,
+}
+
+impl CostModelResolver {
+    pub fn new(http: reqwest::Client) -> Self {
+        Self {
+            http,
+            cache: Default::default(),
+        }
+    }
+
+    pub async fn resolve(
+        &self,
+        url: &Url,
+        indexings: &[DeploymentId],
+    ) -> HashMap<DeploymentId, u128> {
+        let sources = match self.fetch_cost_model_sources(url, indexings).await {
+            Ok(sources) => sources,
+            Err(cost_model_err) => {
+                tracing::debug!(%url, %cost_model_err);
+                return self.cache.lock().clone();
+            }
+        };
+
+        // Only support cost models of the form `default => x;`.
+        let cost_models: HashMap<DeploymentId, u128> = sources
+            .into_iter()
+            .filter_map(|(deployment, src)| Some((deployment, parse_simple_cost_model(&src)?)))
+            .collect();
+
+        *self.cache.lock() = cost_models.clone();
+        cost_models
+    }
+
+    async fn fetch_cost_model_sources(
+        &self,
+        url: &Url,
+        deployments: &[DeploymentId],
+    ) -> anyhow::Result<HashMap<DeploymentId, String>> {
+        let url = url.join("cost").map_err(|_| anyhow!("invalid URL"))?;
+
+        const QUERY: &str = r#"
+            query costModels($deployments: [String!]!) {
+                costModels(deployments: $deployments) {
+                    deployment
+                    model
+                }
+            }
+        "#;
+        let deployments = deployments.iter().map(|item| item.to_string()).collect();
+        #[derive(serde::Deserialize)]
+        #[serde(rename_all = "camelCase")]
+        struct Response {
+            cost_models: Vec<CostModelSource>,
+        }
+        #[derive(serde::Deserialize)]
+        pub struct CostModelSource {
+            pub deployment: DeploymentId,
+            pub model: String,
+        }
+        struct Request {
+            deployments: Vec<String>,
+        }
+        impl IntoDocumentWithVariables for Request {
+            type Variables = serde_json::Value;
+            fn into_document_with_variables(self) -> (Document, Self::Variables) {
+                (
+                    QUERY.into_document(),
+                    serde_json::json!({ "deployments": self.deployments }),
+                )
+            }
+        }
+        let resp = self
+            .http
+            .post(url)
+            .send_graphql::<Response>(Request { deployments })
+            .await??;
+        Ok(resp
+            .cost_models
+            .into_iter()
+            .map(|CostModelSource { deployment, model }| (deployment, model))
+            .collect())
+    }
+}
+
+fn parse_simple_cost_model(src: &str) -> Option<u128> {
+    let (_, rest) = src.split_once("default")?;
+    let (_, rest) = rest.split_once("=>")?;
+    let (consumed, _) = rest.split_once(";")?;
+    let token = consumed.trim();
+    let fee: f64 = token.parse().ok()?;
+    Some((fee * 1e18) as u128)
+}
+
+#[cfg(test)]
+mod test {
+    #[test]
+    fn parse_simple_cost_model() {
+        let tests = [
+            ("default => 0;", 0),
+            ("default => 1;", 1000000000000000000),
+            ("default => 0.00001;", 10000000000000),
+            ("  default  =>  0.004100  ; ", 4100000000000000),
+        ];
+        for (src, expected) in tests {
+            assert_eq!(super::parse_simple_cost_model(src), Some(expected));
+        }
+    }
+}

src/network/internal/state.rs

@@ -5,7 +5,7 @@ use thegraph_core::alloy::primitives::Address;
 use crate::{
     config::BlockedIndexer,
     network::{
-        host_filter::HostFilter, indexer_indexing_cost_model_resolver::CostModelResolver,
+        cost_model::CostModelResolver, host_filter::HostFilter,
         indexer_indexing_poi_blocklist::PoiBlocklist, indexer_indexing_poi_resolver::PoiResolver,
         indexer_indexing_progress_resolver::IndexingProgressResolver,
         version_filter::VersionFilter,