turbo_tasks/
message_queue.rs

1use std::{
2    any::Any,
3    collections::VecDeque,
4    fmt::Display,
5    sync::Arc,
6    time::{Duration, SystemTime},
7};
8
9use dashmap::DashMap;
10use parking_lot::Mutex;
11use serde::Serialize;
12use tokio::sync::{Mutex as TokioMutex, Notify, mpsc};
13
14pub trait CompilationEvent: Sync + Send + Any {
15    fn type_name(&self) -> &'static str;
16    fn severity(&self) -> Severity;
17    fn message(&self) -> String;
18    fn to_json(&self) -> String;
19}
20
21const MAX_QUEUE_SIZE: usize = 256;
22
23type ArcMx<T> = Arc<TokioMutex<T>>;
24type CompilationEventChannel = mpsc::Sender<Arc<dyn CompilationEvent>>;
25
26#[derive(Debug, Clone, Eq, PartialEq, Hash)]
27enum EventChannelType {
28    Global,
29    Type(String),
30}
31
32struct DeliveryStateInner {
33    closed: bool,
34    pending: usize,
35}
36
37struct DeliveryState {
38    inner: Mutex<DeliveryStateInner>,
39    idle: Notify,
40}
41
42pub struct CompilationEventQueue {
43    event_history: ArcMx<VecDeque<Arc<dyn CompilationEvent>>>,
44    subscribers: Arc<DashMap<EventChannelType, Vec<CompilationEventChannel>>>,
45    delivery_state: Arc<DeliveryState>,
46}
47
48impl Default for CompilationEventQueue {
49    fn default() -> Self {
50        let subscribers = DashMap::new();
51        subscribers.insert(
52            EventChannelType::Global,
53            Vec::<CompilationEventChannel>::new(),
54        );
55
56        Self {
57            event_history: Arc::new(TokioMutex::new(VecDeque::with_capacity(MAX_QUEUE_SIZE))),
58            subscribers: Arc::new(subscribers),
59            delivery_state: Arc::new(DeliveryState {
60                inner: Mutex::new(DeliveryStateInner {
61                    closed: false,
62                    pending: 0,
63                }),
64                idle: Notify::new(),
65            }),
66        }
67    }
68}
69
70impl CompilationEventQueue {
71    pub fn send(
72        &self,
73        message: Arc<dyn CompilationEvent>,
74    ) -> Result<(), mpsc::error::SendError<Arc<dyn CompilationEvent>>> {
75        let event_history = self.event_history.clone();
76        let subscribers = self.subscribers.clone();
77        let message_clone = message.clone();
78        let delivery_state = self.delivery_state.clone();
79        let deliver = {
80            let mut state = delivery_state.inner.lock();
81            if state.closed {
82                false
83            } else {
84                state.pending += 1;
85                true
86            }
87        };
88
89        // Spawn a task to handle the async operations
90        tokio::spawn(async move {
91            // Store the message in history
92            let mut history = event_history.lock().await;
93            if history.len() >= MAX_QUEUE_SIZE {
94                history.pop_front();
95            }
96            history.push_back(message_clone.clone());
97            drop(history);
98
99            if deliver {
100                // Send to all active receivers of the same message type
101                if let Some(mut type_subscribers) = subscribers.get_mut(&EventChannelType::Type(
102                    message_clone.type_name().to_owned(),
103                )) {
104                    let mut removal_indices = Vec::new();
105                    for (ix, sender) in type_subscribers.iter().enumerate() {
106                        if sender.send(message_clone.clone()).await.is_err() {
107                            removal_indices.push(ix);
108                        }
109                    }
110
111                    for ix in removal_indices.iter().rev() {
112                        type_subscribers.remove(*ix);
113                    }
114                }
115
116                // Send to all global message subscribers
117                if let Some(mut all_channel) = subscribers.get_mut(&EventChannelType::Global) {
118                    let mut removal_indices = Vec::new();
119                    for (ix, sender) in all_channel.iter_mut().enumerate() {
120                        if sender.send(message_clone.clone()).await.is_err() {
121                            removal_indices.push(ix);
122                        }
123                    }
124
125                    for ix in removal_indices.iter().rev() {
126                        all_channel.remove(*ix);
127                    }
128                }
129            }
130
131            if deliver {
132                let mut state = delivery_state.inner.lock();
133                state.pending -= 1;
134                if state.pending == 0 {
135                    delivery_state.idle.notify_one();
136                }
137            }
138        });
139
140        Ok(())
141    }
142
143    /// Waits until all events sent so far have been delivered to subscribers, then closes all
144    /// subscriber channels. Subscriptions drain their remaining events and then end. Events sent
145    /// afterwards are only recorded in the history, and new subscriptions close after replaying
146    /// it.
147    pub async fn flush_and_close(&self) {
148        loop {
149            let idle = self.delivery_state.idle.notified();
150            {
151                let mut state = self.delivery_state.inner.lock();
152                state.closed = true;
153                if state.pending == 0 {
154                    break;
155                }
156            }
157            idle.await;
158        }
159        // Dropping the senders closes the subscriber channels once they are drained.
160        self.subscribers.clear();
161    }
162
163    pub fn subscribe(
164        &self,
165        event_types: Option<Vec<String>>,
166    ) -> mpsc::Receiver<Arc<dyn CompilationEvent>> {
167        let (tx, rx) = mpsc::channel(MAX_QUEUE_SIZE);
168        let subscribers = self.subscribers.clone();
169        let event_history = self.event_history.clone();
170        let delivery_state = self.delivery_state.clone();
171        let tx_clone = tx.clone();
172
173        // Spawn a task to handle the async operations
174        tokio::spawn(async move {
175            // Store the sender (unless the queue was closed)
176            if let Some(event_types) = event_types {
177                if !delivery_state.inner.lock().closed {
178                    for event_type in event_types.iter() {
179                        let mut type_subscribers = subscribers
180                            .entry(EventChannelType::Type(event_type.clone()))
181                            .or_default();
182                        type_subscribers.push(tx_clone.clone());
183                    }
184                }
185
186                for event in event_history.lock().await.iter() {
187                    if event_types.contains(&event.type_name().to_string()) {
188                        let _ = tx_clone.send(event.clone()).await;
189                    }
190                }
191            } else {
192                if !delivery_state.inner.lock().closed {
193                    let mut global_subscribers =
194                        subscribers.entry(EventChannelType::Global).or_default();
195                    global_subscribers.push(tx_clone.clone());
196                }
197
198                for event in event_history.lock().await.iter() {
199                    let _ = tx_clone.send(event.clone()).await;
200                }
201            }
202            // If the queue was closed, tx_clone was never stored and is dropped here, closing
203            // the receiver after the history replay.
204            if delivery_state.inner.lock().closed {
205                // The queue was closed while subscribing: make sure no sender stored above
206                // lingers (flush_and_close may have cleared the map before we inserted).
207                subscribers.clear();
208            }
209        });
210
211        rx
212    }
213}
214
215#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, Serialize)]
216pub enum Severity {
217    Info,
218    Trace,
219    Warning,
220    Error,
221    Fatal,
222    Event,
223}
224
225impl Display for Severity {
226    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227        match self {
228            Severity::Info => write!(f, "INFO"),
229            Severity::Trace => write!(f, "TRACE"),
230            Severity::Warning => write!(f, "WARNING"),
231            Severity::Error => write!(f, "ERROR"),
232            Severity::Fatal => write!(f, "FATAL"),
233            Severity::Event => write!(f, "EVENT"),
234        }
235    }
236}
237
238#[derive(Debug, Clone, Serialize)]
239/// Compilation event that is used to log the duration of a task
240pub struct TimingEvent {
241    /// Message of the event without the timing information
242    ///
243    /// Example:
244    /// ```rust
245    /// use std::time::Duration;
246    /// use turbo_tasks::message_queue::{CompilationEvent, TimingEvent};
247    ///
248    /// let event = TimingEvent::new("Compiled successfully".to_string(), Duration::from_millis(100));
249    /// let message = event.message();
250    /// assert_eq!(message, "Compiled successfully in 100ms");
251    /// ```
252    pub message: String,
253    /// Duration in milliseconds
254    pub duration: Duration,
255}
256
257impl TimingEvent {
258    pub fn new(message: String, duration: Duration) -> Self {
259        Self { message, duration }
260    }
261}
262
263impl CompilationEvent for TimingEvent {
264    fn type_name(&self) -> &'static str {
265        "TimingEvent"
266    }
267
268    fn severity(&self) -> Severity {
269        Severity::Event
270    }
271
272    fn message(&self) -> String {
273        let duration_secs = self.duration.as_secs_f64();
274        let duration_string = if duration_secs > 120.0 {
275            format!("{:.1}min", duration_secs / 60.0)
276        } else if duration_secs > 40.0 {
277            format!("{duration_secs:.0}s")
278        } else if duration_secs > 2.0 {
279            format!("{duration_secs:.1}s")
280        } else {
281            format!("{:.0}ms", duration_secs * 1000.0)
282        };
283        format!("{} in {}", self.message, duration_string)
284    }
285
286    fn to_json(&self) -> String {
287        serde_json::to_string(self).unwrap()
288    }
289}
290
291#[derive(Debug, Clone, Serialize)]
292pub struct DiagnosticEvent {
293    pub message: String,
294    pub severity: Severity,
295}
296
297impl DiagnosticEvent {
298    pub fn new(severity: Severity, message: String) -> Self {
299        Self { message, severity }
300    }
301}
302
303impl CompilationEvent for DiagnosticEvent {
304    fn type_name(&self) -> &'static str {
305        "DiagnosticEvent"
306    }
307
308    fn severity(&self) -> Severity {
309        self.severity
310    }
311
312    fn message(&self) -> String {
313        self.message.clone()
314    }
315
316    fn to_json(&self) -> String {
317        serde_json::to_string(self).unwrap()
318    }
319}
320
321/// A generic trace event that carries a name, wall-clock timing, and arbitrary attributes.
322/// Forwarded as a `CompilationEvent` to the JS side for inclusion in `.next/trace`.
323#[derive(Debug, Clone, Serialize)]
324#[serde(rename_all = "camelCase")]
325pub struct TraceEvent {
326    pub name: &'static str,
327    pub start_time_ms: f64,
328    pub end_time_ms: f64,
329    /// Should be an array of key value pairs
330    pub attributes: serde_json::Value,
331}
332
333impl TraceEvent {
334    pub fn new(
335        name: &'static str,
336        start_time_ms: f64,
337        end_time_ms: f64,
338        attributes: serde_json::Value,
339    ) -> Self {
340        // basic sanity test
341        debug_assert!(matches!(attributes, serde_json::Value::Array(_)));
342        Self {
343            name,
344            start_time_ms,
345            end_time_ms,
346            attributes,
347        }
348    }
349
350    /// Creates a `TraceEvent` that started at `wall_start` (wall clock) and ran for `duration`.
351    ///
352    /// Prefer passing a duration measured with a monotonic clock (e.g.
353    /// [`std::time::Instant::elapsed`]) over computing it from two [`SystemTime`]s, so that
354    /// wall-clock adjustments don't skew (or negate) the span's duration.
355    pub fn new_with_duration(
356        name: &'static str,
357        wall_start: SystemTime,
358        duration: Duration,
359        attributes: serde_json::Value,
360    ) -> Self {
361        let start_time_ms = wall_start
362            .duration_since(SystemTime::UNIX_EPOCH)
363            .unwrap_or_default()
364            // as_millis_f64 is not stable yet
365            .as_secs_f64()
366            * 1000.0;
367        Self::new(
368            name,
369            start_time_ms,
370            start_time_ms + duration.as_secs_f64() * 1000.0,
371            attributes,
372        )
373    }
374}
375
376impl CompilationEvent for TraceEvent {
377    fn type_name(&self) -> &'static str {
378        "TraceEvent"
379    }
380
381    fn severity(&self) -> Severity {
382        Severity::Event
383    }
384
385    fn message(&self) -> String {
386        let duration_ms = self.end_time_ms - self.start_time_ms;
387        format!("{} in {:.0}ms", self.name, duration_ms)
388    }
389
390    fn to_json(&self) -> String {
391        serde_json::to_string(self).unwrap()
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398
399    #[tokio::test]
400    async fn test_flush_and_close() {
401        let queue = CompilationEventQueue::default();
402        let mut rx = queue.subscribe(None);
403        // Wait until the subscription is registered so delivery is deterministic.
404        while queue
405            .subscribers
406            .get(&EventChannelType::Global)
407            .unwrap()
408            .is_empty()
409        {
410            tokio::task::yield_now().await;
411        }
412
413        queue
414            .send(Arc::new(TimingEvent::new(
415                "test".to_string(),
416                Duration::from_millis(1),
417            )))
418            .unwrap();
419        queue.flush_and_close().await;
420
421        // The event is delivered before the channel closes.
422        let event = rx.recv().await.unwrap();
423        assert_eq!(event.message(), "test in 1ms");
424        assert!(rx.recv().await.is_none());
425
426        // Subscribing after close replays the history and closes immediately.
427        let mut rx2 = queue.subscribe(None);
428        let event = rx2.recv().await.unwrap();
429        assert_eq!(event.message(), "test in 1ms");
430        assert!(rx2.recv().await.is_none());
431    }
432
433    #[test]
434    fn test_trace_event_new_with_duration() {
435        let start = SystemTime::UNIX_EPOCH + Duration::from_secs(10);
436        let event = TraceEvent::new_with_duration(
437            "test",
438            start,
439            Duration::from_millis(1500),
440            serde_json::json!([]),
441        );
442        assert_eq!(event.name, "test");
443        assert_eq!(event.start_time_ms, 10_000.0);
444        assert_eq!(event.end_time_ms, 11_500.0);
445    }
446
447    #[test]
448    fn test_timing_event_string_formatting() {
449        let tests = vec![
450            (Duration::from_nanos(1588), "0ms"),
451            (Duration::from_nanos(1022616), "1ms"),
452            (Duration::from_millis(100), "100ms"),
453            (Duration::from_millis(1000), "1000ms"),
454            (Duration::from_millis(10000), "10.0s"),
455            (Duration::from_millis(20381), "20.4s"),
456            (Duration::from_secs(60), "60s"),
457            (Duration::from_secs(100), "100s"),
458            (Duration::from_secs(125), "2.1min"),
459        ];
460
461        for (duration, expected) in tests {
462            let event = TimingEvent::new("Compiled successfully".to_string(), duration);
463            assert_eq!(
464                event.message(),
465                format!("Compiled successfully in {expected}")
466            );
467        }
468    }
469}
turbo_tasks/message_queue.rs

turbo_tasks/
message_queue.rs
