turbo_trace_server/
server.rs

1use std::{
2    net::{SocketAddr, SocketAddrV4, TcpListener, TcpStream},
3    sync::{Arc, Mutex},
4    thread::spawn,
5};
6
7use anyhow::{Result, bail};
8use serde::{Deserialize, Serialize};
9use tungstenite::{Message, accept};
10
11use crate::{
12    store::SpanId,
13    store_container::StoreContainer,
14    timestamp::Timestamp,
15    u64_string,
16    viewer::{SortMode, Update, ViewLineUpdate, ViewMode, Viewer},
17};
18
19#[derive(Serialize, Debug)]
20#[serde(tag = "type")]
21#[serde(rename_all = "kebab-case")]
22pub enum ServerToClientMessage {
23    ViewLine {
24        #[serde(flatten)]
25        update: ViewLineUpdate,
26    },
27    ViewLinesCount {
28        count: usize,
29        max: u64,
30    },
31    #[serde(rename_all = "camelCase")]
32    QueryResult {
33        #[serde(with = "u64_string")]
34        id: SpanId,
35        is_graph: bool,
36        start: Timestamp,
37        end: Timestamp,
38        duration: Timestamp,
39        cpu: Timestamp,
40        allocations: u64,
41        deallocations: u64,
42        allocation_count: u64,
43        persistent_allocations: u64,
44        args: Vec<(String, String)>,
45        path: Vec<String>,
46        memory_samples: Vec<u64>,
47        memory_pressure_samples: Vec<u8>,
48    },
49}
50
51#[derive(Deserialize, Debug)]
52#[serde(tag = "type")]
53#[serde(rename_all = "kebab-case")]
54pub enum ClientToServerMessage {
55    #[serde(rename_all = "camelCase")]
56    ViewRect {
57        view_rect: ViewRect,
58    },
59    ViewMode {
60        #[serde(with = "u64_string")]
61        id: SpanId,
62        mode: String,
63        inherit: bool,
64    },
65    ResetViewMode {
66        #[serde(with = "u64_string")]
67        id: SpanId,
68    },
69    Query {
70        #[serde(with = "u64_string")]
71        id: SpanId,
72    },
73    Ack,
74    CheckForMoreData,
75}
76
77#[derive(Deserialize, Debug)]
78pub struct Filter {
79    pub op: Op,
80    pub value: u64,
81}
82
83#[derive(Deserialize, Debug)]
84#[serde(rename_all = "snake_case")]
85pub enum Op {
86    Gt,
87    Lt,
88}
89
90#[derive(Deserialize, Debug)]
91#[serde(rename_all = "camelCase")]
92pub struct ViewRect {
93    pub x: u64,
94    pub y: u64,
95    pub width: u64,
96    pub height: u64,
97    pub horizontal_pixels: u64,
98    pub query: String,
99    pub view_mode: String,
100    pub value_mode: String,
101    pub value_filter: Option<Filter>,
102    pub count_filter: Option<Filter>,
103}
104
105struct ConnectionState {
106    store: Arc<StoreContainer>,
107    viewer: Viewer,
108    view_rect: ViewRect,
109    last_update_generation: usize,
110}
111
112pub fn serve(store: Arc<StoreContainer>, port: u16) {
113    let server = TcpListener::bind(SocketAddr::V4(SocketAddrV4::new(
114        std::net::Ipv4Addr::new(127, 0, 0, 1),
115        port,
116    )))
117    .unwrap();
118    for stream in server.incoming() {
119        let store = store.clone();
120
121        spawn(move || {
122            let websocket = accept(stream.unwrap()).unwrap();
123            if let Err(err) = handle_connection(websocket, store) {
124                eprintln!("Error: {err:?}");
125            }
126        });
127    }
128}
129
130fn handle_connection(
131    mut websocket: tungstenite::WebSocket<TcpStream>,
132    store: Arc<StoreContainer>,
133) -> Result<()> {
134    let state = Arc::new(Mutex::new(ConnectionState {
135        store,
136        viewer: Viewer::new(),
137        view_rect: ViewRect {
138            x: 0,
139            y: 0,
140            width: 1,
141            height: 1,
142            horizontal_pixels: 1,
143            query: String::new(),
144            view_mode: "aggregated".to_string(),
145            value_mode: "duration".to_string(),
146            count_filter: None,
147            value_filter: None,
148        },
149        last_update_generation: 0,
150    }));
151    let mut update_skipped = false;
152    let mut ready_for_update = true;
153
154    fn send_update(
155        websocket: &mut tungstenite::WebSocket<TcpStream>,
156        state: &mut ConnectionState,
157        force_send: bool,
158        ready_for_update: &mut bool,
159        update_skipped: &mut bool,
160    ) -> Result<()> {
161        if !*ready_for_update {
162            if force_send {
163                *update_skipped = true;
164            }
165            return Ok(());
166        }
167        let store = state.store.read();
168        if !force_send && state.last_update_generation == store.generation() {
169            return Ok(());
170        }
171        state.last_update_generation = store.generation();
172        let Update {
173            lines: updates,
174            max,
175        } = state.viewer.compute_update(&store, &state.view_rect);
176        let count = updates.len();
177        for update in updates {
178            let message = ServerToClientMessage::ViewLine { update };
179            let message = serde_json::to_string(&message).unwrap();
180            websocket.send(Message::Text(message))?;
181        }
182        let message = ServerToClientMessage::ViewLinesCount { count, max };
183        let message = serde_json::to_string(&message).unwrap();
184        websocket.send(Message::Text(message))?;
185        *ready_for_update = false;
186        Ok(())
187    }
188    loop {
189        match websocket.read()? {
190            Message::Frame(_frame) => {}
191            Message::Text(text) => {
192                let message: ClientToServerMessage = serde_json::from_str(&text)?;
193                let mut state = state.lock().unwrap();
194                match message {
195                    ClientToServerMessage::CheckForMoreData => {
196                        send_update(
197                            &mut websocket,
198                            &mut state,
199                            false,
200                            &mut ready_for_update,
201                            &mut update_skipped,
202                        )?;
203                    }
204                    ClientToServerMessage::ViewRect { view_rect } => {
205                        state.view_rect = view_rect;
206                        send_update(
207                            &mut websocket,
208                            &mut state,
209                            true,
210                            &mut ready_for_update,
211                            &mut update_skipped,
212                        )?;
213                    }
214                    ClientToServerMessage::ViewMode { id, mode, inherit } => {
215                        let (mode, sort_mode) =
216                            if let Some(mode) = mode.strip_suffix("-sorted-by-name") {
217                                (mode, SortMode::Name)
218                            } else if let Some(mode) = mode.strip_suffix("-sorted-by-value") {
219                                (mode, SortMode::Value)
220                            } else if let Some(mode) = mode.strip_suffix("-sorted") {
221                                (mode, SortMode::Value)
222                            } else {
223                                (mode.as_str(), SortMode::ExecutionOrder)
224                            };
225                        match mode {
226                            "raw-spans" => {
227                                state.viewer.set_view_mode(
228                                    id,
229                                    Some((ViewMode::RawSpans { sort_mode }, inherit)),
230                                );
231                            }
232                            "aggregated" => {
233                                state.viewer.set_view_mode(
234                                    id,
235                                    Some((ViewMode::Aggregated { sort_mode }, inherit)),
236                                );
237                            }
238                            "bottom-up" => {
239                                state.viewer.set_view_mode(
240                                    id,
241                                    Some((ViewMode::BottomUp { sort_mode }, inherit)),
242                                );
243                            }
244                            "aggregated-bottom-up" => {
245                                state.viewer.set_view_mode(
246                                    id,
247                                    Some((ViewMode::AggregatedBottomUp { sort_mode }, inherit)),
248                                );
249                            }
250                            _ => {
251                                bail!("unknown view mode: {}", mode)
252                            }
253                        }
254                        send_update(
255                            &mut websocket,
256                            &mut state,
257                            true,
258                            &mut ready_for_update,
259                            &mut update_skipped,
260                        )?;
261                    }
262                    ClientToServerMessage::ResetViewMode { id } => {
263                        state.viewer.set_view_mode(id, None);
264                        send_update(
265                            &mut websocket,
266                            &mut state,
267                            true,
268                            &mut ready_for_update,
269                            &mut update_skipped,
270                        )?;
271                    }
272                    ClientToServerMessage::Query { id } => {
273                        let message = {
274                            let store = state.store.read();
275                            if let Some((span, is_graph)) = store.span(id) {
276                                let root_start = store.root_span().start();
277                                let span_start = span.start() - root_start;
278                                let span_end = span.end() - root_start;
279                                let duration = span.corrected_total_time();
280                                let cpu = span.total_time();
281                                let allocations = span.total_allocations();
282                                let deallocations = span.total_deallocations();
283                                let allocation_count = span.total_allocation_count();
284                                let persistent_allocations = span.total_persistent_allocations();
285                                let args = span
286                                    .args()
287                                    .map(|(k, v)| (k.to_string(), v.to_string()))
288                                    .collect();
289                                let mut path = Vec::new();
290                                let mut current = span;
291                                while let Some(parent) = current.parent() {
292                                    path.push(parent.nice_name().1.to_string());
293                                    current = parent;
294                                }
295                                path.reverse();
296                                let memory_samples =
297                                    store.memory_samples_for_range(span.start(), span.end());
298                                let memory_pressure_samples = store
299                                    .memory_pressure_samples_for_range(span.start(), span.end());
300                                ServerToClientMessage::QueryResult {
301                                    id,
302                                    is_graph,
303                                    start: span_start,
304                                    end: span_end,
305                                    duration,
306                                    cpu,
307                                    allocations,
308                                    deallocations,
309                                    allocation_count,
310                                    persistent_allocations,
311                                    args,
312                                    path,
313                                    memory_samples,
314                                    memory_pressure_samples,
315                                }
316                            } else {
317                                ServerToClientMessage::QueryResult {
318                                    id,
319                                    is_graph: false,
320                                    start: Timestamp::ZERO,
321                                    end: Timestamp::ZERO,
322                                    duration: Timestamp::ZERO,
323                                    cpu: Timestamp::ZERO,
324                                    allocations: 0,
325                                    deallocations: 0,
326                                    allocation_count: 0,
327                                    persistent_allocations: 0,
328                                    args: Vec::new(),
329                                    path: Vec::new(),
330                                    memory_samples: Vec::new(),
331                                    memory_pressure_samples: Vec::new(),
332                                }
333                            }
334                        };
335                        let message = serde_json::to_string(&message).unwrap();
336                        websocket.send(Message::Text(message))?;
337                        send_update(
338                            &mut websocket,
339                            &mut state,
340                            true,
341                            &mut ready_for_update,
342                            &mut update_skipped,
343                        )?;
344
345                        continue;
346                    }
347                    ClientToServerMessage::Ack => {
348                        ready_for_update = true;
349                        if update_skipped {
350                            update_skipped = false;
351                            send_update(
352                                &mut websocket,
353                                &mut state,
354                                true,
355                                &mut ready_for_update,
356                                &mut update_skipped,
357                            )?;
358                        }
359                    }
360                }
361            }
362            Message::Binary(_) => {
363                // This doesn't happen
364            }
365            Message::Close(_) => {
366                return Ok(());
367            }
368            Message::Ping(d) => {
369                websocket.send(Message::Pong(d))?;
370            }
371            Message::Pong(_) => {
372                // thanks for the fish
373            }
374        }
375    }
376}
turbo_trace_server/server.rs

turbo_trace_server/
server.rs
