From 3c7bf52e012be82a888452bd301eac48d24bd102 Mon Sep 17 00:00:00 2001 From: mooncitydev Date: Tue, 19 May 2026 04:26:51 -0700 Subject: [PATCH] fix grpc reconnect after stream death and skip bad protobuf updates clear active_subscription when the yellowstone stream task exits or subscribe fails, so callers are not stuck on "already subscribed" after a disconnect. replace unwrap/expect in grpc pool parsing with validation and skip malformed account or transaction updates instead of panicking the whole stream. --- src/streaming/grpc/pool.rs | 103 +++++++++++++++++++----- src/streaming/yellowstone_grpc.rs | 24 +++++- src/streaming/yellowstone_sub_system.rs | 7 +- 3 files changed, 112 insertions(+), 22 deletions(-) diff --git a/src/streaming/grpc/pool.rs b/src/streaming/grpc/pool.rs index d55f5a7..b0e16e8 100644 --- a/src/streaming/grpc/pool.rs +++ b/src/streaming/grpc/pool.rs @@ -1,5 +1,6 @@ use super::types::{AccountPretty, BlockMetaPretty, TransactionPretty}; use crate::streaming::event_parser::common::high_performance_clock::get_high_perf_clock; +use log::warn; use solana_sdk::{pubkey::Pubkey, signature::Signature}; use std::collections::VecDeque; use std::ops::DerefMut; @@ -93,20 +94,61 @@ pub struct PooledAccountPretty { impl PooledAccountPretty { /// 从 gRPC 更新重置数据 - pub fn reset_from_update(&mut self, account_update: SubscribeUpdateAccount) { - let account_info = account_update.account.unwrap(); + pub fn reset_from_update(&mut self, account_update: SubscribeUpdateAccount) -> bool { + let Some(account_info) = account_update.account else { + warn!("grpc account update missing nested account, skipping"); + return false; + }; self.account.slot = account_update.slot; self.account.signature = if let Some(txn_signature) = account_info.txn_signature { - Signature::try_from(txn_signature.as_slice()).expect("valid signature") + if txn_signature.len() != 64 { + warn!( + "grpc account update has invalid signature length {}, skipping", + txn_signature.len() + ); + return false; + } + match Signature::try_from(txn_signature.as_slice()) { + Ok(sig) => sig, + Err(_) => { + warn!("grpc account update has invalid signature bytes, skipping"); + return false; + } + } } else { Signature::default() }; - self.account.pubkey = - Pubkey::try_from(account_info.pubkey.as_slice()).expect("valid pubkey"); + if account_info.pubkey.len() != 32 { + warn!( + "grpc account update has invalid pubkey length {}, skipping", + account_info.pubkey.len() + ); + return false; + } + self.account.pubkey = match Pubkey::try_from(account_info.pubkey.as_slice()) { + Ok(pubkey) => pubkey, + Err(_) => { + warn!("grpc account update has invalid pubkey bytes, skipping"); + return false; + } + }; self.account.executable = account_info.executable; self.account.lamports = account_info.lamports; - self.account.owner = Pubkey::try_from(account_info.owner.as_slice()).expect("valid pubkey"); + if account_info.owner.len() != 32 { + warn!( + "grpc account update has invalid owner length {}, skipping", + account_info.owner.len() + ); + return false; + } + self.account.owner = match Pubkey::try_from(account_info.owner.as_slice()) { + Ok(owner) => owner, + Err(_) => { + warn!("grpc account update has invalid owner bytes, skipping"); + return false; + } + }; self.account.rent_epoch = account_info.rent_epoch; // 优化数据字段的重用 @@ -119,6 +161,7 @@ impl PooledAccountPretty { } self.account.recv_us = get_high_perf_clock(); + true } } @@ -272,18 +315,35 @@ impl PooledTransactionPretty { &mut self, tx_update: SubscribeUpdateTransaction, block_time: Option, - ) { - let tx = tx_update.transaction.expect("should be defined"); + ) -> bool { + let Some(tx) = tx_update.transaction else { + warn!("grpc transaction update missing nested transaction, skipping"); + return false; + }; + + if tx.signature.len() != 64 { + warn!( + "grpc transaction update has invalid signature length {}, skipping", + tx.signature.len() + ); + return false; + } self.transaction.slot = tx_update.slot; self.transaction.tx_index = Some(tx.index); self.transaction.block_time = block_time; self.transaction.block_hash.clear(); // 重置 block_hash - self.transaction.signature = - Signature::try_from(tx.signature.as_slice()).expect("valid signature"); + self.transaction.signature = match Signature::try_from(tx.signature.as_slice()) { + Ok(signature) => signature, + Err(_) => { + warn!("grpc transaction update has invalid signature bytes, skipping"); + return false; + } + }; self.transaction.is_vote = tx.is_vote; self.transaction.recv_us = get_high_perf_clock(); self.transaction.grpc_tx = tx; + true } } @@ -370,12 +430,17 @@ impl Default for PoolManager { /// 工厂函数用于创建优化的 EventPretty impl EventPrettyPool { /// 创建账户事件 - 使用对象池优化 - pub fn create_account_event_optimized(&self, update: SubscribeUpdateAccount) -> AccountPretty { + pub fn create_account_event_optimized( + &self, + update: SubscribeUpdateAccount, + ) -> Option { let mut pooled_account = self.acquire_account(); - pooled_account.reset_from_update(update); + if !pooled_account.reset_from_update(update) { + return None; + } // 移动数据而不是克隆，避免多余的内存分配 let result = std::mem::replace(pooled_account.deref_mut(), AccountPretty::default()); - result + Some(result) } /// 创建区块事件 - 使用对象池优化 @@ -396,12 +461,14 @@ impl EventPrettyPool { &self, update: SubscribeUpdateTransaction, block_time: Option, - ) -> TransactionPretty { + ) -> Option { let mut pooled_tx = self.acquire_transaction(); - pooled_tx.reset_from_update(update, block_time); + if !pooled_tx.reset_from_update(update, block_time) { + return None; + } // 移动数据而不是克隆 let result = std::mem::replace(pooled_tx.deref_mut(), TransactionPretty::default()); - result + Some(result) } } @@ -414,7 +481,7 @@ pub mod factory { use super::*; /// 使用对象池创建账户事件（推荐用于高性能场景） - pub fn create_account_pretty_pooled(update: SubscribeUpdateAccount) -> AccountPretty { + pub fn create_account_pretty_pooled(update: SubscribeUpdateAccount) -> Option { GLOBAL_POOL_MANAGER.get_event_pool().create_account_event_optimized(update) } @@ -430,7 +497,7 @@ pub mod factory { pub fn create_transaction_pretty_pooled( update: SubscribeUpdateTransaction, block_time: Option, - ) -> TransactionPretty { + ) -> Option { GLOBAL_POOL_MANAGER.get_event_pool().create_transaction_event_optimized(update, block_time) } } diff --git a/src/streaming/yellowstone_grpc.rs b/src/streaming/yellowstone_grpc.rs index d00c7b7..4eb4268 100644 --- a/src/streaming/yellowstone_grpc.rs +++ b/src/streaming/yellowstone_grpc.rs @@ -155,6 +155,7 @@ impl YellowstoneGrpc { return Err(anyhow!("Already subscribed. Use update_subscription() to modify filters")); } + let clear_active_on_drop = ActiveSubscriptionGuard(self.active_subscription.clone()); let mut metrics_handle = None; // 启动自动性能监控（如果启用） if self.config.enable_metrics { @@ -187,7 +188,9 @@ impl YellowstoneGrpc { 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 stream_handle = tokio::spawn(async move { + let _clear_active = ActiveSubscriptionGuard(active_subscription); let mut slot_buffer = SlotBuffer::new(); let mut micro_batch = MicroBatchBuffer::new(); let mut last_slot = 0u64; @@ -218,7 +221,11 @@ impl YellowstoneGrpc { let created_at = msg.created_at; match msg.update_oneof { Some(UpdateOneof::Account(account)) => { - let account_pretty = factory::create_account_pretty_pooled(account); + let Some(account_pretty) = + factory::create_account_pretty_pooled(account) + else { + continue; + }; log::debug!("Received account: {:?}", account_pretty); if let Err(e) = process_grpc_transaction( EventPretty::Account(account_pretty), @@ -248,7 +255,11 @@ impl YellowstoneGrpc { } } Some(UpdateOneof::Transaction(sut)) => { - let transaction_pretty = factory::create_transaction_pretty_pooled(sut, created_at); + let Some(transaction_pretty) = + factory::create_transaction_pretty_pooled(sut, created_at) + else { + continue; + }; log::debug!( "Received transaction: {} at slot {}", transaction_pretty.signature, @@ -343,6 +354,7 @@ impl YellowstoneGrpc { let mut handle_guard = self.subscription_handle.lock().await; *handle_guard = Some(subscription_handle); + std::mem::forget(clear_active_on_drop); Ok(()) } @@ -540,6 +552,14 @@ fn has_buffered_events( } } +struct ActiveSubscriptionGuard(Arc); + +impl Drop for ActiveSubscriptionGuard { + fn drop(&mut self) { + self.0.store(false, Ordering::Release); + } +} + // 实现 Clone trait 以支持模块间共享 impl Clone for YellowstoneGrpc { fn clone(&self) -> Self { diff --git a/src/streaming/yellowstone_sub_system.rs b/src/streaming/yellowstone_sub_system.rs index 4ca9877..52b5ab5 100755 --- a/src/streaming/yellowstone_sub_system.rs +++ b/src/streaming/yellowstone_sub_system.rs @@ -59,8 +59,11 @@ impl YellowstoneGrpc { let created_at = msg.created_at; match msg.update_oneof { Some(UpdateOneof::Transaction(sut)) => { - let transaction_pretty = - factory::create_transaction_pretty_pooled(sut, created_at); + let Some(transaction_pretty) = + factory::create_transaction_pretty_pooled(sut, created_at) + else { + continue; + }; let event_pretty = EventPretty::Transaction(transaction_pretty); if let Err(e) = Self::process_system_transaction(event_pretty, &*callback).await