geyser: support snapshot data

CHANGELOG.md

@@ -12,6 +12,8 @@ The minor version will be incremented upon a breaking change and the patch versi
 
 ### Features
 
+- geyser: support snapshot data ([#180](https://github.com/rpcpool/yellowstone-grpc/pull/180)).
+
 ### Fixes
 
 ### Breaking

yellowstone-grpc-geyser/Cargo.toml

@@ -18,6 +18,7 @@ base64 = "0.21.0"
 bincode = "1.3.3"
 bs58 = "0.4.0"
 clap = { version = "4.3.0", features = ["cargo", "derive"] }
+crossbeam-channel = "0.5.8"
 futures = "0.3.24"
 hyper = { version = "0.14.20", features = ["server"] }
 lazy_static = "1.4.0"

yellowstone-grpc-geyser/config.json

@@ -5,6 +5,7 @@
     },
     "grpc": {
         "address": "0.0.0.0:10000",
+        "snapshot_capacity": null,
         "channel_capacity": "100_000",
         "unary_concurrency_limit": 100,
         "unary_disabled": false,

yellowstone-grpc-geyser/src/config.rs

@@ -62,6 +62,13 @@ impl ConfigLog {
 pub struct ConfigGrpc {
     /// Address of Grpc service.
     pub address: SocketAddr,
+    /// Capacity of the channel used for account from snapshot,
+    /// on reaching the limit Sender block validator startup.
+    #[serde(
+        default = "ConfigGrpc::snapshot_capacity_default",
+        deserialize_with = "deserialize_usize_str_maybe"
+    )]
+    pub snapshot_capacity: Option<usize>,
     /// Capacity of the channel per connection
     #[serde(
         default = "ConfigGrpc::channel_capacity_default",
@@ -83,6 +90,10 @@ pub struct ConfigGrpc {
 }
 
 impl ConfigGrpc {
+    const fn snapshot_capacity_default() -> Option<usize> {
+        None
+    }
+
     const fn channel_capacity_default() -> usize {
         250_000
     }
@@ -304,6 +315,28 @@ where
     }
 }
 
+fn deserialize_usize_str_maybe<'de, D>(deserializer: D) -> Result<Option<usize>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    #[derive(Deserialize)]
+    #[serde(untagged)]
+    enum Value {
+        Integer(usize),
+        String(String),
+    }
+
+    match Option::<Value>::deserialize(deserializer)? {
+        Some(Value::Integer(value)) => Ok(Some(value)),
+        Some(Value::String(value)) => value
+            .replace('_', "")
+            .parse::<usize>()
+            .map(Some)
+            .map_err(de::Error::custom),
+        None => Ok(None),
+    }
+}
+
 fn deserialize_pubkey_set<'de, D>(deserializer: D) -> Result<HashSet<Pubkey>, D::Error>
 where
     D: Deserializer<'de>,

yellowstone-grpc-geyser/src/grpc.rs

@@ -37,7 +37,8 @@ use {
         },
     },
     tokio::{
-        sync::{broadcast, mpsc, Notify, RwLock, Semaphore},
+        sync::{broadcast, mpsc, Mutex, Notify, RwLock, Semaphore},
+        task::yield_now,
         time::{sleep, Duration, Instant},
     },
     tokio_stream::wrappers::ReceiverStream,
