karyon_jsonrpc/client/
subscriptions.rs

1use std::{collections::HashMap, sync::Arc};
2
3use async_channel::{Receiver, Sender};
4#[cfg(feature = "quic")]
5use log::debug;
6use serde_json::{from_value, json, Value};
7
8use karyon_core::async_runtime::lock::RwLock;
9
10use crate::{
11    error::{Error, Result},
12    message::{Notification, NotificationResult, SubscriptionID},
13};
14
15/// Turns a received notification into its subscription id and result.
16/// Must match the server's notification encoder.
17pub type NotificationDecoder = fn(Notification) -> Result<NotificationResult>;
18
19/// Default decoder: `params` is `{"subscription": id, "result": value}`.
20pub(super) fn default_notification_decoder(nt: Notification) -> Result<NotificationResult> {
21    match nt.params {
22        Some(p) => Ok(from_value(p)?),
23        None => Err(Error::InvalidMsg("Invalid notification msg".to_string())),
24    }
25}
26
27/// Parse a raw notification message and decode it with `decoder`.
28#[cfg(feature = "quic")]
29pub(super) fn decode_notification(
30    decoder: NotificationDecoder,
31    msg: Value,
32) -> Result<NotificationResult> {
33    let nt: Notification = from_value(msg)?;
34    debug!("<-- {nt}");
35    decoder(nt)
36}
37
38/// A subscription established when the client's subscribe to a method
39pub struct Subscription {
40    id: SubscriptionID,
41    rx: Receiver<Value>,
42    tx: Sender<Value>,
43}
44
45impl Subscription {
46    fn new(id: SubscriptionID, buffer_size: usize) -> Arc<Self> {
47        let (tx, rx) = async_channel::bounded(buffer_size);
48        Arc::new(Self { tx, id, rx })
49    }
50
51    pub async fn recv(&self) -> Result<Value> {
52        self.rx.recv().await.map_err(|_| Error::SubscriptionClosed)
53    }
54
55    pub fn id(&self) -> SubscriptionID {
56        self.id
57    }
58
59    pub(crate) async fn notify(&self, val: Value) -> Result<()> {
60        if self.tx.is_full() {
61            return Err(Error::SubscriptionBufferFull);
62        }
63        self.tx.send(val).await?;
64        Ok(())
65    }
66
67    fn close(&self) {
68        self.tx.close();
69    }
70}
71
72/// Manages subscriptions for the client.
73pub(crate) struct Subscriptions {
74    subs: RwLock<HashMap<SubscriptionID, Arc<Subscription>>>,
75    sub_buffer_size: usize,
76}
77
78impl Subscriptions {
79    /// Creates a new [`Subscriptions`].
80    pub(super) fn new(sub_buffer_size: usize) -> Arc<Self> {
81        Arc::new(Self {
82            subs: RwLock::new(HashMap::new()),
83            sub_buffer_size,
84        })
85    }
86
87    /// Returns a new [`Subscription`]
88    pub(super) async fn subscribe(&self, id: SubscriptionID) -> Arc<Subscription> {
89        let sub = Subscription::new(id, self.sub_buffer_size);
90        self.subs.write().await.insert(id, sub.clone());
91        sub
92    }
93
94    /// Closes subscription channels and clear the inner map.
95    pub(super) async fn clear(&self) {
96        let mut subs = self.subs.write().await;
97        for sub in subs.values() {
98            sub.close();
99        }
100        subs.clear();
101    }
102
103    /// Unsubscribe from the provided subscription id.
104    pub(super) async fn unsubscribe(&self, id: &SubscriptionID) {
105        if let Some(sub) = self.subs.write().await.remove(id) {
106            sub.close();
107        }
108    }
109
110    /// Notifies the subscription about the given decoded notification.
111    pub(super) async fn notify(&self, nt_res: NotificationResult) -> Result<()> {
112        match self.subs.read().await.get(&nt_res.subscription) {
113            Some(s) => s.notify(nt_res.result.unwrap_or(json!(""))).await?,
114            None => {
115                return Err(Error::InvalidMsg("Unknown notification".to_string()));
116            }
117        }
118
119        Ok(())
120    }
121}
karyon_jsonrpc/client/subscriptions.rs

karyon_jsonrpc/client/
subscriptions.rs
