Grcov report - pubsub.rs

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

//! Provides means to implement a typical Pub/Sub mechanism.

20

//!

21

//! This module provides a type [`Hub`] which can be used both to subscribe,

22

//! and to send the broadcast messages.

23

//!

24

//! The [`Hub`] type is parametrized by two other types:

25

//! - `Message` — the type of a message that shall be delivered to the subscribers;

26

//! - `Registry` — implementation of the subscription/dispatch logic.

27

//!

28

//! A Registry is implemented by defining the following traits:

29

//! - [`Subscribe<K>`];

30

//! - [`Dispatch<M>`];

31

//! - [`Unsubscribe`].

32

//!

33

//! As a result of subscription `Hub::subscribe` method returns an instance of

34

//! [`Receiver<Message,Registry>`]. That can be used as a [`Stream`] to receive the messages.

35

//! Upon drop the [`Receiver<Message, Registry>`] shall unregister itself from the `Hub`.

36

37

use std::{

38

	collections::HashMap,

39

	pin::Pin,

40

	sync::{Arc, Weak},

41

	task::{Context, Poll},

42

};

43

44

use futures::stream::{FusedStream, Stream};

45

// use parking_lot::Mutex;

46

use parking_lot::ReentrantMutex;

47

use std::cell::RefCell;

48

49

use crate::{

50

	id_sequence::SeqID,

51

	mpsc::{TracingUnboundedReceiver, TracingUnboundedSender},

52

};

53

54

#[cfg(test)]

55

mod tests;

56

57

/// Unsubscribe: unregisters a previously created subscription.

58

pub trait Unsubscribe {

59

	/// Remove all registrations of the subscriber with ID `subs_id`.

60

	fn unsubscribe(&mut self, subs_id: SeqID);

61

62

63

/// Subscribe using a key of type `K`

64

pub trait Subscribe<K> {

65

	/// Register subscriber with the ID `subs_id` as having interest to the key `K`.

66

	fn subscribe(&mut self, subs_key: K, subs_id: SeqID);

67

68

69

/// Dispatch a message of type `M`.

70

pub trait Dispatch<M> {

71

	/// The type of the that shall be sent through the channel as a result of such dispatch.

72

	type Item;

73

	/// The type returned by the `dispatch`-method.

74

	type Ret;

75

76

	/// Dispatch the message of type `M`.

77

///

78

	/// The implementation is given an instance of `M` and is supposed to invoke `dispatch` for

79

	/// each matching subscriber, with an argument of type `Self::Item` matching that subscriber.

80

///

81

	/// Note that this does not have to be of the same type with the item that will be sent through

82

	/// to the subscribers. The subscribers will receive a message of type `Self::Item`.

83

	fn dispatch<F>(&mut self, message: M, dispatch: F) -> Self::Ret

84

	where

85

		F: FnMut(&SeqID, Self::Item);

86

87

88

/// A subscription hub.

89

///

90

/// Does the subscription and dispatch.

91

/// The exact subscription and routing behaviour is to be implemented by the Registry (of type `R`).

92

/// The Hub under the hood uses the channel defined in `crate::mpsc` module.

93

#[derive(Debug)]

94

pub struct Hub<M, R> {

95

	tracing_key: &'static str,

96

	shared: Arc<ReentrantMutex<RefCell<Shared<M, R>>>>,

97

98

99

/// The receiving side of the subscription.

100

///

101

/// The messages are delivered as items of a [`Stream`].

102

/// Upon drop this receiver unsubscribes itself from the [`Hub<M, R>`].

103

#[derive(Debug)]

104

pub struct Receiver<M, R>

105

where

106

	R: Unsubscribe,

107

108

	rx: TracingUnboundedReceiver<M>,

109

110

	shared: Weak<ReentrantMutex<RefCell<Shared<M, R>>>>,

111

	subs_id: SeqID,

112

113

114

#[derive(Debug)]

115

struct Shared<M, R> {

116

	id_sequence: crate::id_sequence::IDSequence,

117

	registry: R,

118

	sinks: HashMap<SeqID, TracingUnboundedSender<M>>,

119

120

121

impl<M, R> Hub<M, R>

122

where

123

	R: Unsubscribe,

124

125

	/// Provide access to the registry (for test purposes).

126

	pub fn map_registry_for_tests<MapF, Ret>(&self, map: MapF) -> Ret

127

	where

128

		MapF: FnOnce(&R) -> Ret,

129

130

		let shared_locked = self.shared.lock();

131

		let shared_borrowed = shared_locked.borrow();

132

		map(&shared_borrowed.registry)

133

134

135

136

impl<M, R> Drop for Receiver<M, R>

137

where

138

	R: Unsubscribe,

139

140

