newton-chain-watcher 0.5.2

//! WebSocket event listener for the chain watcher
//!
//! Monitors `NewTaskCreated`, `OperatorAddedToOperatorSet`, and
//! `OperatorRemovedFromOperatorSet` events via WebSocket subscription.
//! Follows the reconnection pattern from `crates/challenger/src/dest/listener.rs`.

use alloy::{
    primitives::Address,
    providers::{Provider, ProviderBuilder},
    rpc::types::Filter,
    sol_types::SolEvent,
    transports::ws::WsConnect,
};
use futures::StreamExt;
use newton_core::{
    config::rpc::{
        ws_reconnect_loop, WsReconnectCallbacks, WEBSOCKET_DROPPED_MAX_RETRY, WEBSOCKET_DROPPED_RETRY_DELAY,
    },
    identity_registry::IdentityRegistry::IdentityBound,
    newton_prover_task_manager::NewtonProverTaskManager::NewTaskCreated,
};
use newton_metric::{inc_chain_watcher_websocket_reconnections, set_chain_watcher_websocket_connected};
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::{error, info};

use crate::event::{ChainEvent, ChainEventType};

// Re-use the transporter's OperatorSet event definitions since they come from
// EigenLayer's AllocationManager (not Newton contracts), so there are no
// generated bindings for them.
alloy::sol! {
    #[allow(missing_docs)]
    struct OperatorSet {
        address avs;
        uint32 id;
    }

    #[allow(missing_docs)]
    event OperatorAddedToOperatorSet(
        address indexed operator,
        OperatorSet operatorSet
    );

    #[allow(missing_docs)]
    event OperatorRemovedFromOperatorSet(
        address indexed operator,
        OperatorSet operatorSet
    );
}

/// Metric callbacks for the chain watcher WebSocket reconnection loop.
struct ChainWatcherWsCallbacks {
    chain_id: u64,
}

impl WsReconnectCallbacks for ChainWatcherWsCallbacks {
    fn on_disconnected(&self) {
        set_chain_watcher_websocket_connected(self.chain_id, false);
    }

    fn on_reconnect_attempt(&self) {
        inc_chain_watcher_websocket_reconnections(self.chain_id);
    }
}

/// listens for on-chain events relevant to the chain watcher
#[derive(Debug)]
pub struct ChainWatcherListener {
    /// chain id being watched
    pub chain_id: u64,
    /// websocket rpc url
    pub ws_url: String,
    /// task manager contract address (for NewTaskCreated events)
    pub task_manager: Address,
    /// allocation manager contract address (for operator set events)
    pub allocation_manager: Address,
    /// identity registry contract address (for IdentityBound events)
    pub identity_registry: Address,
}

impl ChainWatcherListener {
    /// create a new chain watcher listener
    pub fn new(
        chain_id: u64,
        ws_url: String,
        task_manager: Address,
        allocation_manager: Address,
        identity_registry: Address,
    ) -> Self {
        Self {
            chain_id,
            ws_url,
            task_manager,
            allocation_manager,
            identity_registry,
        }
    }

    /// start listening for events, returns a receiver for chain events
    pub async fn start(self, cancel: CancellationToken) -> mpsc::Receiver<ChainEvent> {
        let (tx, rx) = mpsc::channel(100);

        let ws_url = self.ws_url.clone();
        let task_manager = self.task_manager;
        let allocation_manager = self.allocation_manager;
        let identity_registry = self.identity_registry;
        let chain_id = self.chain_id;

        tokio::spawn(async move {
            let callbacks = ChainWatcherWsCallbacks { chain_id };
            let fatal = ws_reconnect_loop(
                "chain watcher",
                WEBSOCKET_DROPPED_MAX_RETRY,
                std::time::Duration::from_secs(WEBSOCKET_DROPPED_RETRY_DELAY),
                || {
                    Self::listen_to_events(
                        ws_url.clone(),
                        task_manager,
                        allocation_manager,
                        identity_registry,
                        chain_id,
                        tx.clone(),
                    )
                },
                callbacks,
            )
            .await;
            if fatal {
                cancel.cancel();
            }
        });

        rx
    }

