sie_sdk/client/
watch.rs

1//! Live status over WebSocket.
2//!
3//! A gateway broadcasts cluster-wide status on `/ws/cluster-status`; a worker broadcasts
4//! its own on `/ws/status`. Which one to open is decided by [`WatchMode`], and `Auto` asks
5//! `/health` first.
6
7use futures_util::StreamExt;
8use tokio_tungstenite::tungstenite::client::IntoClientRequest;
9use tokio_tungstenite::tungstenite::http::HeaderValue;
10use tokio_tungstenite::tungstenite::protocol::Message;
11use url::Url;
12
13use crate::client::{Client, stream::ChunkStream};
14use crate::error::{Error, Result, TransportErrorKind};
15use crate::types::{ClusterStatusMessage, StatusMessage, WorkerStatusMessage};
16
17/// Which status endpoint to watch.
18#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
19pub enum WatchMode {
20    /// Probe `/health` and pick the endpoint that matches.
21    #[default]
22    Auto,
23    /// A gateway's cluster-wide broadcast.
24    Cluster,
25    /// A single worker's own broadcast.
26    Worker,
27}
28
29impl WatchMode {
30    fn path(self) -> &'static str {
31        match self {
32            Self::Cluster | Self::Auto => "/ws/cluster-status",
33            Self::Worker => "/ws/status",
34        }
35    }
36}
37
38impl Client {
39    /// Stream status broadcasts until the connection closes.
40    ///
41    /// With [`WatchMode::Auto`], `/health` is probed once to decide which endpoint to open.
42    pub async fn watch(&self, mode: WatchMode) -> Result<ChunkStream<StatusMessage>> {
43        let resolved = match mode {
44            WatchMode::Auto => {
45                // A worker's /health names itself; anything else is treated as a gateway.
46                match self.health().await {
47                    Ok(health) if health.kind == "worker" => WatchMode::Worker,
48                    _ => WatchMode::Cluster,
49                }
50            }
51            explicit => explicit,
52        };
53
54        let url = self.websocket_url(resolved.path())?;
55        let mut request = url.as_str().into_client_request().map_err(|err| {
56            Error::invalid(format!("could not build the WebSocket request: {err}"))
57        })?;
58
59        // The WebSocket handshake is an ordinary HTTP request, so it carries the same
60        // credentials, under the same origin rule.
61        if let Some(authorization) = self.authorization_header() {
62            request
63                .headers_mut()
64                .insert(reqwest::header::AUTHORIZATION.as_str(), authorization);
65        }
66        if self.websocket_matches_base_origin(&url) {
67            for (name, value) in self.edge_headers() {
68                if let (Ok(name), Ok(value)) = (
69                    tokio_tungstenite::tungstenite::http::HeaderName::from_bytes(
70                        name.as_str().as_bytes(),
71                    ),
72                    HeaderValue::from_bytes(value.as_bytes()),
73                ) {
74                    request.headers_mut().insert(name, value);
75                }
76            }
77        }
78
79        let (socket, response) =
80            tokio_tungstenite::connect_async(request)
81                .await
82                .map_err(|error| match &error {
83                    tokio_tungstenite::tungstenite::Error::Http(response) => Error::Request {
84                        message: format!("WebSocket connection failed: {}", response.status()),
85                        code: None,
86                        status: response.status().as_u16(),
87                        request: None,
88                    },
89                    _ => Error::connection(
90                        TransportErrorKind::Connect,
91                        format!("could not open {url}: {error}"),
92                        error,
93                    ),
94                })?;
95        drop(response);
96
97        Ok(Box::pin(async_stream::try_stream! {
98            let mut socket = socket;
99            while let Some(frame) = socket.next().await {
100                let frame = frame.map_err(|error| {
101                    Error::connection(
102                        TransportErrorKind::MidFlight,
103                        format!("status stream closed: {error}"),
104                        error,
105                    )
106                })?;
107                let payload = match frame {
108                    Message::Text(text) => text.to_string(),
109                    Message::Binary(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
110                    // Ping and pong are answered by the library; a close ends the stream.
111                    Message::Close(_) => return,
112                    _ => continue,
113                };
114                yield decode_status(&payload, resolved)?;
115            }
116        }))
117    }
118
119    /// The `ws`/`wss` counterpart of a path on this client's base URL.
120    fn websocket_url(&self, path: &str) -> Result<Url> {
121        let mut url = self.url(path)?;
122        let scheme = match url.scheme() {
123            "https" => "wss",
124            "http" => "ws",
125            other => {
126                return Err(Error::invalid(format!(
127                    "base_url scheme {other:?} has no WebSocket counterpart"
128                )));
129            }
130        };
131        url.set_scheme(scheme)
132            .map_err(|()| Error::invalid("could not derive the WebSocket URL"))?;
133        Ok(url)
134    }
135}
136
137/// The endpoint decides which shape a payload is: the two share no required field.
138fn decode_status(payload: &str, mode: WatchMode) -> Result<StatusMessage> {
139    match mode {
140        WatchMode::Worker => serde_json::from_str::<WorkerStatusMessage>(payload)
141            .map(|message| StatusMessage::Worker(Box::new(message))),
142        WatchMode::Cluster | WatchMode::Auto => {
143            serde_json::from_str::<ClusterStatusMessage>(payload)
144                .map(|message| StatusMessage::Cluster(Box::new(message)))
145        }
146    }
147    .map_err(|err| Error::decode(format!("malformed status message: {err}")))
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[test]
155    fn websocket_urls_follow_the_base_url_scheme() {
156        let secure = Client::new("https://sie.example.com").unwrap();
157        assert_eq!(
158            secure.websocket_url("/ws/cluster-status").unwrap().as_str(),
159            "wss://sie.example.com/ws/cluster-status"
160        );
161
162        let plain = Client::new("http://localhost:8080").unwrap();
163        assert_eq!(
164            plain.websocket_url("/ws/status").unwrap().as_str(),
165            "ws://localhost:8080/ws/status"
166        );
167    }
168
169    #[test]
170    fn modes_map_to_their_endpoints() {
171        assert_eq!(WatchMode::Cluster.path(), "/ws/cluster-status");
172        assert_eq!(WatchMode::Worker.path(), "/ws/status");
173        assert_eq!(WatchMode::default(), WatchMode::Auto);
174    }
175
176    #[test]
177    fn payloads_decode_into_the_shape_the_endpoint_promises() {
178        let worker = decode_status(
179            r#"{"timestamp": 1.0, "ready": true, "name": "w-1", "machine_profile": "l4",
180                "saturated": false, "gpus": [{"device": "cuda:0", "utilization_pct": 42}]}"#,
181            WatchMode::Worker,
182        )
183        .unwrap();
184        let worker = worker.worker().expect("a worker message");
185        assert_eq!(worker.name, "w-1");
186        assert_eq!(worker.gpus[0].utilization_pct, 42);
187        assert!(worker.pool_name.is_empty());
188
189        let cluster = decode_status(
190            r#"{"timestamp": 2.0, "cluster": {"worker_count": 3, "gpu_count": 6,
191                "models_loaded": 2, "total_qps": 1.5}, "workers": [], "models": []}"#,
192            WatchMode::Cluster,
193        )
194        .unwrap();
195        assert_eq!(cluster.cluster().unwrap().cluster.worker_count, 3);
196        assert!(cluster.worker().is_none());
197    }
198
199    #[test]
200    fn a_malformed_payload_is_a_decode_error() {
201        let err = decode_status("{not json", WatchMode::Cluster).unwrap_err();
202        assert!(
203            err.to_string().contains("malformed status message"),
204            "{err}"
205        );
206    }
207
208    #[tokio::test]
209    async fn watching_an_unreachable_server_fails_rather_than_hanging() {
210        let client = Client::new("http://127.0.0.1:1").unwrap();
211        let Err(err) = client.watch(WatchMode::Cluster).await else {
212            panic!("connecting to a closed port must fail");
213        };
214        assert!(matches!(err, Error::Connection { .. }), "{err:?}");
215    }
216}
sie_sdk/client/watch.rs

sie_sdk/client/
watch.rs
