wowlab_centrifuge/client/
commands.rs

1use std::{sync::atomic::Ordering, time::Duration};
2
3use tokio::sync::{mpsc, oneshot};
4
5use super::{Client, CommandRequest, EVENT_CHANNEL_BUFFER, check_reply};
6use crate::{
7    Subscription,
8    error::Error,
9    proto,
10    subscription::{SubscriptionConfig, SubscriptionInner, SubscriptionState},
11    types::{
12        HistoryResult, PresenceResult, PresenceStats, RpcResult, StreamPosition, SubRefreshResult,
13        SubscribeResult, SubscriptionEvent,
14    },
15};
16
17const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
18
19impl Client {
20    /// Registers a subscription and starts it immediately when connected.
21    /// # Errors
22    /// Returns an error if the channel is already subscribed or the server rejects the subscribe command.
23    pub async fn subscribe(&self, config: SubscriptionConfig) -> Result<Subscription, Error> {
24        let channel = config.channel.clone();
25
26        {
27            let inner = self.inner.read().await;
28
29            if inner.subscriptions.contains_key(&channel) {
30                return Err(Error::already_subscribed(channel));
31            }
32        }
33
34        let (event_tx, event_rx) = mpsc::channel(EVENT_CHANNEL_BUFFER);
35        let sub_inner = SubscriptionInner::new(config, event_tx);
36
37        {
38            let mut inner = self.inner.write().await;
39
40            inner.subscriptions.insert(channel.clone(), sub_inner)
41        };
42
43        if self.is_connected().await {
44            self.subscribe_internal(&channel).await?;
45        }
46
47        Ok(Subscription::new(channel, event_rx))
48    }
49
50    /// Removes the subscription for `channel` from the server and client.
51    /// # Errors
52    /// Returns an error if the client is disconnected or the command fails.
53    pub async fn unsubscribe(&self, channel: &str) -> Result<(), Error> {
54        let cmd = proto::Command {
55            unsubscribe: Some(proto::UnsubscribeRequest {
56                channel: channel.to_string(),
57            }),
58            ..Default::default()
59        };
60
61        self.send_void(cmd).await?;
62
63        {
64            let mut inner = self.inner.write().await;
65
66            inner.subscriptions.remove(channel)
67        };
68
69        Ok(())
70    }
71
72    /// Queries presence information for `channel`.
73    /// # Errors
74    /// Returns an error if the client is disconnected or the command fails.
75    pub async fn presence(&self, channel: &str) -> Result<PresenceResult, Error> {
76        let cmd = proto::Command {
77            presence: Some(proto::PresenceRequest {
78                channel: channel.to_string(),
79            }),
80            ..Default::default()
81        };
82
83        self.send_extract(cmd, "presence", |r| r.presence).await
84    }
85
86    /// Queries aggregate presence statistics for `channel`.
87    /// # Errors
88    /// Returns an error if the client is disconnected or the command fails.
89    pub async fn presence_stats(&self, channel: &str) -> Result<PresenceStats, Error> {
90        let cmd = proto::Command {
91            presence_stats: Some(proto::PresenceStatsRequest {
92                channel: channel.to_string(),
93            }),
94            ..Default::default()
95        };
96
97        self.send_extract(cmd, "presence stats", |r| r.presence_stats)
98            .await
99    }
100
101    /// Queries publication history for `channel`.
102    /// # Errors
103    /// Returns an error if the client is disconnected or the command fails.
104    pub async fn history(
105        &self,
106        channel: &str,
107        limit: i32,
108        since: Option<StreamPosition>,
109        reverse: bool,
110    ) -> Result<HistoryResult, Error> {
111        let cmd = proto::Command {
112            history: Some(proto::HistoryRequest {
113                channel: channel.to_string(),
114                limit,
115                since: since.map(|p| proto::StreamPosition {
116                    offset: p.offset,
117                    epoch: p.epoch,
118                }),
119                reverse,
120            }),
121            ..Default::default()
122        };
123
124        self.send_extract(cmd, "history", |r| r.history).await
125    }
126
127    /// Publishes `data` to `channel`.
128    /// # Errors
129    /// Returns an error if the client is disconnected or the command fails.
130    pub async fn publish(&self, channel: &str, data: Vec<u8>) -> Result<(), Error> {
131        let cmd = proto::Command {
132            publish: Some(proto::PublishRequest {
133                channel: channel.to_string(),
134                data,
135            }),
136            ..Default::default()
137        };
138
139        self.send_void(cmd).await
140    }
141
142    /// Calls the named server RPC method.
143    /// # Errors
144    /// Returns an error if the client is disconnected or the command fails.
145    pub async fn rpc(&self, method: &str, data: Vec<u8>) -> Result<RpcResult, Error> {
146        let cmd = proto::Command {
147            rpc: Some(proto::RpcRequest {
148                method: method.to_string(),
149                data,
150            }),
151            ..Default::default()
152        };
153
154        self.send_extract(cmd, "RPC", |r| r.rpc).await
155    }
156
157    /// Sends asynchronous `data` to the server without publishing it.
158    /// # Errors
159    /// Returns an error if the client is disconnected or the command fails.
160    pub async fn send(&self, data: Vec<u8>) -> Result<(), Error> {
161        let cmd = proto::Command {
162            send: Some(proto::SendRequest { data }),
163            ..Default::default()
164        };
165
166        self.send_void(cmd).await
167    }
168
169    /// Refreshes the subscription token for `channel`.
170    /// # Errors
171    /// Returns an error if the client is disconnected or the command fails.
172    pub async fn sub_refresh(&self, channel: &str, token: &str) -> Result<SubRefreshResult, Error> {
173        let cmd = proto::Command {
174            sub_refresh: Some(proto::SubRefreshRequest {
175                channel: channel.to_string(),
176                token: token.to_string(),
177            }),
178            ..Default::default()
179        };
180
181        self.send_extract(cmd, "sub refresh", |r| r.sub_refresh)
182            .await
183    }
184
185    /// Stops the background connection loop and closes the active transport.
186    pub fn disconnect(&self) {
187        self.shutdown.cancel();
188    }
189
190    pub(super) async fn subscribe_internal(&self, channel: &str) -> Result<SubscribeResult, Error> {
191        let req = {
192            let mut inner = self.inner.write().await;
193            let sub = inner
194                .subscriptions
195                .get_mut(channel)
196                .ok_or_else(|| Error::subscription_not_found(channel))?;
197
198            sub.state = SubscriptionState::Subscribing;
199
200            sub.build_subscribe_request(false)
201        };
202
203        self.emit_subscription_event(channel, SubscriptionEvent::Subscribing)
204            .await;
205
206        let cmd = proto::Command {
207            subscribe: Some(req),
208            ..Default::default()
209        };
210
211        let reply = check_reply(self.send_command(cmd).await?)?;
212
213        let result = reply
214            .subscribe
215            .ok_or_else(|| Error::protocol("No subscribe result"))?;
216
217        let subscribe_result = SubscribeResult::from(result);
218
219        let subscribed = {
220            let mut inner = self.inner.write().await;
221
222            if let Some(sub) = inner.subscriptions.get_mut(channel) {
223                sub.on_subscribed(&subscribe_result);
224
225                true
226            } else {
227                false
228            }
229        };
230
231        if subscribed {
232            self.emit_subscription_event(
233                channel,
234                SubscriptionEvent::Subscribed(subscribe_result.clone()),
235            )
236            .await;
237        }
238
239        Ok(subscribe_result)
240    }
241
242    async fn send_extract<P, T>(
243        &self,
244        cmd: proto::Command,
245        label: &'static str,
246        extract: impl FnOnce(proto::Reply) -> Option<P>,
247    ) -> Result<T, Error>
248    where
249        P: Into<T>,
250    {
251        let reply = check_reply(self.send_command(cmd).await?)?;
252
253        extract(reply)
254            .ok_or_else(|| Error::protocol(format!("No {label} result")))
255            .map(Into::into)
256    }
257
258    async fn send_void(&self, cmd: proto::Command) -> Result<(), Error> {
259        check_reply(self.send_command(cmd).await?)?;
260
261        Ok(())
262    }
263
264    async fn send_command(&self, cmd: proto::Command) -> Result<proto::Reply, Error> {
265        let id = self
266            .inner
267            .read()
268            .await
269            .next_id
270            .fetch_add(1, Ordering::Relaxed);
271        let (reply_tx, reply_rx) = oneshot::channel();
272        let cmd_tx = {
273            let mut inner = self.inner.write().await;
274            let cmd_tx = inner.cmd_tx.clone().ok_or_else(Error::not_connected)?;
275
276            inner.pending.insert(id, reply_tx);
277
278            cmd_tx
279        };
280
281        if let Err(error) = cmd_tx.send(CommandRequest { id, cmd }).await {
282            self.inner.write().await.pending.remove(&id);
283            tracing::debug!(%error, "Command channel closed");
284
285            return Err(Error::not_connected());
286        }
287
288        match tokio::time::timeout(DEFAULT_TIMEOUT, reply_rx).await {
289            Ok(Ok(result)) => result,
290            Ok(Err(error)) => {
291                self.inner.write().await.pending.remove(&id);
292                tracing::debug!(%error, "Command channel closed");
293
294                Err(Error::connection_closed())
295            }
296            Err(error) => {
297                self.inner.write().await.pending.remove(&id);
298                tracing::debug!(?DEFAULT_TIMEOUT, %error, "Command timed out");
299
300                Err(Error::timeout())
301            }
302        }
303    }
304}
wowlab_centrifuge/client/commands.rs

wowlab_centrifuge/client/
commands.rs
