Skip to main content

wowlab_centrifuge/client/
push.rs

1use prost::Message;
2use wowlab_types::sim::FastMap;
3
4use super::{Client, DisconnectAdvice};
5use crate::{
6    error::Error,
7    proto,
8    subscription::SubscriptionState,
9    transport::Transport,
10    types::{ClientEvent, ClientInfo, Publication, SubscribeResult, SubscriptionEvent},
11};
12
13impl Client {
14    pub(super) async fn handle_reply(
15        &self,
16        reply: proto::Reply,
17        transport: &mut Transport,
18    ) -> Result<Option<DisconnectAdvice>, Error> {
19        if reply.id > 0 {
20            let tx = self.inner.write().await.pending.remove(&reply.id);
21
22            if let Some(tx) = tx {
23                let _ = tx.send(Ok(reply));
24            }
25
26            return Ok(None);
27        }
28
29        if let Some(push) = reply.push {
30            let disconnect = self.handle_push(push).await;
31
32            return Ok(disconnect);
33        }
34
35        let send_pong = self.inner.read().await.send_pong;
36
37        if send_pong {
38            let pong = proto::Command::default();
39            let data = pong.encode_length_delimited_to_vec();
40
41            transport.send_raw(data).await?;
42            tracing::trace!("Sent pong");
43        }
44
45        Ok(None)
46    }
47
48    pub(super) async fn build_recovery_subs(&self) -> FastMap<String, proto::SubscribeRequest> {
49        let inner = self.inner.read().await;
50
51        inner
52            .subscriptions
53            .iter()
54            .filter(|(_, sub)| sub.config.recoverable && sub.stream_position.is_some())
55            .map(|(channel, sub)| (channel.clone(), sub.build_subscribe_request(true)))
56            .collect()
57    }
58
59    pub(super) async fn process_server_subs(
60        &self,
61        subs: impl IntoIterator<Item = (String, proto::SubscribeResult)>,
62    ) {
63        for (channel, result) in subs {
64            let subscribe_result = SubscribeResult::from(result);
65            let subscribed = {
66                let mut inner = self.inner.write().await;
67
68                if let Some(sub) = inner.subscriptions.get_mut(&channel) {
69                    sub.on_subscribed(&subscribe_result);
70
71                    true
72                } else {
73                    false
74                }
75            };
76
77            if subscribed {
78                self.emit_subscription_event(
79                    &channel,
80                    SubscriptionEvent::Subscribed(subscribe_result),
81                )
82                .await;
83            }
84        }
85    }
86
87    pub(super) async fn resubscribe_all(&self) {
88        let channels: Vec<String> = {
89            let inner = self.inner.read().await;
90
91            inner
92                .subscriptions
93                .iter()
94                .filter(|(_, sub)| sub.state != SubscriptionState::Subscribed)
95                .map(|(channel, _)| channel.clone())
96                .collect()
97        };
98
99        if channels.is_empty() {
100            tracing::debug!("No subscriptions to resubscribe");
101
102            return;
103        }
104
105        let channel_count = channels.len();
106        let mut succeeded = 0_usize;
107        let mut failed = 0_usize;
108
109        for channel in channels {
110            match self.subscribe_internal(&channel).await {
111                Ok(_) => succeeded += 1,
112                Err(_) => failed += 1,
113            }
114        }
115
116        tracing::info!(
117            channel_count,
118            succeeded,
119            failed,
120            "Resubscription batch completed"
121        );
122    }
123
124    async fn handle_push(&self, push: proto::Push) -> Option<DisconnectAdvice> {
125        let channel = push.channel.clone();
126
127        if let Some(pub_) = push.r#pub {
128            self.handle_push_publication(&channel, pub_).await;
129        } else if let Some(info) = push.join.and_then(|j| j.info) {
130            self.emit_subscription_event(&channel, SubscriptionEvent::Join(ClientInfo::from(info)))
131                .await;
132        } else if let Some(info) = push.leave.and_then(|l| l.info) {
133            self.emit_subscription_event(
134                &channel,
135                SubscriptionEvent::Leave(ClientInfo::from(info)),
136            )
137            .await;
138        } else if let Some(unsub) = push.unsubscribe {
139            self.handle_push_unsubscribe(&channel, unsub).await;
140        } else if let Some(msg) = push.message {
141            self.emit_client_event(ClientEvent::Message(msg.data)).await;
142        } else if let Some(disconnect) = push.disconnect {
143            return Some(DisconnectAdvice {
144                code: disconnect.code,
145                reason: disconnect.reason,
146                reconnect: disconnect.reconnect,
147            });
148        }
149
150        None
151    }
152
153    async fn handle_push_publication(&self, channel: &str, pub_: proto::Publication) {
154        let publication = Publication::from(pub_);
155        let offset = publication.offset;
156
157        {
158            let mut inner = self.inner.write().await;
159
160            if let Some(sub) = inner.subscriptions.get_mut(channel) {
161                sub.update_position(offset);
162            }
163        }
164        self.emit_subscription_event(channel, SubscriptionEvent::Publication(publication))
165            .await;
166    }
167
168    async fn handle_push_unsubscribe(&self, channel: &str, unsub: proto::Unsubscribe) {
169        {
170            let mut inner = self.inner.write().await;
171
172            if let Some(sub) = inner.subscriptions.get_mut(channel) {
173                sub.state = SubscriptionState::Unsubscribed;
174            }
175        }
176        self.emit_subscription_event(
177            channel,
178            SubscriptionEvent::Unsubscribed {
179                code: unsub.code,
180                reason: unsub.reason,
181            },
182        )
183        .await;
184    }
185}