1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
//! Collect [Tracy] profiles in tracing-enabled applications.
//!
//! Assuming the application is well instrumented, this should in practice be a very low effort way
//! to gain great amounts of insight into an application performance.
//!
//! Note, however that Tracy is ultimately a profiling, not an observability, tool. As thus, some
//! of tracing concepts cannot be represented well by Tracy. For instance, out-of-order span
//! entries and exits, are not supported, and neither are spans that are entered and exited on
//! different threads. This crate will attempt to mitigate the problems and retain trace validity
//! at the cost of potentially invalid data. When such a mitigation occurs, trace will contain a
//! message with a note about the problem.
//!
//! Some other caveats to keep in mind:
//!
//! * Only span entries and exits are recorded;
//! * Events show up as messages in Tracy, however Tracy can struggle with large numbers of
//! messages;
//! * Some additional functionality such as plotting and memory allocation profiling is only
//! available as part of the [tracy-client](tracy_client) crate.
//!
//! # Important note
//!
//! Unlike with many other subscriber implementations, simply depending on this crate is sufficient
//! for tracy to be enabled at program startup, even if [`TracyLayer`](TracyLayer) is not
//! registered as a subscriber. While not registering a `TracyLayer` will avoid Tracy from
//! collecting spans, it still broadcasts discovery packets to the local network and exposes the
//! data it collects in the background to that same network. Traces collected by Tracy may include
//! source and assembly code as well.
//!
//! As thus, you may want make sure to only enable the `tracing-tracy` crate conditionally, via the
//! `enable` feature flag provided by this crate.
//!
//! [Tracy]: https://github.com/wolfpld/tracy

use std::{fmt::Write, collections::VecDeque, cell::RefCell};
use tracing_core::{
    field::{Field, Visit},
    span::Id,
    Event, Subscriber,
};
use tracing_subscriber::{
    layer::{Context, Layer},
    registry,
};

use tracy_client::{Span, color_message, message, finish_continuous_frame};

thread_local! {
    /// A stack of spans currently active on the current thread.
    static TRACY_SPAN_STACK: RefCell<VecDeque<(Span, u64)>> =
        RefCell::new(VecDeque::with_capacity(16));
}

/// A tracing layer that collects data in Tracy profiling format.
#[derive(Clone)]
pub struct TracyLayer {
    stack_depth: u16,
}

impl TracyLayer {
    /// Create a new `TracyLayer`.
    ///
    /// Defaults to collecting stack traces.
    pub fn new() -> Self {
        Self { stack_depth: 64 }
    }

    /// Specify the maximum number of stack frames that will be collected.
    ///
    /// Specifying 0 frames will disable stack trace collection.
    pub fn with_stackdepth(mut self, stack_depth: u16) -> Self {
        self.stack_depth = stack_depth;
        self
    }
}

impl Default for TracyLayer {
    fn default() -> Self {
        Self::new()
    }
}

impl<S> Layer<S> for TracyLayer
where
    S: Subscriber + for<'a> registry::LookupSpan<'a>,
{
    fn on_enter(&self, id: &Id, ctx: Context<S>) {
        if let Some(span_data) = ctx.span(id) {
            let metadata = span_data.metadata();
            let file = metadata.file().unwrap_or("<error: not available>");
            let line = metadata.line().unwrap_or(0);
            TRACY_SPAN_STACK.with(|s| {
                s.borrow_mut().push_back((
                    Span::new(metadata.name(), "", file, line, self.stack_depth),
                    id.into_u64()
                ));
            });
        }
    }

    fn on_exit(&self, id: &Id, _: Context<S>) {
        TRACY_SPAN_STACK.with(|s| {
            if let Some((span, span_id)) = s.borrow_mut().pop_back() {
                if id.into_u64() != span_id {
                    color_message(
                        "Tracing spans exited out of order! \
                        Trace may not be accurate for this span stack.",
                        0xFF000000,
                        self.stack_depth,
                    );
                }
                drop(span);
            } else {
                color_message(
                    "Exiting a tracing span, but got nothing on the tracy span stack!",
                    0xFF000000,
                    self.stack_depth,
                );
            }
        });
    }

    fn on_event(&self, event: &Event, _: Context<'_, S>) {
        let mut visitor = TracyEventFieldVisitor {
            dest: String::new(),
            first: true,
            frame_mark: false,
        };
        event.record(&mut visitor);
        if !visitor.first {
            let mut max_len = usize::from(u16::max_value()) - 1;
            if visitor.dest.len() >= max_len {
                while !visitor.dest.is_char_boundary(max_len) {
                    max_len -= 1;
                }
                message(&visitor.dest[..max_len], self.stack_depth);
                color_message(
                    "Message for the previous event was too long, truncated",
                    0xFF000000,
                    self.stack_depth,
                );
            } else {
                message(&visitor.dest, self.stack_depth);
            }
        }
        if visitor.frame_mark {
            finish_continuous_frame!();
        }
    }
}

