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

use crate::file_header::{write_file_header, FILE_MAGIC_EVENT_STREAM};
use crate::raw_event::{RawEvent, Timestamp, TimestampKind};
use crate::serialization::SerializationSink;
use crate::stringtable::{SerializableString, StringId, StringTableBuilder};
use std::error::Error;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Instant;

pub struct ProfilerFiles {
    pub events_file: PathBuf,
    pub string_data_file: PathBuf,
    pub string_index_file: PathBuf,
}

impl ProfilerFiles {
    pub fn new(path_stem: &Path) -> ProfilerFiles {
        ProfilerFiles {
            events_file: path_stem.with_extension("events"),
            string_data_file: path_stem.with_extension("string_data"),
            string_index_file: path_stem.with_extension("string_index"),
        }
    }
}

pub struct Profiler<S: SerializationSink> {
    event_sink: Arc<S>,
    string_table: StringTableBuilder<S>,
    start_time: Instant,
}

impl<S: SerializationSink> Profiler<S> {
    pub fn new(path_stem: &Path) -> Result<Profiler<S>, Box<dyn Error>> {
        let paths = ProfilerFiles::new(path_stem);
        let event_sink = Arc::new(S::from_path(&paths.events_file)?);

        // The first thing in every file we generate must be the file header.
        write_file_header(&*event_sink, FILE_MAGIC_EVENT_STREAM);

        let string_table = StringTableBuilder::new(
            Arc::new(S::from_path(&paths.string_data_file)?),
            Arc::new(S::from_path(&paths.string_index_file)?),
        );

        Ok(Profiler {
            event_sink,
            string_table,
            start_time: Instant::now(),
        })
    }

    #[inline(always)]
    pub fn alloc_string_with_reserved_id<STR: SerializableString + ?Sized>(
        &self,
        id: StringId,
        s: &STR,
    ) -> StringId {
        self.string_table.alloc_with_reserved_id(id, s)
    }

    #[inline(always)]
    pub fn alloc_string<STR: SerializableString + ?Sized>(&self, s: &STR) -> StringId {
        self.string_table.alloc(s)
    }

    pub fn record_event(
        &self,
        event_kind: StringId,
        event_id: StringId,
        thread_id: u64,
        timestamp_kind: TimestampKind,
    ) {
        let duration_since_start = self.start_time.elapsed();
        let nanos_since_start = duration_since_start.as_secs() * 1_000_000_000
            + duration_since_start.subsec_nanos() as u64;
        let timestamp = Timestamp::new(nanos_since_start, timestamp_kind);

        self.event_sink
            .write_atomic(std::mem::size_of::<RawEvent>(), |bytes| {
                debug_assert_eq!(bytes.len(), std::mem::size_of::<RawEvent>());

                let raw_event = RawEvent {
                    event_kind,
                    id: event_id,
                    thread_id,
                    timestamp,
                };

                let raw_event_bytes: &[u8] = unsafe {
                    std::slice::from_raw_parts(
                        &raw_event as *const _ as *const u8,
                        std::mem::size_of::<RawEvent>(),
                    )
                };

                bytes.copy_from_slice(raw_event_bytes);
            });
    }
}