#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-27 04:37:57 +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,365 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2017-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 sp_core::hexdisplay::HexDisplay;
+
+use fnv::{FnvHashMap, FnvHashSet};
+use prometheus_endpoint::{register, CounterVec, Opts, U64};
+
+use sc_utils::{
+	id_sequence::SeqID as SubscriberId,
+	pubsub::{Dispatch, Subscribe, Unsubscribe},
+};
+
+type SubscribersGauge = CounterVec<U64>;
+
+/// A command to subscribe with the specified filters.
+///
+/// Used by the implementation of [`Subscribe<Op>`] trait for [`Registry].
+pub(super) struct SubscribeOp<'a> {
+	pub filter_keys: Option<&'a [StorageKey]>,
+	pub filter_child_keys: Option<&'a [(StorageKey, Option<Vec<StorageKey>>)]>,
+}
+
+#[derive(Debug, Default)]
+pub(super) struct Registry {
+	pub(super) metrics: Option<SubscribersGauge>,
+	pub(super) wildcard_listeners: FnvHashSet<SubscriberId>,
+	pub(super) listeners: HashMap<StorageKey, FnvHashSet<SubscriberId>>,
+	pub(super) child_listeners: HashMap<
+		StorageKey,
+		(HashMap<StorageKey, FnvHashSet<SubscriberId>>, FnvHashSet<SubscriberId>),
+	>,
+	pub(super) sinks: FnvHashMap<SubscriberId, SubscriberSink>,
+}
+
+#[derive(Debug)]
+pub(super) struct SubscriberSink {
+	subs_id: SubscriberId,
+	keys: Keys,
+	child_keys: ChildKeys,
+	was_triggered: bool,
+}
+
+impl Drop for SubscriberSink {
+	fn drop(&mut self) {
+		if !self.was_triggered {
+			log::trace!(
+				target: "storage_notifications",
+				"Listener was never triggered: id={}, keys={:?}, child_keys={:?}",
+				self.subs_id,
+				PrintKeys(&self.keys),
+				PrintChildKeys(&self.child_keys),
+			);
+		}
+	}
+}
+
+impl SubscriberSink {
+	fn new(subs_id: SubscriberId, keys: Keys, child_keys: ChildKeys) -> Self {
+		Self { subs_id, keys, child_keys, was_triggered: false }
+	}
+}
+
+impl Registry {
+	pub(super) fn new(prometheus_registry: Option<PrometheusRegistry>) -> Self {
+		let metrics = prometheus_registry.and_then(|r| {
+			CounterVec::new(
+				Opts::new(
+					"substrate_storage_notification_subscribers",
+					"Number of subscribers in storage notification sytem",
+				),
+				&["action"], // added | removed
+			)
+			.and_then(|g| register(g, &r))
+			.ok()
+		});
+
+		Registry { metrics, ..Default::default() }
+	}
+}
+
+impl Unsubscribe for Registry {
+	fn unsubscribe(&mut self, subs_id: SubscriberId) {
+		self.remove_subscriber(subs_id);
+	}
+}
+
+impl<'a> Subscribe<SubscribeOp<'a>> for Registry {
+	fn subscribe(&mut self, subs_op: SubscribeOp<'a>, subs_id: SubscriberId) {
+		let SubscribeOp { filter_keys, filter_child_keys } = subs_op;
+
+		let keys = Self::listen_from(
+			subs_id,
+			filter_keys.as_ref(),
+			&mut self.listeners,
+			&mut self.wildcard_listeners,
+		);
+
+		let child_keys = filter_child_keys.map(|filter_child_keys| {
+			filter_child_keys
+				.iter()
+				.map(|(c_key, o_keys)| {
+					let (c_listeners, c_wildcards) =
+						self.child_listeners.entry(c_key.clone()).or_default();
+
+					(
+						c_key.clone(),
+						Self::listen_from(
+							subs_id,
+							o_keys.as_ref(),
+							&mut *c_listeners,
+							&mut *c_wildcards,
+						),
+					)
+				})
+				.collect()
+		});
+
+		if let Some(m) = self.metrics.as_ref() {
+			m.with_label_values(&[&"added"]).inc();
+		}
+
+		if self
+			.sinks
+			.insert(subs_id, SubscriberSink::new(subs_id, keys, child_keys))
+			.is_some()
+		{
+			log::warn!("The `subscribe`-method has been passed a non-unique subs_id (in `sc-client-api::notifications`)");
+		}
+	}
+}
+
+impl<'a, Hash, CS, CCS, CCSI> Dispatch<(&'a Hash, CS, CCS)> for Registry
+where
+	Hash: Clone,
+	CS: Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>,
+	CCS: Iterator<Item = (Vec<u8>, CCSI)>,
+	CCSI: Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>,
+{
+	type Item = StorageNotification<Hash>;
+	type Ret = ();
+
+	fn dispatch<F>(&mut self, message: (&'a Hash, CS, CCS), dispatch: F) -> Self::Ret
+	where
+		F: FnMut(&SubscriberId, Self::Item),
+	{
+		let (hash, changeset, child_changeset) = message;
+		self.trigger(hash, changeset, child_changeset, dispatch);
+	}
+}
+
+impl Registry {
+	pub(super) fn trigger<Hash, F>(
+		&mut self,
+		hash: &Hash,
+		changeset: impl Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>,
+		child_changeset: impl Iterator<
+			Item = (Vec<u8>, impl Iterator<Item = (Vec<u8>, Option<Vec<u8>>)>),
+		>,
+		mut dispatch: F,
+	) where
+		Hash: Clone,
+		F: FnMut(&SubscriberId, StorageNotification<Hash>),
+	{
+		let has_wildcard = !self.wildcard_listeners.is_empty();
+
+		// early exit if no listeners
+		if !has_wildcard && self.listeners.is_empty() && self.child_listeners.is_empty() {
+			return
+		}
+
+		let mut subscribers = self.wildcard_listeners.clone();
+		let mut changes = Vec::new();
+		let mut child_changes = Vec::new();
+
+		// Collect subscribers and changes
+		for (k, v) in changeset {
+			let k = StorageKey(k);
+			let listeners = self.listeners.get(&k);
+
+			if let Some(ref listeners) = listeners {
+				subscribers.extend(listeners.iter());
+			}
+
+			if has_wildcard || listeners.is_some() {
+				changes.push((k, v.map(StorageData)));
+			}
+		}
+		for (sk, changeset) in child_changeset {
+			let sk = StorageKey(sk);
+			if let Some((cl, cw)) = self.child_listeners.get(&sk) {
+				let mut changes = Vec::new();
+				for (k, v) in changeset {
+					let k = StorageKey(k);
+					let listeners = cl.get(&k);
+
+					if let Some(ref listeners) = listeners {
+						subscribers.extend(listeners.iter());
+					}
+
+					subscribers.extend(cw.iter());
+
+					if !cw.is_empty() || listeners.is_some() {
+						changes.push((k, v.map(StorageData)));
+					}
+				}
+				if !changes.is_empty() {
+					child_changes.push((sk, changes));
+				}
+			}
+		}
+
+		// Don't send empty notifications
+		if changes.is_empty() && child_changes.is_empty() {
+			return
+		}
+
+		let changes = Arc::<[_]>::from(changes);
+		let child_changes = Arc::<[_]>::from(child_changes);
+
+		// Trigger the events
+		self.sinks.iter_mut().for_each(|(subs_id, sink)| {
+			if subscribers.contains(subs_id) {
+				sink.was_triggered = true;
+
+				let storage_change_set = StorageChangeSet {
+					changes: changes.clone(),
+					child_changes: child_changes.clone(),
+					filter: sink.keys.clone(),
+					child_filters: sink.child_keys.clone(),
+				};
+
+				let notification =
+					StorageNotification { block: hash.clone(), changes: storage_change_set };
+
+				dispatch(subs_id, notification);
+			}
+		});
+	}
+}
+
+impl Registry {
+	fn remove_subscriber(&mut self, subscriber: SubscriberId) -> Option<(Keys, ChildKeys)> {
+		let sink = self.sinks.remove(&subscriber)?;
+
+		Self::remove_subscriber_from(
+			subscriber,
+			&sink.keys,
+			&mut self.listeners,
+			&mut self.wildcard_listeners,
+		);
+		if let Some(child_filters) = &sink.child_keys {
+			for (c_key, filters) in child_filters {
+				if let Some((listeners, wildcards)) = self.child_listeners.get_mut(&c_key) {
+					Self::remove_subscriber_from(
+						subscriber,
+						&filters,
+						&mut *listeners,
+						&mut *wildcards,
+					);
+
+					if listeners.is_empty() && wildcards.is_empty() {
+						self.child_listeners.remove(&c_key);
+					}
+				}
+			}
+		}
+		if let Some(m) = self.metrics.as_ref() {
+			m.with_label_values(&[&"removed"]).inc();
+		}
+
+		Some((sink.keys.clone(), sink.child_keys.clone()))
+	}
+
+	fn remove_subscriber_from(
+		subscriber: SubscriberId,
+		filters: &Keys,
+		listeners: &mut HashMap<StorageKey, FnvHashSet<SubscriberId>>,
+		wildcards: &mut FnvHashSet<SubscriberId>,
+	) {
+		match filters {
+			None => {
+				wildcards.remove(&subscriber);
+			},
+			Some(filters) =>
+				for key in filters.iter() {
+					let remove_key = match listeners.get_mut(key) {
+						Some(ref mut set) => {
+							set.remove(&subscriber);
+							set.is_empty()
+						},
+						None => false,
+					};
+
+					if remove_key {
+						listeners.remove(key);
+					}
+				},
+		}
+	}
+
+	fn listen_from(
+		current_id: SubscriberId,
+		filter_keys: Option<impl AsRef<[StorageKey]>>,
+		listeners: &mut HashMap<StorageKey, FnvHashSet<SubscriberId>>,
+		wildcards: &mut FnvHashSet<SubscriberId>,
+	) -> Keys {
+		match filter_keys {
+			None => {
+				wildcards.insert(current_id);
+				None
+			},
+			Some(keys) => Some(
+				keys.as_ref()
+					.iter()
+					.map(|key| {
+						listeners.entry(key.clone()).or_default().insert(current_id);
+						key.clone()
+					})
+					.collect(),
+			),
+		}
+	}
+}
+
+pub(super) struct PrintKeys<'a>(pub &'a Keys);
+impl<'a> std::fmt::Debug for PrintKeys<'a> {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		if let Some(keys) = self.0 {
+			fmt.debug_list().entries(keys.iter().map(HexDisplay::from)).finish()
+		} else {
+			write!(fmt, "None")
+		}
+	}
+}
+
+pub(super) struct PrintChildKeys<'a>(pub &'a ChildKeys);
+impl<'a> std::fmt::Debug for PrintChildKeys<'a> {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		if let Some(map) = self.0 {
+			fmt.debug_map()
+				.entries(map.iter().map(|(key, values)| (HexDisplay::from(key), PrintKeys(values))))
+				.finish()
+		} else {
+			write!(fmt, "None")
+		}
+	}
+}
@@ -0,0 +1,221 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2017-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 sp_runtime::testing::{Block as RawBlock, ExtrinsicWrapper, H256 as Hash};
+use std::iter::{empty, Empty};
+
+type TestChangeSet = (
+	Vec<(StorageKey, Option<StorageData>)>,
+	Vec<(StorageKey, Vec<(StorageKey, Option<StorageData>)>)>,
+);
+
+impl From<TestChangeSet> for StorageChangeSet {
+	fn from(changes: TestChangeSet) -> Self {
+		// warning hardcoded child trie wildcard to test upon
+		let child_filters = Some(
+			[(StorageKey(vec![4]), None), (StorageKey(vec![5]), None)]
+				.iter()
+				.cloned()
+				.collect(),
+		);
+		StorageChangeSet {
+			changes: From::from(changes.0),
+			child_changes: From::from(changes.1),
+			filter: None,
+			child_filters,
+		}
+	}
+}
+
+impl PartialEq for StorageChangeSet {
+	fn eq(&self, other: &Self) -> bool {
+		self.iter().eq(other.iter())
+	}
+}
+
+type Block = RawBlock<ExtrinsicWrapper<Hash>>;
+
+#[test]
+fn triggering_change_should_notify_wildcard_listeners() {
+	// given
+	let notifications = StorageNotifications::<Block>::new(None);
+	let child_filter = [(StorageKey(vec![4]), None)];
+	let mut recv =
+		futures::executor::block_on_stream(notifications.listen(None, Some(&child_filter[..])));
+
+	// when
+	let changeset = vec![(vec![2], Some(vec![3])), (vec![3], None)];
+	let c_changeset_1 = vec![(vec![5], Some(vec![4])), (vec![6], None)];
+	let c_changeset = vec![(vec![4], c_changeset_1)];
+	notifications.trigger(
+		&Hash::from_low_u64_be(1),
+		changeset.into_iter(),
+		c_changeset.into_iter().map(|(a, b)| (a, b.into_iter())),
+	);
+
+	// then
+	assert_eq!(
+		recv.next().map(StorageNotification::into_fields).unwrap(),
+		(
+			Hash::from_low_u64_be(1),
+			(
+				vec![
+					(StorageKey(vec![2]), Some(StorageData(vec![3]))),
+					(StorageKey(vec![3]), None),
+				],
+				vec![(
+					StorageKey(vec![4]),
+					vec![
+						(StorageKey(vec![5]), Some(StorageData(vec![4]))),
+						(StorageKey(vec![6]), None),
+					]
+				)]
+			)
+				.into()
+		)
+	);
+}
+
+#[test]
+fn should_only_notify_interested_listeners() {
+	// given
+	let notifications = StorageNotifications::<Block>::new(None);
+	let child_filter = [(StorageKey(vec![4]), Some(vec![StorageKey(vec![5])]))];
+	let mut recv1 = futures::executor::block_on_stream(
+		notifications.listen(Some(&[StorageKey(vec![1])]), None),
+	);
+	let mut recv2 = futures::executor::block_on_stream(
+		notifications.listen(Some(&[StorageKey(vec![2])]), None),
+	);
+	let mut recv3 =
+		futures::executor::block_on_stream(notifications.listen(Some(&[]), Some(&child_filter)));
+
+	// when
+	let changeset = vec![(vec![2], Some(vec![3])), (vec![1], None)];
+	let c_changeset_1 = vec![(vec![5], Some(vec![4])), (vec![6], None)];
+
+	let c_changeset = vec![(vec![4], c_changeset_1)];
+	notifications.trigger(
+		&Hash::from_low_u64_be(1),
+		changeset.into_iter(),
+		c_changeset.into_iter().map(|(a, b)| (a, b.into_iter())),
+	);
+
+	// then
+	assert_eq!(
+		recv1.next().map(StorageNotification::into_fields).unwrap(),
+		(Hash::from_low_u64_be(1), (vec![(StorageKey(vec![1]), None),], vec![]).into())
+	);
+	assert_eq!(
+		recv2.next().map(StorageNotification::into_fields).unwrap(),
+		(
+			Hash::from_low_u64_be(1),
+			(vec![(StorageKey(vec![2]), Some(StorageData(vec![3]))),], vec![]).into()
+		)
+	);
+	assert_eq!(
+		recv3.next().map(StorageNotification::into_fields).unwrap(),
+		(
+			Hash::from_low_u64_be(1),
+			(
+				vec![],
+				vec![(
+					StorageKey(vec![4]),
+					vec![(StorageKey(vec![5]), Some(StorageData(vec![4])))]
+				),]
+			)
+				.into()
+		)
+	);
+}
+
+#[test]
+fn should_cleanup_subscribers_if_dropped() {
+	// given
+	let notifications = StorageNotifications::<Block>::new(None);
+	{
+		let child_filter = [(StorageKey(vec![4]), Some(vec![StorageKey(vec![5])]))];
+		let _recv1 = futures::executor::block_on_stream(
+			notifications.listen(Some(&[StorageKey(vec![1])]), None),
+		);
+		let _recv2 = futures::executor::block_on_stream(
+			notifications.listen(Some(&[StorageKey(vec![2])]), None),
+		);
+		let _recv3 = futures::executor::block_on_stream(notifications.listen(None, None));
+		let _recv4 =
+			futures::executor::block_on_stream(notifications.listen(None, Some(&child_filter)));
+		assert_eq!(notifications.map_registry(|r| r.listeners.len()), 2);
+		assert_eq!(notifications.map_registry(|r| r.wildcard_listeners.len()), 2);
+		assert_eq!(notifications.map_registry(|r| r.child_listeners.len()), 1);
+	}
+
+	// when
+	let changeset = vec![(vec![2], Some(vec![3])), (vec![1], None)];
+	let c_changeset = empty::<(_, Empty<_>)>();
+	notifications.trigger(&Hash::from_low_u64_be(1), changeset.into_iter(), c_changeset);
+
+	// then
+	assert_eq!(notifications.map_registry(|r| r.listeners.len()), 0);
+	assert_eq!(notifications.map_registry(|r| r.wildcard_listeners.len()), 0);
+	assert_eq!(notifications.map_registry(|r| r.child_listeners.len()), 0);
+}
+
+#[test]
+fn should_cleanup_subscriber_if_stream_is_dropped() {
+	let notifications = StorageNotifications::<Block>::new(None);
+	let stream = notifications.listen(None, None);
+	assert_eq!(notifications.map_registry(|r| r.sinks.len()), 1);
+	std::mem::drop(stream);
+	assert_eq!(notifications.map_registry(|r| r.sinks.len()), 0);
+}
+
+#[test]
+fn should_not_send_empty_notifications() {
+	// given
+	let mut recv = {
+		let notifications = StorageNotifications::<Block>::new(None);
+		let recv = futures::executor::block_on_stream(notifications.listen(None, None));
+
+		// when
+		let changeset = vec![];
+		let c_changeset = empty::<(_, Empty<_>)>();
+		notifications.trigger(&Hash::from_low_u64_be(1), changeset.into_iter(), c_changeset);
+		recv
+	};
+
+	// then
+	assert_eq!(recv.next().map(StorageNotification::into_fields), None);
+}
+
+impl<B: BlockT> StorageNotifications<B> {
+	fn map_registry<MapF, Ret>(&self, map: MapF) -> Ret
+	where
+		MapF: FnOnce(&Registry) -> Ret,
+	{
+		self.0.map_registry_for_tests(map)
+	}
+}
+
+impl<H> StorageNotification<H> {
+	fn into_fields(self) -> (H, StorageChangeSet) {
+		let Self { block, changes } = self;
+		(block, changes)
+	}
+}