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

@@ -1,6 +1,6 @@
 [package]
 name = "solana-streamer-sdk"
-version = "0.3.5"
+version = "0.3.6"
 edition = "2021"
 authors = ["William <byteblock6@gmail.com>", "sgxiang <sgxiang@gmail.com>", "wei <1415121722@qq.com>"]
 repository = "https://github.com/0xfnzero/solana-streamer"

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<dyn std::error::Error>> {
     )
     .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<dyn std::error::Error>> {
     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
 

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<dyn std::error::Error>> {
     )
     .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<dyn std::error::Error>> {
     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() 方法进行干净关闭，并实现信号处理器以正确清理资源
 
 ## 许可证
 

src/main.rs

@@ -130,6 +130,16 @@ async fn test_grpc() -> Result<(), Box<dyn std::error::Error>> {
     )
     .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<dyn std::error::Error>> {
     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(())
 }
 

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<tokio::task::JoinHandle<()>> {
         // 检查是否启用性能监控
         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)
     }
 
     /// 更新处理次数

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::*;

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<JoinHandle<()>>,
+}
+
+impl SubscriptionHandle {
+    /// Create a new subscription handle
+    pub fn new(
+        stream_handle: JoinHandle<()>,
+        event_handle: JoinHandle<()>,
+        metrics_handle: Option<JoinHandle<()>>,
+    ) -> 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(())
+    }
+}

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<Mutex<PerformanceMetrics>>,
     pub metrics_manager: MetricsManager,
+    pub subscription_handle: Arc<Mutex<Option<SubscriptionHandle>>>,
 }
 
 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(
+        let metrics_manager =
-            metrics.clone(), 
+            MetricsManager::new(metrics.clone(), config_arc, "ShredStream".to_string());
-            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();
+        }
+    }
 }

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<dyn UnifiedEvent>) + 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<Pubkey>,
         event_type_filter: Option<EventTypeFilter>,
         callback: F,
-    ) -> AnyResult<()>
+    ) -> AnyResult<tokio::task::JoinHandle<()>>
     where
         F: Fn(Box<dyn UnifiedEvent>) + 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 {
+        let event_handle = tokio::spawn(async move {
-            if let Err(e) = event_processor
+            while let Some(transaction_with_slot) = rx.next().await {
-                .process_transaction_with_batch(
+                if let Err(e) = event_processor
-                    transaction_with_slot,
+                    .process_transaction_with_batch(
-                    protocols.clone(),
+                        transaction_with_slot,
-                    bot_wallet,
+                        protocols.clone(),
-                    &mut batch_processor,
+                        bot_wallet,
-                    event_type_filter.clone(),
+                        &mut batch_processor,
-                )
+                        event_type_filter.clone(),
-                .await
+                    )
-            {
+                    .await
-                log::error!("Error processing transaction: {e:?}");
+                {
+                    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<Pubkey>,
         event_type_filter: Option<EventTypeFilter>,
         callback: F,
-    ) -> AnyResult<()>
+    ) -> AnyResult<tokio::task::JoinHandle<()>>
     where
         F: Fn(Box<dyn UnifiedEvent>) + 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 {
+        let event_handle = tokio::spawn(async move {
-            if let Err(e) = event_processor
+            while let Some(transaction_with_slot) = rx.next().await {
-                .process_transaction_immediate(
+                if let Err(e) = event_processor
-                    transaction_with_slot,
+                    .process_transaction_immediate(
-                    protocols.clone(),
+                        transaction_with_slot,
-                    bot_wallet,
+                        protocols.clone(),
-                    event_type_filter.clone(),
+                        bot_wallet,
-                    &callback,
+                        event_type_filter.clone(),
-                )
+                        &callback,
-                .await
+                    )
-            {
+                    .await
-                log::error!("Error processing transaction: {e:?}");
+                {
+                    log::error!("Error processing transaction: {e:?}");
+                }
             }
-        }
+        });
 
-        Ok(())
+        Ok(event_handle)
     }
 }

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<String>,
 }
 
-#[derive(Clone)]
 pub struct YellowstoneGrpc {
     pub endpoint: String,
     pub x_token: Option<String>,
@@ -35,6 +34,7 @@ pub struct YellowstoneGrpc {
     pub subscription_manager: SubscriptionManager,
     pub metrics_manager: MetricsManager,
     pub event_processor: EventProcessor,
+    pub subscription_handle: Arc<Mutex<Option<SubscriptionHandle>>>,
 }
 
 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<dyn UnifiedEvent>) + 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<dyn UnifiedEvent>) + 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<Mutex<>>
+        }
+    }
+}
+
 // 实现 Clone trait 以支持模块间共享
 impl Clone for EventProcessor {
     fn clone(&self) -> Self {