solana-streamer/src/streaming/yellowstone_grpc.rs

use futures::{channel::mpsc, StreamExt};
use log::error;
use solana_sdk::pubkey::Pubkey;
use std::sync::Arc;
use tokio::sync::Mutex;
use yellowstone_grpc_proto::geyser::CommitmentLevel;

use crate::common::AnyResult;
use crate::streaming::common::{
    EventBatchProcessor, MetricsManager, PerformanceMetrics, StreamClientConfig,
};
use crate::streaming::event_parser::common::filter::EventTypeFilter;
use crate::streaming::event_parser::{Protocol, UnifiedEvent};
use crate::streaming::grpc::{EventPretty, EventProcessor, StreamHandler, SubscriptionManager};

/// 交易过滤器
pub struct TransactionFilter {
    pub account_include: Vec<String>,
    pub account_exclude: Vec<String>,
    pub account_required: Vec<String>,
}

/// 账户过滤器
pub struct AccountFilter {
    pub account: Vec<String>,
    pub owner: Vec<String>,
}

#[derive(Clone)]
pub struct YellowstoneGrpc {
    pub endpoint: String,
    pub x_token: Option<String>,
    pub config: StreamClientConfig,
    pub metrics: Arc<Mutex<PerformanceMetrics>>,
    pub subscription_manager: SubscriptionManager,
    pub metrics_manager: MetricsManager,
    pub event_processor: EventProcessor,
}

impl YellowstoneGrpc {
    /// 创建客户端，使用默认配置
    pub fn new(endpoint: String, x_token: Option<String>) -> AnyResult<Self> {
        Self::new_with_config(endpoint, x_token, StreamClientConfig::default())
    }

    /// 创建客户端，使用自定义配置
    pub fn new_with_config(
        endpoint: String,
        x_token: Option<String>,
        config: StreamClientConfig,
    ) -> AnyResult<Self> {
        let _ = rustls::crypto::ring::default_provider().install_default().ok();
        let metrics = Arc::new(Mutex::new(PerformanceMetrics::new()));
        let config_arc = Arc::new(config.clone());

        let subscription_manager =
            SubscriptionManager::new(endpoint.clone(), x_token.clone(), config.clone());
        let metrics_manager =
            MetricsManager::new(metrics.clone(), config_arc.clone(), "YellowstoneGrpc".to_string());
        let event_processor = EventProcessor::new(metrics_manager.clone(), config.clone());

        Ok(Self {
            endpoint,
            x_token,
            config,
            metrics,
            subscription_manager,
            metrics_manager,
            event_processor,
        })
    }

    /// 创建高性能客户端
    pub fn new_high_performance(endpoint: String, x_token: Option<String>) -> AnyResult<Self> {
        Self::new_with_config(endpoint, x_token, StreamClientConfig::high_performance())
    }

    /// 创建低延迟客户端
    pub fn new_low_latency(endpoint: String, x_token: Option<String>) -> AnyResult<Self> {
        Self::new_with_config(endpoint, x_token, StreamClientConfig::low_latency())
    }

    /// 创建即时处理客户端
    pub fn new_immediate(endpoint: String, x_token: Option<String>) -> AnyResult<Self> {
        let mut config = StreamClientConfig::low_latency();
        config.enable_metrics = false;
        Self::new_with_config(endpoint, x_token, config)
    }

    /// 获取配置
    pub fn get_config(&self) -> &StreamClientConfig {
        &self.config
    }

    /// 更新配置
    pub fn update_config(&mut self, config: StreamClientConfig) {
        self.config = config;
    }

    /// 获取性能指标
    pub async fn get_metrics(&self) -> PerformanceMetrics {
        self.metrics_manager.get_metrics().await
    }

    /// 打印性能指标
    pub async fn print_metrics(&self) {
        self.metrics_manager.print_metrics().await;
    }

    /// 启用或禁用性能监控
    pub fn set_enable_metrics(&mut self, enabled: bool) {
        self.config.enable_metrics = enabled;
    }

    /// Simplified immediate event subscription (recommended for simple scenarios)
    ///
    /// # Parameters
    /// * `protocols` - List of protocols to monitor
    /// * `bot_wallet` - Optional bot wallet address for filtering related transactions
    /// * `transaction_filter` - Transaction filter specifying accounts to include/exclude
    /// * `account_filter` - Account filter specifying accounts and owners to monitor
    /// * `event_filter` - Optional event filter for further event filtering, no filtering if None
    /// * `commitment` - Optional commitment level, defaults to Confirmed
    /// * `callback` - Event callback function that receives parsed unified events
    ///
    /// # Returns
    /// Returns `AnyResult<()>`, `Ok(())` on success, error information on failure
    pub async fn subscribe_events_immediate<F>(
        &self,
        protocols: Vec<Protocol>,
        bot_wallet: Option<Pubkey>,
        transaction_filter: TransactionFilter,
        account_filter: AccountFilter,
        event_type_filter: Option<EventTypeFilter>,
        commitment: Option<CommitmentLevel>,
        callback: F,
    ) -> AnyResult<()>
    where
        F: Fn(Box<dyn UnifiedEvent>) + Send + Sync + 'static,
    {
        // 启动自动性能监控（如果启用）
        if self.config.enable_metrics {
            self.metrics_manager.start_auto_monitoring().await;
        }

        let transactions = self.subscription_manager.get_subscribe_request_filter(
            transaction_filter.account_include,
            transaction_filter.account_exclude,
            transaction_filter.account_required,
            event_type_filter.clone(),
        );
        let accounts = self.subscription_manager.subscribe_with_account_request(
            account_filter.account,
            account_filter.owner,
            event_type_filter.clone(),
        );

        // 订阅事件
        let (mut subscribe_tx, mut stream) = self
            .subscription_manager
            .subscribe_with_request(transactions, accounts, commitment, event_type_filter.clone())
            .await?;

        // 创建通道，使用配置中的通道大小
        let (mut tx, mut rx) = mpsc::channel::<EventPretty>(self.config.backpressure.channel_size);

        // 启动流处理任务
        let backpressure_strategy = self.config.backpressure.strategy;
        tokio::spawn(async move {
            while let Some(message) = stream.next().await {
                match message {
                    Ok(msg) => {
                        if let Err(e) = StreamHandler::handle_stream_message(
                            msg,
                            &mut tx,
                            &mut subscribe_tx,
                            backpressure_strategy,
                        )
                        .await
                        {
                            error!("Error handling message: {e:?}");
                            break;
                        }
                    }
                    Err(error) => {
                        error!("Stream error: {error:?}");
                        break;
                    }
                }
            }
        });

        // 即时处理交易，无批处理
        let event_processor = self.event_processor.clone();
        tokio::spawn(async move {
            while let Some(event_pretty) = rx.next().await {
                if let Err(e) = event_processor
                    .process_event_transaction_with_metrics(
                        event_pretty,
                        &callback,
                        bot_wallet,
                        protocols.clone(),
                        event_type_filter.clone(),
                    )
                    .await
                {
                    error!("Error processing transaction: {e:?}");
                }
            }
        });

        tokio::signal::ctrl_c().await?;
        Ok(())
    }

