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

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(
-            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();
+        }
+    }
 }

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 {
-            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<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 {
-            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)
     }
 }

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 {