use std::{
hint::spin_loop,
sync::{
Arc,
atomic::{
AtomicU64,
Ordering,
},
},
thread,
};
use qubit_fast_cas::CasCell;
#[cfg(feature = "serde")]
use serde::{
Deserialize,
Deserializer,
};
#[cfg(feature = "serde")]
use crate::validation::{
validate_metrics,
validate_snapshot_counts,
};
use crate::{
MetricError,
internal::OperationState,
};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Metric {
pub(crate) id: Arc<str>,
pub(crate) name: Arc<str>,
pub(crate) total: Option<u64>,
}
impl Metric {
#[must_use]
pub fn new(id: &str, name: &str) -> Self {
Self {
id: Arc::from(id),
name: Arc::from(name),
total: None,
}
}
#[must_use]
pub const fn total(mut self, total: u64) -> Self {
self.total = Some(total);
self
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub const fn configured_total(&self) -> Option<u64> {
self.total
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct MetricDelta {
started: u64,
unclassified: u64,
succeeded: u64,
failed: u64,
cancelled: u64,
}
impl MetricDelta {
#[must_use]
pub const fn new() -> Self {
Self {
started: 0,
unclassified: 0,
succeeded: 0,
failed: 0,
cancelled: 0,
}
}
#[must_use]
pub const fn started(mut self, count: u64) -> Self {
self.started = count;
self
}
#[must_use]
pub const fn unclassified(mut self, count: u64) -> Self {
self.unclassified = count;
self
}
#[must_use]
pub const fn succeeded(mut self, count: u64) -> Self {
self.succeeded = count;
self
}
#[must_use]
pub const fn failed(mut self, count: u64) -> Self {
self.failed = count;
self
}
#[must_use]
pub const fn cancelled(mut self, count: u64) -> Self {
self.cancelled = count;
self
}
}
#[derive(Clone)]
pub struct MetricHandle {
inner: Arc<MetricInner>,
operation_state: Arc<OperationState>,
}
impl MetricHandle {
pub(crate) fn new(
metric: Metric,
operation_state: Arc<OperationState>,
) -> Self {
Self {
inner: Arc::new(MetricInner::new(metric)),
operation_state,
}
}
#[must_use]
pub fn id(&self) -> &str {
self.inner.metric.id()
}
#[must_use]
pub fn name(&self) -> &str {
self.inner.metric.name()
}
pub fn start(&self, count: u64) -> Result<(), MetricError> {
self.apply_delta(MetricDelta::new().started(count))
}
pub fn complete(&self, count: u64) -> Result<(), MetricError> {
self.apply_delta(MetricDelta::new().unclassified(count))
}
pub fn succeed(&self, count: u64) -> Result<(), MetricError> {
self.apply_delta(MetricDelta::new().succeeded(count))
}
pub fn fail(&self, count: u64) -> Result<(), MetricError> {
self.apply_delta(MetricDelta::new().failed(count))
}
pub fn cancel(&self, count: u64) -> Result<(), MetricError> {
self.apply_delta(MetricDelta::new().cancelled(count))
}
pub fn apply_delta(&self, delta: MetricDelta) -> Result<(), MetricError> {
let metric_id = self.id();
let total = self.inner.metric.configured_total();
let _update_guard = self.operation_state.enter_update(metric_id)?;
self.inner.with_update(|counts| {
let mut next = *counts;
apply_delta_to_counts(&mut next, delta, metric_id)?;
next.validate(metric_id, total)?;
*counts = next;
Ok(())
})
}
#[must_use]
pub fn snapshot(&self) -> MetricSnapshot {
let counts = self.inner.snapshot_counts();
MetricSnapshot::from_counts(&self.inner.metric, counts)
}
}
struct MetricInner {
metric: Metric,
gate: CasCell,
active: AtomicU64,
completed_unclassified: AtomicU64,
succeeded: AtomicU64,
failed: AtomicU64,
cancelled: AtomicU64,
}
impl MetricInner {
fn new(metric: Metric) -> Self {
Self {
metric,
gate: CasCell::new(0),
active: AtomicU64::new(0),
completed_unclassified: AtomicU64::new(0),
succeeded: AtomicU64::new(0),
failed: AtomicU64::new(0),
cancelled: AtomicU64::new(0),
}
}
fn with_update<R, F>(&self, mut update: F) -> Result<R, MetricError>
where
F: FnMut(&mut MetricCounts) -> Result<R, MetricError>,
{
let mut attempts = 0;
loop {
let version = self.gate.load();
if version & 1 != 0 {
wait_for_contention(attempts);
attempts += 1;
continue;
}
match self.gate.compare_set(version, version.wrapping_add(1)) {
Ok(()) => {
let _guard = MetricGateGuard::new(
&self.gate,
version.wrapping_add(2),
);
let mut counts = self.read_counts();
let result = update(&mut counts);
if result.is_ok() {
self.write_counts(&counts);
}
return result;
}
Err(_) => {
wait_for_contention(attempts);
attempts += 1;
}
}
}
}
fn read_counts(&self) -> MetricCounts {
MetricCounts {
active: self.active.load(Ordering::Acquire),
completed_unclassified: self
.completed_unclassified
.load(Ordering::Acquire),
succeeded: self.succeeded.load(Ordering::Acquire),
failed: self.failed.load(Ordering::Acquire),
cancelled: self.cancelled.load(Ordering::Acquire),
}
}
fn write_counts(&self, counts: &MetricCounts) {
self.active.store(counts.active, Ordering::Release);
self.completed_unclassified
.store(counts.completed_unclassified, Ordering::Release);
self.succeeded.store(counts.succeeded, Ordering::Release);
self.failed.store(counts.failed, Ordering::Release);
self.cancelled.store(counts.cancelled, Ordering::Release);
}
fn snapshot_counts(&self) -> MetricCounts {
let mut attempts = 0;
loop {
let start = self.gate.load();
if start & 1 != 0 {
wait_for_contention(attempts);
attempts += 1;
continue;
}
let counts = self.read_counts();
if start == self.gate.load() {
return counts;
}
wait_for_contention(attempts);
attempts += 1;
}
}
}
#[derive(Clone, Copy)]
struct MetricCounts {
active: u64,
completed_unclassified: u64,
succeeded: u64,
failed: u64,
cancelled: u64,
}
impl MetricCounts {
fn completed(self) -> Option<u64> {
self.completed_unclassified
.checked_add(self.succeeded)?
.checked_add(self.failed)?
.checked_add(self.cancelled)
}
fn occupied(self) -> Option<u64> {
self.completed()?.checked_add(self.active)
}
fn validate(
self,
metric_id: &str,
total: Option<u64>,
) -> Result<(), MetricError> {
let occupied =
self.occupied().ok_or_else(|| MetricError::CountOverflow {
metric_id: metric_id.into(),
})?;
if let Some(total) = total
&& occupied > total
{
return Err(MetricError::TotalExceeded {
metric_id: metric_id.into(),
total,
attempted: occupied,
});
}
Ok(())
}
}
fn apply_delta_to_counts(
counts: &mut MetricCounts,
delta: MetricDelta,
metric_id: &str,
) -> Result<(), MetricError> {
let terminal_delta = delta
.unclassified
.checked_add(delta.succeeded)
.and_then(|value| value.checked_add(delta.failed))
.and_then(|value| value.checked_add(delta.cancelled))
.ok_or_else(|| MetricError::CountOverflow {
metric_id: metric_id.into(),
})?;
let available_active = counts
.active
.checked_add(delta.started)
.ok_or_else(|| MetricError::CountOverflow {
metric_id: metric_id.into(),
})?;
if terminal_delta > available_active {
return Err(MetricError::InsufficientActive {
metric_id: metric_id.into(),
requested: terminal_delta,
available: available_active,
});
}
counts.active = available_active - terminal_delta;
counts.completed_unclassified = counts
.completed_unclassified
.checked_add(delta.unclassified)
.ok_or_else(|| MetricError::CountOverflow {
metric_id: metric_id.into(),
})?;
counts.succeeded = counts
.succeeded
.checked_add(delta.succeeded)
.ok_or_else(|| MetricError::CountOverflow {
metric_id: metric_id.into(),
})?;
counts.failed =
counts.failed.checked_add(delta.failed).ok_or_else(|| {
MetricError::CountOverflow {
metric_id: metric_id.into(),
}
})?;
counts.cancelled = counts
.cancelled
.checked_add(delta.cancelled)
.ok_or_else(|| MetricError::CountOverflow {
metric_id: metric_id.into(),
})?;
Ok(())
}
struct MetricGateGuard<'gate> {
gate: &'gate CasCell,
next_version: u64,
}
impl<'gate> MetricGateGuard<'gate> {
fn new(gate: &'gate CasCell, next_version: u64) -> Self {
Self { gate, next_version }
}
}
impl Drop for MetricGateGuard<'_> {
fn drop(&mut self) {
self.gate.store(self.next_version);
}
}
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MetricSnapshot {
id: Arc<str>,
name: Arc<str>,
total: Option<u64>,
completed: u64,
active: u64,
succeeded: u64,
failed: u64,
cancelled: u64,
}
impl MetricSnapshot {
fn from_counts(metric: &Metric, counts: MetricCounts) -> Self {
Self {
id: Arc::clone(&metric.id),
name: Arc::clone(&metric.name),
total: metric.total,
completed: counts
.completed()
.expect("validated metric counts must fit in u64"),
active: counts.active,
succeeded: counts.succeeded,
failed: counts.failed,
cancelled: counts.cancelled,
}
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub const fn total(&self) -> Option<u64> {
self.total
}
#[must_use]
pub const fn completed(&self) -> u64 {
self.completed
}
#[must_use]
pub const fn unclassified(&self) -> u64 {
let classified = self
.succeeded
.saturating_add(self.failed)
.saturating_add(self.cancelled);
self.completed.saturating_sub(classified)
}
#[must_use]
pub const fn active(&self) -> u64 {
self.active
}
#[must_use]
pub const fn succeeded(&self) -> u64 {
self.succeeded
}
#[must_use]
pub const fn failed(&self) -> u64 {
self.failed
}
#[must_use]
pub const fn cancelled(&self) -> u64 {
self.cancelled
}
#[must_use]
pub fn completion_fraction(&self) -> Option<f64> {
self.total
.filter(|total| *total > 0)
.map(|total| self.completed as f64 / total as f64)
}
}
#[cfg(feature = "serde")]
#[derive(Deserialize)]
struct MetricSnapshotWire {
id: Arc<str>,
name: Arc<str>,
total: Option<u64>,
completed: u64,
active: u64,
succeeded: u64,
failed: u64,
cancelled: u64,
}
#[cfg(feature = "serde")]
impl<'de> Deserialize<'de> for MetricSnapshot {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = MetricSnapshotWire::deserialize(deserializer)?;
let snapshot = Self {
id: wire.id,
name: wire.name,
total: wire.total,
completed: wire.completed,
active: wire.active,
succeeded: wire.succeeded,
failed: wire.failed,
cancelled: wire.cancelled,
};
let metric = Metric {
id: Arc::clone(&snapshot.id),
name: Arc::clone(&snapshot.name),
total: snapshot.total,
};
validate_metrics(std::slice::from_ref(&metric))
.map_err(serde::de::Error::custom)?;
validate_snapshot_counts(&snapshot)
.map_err(serde::de::Error::custom)?;
Ok(snapshot)
}
}
#[inline]
fn wait_for_contention(attempts: usize) {
if attempts > 0 && attempts.is_multiple_of(16) {
thread::yield_now();
} else {
spin_loop();
}
}