cido 0.2.0

//! This should be made into it's own crate or use something that is already public
use core::future::Future;
use std::{
  cell::{Cell, RefCell},
  collections::HashMap,
  pin::Pin,
  task::{Context, Poll},
  time::{Duration, Instant},
};

use arrayvec::ArrayVec;
use serde::{Deserialize, Serialize};

pub(crate) const METRICS_DEPTH: usize = 8;

#[derive(PartialEq, Eq, Hash, Clone, Debug, PartialOrd, Ord, Serialize)]
#[serde(transparent)]
pub struct MetricsDepthArrayVec {
  inner: ArrayVec<&'static str, METRICS_DEPTH>,
}

impl MetricsDepthArrayVec {
  #[inline]
  pub fn into_inner(self) -> ArrayVec<&'static str, METRICS_DEPTH> {
    self.inner
  }
}

impl std::ops::Deref for MetricsDepthArrayVec {
  type Target = ArrayVec<&'static str, METRICS_DEPTH>;

  fn deref(&self) -> &Self::Target {
    &self.inner
  }
}

impl std::ops::DerefMut for MetricsDepthArrayVec {
  fn deref_mut(&mut self) -> &mut Self::Target {
    &mut self.inner
  }
}

impl From<ArrayVec<&'static str, METRICS_DEPTH>> for MetricsDepthArrayVec {
  fn from(v: ArrayVec<&'static str, METRICS_DEPTH>) -> Self {
    Self { inner: v }
  }
}

impl From<MetricsDepthArrayVec> for ArrayVec<&'static str, METRICS_DEPTH> {
  fn from(v: MetricsDepthArrayVec) -> Self {
    v.inner
  }
}

#[derive(Serialize, Clone, Debug, Default)]
#[serde(transparent)]
pub struct MetricMap {
  pub inner: HashMap<&'static str, FutureMetrics>,
}

impl std::ops::Deref for MetricMap {
  type Target = HashMap<&'static str, FutureMetrics>;

  fn deref(&self) -> &Self::Target {
    &self.inner
  }
}

impl std::ops::DerefMut for MetricMap {
  fn deref_mut(&mut self) -> &mut Self::Target {
    &mut self.inner
  }
}

impl MetricMap {
  fn with_capacity(cap: usize) -> Self {
    Self {
      inner: HashMap::with_capacity(cap),
    }
  }
}

pub type MetricStackMap = HashMap<MetricsDepthArrayVec, FutureMetrics>;

pub fn take_metrics() -> (Option<MetricMap>, Option<MetricStackMap>) {
  STACK.with(|s| {
    let stack = s.inner.borrow();
    if stack.len() > 1 {
      panic!("take_metrics should only be called from a MeterRoot future");
    }
  });
  let metrics = METRICS.with(|m| {
    let mut map = m.borrow_mut();
    let mut new_map = Some(MetricMap::with_capacity(
      map.as_ref().map_or(0, |x| x.len()) + 8,
    ));
    std::mem::swap(&mut *map, &mut new_map);
    new_map
  });

  let metrics_stack = METRICS_STACK.with(|m| {
    let mut map = m.borrow_mut();
    let mut new_map = Some(HashMap::with_capacity(
      map.as_ref().map_or(0, HashMap::len) + 8,
    ));
    std::mem::swap(&mut *map, &mut new_map);
    new_map
  });

  (metrics, metrics_stack)
}

