diff --git a/Cargo.lock b/Cargo.lock index 72b5943d48b3f..e0cc2c1074f80 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9231,6 +9231,7 @@ name = "sc-service-test" version = "2.0.0" dependencies = [ "array-bytes", + "async-channel", "fdlimit", "futures", "log", @@ -9435,13 +9436,14 @@ dependencies = [ name = "sc-utils" version = "4.0.0-dev" dependencies = [ - "backtrace", + "async-channel", "futures", "futures-timer", "lazy_static", "log", "parking_lot 0.12.1", "prometheus", + "sp-arithmetic", "tokio-test", ] diff --git a/client/consensus/common/src/import_queue/basic_queue.rs b/client/consensus/common/src/import_queue/basic_queue.rs index 06122491eedce..6404708151f00 100644 --- a/client/consensus/common/src/import_queue/basic_queue.rs +++ b/client/consensus/common/src/import_queue/basic_queue.rs @@ -118,8 +118,8 @@ impl BasicQueueHandle { } pub fn close(&mut self) { - self.justification_sender.close_channel(); - self.block_import_sender.close_channel(); + self.justification_sender.close(); + self.block_import_sender.close(); } } @@ -597,11 +597,11 @@ mod tests { fn prioritizes_finality_work_over_block_import() { let (result_sender, mut result_port) = buffered_link::buffered_link(100_000); - let (worker, mut finality_sender, mut block_import_sender) = + let (worker, finality_sender, block_import_sender) = BlockImportWorker::new(result_sender, (), Box::new(()), Some(Box::new(())), None); futures::pin_mut!(worker); - let mut import_block = |n| { + let import_block = |n| { let header = Header { parent_hash: Hash::random(), number: n, @@ -612,35 +612,37 @@ mod tests { let hash = header.hash(); - block_on(block_import_sender.send(worker_messages::ImportBlocks( - BlockOrigin::Own, - vec![IncomingBlock { - hash, - header: Some(header), - body: None, - indexed_body: None, - justifications: None, - origin: None, - allow_missing_state: false, - import_existing: false, - state: None, - skip_execution: false, - }], - ))) - .unwrap(); + block_import_sender + .unbounded_send(worker_messages::ImportBlocks( + BlockOrigin::Own, + vec![IncomingBlock { + hash, + header: Some(header), + body: None, + indexed_body: None, + justifications: None, + origin: None, + allow_missing_state: false, + import_existing: false, + state: None, + skip_execution: false, + }], + )) + .unwrap(); hash }; - let mut import_justification = || { + let import_justification = || { let hash = Hash::random(); - block_on(finality_sender.send(worker_messages::ImportJustification( - libp2p::PeerId::random(), - hash, - 1, - (*b"TEST", Vec::new()), - ))) - .unwrap(); + finality_sender + .unbounded_send(worker_messages::ImportJustification( + libp2p::PeerId::random(), + hash, + 1, + (*b"TEST", Vec::new()), + )) + .unwrap(); hash }; diff --git a/client/consensus/common/src/import_queue/buffered_link.rs b/client/consensus/common/src/import_queue/buffered_link.rs index 5c95cbc4260e1..c23a4b0d5d0ab 100644 --- a/client/consensus/common/src/import_queue/buffered_link.rs +++ b/client/consensus/common/src/import_queue/buffered_link.rs @@ -53,7 +53,7 @@ use super::BlockImportResult; /// can be used to buffer commands, and the receiver can be used to poll said commands and transfer /// them to another link. `queue_size_warning` sets the warning threshold of the channel queue size. pub fn buffered_link( - queue_size_warning: i64, + queue_size_warning: usize, ) -> (BufferedLinkSender, BufferedLinkReceiver) { let (tx, rx) = tracing_unbounded("mpsc_buffered_link", queue_size_warning); let tx = BufferedLinkSender { tx }; @@ -166,7 +166,7 @@ impl BufferedLinkReceiver { } /// Close the channel. - pub fn close(&mut self) { + pub fn close(&mut self) -> bool { self.rx.get_mut().close() } } diff --git a/client/service/src/task_manager/mod.rs b/client/service/src/task_manager/mod.rs index 8dc4748b089ee..afccee9033e36 100644 --- a/client/service/src/task_manager/mod.rs +++ b/client/service/src/task_manager/mod.rs @@ -281,7 +281,7 @@ impl SpawnEssentialTaskHandle { let essential_failed = self.essential_failed_tx.clone(); let essential_task = std::panic::AssertUnwindSafe(task).catch_unwind().map(move |_| { log::error!("Essential task `{}` failed. Shutting down service.", name); - let _ = essential_failed.close_channel(); + let _ = essential_failed.close(); }); let _ = self.inner.spawn_inner(name, group, essential_task, task_type); diff --git a/client/service/test/Cargo.toml b/client/service/test/Cargo.toml index f58b8c295ed76..94a844aa7dc0d 100644 --- a/client/service/test/Cargo.toml +++ b/client/service/test/Cargo.toml @@ -12,6 +12,7 @@ repository = "https://github.com/paritytech/substrate/" targets = ["x86_64-unknown-linux-gnu"] [dependencies] +async-channel = "1.8.0" array-bytes = "4.1" fdlimit = "0.2.1" futures = "0.3.21" diff --git a/client/service/test/src/client/mod.rs b/client/service/test/src/client/mod.rs index f8988c5a454c3..c7b19ca8ba27f 100644 --- a/client/service/test/src/client/mod.rs +++ b/client/service/test/src/client/mod.rs @@ -16,6 +16,7 @@ // You should have received a copy of the GNU General Public License // along with this program. If not, see . +use async_channel::TryRecvError; use futures::executor::block_on; use parity_scale_codec::{Decode, Encode, Joiner}; use sc_block_builder::BlockBuilderProvider; @@ -175,16 +176,17 @@ fn finality_notification_check( finalized: &[Hash], stale_heads: &[Hash], ) { - match notifications.try_next() { - Ok(Some(notif)) => { + match notifications.try_recv() { + Ok(notif) => { let stale_heads_expected: HashSet<_> = stale_heads.iter().collect(); let stale_heads: HashSet<_> = notif.stale_heads.iter().collect(); assert_eq!(notif.tree_route.as_ref(), &finalized[..finalized.len() - 1]); assert_eq!(notif.hash, *finalized.last().unwrap()); assert_eq!(stale_heads, stale_heads_expected); }, - Ok(None) => panic!("unexpected notification result, client send channel was closed"), - Err(_) => assert!(finalized.is_empty()), + Err(TryRecvError::Closed) => + panic!("unexpected notification result, client send channel was closed"), + Err(TryRecvError::Empty) => assert!(finalized.is_empty()), } } @@ -983,7 +985,7 @@ fn import_with_justification() { finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[]); finality_notification_check(&mut finality_notifications, &[a3.hash()], &[]); - assert!(finality_notifications.try_next().is_err()); + assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty)); } #[test] @@ -1038,7 +1040,7 @@ fn importing_diverged_finalized_block_should_trigger_reorg() { assert_eq!(client.chain_info().finalized_hash, b1.hash()); finality_notification_check(&mut finality_notifications, &[b1.hash()], &[a2.hash()]); - assert!(finality_notifications.try_next().is_err()); + assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty)); } #[test] @@ -1124,7 +1126,7 @@ fn finalizing_diverged_block_should_trigger_reorg() { finality_notification_check(&mut finality_notifications, &[b1.hash()], &[]); finality_notification_check(&mut finality_notifications, &[b2.hash(), b3.hash()], &[a2.hash()]); - assert!(finality_notifications.try_next().is_err()); + assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty)); } #[test] @@ -1227,7 +1229,7 @@ fn finality_notifications_content() { finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[c1.hash()]); finality_notification_check(&mut finality_notifications, &[d3.hash(), d4.hash()], &[b2.hash()]); - assert!(finality_notifications.try_next().is_err()); + assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty)); } #[test] @@ -1437,7 +1439,7 @@ fn doesnt_import_blocks_that_revert_finality() { finality_notification_check(&mut finality_notifications, &[a3.hash()], &[b2.hash()]); - assert!(finality_notifications.try_next().is_err()); + assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty)); } #[test] diff --git a/client/utils/Cargo.toml b/client/utils/Cargo.toml index e80588453597e..38484285d3065 100644 --- a/client/utils/Cargo.toml +++ b/client/utils/Cargo.toml @@ -10,13 +10,14 @@ description = "I/O for Substrate runtimes" readme = "README.md" [dependencies] -backtrace = "0.3.67" +async-channel = "1.8.0" futures = "0.3.21" futures-timer = "3.0.2" lazy_static = "1.4.0" log = "0.4" parking_lot = "0.12.1" prometheus = { version = "0.13.0", default-features = false } +sp-arithmetic = { version = "6.0.0", default-features = false, path = "../../primitives/arithmetic" } [features] default = ["metered"] diff --git a/client/utils/README.md b/client/utils/README.md index 2da70f09ccbc5..d20fe69efc5ac 100644 --- a/client/utils/README.md +++ b/client/utils/README.md @@ -1,16 +1,11 @@ -Utilities Primitives for Substrate +# Utilities Primitives for Substrate -## Features +This crate provides `mpsc::tracing_unbounded` function that returns wrapper types to +`async_channel::Sender` and `async_channel::Receiver`, which register every +`send`/`received`/`dropped` action happened on the channel. -### metered - -This feature changes the behaviour of the function `mpsc::tracing_unbounded`. -With the disabled feature this function is an alias to `futures::channel::mpsc::unbounded`. -However, when the feature is enabled it creates wrapper types to `UnboundedSender` -and `UnboundedReceiver` to register every `send`/`received`/`dropped` action happened on -the channel. - -Also this feature creates and registers a prometheus vector with name `unbounded_channel_len` and labels: +Also this wrapper creates and registers a prometheus vector with name `unbounded_channel_len` +and labels: | Label | Description | | ------------ | --------------------------------------------- | diff --git a/client/utils/src/lib.rs b/client/utils/src/lib.rs index f0cc1efe6111c..017fc76207d27 100644 --- a/client/utils/src/lib.rs +++ b/client/utils/src/lib.rs @@ -18,17 +18,11 @@ //! Utilities Primitives for Substrate //! -//! # Features +//! This crate provides `mpsc::tracing_unbounded` function that returns wrapper types to +//! `async_channel::Sender` and `async_channel::Receiver`, which register every +//! `send`/`received`/`dropped` action happened on the channel. //! -//! ## metered -//! -//! This feature changes the behaviour of the function `mpsc::tracing_unbounded`. -//! With the disabled feature this function is an alias to `futures::channel::mpsc::unbounded`. -//! However, when the feature is enabled it creates wrapper types to `UnboundedSender` -//! and `UnboundedReceiver` to register every `send`/`received`/`dropped` action happened on -//! the channel. -//! -//! Also this feature creates and registers a prometheus vector with name `unbounded_channel_len` +//! Also this wrapper creates and registers a prometheus vector with name `unbounded_channel_len` //! and labels: //! //! | Label | Description | diff --git a/client/utils/src/metrics.rs b/client/utils/src/metrics.rs index c2b10100f229f..6bbdbe2e2e599 100644 --- a/client/utils/src/metrics.rs +++ b/client/utils/src/metrics.rs @@ -24,7 +24,6 @@ use prometheus::{ Error as PrometheusError, Registry, }; -#[cfg(feature = "metered")] use prometheus::{core::GenericCounterVec, Opts}; lazy_static! { @@ -36,7 +35,6 @@ lazy_static! { .expect("Creating of statics doesn't fail. qed"); } -#[cfg(feature = "metered")] lazy_static! { pub static ref UNBOUNDED_CHANNELS_COUNTER : GenericCounterVec = GenericCounterVec::new( Opts::new("substrate_unbounded_channel_len", "Items in each mpsc::unbounded instance"), @@ -49,8 +47,6 @@ lazy_static! { pub fn register_globals(registry: &Registry) -> Result<(), PrometheusError> { registry.register(Box::new(TOKIO_THREADS_ALIVE.clone()))?; registry.register(Box::new(TOKIO_THREADS_TOTAL.clone()))?; - - #[cfg(feature = "metered")] registry.register(Box::new(UNBOUNDED_CHANNELS_COUNTER.clone()))?; Ok(()) diff --git a/client/utils/src/mpsc.rs b/client/utils/src/mpsc.rs index df45e33ff8ee5..3f783b10060bd 100644 --- a/client/utils/src/mpsc.rs +++ b/client/utils/src/mpsc.rs @@ -16,287 +16,164 @@ // You should have received a copy of the GNU General Public License // along with this program. If not, see . -//! Features to meter unbounded channels - -#[cfg(not(feature = "metered"))] -mod inner { - // just aliased, non performance implications - use futures::channel::mpsc::{self, UnboundedReceiver, UnboundedSender}; - pub type TracingUnboundedSender = UnboundedSender; - pub type TracingUnboundedReceiver = UnboundedReceiver; - - /// Alias `mpsc::unbounded` - pub fn tracing_unbounded( - _key: &'static str, - ) -> (TracingUnboundedSender, TracingUnboundedReceiver) { - mpsc::unbounded() - } +//! Code to meter unbounded channels. + +use crate::metrics::UNBOUNDED_CHANNELS_COUNTER; +use async_channel::{Receiver, Sender, TryRecvError, TrySendError}; +use futures::{ + stream::{FusedStream, Stream}, + task::{Context, Poll}, +}; +use log::error; +use sp_arithmetic::traits::SaturatedConversion; +use std::{ + backtrace::Backtrace, + pin::Pin, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, +}; + +/// Wrapper Type around [`async_channel::Sender`] that increases the global +/// measure when a message is added. +#[derive(Debug)] +pub struct TracingUnboundedSender { + inner: Sender, + name: &'static str, + queue_size_warning: usize, + warning_fired: Arc, + creation_backtrace: Arc, } -#[cfg(feature = "metered")] -mod inner { - // tracing implementation - use crate::metrics::UNBOUNDED_CHANNELS_COUNTER; - use backtrace::Backtrace; - use futures::{ - channel::mpsc::{ - self, SendError, TryRecvError, TrySendError, UnboundedReceiver, UnboundedSender, - }, - sink::Sink, - stream::{FusedStream, Stream}, - task::{Context, Poll}, - }; - use log::error; - use std::{ - pin::Pin, - sync::{ - atomic::{AtomicBool, AtomicI64, Ordering}, - Arc, - }, - }; - - /// Wrapper Type around `UnboundedSender` that increases the global - /// measure when a message is added - #[derive(Debug)] - pub struct TracingUnboundedSender { - inner: UnboundedSender, - name: &'static str, - // To not bother with ordering and possible underflow errors of the unsigned counter - // we just use `i64` and `Ordering::Relaxed`, and perceive `queue_size` as approximate. - // It can turn < 0 though. - queue_size: Arc, - queue_size_warning: i64, - warning_fired: Arc, - creation_backtrace: Arc, - } - - // Strangely, deriving `Clone` requires that `T` is also `Clone`. - impl Clone for TracingUnboundedSender { - fn clone(&self) -> Self { - Self { - inner: self.inner.clone(), - name: self.name, - queue_size: self.queue_size.clone(), - queue_size_warning: self.queue_size_warning, - warning_fired: self.warning_fired.clone(), - creation_backtrace: self.creation_backtrace.clone(), - } +// Strangely, deriving `Clone` requires that `T` is also `Clone`. +impl Clone for TracingUnboundedSender { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + name: self.name, + queue_size_warning: self.queue_size_warning, + warning_fired: self.warning_fired.clone(), + creation_backtrace: self.creation_backtrace.clone(), } } +} - /// Wrapper Type around `UnboundedReceiver` that decreases the global - /// measure when a message is polled - #[derive(Debug)] - pub struct TracingUnboundedReceiver { - inner: UnboundedReceiver, - name: &'static str, - queue_size: Arc, - } - - /// Wrapper around `mpsc::unbounded` that tracks the in- and outflow via - /// `UNBOUNDED_CHANNELS_COUNTER` and warns if the message queue grows - /// above the warning threshold. - pub fn tracing_unbounded( - name: &'static str, - queue_size_warning: i64, - ) -> (TracingUnboundedSender, TracingUnboundedReceiver) { - let (s, r) = mpsc::unbounded(); - let queue_size = Arc::new(AtomicI64::new(0)); - let sender = TracingUnboundedSender { - inner: s, - name, - queue_size: queue_size.clone(), - queue_size_warning, - warning_fired: Arc::new(AtomicBool::new(false)), - creation_backtrace: Arc::new(Backtrace::new_unresolved()), - }; - let receiver = TracingUnboundedReceiver { inner: r, name, queue_size }; - (sender, receiver) - } - - impl TracingUnboundedSender { - /// Proxy function to mpsc::UnboundedSender - pub fn poll_ready(&self, ctx: &mut Context) -> Poll> { - self.inner.poll_ready(ctx) - } - - /// Proxy function to mpsc::UnboundedSender - pub fn is_closed(&self) -> bool { - self.inner.is_closed() - } - - /// Proxy function to mpsc::UnboundedSender - pub fn close_channel(&self) { - self.inner.close_channel() - } - - /// Proxy function to mpsc::UnboundedSender - pub fn disconnect(&mut self) { - self.inner.disconnect() - } - - pub fn start_send(&mut self, msg: T) -> Result<(), SendError> { - // The underlying implementation of [`UnboundedSender::start_send`] is the same as - // [`UnboundedSender::unbounded_send`], so we just reuse the message counting and - // error reporting code from `unbounded_send`. - self.unbounded_send(msg).map_err(TrySendError::into_send_error) - } - - /// Proxy function to mpsc::UnboundedSender - pub fn unbounded_send(&self, msg: T) -> Result<(), TrySendError> { - self.inner.unbounded_send(msg).map(|s| { - UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.name, "send"]).inc(); - - let queue_size = self.queue_size.fetch_add(1, Ordering::Relaxed); - if queue_size == self.queue_size_warning && - self.warning_fired - .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - // `warning_fired` and `queue_size` are not synchronized, so it's possible - // that the warning is fired few times before the `warning_fired` is seen - // by all threads. This seems better than introducing a mutex guarding them. - let mut backtrace = (*self.creation_backtrace).clone(); - backtrace.resolve(); - error!( - "The number of unprocessed messages in channel `{}` reached {}.\n\ - The channel was created at:\n{:?}", - self.name, self.queue_size_warning, backtrace, - ); - } +/// Wrapper Type around [`async_channel::Receiver`] that decreases the global +/// measure when a message is polled. +#[derive(Debug)] +pub struct TracingUnboundedReceiver { + inner: Receiver, + name: &'static str, +} - s - }) - } +/// Wrapper around [`async_channel::unbounded`] that tracks the in- and outflow via +/// `UNBOUNDED_CHANNELS_COUNTER` and warns if the message queue grows +/// above the warning threshold. +pub fn tracing_unbounded( + name: &'static str, + queue_size_warning: usize, +) -> (TracingUnboundedSender, TracingUnboundedReceiver) { + let (s, r) = async_channel::unbounded(); + let sender = TracingUnboundedSender { + inner: s, + name, + queue_size_warning, + warning_fired: Arc::new(AtomicBool::new(false)), + creation_backtrace: Arc::new(Backtrace::force_capture()), + }; + let receiver = TracingUnboundedReceiver { inner: r, name }; + (sender, receiver) +} - /// Proxy function to mpsc::UnboundedSender - pub fn same_receiver(&self, other: &UnboundedSender) -> bool { - self.inner.same_receiver(other) - } +impl TracingUnboundedSender { + /// Proxy function to [`async_channel::Sender`]. + pub fn is_closed(&self) -> bool { + self.inner.is_closed() } - impl TracingUnboundedReceiver { - fn consume(&mut self) { - // consume all items, make sure to reflect the updated count - let mut count = 0; - loop { - if self.inner.is_terminated() { - break - } + /// Proxy function to [`async_channel::Sender`]. + pub fn close(&self) -> bool { + self.inner.close() + } - match self.try_next() { - Ok(Some(..)) => count += 1, - _ => break, - } + /// Proxy function to `async_channel::Sender::try_send`. + pub fn unbounded_send(&self, msg: T) -> Result<(), TrySendError> { + self.inner.try_send(msg).map(|s| { + UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.name, "send"]).inc(); + + if self.inner.len() >= self.queue_size_warning && + self.warning_fired + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + error!( + "The number of unprocessed messages in channel `{}` exceeded {}.\n\ + The channel was created at:\n{}\n + Last message was sent from:\n{}", + self.name, + self.queue_size_warning, + self.creation_backtrace, + Backtrace::force_capture(), + ); } - // and discount the messages - if count > 0 { - UNBOUNDED_CHANNELS_COUNTER - .with_label_values(&[self.name, "dropped"]) - .inc_by(count); - } - } - /// Proxy function to mpsc::UnboundedReceiver - /// that consumes all messages first and updates the counter - pub fn close(&mut self) { - self.consume(); - self.inner.close() - } - - /// Proxy function to mpsc::UnboundedReceiver - /// that discounts the messages taken out - pub fn try_next(&mut self) -> Result, TryRecvError> { - self.inner.try_next().map(|s| { - if s.is_some() { - let _ = self.queue_size.fetch_sub(1, Ordering::Relaxed); - UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.name, "received"]).inc(); - } - s - }) - } + s + }) } +} - impl Drop for TracingUnboundedReceiver { - fn drop(&mut self) { - self.consume(); - } +impl TracingUnboundedReceiver { + /// Proxy function to [`async_channel::Receiver`]. + pub fn close(&mut self) -> bool { + self.inner.close() } - impl Unpin for TracingUnboundedReceiver {} - - impl Stream for TracingUnboundedReceiver { - type Item = T; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let s = self.get_mut(); - match Pin::new(&mut s.inner).poll_next(cx) { - Poll::Ready(msg) => { - if msg.is_some() { - let _ = s.queue_size.fetch_sub(1, Ordering::Relaxed); - UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[s.name, "received"]).inc(); - } - Poll::Ready(msg) - }, - Poll::Pending => Poll::Pending, - } - } + /// Proxy function to [`async_channel::Receiver`] + /// that discounts the messages taken out. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv().map(|s| { + UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.name, "received"]).inc(); + s + }) } +} - impl FusedStream for TracingUnboundedReceiver { - fn is_terminated(&self) -> bool { - self.inner.is_terminated() +impl Drop for TracingUnboundedReceiver { + fn drop(&mut self) { + self.close(); + // the number of messages about to be dropped + let count = self.inner.len(); + // discount the messages + if count > 0 { + UNBOUNDED_CHANNELS_COUNTER + .with_label_values(&[self.name, "dropped"]) + .inc_by(count.saturated_into()); } } +} - impl Sink for TracingUnboundedSender { - type Error = SendError; - - fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - TracingUnboundedSender::poll_ready(&*self, cx) - } - - fn start_send(mut self: Pin<&mut Self>, msg: T) -> Result<(), Self::Error> { - TracingUnboundedSender::start_send(&mut *self, msg) - } +impl Unpin for TracingUnboundedReceiver {} - fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } +impl Stream for TracingUnboundedReceiver { + type Item = T; - fn poll_close( - mut self: Pin<&mut Self>, - _: &mut Context<'_>, - ) -> Poll> { - self.disconnect(); - Poll::Ready(Ok(())) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let s = self.get_mut(); + match Pin::new(&mut s.inner).poll_next(cx) { + Poll::Ready(msg) => { + if msg.is_some() { + UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[s.name, "received"]).inc(); + } + Poll::Ready(msg) + }, + Poll::Pending => Poll::Pending, } } +} - impl Sink for &TracingUnboundedSender { - type Error = SendError; - - fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - TracingUnboundedSender::poll_ready(*self, cx) - } - - fn start_send(self: Pin<&mut Self>, msg: T) -> Result<(), Self::Error> { - self.unbounded_send(msg).map_err(TrySendError::into_send_error) - } - - fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { - // The difference with `TracingUnboundedSender` is intentional. The underlying - // implementation differs for `UnboundedSender` and `&UnboundedSender`: - // the latter closes the channel completely with `close_channel()`, while the former - // only closes this specific sender with `disconnect()`. - self.close_channel(); - Poll::Ready(Ok(())) - } +impl FusedStream for TracingUnboundedReceiver { + fn is_terminated(&self) -> bool { + self.inner.is_terminated() } } - -pub use inner::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender}; diff --git a/client/utils/src/notification.rs b/client/utils/src/notification.rs index 7d866ffc8c019..dabb85d613cc9 100644 --- a/client/utils/src/notification.rs +++ b/client/utils/src/notification.rs @@ -79,7 +79,7 @@ impl NotificationStream { } /// Subscribe to a channel through which the generic payload can be received. - pub fn subscribe(&self, queue_size_warning: i64) -> NotificationReceiver { + pub fn subscribe(&self, queue_size_warning: usize) -> NotificationReceiver { let receiver = self.hub.subscribe((), queue_size_warning); NotificationReceiver { receiver } } diff --git a/client/utils/src/pubsub.rs b/client/utils/src/pubsub.rs index 8136aa0692c26..5293fa42ed94c 100644 --- a/client/utils/src/pubsub.rs +++ b/client/utils/src/pubsub.rs @@ -164,7 +164,7 @@ impl Hub { /// Subscribe to this Hub using the `subs_key: K`. /// /// A subscription with a key `K` is possible if the Registry implements `Subscribe`. - pub fn subscribe(&self, subs_key: K, queue_size_warning: i64) -> Receiver + pub fn subscribe(&self, subs_key: K, queue_size_warning: usize) -> Receiver where R: Subscribe + Unsubscribe, { diff --git a/client/utils/src/pubsub/tests/normal_operation.rs b/client/utils/src/pubsub/tests/normal_operation.rs index d4b614d7a8889..a3ea4f7ddee69 100644 --- a/client/utils/src/pubsub/tests/normal_operation.rs +++ b/client/utils/src/pubsub/tests/normal_operation.rs @@ -37,9 +37,8 @@ fn positive_rx_receives_relevant_messages_and_terminates_upon_hub_drop() { // Hub is disposed. The rx_01 should be over after that. std::mem::drop(hub); - assert!(!rx_01.is_terminated()); - assert_eq!(None, rx_01.next().await); assert!(rx_01.is_terminated()); + assert_eq!(None, rx_01.next().await); }); }