iota_core/
streamer.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::BTreeMap, fmt::Debug, sync::Arc};
6
7use futures::Stream;
8use iota_json_rpc_types::Filter;
9use iota_metrics::{metered_channel::Sender, spawn_monitored_task};
10use iota_sdk_types::ObjectId;
11use iota_types::error::IotaError;
12use parking_lot::RwLock;
13use prometheus_filtered::Registry;
14use tokio::sync::mpsc;
15use tokio_stream::wrappers::ReceiverStream;
16use tracing::{debug, warn};
17
18use crate::subscription_handler::{EVENT_DISPATCH_BUFFER_SIZE, SubscriptionMetrics};
19
20type Subscribers<T, F> = Arc<RwLock<BTreeMap<String, (tokio::sync::mpsc::Sender<T>, F)>>>;
21
22/// The Streamer splits a mpsc channel into multiple mpsc channels using the
23/// subscriber's `Filter<T>` object. Data will be sent to the subscribers in
24/// parallel and the subscription will be dropped if it received a send error.
25pub struct Streamer<T, S, F: Filter<T>> {
26    streamer_queue: Sender<T>,
27    subscribers: Subscribers<S, F>,
28    metrics: Arc<SubscriptionMetrics>,
29    metrics_label: &'static str,
30}
31
32impl<T, S, F> Streamer<T, S, F>
33where
34    S: From<T> + Clone + Debug + Send + Sync + 'static,
35    T: Clone + Send + Sync + 'static,
36    F: Filter<T> + Clone + Send + Sync + 'static + Clone,
37{
38    pub fn spawn(
39        buffer: usize,
40        metrics: Arc<SubscriptionMetrics>,
41        metrics_label: &'static str,
42    ) -> Self {
43        let channel_label = format!("streamer_{metrics_label}");
44        let gauge = if let Some(metrics) = iota_metrics::get_metrics() {
45            metrics
46                .channel_inflight
47                .with_label_values(&[&channel_label])
48        } else {
49            // We call init_metrics very early when starting a node. Therefore when this
50            // happens, it's probably in a test.
51            iota_metrics::init_metrics(&Registry::default());
52            iota_metrics::get_metrics()
53                .unwrap()
54                .channel_inflight
55                .with_label_values(&[&channel_label])
56        };
57
58        let (tx, rx) = iota_metrics::metered_channel::channel(buffer, &gauge);
59        let streamer = Self {
60            streamer_queue: tx,
61            subscribers: Default::default(),
62            metrics: metrics.clone(),
63            metrics_label,
64        };
65        let mut rx = rx;
66        let subscribers = streamer.subscribers.clone();
67        spawn_monitored_task!(async move {
68            while let Some(data) = rx.recv().await {
69                Self::send_to_all_subscribers(
70                    subscribers.clone(),
71                    data,
72                    metrics.clone(),
73                    metrics_label,
74                )
75                .await;
76            }
77        });
78        streamer
79    }
80
81    async fn send_to_all_subscribers(
82        subscribers: Subscribers<S, F>,
83        data: T,
84        metrics: Arc<SubscriptionMetrics>,
85        metrics_label: &'static str,
86    ) {
87        let success_counter = metrics
88            .streaming_success
89            .with_label_values(&[metrics_label]);
90        let failure_counter = metrics
91            .streaming_failure
92            .with_label_values(&[metrics_label]);
93        let subscriber_count = metrics
94            .streaming_active_subscriber_number
95            .with_label_values(&[metrics_label]);
96
97        let to_remove = {
98            let mut to_remove = vec![];
99            let subscribers_snapshot = subscribers.read();
100            subscriber_count.set(subscribers_snapshot.len() as i64);
101
102            for (id, (subscriber, filter)) in subscribers_snapshot.iter() {
103                if !(filter.matches(&data)) {
104                    continue;
105                }
106                let data = data.clone();
107                match subscriber.try_send(data.into()) {
108                    Ok(_) => {
109                        debug!(subscription_id = id, "Streaming data to subscriber.");
110                        success_counter.inc();
111                    }
112                    Err(e) => {
113                        warn!(
114                            subscription_id = id,
115                            "Error when streaming data, removing subscriber. Error: {e}"
116                        );
117                        // It does not matter what the error is - channel full or closed, we remove
118                        // the subscriber. In the case of a full channel,
119                        // this nudges the subscriber to catch up separately and not
120                        // miss any data.
121                        to_remove.push(id.clone());
122                        failure_counter.inc();
123                    }
124                }
125            }
126            to_remove
127        };
128        if !to_remove.is_empty() {
129            let mut subscribers = subscribers.write();
130            for sub in to_remove {
131                subscribers.remove(&sub);
132            }
133        }
134    }
135
136    /// Subscribe to the data stream filtered by the filter object.
137    pub fn subscribe(&self, filter: F) -> impl Stream<Item = S> {
138        let (tx, rx) = mpsc::channel::<S>(EVENT_DISPATCH_BUFFER_SIZE);
139        self.subscribers
140            .write()
141            .insert(ObjectId::random().to_string(), (tx, filter));
142        ReceiverStream::new(rx)
143    }
144
145    pub fn try_send(&self, data: T) -> Result<(), IotaError> {
146        // Only attempt to send if there are active subscribers
147        if self.has_subscribers() {
148            self.streamer_queue.try_send(data).map_err(|e| {
149                self.metrics
150                    .dropped_submissions
151                    .with_label_values(&[self.metrics_label])
152                    .inc();
153
154                IotaError::FailedToDispatchSubscription {
155                    error: e.to_string(),
156                }
157            })
158        } else {
159            // Silently drop the data if no subscribers - this is expected behavior
160            Ok(())
161        }
162    }
163
164    /// Check if there are any active subscribers
165    pub fn has_subscribers(&self) -> bool {
166        !self.subscribers.read().is_empty()
167    }
168}
iota_core/streamer.rs

iota_core/
streamer.rs
