jsonrpsee_core/server/
subscription.rs

1// Copyright 2019-2021 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any
4// person obtaining a copy of this software and associated
5// documentation files (the "Software"), to deal in the
6// Software without restriction, including without
7// limitation the rights to use, copy, modify, merge,
8// publish, distribute, sublicense, and/or sell copies of
9// the Software, and to permit persons to whom the Software
10// is furnished to do so, subject to the following
11// conditions:
12//
13// The above copyright notice and this permission notice
14// shall be included in all copies or substantial portions
15// of the Software.
16//
17// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
18// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
19// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
20// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
21// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
22// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
23// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
24// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
25// DEALINGS IN THE SOFTWARE.
26
27//! Subscription related types and traits for server implementations.
28
29use super::{helpers::MethodSink, MethodResponse, ResponsePayload};
30use crate::{
31    error::SubscriptionError,
32    server::{
33        error::{DisconnectError, PendingSubscriptionAcceptError, SendTimeoutError, TrySendError},
34        rpc_module::ConnectionId,
35    },
36    traits::IdProvider,
37};
38use jsonrpsee_types::{
39    response::SubscriptionPayloadError, ErrorCode, ErrorObjectOwned, Id, SubscriptionId,
40    SubscriptionPayload, SubscriptionResponse,
41};
42use parking_lot::Mutex;
43use rustc_hash::FxHashMap;
44use serde::Serialize;
45use serde_json::value::RawValue;
46use std::{collections::hash_map::Entry, sync::Arc, time::Duration};
47use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
48
49/// Type-alias for subscribers.
50pub type Subscribers = Arc<Mutex<FxHashMap<SubscriptionKey, (MethodSink, mpsc::Receiver<()>)>>>;
51/// Subscription permit.
52pub type SubscriptionPermit = OwnedSemaphorePermit;
53
54/// Convert something into a subscription close notification
55/// before a subscription is terminated.
56pub trait IntoSubscriptionCloseResponse {
57    /// Convert something into a subscription response
58    fn into_response(self) -> SubscriptionCloseResponse;
59}
60
61/// Represents what action that will sent when a subscription callback returns.
62#[derive(Debug)]
63pub enum SubscriptionCloseResponse {
64    /// No further message will be sent.
65    None,
66    /// Send a subscription notification.
67    ///
68    /// The subscription notification has the following format:
69    ///
70    /// ```json
71    /// {
72    ///  "jsonrpc": "2.0",
73    ///  "method": "<method>",
74    ///  "params": {
75    ///    "subscription": "<subscriptionID>",
76    ///    "result": <your msg>
77    ///    }
78    ///  }
79    /// }
80    /// ```
81    Notif(SubscriptionMessage),
82    /// Send a subscription error notification
83    ///
84    /// The error notification has the following format:
85    ///
86    /// ```json
87    /// {
88    ///  "jsonrpc": "2.0",
89    ///  "method": "<method>",
90    ///  "params": {
91    ///    "subscription": "<subscriptionID>",
92    ///    "error": <your msg>
93    ///    }
94    ///  }
95    /// }
96    /// ```
97    NotifErr(SubscriptionError),
98}
99
100impl IntoSubscriptionCloseResponse for Result<(), SubscriptionError> {
101    fn into_response(self) -> SubscriptionCloseResponse {
102        match self {
103            Ok(()) => SubscriptionCloseResponse::None,
104            Err(e) => SubscriptionCloseResponse::NotifErr(e),
105        }
106    }
107}
108
109impl IntoSubscriptionCloseResponse for () {
110    fn into_response(self) -> SubscriptionCloseResponse {
111        SubscriptionCloseResponse::None
112    }
113}
114
115impl IntoSubscriptionCloseResponse for SubscriptionCloseResponse {
116    fn into_response(self) -> Self {
117        self
118    }
119}
120
121/// A complete subscription message or partial subscription message.
122#[derive(Debug, Clone)]
123pub enum SubscriptionMessageInner {
124    /// Complete JSON message.
125    Complete(String),
126    /// Need subscription ID and method name.
127    NeedsData(Box<RawValue>),
128}
129
130/// Subscription message.
131#[derive(Debug, Clone)]
132pub struct SubscriptionMessage(pub(crate) SubscriptionMessageInner);
133
134impl From<Box<RawValue>> for SubscriptionMessage {
135    fn from(json: Box<RawValue>) -> Self {
136        Self(SubscriptionMessageInner::NeedsData(json))
137    }
138}
139
140impl SubscriptionMessage {
141    /// Create a subscription message this is more efficient than [`SubscriptionMessage::from`]
142    /// because it only allocates once.
143    ///
144    /// Fails if the json `result` couldn't be serialized.
145    pub fn new(
146        method: &str,
147        subscription: SubscriptionId,
148        result: &impl Serialize,
149    ) -> Result<Self, serde_json::Error> {
150        let json = serde_json::to_string(&SubscriptionResponse::new(
151            method.into(),
152            SubscriptionPayload { subscription, result },
153        ))?;
154        Ok(Self::from_complete_message(json))
155    }
156
157    pub(crate) fn from_complete_message(msg: String) -> Self {
158        SubscriptionMessage(SubscriptionMessageInner::Complete(msg))
159    }
160}
161
162/// Represent a unique subscription entry based on [`SubscriptionId`] and [`ConnectionId`].
163#[derive(Clone, Debug, PartialEq, Eq, Hash)]
164pub struct SubscriptionKey {
165    pub(crate) conn_id: ConnectionId,
166    pub(crate) sub_id: SubscriptionId<'static>,
167}
168
169/// Represents a subscription until it is unsubscribed.
170// NOTE: The reason why we use `mpsc` here is because it allows `IsUnsubscribed::unsubscribed`
171// to be &self instead of &mut self.
172#[derive(Debug, Clone)]
173pub struct IsUnsubscribed(mpsc::Sender<()>);
174
175impl IsUnsubscribed {
176    /// Returns true if the unsubscribe method has been invoked or the subscription has been
177    /// canceled.
178    ///
179    /// This can be called multiple times as the element in the channel is never
180    /// removed.
181    pub fn is_unsubscribed(&self) -> bool {
182        self.0.is_closed()
183    }
184
185    /// Wrapper over [`tokio::sync::mpsc::Sender::closed`]
186    ///
187    /// # Cancel safety
188    ///
189    /// This method is cancel safe. Once the channel is closed,
190    /// it stays closed forever and all future calls to closed will return immediately.
191    pub async fn unsubscribed(&self) {
192        self.0.closed().await;
193    }
194}
195
196/// Represents a single subscription that is waiting to be accepted or rejected.
197///
198/// If this is dropped without calling `PendingSubscription::reject` or
199/// `PendingSubscriptionSink::accept` a default error is sent out as response to the subscription
200/// call.
201///
202/// Thus, if you want a customized error message then `PendingSubscription::reject` must be called.
203#[derive(Debug)]
204#[must_use = "PendingSubscriptionSink does nothing unless `accept` or `reject` is called"]
205pub struct PendingSubscriptionSink {
206    /// Sink.
207    pub(crate) inner: MethodSink,
208    /// MethodCallback.
209    pub(crate) method: &'static str,
210    /// Shared Mutex of subscriptions for this method.
211    pub(crate) subscribers: Subscribers,
212    /// Unique subscription.
213    pub(crate) uniq_sub: SubscriptionKey,
214    /// ID of the `subscription call` (i.e. not the same as subscription id) which is used
215    /// to reply to subscription method call and must only be used once.
216    pub(crate) id: Id<'static>,
217    /// Sender to answer the subscribe call.
218    pub(crate) subscribe: oneshot::Sender<MethodResponse>,
219    /// Subscription permit.
220    pub(crate) permit: OwnedSemaphorePermit,
221}
222
223impl PendingSubscriptionSink {
224    /// Reject the subscription by responding to the subscription method call with
225    /// the error message from [`jsonrpsee_types::error::ErrorObject`].
226    ///
227    /// # Note
228    ///
229    /// If this is used in the async subscription callback
230    /// the return value is simply ignored because no further notification are propagated
231    /// once reject has been called.
232    pub async fn reject(self, err: impl Into<ErrorObjectOwned>) {
233        let err = MethodResponse::subscription_error(self.id, err.into());
234        _ = self.inner.send(err.to_json()).await;
235        _ = self.subscribe.send(err);
236    }
237
238    /// Attempt to accept the subscription and respond the subscription method call.
239    ///
240    /// # Errors
241    ///
242    /// Fails if the connection is closed, or if the subscription response exceeds the
243    /// `max_response_size`. In the latter case the call is answered with an oversized response
244    /// error and the subscription is not registered.
245    pub async fn accept(self) -> Result<SubscriptionSink, PendingSubscriptionAcceptError> {
246        let (tx, rx) = mpsc::channel(1);
247
248        // The key is reserved before the acknowledgement is sent: an occupied key belongs to a live
249        // subscription of this connection, which must not be replaced.
250        let reserved = match self.subscribers.lock().entry(self.uniq_sub.clone()) {
251            Entry::Vacant(entry) => {
252                entry.insert((self.inner.clone(), rx));
253                true
254            }
255            Entry::Occupied(_) => false,
256        };
257
258        if !reserved {
259            self.reject(ErrorCode::InternalError).await;
260            return Err(PendingSubscriptionAcceptError);
261        }
262
263        // Dropping the sink releases the reservation again if the acknowledgement isn't delivered.
264        let sink = SubscriptionSink {
265            inner: self.inner,
266            method: self.method,
267            subscribers: self.subscribers,
268            uniq_sub: self.uniq_sub,
269            unsubscribe: IsUnsubscribed(tx),
270            _permit: Some(Arc::new(self.permit)),
271        };
272
273        let response = MethodResponse::subscription_response(
274            self.id,
275            ResponsePayload::success_borrowed(&sink.uniq_sub.sub_id),
276            sink.inner.max_response_size() as usize,
277        );
278        let success = response.is_success();
279
280        // TODO: #1052
281        //
282        // Ideally the message should be sent only once.
283        //
284        // The same message is sent twice here because one is sent directly to the transport layer
285        // and the other one is sent internally to accept the subscription.
286        sink.inner.send(response.to_json()).await.map_err(|_| PendingSubscriptionAcceptError)?;
287        self.subscribe.send(response).map_err(|_| PendingSubscriptionAcceptError)?;
288
289        if success {
290            Ok(sink)
291        } else {
292            // The response exceeded the max response size and the call was answered with an error.
293            Err(PendingSubscriptionAcceptError)
294        }
295    }
296
297    /// Get the subscription ID.
298    pub fn subscription_id(&self) -> SubscriptionId<'static> {
299        self.uniq_sub.sub_id.clone()
300    }
301
302    /// Returns connection identifier, which was used to perform pending subscription request
303    pub fn connection_id(&self) -> ConnectionId {
304        self.uniq_sub.conn_id
305    }
306
307    /// Get the capacity of the channel.
308    pub fn capacity(&self) -> usize {
309        self.inner.capacity()
310    }
311
312    /// Get the max capacity of the channel.
313    pub fn max_capacity(&self) -> usize {
314        self.inner.max_capacity()
315    }
316
317    /// Get the method name.
318    pub fn method_name(&self) -> &str {
319        self.method
320    }
321}
322
323/// A handle to an active subscription.
324///
325/// Clones share the same subscription. The subscription is torn down when the last
326/// clone is dropped or the client explicitly unsubscribes.
327#[derive(Debug, Clone)]
328pub struct SubscriptionSink {
329    /// Sink.
330    inner: MethodSink,
331    /// MethodCallback.
332    method: &'static str,
333    /// Shared Mutex of subscriptions for this method.
334    subscribers: Subscribers,
335    /// Unique subscription.
336    uniq_sub: SubscriptionKey,
337    /// A future to that fires once the unsubscribe method has been called.
338    unsubscribe: IsUnsubscribed,
339    /// Subscription permit.
340    _permit: Option<Arc<SubscriptionPermit>>,
341}
342
343impl SubscriptionSink {
344    /// Get the subscription ID.
345    pub fn subscription_id(&self) -> SubscriptionId<'static> {
346        self.uniq_sub.sub_id.clone()
347    }
348
349    /// Get the method name.
350    pub fn method_name(&self) -> &str {
351        self.method
352    }
353
354    /// Get the connection ID.
355    pub fn connection_id(&self) -> ConnectionId {
356        self.uniq_sub.conn_id
357    }
358
359    /// Send out a response on the subscription and wait until there is capacity.
360    ///
361    ///
362    /// Returns
363    /// - `Ok(())` if the message could be sent.
364    /// - `Err(unsent_msg)` if the connection or subscription was closed.
365    ///
366    /// # Cancel safety
367    ///
368    /// This method is cancel-safe and dropping a future loses its spot in the waiting queue.
369    pub async fn send(&self, msg: impl Into<SubscriptionMessage>) -> Result<(), DisconnectError> {
370        let msg = msg.into();
371
372        // Only possible to trigger when the connection is dropped.
373        if self.is_closed() {
374            return Err(DisconnectError(msg));
375        }
376
377        let json = sub_message_to_json(msg, &self.uniq_sub.sub_id, self.method);
378        self.inner.send(json).await
379    }
380
381    /// Similar to `SubscriptionSink::send` but only waits for a limited time.
382    pub async fn send_timeout(
383        &self,
384        msg: impl Into<SubscriptionMessage>,
385        timeout: Duration,
386    ) -> Result<(), SendTimeoutError> {
387        let msg = msg.into();
388
389        // Only possible to trigger when the connection is dropped.
390        if self.is_closed() {
391            return Err(SendTimeoutError::Closed(msg));
392        }
393
394        let json = sub_message_to_json(msg, &self.uniq_sub.sub_id, self.method);
395        self.inner.send_timeout(json, timeout).await
396    }
397
398    /// Attempts to immediately send out the message as JSON string to the subscribers but fails if
399    /// the channel is full or the connection/subscription is closed
400    ///
401    ///
402    /// This differs from [`SubscriptionSink::send`] where it will until there is capacity
403    /// in the channel.
404    pub fn try_send(&mut self, msg: impl Into<SubscriptionMessage>) -> Result<(), TrySendError> {
405        let msg = msg.into();
406
407        // Only possible to trigger when the connection is dropped.
408        if self.is_closed() {
409            return Err(TrySendError::Closed(msg));
410        }
411
412        let json = sub_message_to_json(msg, &self.uniq_sub.sub_id, self.method);
413        self.inner.try_send(json)
414    }
415
416    /// Returns whether the subscription is closed.
417    pub fn is_closed(&self) -> bool {
418        self.inner.is_closed() || !self.is_active_subscription()
419    }
420
421    /// Completes when the subscription has been closed.
422    pub async fn closed(&self) {
423        // Both are cancel-safe thus ok to use select here.
424        tokio::select! {
425            _ = self.inner.closed() => (),
426            _ = self.unsubscribe.unsubscribed() => (),
427        }
428    }
429
430    /// Get the capacity of the subscription.
431    pub fn capacity(&self) -> usize {
432        self.inner.capacity()
433    }
434
435    /// Get the max capacity of the subscription.
436    pub fn max_capacity(&self) -> usize {
437        self.inner.max_capacity()
438    }
439
440    fn is_active_subscription(&self) -> bool {
441        !self.unsubscribe.is_unsubscribed()
442    }
443}
444
445impl Drop for SubscriptionSink {
446    fn drop(&mut self) {
447        // The activity check must happen under the lock: unsubscribing removes the entry (which
448        // marks this sink inactive) under the lock, and the same key can then be registered again.
449        let mut subs = self.subscribers.lock();
450        // Only the last clone should tear down the subscription (see #1622).
451        // The lock serializes concurrent drops so that exactly one `try_unwrap` succeeds.
452        if self.is_active_subscription() &&
453            self._permit.take().is_some_and(|p| Arc::try_unwrap(p).is_ok())
454        {
455            subs.remove(&self.uniq_sub);
456        }
457    }
458}
459
460/// This wraps [`tokio::sync::Semaphore`] and is used to limit the number of subscriptions per
461/// connection.
462#[derive(Debug, Clone)]
463pub struct BoundedSubscriptions {
464    guard: Arc<Semaphore>,
465    max: u32,
466}
467
468impl BoundedSubscriptions {
469    /// Create a new bounded subscription.
470    pub fn new(max_subscriptions: u32) -> Self {
471        Self { guard: Arc::new(Semaphore::new(max_subscriptions as usize)), max: max_subscriptions }
472    }
473
474    /// Attempts to acquire a subscription slot.
475    ///
476    /// Fails if `max_subscriptions` have been exceeded.
477    pub fn acquire(&self) -> Option<SubscriptionPermit> {
478        Arc::clone(&self.guard).try_acquire_owned().ok()
479    }
480
481    /// Get the maximum number of permitted subscriptions.
482    pub const fn max(&self) -> u32 {
483        self.max
484    }
485}
486
487#[derive(Debug)]
488/// Helper struct to manage subscriptions.
489pub struct SubscriptionState<'a> {
490    /// Connection ID
491    pub conn_id: ConnectionId,
492    /// ID provider.
493    pub id_provider: &'a dyn IdProvider,
494    /// Subscription limit
495    pub subscription_permit: SubscriptionPermit,
496}
497
498pub(crate) fn sub_message_to_json(
499    msg: SubscriptionMessage,
500    sub_id: &SubscriptionId,
501    method: &str,
502) -> String {
503    match msg.0 {
504        SubscriptionMessageInner::Complete(msg) => msg,
505        SubscriptionMessageInner::NeedsData(result) => {
506            serde_json::to_string(&SubscriptionResponse::new(
507                method.into(),
508                SubscriptionPayload { subscription: sub_id.clone(), result },
509            ))
510            .expect("Serialize infallible; qed")
511        }
512    }
513}
514
515pub(crate) fn sub_err_to_json(
516    error: SubscriptionError,
517    sub_id: SubscriptionId,
518    method: &str,
519) -> String {
520    serde_json::to_string(&jsonrpsee_types::response::SubscriptionError::new(
521        method.into(),
522        SubscriptionPayloadError { subscription: sub_id, error },
523    ))
524    .expect("Serialize infallible; qed")
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530    use crate::server::{MethodCallback, Methods, RpcModule};
531    use futures_util::FutureExt;
532    use jsonrpsee_types::Params;
533    use serde_json::value::to_raw_value;
534    use std::{
535        sync::atomic::{AtomicBool, Ordering},
536        time::Instant,
537    };
538
539    const SUB_ID: &str = "0xdeadbeef";
540
541    /// Id provider that hands out the same id for every subscription.
542    #[derive(Debug)]
543    struct ConstantId;
544
545    impl IdProvider for ConstantId {
546        fn next_id(&self) -> SubscriptionId<'static> {
547            SUB_ID.to_string().into()
548        }
549    }
550
551    /// Drives the callbacks that `RpcModule::register_subscription_raw` registers.
552    struct Harness {
553        methods: Methods,
554        pending: mpsc::UnboundedReceiver<PendingSubscriptionSink>,
555        permits: BoundedSubscriptions,
556    }
557
558    impl Harness {
559        fn new() -> Self {
560            let (pending_tx, pending) = mpsc::unbounded_channel();
561            let mut module = RpcModule::new(pending_tx);
562            module
563                .register_subscription_raw("sub", "notif", "unsub", |_, pending, ctx| {
564                    ctx.send(pending).unwrap();
565                })
566                .unwrap();
567
568            Self {
569                methods: Methods::clone(&module),
570                pending,
571                permits: BoundedSubscriptions::new(16),
572            }
573        }
574
575        /// Calls the subscribe method with request id `id` and returns the pending subscription
576        /// together with the future that resolves to the answer of the subscribe call.
577        fn subscribe(
578            &mut self,
579            id: u64,
580            conn_id: ConnectionId,
581            sink: &mpsc::Sender<String>,
582        ) -> (PendingSubscriptionSink, futures_util::future::BoxFuture<'static, MethodResponse>)
583        {
584            let state = SubscriptionState {
585                conn_id,
586                id_provider: &ConstantId,
587                subscription_permit: self.permits.acquire().unwrap(),
588            };
589            let Some(MethodCallback::Subscription(subscribe)) = self.methods.method("sub") else {
590                panic!("`sub` is a subscription");
591            };
592            let response =
593                subscribe(Id::Number(id), Params::new(None), MethodSink::new(sink.clone()), state);
594            (self.pending.try_recv().unwrap(), response)
595        }
596
597        /// Calls the unsubscribe method for `SUB_ID` and returns its json answer.
598        fn unsubscribe(&self, id: u64, conn_id: ConnectionId) -> String {
599            call_unsubscribe(&self.methods, SUB_ID, id, conn_id)
600        }
601    }
602
603    struct StopOnDrop<'a>(&'a AtomicBool);
604
605    impl Drop for StopOnDrop<'_> {
606        fn drop(&mut self) {
607            self.0.store(true, Ordering::Relaxed);
608        }
609    }
610
611    fn call_unsubscribe(methods: &Methods, sub_id: &str, id: u64, conn_id: ConnectionId) -> String {
612        let params = format!(r#"["{sub_id}"]"#);
613        let Some(MethodCallback::Unsubscription(unsubscribe)) = methods.method("unsub") else {
614            panic!("`unsub` is an unsubscription");
615        };
616        unsubscribe(Id::Number(id), Params::new(Some(&params)), conn_id, usize::MAX).into_json()
617    }
618
619    #[tokio::test]
620    async fn colliding_subscription_id_keeps_existing_subscription() {
621        let mut harness = Harness::new();
622        let conn_id = ConnectionId(0);
623        let (tx, mut rx) = mpsc::channel(8);
624
625        let (pending, response) = harness.subscribe(1, conn_id, &tx);
626        let first = pending.accept().await.unwrap();
627        let ack = r#"{"jsonrpc":"2.0","id":1,"result":"0xdeadbeef"}"#;
628        assert_eq!(response.await.into_json(), ack);
629        assert_eq!(rx.recv().await.unwrap(), ack);
630
631        // The provider hands out the id again: the second call fails with exactly one error.
632        let (pending, response) = harness.subscribe(2, conn_id, &tx);
633        assert!(matches!(pending.accept().await, Err(PendingSubscriptionAcceptError)));
634        let error =
635            r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32603,"message":"Internal error"}}"#;
636        assert_eq!(response.await.into_json(), error);
637        assert_eq!(rx.recv().await.unwrap(), error);
638        assert!(rx.try_recv().is_err());
639
640        // The first subscription is untouched and still delivers items.
641        assert!(!first.is_closed());
642        first.send(to_raw_value(&"item").unwrap()).await.unwrap();
643        assert_eq!(
644            rx.recv().await.unwrap(),
645            r#"{"jsonrpc":"2.0","method":"notif","params":{"subscription":"0xdeadbeef","result":"item"}}"#
646        );
647
648        // It can be unsubscribed, which removes it.
649        assert_eq!(harness.unsubscribe(3, conn_id), r#"{"jsonrpc":"2.0","id":3,"result":true}"#);
650        assert!(first.is_closed());
651        assert_eq!(harness.unsubscribe(4, conn_id), r#"{"jsonrpc":"2.0","id":4,"result":false}"#);
652    }
653
654    #[tokio::test]
655    async fn same_subscription_id_on_different_connections() {
656        let mut harness = Harness::new();
657        let (tx, _rx) = mpsc::channel(8);
658
659        let (pending, _first_response) = harness.subscribe(1, ConnectionId(0), &tx);
660        let first = pending.accept().await.unwrap();
661        let (pending, _second_response) = harness.subscribe(1, ConnectionId(1), &tx);
662        let second = pending.accept().await.unwrap();
663
664        assert_eq!(
665            harness.unsubscribe(2, ConnectionId(0)),
666            r#"{"jsonrpc":"2.0","id":2,"result":true}"#
667        );
668        assert!(first.is_closed());
669        assert!(!second.is_closed());
670        assert_eq!(
671            harness.unsubscribe(3, ConnectionId(1)),
672            r#"{"jsonrpc":"2.0","id":3,"result":true}"#
673        );
674        assert!(second.is_closed());
675    }
676
677    #[tokio::test]
678    async fn failed_acknowledgement_releases_subscription_id() {
679        let mut harness = Harness::new();
680        let conn_id = ConnectionId(0);
681
682        // The connection is gone, so the acknowledgement can't be delivered.
683        let (tx, rx) = mpsc::channel(8);
684        drop(rx);
685        let (pending, _) = harness.subscribe(1, conn_id, &tx);
686        assert!(matches!(pending.accept().await, Err(PendingSubscriptionAcceptError)));
687        assert_eq!(harness.unsubscribe(2, conn_id), r#"{"jsonrpc":"2.0","id":2,"result":false}"#);
688
689        // The id isn't blocked for later subscriptions.
690        let (tx, _rx) = mpsc::channel(8);
691        let (pending, _response) = harness.subscribe(3, conn_id, &tx);
692        let sink = pending.accept().await.unwrap();
693        assert_eq!(harness.unsubscribe(4, conn_id), r#"{"jsonrpc":"2.0","id":4,"result":true}"#);
694        assert!(sink.is_closed());
695    }
696
697    #[tokio::test]
698    async fn dropping_unsubscribed_sink_keeps_replacement_subscription() {
699        let mut harness = Harness::new();
700        let conn_id = ConnectionId(0);
701        let (tx, _rx) = mpsc::channel(8);
702
703        let (pending, _first_response) = harness.subscribe(1, conn_id, &tx);
704        let first = pending.accept().await.unwrap();
705        assert_eq!(harness.unsubscribe(2, conn_id), r#"{"jsonrpc":"2.0","id":2,"result":true}"#);
706
707        // The same id is registered again while the sink of the old subscription is still alive.
708        let (pending, _second_response) = harness.subscribe(3, conn_id, &tx);
709        let second = pending.accept().await.unwrap();
710
711        drop(first);
712        assert!(!second.is_closed());
713        assert_eq!(harness.unsubscribe(4, conn_id), r#"{"jsonrpc":"2.0","id":4,"result":true}"#);
714        assert!(second.is_closed());
715    }
716
717    /// Drops the sink of a subscription while it is unsubscribed and registered again under the
718    /// same id from another thread. The replacement must survive in every interleaving.
719    #[test]
720    fn dropping_sink_races_with_unsubscribe_and_resubscribe() {
721        let mut harness = Harness::new();
722        let methods = harness.methods.clone();
723        let conn_id = ConnectionId(0);
724        let (tx, mut rx) = mpsc::channel(8);
725        let deadline = Instant::now() + Duration::from_secs(1);
726        let stop = AtomicBool::new(false);
727        let mut removed_replacements = 0;
728
729        std::thread::scope(|scope| {
730            // Stops the contending threads even if an assertion below panics.
731            let _stop = StopOnDrop(&stop);
732
733            // Contend for the registry lock like other connections do.
734            for i in 0..4 {
735                let (methods, stop) = (&methods, &stop);
736                scope.spawn(move || {
737                    while !stop.load(Ordering::Relaxed) {
738                        call_unsubscribe(methods, "other", 0, ConnectionId(1000 + i));
739                    }
740                });
741            }
742
743            while Instant::now() < deadline {
744                while rx.try_recv().is_ok() {}
745
746                let (pending, _first_response) = harness.subscribe(1, conn_id, &tx);
747                let first = pending.accept().now_or_never().unwrap().unwrap();
748                let (pending, _second_response) = harness.subscribe(2, conn_id, &tx);
749
750                let go = AtomicBool::new(false);
751                let second = std::thread::scope(|race| {
752                    let (methods, go) = (&methods, &go);
753                    let dropper = race.spawn(move || {
754                        while !go.load(Ordering::Acquire) {
755                            std::hint::spin_loop();
756                        }
757                        drop(first);
758                    });
759                    let replacer = race.spawn(move || {
760                        while !go.load(Ordering::Acquire) {
761                            std::hint::spin_loop();
762                        }
763                        call_unsubscribe(methods, SUB_ID, 3, conn_id);
764                        pending.accept().now_or_never().unwrap().unwrap()
765                    });
766                    go.store(true, Ordering::Release);
767                    dropper.join().unwrap();
768                    replacer.join().unwrap()
769                });
770
771                if second.is_closed() {
772                    removed_replacements += 1;
773                }
774            }
775        });
776
777        assert_eq!(removed_replacements, 0);
778    }
779}
jsonrpsee_core/server/subscription.rs

jsonrpsee_core/server/
subscription.rs
