Fix the metered unbounded sender/recievers (#6246)

* Fix the metered unbounded sender/recievers

* Use a counter instead

* Update client/rpc/src/system/tests.rs

* Add an is_terminated check

* Add FusedStream impl

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/primitives/utils/src/metrics.rs

@@ -23,8 +23,8 @@ use prometheus::{
 	core::{ AtomicU64, GenericGauge, GenericCounter },
 };
 
-#[cfg(features = "metered")]
-use prometheus::{core::GenericGaugeVec, Opts};
+#[cfg(feature = "metered")]
+use prometheus::{core::GenericCounterVec, Opts};
 
 
 lazy_static! {
@@ -37,9 +37,9 @@ lazy_static! {
 	).expect("Creating of statics doesn't fail. qed");
 }
 
-#[cfg(features = "metered")]
+#[cfg(feature = "metered")]
 lazy_static! {
-	pub static ref UNBOUNDED_CHANNELS_COUNTER : GenericGaugeVec<AtomicU64> = GenericGaugeVec::new(
+	pub static ref UNBOUNDED_CHANNELS_COUNTER : GenericCounterVec<AtomicU64> = GenericCounterVec::new(
 		Opts::new("unbounded_channel_len", "Items in each mpsc::unbounded instance"),
 		&["entity", "action"] // 'name of channel, send|received|dropped
 	).expect("Creating of statics doesn't fail. qed");
@@ -52,7 +52,7 @@ 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(features = "metered")]
+	#[cfg(feature = "metered")]
 	registry.register(Box::new(UNBOUNDED_CHANNELS_COUNTER.clone()))?;
 
 	Ok(())

substrate/primitives/utils/src/mpsc.rs

@@ -17,7 +17,7 @@
 
 //! Features to meter unbounded channels
 
-#[cfg(not(features = "metered"))]
+#[cfg(not(feature = "metered"))]
 mod inner {
 	// just aliased, non performance implications
 	use futures::channel::mpsc::{self, UnboundedReceiver, UnboundedSender};
@@ -31,22 +31,29 @@ mod inner {
 }
 
 
-#[cfg(features = "metered")]
+#[cfg(feature = "metered")]
 mod inner {
 	//tracing implementation
 	use futures::channel::mpsc::{self,
 		UnboundedReceiver, UnboundedSender,
 		TryRecvError, TrySendError, SendError
 	};
-	use futures::{sink::Sink, task::{Poll, Context}, stream::Stream};
+	use futures::{sink::Sink, task::{Poll, Context}, stream::{Stream, FusedStream}};
 	use std::pin::Pin;
 	use crate::metrics::UNBOUNDED_CHANNELS_COUNTER;
 
 	/// Wrapper Type around `UnboundedSender` that increases the global
 	/// measure when a message is added
-	#[derive(Debug, Clone)]
+	#[derive(Debug)]
 	pub struct TracingUnboundedSender<T>(&'static str, UnboundedSender<T>);
 
+	// Strangely, deriving `Clone` requires that `T` is also `Clone`.
+	impl<T> Clone for TracingUnboundedSender<T> {
+		fn clone(&self) -> Self {
+			Self(self.0, self.1.clone())
+		}
+	}
+
 	/// Wrapper Type around `UnboundedReceiver` that decreases the global
 	/// measure when a message is polled
 	#[derive(Debug)]
@@ -88,7 +95,7 @@ mod inner {
 		/// Proxy function to mpsc::UnboundedSender
 		pub fn unbounded_send(&self, msg: T) -> Result<(), TrySendError<T>> {
 			self.1.unbounded_send(msg).map(|s|{
-				UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, &"send"]).incr();
+				UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, &"send"]).inc();
 				s
 			})
 		}
@@ -104,13 +111,19 @@ mod inner {
 		fn consume(&mut self) {
 			// consume all items, make sure to reflect the updated count
 			let mut count = 0;
-			while let Ok(Some(..)) = self.try_next() {
-				count += 1;
-			}
+			loop {
+				if self.1.is_terminated() {
+					break;
+				}
 
+				match self.try_next() {
+					Ok(Some(..)) => count += 1,
+					_ => break
+				}
+			}
 			// and discount the messages
 			if count > 0 {
-				UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, &"dropped"]).incr_by(count);
+				UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, &"dropped"]).inc_by(count);
 			}
 
 		}
@@ -127,7 +140,7 @@ mod inner {
 		pub fn try_next(&mut self) -> Result<Option<T>, TryRecvError> {
 			self.1.try_next().map(|s| {
 				if s.is_some() {
-					UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, &"received"]).incr();
+					UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, &"received"]).inc();
 				}
 				s
 			})
@@ -153,7 +166,7 @@ mod inner {
 			match Pin::new(&mut s.1).poll_next(cx) {
 				Poll::Ready(msg) => {
 					if msg.is_some() {
-						UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.0, "received"]).incr();
+						UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[s.0, "received"]).inc();
 				   	}
 					Poll::Ready(msg)
 				}
@@ -164,6 +177,11 @@ mod inner {
 		}
 	}
 
+	impl<T> FusedStream for TracingUnboundedReceiver<T> {
+		fn is_terminated(&self) -> bool {
+			self.1.is_terminated()
+		}
+	}
 
 	impl<T> Sink<T> for TracingUnboundedSender<T> {
 		type Error = SendError;