#10576: generic utility to unsubscribe from broadcast upon drop of the rx-side. (#10708)

* #10576: refactor `sc-utils::notification` and `sc-client-api::notifications`, so that they use common subscribe/unsubscribe routines * Add some docs. Reorganise `sc-utils::notification` * `sc-clent-api::notifications` and `sc-utils::notification` — ensure the SubscriptionGuard is dropped before the Rx-channel * `sc-utils::pubsub::SubscriptionGuard` make it a bit more ergonomic. Let the `Rx` to be put inside of the `SubscriptionGuard`, so that the latter shall guarantee the order: - first unsubscribe; - then drop the `Rx`. * Being less zealous with splitting the modules into little pieces * rework pubsub: the concrete usage should only define a good registry type * sc-client-api::notifications: make it comply with the reworked pubsub * cargo fmt * make sc-client-api tests work * Address the review notes * cargo fmt * Describe the behaviour of pubsub registry * Doc-comments for module `sc-utils::pubsub` * Fix: it used to send notifications regardless of the filter setup during subscription * `sc-client-api::StorageNotifications` the API does not have to require mut-self-reference. As a result `sc-service::Client` does not have to wrap its `storage_notifications` into a Mutex. * cargo fmt * Several changes addressing the notes by @bckhr. - Remove the `impl Default for StorageNotifications<Block>`; - no need for groupping the `remove_from` and `listen_from` into a separate `helpers` module; - remove unnecessary import `use registry::SubscribeOp`. * Add a doc-comment to the `sc-client::notifications::SubscribeOp` * As per @bkchr note on the unproven assertion: behave gracefully upon receiving a duplicate subscription-ID. * sc-utils::pubsub: log when a registry yields an ID that does point to an existing sink * `sc-utils::notifications`: payload materialized lazily * Update Cargo.lock (after adding `log` as a dependency to the `sc-utils`) * `sc-client-api::notifications`: introduce a struct (instead of a type def) for the notification message * Get rid of `sc-utils::pubsub::Channel` trait (instead just use the `sc-utils::mpsc`) * The SubsID is no more generic: the fact it is a `Copy` is known — no need to pass it by ref * sc-utils::pubsub internals do not have to be generic over the channel type * Rename Hub::dispatch into Hub::send * That method was unnecessary (`SubscriberSink::render_notification`) * cargo fmt * No need for a separate UnsubscribeGuard type * Ditch the type-def of SubsID in the sc-utils::pubsub, instead — just use the crate::id_sequence::SeqID * Return the <Registry as Dispatch>::Ret when sending an item * Make the `Hub<M, R>::lock_registry(...)` method more ergonomic * cargo doc links * cargo doc links * Use a simpler name for the type * cargo doc links * Derive `Default` rather than implement it * Derive `Default` rather than implement it * Remove an unnecessary usage of type_name * Define a more cautious order between sinks.remove->registry.unsubscribe and registry.subscribe->sinks.insert * Hub: lock_registry_for_tests->map_registry_for_tests — a safer choice for a public API * Replace Mutex over the shared Registry with a ReentrableMutex+RefCell * sc-utils::pubsub: add tests for a panicking registry * Add the missing copyright headers * Arc<Vec<_>> -> Arc<[_]>
2026-04-28 09:47:56 +00:00 · 2022-02-28 23:51:35 +02:00
parent 89b5dafc29
commit 4009148889
19 changed files with 1609 additions and 677 deletions
@@ -0,0 +1,123 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use futures::StreamExt;
+use tokio_test::block_on;
+
+use super::*;
+
+mod normal_operation;
+mod panicking_registry;
+
+const TK: &str = "a_tracing_key";
+
+type Message = u64;
+type TestHub = Hub<Message, Registry>;
+type TestReceiver = Receiver<Message, Registry>;
+
+#[derive(Default)]
+struct Registry {
+	subscribers: HashMap<SeqID, SubsKey>,
+}
+
+struct SubsKey {
+	_receiver: Option<TestReceiver>,
+	panic: SubsKeyPanic,
+}
+
+impl SubsKey {
+	fn new() -> Self {
+		Self { _receiver: None, panic: SubsKeyPanic::None }
+	}
+	fn with_receiver(self, receiver: TestReceiver) -> Self {
+		Self { _receiver: Some(receiver), ..self }
+	}
+	fn with_panic(self, panic: SubsKeyPanic) -> Self {
+		Self { panic, ..self }
+	}
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum SubsKeyPanic {
+	None,
+
+	OnSubscribePanicBefore,
+	OnSubscribePanicAfter,
+
+	OnUnsubscribePanicBefore,
+	OnUnsubscribePanicAfter,
+
+	OnDispatchPanicBefore,
+	OnDispatchPanicAfter,
+}
+
+impl<M> Hub<M, Registry> {
+	fn subs_count(&self) -> usize {
+		self.map_registry_for_tests(|r| r.subscribers.len())
+	}
+	fn sink_count(&self) -> usize {
+		self.shared.lock().borrow().sinks.len()
+	}
+}
+
+impl Subscribe<SubsKey> for Registry {
+	fn subscribe(&mut self, subs_key: SubsKey, subs_id: SeqID) {
+		let sk_panic = subs_key.panic;
+
+		if sk_panic == SubsKeyPanic::OnSubscribePanicBefore {
+			panic!("on-subscribe-panic-before")
+		}
+		self.subscribers.insert(subs_id, subs_key);
+		if sk_panic == SubsKeyPanic::OnSubscribePanicAfter {
+			panic!("on-subscribe-panic-after")
+		}
+	}
+}
+impl Unsubscribe for Registry {
+	fn unsubscribe(&mut self, subs_id: SeqID) {
+		let sk_panic =
+			self.subscribers.get(&subs_id).map(|sk| sk.panic).unwrap_or(SubsKeyPanic::None);
+
+		if sk_panic == SubsKeyPanic::OnUnsubscribePanicBefore {
+			panic!("on-unsubscribe-panic-before")
+		}
+		self.subscribers.remove(&subs_id);
+		if sk_panic == SubsKeyPanic::OnUnsubscribePanicAfter {
+			panic!("on-unsubscribe-panic-after")
+		}
+	}
+}
+impl Dispatch<Message> for Registry {
+	type Item = Message;
+	type Ret = ();
+
+	fn dispatch<F>(&mut self, message: Message, mut dispatch: F) -> Self::Ret
+	where
+		F: FnMut(&SeqID, Self::Item),
+	{
+		self.subscribers.iter().for_each(|(id, subs_key)| {
+			if subs_key.panic == SubsKeyPanic::OnDispatchPanicBefore {
+				panic!("on-dispatch-panic-before")
+			}
+			dispatch(id, message);
+			if subs_key.panic == SubsKeyPanic::OnDispatchPanicAfter {
+				panic!("on-dispatch-panic-after")
+			}
+		});
+	}
+}
@@ -0,0 +1,88 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use super::*;
+
+#[test]
+fn positive_rx_receives_relevant_messages_and_terminates_upon_hub_drop() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+		assert_eq!(hub.subs_count(), 0);
+
+		// No subscribers yet. That message is not supposed to get to anyone.
+		hub.send(0);
+
+		let mut rx_01 = hub.subscribe(SubsKey::new());
+		assert_eq!(hub.subs_count(), 1);
+
+		// That message is sent after subscription. Should be delivered into rx_01.
+		hub.send(1);
+		assert_eq!(Some(1), rx_01.next().await);
+
+		// 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());
+	});
+}
+
+#[test]
+fn positive_subs_count_is_correct_upon_drop_of_rxs() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+		assert_eq!(hub.subs_count(), 0);
+
+		let rx_01 = hub.subscribe(SubsKey::new());
+		assert_eq!(hub.subs_count(), 1);
+		let rx_02 = hub.subscribe(SubsKey::new());
+		assert_eq!(hub.subs_count(), 2);
+
+		std::mem::drop(rx_01);
+		assert_eq!(hub.subs_count(), 1);
+		std::mem::drop(rx_02);
+		assert_eq!(hub.subs_count(), 0);
+	});
+}
+
+#[test]
+fn positive_subs_count_is_correct_upon_drop_of_rxs_on_cloned_hubs() {
+	block_on(async {
+		let hub_01 = TestHub::new(TK);
+		let hub_02 = hub_01.clone();
+		assert_eq!(hub_01.subs_count(), 0);
+		assert_eq!(hub_02.subs_count(), 0);
+
+		let rx_01 = hub_02.subscribe(SubsKey::new());
+		assert_eq!(hub_01.subs_count(), 1);
+		assert_eq!(hub_02.subs_count(), 1);
+
+		let rx_02 = hub_02.subscribe(SubsKey::new());
+		assert_eq!(hub_01.subs_count(), 2);
+		assert_eq!(hub_02.subs_count(), 2);
+
+		std::mem::drop(rx_01);
+		assert_eq!(hub_01.subs_count(), 1);
+		assert_eq!(hub_02.subs_count(), 1);
+
+		std::mem::drop(rx_02);
+		assert_eq!(hub_01.subs_count(), 0);
+		assert_eq!(hub_02.subs_count(), 0);
+	});
+}
@@ -0,0 +1,248 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use super::*;
+
+use std::panic::{catch_unwind, AssertUnwindSafe};
+
+fn assert_hub_props(hub: &TestHub, sinks_count: usize, subs_count: usize) {
+	assert_eq!(hub.sink_count(), sinks_count);
+	assert_eq!(hub.subs_count(), subs_count);
+}
+
+#[test]
+fn t01() {
+	let hub = TestHub::new(TK);
+	assert_hub_props(&hub, 0, 0);
+
+	let rx_01 = hub.subscribe(SubsKey::new());
+	assert_hub_props(&hub, 1, 1);
+
+	std::mem::drop(rx_01);
+	assert_hub_props(&hub, 0, 0);
+}
+
+#[test]
+fn t02() {
+	block_on(async {
+		// Create a Hub
+		let hub = TestHub::new(TK);
+		assert_hub_props(&hub, 0, 0);
+
+		// Subscribe rx-01
+		let rx_01 = hub.subscribe(SubsKey::new());
+		assert_hub_props(&hub, 1, 1);
+
+		// Subscribe rx-02 so that its unsubscription will lead to an attempt to drop rx-01 in the
+		// middle of unsubscription of rx-02
+		let rx_02 = hub.subscribe(SubsKey::new().with_receiver(rx_01));
+		assert_hub_props(&hub, 2, 2);
+
+		// Subscribe rx-03 in order to see that it will receive messages after the unclean
+		// unsubscription
+		let mut rx_03 = hub.subscribe(SubsKey::new());
+		assert_hub_props(&hub, 3, 3);
+
+		// drop rx-02 leads to an attempt to unsubscribe rx-01
+		assert!(catch_unwind(AssertUnwindSafe(move || {
+			std::mem::drop(rx_02);
+		}))
+		.is_err());
+
+		// One of the rxes could not unsubscribe
+		assert_hub_props(&hub, 2, 2);
+
+		// Subscribe rx-04 in order to see that it will receive messages after the unclean
+		// unsubscription
+		let mut rx_04 = hub.subscribe(SubsKey::new());
+		assert_hub_props(&hub, 3, 3);
+
+		hub.send(2);
+
+		// The messages are still received
+		assert_eq!(rx_03.next().await, Some(2));
+		assert_eq!(rx_04.next().await, Some(2));
+
+		// Perform a clean unsubscription
+		std::mem::drop(rx_04);
+
+		hub.send(3);
+
+		// The messages are still received
+		assert_eq!(rx_03.next().await, Some(3));
+
+		std::mem::drop(rx_03);
+
+		hub.send(4);
+
+		// The stuck subscription is still there
+		assert_hub_props(&hub, 1, 1);
+	});
+}
+
+async fn add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(hub: &TestHub) {
+	let rx_01 = hub.subscribe(SubsKey::new());
+	let rx_02 = hub.subscribe(SubsKey::new());
+
+	hub.send(1);
+	hub.send(2);
+	hub.send(3);
+
+	assert_eq!(rx_01.take(3).collect::<Vec<_>>().await, vec![1, 2, 3]);
+
+	hub.send(4);
+	hub.send(5);
+	hub.send(6);
+
+	assert_eq!(rx_02.take(6).collect::<Vec<_>>().await, vec![1, 2, 3, 4, 5, 6]);
+}
+
+#[test]
+fn t03() {
+	block_on(async {
+		// Create a Hub
+		let hub = TestHub::new(TK);
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+
+		assert!(catch_unwind(AssertUnwindSafe(
+			|| hub.subscribe(SubsKey::new().with_panic(SubsKeyPanic::OnSubscribePanicBefore))
+		))
+		.is_err());
+
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+	});
+}
+
+#[test]
+fn t04() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+
+		assert!(catch_unwind(AssertUnwindSafe(
+			|| hub.subscribe(SubsKey::new().with_panic(SubsKeyPanic::OnSubscribePanicAfter))
+		))
+		.is_err());
+
+		// the registry has panicked after it has added a subs-id into its internal storage — the
+		// sinks do not leak, although the subscriptions storage contains some garbage
+		assert_hub_props(&hub, 0, 1);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 1);
+	})
+}
+
+#[test]
+fn t05() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+
+		let rx_01 =
+			hub.subscribe(SubsKey::new().with_panic(SubsKeyPanic::OnUnsubscribePanicBefore));
+
+		assert_hub_props(&hub, 1, 1);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 1, 1);
+
+		assert!(catch_unwind(AssertUnwindSafe(move || std::mem::drop(rx_01))).is_err());
+
+		// the registry has panicked on-unsubscribe before it removed the subs-id from its internal
+		// storage — the sinks do not leak, although the subscriptions storage contains some garbage
+		assert_hub_props(&hub, 0, 1);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 1);
+	})
+}
+
+#[test]
+fn t06() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+
+		let rx_01 = hub.subscribe(SubsKey::new().with_panic(SubsKeyPanic::OnUnsubscribePanicAfter));
+
+		assert_hub_props(&hub, 1, 1);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 1, 1);
+
+		assert!(catch_unwind(AssertUnwindSafe(move || std::mem::drop(rx_01))).is_err());
+
+		// the registry has panicked on-unsubscribe after it removed the subs-id from its internal
+		// storage — the sinks do not leak, the subscriptions storage does not contain any garbage
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+	})
+}
+
+#[test]
+fn t07() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+
+		let rx_01 = hub.subscribe(SubsKey::new().with_panic(SubsKeyPanic::OnDispatchPanicBefore));
+		assert_hub_props(&hub, 1, 1);
+		assert!(catch_unwind(AssertUnwindSafe(|| hub.send(1))).is_err());
+		assert_hub_props(&hub, 1, 1);
+
+		std::mem::drop(rx_01);
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+	})
+}
+
+#[test]
+fn t08() {
+	block_on(async {
+		let hub = TestHub::new(TK);
+
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+
+		let rx_01 = hub.subscribe(SubsKey::new().with_panic(SubsKeyPanic::OnDispatchPanicAfter));
+		assert_hub_props(&hub, 1, 1);
+		assert!(catch_unwind(AssertUnwindSafe(|| hub.send(1))).is_err());
+		assert_hub_props(&hub, 1, 1);
+
+		std::mem::drop(rx_01);
+		assert_hub_props(&hub, 0, 0);
+		add_some_subscribers_see_that_messages_are_delivered_and_unsubscribe(&hub).await;
+		assert_hub_props(&hub, 0, 0);
+	})
+}