vercel/next.js · #98426

Fix startTime and shutdown delivery of forwarded Turbopack trace spans

mischnic · merged Sep 22, 20265 files · 182 + / 38 −
packages/next/src/build/turbopack-build/impl.ts8 + / 5 −
@@ -153,12 +153,15 @@ export async function turbopackBuild(telemetry: Telemetry): Promise<{     signal: shutdownController.signal,   })   const runShutdown = async () => {-    // Shutdown may trigger final compilation events (e.g. persistence,-    // compaction trace spans).  This is the last chance to capture them.-    // After shutdown resolves we abort the signal to close the iterator-    // and drain any remaining buffered events.-     await project.shutdown()+    // Shutdown flushes and closes the compilation event queue, so the+    // subscription ends once final events (e.g. persistence, compaction trace+    // spans) have been delivered. The timeout is only a backstop against a+    // subscription that never closes.+    await Promise.race([+      compilationEvents,+      new Promise((resolve) => setTimeout(resolve, 10_000).unref()),+    ])     shutdownController.abort()     await compilationEvents   }
packages/next/src/trace/trace.test.ts20 + / 0 −
@@ -147,6 +147,26 @@ describe('Trace', () => {       ])     }) +    it('records manualTraceChild spans with their own start time', () => {+      const root = trace('root-span')+      // A span that started one minute ago and ran for 1.5s, as happens when+      // the event is only recorded after the work completes (e.g. Turbopack+      // compilation events forwarded from Rust).+      const startMs = Date.now() - 60_000+      const stopMs = startMs + 1500+      const msToNs = (ms: number) => BigInt(ms) * BigInt(1_000_000)+      root.manualTraceChild('manual-child', msToNs(startMs), msToNs(stopMs))+      root.stop()++      const events = getTraceEvents()+      const child = events.find((e) => e.name === 'manual-child')+      expect(child).toBeDefined()+      // startTime must be the span's actual start, not the time the event was+      // recorded (which is when manualTraceChild is called).+      expect(child!.startTime).toEqual(startMs)+      expect(child!.duration).toEqual(1_500_000)+    })+     // Builds flush repeatedly -- after each compilation, and again at the     // end -- so every flush must leave the file usable.     it('writes buffered spans on each flush', async () => {
packages/next/src/trace/trace.ts11 + / 1 −
@@ -43,10 +43,14 @@ export class Span {     parentId,     attrs,     startTime,+    wallStartTimeMs,   }: {     name: string     parentId?: SpanId     startTime?: bigint+    // Wall-clock start time in milliseconds since epoch, matching `startTime`.+    // Defaults to the construction time.+    wallStartTimeMs?: number     attrs?: Attributes   }) {     this.name = name@@ -61,7 +65,7 @@ export class Span {     // `These times are relative to an arbitrary time in the past, and not related to the time of day and therefore not subject to clock drift`     // https://nodejs.org/api/process.html#processhrtimetime     // Capturing current datetime as additional metadata for external reconstruction.-    this.now = Date.now()+    this.now = wallStartTimeMs ?? Date.now()   }    // Durations are reported as microseconds. This gives 1000x the precision@@ -118,6 +122,12 @@ export class Span {       parentId: this.id,       attrs,       startTime: startTime ? startTime + correction : process.hrtime.bigint(),+      // The span is constructed when the event is recorded, which may be long+      // after the work started (e.g. compilation events forwarded from+      // Turbopack), so derive the wall-clock start from the given epoch time.+      wallStartTimeMs: startTime+        ? Number(startTime / NUM_OF_MILLISEC_IN_NANOSEC)+        : undefined,     })     span.stop(stopTime ? stopTime + correction : process.hrtime.bigint())   }
turbopack/crates/turbo-tasks/src/manager.rs4 + / 0 −
@@ -1419,6 +1419,10 @@ impl<B: Backend + 'static> TurboTasks<B> {                 }             }             self.backend.stop(self);+            // Deliver compilation events sent during shutdown (e.g. the persistence trace span)+            // to subscribers before returning, then close the queue so subscriptions end after+            // draining.+            self.compilation_events.flush_and_close().await;         })         .await;     }
turbopack/crates/turbo-tasks/src/message_queue.rs139 + / 32 −
@@ -1,8 +1,9 @@ use std::{any::Any, collections::VecDeque, fmt::Display, sync::Arc, time::Duration};  use dashmap::DashMap;+use parking_lot::Mutex; use serde::Serialize;-use tokio::sync::{Mutex, mpsc};+use tokio::sync::{Mutex as TokioMutex, Notify, mpsc};  pub trait CompilationEvent: Sync + Send + Any {     fn type_name(&self) -> &'static str;@@ -13,7 +14,7 @@ pub trait CompilationEvent: Sync + Send + Any {  const MAX_QUEUE_SIZE: usize = 256; -type ArcMx<T> = Arc<Mutex<T>>;+type ArcMx<T> = Arc<TokioMutex<T>>; type CompilationEventChannel = mpsc::Sender<Arc<dyn CompilationEvent>>;  #[derive(Debug, Clone, Eq, PartialEq, Hash)]@@ -22,9 +23,20 @@ enum EventChannelType {     Type(String), } +struct DeliveryStateInner {+    closed: bool,+    pending: usize,+}++struct DeliveryState {+    inner: Mutex<DeliveryStateInner>,+    idle: Notify,+}+ pub struct CompilationEventQueue {     event_history: ArcMx<VecDeque<Arc<dyn CompilationEvent>>>,     subscribers: Arc<DashMap<EventChannelType, Vec<CompilationEventChannel>>>,+    delivery_state: Arc<DeliveryState>, }  impl Default for CompilationEventQueue {@@ -36,8 +48,15 @@ impl Default for CompilationEventQueue {         );          Self {-            event_history: Arc::new(Mutex::new(VecDeque::with_capacity(MAX_QUEUE_SIZE))),+            event_history: Arc::new(TokioMutex::new(VecDeque::with_capacity(MAX_QUEUE_SIZE))),             subscribers: Arc::new(subscribers),+            delivery_state: Arc::new(DeliveryState {+                inner: Mutex::new(DeliveryStateInner {+                    closed: false,+                    pending: 0,+                }),+                idle: Notify::new(),+            }),         }     } }@@ -50,6 +69,16 @@ impl CompilationEventQueue {         let event_history = self.event_history.clone();         let subscribers = self.subscribers.clone();         let message_clone = message.clone();+        let delivery_state = self.delivery_state.clone();+        let deliver = {+            let mut state = delivery_state.inner.lock();+            if state.closed {+                false+            } else {+                state.pending += 1;+                true+            }+        };          // Spawn a task to handle the async operations         tokio::spawn(async move {@@ -59,58 +88,93 @@ impl CompilationEventQueue {                 history.pop_front();             }             history.push_back(message_clone.clone());+            drop(history);++            if deliver {+                // Send to all active receivers of the same message type+                if let Some(mut type_subscribers) = subscribers.get_mut(&EventChannelType::Type(+                    message_clone.type_name().to_owned(),+                )) {+                    let mut removal_indices = Vec::new();+                    for (ix, sender) in type_subscribers.iter().enumerate() {+                        if sender.send(message_clone.clone()).await.is_err() {+                            removal_indices.push(ix);+                        }+                    } -            // Send to all active receivers of the same message type-            if let Some(mut type_subscribers) = subscribers.get_mut(&EventChannelType::Type(-                message_clone.type_name().to_owned(),-            )) {-                let mut removal_indices = Vec::new();-                for (ix, sender) in type_subscribers.iter().enumerate() {-                    if sender.send(message_clone.clone()).await.is_err() {-                        removal_indices.push(ix);+                    for ix in removal_indices.iter().rev() {+                        type_subscribers.remove(*ix);                     }                 } -                for ix in removal_indices.iter().rev() {-                    type_subscribers.remove(*ix);-                }-            }+                // Send to all global message subscribers+                if let Some(mut all_channel) = subscribers.get_mut(&EventChannelType::Global) {+                    let mut removal_indices = Vec::new();+                    for (ix, sender) in all_channel.iter_mut().enumerate() {+                        if sender.send(message_clone.clone()).await.is_err() {+                            removal_indices.push(ix);+                        }+                    } -            // Send to all global message subscribers-            let mut all_channel = subscribers.get_mut(&EventChannelType::Global).unwrap();-            let mut removal_indices = Vec::new();-            for (ix, sender) in all_channel.iter_mut().enumerate() {-                if sender.send(message_clone.clone()).await.is_err() {-                    removal_indices.push(ix);+                    for ix in removal_indices.iter().rev() {+                        all_channel.remove(*ix);+                    }                 }             } -            for ix in removal_indices.iter().rev() {-                all_channel.remove(*ix);+            if deliver {+                let mut state = delivery_state.inner.lock();+                state.pending -= 1;+                if state.pending == 0 {+                    delivery_state.idle.notify_one();+                }             }         });          Ok(())     } +    /// Waits until all events sent so far have been delivered to subscribers, then closes all+    /// subscriber channels. Subscriptions drain their remaining events and then end. Events sent+    /// afterwards are only recorded in the history, and new subscriptions close after replaying+    /// it.+    pub async fn flush_and_close(&self) {+        loop {+            let idle = self.delivery_state.idle.notified();+            {+                let mut state = self.delivery_state.inner.lock();+                state.closed = true;+                if state.pending == 0 {+                    break;+                }+            }+            idle.await;+        }+        // Dropping the senders closes the subscriber channels once they are drained.+        self.subscribers.clear();+    }+     pub fn subscribe(         &self,         event_types: Option<Vec<String>>,     ) -> mpsc::Receiver<Arc<dyn CompilationEvent>> {         let (tx, rx) = mpsc::channel(MAX_QUEUE_SIZE);         let subscribers = self.subscribers.clone();         let event_history = self.event_history.clone();+        let delivery_state = self.delivery_state.clone();         let tx_clone = tx.clone();          // Spawn a task to handle the async operations         tokio::spawn(async move {-            // Store the sender+            // Store the sender (unless the queue was closed)             if let Some(event_types) = event_types {-                for event_type in event_types.iter() {-                    let mut type_subscribers = subscribers-                        .entry(EventChannelType::Type(event_type.clone()))-                        .or_default();-                    type_subscribers.push(tx_clone.clone());+                if !delivery_state.inner.lock().closed {+                    for event_type in event_types.iter() {+                        let mut type_subscribers = subscribers+                            .entry(EventChannelType::Type(event_type.clone()))+                            .or_default();+                        type_subscribers.push(tx_clone.clone());+                    }                 }                  for event in event_history.lock().await.iter() {@@ -119,14 +183,23 @@ impl CompilationEventQueue {                     }                 }             } else {-                let mut global_subscribers =-                    subscribers.entry(EventChannelType::Global).or_default();-                global_subscribers.push(tx_clone.clone());+                if !delivery_state.inner.lock().closed {+                    let mut global_subscribers =+                        subscribers.entry(EventChannelType::Global).or_default();+                    global_subscribers.push(tx_clone.clone());+                }                  for event in event_history.lock().await.iter() {                     let _ = tx_clone.send(event.clone()).await;                 }             }+            // If the queue was closed, tx_clone was never stored and is dropped here, closing+            // the receiver after the history replay.+            if delivery_state.inner.lock().closed {+                // The queue was closed while subscribing: make sure no sender stored above+                // lingers (flush_and_close may have cleared the map before we inserted).+                subscribers.clear();+            }         });          rx@@ -292,6 +365,40 @@ impl CompilationEvent for TraceEvent { mod tests {     use super::*; +    #[tokio::test]+    async fn test_flush_and_close() {+        let queue = CompilationEventQueue::default();+        let mut rx = queue.subscribe(None);+        // Wait until the subscription is registered so delivery is deterministic.+        while queue+            .subscribers+            .get(&EventChannelType::Global)+            .unwrap()+            .is_empty()+        {+            tokio::task::yield_now().await;+        }++        queue+            .send(Arc::new(TimingEvent::new(+                "test".to_string(),+                Duration::from_millis(1),+            )))+            .unwrap();+        queue.flush_and_close().await;++        // The event is delivered before the channel closes.+        let event = rx.recv().await.unwrap();+        assert_eq!(event.message(), "test in 1ms");+        assert!(rx.recv().await.is_none());++        // Subscribing after close replays the history and closes immediately.+        let mut rx2 = queue.subscribe(None);+        let event = rx2.recv().await.unwrap();+        assert_eq!(event.message(), "test in 1ms");+        assert!(rx2.recv().await.is_none());+    }+     #[test]     fn test_timing_event_string_formatting() {         let tests = vec![