rainier_websocket/
socket.rs

1//! The handle a handler holds — [`Socket`].
2
3use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
4use std::sync::Arc;
5
6use serde::Serialize;
7use tokio::sync::mpsc;
8
9use rainier_support::{Error, Result};
10
11use crate::message::Message;
12
13/// A unique id for one connection, for the life of the process.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
15pub struct SocketId(u64);
16
17impl SocketId {
18    /// The next id. Monotonic, so ordering by it is arrival order.
19    pub fn next() -> Self {
20        static NEXT: AtomicU64 = AtomicU64::new(1);
21        Self(NEXT.fetch_add(1, Ordering::Relaxed))
22    }
23
24    /// The id as a number, for logs and for keying a map.
25    pub fn get(self) -> u64 {
26        self.0
27    }
28}
29
30impl std::fmt::Display for SocketId {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "{}", self.0)
33    }
34}
35
36/// This process's share of a socket identity.
37///
38/// A [`SocketId`] is a counter, so replica A and replica B both have a socket
39/// `7`. That is fine while the only thing reading the number is the process
40/// that issued it, and wrong the moment a broadcast crosses between them: "do
41/// not send this to socket 7" would silence a different browser on every other
42/// replica.
43///
44/// So [`Socket::identity`] pairs the counter with this, which is generated once
45/// per process from the clock. Two processes agreeing on it would need to have
46/// started in the same microsecond.
47pub fn instance_id() -> &'static str {
48    static INSTANCE: std::sync::OnceLock<String> = std::sync::OnceLock::new();
49
50    INSTANCE.get_or_init(|| {
51        let micros = std::time::SystemTime::now()
52            .duration_since(std::time::UNIX_EPOCH)
53            .map(|since| since.as_micros())
54            .unwrap_or_default();
55
56        format!("{micros:x}")
57    })
58}
59
60/// The socket an identity names, if it is one of **this** process's.
61///
62/// `None` for an identity minted by another replica, which is the answer that
63/// makes a fan-out correct: that socket is not here, so there is nothing to
64/// skip and everyone here should be told.
65///
66/// ```
67/// use rainier_websocket::socket::{socket_from_identity, instance_id};
68///
69/// assert_eq!(socket_from_identity(&format!("{}.7", instance_id())).map(|id| id.get()), Some(7));
70/// assert!(socket_from_identity("some-other-replica.7").is_none());
71/// assert!(socket_from_identity("nonsense").is_none());
72/// ```
73pub fn socket_from_identity(identity: &str) -> Option<SocketId> {
74    let (instance, id) = identity.rsplit_once('.')?;
75
76    if instance != instance_id() {
77        return None;
78    }
79
80    id.parse().ok().map(SocketId)
81}
82
83/// One connected client.
84///
85/// Cheap to clone and safe to keep: sending goes through a channel, so a handle
86/// stored in a room registry does not hold the connection's task or its
87/// buffers. Sending to a socket that has gone away is an error, not a panic.
88///
89/// ```ignore
90/// socket.send("hello").await?;
91/// socket.send_json(&Update { count: 3 }).await?;
92/// socket.close().await;
93/// ```
94#[derive(Debug, Clone)]
95pub struct Socket {
96    id: SocketId,
97    path: Arc<str>,
98    params: Arc<Vec<(String, String)>>,
99    outbound: mpsc::UnboundedSender<Outbound>,
100    closed: Arc<AtomicBool>,
101}
102
103/// What the connection task pulls off the queue.
104#[derive(Debug)]
105pub enum Outbound {
106    /// Send this frame.
107    Send(Message),
108    /// Close the connection, with an optional reason.
109    Close(Option<String>),
110}
111
112impl Socket {
113    /// Build a handle over `outbound`.
114    ///
115    /// The transport owns the receiving half; a handler is only ever given
116    /// this side.
117    pub fn new(
118        id: SocketId,
119        path: impl Into<Arc<str>>,
120        params: Vec<(String, String)>,
121        outbound: mpsc::UnboundedSender<Outbound>,
122    ) -> Self {
123        Self {
124            id,
125            path: path.into(),
126            params: Arc::new(params),
127            outbound,
128            closed: Arc::new(AtomicBool::new(false)),
129        }
130    }
131
132    /// This connection's id.
133    pub fn id(&self) -> SocketId {
134        self.id
135    }
136
137    /// This connection's id, qualified so it is unique across replicas.
138    ///
139    /// **This is what to send the browser**, and what it should echo back in
140    /// `X-Socket-ID` so [`to_others`](../rainier_broadcast/struct.Broadcasting.html#method.event_except)
141    /// can skip it:
142    ///
143    /// ```ignore
144    /// async fn on_connect(&self, socket: &Socket) -> Result<()> {
145    ///     socket.send_json(&json!({ "socket_id": socket.identity() })).await
146    /// }
147    /// ```
148    ///
149    /// Sending the bare [`id`](Self::id) instead works on one replica and
150    /// fails quietly on two: every replica has a socket `7`, so "everyone
151    /// except 7" would silence an unrelated browser on each of the others.
152    pub fn identity(&self) -> String {
153        format!("{}.{}", instance_id(), self.id.get())
154    }
155
156    /// The path it connected to.
157    pub fn path(&self) -> &str {
158        &self.path
159    }
160
161    /// A route parameter captured from the path — `rooms/{room}` gives
162    /// `room`.
163    pub fn param(&self, name: &str) -> Option<&str> {
164        self.params.iter().find(|(key, _)| key == name).map(|(_, value)| value.as_str())
165    }
166
167    /// A route parameter, parsed.
168    ///
169    /// A **bad request** when it will not parse: the value came from the URL
170    /// the client asked for.
171    pub fn parse_param<T>(&self, name: &str) -> Result<T>
172    where
173        T: std::str::FromStr,
174        T::Err: std::fmt::Display,
175    {
176        let raw = self
177            .param(name)
178            .ok_or_else(|| Error::internal(format!("the route captured no `{name}`")))?;
179
180        raw.parse::<T>()
181            .map_err(|e| Error::bad_request(format!("`{name}` in the path is invalid: {e}")))
182    }
183
184    /// Every captured parameter.
185    pub fn params(&self) -> &[(String, String)] {
186        &self.params
187    }
188
189    /// Send a frame.
190    ///
191    /// Queued rather than written: a handler that awaited the socket would
192    /// block on a slow client, and one slow client must not hold up the task
193    /// reading everyone else's messages.
194    pub fn send(&self, message: impl Into<Message>) -> Result<()> {
195        self.push(Outbound::Send(message.into()))
196    }
197
198    /// Send a value as a JSON text frame.
199    pub fn send_json(&self, value: &impl Serialize) -> Result<()> {
200        self.send(Message::json(value)?)
201    }
202
203    /// Close the connection.
204    ///
205    /// Idempotent: closing twice is not an error, because two halves of a
206    /// handler deciding to close at once is normal.
207    pub fn close(&self) -> Result<()> {
208        if self.closed.swap(true, Ordering::SeqCst) {
209            return Ok(());
210        }
211        self.push(Outbound::Close(None))
212    }
213
214    /// Close with a reason the peer will see.
215    pub fn close_with(&self, reason: impl Into<String>) -> Result<()> {
216        if self.closed.swap(true, Ordering::SeqCst) {
217            return Ok(());
218        }
219        self.push(Outbound::Close(Some(reason.into())))
220    }
221
222    /// Whether this handle has been asked to close, or the connection has
223    /// gone.
224    pub fn is_closed(&self) -> bool {
225        self.closed.load(Ordering::SeqCst) || self.outbound.is_closed()
226    }
227
228    fn push(&self, outbound: Outbound) -> Result<()> {
229        self.outbound.send(outbound).map_err(|_| {
230            // The peer disconnected. Routine — a registry holding stale
231            // handles finds out this way — so it is an error to handle, not a
232            // panic and not a silent success.
233            Error::internal(format!("socket {} has gone away", self.id))
234        })
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    fn socket() -> (Socket, mpsc::UnboundedReceiver<Outbound>) {
243        let (tx, rx) = mpsc::unbounded_channel();
244        (Socket::new(SocketId::next(), "/ws/rooms/7", vec![("room".into(), "7".into())], tx), rx)
245    }
246
247    #[test]
248    fn an_identity_is_the_process_and_the_socket() {
249        let (first, _rx) = socket();
250        let (second, _rx2) = socket();
251
252        assert_ne!(first.identity(), second.identity(), "two sockets, two identities");
253
254        // Both halves are there, and the process half is the same for both:
255        // it is what stops replica B skipping *its* socket 7 when replica A
256        // said "everyone except 7".
257        assert!(first.identity().ends_with(&format!(".{}", first.id().get())));
258        assert_eq!(
259            first.identity().split('.').next(),
260            second.identity().split('.').next(),
261            "one process, one instance id"
262        );
263    }
264
265    #[test]
266    fn a_send_is_queued_not_written() {
267        let (socket, mut rx) = socket();
268
269        socket.send("hello").unwrap();
270
271        match rx.try_recv().unwrap() {
272            Outbound::Send(Message::Text(text)) => assert_eq!(text, "hello"),
273            other => panic!("{other:?}"),
274        }
275    }
276
277    #[test]
278    fn route_parameters_come_through() {
279        let (socket, _rx) = socket();
280
281        assert_eq!(socket.param("room"), Some("7"));
282        assert_eq!(socket.parse_param::<u64>("room").unwrap(), 7);
283        assert_eq!(socket.param("nope"), None);
284    }
285
286    #[test]
287    fn a_parameter_that_will_not_parse_is_a_bad_request() {
288        let (tx, _rx) = mpsc::unbounded_channel();
289        let socket =
290            Socket::new(SocketId::next(), "/ws/rooms/x", vec![("room".into(), "x".into())], tx);
291
292        assert_eq!(socket.parse_param::<u64>("room").unwrap_err().status(), 400);
293    }
294
295    #[test]
296    fn closing_twice_is_not_an_error() {
297        let (socket, mut rx) = socket();
298
299        socket.close().unwrap();
300        socket.close().unwrap();
301
302        assert!(matches!(rx.try_recv(), Ok(Outbound::Close(None))));
303        assert!(rx.try_recv().is_err(), "only one close should be queued");
304        assert!(socket.is_closed());
305    }
306
307    #[test]
308    fn sending_to_a_socket_that_went_away_is_an_error_not_a_panic() {
309        let (socket, rx) = socket();
310        drop(rx);
311
312        let err = socket.send("hello").unwrap_err();
313        assert!(err.message().contains("gone away"), "{}", err.message());
314        assert!(socket.is_closed());
315    }
316
317    #[test]
318    fn ids_are_unique_and_ordered() {
319        let first = SocketId::next();
320        let second = SocketId::next();
321
322        assert!(second > first);
323    }
324}
rainier_websocket/socket.rs

rainier_websocket/
socket.rs
