use crate::common::AnyResult; use crate::streaming::common::{ parse_grpc_transaction_events, process_grpc_transaction, transaction_metrics_callback, MetricsManager, MicroBatchBuffer, PerformanceMetrics, SlotBuffer, StreamClientConfig, SubscriptionHandle, }; use crate::streaming::event_parser::common::filter::EventTypeFilter; use crate::streaming::event_parser::{DexEvent, Protocol}; use crate::streaming::grpc::pool::factory; use crate::streaming::grpc::{EventPretty, SubscriptionManager, TransactionPretty}; use anyhow::anyhow; use futures::channel::mpsc; use futures::{SinkExt, StreamExt}; use log::error; use sol_parser_sdk::grpc::OrderMode; use solana_sdk::pubkey::Pubkey; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::Mutex; use tokio::time::{Duration, Instant}; use yellowstone_grpc_proto::geyser::subscribe_update::UpdateOneof; use yellowstone_grpc_proto::geyser::{ CommitmentLevel, SubscribeRequest, SubscribeRequestFilterAccountsFilter, SubscribeRequestPing, }; /// 交易过滤器 #[derive(Debug, Clone)] pub struct TransactionFilter { pub account_include: Vec, pub account_exclude: Vec, pub account_required: Vec, } /// 账户过滤器 #[derive(Debug, Clone)] pub struct AccountFilter { pub account: Vec, pub owner: Vec, pub filters: Vec, } pub struct YellowstoneGrpc { pub endpoint: String, pub x_token: Option, pub config: StreamClientConfig, pub subscription_manager: SubscriptionManager, pub subscription_handle: Arc>>, // Dynamic subscription management fields pub active_subscription: Arc, pub control_tx: Arc>>>, pub current_request: Arc>>, pub event_type_filter: Arc>>, } impl YellowstoneGrpc { /// 创建客户端,使用默认配置 pub fn new(endpoint: String, x_token: Option) -> AnyResult { Self::new_with_config(endpoint, x_token, StreamClientConfig::default()) } /// 创建客户端,使用自定义配置 pub fn new_with_config( endpoint: String, x_token: Option, config: StreamClientConfig, ) -> AnyResult { sol_parser_sdk::warmup_parser(); let _ = rustls::crypto::ring::default_provider().install_default().ok(); let subscription_manager = SubscriptionManager::new(endpoint.clone(), x_token.clone(), config.clone()); MetricsManager::init(config.enable_metrics); Ok(Self { endpoint, x_token, config, subscription_manager, subscription_handle: Arc::new(Mutex::new(None)), active_subscription: Arc::new(AtomicBool::new(false)), control_tx: Arc::new(tokio::sync::Mutex::new(None)), current_request: Arc::new(tokio::sync::RwLock::new(None)), event_type_filter: Arc::new(tokio::sync::RwLock::new(None)), }) } /// 获取配置 pub fn get_config(&self) -> &StreamClientConfig { &self.config } /// 更新配置 pub fn update_config(&mut self, config: StreamClientConfig) { self.config = config; } /// 获取性能指标 pub fn get_metrics(&self) -> PerformanceMetrics { MetricsManager::global().get_metrics() } /// 打印性能指标 pub fn print_metrics(&self) { MetricsManager::global().print_metrics(); } /// 启用或禁用性能监控 pub fn set_enable_metrics(&mut self, enabled: bool) { self.config.enable_metrics = enabled; } /// 停止当前订阅 pub async fn stop(&self) { let mut handle_guard = self.subscription_handle.lock().await; if let Some(handle) = handle_guard.take() { handle.stop(); } *self.control_tx.lock().await = None; *self.current_request.write().await = None; self.active_subscription.store(false, Ordering::Release); } /// 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 #[allow(clippy::too_many_arguments)] pub async fn subscribe_events_immediate( &self, protocols: Vec, bot_wallet: Option, transaction_filter: Vec, account_filter: Vec, event_type_filter: Option, commitment: Option, callback: F, ) -> AnyResult<()> where F: Fn(DexEvent) + Send + Sync + 'static, { *self.event_type_filter.write().await = event_type_filter.clone(); if self .active_subscription .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) .is_err() { return Err(anyhow!("Already subscribed. Use update_subscription() to modify filters")); } let mut start_guard = SubscriptionStartGuard::new(self.active_subscription.clone()); // 启动自动性能监控(如果启用) if self.config.enable_metrics { start_guard.set_metrics_handle(MetricsManager::global().start_auto_monitoring().await); } let transactions = self .subscription_manager .get_subscribe_request_filter(transaction_filter, event_type_filter.as_ref()); let accounts = self .subscription_manager .subscribe_with_account_request(account_filter, event_type_filter.as_ref()); // 订阅事件 let (subscribe_tx, mut stream, subscribe_request) = self .subscription_manager .subscribe_with_request(transactions, accounts, commitment, event_type_filter.as_ref()) .await?; // 用 Arc> 包装 subscribe_tx 以支持多线程共享 let subscribe_tx = Arc::new(Mutex::new(subscribe_tx)); *self.current_request.write().await = Some(subscribe_request); let (control_tx, mut control_rx) = mpsc::channel(100); *self.control_tx.lock().await = Some(control_tx); // Wrap callback once before the async block let callback = Arc::new(callback); let transaction_callback = transaction_metrics_callback(callback.clone()); let order_mode = self.config.order_mode; let order_timeout_ms = self.config.order_timeout_ms; let micro_batch_us = self.config.micro_batch_us; let active_subscription = self.active_subscription.clone(); let control_tx_cleanup = self.control_tx.clone(); let current_request_cleanup = self.current_request.clone(); let metrics_handle = start_guard.disarm(); let stream_handle = tokio::spawn(async move { let cleanup = SubscriptionCleanupGuard::new( active_subscription, control_tx_cleanup, current_request_cleanup, metrics_handle, ); let mut slot_buffer = SlotBuffer::new(); let mut micro_batch = MicroBatchBuffer::new(); let mut last_slot = 0u64; let check_interval = match order_mode { OrderMode::MicroBatch => Duration::from_micros(micro_batch_us.max(1)), _ => Duration::from_millis((order_timeout_ms / 2).max(1)), }; let mut next_check = Instant::now() + check_interval; loop { if has_buffered_events(order_mode, &slot_buffer, µ_batch) { flush_ordered_timeouts( order_mode, &mut slot_buffer, &mut micro_batch, transaction_callback.clone(), order_timeout_ms, micro_batch_us, &mut next_check, check_interval, ); } tokio::select! { message = stream.next() => { match message { Some(Ok(msg)) => { let created_at = msg.created_at; match msg.update_oneof { Some(UpdateOneof::Account(account)) => { let Some(account_pretty) = factory::try_create_account_pretty_pooled(account) else { continue; }; log::debug!("Received account: {:?}", account_pretty); if let Err(e) = process_grpc_transaction( EventPretty::Account(account_pretty), &protocols, event_type_filter.as_ref(), callback.clone(), bot_wallet, ) .await { error!("Error processing account event: {e:?}"); } } Some(UpdateOneof::BlockMeta(sut)) => { let block_meta_pretty = factory::create_block_meta_pretty_pooled(sut, created_at); log::debug!("Received block meta: {:?}", block_meta_pretty); if let Err(e) = process_grpc_transaction( EventPretty::BlockMeta(block_meta_pretty), &protocols, event_type_filter.as_ref(), callback.clone(), bot_wallet, ) .await { error!("Error processing block meta event: {e:?}"); } } Some(UpdateOneof::Transaction(sut)) => { let Some(transaction_pretty) = factory::try_create_transaction_pretty_pooled(sut, created_at) else { continue; }; log::debug!( "Received transaction: {} at slot {}", transaction_pretty.signature, transaction_pretty.slot ); handle_ordered_transaction( transaction_pretty, &protocols, event_type_filter.as_ref(), transaction_callback.clone(), bot_wallet, order_mode, &mut slot_buffer, &mut micro_batch, &mut last_slot, micro_batch_us, ); } Some(UpdateOneof::Ping(_)) => { // 只在需要时获取锁,并立即释放 if let Ok(mut tx_guard) = subscribe_tx.try_lock() { let _ = tx_guard .send(SubscribeRequest { ping: Some(SubscribeRequestPing { id: 1 }), ..Default::default() }) .await; } let ts = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); log::debug!("service is ping: {}", ts); } Some(UpdateOneof::Pong(_)) => { let ts = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); log::debug!("service is pong: {}", ts); } _ => { log::debug!("Received other message type"); } } } Some(Err(error)) => { error!("Stream error: {error:?}"); flush_ordered_on_disconnect( order_mode, &mut slot_buffer, &mut micro_batch, transaction_callback.clone(), ); break; } None => { flush_ordered_on_disconnect( order_mode, &mut slot_buffer, &mut micro_batch, transaction_callback.clone(), ); break; } } } Some(update) = control_rx.next() => { if let Err(e) = subscribe_tx.lock().await.send(update).await { error!("Failed to send subscription update: {}", e); flush_ordered_on_disconnect( order_mode, &mut slot_buffer, &mut micro_batch, transaction_callback.clone(), ); break; } } _ = tokio::time::sleep_until(next_check), if has_buffered_events(order_mode, &slot_buffer, µ_batch) => { flush_ordered_timeouts( order_mode, &mut slot_buffer, &mut micro_batch, transaction_callback.clone(), order_timeout_ms, micro_batch_us, &mut next_check, check_interval, ); } } } cleanup.finish().await; }); // 保存订阅句柄 let subscription_handle = SubscriptionHandle::new(stream_handle, None, None); let mut handle_guard = self.subscription_handle.lock().await; *handle_guard = Some(subscription_handle); Ok(()) } /// Update subscription filters at runtime without reconnection /// /// # Parameters /// * `transaction_filter` - New transaction filter to apply /// * `account_filter` - New account filter to apply /// /// # Returns /// Returns `AnyResult<()>` on success, error on failure pub async fn update_subscription( &self, transaction_filter: Vec, account_filter: Vec, ) -> AnyResult<()> { let mut control_sender = { let control_guard = self.control_tx.lock().await; if !self.active_subscription.load(Ordering::Acquire) { return Err(anyhow!("No active subscription to update")); } control_guard .as_ref() .ok_or_else(|| anyhow!("No active subscription to update"))? .clone() }; let mut request = self .current_request .read() .await .as_ref() .ok_or_else(|| anyhow!("No active subscription"))? .clone(); request.transactions = self .subscription_manager .get_subscribe_request_filter( transaction_filter, self.event_type_filter.read().await.as_ref(), ) .unwrap_or_default(); request.accounts = self .subscription_manager .subscribe_with_account_request( account_filter, self.event_type_filter.read().await.as_ref(), ) .unwrap_or_default(); control_sender .send(request.clone()) .await .map_err(|e| anyhow!("Failed to send update: {}", e))?; *self.current_request.write().await = Some(request); Ok(()) } } #[allow(clippy::too_many_arguments)] fn handle_ordered_transaction( transaction_pretty: TransactionPretty, protocols: &[Protocol], event_type_filter: Option<&EventTypeFilter>, callback: Arc, bot_wallet: Option, order_mode: OrderMode, slot_buffer: &mut SlotBuffer, micro_batch: &mut MicroBatchBuffer, last_slot: &mut u64, micro_batch_us: u64, ) { let fallback_slot = transaction_pretty.slot; let fallback_tx_index = transaction_pretty.tx_index.unwrap_or(0); let recv_us = transaction_pretty.recv_us; let events = parse_grpc_transaction_events(transaction_pretty, protocols, event_type_filter, bot_wallet); match order_mode { OrderMode::Unordered => { for event in events { callback(event); } } OrderMode::Ordered => { if fallback_slot > *last_slot && *last_slot > 0 { deliver_events(callback.clone(), slot_buffer.flush_before(fallback_slot)); } *last_slot = fallback_slot; for event in events { let (slot, tx_index) = event_order_key(&event, fallback_slot, fallback_tx_index); slot_buffer.push(slot, tx_index, event); } } OrderMode::StreamingOrdered => { // Events parsed from one transaction share its (slot, tx_index); push them as a // single group so the streaming watermark advances per transaction and no event of // a multi-event transaction is dropped. for ((slot, tx_index), group) in group_events_by_order_key(events, fallback_slot, fallback_tx_index) { deliver_events(callback.clone(), slot_buffer.push_streaming(slot, tx_index, group)); } } OrderMode::MicroBatch => { for event in events { let (slot, tx_index) = event_order_key(&event, fallback_slot, fallback_tx_index); if micro_batch.push(slot, tx_index, event, recv_us, micro_batch_us) { deliver_events(callback.clone(), micro_batch.flush()); } } } } } #[allow(clippy::too_many_arguments)] fn flush_ordered_timeouts( order_mode: OrderMode, slot_buffer: &mut SlotBuffer, micro_batch: &mut MicroBatchBuffer, callback: Arc, order_timeout_ms: u64, micro_batch_us: u64, next_check: &mut Instant, check_interval: Duration, ) { if Instant::now() < *next_check { return; } *next_check = Instant::now() + check_interval; match order_mode { OrderMode::Ordered => { if slot_buffer.should_timeout(order_timeout_ms) { deliver_events(callback, slot_buffer.flush_all()); } } OrderMode::StreamingOrdered => { if slot_buffer.should_timeout(order_timeout_ms) { deliver_events(callback, slot_buffer.flush_streaming_timeout()); } } OrderMode::MicroBatch => { let now_us = crate::streaming::event_parser::common::high_performance_clock::get_high_perf_clock( ); if micro_batch.should_flush(now_us, micro_batch_us) { deliver_events(callback, micro_batch.flush()); } } OrderMode::Unordered => {} } } fn flush_ordered_on_disconnect( order_mode: OrderMode, slot_buffer: &mut SlotBuffer, micro_batch: &mut MicroBatchBuffer, callback: Arc, ) { match order_mode { OrderMode::Ordered => deliver_events(callback, slot_buffer.flush_all()), OrderMode::StreamingOrdered => { deliver_events(callback, slot_buffer.flush_streaming_timeout()) } OrderMode::MicroBatch => deliver_events(callback, micro_batch.flush()), OrderMode::Unordered => {} } } #[inline] fn event_order_key(event: &DexEvent, fallback_slot: u64, fallback_tx_index: u64) -> (u64, u64) { let metadata = event.metadata(); (metadata.slot.max(fallback_slot), metadata.tx_index.unwrap_or(fallback_tx_index)) } /// Group consecutive events that share the same `(slot, tx_index)` order key, preserving order. /// Events of a single transaction carry the same key, so this normally yields one group, but it /// stays correct if a transaction ever spans multiple keys. fn group_events_by_order_key( events: Vec, fallback_slot: u64, fallback_tx_index: u64, ) -> Vec<((u64, u64), Vec)> { let mut groups: Vec<((u64, u64), Vec)> = Vec::new(); for event in events { let key = event_order_key(&event, fallback_slot, fallback_tx_index); match groups.last_mut() { Some((last_key, group)) if *last_key == key => group.push(event), _ => groups.push((key, vec![event])), } } groups } #[inline] fn deliver_events(callback: Arc, events: Vec) { for event in events { callback(event); } } #[inline] fn has_buffered_events( order_mode: OrderMode, slot_buffer: &SlotBuffer, micro_batch: &MicroBatchBuffer, ) -> bool { match order_mode { OrderMode::Ordered | OrderMode::StreamingOrdered => !slot_buffer.is_empty(), OrderMode::MicroBatch => !micro_batch.is_empty(), OrderMode::Unordered => false, } } struct SubscriptionStartGuard { active_subscription: Arc, metrics_handle: Option>, disarmed: bool, } impl SubscriptionStartGuard { fn new(active_subscription: Arc) -> Self { Self { active_subscription, metrics_handle: None, disarmed: false } } fn set_metrics_handle(&mut self, metrics_handle: Option>) { self.metrics_handle = metrics_handle; } fn disarm(&mut self) -> Option> { self.disarmed = true; self.metrics_handle.take() } } impl Drop for SubscriptionStartGuard { fn drop(&mut self) { if !self.disarmed { self.active_subscription.store(false, Ordering::Release); if let Some(handle) = self.metrics_handle.take() { handle.abort(); } } } } struct SubscriptionCleanupGuard { active_subscription: Arc, control_tx: Arc>>>, current_request: Arc>>, metrics_handle: Option>, finished: bool, } impl SubscriptionCleanupGuard { fn new( active_subscription: Arc, control_tx: Arc>>>, current_request: Arc>>, metrics_handle: Option>, ) -> Self { Self { active_subscription, control_tx, current_request, metrics_handle, finished: false } } async fn finish(mut self) { *self.control_tx.lock().await = None; *self.current_request.write().await = None; self.active_subscription.store(false, Ordering::Release); if let Some(handle) = self.metrics_handle.take() { handle.abort(); } self.finished = true; } } impl Drop for SubscriptionCleanupGuard { fn drop(&mut self) { if !self.finished { self.active_subscription.store(false, Ordering::Release); if let Some(handle) = self.metrics_handle.take() { handle.abort(); } } } } // 实现 Clone trait 以支持模块间共享 impl Clone for YellowstoneGrpc { fn clone(&self) -> Self { Self { endpoint: self.endpoint.clone(), x_token: self.x_token.clone(), config: self.config.clone(), subscription_manager: self.subscription_manager.clone(), subscription_handle: self.subscription_handle.clone(), // 共享同一个 Arc> active_subscription: self.active_subscription.clone(), control_tx: self.control_tx.clone(), event_type_filter: self.event_type_filter.clone(), current_request: self.current_request.clone(), } } }