geyser: wrap message into Box in snapshot channel

CHANGELOG.md

@@ -14,6 +14,8 @@ The minor version will be incremented upon a breaking change and the patch versi
 
 ### Features
 
+- geyser: wrap message into `Box` in snapshot channel ([#418](https://github.com/rpcpool/yellowstone-grpc/pull/418))
+
 ### Breaking
 
 ## 2024-08-26

yellowstone-grpc-geyser/src/grpc.rs

@@ -728,7 +728,7 @@ pub struct GrpcService {
     config_filters: Arc<ConfigGrpcFilters>,
     blocks_meta: Option<BlockMetaStorage>,
     subscribe_id: AtomicUsize,
-    snapshot_rx: Mutex<Option<crossbeam_channel::Receiver<Option<Message>>>>,
+    snapshot_rx: Mutex<Option<crossbeam_channel::Receiver<Box<Message>>>>,
     broadcast_tx: broadcast::Sender<(CommitmentLevel, Arc<Vec<Arc<Message>>>)>,
     debug_clients_tx: Option<mpsc::UnboundedSender<DebugClientMessage>>,
 }
@@ -741,7 +741,7 @@ impl GrpcService {
         debug_clients_tx: Option<mpsc::UnboundedSender<DebugClientMessage>>,
         is_reload: bool,
     ) -> anyhow::Result<(
-        Option<crossbeam_channel::Sender<Option<Message>>>,
+        Option<crossbeam_channel::Sender<Box<Message>>>,
         mpsc::UnboundedSender<Arc<Message>>,
         Arc<Notify>,
     )> {
@@ -1125,7 +1125,7 @@ impl GrpcService {
         config_filters: Arc<ConfigGrpcFilters>,
         stream_tx: mpsc::Sender<TonicResult<SubscribeUpdate>>,
         mut client_rx: mpsc::UnboundedReceiver<Option<Filter>>,
-        mut snapshot_rx: Option<crossbeam_channel::Receiver<Option<Message>>>,
+        mut snapshot_rx: Option<crossbeam_channel::Receiver<Box<Message>>>,
         mut messages_rx: broadcast::Receiver<(CommitmentLevel, Arc<Vec<Arc<Message>>>)>,
         debug_client_tx: Option<mpsc::UnboundedSender<DebugClientMessage>>,
         drop_client: impl FnOnce(),
@@ -1256,7 +1256,7 @@ impl GrpcService {
         endpoint: &str,
         stream_tx: &mpsc::Sender<TonicResult<SubscribeUpdate>>,
         client_rx: &mut mpsc::UnboundedReceiver<Option<Filter>>,
-        snapshot_rx: crossbeam_channel::Receiver<Option<Message>>,
+        snapshot_rx: crossbeam_channel::Receiver<Box<Message>>,
         is_alive: &mut bool,
         filter: &mut Filter,
     ) {
@@ -1292,18 +1292,14 @@ impl GrpcService {
             let message = match snapshot_rx.try_recv() {
                 Ok(message) => {
                     MESSAGE_QUEUE_SIZE.dec();
-                    match message {
-                        Some(message) => message,
-                        None => break,
-                    }
+                    message
                 }
                 Err(crossbeam_channel::TryRecvError::Empty) => {
                     sleep(Duration::from_millis(1)).await;
                     continue;
                 }
                 Err(crossbeam_channel::TryRecvError::Disconnected) => {
-                    error!("client #{id}: snapshot channel disconnected");
-                    *is_alive = false;
+                    info!("client #{id}: end of startup");
                     break;
                 }
             };

yellowstone-grpc-geyser/src/plugin.rs

@@ -13,7 +13,7 @@ use {
         concat, env,
         sync::{
             atomic::{AtomicBool, Ordering},
-            Arc,
+            Arc, Mutex,
         },
         time::Duration,
     },
@@ -26,7 +26,7 @@ use {
 #[derive(Debug)]
 pub struct PluginInner {
     runtime: Runtime,
-    snapshot_channel: Option<crossbeam_channel::Sender<Option<Message>>>,
+    snapshot_channel: Mutex<Option<crossbeam_channel::Sender<Box<Message>>>>,
     snapshot_channel_closed: AtomicBool,
     grpc_channel: mpsc::UnboundedSender<Arc<Message>>,
     grpc_shutdown: Arc<Notify>,
@@ -101,7 +101,7 @@ impl GeyserPlugin for Plugin {
 
         self.inner = Some(PluginInner {
             runtime,
-            snapshot_channel,
+            snapshot_channel: Mutex::new(snapshot_channel),
             snapshot_channel_closed: AtomicBool::new(false),
             grpc_channel,
             grpc_shutdown,
@@ -137,10 +137,10 @@ impl GeyserPlugin for Plugin {
                 ReplicaAccountInfoVersions::V0_0_3(info) => info,
             };
 
-            let message = Message::Account((account, slot, is_startup).into());
             if is_startup {
-                if let Some(channel) = &inner.snapshot_channel {
-                    match channel.send(Some(message)) {
+                if let Some(channel) = inner.snapshot_channel.lock().unwrap().as_ref() {
+                    let message = Message::Account((account, slot, is_startup).into());
+                    match channel.send(Box::new(message)) {
                         Ok(()) => MESSAGE_QUEUE_SIZE.inc(),
                         Err(_) => {
                             if !inner.snapshot_channel_closed.swap(true, Ordering::Relaxed) {
@@ -152,6 +152,7 @@ impl GeyserPlugin for Plugin {
                     }
                 }
             } else {
+                let message = Message::Account((account, slot, is_startup).into());
                 inner.send_message(message);
             }
 
@@ -161,12 +162,7 @@ impl GeyserPlugin for Plugin {
 
     fn notify_end_of_startup(&self) -> PluginResult<()> {
         self.with_inner(|inner| {
-            if let Some(channel) = &inner.snapshot_channel {
-                match channel.send(None) {
-                    Ok(()) => MESSAGE_QUEUE_SIZE.inc(),
-                    Err(_) => panic!("failed to send message to startup queue: channel closed"),
-                }
-            }
+            let _snapshot_channel = inner.snapshot_channel.lock().unwrap().take();
             Ok(())
         })
     }