From 9ec7194d2a5ec956cc43df49af3e76cd638361b2 Mon Sep 17 00:00:00 2001 From: ysq Date: Thu, 21 Aug 2025 15:47:31 +0800 Subject: [PATCH] feat: add subscription management and graceful shutdown support (v0.3.6) - Add SubscriptionHandle for managing stream lifecycle - Implement graceful shutdown with stop() methods - Update documentation with new features - Bump version to 0.3.6 --- Cargo.toml | 2 +- README.md | 26 +++++++- README_CN.md | 26 +++++++- src/main.rs | 20 ++++++ src/streaming/common/metrics.rs | 7 +- src/streaming/common/mod.rs | 2 + src/streaming/common/subscription.rs | 38 +++++++++++ src/streaming/shred/connection.rs | 27 +++++--- src/streaming/shred_stream.rs | 95 ++++++++++++++++------------ src/streaming/yellowstone_grpc.rs | 63 +++++++++++++++--- 10 files changed, 238 insertions(+), 68 deletions(-) create mode 100644 src/streaming/common/subscription.rs diff --git a/Cargo.toml b/Cargo.toml index 790aa5d..6dbd2e9 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "solana-streamer-sdk" -version = "0.3.5" +version = "0.3.6" edition = "2021" authors = ["William ", "sgxiang ", "wei <1415121722@qq.com>"] repository = "https://github.com/0xfnzero/solana-streamer" diff --git a/README.md b/README.md index f3db4bf..98272a8 100755 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ A lightweight Rust library for real-time event streaming from Solana DEX trading 15. **Backpressure Handling**: Supports blocking, dropping, retrying, ordered, and other backpressure strategies 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 ## Installation @@ -44,14 +45,14 @@ Add the dependency to your `Cargo.toml`: ```toml # Add to your Cargo.toml -solana-streamer-sdk = { path = "./solana-streamer", version = "0.3.5" } +solana-streamer-sdk = { path = "./solana-streamer", version = "0.3.6" } ``` ### Use crates.io ```toml # Add to your Cargo.toml -solana-streamer-sdk = "0.3.5" +solana-streamer-sdk = "0.3.6" ``` ## Usage Examples @@ -206,6 +207,16 @@ async fn test_grpc() -> Result<(), Box> { ) .await?; + // Support stop method, test code - stop after 1000 seconds asynchronously + let grpc_clone = grpc.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(1000)).await; + grpc_clone.stop().await; + }); + + println!("Waiting for Ctrl+C to stop..."); + tokio::signal::ctrl_c().await?; + Ok(()) } @@ -239,6 +250,16 @@ async fn test_shreds() -> Result<(), Box> { println!("Listening for events, press Ctrl+C to stop..."); shred_stream.shredstream_subscribe(protocols, None, event_type_filter, callback).await?; + // Support stop method, test code - stop after 1000 seconds asynchronously + let shred_clone = shred_stream.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(1000)).await; + shred_clone.stop().await; + }); + + println!("Waiting for Ctrl+C to stop..."); + tokio::signal::ctrl_c().await?; + Ok(()) } @@ -525,6 +546,7 @@ MIT License 4. **Error Handling**: Robust error handling for network issues and service interruptions 5. **Batch Processing Optimization**: Use batch processing to reduce callback overhead and improve throughput 6. **Performance Monitoring**: Enable performance monitoring to identify bottlenecks and optimization opportunities +7. **Graceful Shutdown**: Use the stop() method for clean shutdown and implement signal handlers for proper resource cleanup ## Important Notes diff --git a/README_CN.md b/README_CN.md index 21fefa4..49e76ee 100644 --- a/README_CN.md +++ b/README_CN.md @@ -28,6 +28,7 @@ 15. **背压处理**: 支持阻塞、丢弃、重试、有序等多种背压策略 16. **运行时配置更新**: 支持在运行时动态更新配置参数 17. **全函数性能监控**: 所有subscribe_events函数都支持性能监控，自动收集和报告性能指标 +18. **优雅关闭**: 支持编程式 stop() 方法进行干净的关闭 ## 安装 @@ -44,14 +45,14 @@ git clone https://github.com/0xfnzero/solana-streamer ```toml # 添加到您的 Cargo.toml -solana-streamer-sdk = { path = "./solana-streamer", version = "0.3.5" } +solana-streamer-sdk = { path = "./solana-streamer", version = "0.3.6" } ``` ### 使用 crates.io ```toml # 添加到您的 Cargo.toml -solana-streamer-sdk = "0.3.5" +solana-streamer-sdk = "0.3.6" ``` ## 使用示例 @@ -206,6 +207,16 @@ async fn test_grpc() -> Result<(), Box> { ) .await?; + // 支持 stop 方法，测试代码 - 异步1000秒之后停止 + let grpc_clone = grpc.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(1000)).await; + grpc_clone.stop().await; + }); + + println!("Waiting for Ctrl+C to stop..."); + tokio::signal::ctrl_c().await?; + Ok(()) } @@ -239,6 +250,16 @@ async fn test_shreds() -> Result<(), Box> { println!("Listening for events, press Ctrl+C to stop..."); shred_stream.shredstream_subscribe(protocols, None, event_type_filter, callback).await?; + // 支持 stop 方法，测试代码 - 异步1000秒之后停止 + let shred_clone = shred_stream.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(1000)).await; + shred_clone.stop().await; + }); + + println!("Waiting for Ctrl+C to stop..."); + tokio::signal::ctrl_c().await?; + Ok(()) } @@ -516,6 +537,7 @@ src/ 4. **错误处理**: 对网络问题和服务中断进行健壮的错误处理 5. **批处理优化**: 使用批处理减少回调开销，提高吞吐量 6. **性能监控**: 启用性能监控以识别瓶颈和优化机会 +7. **优雅关闭**: 使用 stop() 方法进行干净关闭，并实现信号处理器以正确清理资源 ## 许可证 diff --git a/src/main.rs b/src/main.rs index 5969861..dd110cd 100755 --- a/src/main.rs +++ b/src/main.rs @@ -130,6 +130,16 @@ async fn test_grpc() -> Result<(), Box> { ) .await?; + // 支持 stop 方法，测试代码 - 异步1000秒之后停止 + let grpc_clone = grpc.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(1000)).await; + grpc_clone.stop().await; + }); + + println!("Waiting for Ctrl+C to stop..."); + tokio::signal::ctrl_c().await?; + Ok(()) } @@ -163,6 +173,16 @@ async fn test_shreds() -> Result<(), Box> { println!("Listening for events, press Ctrl+C to stop..."); shred_stream.shredstream_subscribe(protocols, None, event_type_filter, callback).await?; + // 支持 stop 方法，测试代码 - 异步1000秒之后停止 + let shred_clone = shred_stream.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_secs(1000)).await; + shred_clone.stop().await; + }); + + println!("Waiting for Ctrl+C to stop..."); + tokio::signal::ctrl_c().await?; + Ok(()) } diff --git a/src/streaming/common/metrics.rs b/src/streaming/common/metrics.rs index 26126df..2ccdebe 100644 --- a/src/streaming/common/metrics.rs +++ b/src/streaming/common/metrics.rs @@ -199,14 +199,14 @@ impl MetricsManager { } /// 启动自动性能监控任务 - pub async fn start_auto_monitoring(&self) { + pub async fn start_auto_monitoring(&self) -> Option> { // 检查是否启用性能监控 if !self.config.enable_metrics { - return; // 如果未启用性能监控，不启动监控任务 + return None; // 如果未启用性能监控，不启动监控任务 } let metrics_manager = self.clone(); - tokio::spawn(async move { + let handle = tokio::spawn(async move { let mut interval = tokio::time::interval(tokio::time::Duration::from_secs( DEFAULT_METRICS_PRINT_INTERVAL_SECONDS, )); @@ -215,6 +215,7 @@ impl MetricsManager { metrics_manager.print_metrics().await; } }); + Some(handle) } /// 更新处理次数 diff --git a/src/streaming/common/mod.rs b/src/streaming/common/mod.rs index 1132b0e..bbbc5e1 100644 --- a/src/streaming/common/mod.rs +++ b/src/streaming/common/mod.rs @@ -3,9 +3,11 @@ pub mod config; pub mod metrics; pub mod batch; pub mod constants; +pub mod subscription; // 重新导出主要类型 pub use config::*; pub use metrics::*; pub use batch::*; pub use constants::*; +pub use subscription::*; diff --git a/src/streaming/common/subscription.rs b/src/streaming/common/subscription.rs new file mode 100644 index 0000000..ca29cbb --- /dev/null +++ b/src/streaming/common/subscription.rs @@ -0,0 +1,38 @@ +use tokio::task::JoinHandle; + +/// Subscription handle for managing and stopping subscriptions +pub struct SubscriptionHandle { + stream_handle: JoinHandle<()>, + event_handle: JoinHandle<()>, + metrics_handle: Option>, +} + +impl SubscriptionHandle { + /// Create a new subscription handle + pub fn new( + stream_handle: JoinHandle<()>, + event_handle: JoinHandle<()>, + metrics_handle: Option>, + ) -> Self { + Self { stream_handle, event_handle, metrics_handle } + } + + /// Stop subscription and abort all related tasks + pub fn stop(self) { + self.stream_handle.abort(); + self.event_handle.abort(); + if let Some(handle) = self.metrics_handle { + handle.abort(); + } + } + + /// Asynchronously wait for all tasks to complete + pub async fn join(self) -> Result<(), tokio::task::JoinError> { + let _ = self.stream_handle.await; + let _ = self.event_handle.await; + if let Some(handle) = self.metrics_handle { + let _ = handle.await; + } + Ok(()) + } +} diff --git a/src/streaming/shred/connection.rs b/src/streaming/shred/connection.rs index ea08b39..8a27a8a 100644 --- a/src/streaming/shred/connection.rs +++ b/src/streaming/shred/connection.rs @@ -3,10 +3,10 @@ use tokio::sync::Mutex; use tonic::transport::Channel; use crate::common::AnyResult; -use crate::streaming::common::{ - MetricsManager, PerformanceMetrics, StreamClientConfig, -}; use crate::protos::shredstream::shredstream_proxy_client::ShredstreamProxyClient; +use crate::streaming::common::{ + MetricsManager, PerformanceMetrics, StreamClientConfig, SubscriptionHandle, +}; /// ShredStream gRPC 客户端 #[derive(Clone)] @@ -15,6 +15,7 @@ pub struct ShredStreamGrpc { pub config: StreamClientConfig, pub metrics: Arc>, pub metrics_manager: MetricsManager, + pub subscription_handle: Arc>>, } impl ShredStreamGrpc { @@ -28,18 +29,16 @@ impl ShredStreamGrpc { let shredstream_client = ShredstreamProxyClient::connect(endpoint.clone()).await?; let metrics = Arc::new(Mutex::new(PerformanceMetrics::new())); let config_arc = Arc::new(config.clone()); - - let metrics_manager = MetricsManager::new( - metrics.clone(), - config_arc, - "ShredStream".to_string() - ); - + + let metrics_manager = + MetricsManager::new(metrics.clone(), config_arc, "ShredStream".to_string()); + Ok(Self { shredstream_client: Arc::new(shredstream_client), config, metrics, metrics_manager, + subscription_handle: Arc::new(Mutex::new(None)), }) } @@ -82,4 +81,12 @@ impl ShredStreamGrpc { pub async fn start_auto_metrics_monitoring(&self) { self.metrics_manager.start_auto_monitoring().await; } + + /// 停止当前订阅 + pub async fn stop(&self) { + let mut handle_guard = self.subscription_handle.lock().await; + if let Some(handle) = handle_guard.take() { + handle.stop(); + } + } } diff --git a/src/streaming/shred_stream.rs b/src/streaming/shred_stream.rs index b121293..356e774 100755 --- a/src/streaming/shred_stream.rs +++ b/src/streaming/shred_stream.rs @@ -1,7 +1,7 @@ use solana_sdk::pubkey::Pubkey; use crate::common::AnyResult; -use crate::streaming::common::EventBatchProcessor; +use crate::streaming::common::{EventBatchProcessor, SubscriptionHandle}; use crate::streaming::event_parser::common::filter::EventTypeFilter; use crate::streaming::event_parser::{Protocol, UnifiedEvent}; use crate::streaming::shred::{ShredEventProcessor, ShredStreamHandler, TransactionWithSlot}; @@ -20,27 +20,38 @@ impl ShredStreamGrpc { where F: Fn(Box) + Send + Sync + 'static, { + // 如果已有活跃订阅，先停止它 + self.stop().await; + + let mut metrics_handle = None; // 启动自动性能监控（如果启用） if self.config.enable_metrics { - self.metrics_manager.start_auto_monitoring().await; + metrics_handle = self.metrics_manager.start_auto_monitoring().await; } // 启动流处理 let client = (*self.shredstream_client).clone(); - let (_stream_task, rx) = ShredStreamHandler::start_stream_processing( + let (stream_task, rx) = ShredStreamHandler::start_stream_processing( client, self.config.backpressure.channel_size, ) .await?; - // 根据配置选择处理模式 - if self.config.batch.enabled { + // 根据配置选择处理模式并获取事件处理任务句柄 + let event_handle = if self.config.batch.enabled { // 批处理模式 - self.process_with_batch(rx, protocols, bot_wallet, event_type_filter, callback).await + self.process_with_batch(rx, protocols, bot_wallet, event_type_filter, callback).await? } else { // 即时处理模式 - self.process_immediate(rx, protocols, bot_wallet, event_type_filter, callback).await - } + self.process_immediate(rx, protocols, bot_wallet, event_type_filter, callback).await? + }; + + // 保存订阅句柄 + let subscription_handle = SubscriptionHandle::new(stream_task, event_handle, metrics_handle); + let mut handle_guard = self.subscription_handle.lock().await; + *handle_guard = Some(subscription_handle); + + Ok(()) } /// 批处理模式 @@ -51,7 +62,7 @@ impl ShredStreamGrpc { bot_wallet: Option, event_type_filter: Option, callback: F, - ) -> AnyResult<()> + ) -> AnyResult> where F: Fn(Box) + Send + Sync + 'static, { @@ -74,25 +85,27 @@ impl ShredStreamGrpc { let event_processor = ShredEventProcessor::new(self.metrics_manager.clone(), self.config.clone()); - while let Some(transaction_with_slot) = rx.next().await { - if let Err(e) = event_processor - .process_transaction_with_batch( - transaction_with_slot, - protocols.clone(), - bot_wallet, - &mut batch_processor, - event_type_filter.clone(), - ) - .await - { - log::error!("Error processing transaction: {e:?}"); + let event_handle = tokio::spawn(async move { + while let Some(transaction_with_slot) = rx.next().await { + if let Err(e) = event_processor + .process_transaction_with_batch( + transaction_with_slot, + protocols.clone(), + bot_wallet, + &mut batch_processor, + event_type_filter.clone(), + ) + .await + { + log::error!("Error processing transaction: {e:?}"); + } } - } - // 处理剩余的事件 - batch_processor.flush(); + // 处理剩余的事件 + batch_processor.flush(); + }); - Ok(()) + Ok(event_handle) } /// 即时处理模式 @@ -103,7 +116,7 @@ impl ShredStreamGrpc { bot_wallet: Option, event_type_filter: Option, callback: F, - ) -> AnyResult<()> + ) -> AnyResult> where F: Fn(Box) + Send + Sync + 'static, { @@ -113,21 +126,23 @@ impl ShredStreamGrpc { let event_processor = ShredEventProcessor::new(self.metrics_manager.clone(), self.config.clone()); - while let Some(transaction_with_slot) = rx.next().await { - if let Err(e) = event_processor - .process_transaction_immediate( - transaction_with_slot, - protocols.clone(), - bot_wallet, - event_type_filter.clone(), - &callback, - ) - .await - { - log::error!("Error processing transaction: {e:?}"); + let event_handle = tokio::spawn(async move { + while let Some(transaction_with_slot) = rx.next().await { + if let Err(e) = event_processor + .process_transaction_immediate( + transaction_with_slot, + protocols.clone(), + bot_wallet, + event_type_filter.clone(), + &callback, + ) + .await + { + log::error!("Error processing transaction: {e:?}"); + } } - } + }); - Ok(()) + Ok(event_handle) } } diff --git a/src/streaming/yellowstone_grpc.rs b/src/streaming/yellowstone_grpc.rs index 0ceff0c..f88e88b 100644 --- a/src/streaming/yellowstone_grpc.rs +++ b/src/streaming/yellowstone_grpc.rs @@ -7,7 +7,7 @@ use yellowstone_grpc_proto::geyser::CommitmentLevel; use crate::common::AnyResult; use crate::streaming::common::{ - EventBatchProcessor, MetricsManager, PerformanceMetrics, StreamClientConfig, + EventBatchProcessor, MetricsManager, PerformanceMetrics, StreamClientConfig, SubscriptionHandle, }; use crate::streaming::event_parser::common::filter::EventTypeFilter; use crate::streaming::event_parser::{Protocol, UnifiedEvent}; @@ -26,7 +26,6 @@ pub struct AccountFilter { pub owner: Vec, } -#[derive(Clone)] pub struct YellowstoneGrpc { pub endpoint: String, pub x_token: Option, @@ -35,6 +34,7 @@ pub struct YellowstoneGrpc { pub subscription_manager: SubscriptionManager, pub metrics_manager: MetricsManager, pub event_processor: EventProcessor, + pub subscription_handle: Arc>>, } impl YellowstoneGrpc { @@ -67,6 +67,7 @@ impl YellowstoneGrpc { subscription_manager, metrics_manager, event_processor, + subscription_handle: Arc::new(Mutex::new(None)), }) } @@ -112,6 +113,14 @@ impl YellowstoneGrpc { 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(); + } + } + /// Simplified immediate event subscription (recommended for simple scenarios) /// /// # Parameters @@ -138,9 +147,13 @@ impl YellowstoneGrpc { where F: Fn(Box) + Send + Sync + 'static, { + // 如果已有活跃订阅，先停止它 + self.stop().await; + + let mut metrics_handle = None; // 启动自动性能监控（如果启用） if self.config.enable_metrics { - self.metrics_manager.start_auto_monitoring().await; + metrics_handle = self.metrics_manager.start_auto_monitoring().await; } let transactions = self.subscription_manager.get_subscribe_request_filter( @@ -166,7 +179,7 @@ impl YellowstoneGrpc { // 启动流处理任务 let backpressure_strategy = self.config.backpressure.strategy; - tokio::spawn(async move { + let stream_handle = tokio::spawn(async move { while let Some(message) = stream.next().await { match message { Ok(msg) => { @@ -192,7 +205,7 @@ impl YellowstoneGrpc { // 即时处理交易，无批处理 let event_processor = self.event_processor.clone(); - tokio::spawn(async move { + let event_handle = tokio::spawn(async move { while let Some(event_pretty) = rx.next().await { if let Err(e) = event_processor .process_event_transaction_with_metrics( @@ -209,7 +222,12 @@ impl YellowstoneGrpc { } }); - tokio::signal::ctrl_c().await?; + // 保存订阅句柄 + let subscription_handle = + SubscriptionHandle::new(stream_handle, event_handle, metrics_handle); + let mut handle_guard = self.subscription_handle.lock().await; + *handle_guard = Some(subscription_handle); + Ok(()) } @@ -245,9 +263,13 @@ impl YellowstoneGrpc { where F: Fn(Box) + Send + Sync + 'static, { + // 如果已有活跃订阅，先停止它 + self.stop().await; + + let mut metrics_handle = None; // 启动自动性能监控（如果启用） if self.config.enable_metrics { - self.metrics_manager.start_auto_monitoring().await; + metrics_handle = self.metrics_manager.start_auto_monitoring().await; } let transactions = self.subscription_manager.get_subscribe_request_filter( @@ -286,7 +308,7 @@ impl YellowstoneGrpc { // Start task to process the stream let backpressure_strategy = self.config.backpressure.strategy; - tokio::spawn(async move { + let stream_handle = tokio::spawn(async move { while let Some(message) = stream.next().await { match message { Ok(msg) => { @@ -312,7 +334,7 @@ impl YellowstoneGrpc { // Process transactions with batch processing let event_processor = self.event_processor.clone(); - tokio::spawn(async move { + let event_handle = tokio::spawn(async move { while let Some(event_pretty) = rx.next().await { if let Err(e) = event_processor .process_event_transaction_with_batch( @@ -332,11 +354,32 @@ impl YellowstoneGrpc { batch_processor.flush(); }); - tokio::signal::ctrl_c().await?; + // 保存订阅句柄 + let subscription_handle = + SubscriptionHandle::new(stream_handle, event_handle, metrics_handle); + let mut handle_guard = self.subscription_handle.lock().await; + *handle_guard = Some(subscription_handle); + Ok(()) } } +// 实现 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(), + metrics: self.metrics.clone(), + subscription_manager: self.subscription_manager.clone(), + metrics_manager: self.metrics_manager.clone(), + event_processor: self.event_processor.clone(), + subscription_handle: self.subscription_handle.clone(), // 共享同一个 Arc> + } + } +} + // 实现 Clone trait 以支持模块间共享 impl Clone for EventProcessor { fn clone(&self) -> Self {