use core::cell::RefCell;
use std::sync::Arc;
use std::time::Instant;
use super::context;
use super::shard::Shared;
use super::tuning::{TaskHooks, TaskMeta};
struct Note {
previous: Option<u64>,
began: Option<(Instant, *const Shared)>,
}
std::thread_local! {
static POLLS: RefCell<Vec<Note>> = const { RefCell::new(Vec::new()) };
}
pub(crate) fn compose(hooks: &mut TaskHooks, long_poll: bool, latencies: bool) {
if !long_poll && !latencies {
return;
}
let (before, after) = (hooks.before_poll.take(), hooks.after_poll.take());
hooks.before_poll = Some(Arc::new(move |meta: &TaskMeta<'_>| {
if let Some(before) = &before {
before(meta);
}
begin(meta.id().0.get(), long_poll, latencies);
}));
hooks.after_poll = Some(Arc::new(move |meta: &TaskMeta<'_>| {
end();
if let Some(after) = &after {
after(meta);
}
}));
}
fn begin(id: u64, long_poll: bool, latencies: bool) {
let previous = if long_poll {
context::with_slot(|slot| slot.map(|slot| slot.note_task(id)))
} else {
None
};
let began = if latencies {
context::current(|shard| shard.map(|s| (Instant::now(), Arc::as_ptr(&s.shared))))
} else {
None
};
let _ = POLLS.try_with(|polls| polls.borrow_mut().push(Note { previous, began }));
}
fn end() {
let Some(note) = POLLS
.try_with(|polls| polls.borrow_mut().pop())
.ok()
.flatten()
else {
return;
};
if let Some(previous) = note.previous {
context::with_slot(|slot| slot.map(|slot| slot.note_task(previous)));
}
if let Some((began, shared)) = note.began {
unsafe { &*shared }.counters.poll_time(began.elapsed());
}
}