autoconnect_common/
registry.rs

1use dashmap::DashMap;
2use futures::channel::mpsc;
3use uuid::Uuid;
4
5use autopush_common::errors::{ApcErrorKind, Result};
6use autopush_common::notification::Notification;
7
8use crate::protocol::ServerNotification;
9
10/// Default capacity for bounded notification channels per client.
11const DEFAULT_CHANNEL_CAPACITY: usize = 128;
12
13/// A connected Websocket client.
14#[derive(Debug)]
15struct RegisteredClient {
16    /// The user agent's unique ID.
17    #[allow(dead_code)]
18    pub uaid: Uuid,
19    /// The local ID, used to potentially distinquish multiple UAID connections.
20    pub uid: Uuid,
21    /// The inbound channel for delivery of locally routed Push Notifications
22    pub tx: mpsc::Sender<ServerNotification>,
23}
24
25/// Contains a mapping of UAID to the associated RegisteredClient.
26pub struct ClientRegistry {
27    clients: DashMap<Uuid, RegisteredClient>,
28    /// Maximum number of buffered notifications per client before backpressure.
29    channel_capacity: usize,
30}
31
32impl Default for ClientRegistry {
33    fn default() -> Self {
34        Self {
35            clients: DashMap::new(),
36            channel_capacity: DEFAULT_CHANNEL_CAPACITY,
37        }
38    }
39}
40
41impl ClientRegistry {
42    /// Create a new ClientRegistry with a configurable channel capacity.
43    pub fn with_channel_capacity(channel_capacity: usize) -> Self {
44        Self {
45            clients: DashMap::new(),
46            channel_capacity,
47        }
48    }
49
50    /// Informs this server that a new `client` has connected
51    ///
52    /// For now just registers internal state by keeping track of the `client`,
53    /// namely its channel to send notifications back.
54    pub fn connect(&self, uaid: Uuid, uid: Uuid) -> mpsc::Receiver<ServerNotification> {
55        trace!("ClientRegistry::connect");
56        let (tx, snotif_stream) = mpsc::channel(self.channel_capacity);
57        let client = RegisteredClient { uaid, uid, tx };
58        if let Some(old_client) = self.clients.insert(uaid, client) {
59            // Drop existing connection
60            let result = old_client
61                .tx
62                .clone()
63                .try_send(ServerNotification::Disconnect);
64            if result.is_ok() {
65                debug!("ClientRegistry::connect Ghosting client, new one wants to connect");
66            }
67        }
68        snotif_stream
69    }
70
71    /// A notification has come for the uaid
72    pub fn notify(&self, uaid: Uuid, notif: Notification) -> Result<()> {
73        trace!("ClientRegistry::notify");
74        let Some(client) = self.clients.get(&uaid) else {
75            return Err(ApcErrorKind::ClientNotConnected.into());
76        };
77        debug!("ClientRegistry::notify Found a client to deliver a notification to");
78        match client
79            .tx
80            .clone()
81            .try_send(ServerNotification::Notification(notif))
82        {
83            Ok(()) => {
84                debug!("ClientRegistry::notify Dropped notification in queue");
85                Ok(())
86            }
87            Err(e) if e.is_full() => Err(ApcErrorKind::ChannelFull.into()),
88            Err(_) => Err(ApcErrorKind::ClientNotConnected.into()),
89        }
90    }
91
92    /// A check for notification command has come for the uaid
93    pub fn check_storage(&self, uaid: Uuid) -> Result<()> {
94        trace!("ClientRegistry::check_storage");
95        let Some(client) = self.clients.get(&uaid) else {
96            return Err(ApcErrorKind::ClientNotConnected.into());
97        };
98        match client.tx.clone().try_send(ServerNotification::CheckStorage) {
99            Ok(()) => {
100                debug!("ClientRegistry::check_storage Told client to check storage");
101                Ok(())
102            }
103            Err(e) if e.is_full() => Err(ApcErrorKind::ChannelFull.into()),
104            Err(_) => Err(ApcErrorKind::ClientNotConnected.into()),
105        }
106    }
107
108    /// The client specified by `uaid` has disconnected.
109    pub fn disconnect(&self, uaid: &Uuid, uid: &Uuid) -> Result<()> {
110        trace!("ClientRegistry::disconnect");
111        let client_exists = self
112            .clients
113            .get(uaid)
114            .is_some_and(|client| client.uid == *uid);
115        if client_exists {
116            self.clients
117                .remove(uaid)
118                .ok_or_else(|| ApcErrorKind::GeneralError("Couldn't remove client".into()))?;
119            return Ok(());
120        }
121        Err(ApcErrorKind::ClientNotConnected.into())
122    }
123
124    pub fn count(&self) -> usize {
125        self.clients.len()
126    }
127}