use serde::Deserialize;
use snap::read::FrameDecoder;
use std::io::{self, Read};
use std::time::Duration;
#[cfg(feature = "profile")]
use {
bumpalo::Bump as Arena,
chrono::offset::Utc,
flume::Sender,
once_cell::sync::Lazy,
serde::Serialize,
snap::write::FrameEncoder,
std::env,
std::fs::File,
std::io::Write,
std::mem,
std::sync::Mutex,
std::thread::{self, JoinHandle},
std::time::Instant,
};
#[cfg(feature = "profile")]
static PROFILER: Lazy<Profiler> = Lazy::new(|| {
let (sender, reciever) = flume::unbounded();
let thread = thread::Builder::new()
.name("wtf-profiler".to_string())
.spawn(move || {
fn create_file() -> Option<FrameEncoder<File>> {
let program_name = env::current_exe().ok()?;
let program_name = program_name.file_name()?.to_str()?;
let timestamp = Utc::now().format("%F-%T");
let file = File::create(format!("{}-{}.wtf", program_name, timestamp)).ok()?;
Some(FrameEncoder::new(file))
}
let mut file = create_file().expect("WTF: Failed to create file for profile");
#[derive(Serialize)]
struct TaskDataS<'a> {
name: &'a str,
duration: Duration,
subtasks: Vec<&'a Self>,
}
let mut task_arena = Arena::with_capacity(mem::size_of::<TaskDataS>() * 1000);
let mut frame_number: usize = 0;
let mut frame = TaskDataS {
name: "",
duration: Duration::default(),
subtasks: Vec::with_capacity(10),
};
let mut parent_stack: Vec<*mut TaskDataS> = vec![&mut frame];
loop {
let msg = reciever.recv_timeout(Duration::from_millis(100));
match msg {
Ok(ProfilerMessage::TaskStart { name }) => {
let task = TaskDataS {
name,
duration: Duration::default(),
subtasks: Vec::new(),
};
let task_mut_ref: *mut TaskDataS = task_arena.alloc(task);
let task_ref = unsafe { &*task_mut_ref };
let parent = *parent_stack.last().unwrap();
let parent_subtasks = unsafe { &mut (*parent).subtasks };
parent_subtasks.push(task_ref);
parent_stack.push(task_mut_ref);
}
Ok(ProfilerMessage::TaskEnd { elapsed }) => {
let task = *parent_stack.last().unwrap();
unsafe { (*task).duration = elapsed }
if parent_stack.len() == 1 {
frame_number += 1;
let frame_name = format!("Frame #{}", frame_number);
frame.name = unsafe { mem::transmute(frame_name.as_str()) };
bincode::serialize_into(&mut file, &frame)
.expect("WTF: Failed to write data to file");
frame.name = "";
frame.subtasks.clear();
task_arena.reset();
} else {
parent_stack.pop().unwrap();
}
}
_ => {
if PROFILER
.thread
.lock()
.expect("WTF: Failed to acquire thread lock")
.is_none()
{
break;
}
}
}
}
file.flush().expect("WTF: Failed to write data to file");
})
.expect("WTF: Failed to spawn data writing thread");
let thread = Mutex::new(Some(thread));
Profiler { sender, thread }
});
pub struct Profiler {
#[cfg(feature = "profile")]
sender: Sender<ProfilerMessage>,
#[cfg(feature = "profile")]
thread: Mutex<Option<JoinHandle<()>>>,
}
#[cfg(feature = "profile")]
#[doc(hidden)]
pub type ProfilingReturnType = TaskRecord;
#[cfg(not(feature = "profile"))]
#[doc(hidden)]
pub type ProfilingReturnType = ();
impl Profiler {
#[must_use = "Must assign to a variable: \"_profile = Profiler::new_frame()\""]
pub fn new_frame() -> ProfilingReturnType {
#[cfg(feature = "profile")]
TaskRecord {
start: Instant::now(),
}
}
#[must_use = "Must assign to a variable: \"_profile = Profiler::profile_task()\""]
#[allow(unused_variables)]
pub fn profile_task(name: &'static str) -> ProfilingReturnType {
#[cfg(feature = "profile")]
{
PROFILER
.sender
.send(ProfilerMessage::TaskStart { name })
.expect("WTF: Failed to send task across a thread");
TaskRecord {
start: Instant::now(),
}
}
}
pub fn end_profiling() {
#[cfg(feature = "profile")]
{
let mut profiler_lock = PROFILER
.thread
.lock()
.expect("WTF: Failed to acquire thread lock");
let thread = profiler_lock
.take()
.expect("WTF: Profiler::end_profiling() has already been called once");
drop(profiler_lock);
thread
.join()
.expect("WTF: Failed to join data writing thread");
}
}
}
pub type ProfileData = Box<[TaskData]>;
pub fn read_profile_data<R: Read>(reader: R) -> Result<ProfileData, bincode::Error> {
let mut reader = FrameDecoder::new(reader);
let mut frames = Vec::new();
loop {
let frame = bincode::deserialize_from(&mut reader);
let frame = match frame.map_err(|err| *err) {
Ok(task) => task,
Err(bincode::ErrorKind::Io(err)) if err.kind() == io::ErrorKind::UnexpectedEof => break,
Err(err) => return Err(Box::new(err)),
};
frames.push(frame);
}
Ok(frames.into_boxed_slice())
}
#[derive(Clone, Deserialize)]
pub struct TaskData {
pub name: Box<str>,
pub duration: Duration,
pub subtasks: Box<[Self]>,
}
#[cfg(feature = "profile")]
enum ProfilerMessage {
TaskStart { name: &'static str },
TaskEnd { elapsed: Duration },
}
#[cfg(feature = "profile")]
pub struct TaskRecord {
start: Instant,
}
#[cfg(feature = "profile")]
impl Drop for TaskRecord {
fn drop(&mut self) {
let msg = ProfilerMessage::TaskEnd {
elapsed: self.start.elapsed(),
};
PROFILER
.sender
.send(msg)
.expect("WTF: Failed to send task across a thread");
}
}