fix: initialize allocations monitor (#428)

Author: gusinacio

common/src/allocations/monitor.rs

@@ -8,15 +8,11 @@ use std::{
 };
 
 use super::Allocation;
-use crate::prelude::SubgraphClient;
+use crate::{prelude::SubgraphClient, watcher::new_watcher};
 use alloy::primitives::{TxHash, B256, U256};
 use graphql_client::GraphQLQuery;
 use thegraph_core::{Address, DeploymentId};
-use tokio::{
-    sync::watch::{self, Receiver},
-    time::{self, sleep},
-};
-use tracing::warn;
+use tokio::sync::watch::Receiver;
 
 type BigInt = U256;
 type Bytes = B256;
@@ -58,43 +54,21 @@ impl TryFrom<allocations_query::AllocationFragment> for Allocation {
 }
 
 /// An always up-to-date list of an indexer's active and recently closed allocations.
-pub fn indexer_allocations(
+pub async fn indexer_allocations(
     network_subgraph: &'static SubgraphClient,
     indexer_address: Address,
     interval: Duration,
     recently_closed_allocation_buffer: Duration,
-) -> Receiver<HashMap<Address, Allocation>> {
-    let (tx, rx) = watch::channel(HashMap::new());
-    tokio::spawn(async move {
-        // Refresh indexer allocations every now and then
-        let mut time_interval = time::interval(interval);
-        time_interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
-        loop {
-            time_interval.tick().await;
-            let result = get_allocations(
-                network_subgraph,
-                indexer_address,
-                recently_closed_allocation_buffer,
-            )
-            .await;
-            match result {
-                Ok(allocations) => {
-                    tx.send(allocations)
-                        .expect("Failed to update indexer_allocations channel");
-                }
-                Err(err) => {
-                    warn!(
-                        "Failed to fetch active or recently closed allocations for indexer {:?}: {}",
-                        indexer_address, err
-                    );
-
-                    // Sleep for a bit before we retry
-                    sleep(interval.div_f32(2.0)).await;
-                }
-            }
-        }
-    });
-    rx
+) -> anyhow::Result<Receiver<HashMap<Address, Allocation>>> {
+    new_watcher(interval, move || async move {
+        get_allocations(
+            network_subgraph,
+            indexer_address,
+            recently_closed_allocation_buffer,
+        )
+        .await
+    })
+    .await
 }
 
 pub async fn get_allocations(

common/src/attestations/dispute_manager.rs

@@ -2,13 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::subgraph_client::SubgraphClient;
+use crate::watcher::new_watcher;
 use alloy::primitives::Address;
 use anyhow::Error;
 use graphql_client::GraphQLQuery;
 use std::time::Duration;
-use tokio::sync::watch::{self, Receiver};
-use tokio::time::{self, sleep};
-use tracing::warn;
+use tokio::sync::watch::Receiver;
 
 type Bytes = Address;
 
@@ -21,46 +20,26 @@ type Bytes = Address;
 )]
 struct DisputeManager;
 
-pub fn dispute_manager(
+pub async fn dispute_manager(
     network_subgraph: &'static SubgraphClient,
     interval: Duration,
-) -> Receiver<Option<Address>> {
-    let (tx, rx) = watch::channel(None);
-    tokio::spawn(async move {
-        let mut time_interval = time::interval(interval);
-        time_interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
-        loop {
-            time_interval.tick().await;
-
-            let result = async {
-                let response = network_subgraph
-                    .query::<DisputeManager, _>(dispute_manager::Variables {})
-                    .await?;
-                response?
-                    .graph_network
-                    .map(|network| network.dispute_manager)
-                    .ok_or_else(|| Error::msg("Network 1 not found in network subgraph"))
-            }
-            .await;
-
-            match result {
-                Ok(address) => tx
-                    .send(Some(address))
-                    .expect("Failed to update dispute_manager channel"),
-                Err(err) => {
-                    warn!("Failed to query dispute manager for network: {}", err);
-                    // Sleep for a bit before we retry
-                    sleep(interval.div_f32(2.0)).await;
-                }
-            }
-        }
-    });
-    rx
+) -> anyhow::Result<Receiver<Address>> {
+    new_watcher(interval, move || async move {
+        let response = network_subgraph
+            .query::<DisputeManager, _>(dispute_manager::Variables {})
+            .await?;
+        response?
+            .graph_network
+            .map(|network| network.dispute_manager)
+            .ok_or_else(|| Error::msg("Network 1 not found in network subgraph"))
+    })
+    .await
 }
 
 #[cfg(test)]
 mod test {
     use serde_json::json;
+    use tokio::time::sleep;
     use wiremock::{
         matchers::{method, path},
         Mock, MockServer, ResponseTemplate,
@@ -109,9 +88,11 @@ mod test {
     async fn test_parses_dispute_manager_from_network_subgraph_correctly() {
         let (network_subgraph, _mock_server) = setup_mock_network_subgraph().await;
 
-        let dispute_manager = dispute_manager(network_subgraph, Duration::from_secs(60));
+        let dispute_manager = dispute_manager(network_subgraph, Duration::from_secs(60))
+            .await
+            .unwrap();
         sleep(Duration::from_millis(50)).await;
         let result = *dispute_manager.borrow();
-        assert_eq!(result.unwrap(), *DISPUTE_MANAGER_ADDRESS);
+        assert_eq!(result, *DISPUTE_MANAGER_ADDRESS);
     }
 }

common/src/attestations/signers.rs

@@ -2,97 +2,59 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use bip39::Mnemonic;
-use std::collections::HashMap;
 use std::sync::Arc;
+use std::{collections::HashMap, sync::Mutex};
 use thegraph_core::{Address, ChainId};
-use tokio::{
-    select,
-    sync::{
-        watch::{self, Receiver},
-        Mutex,
-    },
-};
+use tokio::sync::watch::Receiver;
 use tracing::warn;
 
-use crate::prelude::{Allocation, AttestationSigner};
+use crate::{
+    prelude::{Allocation, AttestationSigner},
+    watcher::join_and_map_watcher,
+};
 
 /// An always up-to-date list of attestation signers, one for each of the indexer's allocations.
-pub async fn attestation_signers(
-    mut indexer_allocations_rx: Receiver<HashMap<Address, Allocation>>,
+pub fn attestation_signers(
+    indexer_allocations_rx: Receiver<HashMap<Address, Allocation>>,
     indexer_mnemonic: Mnemonic,
     chain_id: ChainId,
-    mut dispute_manager_rx: Receiver<Option<Address>>,
+    dispute_manager_rx: Receiver<Address>,
 ) -> Receiver<HashMap<Address, AttestationSigner>> {
     let attestation_signers_map: &'static Mutex<HashMap<Address, AttestationSigner>> =
         Box::leak(Box::new(Mutex::new(HashMap::new())));
-    let indexer_mnemonic = indexer_mnemonic.to_string();
+    let indexer_mnemonic = Arc::new(indexer_mnemonic.to_string());
 
-    let starter_signers_map = modify_sigers(
-        Arc::new(indexer_mnemonic.clone()),
-        chain_id,
-        attestation_signers_map,
-        indexer_allocations_rx.clone(),
-        dispute_manager_rx.clone(),
+    join_and_map_watcher(
+        indexer_allocations_rx,
+        dispute_manager_rx,
+        move |(allocation, dispute)| {
+            let indexer_mnemonic = indexer_mnemonic.clone();
+            modify_sigers(
+                &indexer_mnemonic,
+                chain_id,
+                attestation_signers_map,
+                &allocation,
+                &dispute,
+            )
+        },
     )
-    .await;
-
-    // Whenever the indexer's active or recently closed allocations change, make sure
-    // we have attestation signers for all of them.
-    let (signers_tx, signers_rx) = watch::channel(starter_signers_map);
-    tokio::spawn(async move {
-        loop {
-            let updated_signers = select! {
-                Ok(())= indexer_allocations_rx.changed() =>{
-                    modify_sigers(
-                        Arc::new(indexer_mnemonic.clone()),
-                        chain_id,
-                        attestation_signers_map,
-                        indexer_allocations_rx.clone(),
-                        dispute_manager_rx.clone(),
-                    ).await
-                },
-                Ok(())= dispute_manager_rx.changed() =>{
-                    modify_sigers(
-                        Arc::new(indexer_mnemonic.clone()),
-                        chain_id,
-                        attestation_signers_map,
-                        indexer_allocations_rx.clone(),
-                        dispute_manager_rx.clone()
-                    ).await
-                },
-                else=>{
-                    // Something is wrong.
-                    panic!("dispute_manager_rx or allocations_rx was dropped");
-                }
-            };
-            signers_tx
-                .send(updated_signers)
-                .expect("Failed to update signers channel");
-        }
-    });
-
-    signers_rx
 }
-async fn modify_sigers(
-    indexer_mnemonic: Arc<String>,
+fn modify_sigers(
+    indexer_mnemonic: &str,
     chain_id: ChainId,
     attestation_signers_map: &'static Mutex<HashMap<Address, AttestationSigner>>,
-    allocations_rx: Receiver<HashMap<Address, Allocation>>,
-    dispute_manager_rx: Receiver<Option<Address>>,
+    allocations: &HashMap<Address, Allocation>,
+    dispute_manager: &Address,
 ) -> HashMap<thegraph_core::Address, AttestationSigner> {
-    let mut signers = attestation_signers_map.lock().await;
-    let allocations = allocations_rx.borrow().clone();
-    let Some(dispute_manager) = *dispute_manager_rx.borrow() else {
-        return signers.clone();
-    };
+    let mut signers = attestation_signers_map.lock().unwrap();
     // Remove signers for allocations that are no longer active or recently closed
     signers.retain(|id, _| allocations.contains_key(id));
 
     // Create signers for new allocations
     for (id, allocation) in allocations.iter() {
         if !signers.contains_key(id) {
             let signer =
-                AttestationSigner::new(&indexer_mnemonic, allocation, chain_id, dispute_manager);
+                AttestationSigner::new(indexer_mnemonic, allocation, chain_id, *dispute_manager);
             match signer {
                 Ok(signer) => {
                     signers.insert(*id, signer);
@@ -113,24 +75,22 @@ async fn modify_sigers(
 
 #[cfg(test)]
 mod tests {
+    use tokio::sync::watch;
+
     use crate::test_vectors::{DISPUTE_MANAGER_ADDRESS, INDEXER_ALLOCATIONS, INDEXER_MNEMONIC};
 
     use super::*;
 
     #[tokio::test]
     async fn test_attestation_signers_update_with_allocations() {
         let (allocations_tx, allocations_rx) = watch::channel(HashMap::new());
-        let (dispute_manager_tx, dispute_manager_rx) = watch::channel(None);
-        dispute_manager_tx
-            .send(Some(*DISPUTE_MANAGER_ADDRESS))
-            .unwrap();
+        let (_, dispute_manager_rx) = watch::channel(*DISPUTE_MANAGER_ADDRESS);
         let mut signers = attestation_signers(
             allocations_rx,
             INDEXER_MNEMONIC.clone(),
             1,
             dispute_manager_rx,
-        )
-        .await;
+        );
 
         // Test that an empty set of allocations leads to an empty set of signers
         allocations_tx.send(HashMap::new()).unwrap();

common/src/indexer_service/http/indexer_service.rs

@@ -226,7 +226,9 @@ impl IndexerService {
         )));
 
         // Identify the dispute manager for the configured network
-        let dispute_manager = dispute_manager(network_subgraph, Duration::from_secs(3600));
+        let dispute_manager = dispute_manager(network_subgraph, Duration::from_secs(3600))
+            .await
+            .expect("Failed to initialize dispute manager");
 
         // Monitor the indexer's own allocations
         let allocations = indexer_allocations(
@@ -243,7 +245,9 @@ impl IndexerService {
                 .subgraphs
                 .network
                 .recently_closed_allocation_buffer_secs,
-        );
+        )
+        .await
+        .expect("Failed to initialize indexer_allocations watcher");
 
         // Maintain an up-to-date set of attestation signers, one for each
         // allocation
@@ -252,8 +256,7 @@ impl IndexerService {
             options.config.indexer.operator_mnemonic.clone(),
             options.config.blockchain.chain_id as u64,
             dispute_manager,
-        )
-        .await;
+        );
 
         let escrow_subgraph: &'static SubgraphClient = Box::leak(Box::new(SubgraphClient::new(
             http_client,

common/src/lib.rs

@@ -10,6 +10,7 @@ pub mod graphql;
 pub mod indexer_service;
 pub mod subgraph_client;
 pub mod tap;
+pub mod watcher;
 
 #[cfg(test)]
 mod test_vectors;