    async fn listen_to_events(
        ws_url: String,
        task_manager: Address,
        allocation_manager: Address,
        identity_registry: Address,
        chain_id: u64,
        tx: mpsc::Sender<ChainEvent>,
    ) -> eyre::Result<()> {
        let ws = WsConnect::new(ws_url);
        let provider = ProviderBuilder::new()
            .disable_recommended_fillers()
            .connect_ws(ws)
            .await?;

        set_chain_watcher_websocket_connected(chain_id, true);

        // Subscribe to NewTaskCreated from the TaskManager
        let task_filter = Filter::new()
            .address(task_manager)
            .event_signature(NewTaskCreated::SIGNATURE_HASH);

        // Subscribe to operator set events from the AllocationManager
        let operator_filter = Filter::new().address(allocation_manager).events([
            OperatorAddedToOperatorSet::SIGNATURE,
            OperatorRemovedFromOperatorSet::SIGNATURE,
        ]);

        // Subscribe to IdentityBound events from the IdentityRegistry
        let identity_filter = Filter::new()
            .address(identity_registry)
            .event_signature(IdentityBound::SIGNATURE_HASH);

        let task_sub = provider.subscribe_logs(&task_filter).await?;
        let operator_sub = provider.subscribe_logs(&operator_filter).await?;
        let identity_sub = provider.subscribe_logs(&identity_filter).await?;
        let mut stream = futures::stream::select_all([
            task_sub.into_stream(),
            operator_sub.into_stream(),
            identity_sub.into_stream(),
        ]);

        set_chain_watcher_websocket_connected(chain_id, true);
        info!(chain_id, "listening to chain watcher events");

        while let Some(log) = stream.next().await {
            let block_number = log.block_number.unwrap_or(0);
            let tx_hash = log.transaction_hash.unwrap_or_default();

            if let Ok(decoded) = log.log_decode::<NewTaskCreated>() {
                let NewTaskCreated { taskId, task, .. } = decoded.inner.data;
                info!(chain_id, task_id = %taskId, "new task created event");

                if let Err(e) = tx
                    .send(ChainEvent {
                        chain_id,
                        event_type: ChainEventType::DirectOnchainTask {
                            task_id: taskId,
                            task: Box::new(task),
                        },
                        block_number,
                        tx_hash,
                    })
                    .await
                {
                    error!("failed to send new task event: {e}");
                }
            } else if let Ok(decoded) = log.log_decode::<OperatorAddedToOperatorSet>() {
                let OperatorAddedToOperatorSet { operator, operatorSet } = decoded.inner.data;
                info!(
                    chain_id,
                    %operator,
                    operator_set_id = operatorSet.id,
                    "operator added to set"
                );

                if let Err(e) = tx
                    .send(ChainEvent {
                        chain_id,
                        event_type: ChainEventType::OperatorAdded {
                            operator,
                            operator_set_avs: operatorSet.avs,
                            operator_set_id: operatorSet.id,
                        },
                        block_number,
                        tx_hash,
                    })
                    .await
                {
                    error!("failed to send operator added event: {e}");
                }
            } else if let Ok(decoded) = log.log_decode::<OperatorRemovedFromOperatorSet>() {
                let OperatorRemovedFromOperatorSet { operator, operatorSet } = decoded.inner.data;
                info!(
                    chain_id,
                    %operator,
                    operator_set_id = operatorSet.id,
                    "operator removed from set"
                );

                if let Err(e) = tx
                    .send(ChainEvent {
                        chain_id,
                        event_type: ChainEventType::OperatorRemoved {
                            operator,
                            operator_set_avs: operatorSet.avs,
                            operator_set_id: operatorSet.id,
                        },
                        block_number,
                        tx_hash,
                    })
                    .await
                {
                    error!("failed to send operator removed event: {e}");
                }
            } else if let Ok(decoded) = log.log_decode::<IdentityBound>() {
                let IdentityBound {
                    identityOwner,
                    identityDomain,
                    identityData,
                } = decoded.inner.data;
                info!(chain_id, %identityOwner, "identity data bound event");

                if let Err(e) = tx
                    .send(ChainEvent {
                        chain_id,
                        event_type: ChainEventType::IdentityDataBound {
                            identity_owner: identityOwner,
                            identity_domain: identityDomain,
                            data_ref_id: identityData,
                        },
                        block_number,
                        tx_hash,
                    })
                    .await
                {
                    error!("failed to send identity data bound event: {e}");
                }
            }
        }

        Ok(())
    }
}