From 1e6841149d844a08b7f297356d651812d5d354c9 Mon Sep 17 00:00:00 2001 From: ioxde <228087182+ioxde@users.noreply.github.com> Date: Mon, 25 Aug 2025 17:09:14 -0700 Subject: [PATCH] feat: add dynamic subscription management with runtime filter updates - Add update_subscription() for runtime filter updates without reconnection - BREAKING: Return subscription request from subscribe_with_request() - Implement Default trait for EventTypeFilter - Add comprehensive dynamic_subscription example --- README.md | 41 ++ examples/dynamic_subscription.rs | 468 ++++++++++++++++++++ src/streaming/event_parser/common/filter.rs | 2 +- src/streaming/grpc/subscription.rs | 5 +- src/streaming/yellowstone_grpc.rs | 204 +++++++-- src/streaming/yellowstone_sub_system.rs | 2 +- 6 files changed, 677 insertions(+), 45 deletions(-) create mode 100644 examples/dynamic_subscription.rs diff --git a/README.md b/README.md index cc73cf0..db44d9c 100755 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ A lightweight Rust library for real-time event streaming from Solana DEX trading 16. **Runtime Configuration Updates**: Supports dynamic configuration parameter updates at runtime 17. **Full Function Performance Monitoring**: All subscribe_events functions support performance monitoring, automatically collecting and reporting performance metrics 18. **Graceful Shutdown**: Support for programmatic stop() method for clean shutdown +19. **Dynamic Subscription Management**: Runtime filter updates without reconnection, enabling adaptive monitoring strategies ## Installation @@ -72,6 +73,20 @@ This example demonstrates: The example uses a predefined transaction signature and shows how to extract protocol-specific events from the transaction data. +### Dynamic Subscription Management Example + +Test runtime filter updates without reconnection: + +```bash +cargo run --example dynamic_subscription +``` + +This example demonstrates: +- Creating initial subscriptions with specific protocol filters +- Updating subscription filters at runtime without reconnection +- Single subscription enforcement and proper error handling +- Clean shutdown and resource management + ### Advanced Usage - Complete Example ```rust @@ -470,6 +485,32 @@ let event_type_filter = Some(EventTypeFilter { }); ``` +## Dynamic Subscription Management + +Update subscription filters at runtime without reconnecting to the stream. + +```rust +// Update filters on existing subscription +grpc.update_subscription( + TransactionFilter { + account_include: vec!["new_program_id".to_string()], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, +).await?; +``` + +- **No Reconnection**: Filter changes apply immediately without closing the stream +- **Atomic Updates**: Both transaction and account filters updated together +- **Single Subscription**: One active subscription per client instance +- **Compatible**: Works with both immediate and advanced subscription methods + +Note: Multiple subscription attempts on the same client return an error. + ## Supported Protocols - **PumpFun**: Primary meme coin trading platform diff --git a/examples/dynamic_subscription.rs b/examples/dynamic_subscription.rs new file mode 100644 index 0000000..fcb7365 --- /dev/null +++ b/examples/dynamic_subscription.rs @@ -0,0 +1,468 @@ +use anyhow::Result; +use solana_streamer_sdk::streaming::yellowstone_grpc::{AccountFilter, TransactionFilter, YellowstoneGrpc}; +use solana_streamer_sdk::streaming::event_parser::Protocol; +use solana_streamer_sdk::streaming::event_parser::common::filter::EventTypeFilter; +use solana_streamer_sdk::streaming::event_parser::common::types::EventType; +use solana_sdk::signature::{Keypair, Signer}; +use std::sync::{Arc, Mutex}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; +use tokio::time::sleep; + +const PUMPFUN_PROGRAM_ID: &str = "6EF8rrecthR5Dkzon8Nwu78hRvfCKubJ14M5uBEwF6P"; +const RAYDIUM_CPMM_PROGRAM_ID: &str = "CPMMoo8L3F4NbTegBCKVNunggL7H1ZpdTHKxQB5qKP1C"; + +const GRPC_ENDPOINT: &str = "https://solana-yellowstone-grpc.publicnode.com:443"; +const API_KEY: Option<&str> = None; +const MONITORING_DURATION_SECS: u64 = 2; + +/// Demonstrates dynamic subscription updates and filter changes in real-time +#[tokio::main] +async fn main() -> Result<()> { + env_logger::init(); + + println!("Connecting to Yellowstone gRPC at {}", GRPC_ENDPOINT); + let client = Arc::new(YellowstoneGrpc::new( + GRPC_ENDPOINT.to_string(), + API_KEY.map(|s| s.to_string()) + )?); + + let event_counter = Arc::new(AtomicU64::new(0)); + let counter = event_counter.clone(); + + let callback = move |event: Box| { + let count = counter.fetch_add(1, Ordering::Relaxed); + + let protocol = match event.event_type() { + EventType::PumpFunBuy | EventType::PumpFunSell => "PumpFun", + EventType::RaydiumCpmmSwapBaseInput | EventType::RaydiumCpmmSwapBaseOutput => "RaydiumCpmm", + _ => "Unknown" + }; + + println!("Event #{}: {:11} - {:.8}...", count + 1, protocol, event.signature()); + }; + + println!("\n=== Phase 1: PumpFun only ==="); + let pumpfun_filter = TransactionFilter { + account_include: vec![PUMPFUN_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }; + + let account_filter = AccountFilter { + account: vec![], + owner: vec![], + }; + let trade_event_filter = EventTypeFilter { + include: vec![ + EventType::PumpFunBuy, + EventType::PumpFunSell, + EventType::RaydiumCpmmSwapBaseInput, + EventType::RaydiumCpmmSwapBaseOutput, + ], + }; + + if let Err(e) = client.subscribe_events_immediate( + vec![Protocol::PumpFun, Protocol::RaydiumCpmm], + None, + pumpfun_filter, + account_filter, + Some(trade_event_filter), + None, + callback, + ).await { + println!("Failed to create subscription: {}", e); + return Ok(()); + } + + println!("Subscribed to PumpFun transactions with trade event filters, monitoring for {}s...", MONITORING_DURATION_SECS); + sleep(Duration::from_secs(MONITORING_DURATION_SECS)).await; + let phase1_count = event_counter.load(Ordering::Relaxed); + println!("Phase 1: {} events", phase1_count); + + println!("\n=== Phase 2: PumpFun + RaydiumCpmm ==="); + let multi_protocol_filter = TransactionFilter { + account_include: vec![ + PUMPFUN_PROGRAM_ID.to_string(), + RAYDIUM_CPMM_PROGRAM_ID.to_string(), + ], + account_exclude: vec![], + account_required: vec![], + }; + + if let Err(e) = client.update_subscription( + multi_protocol_filter, + AccountFilter { + account: vec![], + owner: vec![], + }, + ).await { + println!("Failed to update subscription: {}", e); + return Ok(()); + } + + println!("Updated to PumpFun + RaydiumCpmm transactions, monitoring for {}s...", MONITORING_DURATION_SECS); + sleep(Duration::from_secs(MONITORING_DURATION_SECS)).await; + let phase2_count = event_counter.load(Ordering::Relaxed); + println!("Phase 2: {} events", phase2_count - phase1_count); + + println!("\n=== Phase 3: RaydiumCpmm only ==="); + let raydium_cpmm_filter = TransactionFilter { + account_include: vec![RAYDIUM_CPMM_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }; + + if let Err(e) = client.update_subscription( + raydium_cpmm_filter, + AccountFilter { + account: vec![], + owner: vec![], + }, + ).await { + println!("Failed to update subscription: {}", e); + return Ok(()); + } + + println!("Updated to RaydiumCpmm transactions only, monitoring for {}s...", MONITORING_DURATION_SECS); + sleep(Duration::from_secs(MONITORING_DURATION_SECS)).await; + let phase3_count = event_counter.load(Ordering::Relaxed); + println!("Phase 3: {} events", phase3_count - phase2_count); + + println!("\n=== Phase 4: Back to PumpFun only ==="); + let pumpfun_only_filter = TransactionFilter { + account_include: vec![PUMPFUN_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }; + + if let Err(e) = client.update_subscription( + pumpfun_only_filter, + AccountFilter { + account: vec![], + owner: vec![], + }, + ).await { + println!("Failed to update subscription: {}", e); + return Ok(()); + } + + println!("Updated to PumpFun transactions only, monitoring for {}s...", MONITORING_DURATION_SECS); + sleep(Duration::from_secs(MONITORING_DURATION_SECS)).await; + let phase4_count = event_counter.load(Ordering::Relaxed); + println!("Phase 4: {} events", phase4_count - phase3_count); + + println!("\n=== Phase 5: All events ==="); + let empty_filter = TransactionFilter { + account_include: vec![], + account_exclude: vec![], + account_required: vec![], + }; + + if let Err(e) = client.update_subscription( + empty_filter, + AccountFilter { + account: vec![], + owner: vec![], + }, + ).await { + println!("Failed to update subscription: {}", e); + return Ok(()); + } + + println!("Updated to all transactions (no filters), monitoring for {}s...", MONITORING_DURATION_SECS); + sleep(Duration::from_secs(MONITORING_DURATION_SECS)).await; + let phase5_count = event_counter.load(Ordering::Relaxed); + println!("Phase 5: {} events", phase5_count - phase4_count); + + println!("\n=== Phase 6: Silence ==="); + + let random_keypair_1 = Keypair::new(); + let random_keypair_2 = Keypair::new(); + let random_pubkey_1 = random_keypair_1.pubkey(); + let random_pubkey_2 = random_keypair_2.pubkey(); + + let silence_filter = TransactionFilter { + account_include: vec![], + account_exclude: vec![], + account_required: vec![ + random_pubkey_1.to_string(), + random_pubkey_2.to_string(), + ], + }; + + if let Err(e) = client.update_subscription( + silence_filter, + AccountFilter { + account: vec![], + owner: vec![], + }, + ).await { + println!("Failed to update subscription: {}", e); + return Ok(()); + } + + println!("Updated to random addresses (expecting silence), monitoring for 3s..."); + let before_silence = event_counter.load(Ordering::Relaxed); + let start_time = Instant::now(); + let last_event_time = Arc::new(Mutex::new(start_time)); + let last_event_time_clone = last_event_time.clone(); + + let mut last_count = before_silence; + for _ in 0..6 { + sleep(Duration::from_millis(500)).await; + let current_count = event_counter.load(Ordering::Relaxed); + if current_count > last_count { + if let Ok(mut time) = last_event_time_clone.lock() { + *time = Instant::now(); + } + last_count = current_count; + } + } + + let final_count = event_counter.load(Ordering::Relaxed); + let events_during_silence = final_count - before_silence; + + if events_during_silence == 0 { + println!("Phase 6: 0 events (immediate filter application)"); + } else if let Ok(last_time) = last_event_time.lock() { + let propagation_time = last_time.duration_since(start_time); + println!("Phase 6: {} events during propagation, filter took {}ms", + events_during_silence, propagation_time.as_millis()); + } + + println!("\n=== Phase 7: Shutdown ==="); + + let shutdown_client = Arc::new(YellowstoneGrpc::new( + GRPC_ENDPOINT.to_string(), + API_KEY.map(|s| s.to_string()) + )?); + + let shutdown_event_counter = Arc::new(AtomicU64::new(0)); + let shutdown_counter = shutdown_event_counter.clone(); + let shutdown_callback = move |_event: Box| { + shutdown_counter.fetch_add(1, Ordering::Relaxed); + }; + + if let Err(e) = shutdown_client.subscribe_events_immediate( + vec![Protocol::PumpFun, Protocol::RaydiumCpmm], + None, + TransactionFilter { + account_include: vec![], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, + None, + None, + shutdown_callback, + ).await { + println!("Failed to subscribe shutdown client: {}", e); + return Ok(()); + } + + sleep(Duration::from_millis(1000)).await; + let pre_stop_count = shutdown_event_counter.load(Ordering::Relaxed); + println!("Received {} events before stop", pre_stop_count); + + let stop_time = Instant::now(); + shutdown_client.stop().await; + let shutdown_duration = stop_time.elapsed(); + println!("stop() completed in {:.1}ms", shutdown_duration.as_millis()); + + let post_stop_count = shutdown_event_counter.load(Ordering::Relaxed); + let during_stop = post_stop_count - pre_stop_count; + if during_stop > 0 { + println!(" {} events received during stop()", during_stop); + } + + let last_event_time = Arc::new(Mutex::new(stop_time)); + let last_event_time_clone = last_event_time.clone(); + let mut last_count = post_stop_count; + + for _ in 0..20 { + sleep(Duration::from_millis(100)).await; + let current_count = shutdown_event_counter.load(Ordering::Relaxed); + if current_count > last_count { + if let Ok(mut time) = last_event_time_clone.lock() { + *time = Instant::now(); + } + last_count = current_count; + } + } + + let final_count = shutdown_event_counter.load(Ordering::Relaxed); + let after_stop = final_count - post_stop_count; + + if after_stop == 0 { + println!("Phase 7: Clean shutdown - no events after stop()"); + } else if let Ok(last_time) = last_event_time.lock() { + let post_stop_duration = last_time.duration_since(stop_time); + let silence_duration = Instant::now().duration_since(*last_time); + println!("Phase 7: {} events arrived up to {}ms after stop(), then silent for {}ms", + after_stop, post_stop_duration.as_millis(), silence_duration.as_millis()); + } + + println!("\n=== Subscription enforcement ==="); + + let test_callback = |_event: Box| {}; + + match client.subscribe_events_immediate( + vec![Protocol::RaydiumCpmm], + None, + TransactionFilter { + account_include: vec![RAYDIUM_CPMM_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, + None, + None, + test_callback, + ).await { + Ok(_) => println!("ERROR: Same client created second subscription"), + Err(e) if e.to_string().contains("Already subscribed") => { + println!("✓ Single subscription enforcement working"); + }, + Err(e) => println!("Unexpected error: {}", e), + } + + let client2 = Arc::new(YellowstoneGrpc::new( + GRPC_ENDPOINT.to_string(), + API_KEY.map(|s| s.to_string()) + )?); + + let client2_counter = Arc::new(AtomicU64::new(0)); + let counter2 = client2_counter.clone(); + let client2_callback = move |_event: Box| { + counter2.fetch_add(1, Ordering::Relaxed); + }; + + match client2.subscribe_events_immediate( + vec![Protocol::RaydiumCpmm], + None, + TransactionFilter { + account_include: vec![RAYDIUM_CPMM_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, + None, + None, + client2_callback, + ).await { + Ok(_) => { + sleep(Duration::from_millis(500)).await; + let count = client2_counter.load(Ordering::Relaxed); + println!("✓ Second client: {} events", count); + client2.stop().await; + }, + Err(e) => println!("ERROR: Second client failed: {}", e), + } + + println!("\n=== Advanced subscription enforcement ==="); + + let test_callback_advanced = |_event: Box| {}; + + let client3 = Arc::new(YellowstoneGrpc::new( + GRPC_ENDPOINT.to_string(), + API_KEY.map(|s| s.to_string()) + )?); + + // First subscription should succeed + match client3.subscribe_events_advanced( + vec![Protocol::RaydiumCpmm], + None, + TransactionFilter { + account_include: vec![RAYDIUM_CPMM_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, + None, + None, + test_callback_advanced, + ).await { + Ok(_) => { + // Second subscription attempt on same client should fail + match client3.subscribe_events_advanced( + vec![Protocol::RaydiumCpmm], + None, + TransactionFilter { + account_include: vec![RAYDIUM_CPMM_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, + None, + None, + |_| {}, + ).await { + Ok(_) => println!("ERROR: Same client created second advanced subscription"), + Err(e) if e.to_string().contains("Already subscribed") => { + println!("✓ Advanced single subscription enforcement working"); + }, + Err(e) => println!("Unexpected error: {}", e), + } + }, + Err(e) => println!("ERROR: First advanced subscription failed: {}", e), + } + + // Test that a second client can subscribe using advanced method + let client4 = Arc::new(YellowstoneGrpc::new( + GRPC_ENDPOINT.to_string(), + API_KEY.map(|s| s.to_string()) + )?); + + let client4_counter = Arc::new(AtomicU64::new(0)); + let counter4 = client4_counter.clone(); + let client4_callback = move |_event: Box| { + counter4.fetch_add(1, Ordering::Relaxed); + }; + + match client4.subscribe_events_advanced( + vec![Protocol::RaydiumCpmm], + None, + TransactionFilter { + account_include: vec![RAYDIUM_CPMM_PROGRAM_ID.to_string()], + account_exclude: vec![], + account_required: vec![], + }, + AccountFilter { + account: vec![], + owner: vec![], + }, + None, + None, + client4_callback, + ).await { + Ok(_) => { + sleep(Duration::from_millis(500)).await; + let count = client4_counter.load(Ordering::Relaxed); + println!("✓ Second client (advanced): {} events", count); + client4.stop().await; + }, + Err(e) => println!("ERROR: Second client (advanced) failed: {}", e), + } + + client3.stop().await; + + client.stop().await; + + Ok(()) +} diff --git a/src/streaming/event_parser/common/filter.rs b/src/streaming/event_parser/common/filter.rs index 8fa3296..2ac0743 100644 --- a/src/streaming/event_parser/common/filter.rs +++ b/src/streaming/event_parser/common/filter.rs @@ -2,7 +2,7 @@ use crate::streaming::event_parser::common::{ types::EventType, ACCOUNT_EVENT_TYPES, BLOCK_EVENT_TYPES, }; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct EventTypeFilter { pub include: Vec, } diff --git a/src/streaming/grpc/subscription.rs b/src/streaming/grpc/subscription.rs index 2efdddd..15579fa 100644 --- a/src/streaming/grpc/subscription.rs +++ b/src/streaming/grpc/subscription.rs @@ -49,6 +49,7 @@ impl SubscriptionManager { ) -> AnyResult<( impl Sink, impl Stream>, + SubscribeRequest, )> { let blocks_meta = if event_type_filter.is_some() && event_type_filter.as_ref().unwrap().include_block_event() @@ -71,8 +72,8 @@ impl SubscriptionManager { ..Default::default() }; let mut client = self.connect().await?; - let (sink, stream) = client.subscribe_with_request(Some(subscribe_request)).await?; - Ok((sink, stream)) + let (sink, stream) = client.subscribe_with_request(Some(subscribe_request.clone())).await?; + Ok((sink, stream, subscribe_request)) } /// 创建账户订阅请求并返回流 diff --git a/src/streaming/yellowstone_grpc.rs b/src/streaming/yellowstone_grpc.rs index f88e88b..d5ffdd3 100644 --- a/src/streaming/yellowstone_grpc.rs +++ b/src/streaming/yellowstone_grpc.rs @@ -1,10 +1,12 @@ -use futures::{channel::mpsc, StreamExt}; +use futures::{channel::mpsc, SinkExt, StreamExt}; use log::error; use solana_sdk::pubkey::Pubkey; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tokio::sync::Mutex; -use yellowstone_grpc_proto::geyser::CommitmentLevel; +use tokio::sync::{Mutex, RwLock}; +use yellowstone_grpc_proto::geyser::{CommitmentLevel, SubscribeRequest}; +use anyhow::anyhow; use crate::common::AnyResult; use crate::streaming::common::{ EventBatchProcessor, MetricsManager, PerformanceMetrics, StreamClientConfig, SubscriptionHandle, @@ -14,6 +16,7 @@ use crate::streaming::event_parser::{Protocol, UnifiedEvent}; use crate::streaming::grpc::{EventPretty, EventProcessor, StreamHandler, SubscriptionManager}; /// 交易过滤器 +#[derive(Debug, Clone)] pub struct TransactionFilter { pub account_include: Vec, pub account_exclude: Vec, @@ -21,6 +24,7 @@ pub struct TransactionFilter { } /// 账户过滤器 +#[derive(Debug, Clone)] pub struct AccountFilter { pub account: Vec, pub owner: Vec, @@ -35,6 +39,10 @@ pub struct YellowstoneGrpc { pub metrics_manager: MetricsManager, pub event_processor: EventProcessor, pub subscription_handle: Arc>>, + // Dynamic subscription management fields + pub active_subscription: Arc, + pub control_tx: Arc>>>, + pub current_request: Arc>>, } impl YellowstoneGrpc { @@ -68,6 +76,9 @@ impl YellowstoneGrpc { metrics_manager, event_processor, subscription_handle: Arc::new(Mutex::new(None)), + active_subscription: Arc::new(AtomicBool::new(false)), + control_tx: Arc::new(Mutex::new(None)), + current_request: Arc::new(RwLock::new(None)), }) } @@ -119,6 +130,10 @@ impl YellowstoneGrpc { 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) @@ -147,8 +162,14 @@ impl YellowstoneGrpc { where F: Fn(Box) + Send + Sync + 'static, { - // 如果已有活跃订阅，先停止它 - self.stop().await; + 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 metrics_handle = None; // 启动自动性能监控（如果启用） @@ -169,35 +190,53 @@ impl YellowstoneGrpc { ); // 订阅事件 - let (mut subscribe_tx, mut stream) = self + let (mut subscribe_tx, mut stream, subscribe_request) = self .subscription_manager .subscribe_with_request(transactions, accounts, commitment, event_type_filter.clone()) .await?; + *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); + // 创建通道，使用配置中的通道大小 let (mut tx, mut rx) = mpsc::channel::(self.config.backpressure.channel_size); // 启动流处理任务 let backpressure_strategy = self.config.backpressure.strategy; let stream_handle = 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; + loop { + tokio::select! { + message = stream.next() => { + match message { + Some(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; + } + } + Some(Err(error)) => { + error!("Stream error: {error:?}"); + break; + } + None => break, } } - Err(error) => { - error!("Stream error: {error:?}"); - break; + + Some(update) = control_rx.next() => { + if let Err(e) = subscribe_tx.send(update).await { + error!("Failed to send subscription update: {}", e); + break; + } } } } @@ -263,8 +302,15 @@ impl YellowstoneGrpc { where F: Fn(Box) + Send + Sync + 'static, { - // 如果已有活跃订阅，先停止它 - self.stop().await; + // Single subscription enforcement + 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 metrics_handle = None; // 启动自动性能监控（如果启用） @@ -285,11 +331,17 @@ impl YellowstoneGrpc { ); // Subscribe to events - let (mut subscribe_tx, mut stream) = self + let (mut subscribe_tx, mut stream, subscribe_request) = self .subscription_manager .subscribe_with_request(transactions, accounts, commitment, event_type_filter.clone()) .await?; + *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); + // Create channel let (mut tx, mut rx) = mpsc::channel::(self.config.backpressure.channel_size); @@ -309,24 +361,36 @@ impl YellowstoneGrpc { // Start task to process the stream let backpressure_strategy = self.config.backpressure.strategy; let stream_handle = 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; + loop { + tokio::select! { + message = stream.next() => { + match message { + Some(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; + } + } + Some(Err(error)) => { + error!("Stream error: {error:?}"); + break; + } + None => break, } } - Err(error) => { - error!("Stream error: {error:?}"); - break; + + Some(update) = control_rx.next() => { + if let Err(e) = subscribe_tx.send(update).await { + error!("Failed to send subscription update: {}", e); + break; + } } } } @@ -362,6 +426,61 @@ impl YellowstoneGrpc { 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: TransactionFilter, + account_filter: AccountFilter, + ) -> 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.account_include, + transaction_filter.account_exclude, + transaction_filter.account_required, + None, + ) + .unwrap_or_default(); + + request.accounts = self.subscription_manager + .subscribe_with_account_request( + account_filter.account, + account_filter.owner, + None, + ) + .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(()) + } } // 实现 Clone trait 以支持模块间共享 @@ -376,6 +495,9 @@ impl Clone for YellowstoneGrpc { metrics_manager: self.metrics_manager.clone(), event_processor: self.event_processor.clone(), subscription_handle: self.subscription_handle.clone(), // 共享同一个 Arc> + active_subscription: self.active_subscription.clone(), + control_tx: self.control_tx.clone(), + current_request: self.current_request.clone(), } } } diff --git a/src/streaming/yellowstone_sub_system.rs b/src/streaming/yellowstone_sub_system.rs index ebb4aa9..c47d36f 100755 --- a/src/streaming/yellowstone_sub_system.rs +++ b/src/streaming/yellowstone_sub_system.rs @@ -47,7 +47,7 @@ impl YellowstoneGrpc { addrs, None, ); - let (mut subscribe_tx, mut stream) = self + let (mut subscribe_tx, mut stream, _) = self .subscription_manager .subscribe_with_request(transactions, None, None, None) .await?;