/// Allows for measuring execution time for each future individually and within a call stack
///
/// We don't expose these metrics yet, but will in the future
///
/// This trait should be extracted into it's own library or we need to migrate to a library
/// that provides something similar
pub trait FutureMeter<'a>: Sized {
  /// Add an entry in the call stack for this future
  fn meter(self, name: &'static str) -> Meter<Self> {
    self.meter_with_logging(name, Duration::ZERO)
  }

  /// Similar to `meter_with_logging` but specifies interval in milliseconds
  fn meter_ms(self, name: &'static str, log_interval_ms: u64) -> Meter<Self> {
    self.meter_with_logging(name, Duration::from_millis(log_interval_ms))
  }

  /// Add an entry in the call stack for this future and log when it has been running for more than
  /// `log_interval` and when it finishes after `log_interval`
  fn meter_with_logging(self, name: &'static str, log_interval: Duration) -> Meter<Self> {
    let now = Instant::now();
    // let name = ArrayString::from(&name[..NAME_LEN]).unwrap();
    Meter {
      new: true,
      name,
      created: now,
      last_logged: now,
      log_interval,
      metrics: FutureMetrics::new(now),
      f: self,
    }
  }

  /// Like `meter` but creates a root that other futures can use
  ///
  /// Multiple roots are ok to use, but creating a new one will disassociate the existing call stack
  ///
  /// This needs to be used for all futures that are at the base of the call stack that will be measured.
  fn meter_root(self, name: &'static str) -> MeterRoot<Self> {
    self.meter_root_with_logging(name, Duration::ZERO)
  }

  /// Like `meter_ms` but creates a root that other futures can use
  ///
  /// Multiple roots are ok to use, but creating a new one will disassociate the existing call stack
  ///
  /// This needs to be used for all futures that are at the base of the call stack that will be measured.
  fn meter_root_ms(self, name: &'static str, log_interval_ms: u64) -> MeterRoot<Self> {
    self.meter_root_with_logging(name, Duration::from_millis(log_interval_ms))
  }

  /// Like `meter_with_logging` but creates a root that other futures can use
  ///
  /// Multiple roots are ok to use, but creating a new one will disassociate the existing call stack
  ///
  /// This needs to be used for all futures that are at the base of the call stack that will be measured.
  fn meter_root_with_logging(self, name: &'static str, log_interval: Duration) -> MeterRoot<Self> {
    MeterRoot {
      metrics_stack: Some(HashMap::with_capacity(32)),
      metrics: Some(MetricMap::with_capacity(32)),
      has_own_metrics: true,
      inner: self.meter_with_logging(name, log_interval),
    }
  }
}

impl<'a, F: Future + Sized> FutureMeter<'a> for F {}

pub struct MeterRoot<F> {
  metrics_stack: Option<MetricStackMap>,
  metrics: Option<MetricMap>,
  has_own_metrics: bool,
  inner: Meter<F>,
}

pub struct MeterRootOutput<O> {
  metrics_stack: MetricStackMap,
  metrics: MetricMap,
  inner: O,
}

impl<O> MeterRootOutput<O> {
  /// Use this if you do not want to aggregate values into the existing MeterRoot
  pub fn into_inner(self) -> O {
    self.inner
  }
  /// Use this if you do want to aggregate values into the existing MeterRoot from this one
  pub fn aggregate_values(self) -> O {
    METRICS.with(|m| {
      let mut ref_mut = m.borrow_mut();
      let Some(map) = ref_mut.as_mut() else {
        return;
      };
      for (key, val) in self.metrics.inner {
        match map.entry(key) {
          std::collections::hash_map::Entry::Occupied(o) => {
            o.into_mut().sum(val);
          }
          std::collections::hash_map::Entry::Vacant(v) => {
            v.insert(val);
          }
        }
      }
    });
    METRICS_STACK.with(|m| {
      let mut ref_mut = m.borrow_mut();
      let Some(map) = ref_mut.as_mut() else {
        return;
      };
      for (key, val) in self.metrics_stack {
        match map.entry(key) {
          std::collections::hash_map::Entry::Occupied(o) => {
            o.into_mut().sum(val);
          }
          std::collections::hash_map::Entry::Vacant(v) => {
            v.insert(val);
          }
        }
      }
    });
    self.inner
  }
}

pub struct Meter<F> {
  new: bool,
  name: &'static str,
  created: Instant,
  last_logged: Instant,
  log_interval: Duration,
  metrics: FutureMetrics,
  f: F,
}

#[serde_with::serde_as]
#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
pub struct FutureMetrics {
  entered: u64,
  called: u64,
  #[serde(skip)]
  recent_aggregation: Option<Instant>,
  #[serde_as(as = "serde_with::DurationNanoSeconds")]
  processing: Duration,
  #[serde_as(as = "serde_with::DurationNanoSeconds")]
  total: Duration,
}

