linera_core/
notifier.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use linera_base::identifiers::ChainId;
7use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
8use tracing::trace;
9
10use crate::worker;
11
12// TODO(#2171): replace this with a Tokio broadcast channel
13
14/// A `Notifier` holds references to clients waiting to receive notifications
15/// from the validator.
16/// Clients will be evicted if their connections are terminated.
17pub struct ChannelNotifier<N> {
18    inner: papaya::HashMap<ChainId, Vec<UnboundedSender<N>>>,
19}
20
21impl<N> Default for ChannelNotifier<N> {
22    fn default() -> Self {
23        Self {
24            inner: papaya::HashMap::default(),
25        }
26    }
27}
28
29impl<N> ChannelNotifier<N> {
30    /// Registers a sender for notifications on the given chain IDs.
31    pub fn add_sender(&self, chain_ids: Vec<ChainId>, sender: &UnboundedSender<N>) {
32        let pinned = self.inner.pin();
33        for id in chain_ids {
34            pinned.update_or_insert_with(
35                id,
36                |senders| senders.iter().cloned().chain([sender.clone()]).collect(),
37                || vec![sender.clone()],
38            );
39        }
40    }
41
42    /// Creates a subscription given a collection of chain IDs and a sender to the client.
43    pub fn subscribe(&self, chain_ids: Vec<ChainId>) -> UnboundedReceiver<N> {
44        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
45        self.add_sender(chain_ids, &tx);
46        rx
47    }
48
49    /// Creates a subscription given a collection of chain IDs and a sender to the client.
50    /// Immediately posts a first notification as an ACK.
51    pub fn subscribe_with_ack(&self, chain_ids: Vec<ChainId>, ack: N) -> UnboundedReceiver<N> {
52        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
53        self.add_sender(chain_ids, &tx);
54        tx.send(ack)
55            .expect("pushing to a new channel should succeed");
56        rx
57    }
58}
59
60impl<N> ChannelNotifier<N>
61where
62    N: Clone,
63{
64    /// Notifies all the clients waiting for a notification from a given chain.
65    pub fn notify_chain(&self, chain_id: &ChainId, notification: &N) {
66        let pinned = self.inner.pin();
67
68        // Read senders outside of `compute` to avoid side effects in a
69        // retriable closure. papaya's `compute` may call its closure
70        // multiple times on CAS contention, so `send()` must not happen
71        // inside it.
72        let Some(senders) = pinned.get(chain_id).cloned() else {
73            trace!("Chain {chain_id} has no subscribers.");
74            return;
75        };
76
77        // Send notifications (side effect — must happen exactly once).
78        let mut has_dead = false;
79        for sender in &senders {
80            if sender.send(notification.clone()).is_err() {
81                has_dead = true;
82            }
83        }
84
85        // Clean up dead senders. The closure is pure: `is_closed()` is
86        // idempotent and has no side effects, so retries are safe.
87        if has_dead {
88            pinned.compute(*chain_id, |entry| {
89                let Some((_key, current_senders)) = entry else {
90                    return papaya::Operation::Abort(());
91                };
92                let live: Vec<_> = current_senders
93                    .iter()
94                    .filter(|s| !s.is_closed())
95                    .cloned()
96                    .collect();
97                if live.is_empty() {
98                    trace!("No more subscribers for chain {chain_id}. Removing entry.");
99                    papaya::Operation::Remove
100                } else {
101                    papaya::Operation::Insert(live)
102                }
103            });
104        }
105    }
106}
107
108/// A type that can dispatch notifications produced by the worker.
109pub trait Notifier: Clone + Send + 'static {
110    /// Dispatches the given notifications to the interested clients.
111    fn notify(&self, notifications: &[worker::Notification]);
112}
113
114impl Notifier for Arc<ChannelNotifier<worker::Notification>> {
115    fn notify(&self, notifications: &[worker::Notification]) {
116        for notification in notifications {
117            self.notify_chain(&notification.chain_id, notification);
118        }
119    }
120}
121
122impl Notifier for () {
123    fn notify(&self, _notifications: &[worker::Notification]) {}
124}
125
126#[cfg(with_testing)]
127impl Notifier for Arc<std::sync::Mutex<Vec<worker::Notification>>> {
128    fn notify(&self, notifications: &[worker::Notification]) {
129        let mut guard = self.lock().unwrap();
130        guard.extend(notifications.iter().cloned())
131    }
132}
133
134#[cfg(test)]
135/// Unit tests for the [`Notifier`].
136pub mod tests {
137    use std::{
138        sync::{atomic::Ordering, Arc},
139        time::Duration,
140    };
141
142    use linera_execution::test_utils::dummy_chain_description;
143
144    use super::*;
145
146    #[test]
147    fn test_concurrent() {
148        let notifier = ChannelNotifier::default();
149
150        let chain_a = dummy_chain_description(0).id();
151        let chain_b = dummy_chain_description(1).id();
152
153        let a_rec = Arc::new(std::sync::atomic::AtomicUsize::new(0));
154        let b_rec = Arc::new(std::sync::atomic::AtomicUsize::new(0));
155        let a_b_rec = Arc::new(std::sync::atomic::AtomicUsize::new(0));
156
157        let mut rx_a = notifier.subscribe(vec![chain_a]);
158        let mut rx_b = notifier.subscribe(vec![chain_b]);
159        let mut rx_a_b = notifier.subscribe(vec![chain_a, chain_b]);
160
161        let a_rec_clone = a_rec.clone();
162        let b_rec_clone = b_rec.clone();
163        let a_b_rec_clone = a_b_rec.clone();
164
165        let notifier = Arc::new(notifier);
166
167        std::thread::spawn(move || {
168            while rx_a.blocking_recv().is_some() {
169                a_rec_clone.fetch_add(1, Ordering::Relaxed);
170            }
171        });
172
173        std::thread::spawn(move || {
174            while rx_b.blocking_recv().is_some() {
175                b_rec_clone.fetch_add(1, Ordering::Relaxed);
176            }
177        });
178
179        std::thread::spawn(move || {
180            while rx_a_b.blocking_recv().is_some() {
181                a_b_rec_clone.fetch_add(1, Ordering::Relaxed);
182            }
183        });
184
185        const NOTIFICATIONS_A: usize = 500;
186        const NOTIFICATIONS_B: usize = 700;
187
188        let a_notifier = notifier.clone();
189        let handle_a = std::thread::spawn(move || {
190            for _ in 0..NOTIFICATIONS_A {
191                a_notifier.notify_chain(&chain_a, &());
192            }
193        });
194
195        let handle_b = std::thread::spawn(move || {
196            for _ in 0..NOTIFICATIONS_B {
197                notifier.notify_chain(&chain_b, &());
198            }
199        });
200
201        // finish sending all the messages
202        handle_a.join().unwrap();
203        handle_b.join().unwrap();
204
205        // give some time for the messages to be received.
206        std::thread::sleep(Duration::from_millis(100));
207
208        assert_eq!(a_rec.load(Ordering::Relaxed), NOTIFICATIONS_A);
209        assert_eq!(b_rec.load(Ordering::Relaxed), NOTIFICATIONS_B);
210        assert_eq!(
211            a_b_rec.load(Ordering::Relaxed),
212            NOTIFICATIONS_A + NOTIFICATIONS_B
213        );
214    }
215
216    #[test]
217    fn test_eviction() {
218        let notifier = ChannelNotifier::default();
219
220        let chain_a = dummy_chain_description(0).id();
221        let chain_b = dummy_chain_description(1).id();
222        let chain_c = dummy_chain_description(2).id();
223        let chain_d = dummy_chain_description(3).id();
224
225        // Chain A -> Notify A, Notify B
226        // Chain B -> Notify A, Notify B
227        // Chain C -> Notify C
228        // Chain D -> Notify A, Notify B, Notify C, Notify D
229
230        let mut rx_a = notifier.subscribe(vec![chain_a, chain_b, chain_d]);
231        let mut rx_b = notifier.subscribe(vec![chain_a, chain_b, chain_d]);
232        let mut rx_c = notifier.subscribe(vec![chain_c, chain_d]);
233        let mut rx_d = notifier.subscribe(vec![chain_d]);
234
235        assert_eq!(notifier.inner.len(), 4);
236
237        rx_c.close();
238        notifier.notify_chain(&chain_c, &());
239        assert_eq!(notifier.inner.len(), 3);
240
241        rx_a.close();
242        notifier.notify_chain(&chain_a, &());
243        assert_eq!(notifier.inner.len(), 3);
244
245        rx_b.close();
246        notifier.notify_chain(&chain_b, &());
247        assert_eq!(notifier.inner.len(), 2);
248
249        notifier.notify_chain(&chain_a, &());
250        assert_eq!(notifier.inner.len(), 1);
251
252        rx_d.close();
253        notifier.notify_chain(&chain_d, &());
254        assert_eq!(notifier.inner.len(), 0);
255    }
256}
linera_core/notifier.rs

linera_core/
notifier.rs
