1
// This file is part of Substrate.
2

            
3
// Copyright (C) Parity Technologies (UK) Ltd.
4
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5

            
6
// This program is free software: you can redistribute it and/or modify
7
// it under the terms of the GNU General Public License as published by
8
// the Free Software Foundation, either version 3 of the License, or
9
// (at your option) any later version.
10

            
11
// This program is distributed in the hope that it will be useful,
12
// but WITHOUT ANY WARRANTY; without even the implied warranty of
13
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
// GNU General Public License for more details.
15

            
16
// You should have received a copy of the GNU General Public License
17
// along with this program. If not, see <https://www.gnu.org/licenses/>.
18

            
19
//! Code to meter unbounded channels.
20

            
21
pub use async_channel::{TryRecvError, TrySendError};
22

            
23
use crate::metrics::{
24
	DROPPED_LABEL, RECEIVED_LABEL, SENT_LABEL, UNBOUNDED_CHANNELS_COUNTER, UNBOUNDED_CHANNELS_SIZE,
25
};
26
use async_channel::{Receiver, Sender};
27
use futures::{
28
	stream::{FusedStream, Stream},
29
	task::{Context, Poll},
30
};
31
use log::error;
32
use sp_arithmetic::traits::SaturatedConversion;
33
use std::{
34
	backtrace::Backtrace,
35
	pin::Pin,
36
	sync::{
37
		atomic::{AtomicBool, Ordering},
38
		Arc,
39
	},
40
};
41

            
42
/// Wrapper Type around [`async_channel::Sender`] that increases the global
43
/// measure when a message is added.
44
#[derive(Debug)]
45
pub struct TracingUnboundedSender<T> {
46
	inner: Sender<T>,
47
	name: &'static str,
48
	queue_size_warning: usize,
49
	warning_fired: Arc<AtomicBool>,
50
	creation_backtrace: Arc<Backtrace>,
51
}
52

            
53
// Strangely, deriving `Clone` requires that `T` is also `Clone`.
54
impl<T> Clone for TracingUnboundedSender<T> {
55
	fn clone(&self) -> Self {
56
		Self {
57
			inner: self.inner.clone(),
58
			name: self.name,
59
			queue_size_warning: self.queue_size_warning,
60
			warning_fired: self.warning_fired.clone(),
61
			creation_backtrace: self.creation_backtrace.clone(),
62
		}
63
	}
64
}
65

            
66
/// Wrapper Type around [`async_channel::Receiver`] that decreases the global
67
/// measure when a message is polled.
68
#[derive(Debug)]
69
pub struct TracingUnboundedReceiver<T> {
70
	inner: Receiver<T>,
71
	name: &'static str,
72
}
73

            
74
/// Wrapper around [`async_channel::unbounded`] that tracks the in- and outflow via
75
/// `UNBOUNDED_CHANNELS_COUNTER` and warns if the message queue grows
76
/// above the warning threshold.
77
pub fn tracing_unbounded<T>(
78
	name: &'static str,
79
	queue_size_warning: usize,
80
) -> (TracingUnboundedSender<T>, TracingUnboundedReceiver<T>) {
81
	let (s, r) = async_channel::unbounded();
82
	let sender = TracingUnboundedSender {
83
		inner: s,
84
		name,
85
		queue_size_warning,
86
		warning_fired: Arc::new(AtomicBool::new(false)),
87
		creation_backtrace: Arc::new(Backtrace::force_capture()),
88
	};
89
	let receiver = TracingUnboundedReceiver { inner: r, name: name.into() };
90
	(sender, receiver)
91
}
92

            
93
impl<T> TracingUnboundedSender<T> {
94
	/// Proxy function to [`async_channel::Sender`].
95
	pub fn is_closed(&self) -> bool {
96
		self.inner.is_closed()
97
	}
98

            
99
	/// Proxy function to [`async_channel::Sender`].
100
	pub fn close(&self) -> bool {
101
		self.inner.close()
102
	}
103

            
104
	/// Proxy function to `async_channel::Sender::try_send`.
105
	pub fn unbounded_send(&self, msg: T) -> Result<(), TrySendError<T>> {
106
		self.inner.try_send(msg).map(|s| {
107
			UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.name, SENT_LABEL]).inc();
108
			UNBOUNDED_CHANNELS_SIZE
109
				.with_label_values(&[self.name])
110
				.set(self.inner.len().saturated_into());
111

            
112
			if self.inner.len() >= self.queue_size_warning &&
113
				self.warning_fired
114
					.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
115
					.is_ok()
116
			{
117
				error!(
118
					"The number of unprocessed messages in channel `{}` exceeded {}.\n\
119
					 The channel was created at:\n{}\n
120
					 Last message was sent from:\n{}",
121
					self.name,
122
					self.queue_size_warning,
123
					self.creation_backtrace,
124
					Backtrace::force_capture(),
125
				);
126
			}
127

            
128
			s
129
		})