    /// Advanced event subscription with batch processing and backpressure handling
    ///
    /// # Parameters
    /// * `protocols` - List of protocols to monitor
    /// * `bot_wallet` - Optional bot wallet address for filtering related transactions
    /// * `transaction_filter` - Transaction filter specifying accounts to include/exclude
    /// * `account_filter` - Account filter specifying accounts and owners to monitor
    /// * `event_filter` - Optional event filter for further event filtering, no filtering if None
    /// * `commitment` - Optional commitment level, defaults to Confirmed
    /// * `callback` - Event callback function that receives parsed unified events
    ///
    /// # Features
    /// * Batch processing for improved throughput
    /// * Backpressure handling to prevent memory overflow
    /// * Automatic performance monitoring (if enabled)
    /// * Configurable batch size and timeout
    ///
    /// # Returns
    /// Returns `AnyResult<()>`, `Ok(())` on success, error information on failure
    pub async fn subscribe_events_advanced<F>(
        &self,
        protocols: Vec<Protocol>,
        bot_wallet: Option<Pubkey>,
        transaction_filter: TransactionFilter,
        account_filter: AccountFilter,
        event_type_filter: Option<EventTypeFilter>,
        commitment: Option<CommitmentLevel>,
        callback: F,
    ) -> AnyResult<()>
    where
        F: Fn(Box<dyn UnifiedEvent>) + Send + Sync + 'static,
    {
        // 启动自动性能监控（如果启用）
        if self.config.enable_metrics {
            self.metrics_manager.start_auto_monitoring().await;
        }

        let transactions = self.subscription_manager.get_subscribe_request_filter(
            transaction_filter.account_include,
            transaction_filter.account_exclude,
            transaction_filter.account_required,
            event_type_filter.clone(),
        );
        let accounts = self.subscription_manager.subscribe_with_account_request(
            account_filter.account,
            account_filter.owner,
            event_type_filter.clone(),
        );

        // Subscribe to events
        let (mut subscribe_tx, mut stream) = self
            .subscription_manager
            .subscribe_with_request(transactions, accounts, commitment, event_type_filter.clone())
            .await?;

        // Create channel
        let (mut tx, mut rx) = mpsc::channel::<EventPretty>(self.config.backpressure.channel_size);

        // 创建批处理器，将单个事件回调转换为批量回调
        let batch_callback = move |events: Vec<Box<dyn UnifiedEvent>>| {
            for event in events {
                callback(event);
            }
        };

        let mut batch_processor = EventBatchProcessor::new(
            batch_callback,
            self.config.batch.batch_size,
            self.config.batch.batch_timeout_ms,
        );

        // Start task to process the stream
        let backpressure_strategy = self.config.backpressure.strategy;
        tokio::spawn(async move {
            while let Some(message) = stream.next().await {
                match message {
                    Ok(msg) => {
                        if let Err(e) = StreamHandler::handle_stream_message(
                            msg,
                            &mut tx,
                            &mut subscribe_tx,
                            backpressure_strategy,
                        )
                        .await
                        {
                            error!("Error handling message: {e:?}");
                            break;
                        }
                    }
                    Err(error) => {
                        error!("Stream error: {error:?}");
                        break;
                    }
                }
            }
        });

        // Process transactions with batch processing
        let event_processor = self.event_processor.clone();
        tokio::spawn(async move {
            while let Some(event_pretty) = rx.next().await {
                if let Err(e) = event_processor
                    .process_event_transaction_with_batch(
                        event_pretty,
                        &mut batch_processor,
                        bot_wallet,
                        protocols.clone(),
                        event_type_filter.clone(),
                    )
                    .await
                {
                    error!("Error processing transaction: {e:?}");
                }
            }

            // 处理剩余的事件
            batch_processor.flush();
        });

        tokio::signal::ctrl_c().await?;
        Ok(())
    }
}

// 实现 Clone trait 以支持模块间共享
impl Clone for EventProcessor {
    fn clone(&self) -> Self {
        Self {
            metrics_manager: self.metrics_manager.clone(),
            config: self.config.clone(),
            parser_cache: self.parser_cache.clone(),
        }
    }
}