impl FutureMetrics {
  fn new(start: Instant) -> Self {
    Self {
      entered: 0,
      called: 0,
      recent_aggregation: Some(start),
      processing: Duration::ZERO,
      total: Duration::ZERO,
    }
  }
  pub fn total(&self) -> Duration {
    self.total
  }
  fn aggregate(&mut self, now: Instant, processing_start: Instant, is_ready: bool, new: bool) {
    self.entered += 1;
    self.called += u64::from(new);
    self.processing += now - processing_start;
    self.total += now - self.recent_aggregation.unwrap_or(processing_start);
    self.recent_aggregation = (!is_ready).then_some(now);
  }
  fn sum(&mut self, other: Self) {
    self.entered += other.entered;
    self.called += other.called;
    self.recent_aggregation = self.recent_aggregation.max(other.recent_aggregation);
    self.total += other.total;
    self.processing += other.processing;
  }
}

impl core::fmt::Display for FutureMetrics {
  fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
    let Self {
      entered,
      called,
      processing,
      total,
      recent_aggregation: _,
    } = *self;
    write!(
      f,
      "called: {called:9}, entered: {entered:9}, processing: {processing:?}, total: {total:?}"
    )
  }
}

impl<F> MeterRoot<F> {
  fn swap(&mut self, assert: bool) {
    METRICS.with(|m| {
      let mut map = m.borrow_mut();
      if assert {
        debug_assert!(
          self.metrics.is_some(),
          "MeterRoot should have metrics when not being polled"
        );
      }
      std::mem::swap(&mut self.metrics, &mut *map);
    });
    METRICS_STACK.with(|m| {
      let mut map = m.borrow_mut();
      if assert {
        debug_assert!(
          self.metrics_stack.is_some(),
          "MeterRoot should have metrics_stack when not being polled"
        );
      }
      std::mem::swap(&mut self.metrics_stack, &mut *map);
    });
    self.has_own_metrics = !self.has_own_metrics;
  }
}

impl<F> Drop for MeterRoot<F> {
  fn drop(&mut self) {
    if !self.has_own_metrics {
      self.swap(false);
    }
  }
}

impl<F: Future> Future for MeterRoot<F> {
  type Output = MeterRootOutput<<F as Future>::Output>;

  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    let this = unsafe { self.get_unchecked_mut() };
    POLL_START.with(|p| p.set(Instant::now()));
    this.swap(true);
    let f = unsafe { Pin::new_unchecked(&mut this.inner) };
    let poll = f.poll(cx);
    // cleanup no matter what the result of poll is.
    this.swap(false);
    poll.map(|o| MeterRootOutput {
      metrics_stack: this.metrics_stack.take().unwrap_or_default(),
      metrics: this.metrics.take().unwrap_or_default(),
      inner: o,
    })
  }
}