130
	}
131

            
132
	/// The number of elements in the channel (proxy function to [`async_channel::Sender`]).
133
	pub fn len(&self) -> usize {
134
		self.inner.len()
135
	}
136
}
137

            
138
impl<T> TracingUnboundedReceiver<T> {
139
	/// Proxy function to [`async_channel::Receiver`].
140
	pub fn close(&mut self) -> bool {
141
		self.inner.close()
142
	}
143

            
144
	/// Proxy function to [`async_channel::Receiver`]
145
	/// that discounts the messages taken out.
146
	pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
147
		self.inner.try_recv().map(|s| {
148
			UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[self.name, RECEIVED_LABEL]).inc();
149
			UNBOUNDED_CHANNELS_SIZE
150
				.with_label_values(&[self.name])
151
				.set(self.inner.len().saturated_into());
152
			s
153
		})
154
	}
155

            
156
	/// The number of elements in the channel (proxy function to [`async_channel::Receiver`]).
157
	pub fn len(&self) -> usize {
158
		self.inner.len()
159
	}
160

            
161
	/// The name of this receiver
162
	pub fn name(&self) -> &'static str {
163
		self.name
164
	}
165
}
166

            
167
impl<T> Drop for TracingUnboundedReceiver<T> {
168
	fn drop(&mut self) {
169
		// Close the channel to prevent any further messages to be sent into the channel
170
		self.close();
171
		// The number of messages about to be dropped
172
		let count = self.inner.len();
173
		// Discount the messages
174
		if count > 0 {
175
			UNBOUNDED_CHANNELS_COUNTER
176
				.with_label_values(&[self.name, DROPPED_LABEL])
177
				.inc_by(count.saturated_into());
178
		}
179
		// Reset the size metric to 0
180
		UNBOUNDED_CHANNELS_SIZE.with_label_values(&[self.name]).set(0);
181
		// Drain all the pending messages in the channel since they can never be accessed,
182
		// this can be removed once https://github.com/smol-rs/async-channel/issues/23 is
183
		// resolved
184
		while let Ok(_) = self.inner.try_recv() {}
185
	}
186
}
187

            
188
impl<T> Unpin for TracingUnboundedReceiver<T> {}
189

            
190
impl<T> Stream for TracingUnboundedReceiver<T> {
191
	type Item = T;
192

            
193
	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
194
		let s = self.get_mut();
195
		match Pin::new(&mut s.inner).poll_next(cx) {
196
			Poll::Ready(msg) => {
197
				if msg.is_some() {
198
					UNBOUNDED_CHANNELS_COUNTER.with_label_values(&[s.name, RECEIVED_LABEL]).inc();
199
					UNBOUNDED_CHANNELS_SIZE
200
						.with_label_values(&[s.name])
201
						.set(s.inner.len().saturated_into());
202
				}
203
				Poll::Ready(msg)
204
			},
205
			Poll::Pending => Poll::Pending,
206
		}
207
	}
208
}
209

            
210
impl<T> FusedStream for TracingUnboundedReceiver<T> {
211
	fn is_terminated(&self) -> bool {
212
		self.inner.is_terminated()
213
	}
214
}
215

            
216
#[cfg(test)]
217
mod tests {
218
	use super::tracing_unbounded;
219
	use async_channel::{self, RecvError, TryRecvError};
220

            
221
	#[test]
222
	fn test_tracing_unbounded_receiver_drop() {
223
		let (tracing_unbounded_sender, tracing_unbounded_receiver) =
224
			tracing_unbounded("test-receiver-drop", 10);
225
		let (tx, rx) = async_channel::unbounded::<usize>();
226

            
227
		tracing_unbounded_sender.unbounded_send(tx).unwrap();
228
		drop(tracing_unbounded_receiver);
229

            
230
		assert_eq!(rx.try_recv(), Err(TryRecvError::Closed));
231
		assert_eq!(rx.recv_blocking(), Err(RecvError));
232
	}
233
}