Send notifications from the proxy in batches

Author: ndr-ds

linera-core/src/client/chain_client/mod.rs

@@ -442,9 +442,11 @@ impl<Env: Environment> ChainClient<Env> {
     /// Subscribes to notifications from the specified chain.
     #[instrument(level = "trace")]
     pub fn subscribe_to(&self, chain_id: ChainId) -> Result<NotificationStream, LocalNodeError> {
-        Ok(Box::pin(UnboundedReceiverStream::new(
-            self.client.notifier.subscribe(vec![chain_id]),
-        )))
+        // Flatten the batched notifications into individual items
+        Ok(Box::pin(
+            UnboundedReceiverStream::new(self.client.notifier.subscribe(vec![chain_id]))
+                .flat_map(stream::iter),
+        ))
     }
 
     /// Returns the storage client used by this client's local node.

linera-core/src/notifier.rs

@@ -14,8 +14,9 @@ use crate::worker;
 /// A `Notifier` holds references to clients waiting to receive notifications
 /// from the validator.
 /// Clients will be evicted if their connections are terminated.
+/// Notifications are sent in batches (Vec<N>) to reduce channel operations.
 pub struct ChannelNotifier<N> {
-    inner: papaya::HashMap<ChainId, Vec<UnboundedSender<N>>>,
+    inner: papaya::HashMap<ChainId, Vec<UnboundedSender<Vec<N>>>>,
 }
 
 impl<N> Default for ChannelNotifier<N> {
@@ -27,7 +28,7 @@ impl<N> Default for ChannelNotifier<N> {
 }
 
 impl<N> ChannelNotifier<N> {
-    fn add_sender(&self, chain_ids: Vec<ChainId>, sender: &UnboundedSender<N>) {
+    fn add_sender(&self, chain_ids: Vec<ChainId>, sender: &UnboundedSender<Vec<N>>) {
         let pinned = self.inner.pin();
         for id in chain_ids {
             pinned.update_or_insert_with(
@@ -39,18 +40,19 @@ impl<N> ChannelNotifier<N> {
     }
 
     /// Creates a subscription given a collection of chain IDs and a sender to the client.
-    pub fn subscribe(&self, chain_ids: Vec<ChainId>) -> UnboundedReceiver<N> {
+    /// Returns a receiver that yields batches of notifications (Vec<N>).
+    pub fn subscribe(&self, chain_ids: Vec<ChainId>) -> UnboundedReceiver<Vec<N>> {
         let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
         self.add_sender(chain_ids, &tx);
         rx
     }
 
     /// Creates a subscription given a collection of chain IDs and a sender to the client.
-    /// Immediately posts a first notification as an ACK.
-    pub fn subscribe_with_ack(&self, chain_ids: Vec<ChainId>, ack: N) -> UnboundedReceiver<N> {
+    /// Immediately posts a first notification batch as an ACK.
+    pub fn subscribe_with_ack(&self, chain_ids: Vec<ChainId>, ack: N) -> UnboundedReceiver<Vec<N>> {
         let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
         self.add_sender(chain_ids, &tx);
-        tx.send(ack)
+        tx.send(vec![ack])
             .expect("pushing to a new channel should succeed");
         rx
     }
@@ -61,15 +63,19 @@ where
     N: Clone,
 {
     /// Notifies all the clients waiting for a notification from a given chain.
-    pub fn notify_chain(&self, chain_id: &ChainId, notification: &N) {
+    /// Sends a batch of notifications to all subscribers of the chain.
+    pub fn notify_chain(&self, chain_id: &ChainId, notifications: Vec<N>) {
+        if notifications.is_empty() {
+            return;
+        }
         self.inner.pin().compute(*chain_id, |senders| {
             let Some((_key, senders)) = senders else {
                 trace!("Chain {chain_id} has no subscribers.");
                 return papaya::Operation::Abort(());
             };
             let live_senders = senders
                 .iter()
-                .filter(|sender| sender.send(notification.clone()).is_ok())
+                .filter(|sender| sender.send(notifications.clone()).is_ok())
                 .cloned()
                 .collect::<Vec<_>>();
             if live_senders.is_empty() {
@@ -87,8 +93,17 @@ pub trait Notifier: Clone + Send + 'static {
 
 impl Notifier for Arc<ChannelNotifier<worker::Notification>> {
     fn notify(&self, notifications: &[worker::Notification]) {
+        // Group notifications by chain_id to send them in batches
+        let mut by_chain: std::collections::HashMap<ChainId, Vec<worker::Notification>> =
+            std::collections::HashMap::new();
         for notification in notifications {
-            self.notify_chain(&notification.chain_id, notification);
+            by_chain
+                .entry(notification.chain_id)
+                .or_default()
+                .push(notification.clone());
+        }
+        for (chain_id, chain_notifications) in by_chain {
+            self.notify_chain(&chain_id, chain_notifications);
         }
     }
 }
@@ -138,20 +153,20 @@ pub mod tests {
         let notifier = Arc::new(notifier);
 
         std::thread::spawn(move || {
-            while rx_a.blocking_recv().is_some() {
-                a_rec_clone.fetch_add(1, Ordering::Relaxed);
+            while let Some(batch) = rx_a.blocking_recv() {
+                a_rec_clone.fetch_add(batch.len(), Ordering::Relaxed);
             }
         });
 
         std::thread::spawn(move || {
-            while rx_b.blocking_recv().is_some() {
-                b_rec_clone.fetch_add(1, Ordering::Relaxed);
+            while let Some(batch) = rx_b.blocking_recv() {
+                b_rec_clone.fetch_add(batch.len(), Ordering::Relaxed);
             }
         });
 
         std::thread::spawn(move || {
-            while rx_a_b.blocking_recv().is_some() {
-                a_b_rec_clone.fetch_add(1, Ordering::Relaxed);
+            while let Some(batch) = rx_a_b.blocking_recv() {
+                a_b_rec_clone.fetch_add(batch.len(), Ordering::Relaxed);
             }
         });
 
@@ -161,13 +176,13 @@ pub mod tests {
         let a_notifier = notifier.clone();
         let handle_a = std::thread::spawn(move || {
             for _ in 0..NOTIFICATIONS_A {
-                a_notifier.notify_chain(&chain_a, &());
+                a_notifier.notify_chain(&chain_a, vec![()]);
             }
         });
 
         let handle_b = std::thread::spawn(move || {
             for _ in 0..NOTIFICATIONS_B {
-                notifier.notify_chain(&chain_b, &());
+                notifier.notify_chain(&chain_b, vec![()]);
             }
         });
 
@@ -208,22 +223,22 @@ pub mod tests {
         assert_eq!(notifier.inner.len(), 4);
 
         rx_c.close();
-        notifier.notify_chain(&chain_c, &());
+        notifier.notify_chain(&chain_c, vec![()]);
         assert_eq!(notifier.inner.len(), 3);
 
         rx_a.close();
-        notifier.notify_chain(&chain_a, &());
+        notifier.notify_chain(&chain_a, vec![()]);
         assert_eq!(notifier.inner.len(), 3);
 
         rx_b.close();
-        notifier.notify_chain(&chain_b, &());
+        notifier.notify_chain(&chain_b, vec![()]);
         assert_eq!(notifier.inner.len(), 2);
 
-        notifier.notify_chain(&chain_a, &());
+        notifier.notify_chain(&chain_a, vec![()]);
         assert_eq!(notifier.inner.len(), 1);
 
         rx_d.close();
-        notifier.notify_chain(&chain_d, &());
+        notifier.notify_chain(&chain_d, vec![()]);
         assert_eq!(notifier.inner.len(), 0);
     }
 }

linera-core/src/unit_tests/test_utils.rs

@@ -12,6 +12,8 @@ use async_trait::async_trait;
 use futures::{
     future::Either,
     lock::{Mutex, MutexGuard},
+    stream,
+    stream::StreamExt as _,
     Future,
 };
 use linera_base::{
@@ -492,7 +494,9 @@ where
     ) -> Result<(), Result<NotificationStream, NodeError>> {
         let validator = self.client.lock().await;
         let rx = validator.notifier.subscribe(chains);
-        let stream: NotificationStream = Box::pin(UnboundedReceiverStream::new(rx));
+        // Flatten the batched notifications into individual items
+        let stream: NotificationStream =
+            Box::pin(UnboundedReceiverStream::new(rx).flat_map(stream::iter));
         sender.send(Ok(stream))
     }
 

linera-service/src/proxy/grpc.rs

@@ -15,7 +15,7 @@ use std::{
 
 use anyhow::Result;
 use async_trait::async_trait;
-use futures::{future::BoxFuture, FutureExt as _};
+use futures::{future::BoxFuture, stream, stream::BoxStream, stream::StreamExt as _, FutureExt as _};
 use linera_base::identifiers::ChainId;
 use linera_core::{
     data_types::{CertificatesByHeightRequest, ChainInfo, ChainInfoQuery},
@@ -383,7 +383,7 @@ impl<S> ValidatorNode for GrpcProxy<S>
 where
     S: Storage + Clone + Send + Sync + 'static,
 {
-    type SubscribeStream = UnboundedReceiverStream<Result<Notification, Status>>;
+    type SubscribeStream = BoxStream<'static, Result<Notification, Status>>;
 
     #[instrument(skip_all, err(Display), fields(method = "handle_block_proposal"))]
     async fn handle_block_proposal(
@@ -482,7 +482,9 @@ where
             .0
             .notifier
             .subscribe_with_ack(chain_ids, Ok(Notification::default()));
-        Ok(Response::new(UnboundedReceiverStream::new(rx)))
+        // Flatten the batched notifications into individual items for the gRPC stream
+        let flattened = UnboundedReceiverStream::new(rx).flat_map(stream::iter);
+        Ok(Response::new(Box::pin(flattened)))
     }
 
     #[instrument(skip_all, err(Display))]
@@ -855,13 +857,23 @@ where
         &self,
         request: Request<NotificationBatch>,
     ) -> Result<Response<()>, Status> {
+        // Group notifications by chain_id to send them in batches
+        let mut by_chain: std::collections::HashMap<
+            ChainId,
+            Vec<Result<Notification, Status>>,
+        > = std::collections::HashMap::new();
+
         for notification in request.into_inner().notifications {
             let chain_id = notification
                 .chain_id
                 .clone()
                 .ok_or_else(|| Status::invalid_argument("Missing field: chain_id."))?
                 .try_into()?;
-            self.0.notifier.notify_chain(&chain_id, &Ok(notification));
+            by_chain.entry(chain_id).or_default().push(Ok(notification));
+        }
+
+        for (chain_id, notifications) in by_chain {
+            self.0.notifier.notify_chain(&chain_id, notifications);
         }
         Ok(Response::new(()))
     }