neo3/neo_clients/rpc/
pubsub.rs

1use std::{
2	collections::VecDeque,
3	marker::PhantomData,
4	pin::Pin,
5	task::{Context, Poll},
6};
7
8use crate::neo_clients::{JsonRpcProvider, RpcClient};
9use futures_util::stream::Stream;
10use pin_project::{pin_project, pinned_drop};
11use primitive_types::U256;
12use serde::de::DeserializeOwned;
13use serde_json::value::RawValue;
14use tracing::error;
15
16/// A transport implementation supporting pub sub subscriptions.
17pub trait PubsubClient: JsonRpcProvider {
18	/// The type of stream this transport returns
19	type NotificationStream: futures_core::Stream<Item = Box<RawValue>> + Send + Unpin;
20
21	/// Add a subscription to this transport
22	fn subscribe<T: Into<U256>>(&self, id: T) -> Result<Self::NotificationStream, Self::Error>;
23
24	/// Remove a subscription from this transport
25	fn unsubscribe<T: Into<U256>>(&self, id: T) -> Result<(), Self::Error>;
26}
27
28#[must_use = "subscriptions do nothing unless you stream them"]
29#[pin_project(PinnedDrop)]
30/// Streams data from an installed filter via `neo_subscribe`
31pub struct SubscriptionStream<'a, P: PubsubClient, R: DeserializeOwned> {
32	/// The subscription's installed id on the neo node
33	pub id: U256,
34
35	loaded_elements: VecDeque<R>,
36
37	pub(crate) provider: &'a RpcClient<P>,
38
39	#[pin]
40	rx: P::NotificationStream,
41
42	ret: PhantomData<R>,
43}
44
45impl<'a, P, R> SubscriptionStream<'a, P, R>
46where
47	P: PubsubClient,
48	R: DeserializeOwned,
49{
50	/// Creates a new subscription stream for the provided subscription id.
51	///
52	/// ### Note
53	/// Most providers treat `SubscriptionStream` IDs as global singletons.
54	/// Instantiating this directly with a known ID will likely cause any
55	/// existing streams with that ID to end. To avoid this, start a new stream
56	/// using the provider's subscription method instead of `SubscriptionStream::new`.
57	pub fn new(id: U256, provider: &'a RpcClient<P>) -> Result<Self, P::Error> {
58		// Call the underlying PubsubClient's subscribe
59		let rx = provider.as_ref().subscribe(id)?;
60		Ok(Self { id, provider, rx, ret: PhantomData, loaded_elements: VecDeque::new() })
61	}
62
63	/// Set the loaded elements buffer. This buffer contains logs waiting for
64	/// the consumer to read. Setting the buffer can be used to add logs
65	/// without receiving them from the RPC node
66	///
67	/// ### Warning
68	///
69	/// Setting the buffer will drop any logs in the current buffer.
70	pub fn set_loaded_elements(&mut self, loaded_elements: VecDeque<R>) {
71		self.loaded_elements = loaded_elements;
72	}
73}
74
75// Each subscription item is a serde_json::Value which must be decoded to the
76// subscription's return type.
77// Note: While this could potentially be replaced with an `rx.map` in the constructor,
78// the current implementation allows for better error handling and logging of
79// deserialization failures without dropping the entire stream.
80impl<'a, P, R> Stream for SubscriptionStream<'a, P, R>
81where
82	P: PubsubClient,
83	R: DeserializeOwned,
84{
85	type Item = R;
86
87	fn poll_next(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
88		if !self.loaded_elements.is_empty() {
89			let next_element = self.get_mut().loaded_elements.pop_front();
90			return Poll::Ready(next_element);
91		}
92
93		let mut this = self.project();
94		loop {
95			return match futures_util::ready!(this.rx.as_mut().poll_next(ctx)) {
96				Some(item) => match serde_json::from_str(item.get()) {
97					Ok(res) => Poll::Ready(Some(res)),
98					Err(err) => {
99						error!("failed to deserialize item {:?}", err);
100						continue;
101					},
102				},
103				None => Poll::Ready(None),
104			};
105		}
106	}
107}
108
109#[pinned_drop]
110impl<P, R> PinnedDrop for SubscriptionStream<'_, P, R>
111where
112	P: PubsubClient,
113	R: DeserializeOwned,
114{
115	fn drop(self: Pin<&mut Self>) {
116		// on drop it removes the handler from the websocket so that it stops
117		// getting populated. We need to call `unsubscribe` explicitly to cancel
118		// the subscription
119		let _ = (*self.provider).as_ref().unsubscribe(self.id);
120	}
121}
neo3/neo_clients/rpc/pubsub.rs

neo3/neo_clients/rpc/
pubsub.rs
