newton_prover_aggregator/

lib.rs

//! Aggregator crate

#[macro_use]
extern crate newton_prover_core;

/// Config
pub mod config;
/// Median-based consensus logic for normalizing operator responses
pub mod consensus;
/// Core aggregator logic shared between implementations
pub mod core;
/// Aggregator error
pub mod error;
/// RPC server
pub mod rpc_server;
use alloy::{
    dyn_abi::SolType,
    providers::{Provider, ProviderBuilder, WsConnect},
    rpc::types::Filter,
    sol_types::{SolEvent, SolValue},
};
use ark_ec::AffineRepr;
pub use core::AggregatorCore;
use eigensdk::client_avsregistry::reader::AvsRegistryChainReader;
pub use error::{AggregatorCoreError, AggregatorError};
use futures_util::StreamExt;
use jsonrpsee::{
    server::{RpcModule, Server},
    types::ErrorObject,
};
pub use newton_prover_chainio::bls::BlsAggregationServiceResponse;
use newton_prover_chainio::bls::ServiceHandle;
use newton_prover_core::{
    common::chain::get_block_duration_ms,
    config::{key::EcdsaKey, rpc::ChainRpcProviderConfig, NewtonAvsConfig},
    r#newton_prover_task_manager::{
        INewtonProverTaskManager::{Task, TaskResponse as BindingTaskResponse},
        NewtonProverTaskManager::NewTaskCreated,
    },
    TaskId,
};
pub use rpc_server::{
    ConsensusCommitRequest, ConsensusCommitResponse, ConsensusData, ConsensusPrepareRequest, ConsensusPrepareResponse,
    FieldAdjustment, OperatorErrorResponse, SignConsensusRequest, SignConsensusResponse, SignedTaskResponse,
    TaskResponse,
};
// Backward compatibility re-exports (deprecated)
#[allow(deprecated)]
pub use rpc_server::{
    EvaluateAndSignRequest, EvaluateAndSignResponse, FetchPolicyDataRequest, UnsignedPolicyDataResponse,
};
use std::{
    collections::HashMap,
    net::{SocketAddr, ToSocketAddrs},
    sync::Arc,
};
use tokio::{
    task::JoinHandle,
    time::{sleep, Duration},
};
use tracing::{error, info, instrument, warn, Instrument};

use crate::config::AggregatorConfig;
use newton_prover_chainio::avs::writer::AvsWriter;

// Removed TASK_CHALLENGE_WINDOW_BLOCK constant - now fetched dynamically from contract

/// Aggregator
#[allow(missing_debug_implementations)]
pub struct Aggregator {
    /// Socket address
    port_address: String,
    /// AVS writer
    avs_writer: AvsWriter,
    /// HashMap to store tasks
    tasks: HashMap<TaskId, Task>,
    /// Core aggregator logic shared with gateway aggregator
    core: Arc<AggregatorCore>,
    /// RPC provider configuration for all chains
    rpc: ChainRpcProviderConfig,
    /// Chain ID for RPC lookups (destination chain in multichain mode)
    chain_id: u64,
    /// Cached task response window block (fetched once at initialization)
    task_response_window_block: u64,
    /// Timeout for waiting for aggregation responses (in seconds)
    aggregation_timeout_secs: u64,
}

