turbo_tasks/
message_queue.rs

1use std::{any::Any, collections::VecDeque, fmt::Display, sync::Arc, time::Duration};
2
3use dashmap::DashMap;
4use serde::Serialize;
5use tokio::sync::{Mutex, mpsc};
6
7pub trait CompilationEvent: Sync + Send + Any {
8    fn type_name(&self) -> &'static str;
9    fn severity(&self) -> Severity;
10    fn message(&self) -> String;
11    fn to_json(&self) -> String;
12}
13
14const MAX_QUEUE_SIZE: usize = 256;
15
16type ArcMx<T> = Arc<Mutex<T>>;
17type CompilationEventChannel = mpsc::Sender<Arc<dyn CompilationEvent>>;
18
19#[derive(Debug, Clone, Eq, PartialEq, Hash)]
20enum EventChannelType {
21    Global,
22    Type(String),
23}
24
25pub struct CompilationEventQueue {
26    event_history: ArcMx<VecDeque<Arc<dyn CompilationEvent>>>,
27    subscribers: Arc<DashMap<EventChannelType, Vec<CompilationEventChannel>>>,
28}
29
30impl Default for CompilationEventQueue {
31    fn default() -> Self {
32        let subscribers = DashMap::new();
33        subscribers.insert(
34            EventChannelType::Global,
35            Vec::<CompilationEventChannel>::new(),
36        );
37
38        Self {
39            event_history: Arc::new(Mutex::new(VecDeque::with_capacity(MAX_QUEUE_SIZE))),
40            subscribers: Arc::new(subscribers),
41        }
42    }
43}
44
45impl CompilationEventQueue {
46    pub fn send(
47        &self,
48        message: Arc<dyn CompilationEvent>,
49    ) -> Result<(), mpsc::error::SendError<Arc<dyn CompilationEvent>>> {
50        let event_history = self.event_history.clone();
51        let subscribers = self.subscribers.clone();
52        let message_clone = message.clone();
53
54        // Spawn a task to handle the async operations
55        tokio::spawn(async move {
56            // Store the message in history
57            let mut history = event_history.lock().await;
58            if history.len() >= MAX_QUEUE_SIZE {
59                history.pop_front();
60            }
61            history.push_back(message_clone.clone());
62
63            // Send to all active receivers of the same message type
64            if let Some(mut type_subscribers) = subscribers.get_mut(&EventChannelType::Type(
65                message_clone.type_name().to_owned(),
66            )) {
67                let mut removal_indices = Vec::new();
68                for (ix, sender) in type_subscribers.iter().enumerate() {
69                    if sender.send(message_clone.clone()).await.is_err() {
70                        removal_indices.push(ix);
71                    }
72                }
73
74                for ix in removal_indices.iter().rev() {
75                    type_subscribers.remove(*ix);
76                }
77            }
78
79            // Send to all global message subscribers
80            let mut all_channel = subscribers.get_mut(&EventChannelType::Global).unwrap();
81            let mut removal_indices = Vec::new();
82            for (ix, sender) in all_channel.iter_mut().enumerate() {
83                if sender.send(message_clone.clone()).await.is_err() {
84                    removal_indices.push(ix);
85                }
86            }
87
88            for ix in removal_indices.iter().rev() {
89                all_channel.remove(*ix);
90            }
91        });
92
93        Ok(())
94    }
95
96    pub fn subscribe(
97        &self,
98        event_types: Option<Vec<String>>,
99    ) -> mpsc::Receiver<Arc<dyn CompilationEvent>> {
100        let (tx, rx) = mpsc::channel(MAX_QUEUE_SIZE);
101        let subscribers = self.subscribers.clone();
102        let event_history = self.event_history.clone();
103        let tx_clone = tx.clone();
104
105        // Spawn a task to handle the async operations
106        tokio::spawn(async move {
107            // Store the sender
108            if let Some(event_types) = event_types {
109                for event_type in event_types.iter() {
110                    let mut type_subscribers = subscribers
111                        .entry(EventChannelType::Type(event_type.clone()))
112                        .or_default();
113                    type_subscribers.push(tx_clone.clone());
114                }
115
116                for event in event_history.lock().await.iter() {
117                    if event_types.contains(&event.type_name().to_string()) {
118                        let _ = tx_clone.send(event.clone()).await;
119                    }
120                }
121            } else {
122                let mut global_subscribers =
123                    subscribers.entry(EventChannelType::Global).or_default();
124                global_subscribers.push(tx_clone.clone());
125
126                for event in event_history.lock().await.iter() {
127                    let _ = tx_clone.send(event.clone()).await;
128                }
129            }
130        });
131
132        rx
133    }
134}
135
136#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, Serialize)]
137pub enum Severity {
138    Info,
139    Trace,
140    Warning,
141    Error,
142    Fatal,
143    Event,
144}
145
146impl Display for Severity {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        match self {
149            Severity::Info => write!(f, "INFO"),
150            Severity::Trace => write!(f, "TRACE"),
151            Severity::Warning => write!(f, "WARNING"),
152            Severity::Error => write!(f, "ERROR"),
153            Severity::Fatal => write!(f, "FATAL"),
154            Severity::Event => write!(f, "EVENT"),
155        }
156    }
157}
158
159#[derive(Debug, Clone, Serialize)]
160/// Compilation event that is used to log the duration of a task
161pub struct TimingEvent {
162    /// Message of the event without the timing information
163    ///
164    /// Example:
165    /// ```rust
166    /// use std::time::Duration;
167    /// use turbo_tasks::message_queue::{CompilationEvent, TimingEvent};
168    ///
169    /// let event = TimingEvent::new("Compiled successfully".to_string(), Duration::from_millis(100));
170    /// let message = event.message();
171    /// assert_eq!(message, "Compiled successfully in 100ms");
172    /// ```
173    pub message: String,
174    /// Duration in milliseconds
175    pub duration: Duration,
176}
177
178impl TimingEvent {
179    pub fn new(message: String, duration: Duration) -> Self {
180        Self { message, duration }
181    }
182}
183
184impl CompilationEvent for TimingEvent {
185    fn type_name(&self) -> &'static str {
186        "TimingEvent"
187    }
188
189    fn severity(&self) -> Severity {
190        Severity::Event
191    }
192
193    fn message(&self) -> String {
194        let duration_secs = self.duration.as_secs_f64();
195        let duration_string = if duration_secs > 120.0 {
196            format!("{:.1}min", duration_secs / 60.0)
197        } else if duration_secs > 40.0 {
198            format!("{duration_secs:.0}s")
199        } else if duration_secs > 2.0 {
200            format!("{duration_secs:.1}s")
201        } else {
202            format!("{:.0}ms", duration_secs * 1000.0)
203        };
204        format!("{} in {}", self.message, duration_string)
205    }
206
207    fn to_json(&self) -> String {
208        serde_json::to_string(self).unwrap()
209    }
210}
211
212#[derive(Debug, Clone, Serialize)]
213pub struct DiagnosticEvent {
214    pub message: String,
215    pub severity: Severity,
216}
217
218impl DiagnosticEvent {
219    pub fn new(severity: Severity, message: String) -> Self {
220        Self { message, severity }
221    }
222}
223
224impl CompilationEvent for DiagnosticEvent {
225    fn type_name(&self) -> &'static str {
226        "DiagnosticEvent"
227    }
228
229    fn severity(&self) -> Severity {
230        self.severity
231    }
232
233    fn message(&self) -> String {
234        self.message.clone()
235    }
236
237    fn to_json(&self) -> String {
238        serde_json::to_string(self).unwrap()
239    }
240}
241
242/// A generic trace event that carries a name, wall-clock timing, and arbitrary attributes.
243/// Forwarded as a `CompilationEvent` to the JS side for inclusion in `.next/trace`.
244#[derive(Debug, Clone, Serialize)]
245#[serde(rename_all = "camelCase")]
246pub struct TraceEvent {
247    pub name: &'static str,
248    pub start_time_ms: f64,
249    pub end_time_ms: f64,
250    /// Should be an array of key value pairs
251    pub attributes: serde_json::Value,
252}
253
254impl TraceEvent {
255    pub fn new(
256        name: &'static str,
257        start_time_ms: f64,
258        end_time_ms: f64,
259        attributes: serde_json::Value,
260    ) -> Self {
261        // basic sanity test
262        debug_assert!(matches!(attributes, serde_json::Value::Array(_)));
263        Self {
264            name,
265            start_time_ms,
266            end_time_ms,
267            attributes,
268        }
269    }
270}
271
272impl CompilationEvent for TraceEvent {
273    fn type_name(&self) -> &'static str {
274        "TraceEvent"
275    }
276
277    fn severity(&self) -> Severity {
278        Severity::Event
279    }
280
281    fn message(&self) -> String {
282        let duration_ms = self.end_time_ms - self.start_time_ms;
283        format!("{} in {:.0}ms", self.name, duration_ms)
284    }
285
286    fn to_json(&self) -> String {
287        serde_json::to_string(self).unwrap()
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn test_timing_event_string_formatting() {
297        let tests = vec![
298            (Duration::from_nanos(1588), "0ms"),
299            (Duration::from_nanos(1022616), "1ms"),
300            (Duration::from_millis(100), "100ms"),
301            (Duration::from_millis(1000), "1000ms"),
302            (Duration::from_millis(10000), "10.0s"),
303            (Duration::from_millis(20381), "20.4s"),
304            (Duration::from_secs(60), "60s"),
305            (Duration::from_secs(100), "100s"),
306            (Duration::from_secs(125), "2.1min"),
307        ];
308
309        for (duration, expected) in tests {
310            let event = TimingEvent::new("Compiled successfully".to_string(), duration);
311            assert_eq!(
312                event.message(),
313                format!("Compiled successfully in {expected}")
314            );
315        }
316    }
317}
turbo_tasks/message_queue.rs

turbo_tasks/
message_queue.rs