@@ -675,15 +676,21 @@ pub struct GrpcService {
     config: ConfigGrpc,
     blocks_meta: Option<BlockMetaStorage>,
     subscribe_id: AtomicUsize,
+    snapshot_rx: Mutex<Option<crossbeam_channel::Receiver<Option<Message>>>>,
     broadcast_tx: broadcast::Sender<(CommitmentLevel, Arc<Vec<Message>>)>,
 }
 
 impl GrpcService {
+    #[allow(clippy::type_complexity)]
     pub fn create(
         config: ConfigGrpc,
         block_fail_action: ConfigBlockFailAction,
     ) -> Result<
-        (mpsc::UnboundedSender<Message>, Arc<Notify>),
+        (
+            Option<crossbeam_channel::Sender<Option<Message>>>,
+            mpsc::UnboundedSender<Message>,
+            Arc<Notify>,
+        ),
         Box<dyn std::error::Error + Send + Sync>,
     > {
         // Bind service address
@@ -693,6 +700,15 @@ impl GrpcService {
             Some(Duration::from_secs(20)), // tcp_keepalive
         )?;
 
+        // Snapshot channel
+        let (snapshot_tx, snapshot_rx) = match config.snapshot_capacity {
+            Some(cap) => {
+                let (tx, rx) = crossbeam_channel::bounded(cap);
+                (Some(tx), Some(rx))
+            }
+            None => (None, None),
+        };
+
         // Blocks meta storage
         let (blocks_meta, blocks_meta_tx) = if config.unary_disabled {
             (None, None)
@@ -710,6 +726,7 @@ impl GrpcService {
             config,
             blocks_meta,
             subscribe_id: AtomicUsize::new(0),
+            snapshot_rx: Mutex::new(snapshot_rx),
             broadcast_tx: broadcast_tx.clone(),
         })
         .accept_compressed(CompressionEncoding::Gzip)
@@ -740,7 +757,7 @@ impl GrpcService {
                 .await
         });
 
-        Ok((messages_tx, shutdown))
+        Ok((snapshot_tx, messages_tx, shutdown))
     }
 
     async fn geyser_loop(
@@ -1018,65 +1035,118 @@ impl GrpcService {
         mut filter: Filter,
         stream_tx: mpsc::Sender<TonicResult<SubscribeUpdate>>,
         mut client_rx: mpsc::UnboundedReceiver<Option<Filter>>,
+        mut snapshot_rx: Option<crossbeam_channel::Receiver<Option<Message>>>,
         mut messages_rx: broadcast::Receiver<(CommitmentLevel, Arc<Vec<Message>>)>,
         drop_client: impl FnOnce(),
     ) {
         CONNECTIONS_TOTAL.inc();
         info!("client #{id}: new");
-        'outer: loop {
-            tokio::select! {
-                message = client_rx.recv() => {
-                    match message {
-                        Some(Some(filter_new)) => {
-                            filter = filter_new;
-                            info!("client #{id}: filter updated");
-                        }
-                        Some(None) => {
-                            break 'outer;
-                        },
-                        None => {
-                            break 'outer;
+
+        let mut is_alive = true;
+        if let Some(snapshot_rx) = snapshot_rx.take() {
+            info!("client #{id}: going to receive snapshot data");
+
+            // we start with default filter, for snapshot we need wait actual filter first
+            match client_rx.recv().await {
+                Some(Some(filter_new)) => {
+                    filter = filter_new;
+                    info!("client #{id}: filter updated");
+                }
+                Some(None) => {
+                    is_alive = false;
+                }
+                None => {
+                    is_alive = false;
+                }
+            };
+
+            while is_alive {
+                let message = match snapshot_rx.try_recv() {
+                    Ok(message) => {
+                        MESSAGE_QUEUE_SIZE.dec();
+                        match message {
+                            Some(message) => message,
+                            None => break,
                         }
                     }
+                    Err(crossbeam_channel::TryRecvError::Empty) => {
+                        yield_now().await;
+                        continue;
+                    }
+                    Err(crossbeam_channel::TryRecvError::Disconnected) => {
+                        error!("client #{id}: snapshot channel disconnected");
+                        is_alive = false;
+                        break;
+                    }
+                };
+
+                for message in filter.get_update(&message) {
+                    if stream_tx.send(Ok(message)).await.is_err() {
+                        error!("client #{id}: stream closed");
+                        is_alive = false;
+                        break;
+                    }
                 }
-                message = messages_rx.recv() => {
-                    match message {
-                        Ok((commitment, messages)) => {
-                            if commitment == filter.get_commitment_level() {
-                                for message in messages.iter() {
-                                    for message in filter.get_update(message) {
-                                        match stream_tx.try_send(Ok(message)) {
-                                            Ok(()) => {}
-                                            Err(mpsc::error::TrySendError::Full(_)) => {
-                                                error!("client #{id}: lagged to send update");
-                                                tokio::spawn(async move {
-                                                    let _ = stream_tx.send(Err(Status::internal("lagged"))).await;
-                                                });
-                                                break 'outer;
-                                            }
-                                            Err(mpsc::error::TrySendError::Closed(_)) => {
-                                                error!("client #{id}: stream closed");
-                                                break 'outer;
-                                            }
+            }
+        }
+
+        if is_alive {
+            'outer: loop {
+                tokio::select! {
+                    message = client_rx.recv() => {
+                        match message {
+                            Some(Some(filter_new)) => {
+                                filter = filter_new;
+                                info!("client #{id}: filter updated");
+                            }
+                            Some(None) => {
+                                break 'outer;
+                            },
+                            None => {
+                                break 'outer;
+                            }
+                        }
+                    }
+                    message = messages_rx.recv() => {
+                        let (commitment, messages) = match message {
+                            Ok((commitment, messages)) => (commitment, messages),
+                            Err(broadcast::error::RecvError::Closed) => {
+                                break 'outer;
+                            },
+                            Err(broadcast::error::RecvError::Lagged(_)) => {
+                                info!("client #{id}: lagged to receive geyser messages");
+                                tokio::spawn(async move {
+                                    let _ = stream_tx.send(Err(Status::internal("lagged"))).await;
+                                });
+                                break 'outer;
+                            }
+                        };
+
+                        if commitment == filter.get_commitment_level() {
+                            for message in messages.iter() {
+                                for message in filter.get_update(message) {
+                                    match stream_tx.try_send(Ok(message)) {
+                                        Ok(()) => {}
+                                        Err(mpsc::error::TrySendError::Full(_)) => {
+                                            error!("client #{id}: lagged to send update");
+                                            tokio::spawn(async move {
+                                                let _ = stream_tx.send(Err(Status::internal("lagged"))).await;
+                                            });
+                                            break 'outer;
+                                        }
+                                        Err(mpsc::error::TrySendError::Closed(_)) => {
+                                            error!("client #{id}: stream closed");
+                                            break 'outer;
                                         }
                                     }
                                 }
                             }
                         }
-                        Err(broadcast::error::RecvError::Closed) => {
-                            break 'outer;
-                        },
-                        Err(broadcast::error::RecvError::Lagged(_)) => {
-                            info!("client #{id}: lagged to receive geyser messages");
-                            tokio::spawn(async move {
-                                let _ = stream_tx.send(Err(Status::internal("lagged"))).await;
-                            });
-                            break 'outer;
-                        }
                     }
                 }
             }
         }
+
         info!("client #{id}: removed");
         CONNECTIONS_TOTAL.dec();
         drop_client();
@@ -1189,6 +1259,7 @@ impl Geyser for GrpcService {
             filter,
             stream_tx,
             client_rx,
+            self.snapshot_rx.lock().await.take(),
             self.broadcast_tx.subscribe(),
             move || {
                 notify_exit1.notify_one();