impl Aggregator {
    /// Creates a new aggregator
    ///
    /// # Arguments
    ///
    /// * [`NewtonAvsConfig<AggregatorConfig>`] - The configuration for the aggregator
    ///
    /// # Returns
    ///
    /// * `Self` - The aggregator
    ///
    /// # Errors
    ///
    /// * `AggregatorError` - The error that occurred
    pub async fn new(config: NewtonAvsConfig<AggregatorConfig>) -> Result<Self, AggregatorError> {
        let chain_id = config.chain_id;
        let source_chain_id = config.get_source_chain_id();

        // Get RPC config for the current chain (destination chain in multichain mode)
        let rpc_config = config.rpc.get_or_err(chain_id)?.clone();

        // Source chain: where EigenLayer AVS contracts (OperatorRegistry, etc.) are deployed
        let source_rpc = config.rpc.get_or_err(source_chain_id)?.clone();

        let avs_registry_chain_reader = AvsRegistryChainReader::new(
            config.contracts.avs.operator_registry,
            config.contracts.avs.operator_state_retriever,
            source_rpc.http.clone(),
        )
        .await?;

        let avs_writer = AvsWriter::new(
            config.contracts.avs.newton_prover_task_manager,
            source_rpc.http.clone(),
            config.service.signer.clone(),
        )
        .await?;

        let aggregator_address = newton_prover_core::keys::load_ecdsa(&config.service.signer)
            .map_err(|e| AggregatorError::KeyError(e.to_string()))?
            .address();
        newton_prover_metrics::set_aggregator_address(&aggregator_address.to_string());

        let core = AggregatorCore::new(
            avs_registry_chain_reader,
            config.contracts.avs.operator_registry,
            config.contracts.avs.operator_state_retriever,
            config.is_destination_chain(),
            config.rpc.clone(),
            source_chain_id,
            chain_id,
            config.contracts.avs.newton_prover_task_manager,
        )
        .await
        .map_err(|e| AggregatorError::KeyError(e.to_string()))?;

        // Fetch task response window block once at initialization (cached for reuse)
        let task_response_window_block =
            newton_prover_core::common::chain::get_task_response_window_blocks(&rpc_config.http)
                .await
                .map_err(|e| AggregatorError::KeyError(format!("Failed to get task response window: {}", e)))?
                as u64;

        Ok(Self {
            port_address: config.service.ip_address.clone(),
            avs_writer,
            tasks: HashMap::new(),
            core: Arc::new(core),
            rpc: config.rpc.clone(),
            chain_id,
            task_response_window_block,
            aggregation_timeout_secs: config.service.aggregation_timeout_secs,
        })
    }

    /// Starts the aggregator service with three tasks: one for the server,
    /// one for processing tasks and one for processing aggregator responses.
    ///
    /// # Arguments
    ///
    /// * `ws_rpc_url` - The websocket RPC URL
    ///
    /// # Returns
    ///
    /// * `eyre::Result<()>` - The result of the operation
    ///
    /// # Errors
    ///
    /// * The error that occurred
    pub async fn start(self, ws_rpc_url: String) -> eyre::Result<()> {
        info!("Starting aggregator");

        // Initialize Prometheus metrics exporter on port 9092
        let metrics_addr: SocketAddr = "0.0.0.0:9092".parse()?;
        let _metrics_handle = newton_prover_metrics::prometheus::init_prometheus_exporter(metrics_addr).await?;
        info!("Prometheus metrics server started on {}", metrics_addr);

        let Self {
            avs_writer,
            port_address,
            tasks,
            core,
            rpc: _,
            chain_id: _,
            task_response_window_block,
            aggregation_timeout_secs,
        } = self;

        let tasks = Arc::new(tokio::sync::Mutex::new(tasks));
        let core_clone = Arc::clone(&core);

        // Spawn three tasks: one for the server, one for processing tasks, and one for processing aggregator responses
        // 1) Process signatures
        let server_handle = Self::start_server(port_address, Arc::clone(&core)).await?;
        // 2) Process tasks
        let process_handle = tokio::spawn(Self::process_tasks(
            ws_rpc_url.clone(),
            task_response_window_block,
            Arc::clone(&tasks),
            Arc::clone(&core),
        ));
        // 3) Process aggregator responses
        let responses_handle = tokio::spawn(Self::process_aggregator_responses(
            Arc::clone(&tasks),
            Arc::clone(&core),
            avs_writer,
            aggregation_timeout_secs,
        ));

        // Wait for the tasks to complete and handle potential errors
        let (_server_result, process_result, responses_result) =
            tokio::try_join!(server_handle, process_handle, responses_handle)
                .inspect_err(|e| error!("Error in task execution: {e:?}"))
                .map_err(|e| eyre::eyre!("Task execution failed {e}"))?;

        process_result?;
        responses_result?;

        Ok(())
    }

