refactor: use tokio::watch for attestation_signers (#387)

Author: taslimmuhammed

common/src/attestations/dispute_manager.rs

@@ -44,12 +44,9 @@ pub fn dispute_manager(
             .await;
 
             match result {
-                Ok(address) => {
-                    if tx.send(Some(address)).is_err() {
-                        // stopping
-                        break;
-                    }
-                }
+                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

common/src/attestations/signers.rs

@@ -1,85 +1,100 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use eventuals::{Eventual, EventualExt, EventualWriter};
+use eventuals::{Eventual, EventualExt};
 use std::collections::HashMap;
 use std::sync::Arc;
 use thegraph_core::{Address, ChainId};
-use tokio::sync::watch;
 use tokio::{
     select,
-    sync::{watch::Receiver, Mutex},
+    sync::{
+        watch::{self, Receiver},
+        Mutex,
+    },
 };
 use tracing::warn;
 
 use crate::prelude::{Allocation, AttestationSigner};
 
 /// An always up-to-date list of attestation signers, one for each of the indexer's allocations.
-pub fn attestation_signers(
+pub async fn attestation_signers(
     indexer_allocations: Eventual<HashMap<Address, Allocation>>,
     indexer_mnemonic: String,
     chain_id: ChainId,
     mut dispute_manager_rx: Receiver<Option<Address>>,
-) -> Eventual<HashMap<Address, AttestationSigner>> {
+) -> Receiver<HashMap<Address, AttestationSigner>> {
     let attestation_signers_map: &'static Mutex<HashMap<Address, AttestationSigner>> =
         Box::leak(Box::new(Mutex::new(HashMap::new())));
 
+    // Actively listening to indexer_allocations to update allocations channel
+    // Temporary fix until the indexer_allocations is migrated to tokio watch
+    let (allocations_tx, mut allocations_rx) =
+        watch::channel(indexer_allocations.value_immediate().unwrap_or_default());
+    indexer_allocations
+        .pipe(move |allocatons| {
+            allocations_tx
+                .send(allocatons)
+                .expect("Failed to update allocations channel");
+        })
+        .forever();
+
+    let starter_signers_map = modify_sigers(
+        Arc::new(indexer_mnemonic.clone()),
+        chain_id,
+        attestation_signers_map,
+        allocations_rx.clone(),
+        dispute_manager_rx.clone(),
+    )
+    .await;
+
     // Whenever the indexer's active or recently closed allocations change, make sure
     // we have attestation signers for all of them.
-    let (mut signers_writer, signers_reader) =
-        Eventual::<HashMap<Address, AttestationSigner>>::new();
-
+    let (signers_tx, signers_rx) = watch::channel(starter_signers_map);
     tokio::spawn(async move {
-        // Listening to the allocation eventual and converting them to reciever.
-        // Using pipe for updation.
-        // For temporary pupose only.
-        let (allocations_tx, mut allocations_rx) =
-            watch::channel(indexer_allocations.value().await.unwrap());
-        let _p1 = indexer_allocations.pipe(move |allocatons| {
-            let _ = allocations_tx.send(allocatons);
-        });
-
         loop {
-            select! {
-                Ok(_)= allocations_rx.changed() =>{
+            let updated_signers = select! {
+                Ok(())= allocations_rx.changed() =>{
                     modify_sigers(
                         Arc::new(indexer_mnemonic.clone()),
                         chain_id,
                         attestation_signers_map,
                         allocations_rx.clone(),
                         dispute_manager_rx.clone(),
-                        &mut signers_writer).await;
+                    ).await
                 },
-                Ok(_)= dispute_manager_rx.changed() =>{
-                    modify_sigers(Arc::new(indexer_mnemonic.clone()),
-                    chain_id,
-                    attestation_signers_map,
-                    allocations_rx.clone(),
-                    dispute_manager_rx.clone(),
-                    &mut signers_writer).await;
+                Ok(())= dispute_manager_rx.changed() =>{
+                    modify_sigers(
+                        Arc::new(indexer_mnemonic.clone()),
+                        chain_id,
+                        attestation_signers_map,
+                        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_reader
+    signers_rx
 }
 async fn modify_sigers(
     indexer_mnemonic: Arc<String>,
     chain_id: ChainId,
     attestation_signers_map: &'static Mutex<HashMap<Address, AttestationSigner>>,
     allocations_rx: Receiver<HashMap<Address, Allocation>>,
     dispute_manager_rx: Receiver<Option<Address>>,
-    signers_writer: &mut EventualWriter<HashMap<Address, AttestationSigner>>,
-) {
+) -> 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;
+        return signers.clone();
     };
     // Remove signers for allocations that are no longer active or recently closed
     signers.retain(|id, _| allocations.contains_key(id));
@@ -101,7 +116,7 @@ async fn modify_sigers(
         }
     }
 
-    signers_writer.write(signers.clone());
+    signers.clone()
 }
 
 #[cfg(test)]
@@ -115,29 +130,30 @@ mod tests {
     #[tokio::test]
     async fn test_attestation_signers_update_with_allocations() {
         let (mut allocations_writer, allocations) = Eventual::<HashMap<Address, Allocation>>::new();
-        let (dispute_manager_writer, dispute_manager) = watch::channel(None);
-
-        dispute_manager_writer
+        let (dispute_manager_tx, dispute_manager_rx) = watch::channel(None);
+        dispute_manager_tx
             .send(Some(*DISPUTE_MANAGER_ADDRESS))
             .unwrap();
-
-        let signers = attestation_signers(
+        let mut signers = attestation_signers(
             allocations,
             (*INDEXER_OPERATOR_MNEMONIC).to_string(),
             1,
-            dispute_manager,
-        );
-        let mut signers = signers.subscribe();
+            dispute_manager_rx,
+        )
+        .await;
 
         // Test that an empty set of allocations leads to an empty set of signers
         allocations_writer.write(HashMap::new());
-        let latest_signers = signers.next().await.unwrap();
+        signers.changed().await.unwrap();
+        let latest_signers = signers.borrow().clone();
         assert_eq!(latest_signers, HashMap::new());
 
         // Test that writing our set of test allocations results in corresponding signers for all of them
         allocations_writer.write((*INDEXER_ALLOCATIONS).clone());
-        let latest_signers = signers.next().await.unwrap();
+        signers.changed().await.unwrap();
+        let latest_signers = signers.borrow().clone();
         assert_eq!(latest_signers.len(), INDEXER_ALLOCATIONS.len());
+
         for signer_allocation_id in latest_signers.keys() {
             assert!(INDEXER_ALLOCATIONS
                 .keys()

common/src/indexer_service/http/indexer_service.rs

@@ -28,6 +28,7 @@ use thegraph_core::{Address, Attestation, DeploymentId};
 use thiserror::Error;
 use tokio::net::TcpListener;
 use tokio::signal;
+use tokio::sync::watch::Receiver;
 use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
 use tower_http::{cors, cors::CorsLayer, normalize_path::NormalizePath, trace::TraceLayer};
 use tracing::error;
@@ -182,7 +183,7 @@ where
     I: IndexerServiceImpl + Sync + Send + 'static,
 {
     pub config: IndexerServiceConfig,
-    pub attestation_signers: Eventual<HashMap<Address, AttestationSigner>>,
+    pub attestation_signers: Receiver<HashMap<Address, AttestationSigner>>,
     pub tap_manager: Manager<IndexerTapContext>,
     pub service_impl: Arc<I>,
 
@@ -248,7 +249,8 @@ impl IndexerService {
             options.config.indexer.operator_mnemonic.clone(),
             options.config.graph_network.chain_id,
             dispute_manager,
-        );
+        )
+        .await;
 
         let escrow_subgraph: &'static SubgraphClient = Box::leak(Box::new(SubgraphClient::new(
             http_client,

common/src/indexer_service/http/request_handler.rs

@@ -155,12 +155,9 @@ where
         .map_err(IndexerServiceError::ReceiptError)?;
 
     // Check if we have an attestation signer for the allocation the receipt was created for
-    let signers = state
+    let signer = state
         .attestation_signers
-        .value_immediate()
-        .ok_or_else(|| IndexerServiceError::ServiceNotReady)?;
-
-    let signer = signers
+        .borrow()
         .get(&allocation_id)
         .cloned()
         .ok_or_else(|| (IndexerServiceError::NoSignerForAllocation(allocation_id)))?;