#![no_std]
#![warn(missing_docs)]
use core::{future::Future, pin::Pin, task::Poll};
pub mod instruments;
pub use instruments::*;
pub trait TraceFuture: Future
where
Self: Sized,
{
fn trace_task<I: Instrument>(self, instrument: &mut I) -> TraceTaskFuture<'_, Self, I> {
TraceTaskFuture {
fut: self,
instrument,
polled_once: false,
}
}
fn trace_poll<I: Instrument>(self, instrument: &mut I) -> TracePollFuture<'_, Self, I> {
TracePollFuture {
fut: self,
instrument,
}
}
fn trace_task_and_poll<'a, I1: Instrument, I2: Instrument>(
self,
task_instrument: &'a mut I1,
poll_instrument: &'a mut I2,
) -> TraceTaskAndPollFuture<'a, Self, I1, I2> {
TraceTaskAndPollFuture {
fut: self,
task_instrument,
poll_instrument,
polled_once: false,
}
}
}
impl<F: Future> TraceFuture for F {}
pub trait Instrument {
fn on_enter(&mut self);
fn on_exit(&mut self);
}
pin_project_lite::pin_project! {
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[doc(hidden)]
pub struct TraceTaskFuture<'a, F, I>
where
F: Future,
I: Instrument,
{
#[pin]
fut: F,
instrument: &'a mut I,
polled_once: bool
}
}
impl<'p, F, P> Future for TraceTaskFuture<'p, F, P>
where
F: Future,
P: Instrument,
{
type Output = F::Output;
fn poll(
self: Pin<&mut Self>,
cx: &mut core::task::Context<'_>,
) -> core::task::Poll<Self::Output> {
let this = self.project();
if !*this.polled_once {
this.instrument.on_enter();
}
*this.polled_once = true;
let poll_result = this.fut.poll(cx);
match poll_result {
Poll::Ready(c) => {
this.instrument.on_exit();
Poll::Ready(c)
}
Poll::Pending => Poll::Pending,
}
}
}
pin_project_lite::pin_project! {
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[doc(hidden)]
pub struct TracePollFuture<'a, F, I>
where
F: Future,
I: Instrument,
{
#[pin]
fut: F,
instrument: &'a mut I,
}
}
impl<'p, F, I> Future for TracePollFuture<'p, F, I>
where
F: Future,
I: Instrument,
{
type Output = F::Output;
fn poll(
self: Pin<&mut Self>,
cx: &mut core::task::Context<'_>,
) -> core::task::Poll<Self::Output> {
let this = self.project();
this.instrument.on_enter();
let poll_result = this.fut.poll(cx);
this.instrument.on_exit();
match poll_result {
Poll::Ready(c) => Poll::Ready(c),
Poll::Pending => Poll::Pending,
}
}
}
pin_project_lite::pin_project! {
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[doc(hidden)]
pub struct TraceTaskAndPollFuture<'a, F, T, P>
where
F: Future,
T: Instrument,
P: Instrument
{
#[pin]
fut: F,
task_instrument: &'a mut T,
poll_instrument: &'a mut P,
polled_once: bool,
}
}
impl<'p, F, I1, I2> Future for TraceTaskAndPollFuture<'p, F, I1, I2>
where
F: Future,
I1: Instrument,
I2: Instrument,
{
type Output = F::Output;
fn poll(
self: Pin<&mut Self>,
cx: &mut core::task::Context<'_>,
) -> core::task::Poll<Self::Output> {
let this = self.project();
if !*this.polled_once {
this.task_instrument.on_enter();
}
*this.polled_once = true;
this.poll_instrument.on_enter();
let poll_result = this.fut.poll(cx);
this.poll_instrument.on_exit();
match poll_result {
Poll::Ready(c) => {
this.task_instrument.on_exit();
Poll::Ready(c)
}
Poll::Pending => Poll::Pending,
}
}
}