struct TracyEventFieldVisitor {
    dest: String,
    frame_mark: bool,
    first: bool,
}

impl Visit for TracyEventFieldVisitor {
    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
        // FIXME: this is a very crude formatter, but we don’t have
        // an easy way to do anything better...
        if self.first {
            let _ = write!(&mut self.dest, "{} = {:?}", field.name(), value);
            self.first = false;
        } else {
            let _ = write!(&mut self.dest, ", {} = {:?}", field.name(), value);
        }
    }

    fn record_bool(&mut self, field: &Field, value: bool) {
        match (value, field.name()) {
            (true, "tracy.frame_mark") => self.frame_mark = true,
            _ => self.record_debug(field, &value),
        }
    }
}

#[cfg(test)]
mod tests {
    use futures::future::join_all;
    use tracing_attributes::instrument;
    use tracing::{debug, event, info, span, info_span, Level};
    use tracing_subscriber::layer::SubscriberExt;

    fn setup_subscriber() {
        static ONCE: std::sync::Once = std::sync::Once::new();
        ONCE.call_once(|| {
            tracing::subscriber::set_global_default(
                tracing_subscriber::registry().with(super::TracyLayer::new()),
            )
            .unwrap();
        });
    }

    #[test]
    fn it_works() {
        setup_subscriber();
        let span = span!(Level::TRACE, "a sec");
        let _enter = span.enter();
        event!(Level::INFO, "EXPLOSION!");
    }

    #[test]
    fn it_works_2() {
        setup_subscriber();
        let span = span!(Level::TRACE, "2 secs");
        let _enter = span.enter();
        event!(
            Level::INFO,
            message = "DOUBLE THE EXPLOSION!",
            tracy.frame_mark = true
        );
    }

    #[test]
    fn multiple_entries() {
        setup_subscriber();
        let span = span!(Level::INFO, "multiple_entries");
        span.in_scope(|| {});
        span.in_scope(|| {});

        let span = span!(Level::INFO, "multiple_entries 2");
        span.in_scope(|| {
            span.in_scope(|| {})
        });
    }

    #[test]
    fn out_of_order() {
        setup_subscriber();
        let span1 = span!(Level::INFO, "out of order exits 1");
        let span2 = span!(Level::INFO, "out of order exits 2");
        let span3 = span!(Level::INFO, "out of order exits 3");
        let entry1 = span1.enter();
        let entry2 = span2.enter();
        let entry3 = span3.enter();
        drop(entry2);
        drop(entry3);
        drop(entry1);
    }

    #[test]
    fn exit_in_different_thread() {
        setup_subscriber();
        let span = Box::leak(Box::new(span!(Level::INFO, "exit in different thread")));
        let entry = span.enter();
        let thread = std::thread::spawn(|| drop(entry));
        thread.join().unwrap();
    }

    #[instrument]
    async fn parent_task(subtasks: usize) {
        info!("spawning subtasks...");
        let subtasks = (1..=subtasks)
            .map(|number| {
                debug!(message = "creating subtask;", number);
                subtask(number)
            })
            .collect::<Vec<_>>();

        let result = join_all(subtasks).await;

        debug!("all subtasks completed");
        let sum: usize = result.into_iter().sum();
        info!(sum);
    }

    #[instrument]
    async fn subtask(number: usize) -> usize {
        info!("sleeping in subtask {}...", number);
        tokio::time::delay_for(std::time::Duration::from_millis(10)).await;
        info!("sleeping in subtask {}...", number);
        tokio::time::delay_for(std::time::Duration::from_millis(number as _)).await;
        info!("sleeping in subtask {}...", number);
        tokio::time::delay_for(std::time::Duration::from_millis(10)).await;
        number
    }

    // Test based on the spawny_things example from the tracing repository.
    #[tokio::test]
    async fn async_futures() {
        setup_subscriber();
        parent_task(5).await;
    }

    #[test]
    fn message_too_long() {
        setup_subscriber();
        info!("{}", "a".repeat(u16::max_value().into()));
    }

    #[test]
    fn long_span_data() {
        setup_subscriber();
        let data = "c".repeat(u16::max_value().into());
        info_span!("some span name", "{}", data).in_scope(|| {});
    }
}