use crate::metrics::Observer;
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
enum State {
Passive,
Polling,
}
pub struct ObservedFuture<Obs, Fut>
where
Obs: Observer<Fut::Output>,
Fut: Future,
{
state: State,
observer: Obs,
inner: Fut,
}
impl<Obs, Fut> Future for ObservedFuture<Obs, Fut>
where
Obs: Observer<Fut::Output>,
Fut: Future,
{
type Output = Fut::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = unsafe { self.get_unchecked_mut() };
let inner = unsafe { Pin::new_unchecked(&mut this.inner) };
if matches!(this.state, State::Passive) {
this.state = State::Polling;
this.observer.on_first_poll();
}
this.observer.on_poll();
match inner.poll(cx) {
Poll::Ready(output) => {
this.state = State::Passive;
this.observer.on_poll_ready(&output);
this.observer.on_finish(Some(&output));
Poll::Ready(output)
}
Poll::Pending => Poll::Pending,
}
}
}
impl<Obs, Fut> Drop for ObservedFuture<Obs, Fut>
where
Obs: Observer<Fut::Output>,
Fut: Future,
{
fn drop(&mut self) {
if matches!(self.state, State::Polling) {
self.observer.on_finish(None);
}
self.observer.on_drop();
}
}
pub trait ObservedFutureExt: Future + Sized {
fn observe<Obs: Observer<Self::Output>>(self, observer: Obs) -> ObservedFuture<Obs, Self> {
ObservedFuture {
state: State::Passive,
observer,
inner: self,
}
}
}
impl<F: Future + Sized> ObservedFutureExt for F {}