impl<F: Future> Future for Meter<F> {
  type Output = <F as Future>::Output;
  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    // Safety: we never move this
    let this = unsafe { self.get_unchecked_mut() };
    let new = this.new;
    this.new = false;
    let guard = STACK.with(|s| s.push(this.name));
    // Safety: we never move f
    let f = unsafe { Pin::new_unchecked(&mut this.f) };
    match f.poll(cx) {
      Poll::Pending => {
        STACK.with(|stack| {
          if stack.inner.borrow().len() == guard.index + 1 {
            // this is the one initiating the pending call, so set the instant
            PENDING_RETURNED.with(|i| i.set(Instant::now()))
          }
        });
        let now = PENDING_RETURNED.with(Cell::get);
        let poll_start = POLL_START.with(Cell::get).max(this.created);
        this.metrics.aggregate(now, poll_start, false, new);
        if this.log_interval > Duration::ZERO && now - this.last_logged >= this.log_interval {
          this.last_logged = now;
          STACK.with(|stack| {
            let stack = stack.inner.borrow();
            tracing::trace!(
              stack = ?stack.as_slice(),
              total = ?this.metrics.total,
              processing = ?this.metrics.processing,
              entered = %this.metrics.entered,
              "Future working",
            );
          })
        }
        METRICS.with(|m| {
          let mut map = m.borrow_mut();
          let Some(map) = map.as_mut() else {
            return;
          };
          let metrics = map
            .entry(this.name)
            .or_insert_with(|| FutureMetrics::new(this.created));
          metrics.aggregate(now, poll_start, false, new);
        });
        // skip root
        if guard.index > 0 {
          METRICS_STACK.with(|m| {
            let mut map = m.borrow_mut();
            let Some(map) = map.as_mut() else {
              return;
            };
            STACK.with(|stack| {
              let stack = stack.inner.borrow();
              let array = stack[stack.len().saturating_sub(METRICS_DEPTH)..]
                .iter()
                .copied()
                .collect::<ArrayVec<_, METRICS_DEPTH>>();
              let metrics = map
                .entry(array.into())
                .or_insert_with(|| FutureMetrics::new(this.created));
              metrics.aggregate(now, poll_start, false, new);
            });
          });
        }
        Poll::Pending
      }
      Poll::Ready(o) => {
        let now = Instant::now();
        let poll_start = this.created.max(POLL_START.with(Cell::get));
        this.metrics.aggregate(now, poll_start, true, new);
        if now - this.last_logged > this.log_interval {
          STACK.with(|stack| {
            let stack = stack.inner.borrow();
            tracing::debug!(
              stack = ?stack.as_slice(),
              total = ?this.metrics.total,
              processing = ?this.metrics.processing,
              entered = %this.metrics.entered,
              "Future complete",
            );
          });
        }
        METRICS.with(|m| {
          let mut map = m.borrow_mut();
          let Some(map) = map.as_mut() else {
            return;
          };
          let metrics = map
            .entry(this.name)
            .or_insert_with(|| FutureMetrics::new(this.created));
          metrics.aggregate(now, poll_start, true, new);
        });
        // skip root
        if guard.index > 0 {
          METRICS_STACK.with(|m| {
            let mut map = m.borrow_mut();
            let Some(map) = map.as_mut() else {
              return;
            };
            STACK.with(|stack| {
              let stack = stack.inner.borrow();
              let array = stack[stack.len().saturating_sub(METRICS_DEPTH)..]
                .iter()
                .copied()
                .collect::<ArrayVec<_, METRICS_DEPTH>>();
              let metrics = map
                .entry(array.into())
                .or_insert_with(|| FutureMetrics::new(this.created));
              metrics.aggregate(now, poll_start, true, new);
            });
          });
        }
        Poll::Ready(o)
      }
    }
  }
}

thread_local! {
  static STACK: OwnedStack<&'static str> = OwnedStack::with_capacity(20);
  static POLL_START: Cell<Instant> = Cell::new(Instant::now());
  static PENDING_RETURNED: Cell<Instant> = Cell::new(Instant::now());
  static METRICS_STACK: RefCell<Option<MetricStackMap>> = const { RefCell::new(None) };
  static METRICS: RefCell<Option<MetricMap>> = const { RefCell::new(None) };
}

pub struct OwnedStackGuard<T> {
  index: usize,
  _marker: std::marker::PhantomData<fn() -> T>,
}

impl<T> Drop for OwnedStackGuard<T> {
  fn drop(&mut self) {
    STACK.with(|stack| {
      let mut stack = stack.inner.borrow_mut();
      assert_eq!(stack.len(), self.index + 1);
      let _ = stack.pop();
    });
  }
}

#[derive(Default, Debug)]
#[must_use]
pub struct OwnedStack<T> {
  inner: RefCell<Vec<T>>,
}

impl<T> OwnedStack<T> {
  pub fn with_capacity(capacity: usize) -> Self {
    Self {
      inner: RefCell::new(Vec::with_capacity(capacity)),
    }
  }
  pub fn push(&self, t: T) -> OwnedStackGuard<T> {
    let mut inner = self.inner.borrow_mut();
    let index = inner.len();
    inner.push(t);
    OwnedStackGuard {
      index,
      _marker: std::marker::PhantomData,
    }
  }
}