pure_stage/
trace_buffer.rs

1#![expect(clippy::borrowed_box)]
2// Copyright 2025 PRAGMA
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! This module contains the [`TraceBuffer`] type, which is used to record the trace of a simulation.
17
18use std::{
19    borrow::Borrow,
20    collections::VecDeque,
21    fmt::{Debug, Display, Formatter},
22    sync::Arc,
23    time::Duration,
24};
25
26use cbor4ii::serde::from_slice;
27use parking_lot::Mutex;
28
29use crate::{
30    Effect, ExternalEffect, Instant, Name, ScheduleId, SendData,
31    effect::{StageEffect, StageResponse},
32    serde::to_cbor,
33};
34
35/// A buffer for recording the trace of a simulation.
36///
37/// The buffer has a bounded size and will drop the oldest entries when it is full.
38/// Attempting to push an entry that would exceed the size will try to free up space, but it will
39/// retain at least `min_entries`; if this does not suffice, the new entry will be dropped.
40///
41/// Each message in the buffer is a CBOR-encoded tuple of `(Instant, TraceEntry)`.
42#[derive(Default)]
43pub struct TraceBuffer {
44    messages: VecDeque<Vec<u8>>,
45    min_entries: usize,
46    max_size: usize,
47    used_size: usize,
48    dropped_messages: usize,
49    fetch_replay: Option<std::vec::IntoIter<TraceEntry>>,
50}
51
52/// An owned entry can be serialized and deserialized.
53///
54/// However, we need the entries in the continuation of the running program and we cannot clone,
55/// so there is a variant of this enum with the very same structure the captures only references
56/// to work around this problem: `TraceEntryRef`.
57#[derive(PartialEq, serde::Serialize, serde::Deserialize)]
58pub enum TraceEntry {
59    Suspend(Effect),
60    Resume {
61        stage: Name,
62        response: StageResponse,
63    },
64    Clock(Instant),
65    Input {
66        stage: Name,
67        #[serde(with = "crate::serde::serialize_send_data")]
68        input: Box<dyn SendData>,
69    },
70    State {
71        stage: Name,
72        #[serde(with = "crate::serde::serialize_send_data")]
73        state: Box<dyn SendData>,
74    },
75    Terminated {
76        stage: Name,
77        reason: TerminationReason,
78    },
79}
80
81#[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)]
82pub enum TerminationReason {
83    Voluntary,
84    Supervision(Name),
85    Aborted,
86}
87
88impl TraceEntry {
89    pub fn to_json(&self) -> serde_json::Value {
90        match &self {
91            TraceEntry::Suspend(effect) => serde_json::json!({
92                "type": "suspend",
93                "effect": effect.to_json(),
94            }),
95            TraceEntry::Resume { stage, response } => serde_json::json!({
96                "type": "resume",
97                "stage": stage,
98                "response": response,
99            }),
100            TraceEntry::Clock(instant) => serde_json::json!({
101            "type": "clock",
102            "instant": instant,
103            }),
104            TraceEntry::Input { stage, input } => serde_json::json!({
105            "type": "input",
106            "stage": stage,
107            "input": input.to_string(),
108            }),
109            TraceEntry::State { stage, state } => serde_json::json!({
110            "type": "state",
111            "stage": stage,
112            "state": state.to_string(),
113            }),
114            TraceEntry::Terminated { stage, reason } => serde_json::json!({
115            "type": "terminated",
116            "stage": stage,
117            "reason": reason,
118            }),
119        }
120    }
121}
122
123impl Debug for TraceEntry {
124    /// This debug instance does not output the runnable field in the Resume case.
125    /// That field is only useful for display purposes.
126    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
127        match self {
128            TraceEntry::Suspend(effect) => f.debug_tuple("Suspend").field(&effect).finish(),
129            TraceEntry::Resume { stage, response, .. } => {
130                f.debug_struct("Resume").field("stage", stage).field("response", response).finish()
131            }
132            TraceEntry::Clock(instant) => f.debug_tuple("Clock").field(instant).finish(),
133            TraceEntry::Input { stage, input } => {
134                f.debug_struct("Input").field("stage", stage).field("input", input).finish()
135            }
136            TraceEntry::State { stage, state } => {
137                f.debug_struct("State").field("stage", stage).field("state", state).finish()
138            }
139            TraceEntry::Terminated { stage, reason } => {
140                f.debug_struct("Terminated").field("stage", stage).field("reason", reason).finish()
141            }
142        }
143    }
144}
145
146impl Display for TraceEntry {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        match self {
149            TraceEntry::Suspend(effect) => write!(f, "suspend {stage}", stage = effect.at_stage()),
150            TraceEntry::Resume { stage, response } => {
151                write!(
152                    f,
153                    "resume {stage}{response}",
154                    stage = stage.as_str(),
155                    response = match response {
156                        StageResponse::Unit => "".to_string(),
157                        other @ StageResponse::ClockResponse(_)
158                        | other @ StageResponse::WaitResponse(_)
159                        | other @ StageResponse::CallResponse(_)
160                        | other @ StageResponse::CancelScheduleResponse(_)
161                        | other @ StageResponse::ExternalResponse(_)
162                        | other @ StageResponse::AddStageResponse(_)
163                        | other @ StageResponse::ContramapResponse(_) => format!(" -> {other}"),
164                    },
165                )
166            }
167            TraceEntry::Clock(instant) => write!(f, "clock {instant}"),
168            TraceEntry::Input { stage, input } => {
169                write!(
170                    f,
171                    "input {stage} -> {input}",
172                    stage = stage.as_str(),
173                    input = input.as_send_data_value().borrow()
174                )
175            }
176            TraceEntry::State { stage, state } => {
177                write!(
178                    f,
179                    "state {stage} -> {state}",
180                    stage = stage.as_str(),
181                    state = state.as_send_data_value().borrow()
182                )
183            }
184            TraceEntry::Terminated { stage, reason } => {
185                write!(f, "terminated {stage} {reason:?}")
186            }
187        }
188    }
189}
190
191impl TraceEntry {
192    /// Construct a suspend entry.
193    pub fn suspend(effect: Effect) -> Self {
194        Self::Suspend(effect)
195    }
196
197    /// Construct a resume entry.
198    pub fn resume(stage: impl AsRef<str>, response: StageResponse) -> Self {
199        Self::Resume { stage: Name::from(stage.as_ref()), response }
200    }
201
202    /// Construct a clock entry.
203    pub fn clock(instant: Instant) -> Self {
204        Self::Clock(instant)
205    }
206
207    /// Construct an input entry.
208    pub fn input(stage: impl AsRef<str>, input: Box<dyn SendData>) -> Self {
209        Self::Input { stage: Name::from(stage.as_ref()), input }
210    }
211
212    /// Construct a state entry.
213    pub fn state(stage: impl AsRef<str>, state: Box<dyn SendData>) -> Self {
214        Self::State { stage: Name::from(stage.as_ref()), state }
215    }
216
217    /// Construct a terminated entry.
218    pub fn terminated(stage: impl AsRef<str>, reason: TerminationReason) -> Self {
219        Self::Terminated { stage: Name::from(stage.as_ref()), reason }
220    }
221
222    pub fn at_stage<'a, 'b: 'a>(&'b self) -> Option<&'a Name> {
223        match self {
224            TraceEntry::Suspend(effect) => Some(effect.at_stage()),
225            TraceEntry::Resume { stage, .. } => Some(stage),
226            TraceEntry::Clock(..) => None,
227            TraceEntry::Input { stage, .. } => Some(stage),
228            TraceEntry::State { stage, .. } => Some(stage),
229            TraceEntry::Terminated { stage, .. } => Some(stage),
230        }
231    }
232}
233
234/// A non-owning variant of [`TraceEntry`] that allows serializing an entry without consuming it.
235///
236/// This structure cannot be deserialized, use the owned version for that.
237#[derive(Debug, PartialEq, serde::Serialize)]
238enum TraceEntryRef<'a> {
239    Suspend(&'a Effect),
240    Resume { stage: &'a Name, response: &'a StageResponse },
241    Clock(Instant),
242    Input { stage: &'a Name, input: &'a Box<dyn SendData> },
243    State { stage: &'a Name, state: &'a Box<dyn SendData> },
244    Terminated { stage: &'a Name, reason: TerminationReasonRef<'a> },
245}
246
247#[derive(Debug, PartialEq, serde::Serialize)]
248enum TerminationReasonRef<'a> {
249    Voluntary,
250    Supervision(&'a Name),
251    Aborted,
252}
253
254impl<'a> From<&'a TerminationReason> for TerminationReasonRef<'a> {
255    fn from(reason: &'a TerminationReason) -> Self {
256        match reason {
257            TerminationReason::Voluntary => TerminationReasonRef::Voluntary,
258            TerminationReason::Supervision(child) => TerminationReasonRef::Supervision(child),
259            TerminationReason::Aborted => TerminationReasonRef::Aborted,
260        }
261    }
262}
263
264/// Helper struct that has the same serialization format as TraceEntry but doesn’t require owned effect data.
265#[derive(serde::Serialize)]
266enum TraceEntryRefRef<'a> {
267    Suspend(EffectRef<'a>),
268    Resume { stage: &'a Name, response: StageResponseRef<'a> },
269}
270
271/// Helper struct that has the same serialization format as Effect but doesn’t require owned effect data.
272#[derive(serde::Serialize)]
273enum EffectRef<'a> {
274    Receive { at_stage: &'a Name },
275    Send { from: &'a Name, to: &'a Name, msg: &'a dyn SendData },
276    Call { from: &'a Name, to: &'a Name, duration: Duration, msg: &'a dyn SendData },
277    Clock { at_stage: &'a Name },
278    Wait { at_stage: &'a Name, duration: Duration },
279    Schedule { at_stage: &'a Name, msg: &'a dyn SendData, id: ScheduleId },
280    CancelSchedule { at_stage: &'a Name, id: ScheduleId },
281    External { at_stage: &'a Name, effect: &'a dyn crate::ExternalEffect },
282    Terminate { at_stage: &'a Name },
283    AddStage { at_stage: &'a Name, name: &'a Name },
284    WireStage { at_stage: &'a Name, name: &'a Name, initial_state: &'a dyn SendData, tombstone: &'a dyn SendData },
285    Contramap { at_stage: &'a Name, original: &'a Name, new_name: &'a Name },
286}
287
288impl<'a> EffectRef<'a> {
289    fn from(at_stage: &'a Name, effect: &'a StageEffect<Box<dyn SendData>>) -> Option<Self> {
290        Some(match effect {
291            StageEffect::Receive => EffectRef::Receive { at_stage },
292            StageEffect::Send(to, _call, msg) => EffectRef::Send { from: at_stage, to, msg: &**msg },
293            StageEffect::Call(..) => return None,
294            StageEffect::Clock => EffectRef::Clock { at_stage },
295            StageEffect::Wait(duration) => EffectRef::Wait { at_stage, duration: *duration },
296            StageEffect::Schedule(msg, id) => EffectRef::Schedule { at_stage, msg: &**msg, id: *id },
297            StageEffect::CancelSchedule(id) => EffectRef::CancelSchedule { at_stage, id: *id },
298            StageEffect::External(effect) => EffectRef::External { at_stage, effect: &**effect },
299            StageEffect::Terminate => EffectRef::Terminate { at_stage },
300            StageEffect::AddStage(name) => EffectRef::AddStage { at_stage, name },
301            StageEffect::WireStage(name, _transition, initial_state, tombstone) => {
302                EffectRef::WireStage { at_stage, name, initial_state: &**initial_state, tombstone: &**tombstone }
303            }
304            StageEffect::Contramap { original, new_name, transform: _ } => {
305                EffectRef::Contramap { at_stage, original, new_name }
306            }
307        })
308    }
309}
310
311/// Helper struct that has the same serialization format as StageResponse but doesn’t require owned response data.
312#[derive(serde::Serialize)]
313enum StageResponseRef<'a> {
314    ExternalResponse(&'a dyn SendData),
315}
316
317impl TraceBuffer {
318    /// Create a new shareable trace buffer.
319    ///
320    /// This is used to retain a reference to the trace buffer while the simulation is running.
321    pub fn new_shared(min_entries: usize, max_size: usize) -> Arc<Mutex<Self>> {
322        Arc::new(Mutex::new(Self::new(min_entries, max_size)))
323    }
324
325    /// Create a new trace buffer with the given size limits.
326    ///
327    /// Set `min_entries` and `max_size` to zero to obtain a buffer that will drop all entries
328    /// without allocating memory or otherwise wasting much time.
329    pub fn new(min_entries: usize, max_size: usize) -> Self {
330        Self { messages: VecDeque::new(), min_entries, max_size, used_size: 0, dropped_messages: 0, fetch_replay: None }
331    }
332
333    /// Push an effect to the trace buffer.
334    pub fn push_suspend(&mut self, effect: &Effect) {
335        self.push(TraceEntryRef::Suspend(effect));
336    }
337
338    pub fn push_suspend_external(&mut self, at_stage: &Name, effect: &dyn crate::ExternalEffect) {
339        self.push(TraceEntryRefRef::Suspend(EffectRef::External { at_stage, effect }));
340    }
341
342    pub fn push_suspend_call(&mut self, at_stage: &Name, to: &Name, duration: Duration, msg: &dyn SendData) {
343        self.push(TraceEntryRefRef::Suspend(EffectRef::Call { from: at_stage, to, duration, msg }));
344    }
345
346    pub fn push_suspend_ref(&mut self, at_stage: &Name, effect: &StageEffect<Box<dyn SendData>>) {
347        if let Some(effect) = EffectRef::from(at_stage, effect) {
348            self.push(TraceEntryRefRef::Suspend(effect));
349        }
350    }
351
352    /// Push a resume event to the trace buffer.
353    pub fn push_resume(&mut self, stage: &Name, response: &StageResponse) {
354        self.push(TraceEntryRef::Resume { stage, response });
355    }
356
357    pub fn push_resume_external(&mut self, stage: &Name, response: &dyn SendData) {
358        self.push(TraceEntryRefRef::Resume { stage, response: StageResponseRef::ExternalResponse(response) });
359    }
360
361    /// Push a clock update to the trace buffer.
362    pub fn push_clock(&mut self, instant: Instant) {
363        self.push(TraceEntryRef::Clock(instant));
364    }
365
366    /// Push a receive event to the trace buffer.
367    pub fn push_input(&mut self, stage: &Name, input: &Box<dyn SendData>) {
368        self.push(TraceEntryRef::Input { stage, input });
369    }
370
371    /// Push a state update to the trace buffer.
372    ///
373    /// This happens every time polling a stage yields a `Poll::Ready(Ok(...))`, i.e. as soon as the next
374    /// stage state has been computed.
375    pub fn push_state(&mut self, stage: &Name, state: &Box<dyn SendData>) {
376        self.push(TraceEntryRef::State { stage, state });
377    }
378
379    pub fn push_terminated(&mut self, stage: &Name, reason: TerminationReason) {
380        self.push(TraceEntryRef::Terminated { stage, reason: (&reason).into() });
381    }
382
383    pub fn push_terminated_voluntary(&mut self, stage: &Name) {
384        self.push(TraceEntryRef::Terminated { stage, reason: TerminationReasonRef::Voluntary });
385    }
386
387    pub fn push_terminated_supervision(&mut self, stage: &Name, child: &Name) {
388        self.push(TraceEntryRef::Terminated { stage, reason: TerminationReasonRef::Supervision(child) });
389    }
390
391    pub fn push_terminated_aborted(&mut self, stage: &Name) {
392        self.push(TraceEntryRef::Terminated { stage, reason: TerminationReasonRef::Aborted });
393    }
394
395    fn push<T: serde::Serialize>(&mut self, msg: T) {
396        let msg = to_cbor(&(Instant::now(), msg));
397
398        if self.max_size == 0 {
399            return;
400        }
401
402        self.used_size += msg.len();
403
404        #[expect(clippy::expect_used)]
405        while self.used_size > self.max_size && self.messages.len() > self.min_entries {
406            self.used_size -= self.messages.pop_front().expect("messages is definitely not empty").len();
407            self.dropped_messages += 1;
408        }
409
410        // the loop above keeps the number of entries at least at min_entries, so if we can
411        // push a new one later we can remove one more
412        let one_more = self.messages.front().map(|m| m.len()).unwrap_or(0);
413        if self.used_size > self.max_size && self.used_size - one_more <= self.max_size {
414            self.used_size -= one_more;
415            self.messages.pop_front();
416        }
417
418        if self.used_size > self.max_size {
419            self.used_size -= msg.len();
420            self.dropped_messages += 1;
421            return;
422        }
423
424        self.messages.push_back(msg);
425    }
426
427    /// Iterate over the entries in the trace buffer.
428    pub fn iter(&self) -> impl Iterator<Item = &[u8]> {
429        self.messages.iter().map(|m| m.as_slice())
430    }
431
432    /// Iterate over the deserialized entries in the trace buffer.
433    pub fn iter_entries(&self) -> impl Iterator<Item = (Instant, TraceEntry)> {
434        #[expect(clippy::expect_used)]
435        self.messages.iter().map(|m| from_slice(m).expect("trace buffer is not supposed to contain invalid CBOR"))
436    }
437
438    /// Take the entries from the trace buffer, leaving it empty.
439    pub fn take(&mut self) -> Vec<Vec<u8>> {
440        std::mem::take(&mut self.messages).into()
441    }
442
443    /// Hydrate the trace buffer (stored in CBOR format) into a vector of entries.
444    #[expect(clippy::expect_used)]
445    pub fn hydrate(&self) -> Vec<(Instant, TraceEntry)> {
446        self.messages
447            .iter()
448            .map(|m| from_slice(m).expect("trace buffer is not supposed to contain invalid CBOR"))
449            .collect()
450    }
451
452    /// Hydrate the trace buffer (stored in CBOR format) into a vector of entries.
453    #[expect(clippy::expect_used)]
454    pub fn hydrate_without_timestamps(&self) -> Vec<TraceEntry> {
455        self.messages
456            .iter()
457            .map(|m| {
458                from_slice::<(Instant, TraceEntry)>(m).expect("trace buffer is not supposed to contain invalid CBOR").1
459            })
460            .collect()
461    }
462
463    /// Clear the trace buffer.
464    pub fn clear(&mut self) {
465        self.messages.clear();
466        self.used_size = 0;
467        self.dropped_messages = 0;
468    }
469
470    /// Get the number of entries in the trace buffer.
471    pub fn len(&self) -> usize {
472        self.messages.len()
473    }
474
475    /// Check if the trace buffer is empty.
476    pub fn is_empty(&self) -> bool {
477        self.messages.is_empty()
478    }
479
480    /// Get the total size of the entries in the trace buffer.
481    pub fn used_size(&self) -> usize {
482        self.used_size
483    }
484
485    /// Get the number of messages that have been dropped from the trace buffer.
486    pub fn dropped_messages(&self) -> usize {
487        self.dropped_messages
488    }
489
490    pub fn set_fetch_replay(&mut self, replay: std::vec::IntoIter<TraceEntry>) {
491        self.fetch_replay = Some(replay);
492    }
493
494    pub fn take_fetch_replay(&mut self) -> Option<std::vec::IntoIter<TraceEntry>> {
495        self.fetch_replay.take()
496    }
497
498    pub fn fetch_replay_mut(&mut self) -> Option<&mut std::vec::IntoIter<TraceEntry>> {
499        self.fetch_replay.as_mut()
500    }
501
502    pub fn drop_guard(this: &Arc<Mutex<Self>>) -> DropGuard {
503        DropGuard { buffer: this.clone(), active: true }
504    }
505}
506
507pub struct DropGuard {
508    buffer: Arc<Mutex<TraceBuffer>>,
509    active: bool,
510}
511
512impl DropGuard {
513    pub fn defuse(mut self) {
514        self.active = false;
515    }
516}
517
518impl Drop for DropGuard {
519    fn drop(&mut self) {
520        if !self.active {
521            return;
522        }
523        eprintln!("Dropping trace buffer");
524        for entry in self.buffer.lock().iter() {
525            match from_slice::<(Instant, TraceEntry)>(entry) {
526                Ok((instant, entry)) => eprintln!("{instant} {entry:?}"),
527                Err(error) => eprintln!("error deserializing trace entry: {error:?}"),
528            };
529        }
530        eprintln!("Dropped trace buffer");
531    }
532}
533
534#[allow(clippy::wildcard_enum_match_arm, clippy::panic)]
535pub fn find_next_external_suspend(
536    fetch_replay: &mut std::vec::IntoIter<TraceEntry>,
537    at_stage: &Name,
538) -> Option<Box<dyn ExternalEffect>> {
539    find_next(
540        fetch_replay,
541        |entry| entry.at_stage() == Some(at_stage),
542        |entry| match entry {
543            TraceEntry::Suspend(Effect::External { effect, .. }) => effect,
544            entry => panic!("unexpected trace entry when finding next external suspend: {:?}", entry),
545        },
546    )
547}
548
549#[allow(clippy::wildcard_enum_match_arm, clippy::panic)]
550pub fn find_next_external_resume(
551    fetch_replay: &mut std::vec::IntoIter<TraceEntry>,
552    at_stage: &Name,
553) -> Option<Box<dyn SendData>> {
554    find_next(
555        fetch_replay,
556        |entry| entry.at_stage() == Some(at_stage),
557        |entry| match entry {
558            TraceEntry::Resume { response: StageResponse::ExternalResponse(response), .. } => response,
559            entry => panic!("unexpected trace entry when finding next external resume: {:?}", entry),
560        },
561    )
562}
563
564fn find_next<T>(
565    fetch_replay: &mut std::vec::IntoIter<TraceEntry>,
566    predicate: impl Fn(&TraceEntry) -> bool,
567    extract: impl FnOnce(TraceEntry) -> T,
568) -> Option<T> {
569    let log = fetch_replay.as_mut_slice();
570    let idx = log.iter().position(predicate)?;
571    log[..=idx].rotate_right(1);
572    fetch_replay.next().map(extract)
573}
574
575#[cfg(test)]
576mod tests {
577    use tokio::sync::mpsc;
578
579    use super::*;
580    use crate::OutputEffect;
581
582    #[test]
583    fn test_serialization() {
584        let effect = OutputEffect::new(Name::from("test"), 42u32, mpsc::channel(1).0);
585        let trr = TraceEntryRefRef::Suspend(EffectRef::External { at_stage: &Name::from("test"), effect: &effect });
586        let rr = to_cbor(&trr);
587        let json = serde_json::to_string(&trr).unwrap();
588        let r = to_cbor(&TraceEntryRef::Suspend(&Effect::External {
589            at_stage: Name::from("test"),
590            effect: Box::new(effect),
591        }));
592        let effect = OutputEffect::new(Name::from("test"), 42u32, mpsc::channel(1).0);
593        let t =
594            to_cbor(&TraceEntry::suspend(Effect::External { at_stage: Name::from("test"), effect: Box::new(effect) }));
595        assert_eq!(rr, r);
596        assert_eq!(rr, t);
597
598        assert_eq!(
599            json,
600            r#"{"Suspend":{"External":{"at_stage":"test","effect":{"typetag":"pure_stage::output::OutputEffect<u32>","value":{"name":"test","msg":42,"sender":{}}}}}}"#
601        );
602    }
603}
pure_stage/trace_buffer.rs

pure_stage/
trace_buffer.rs