	fn drop(&mut self) {

141

		if let Some(shared) = self.shared.upgrade() {

142

			shared.lock().borrow_mut().unsubscribe(self.subs_id);

143

144

145

146

147

impl<M, R> Hub<M, R> {

148

	/// Create a new instance of Hub (with default value for the Registry).

149

	pub fn new(tracing_key: &'static str) -> Self

150

	where

151

		R: Default,

152

153

		Self::new_with_registry(tracing_key, Default::default())

154

155

156

	/// Create a new instance of Hub over the initialized Registry.

157

	pub fn new_with_registry(tracing_key: &'static str, registry: R) -> Self {

158

		let shared =

159

			Shared { registry, sinks: Default::default(), id_sequence: Default::default() };

160

		let shared = Arc::new(ReentrantMutex::new(RefCell::new(shared)));

161

		Self { tracing_key, shared }

162

163

164

	/// Subscribe to this Hub using the `subs_key: K`.

165

///

166

	/// A subscription with a key `K` is possible if the Registry implements `Subscribe<K>`.

167

	pub fn subscribe<K>(&self, subs_key: K, queue_size_warning: usize) -> Receiver<M, R>

168

	where

169

		R: Subscribe<K> + Unsubscribe,

170

171

		let shared_locked = self.shared.lock();

172

		let mut shared_borrowed = shared_locked.borrow_mut();

173

174

		let subs_id = shared_borrowed.id_sequence.next_id();

175

176

		// The order (registry.subscribe then sinks.insert) is important here:

177

		// assuming that `Subscribe<K>::subscribe` can panic, it is better to at least

178

		// have the sink disposed.

179

		shared_borrowed.registry.subscribe(subs_key, subs_id);

180

181

		let (tx, rx) = crate::mpsc::tracing_unbounded(self.tracing_key, queue_size_warning);

182

		assert!(shared_borrowed.sinks.insert(subs_id, tx).is_none(), "Used IDSequence to create another ID. Should be unique until u64 is overflowed. Should be unique.");

183

184

		Receiver { shared: Arc::downgrade(&self.shared), subs_id, rx }

185

186

187

	/// Send the message produced with `Trigger`.

188

///

189

	/// This is possible if the registry implements `Dispatch<Trigger, Item = M>`.

190

	pub fn send<Trigger>(&self, trigger: Trigger) -> <R as Dispatch<Trigger>>::Ret

191

	where

192

		R: Dispatch<Trigger, Item = M>,

193

194

		let shared_locked = self.shared.lock();

195

		let mut shared_borrowed = shared_locked.borrow_mut();

196

		let (registry, sinks) = shared_borrowed.get_mut();

197

198

		registry.dispatch(trigger, |subs_id, item| {

199

			if let Some(tx) = sinks.get_mut(subs_id) {

200

				if let Err(send_err) = tx.unbounded_send(item) {

201

					log::warn!("Sink with SubsID = {} failed to perform unbounded_send: {} ({} as Dispatch<{}, Item = {}>::dispatch(...))", subs_id, send_err, std::any::type_name::<R>(),

202

					std::any::type_name::<Trigger>(),

203

					std::any::type_name::<M>());

204

205

			} else {

206

				log::warn!(

207

					"No Sink for SubsID = {} ({} as Dispatch<{}, Item = {}>::dispatch(...))",

208

					subs_id,

209

					std::any::type_name::<R>(),

210

					std::any::type_name::<Trigger>(),

211

					std::any::type_name::<M>(),

212

);

213

214

})

215

216

217

218

impl<M, R> Shared<M, R> {

219

	fn get_mut(&mut self) -> (&mut R, &mut HashMap<SeqID, TracingUnboundedSender<M>>) {

220

		(&mut self.registry, &mut self.sinks)

221

222

223

	fn unsubscribe(&mut self, subs_id: SeqID)

224

	where

225

		R: Unsubscribe,

226

227

		// The order (sinks.remove then registry.unsubscribe) is important here:

228

		// assuming that `Unsubscribe::unsubscribe` can panic, it is better to at least

229

		// have the sink disposed.

230

		self.sinks.remove(&subs_id);

231

		self.registry.unsubscribe(subs_id);

232

233

234

235

impl<M, R> Clone for Hub<M, R> {

236

	fn clone(&self) -> Self {

237

		Self { tracing_key: self.tracing_key, shared: self.shared.clone() }

238

239

240

241

impl<M, R> Unpin for Receiver<M, R> where R: Unsubscribe {}

242

243

impl<M, R> Stream for Receiver<M, R>

244

where

245

	R: Unsubscribe,

246

247

	type Item = M;

248

249

	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {

250

		Pin::new(&mut self.get_mut().rx).poll_next(cx)

251

252

253

254

impl<Ch, R> FusedStream for Receiver<Ch, R>

255

where

256

	R: Unsubscribe,

257

258

	fn is_terminated(&self) -> bool {

259

		self.rx.is_terminated()

260

261