    /// Starts the RPC server and returns a future that ends once the server is stopped
    async fn start_server(
        port_address: String,
        core: Arc<AggregatorCore>,
    ) -> eyre::Result<JoinHandle<()>, AggregatorError> {
        // See https://github.com/paritytech/jsonrpsee/blob/42461391fee47c94d42c4a7303355525291df9f6/examples/examples/cors_server.rs
        let mut module = RpcModule::new(core);
        module
            .register_async_method("process_signed_task_response", |params, ctx, _| async move {
                let core = ctx.as_ref();
                let signed_task_response = params
                    .parse::<SignedTaskResponse>()
                    .map_err(|err| ErrorObject::owned(0, err.to_string(), None::<()>))?;

                // Call the process_signed_response function from AggregatorCore
                let result = core.process_signed_response(signed_task_response).await;

                match result {
                    Ok(()) => Ok(true),
                    Err(err) => Err(ErrorObject::owned(0, err.to_string(), None::<()>)),
                }
            })
            .expect("method name is unique");
        let socket: SocketAddr = port_address
            .to_socket_addrs()
            .map_err(AggregatorError::IOError)?
            .next()
            .ok_or_else(|| {
                AggregatorError::IOError(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "No addresses found for hostname",
                ))
            })?;
        let middleware = tower::ServiceBuilder::new();
        let server = Server::builder().set_http_middleware(middleware).build(socket).await?;

        let handle = server.start(module);

        info!("Server running at {socket}");

        Ok(tokio::spawn(handle.stopped()))
    }

    /// Processes new tasks created in the task manager contract
    #[instrument(skip_all)]
    async fn process_tasks(
        ws_rpc_url: String,
        task_response_window_block: u64,
        tasks: Arc<tokio::sync::Mutex<HashMap<TaskId, Task>>>,
        core: Arc<AggregatorCore>,
    ) -> eyre::Result<()> {
        let ws = WsConnect::new(ws_rpc_url.clone());
        let provider = ProviderBuilder::new()
            .disable_recommended_fillers()
            .connect_ws(ws)
            .await?;
        let chain_id = provider.get_chain_id().await?;

        let filter = Filter::new().event_signature(NewTaskCreated::SIGNATURE_HASH);
        let sub = provider.subscribe_logs(&filter).await?;
        let mut stream = sub.into_stream();

        while let Some(log) = stream.next().await {
            let NewTaskCreated { taskId, task, state: _ } = log.log_decode()?.inner.data;

            let span = tracing::info_span!(
                "process_task",
                task_id = %taskId,
                policy_client = %task.policyClient
            );

            async {
                info!("aggregator: new task created");

                tasks.lock().await.insert(taskId, task.clone());

                let mut quorum_nums: Vec<u8> = vec![];
                for val in task.quorumNumbers.iter() {
                    quorum_nums.push(*val);
                }
                info!(
                    quorum_nums = %hex!(&quorum_nums),
                    "aggregator: initialize_task quorum_nums"
                );

                let block_time_ms = get_block_duration_ms(chain_id).unwrap_or(12000);
                let time_to_expiry = tokio::time::Duration::from_millis(task_response_window_block * block_time_ms);

                info!("aggregator: initializing task");
                core.initialize_task(
                    taskId,
                    task.taskCreatedBlock.into(),
                    quorum_nums,
                    task.quorumThresholdPercentage.try_into()?,
                    time_to_expiry,
                    0, // Expected operator count unknown for standalone aggregator (disables early exit)
                )
                .await
                .map_err(|e| {
                    error!("Failed to initialize task {}: {:?}", taskId, e);
                    eyre::eyre!(e)
                })?;

                Ok::<(), eyre::Error>(())
            }
            .instrument(span)
            .await?;
        }

        Ok(())
    }

    /// Processes BLS Aggregator Service responses
    ///
    /// With per-task channels, we iterate over active tasks and wait for their responses.
    /// This eliminates mutex contention and response stealing.
    async fn process_aggregator_responses(
        tasks: Arc<tokio::sync::Mutex<HashMap<TaskId, Task>>>,
        core: Arc<AggregatorCore>,
        avs_writer: AvsWriter,
        aggregation_timeout_secs: u64,
    ) -> eyre::Result<()> {
        loop {
            // Get list of active task IDs
            let active_task_ids: Vec<TaskId> = {
                let tasks_guard = tasks.lock().await;
                tasks_guard.keys().cloned().collect()
            };

            if active_task_ids.is_empty() {
                // No active tasks, wait a bit before checking again
                tokio::time::sleep(Duration::from_millis(100)).await;
                continue;
            }

            // Filter to only tasks that have receivers (not already consumed)
            // Check receivers efficiently to avoid unnecessary wait_for_aggregation calls
            let tasks_with_receivers: Vec<TaskId> = {
                let mut tasks_with_receivers = Vec::new();
                for task_id in active_task_ids {
                    if core.has_task_receiver(task_id).await {
                        tasks_with_receivers.push(task_id);
                    }
                }
                tasks_with_receivers
            };

            if tasks_with_receivers.is_empty() {
                // No tasks with receivers, wait a bit before checking again
                tokio::time::sleep(Duration::from_millis(200)).await;
                continue;
            }

            // Use a longer timeout (1-2 seconds) to reduce polling overhead
            // This is more efficient than 100ms polling and reduces latency
            // We process tasks sequentially but with a reasonable timeout per task
            let mut found_response = false;
            for task_id in tasks_with_receivers {
                // Double-check receiver still exists (may have been consumed by another call)
                if !core.has_task_receiver(task_id).await {
                    continue;
                }

                // Use configurable timeout to accommodate destination chain operations
                // (certificate verifier queries add ~500ms-1s of latency)
                let timeout = Duration::from_secs(aggregation_timeout_secs);
                match core.wait_for_aggregation(task_id, timeout).await {
                    Ok(service_response) => {
                        found_response = true;
                        info!("Received aggregated response for task_id: {}", service_response.task_id);

                        let task = match tasks.lock().await.get(&service_response.task_id).cloned() {
                            Some(task) => task,
                            None => {
                                info!("Task not found for task_id: {}", service_response.task_id);
                                continue;
                            }
                        };

                        let task_response = core.task_states.get(&task.taskId).and_then(|state| {
                            state
                                .task_responses
                                .get(&service_response.task_response_digest)
                                .cloned()
                        });

                        if let Some(task_response) = task_response {
                            if let Err(err) = Self::submit_aggregated_response_with_retry(
                                &core,
                                &avs_writer,
                                &task,
                                &task_response,
                                &service_response,
                            )
                            .await
                            {
                                error!(
                                    task_id = hex!(task.taskId),
                                    error = ?err,
                                    "Failed to submit aggregated response to contract"
                                );
                            }
                        } else {
                            info!(
                                "Not found task_response for task_id: {}",
                                hex!(service_response.task_id)
                            );
                        }
                        break; // Process one response per loop iteration
                    }
                    Err(crate::error::AggregatorCoreError::Timeout { .. }) => {
                        // Timeout is expected, continue to next task
                        continue;
                    }
                    Err(crate::error::AggregatorCoreError::TaskNotInitialized { .. }) => {
                        // Receiver already consumed (by gateway or previous call) - skip this task
                        continue;
                    }
                    Err(crate::error::AggregatorCoreError::InvalidTimeoutDuration(_)) => {
                        // Invalid timeout - this shouldn't happen, but handle gracefully
                        warn!(task_id = %task_id, "Invalid timeout duration - skipping task");
                        continue;
                    }
                    Err(e) => {
                        error!(task_id = %task_id, "Error receiving aggregated response: {:?}", e);
                        continue;
                    }
                }
            }

            if !found_response {
                // No responses found, wait a bit before checking again
                // Reduced from 50ms to avoid excessive CPU usage when no tasks are active
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
        }
    }

    async fn submit_aggregated_response_with_retry(
        core: &AggregatorCore,
        avs_writer: &AvsWriter,
        task: &Task,
        task_response: &BindingTaskResponse,
        service_response: &BlsAggregationServiceResponse,
    ) -> Result<(), AggregatorError> {
        const MAX_RETRIES: usize = 3;
        let mut attempt: usize = 0;
        let mut delay = Duration::from_millis(500);

        loop {
            attempt += 1;

            let result = core
                .submit_aggregated_response(
                    avs_writer,
                    task.clone(),
                    task_response.clone(),
                    service_response.clone(),
                )
                .await
                .map(|_| ())
                .map_err(|_| AggregatorError::SendAggregatedResponseError);

            match result {
                Ok(()) => return Ok(()),
                Err(err) => {
                    let err_str = err.to_string();
                    let should_retry =
                        err_str.contains("SendAggregatedResponseError") || err_str.contains("Failed to submit");
                    if !should_retry || attempt >= MAX_RETRIES {
                        return Err(AggregatorError::SendAggregatedResponseError);
                    }

                    warn!(
                        task_id = hex!(service_response.task_id),
                        attempt,
                        retry_in_secs = delay.as_secs(),
                        "Failed to submit aggregated response; retrying"
                    );

                    sleep(delay).await;
                    delay = (delay.saturating_mul(2)).min(Duration::from_millis(2000));
                }
            }
        }
    }

    fn should_retry_contract_submission(err: &AggregatorError) -> bool {
        matches!(err, AggregatorError::SendAggregatedResponseError)
    }
}