use std::fmt;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use super::{DurabilityError, DurableStore, ScanCursor, ScanPage, StoredEntry};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct StoreCounts {
pub apply_routed_batch_calls: u64,
pub apply_routed_batch_items: u64,
pub get_routed_calls: u64,
pub get_routed_rows: u64,
pub append_calls: u64,
pub append_rows: u64,
pub append_bytes: u64,
pub append_all_calls: u64,
pub append_all_rows: u64,
pub append_all_bytes: u64,
pub read_from_calls: u64,
pub read_from_rows: u64,
pub read_at_calls: u64,
pub read_at_rows: u64,
pub cas_calls: u64,
pub read_value_calls: u64,
pub read_stream_next_seq_calls: u64,
pub scan_page_calls: u64,
pub scan_page_rows: u64,
pub flushes: u64,
pub truncate_calls: u64,
pub truncate_rows: u64,
}
impl StoreCounts {
#[must_use]
pub const fn since(&self, earlier: &Self) -> Self {
Self {
apply_routed_batch_calls: self
.apply_routed_batch_calls
.saturating_sub(earlier.apply_routed_batch_calls),
apply_routed_batch_items: self
.apply_routed_batch_items
.saturating_sub(earlier.apply_routed_batch_items),
get_routed_calls: self
.get_routed_calls
.saturating_sub(earlier.get_routed_calls),
get_routed_rows: self.get_routed_rows.saturating_sub(earlier.get_routed_rows),
append_calls: self.append_calls.saturating_sub(earlier.append_calls),
append_rows: self.append_rows.saturating_sub(earlier.append_rows),
append_bytes: self.append_bytes.saturating_sub(earlier.append_bytes),
append_all_calls: self
.append_all_calls
.saturating_sub(earlier.append_all_calls),
append_all_rows: self.append_all_rows.saturating_sub(earlier.append_all_rows),
append_all_bytes: self
.append_all_bytes
.saturating_sub(earlier.append_all_bytes),
read_from_calls: self.read_from_calls.saturating_sub(earlier.read_from_calls),
read_from_rows: self.read_from_rows.saturating_sub(earlier.read_from_rows),
read_at_calls: self.read_at_calls.saturating_sub(earlier.read_at_calls),
read_at_rows: self.read_at_rows.saturating_sub(earlier.read_at_rows),
cas_calls: self.cas_calls.saturating_sub(earlier.cas_calls),
read_value_calls: self
.read_value_calls
.saturating_sub(earlier.read_value_calls),
read_stream_next_seq_calls: self
.read_stream_next_seq_calls
.saturating_sub(earlier.read_stream_next_seq_calls),
scan_page_calls: self.scan_page_calls.saturating_sub(earlier.scan_page_calls),
scan_page_rows: self.scan_page_rows.saturating_sub(earlier.scan_page_rows),
flushes: self.flushes.saturating_sub(earlier.flushes),
truncate_calls: self.truncate_calls.saturating_sub(earlier.truncate_calls),
truncate_rows: self.truncate_rows.saturating_sub(earlier.truncate_rows),
}
}
#[must_use]
pub const fn plus(&self, other: &Self) -> Self {
Self {
apply_routed_batch_calls: self
.apply_routed_batch_calls
.saturating_add(other.apply_routed_batch_calls),
apply_routed_batch_items: self
.apply_routed_batch_items
.saturating_add(other.apply_routed_batch_items),
get_routed_calls: self.get_routed_calls.saturating_add(other.get_routed_calls),
get_routed_rows: self.get_routed_rows.saturating_add(other.get_routed_rows),
append_calls: self.append_calls.saturating_add(other.append_calls),
append_rows: self.append_rows.saturating_add(other.append_rows),
append_bytes: self.append_bytes.saturating_add(other.append_bytes),
append_all_calls: self.append_all_calls.saturating_add(other.append_all_calls),
append_all_rows: self.append_all_rows.saturating_add(other.append_all_rows),
append_all_bytes: self.append_all_bytes.saturating_add(other.append_all_bytes),
read_from_calls: self.read_from_calls.saturating_add(other.read_from_calls),
read_from_rows: self.read_from_rows.saturating_add(other.read_from_rows),
read_at_calls: self.read_at_calls.saturating_add(other.read_at_calls),
read_at_rows: self.read_at_rows.saturating_add(other.read_at_rows),
cas_calls: self.cas_calls.saturating_add(other.cas_calls),
read_value_calls: self.read_value_calls.saturating_add(other.read_value_calls),
read_stream_next_seq_calls: self
.read_stream_next_seq_calls
.saturating_add(other.read_stream_next_seq_calls),
scan_page_calls: self.scan_page_calls.saturating_add(other.scan_page_calls),
scan_page_rows: self.scan_page_rows.saturating_add(other.scan_page_rows),
flushes: self.flushes.saturating_add(other.flushes),
truncate_calls: self.truncate_calls.saturating_add(other.truncate_calls),
truncate_rows: self.truncate_rows.saturating_add(other.truncate_rows),
}
}
#[must_use]
pub const fn named(&self) -> [(&'static str, u64); 22] {
[
("append_calls", self.append_calls),
("append_rows", self.append_rows),
("append_bytes", self.append_bytes),
("append_all_calls", self.append_all_calls),
("append_all_rows", self.append_all_rows),
("append_all_bytes", self.append_all_bytes),
("read_from_calls", self.read_from_calls),
("read_from_rows", self.read_from_rows),
("read_at_calls", self.read_at_calls),
("read_at_rows", self.read_at_rows),
("cas_calls", self.cas_calls),
("read_value_calls", self.read_value_calls),
(
"read_stream_next_seq_calls",
self.read_stream_next_seq_calls,
),
("scan_page_calls", self.scan_page_calls),
("scan_page_rows", self.scan_page_rows),
("flushes", self.flushes),
("truncate_calls", self.truncate_calls),
("truncate_rows", self.truncate_rows),
("apply_routed_batch_calls", self.apply_routed_batch_calls),
("apply_routed_batch_items", self.apply_routed_batch_items),
("get_routed_calls", self.get_routed_calls),
("get_routed_rows", self.get_routed_rows),
]
}
}
impl fmt::Display for StoreCounts {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
for (name, value) in self.named() {
writeln!(formatter, "{name} {value}")?;
}
Ok(())
}
}
pub trait StoreCountObserver: fmt::Debug + Send + Sync {
fn changed(&self);
fn flushed(&self, store: &str, batch_rows: Option<u64>);
}
#[derive(Debug)]
struct ObserverBinding {
name: String,
observer: Arc<dyn StoreCountObserver>,
}
static OBSERVER_INSTALL: Mutex<()> = Mutex::new(());
#[derive(Debug, Default)]
struct Counters {
apply_routed_batch_calls: AtomicU64,
apply_routed_batch_items: AtomicU64,
get_routed_calls: AtomicU64,
get_routed_rows: AtomicU64,
append_calls: AtomicU64,
append_rows: AtomicU64,
append_bytes: AtomicU64,
append_all_calls: AtomicU64,
append_all_rows: AtomicU64,
append_all_bytes: AtomicU64,
read_from_calls: AtomicU64,
read_from_rows: AtomicU64,
read_at_calls: AtomicU64,
read_at_rows: AtomicU64,
cas_calls: AtomicU64,
read_value_calls: AtomicU64,
read_stream_next_seq_calls: AtomicU64,
scan_page_calls: AtomicU64,
scan_page_rows: AtomicU64,
flushes: AtomicU64,
truncate_calls: AtomicU64,
truncate_rows: AtomicU64,
observer: OnceLock<ObserverBinding>,
}
impl Counters {
fn add(&self, counter: &AtomicU64, amount: u64) {
if amount == 0 {
return;
}
counter.fetch_add(amount, Ordering::Relaxed);
if let Some(binding) = self.observer.get() {
binding.observer.changed();
}
}
fn flushed(&self) {
if let Some(binding) = self.observer.get() {
binding.observer.flushed(&binding.name, None);
}
}
fn snapshot(&self) -> StoreCounts {
let read = |counter: &AtomicU64| counter.load(Ordering::Relaxed);
StoreCounts {
apply_routed_batch_calls: read(&self.apply_routed_batch_calls),
apply_routed_batch_items: read(&self.apply_routed_batch_items),
get_routed_calls: read(&self.get_routed_calls),
get_routed_rows: read(&self.get_routed_rows),
append_calls: read(&self.append_calls),
append_rows: read(&self.append_rows),
append_bytes: read(&self.append_bytes),
append_all_calls: read(&self.append_all_calls),
append_all_rows: read(&self.append_all_rows),
append_all_bytes: read(&self.append_all_bytes),
read_from_calls: read(&self.read_from_calls),
read_from_rows: read(&self.read_from_rows),
read_at_calls: read(&self.read_at_calls),
read_at_rows: read(&self.read_at_rows),
cas_calls: read(&self.cas_calls),
read_value_calls: read(&self.read_value_calls),
read_stream_next_seq_calls: read(&self.read_stream_next_seq_calls),
scan_page_calls: read(&self.scan_page_calls),
scan_page_rows: read(&self.scan_page_rows),
flushes: read(&self.flushes),
truncate_calls: read(&self.truncate_calls),
truncate_rows: read(&self.truncate_rows),
}
}
}
#[derive(Clone, Debug, Default)]
pub struct StoreCountsHandle {
stores: Vec<Arc<Counters>>,
}
impl StoreCountsHandle {
#[must_use]
pub fn store_count(&self) -> usize {
self.stores.len()
}
pub fn install_observer(
&self,
names: &[&str],
observer: Arc<dyn StoreCountObserver>,
) -> Result<(), DurabilityError> {
if self.stores.is_empty() || names.len() != self.stores.len() {
return Err(observer_error(
"one name is required for each counted store",
));
}
if names.iter().any(|name| name.trim().is_empty()) {
return Err(observer_error("store name is empty"));
}
for (index, counters) in self.stores.iter().enumerate() {
if self.stores[..index]
.iter()
.any(|earlier| Arc::ptr_eq(earlier, counters))
{
return Err(observer_error(
"the handle contains a duplicate counted store",
));
}
}
let installation = OBSERVER_INSTALL
.lock()
.map_err(|_| observer_error("installation lock poisoned"))?;
if self
.stores
.iter()
.any(|counters| counters.observer.get().is_some())
{
return Err(observer_error("observer already installed"));
}
for ((counters, name), observer) in self
.stores
.iter()
.zip(names)
.zip(std::iter::repeat_n(observer, self.stores.len()))
{
counters
.observer
.set(ObserverBinding {
name: (*name).to_owned(),
observer,
})
.map_err(|_| observer_error("observer installation invariant failed"))?;
}
drop(installation);
Ok(())
}
#[must_use]
pub fn with(mut self, other: &Self) -> Self {
self.stores.extend(other.stores.iter().cloned());
self
}
#[must_use]
pub fn snapshot(&self) -> StoreCounts {
self.stores
.iter()
.fold(StoreCounts::default(), |total, counters| {
total.plus(&counters.snapshot())
})
}
}
static STORES_OPENED: AtomicU64 = AtomicU64::new(0);
#[must_use]
pub fn stores_opened() -> u64 {
STORES_OPENED.load(Ordering::Relaxed)
}
pub fn record_store_opened() {
add(&STORES_OPENED, 1);
}
fn add(counter: &AtomicU64, amount: u64) {
counter.fetch_add(amount, Ordering::Relaxed);
}
fn observer_error(message: &str) -> DurabilityError {
DurabilityError::ConfigError(format!("store count observer: {message}"))
}
fn widen(amount: usize) -> u64 {
u64::try_from(amount).unwrap_or(u64::MAX)
}
#[derive(Debug)]
pub struct CountingStore<S> {
inner: S,
counters: Arc<Counters>,
}
impl<S: DurableStore> CountingStore<S> {
#[must_use]
pub fn new(inner: S) -> Self {
Self {
inner,
counters: Arc::default(),
}
}
#[must_use]
pub fn handle(&self) -> StoreCountsHandle {
StoreCountsHandle {
stores: vec![Arc::clone(&self.counters)],
}
}
#[must_use]
pub fn snapshot(&self) -> StoreCounts {
self.counters.snapshot()
}
}
#[async_trait::async_trait]
impl<S: DurableStore> DurableStore for CountingStore<S> {
async fn apply_routed_batch(
&self,
route_key: &str,
items: Vec<super::RoutedBatchItem>,
) -> Result<(), DurabilityError> {
self.counters
.add(&self.counters.apply_routed_batch_calls, 1);
let count = widen(items.len());
self.inner.apply_routed_batch(route_key, items).await?;
self.counters
.add(&self.counters.apply_routed_batch_items, count);
Ok(())
}
async fn get_routed(
&self,
route_key: &str,
key: &[u8],
) -> Result<Option<Vec<u8>>, DurabilityError> {
self.counters.add(&self.counters.get_routed_calls, 1);
let value = self.inner.get_routed(route_key, key).await?;
if value.is_some() {
self.counters.add(&self.counters.get_routed_rows, 1);
}
Ok(value)
}
async fn append(
&self,
stream_key: &str,
payload: Vec<u8>,
expected_seq: u64,
) -> Result<u64, DurabilityError> {
self.counters.add(&self.counters.append_calls, 1);
let bytes = widen(payload.len());
let assigned = self.inner.append(stream_key, payload, expected_seq).await?;
self.counters.add(&self.counters.append_rows, 1);
self.counters.add(&self.counters.append_bytes, bytes);
Ok(assigned)
}
async fn read_from(
&self,
stream_key: &str,
offset: u64,
limit: usize,
) -> Result<Vec<StoredEntry>, DurabilityError> {
self.counters.add(&self.counters.read_from_calls, 1);
let entries = self.inner.read_from(stream_key, offset, limit).await?;
self.counters
.add(&self.counters.read_from_rows, widen(entries.len()));
Ok(entries)
}
async fn append_all(
&self,
stream_key: &str,
payloads: Vec<Vec<u8>>,
expected_seq: u64,
) -> Result<u64, DurabilityError> {
self.counters.add(&self.counters.append_all_calls, 1);
let rows = widen(payloads.len());
let bytes = payloads
.iter()
.map(|payload| widen(payload.len()))
.fold(0_u64, u64::saturating_add);
let next = self
.inner
.append_all(stream_key, payloads, expected_seq)
.await?;
self.counters.add(&self.counters.append_all_rows, rows);
self.counters.add(&self.counters.append_all_bytes, bytes);
Ok(next)
}
async fn read_at(
&self,
stream_key: &str,
sequence: u64,
) -> Result<Option<StoredEntry>, DurabilityError> {
self.counters.add(&self.counters.read_at_calls, 1);
let entry = self.inner.read_at(stream_key, sequence).await?;
if entry.is_some() {
self.counters.add(&self.counters.read_at_rows, 1);
}
Ok(entry)
}
async fn cas(&self, key: &str, old_value: u64, new_value: u64) -> Result<(), DurabilityError> {
self.counters.add(&self.counters.cas_calls, 1);
self.inner.cas(key, old_value, new_value).await
}
async fn read_stream_next_seq(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
self.counters
.add(&self.counters.read_stream_next_seq_calls, 1);
self.inner.read_stream_next_seq(stream_key).await
}
async fn read_value(&self, key: &str) -> Result<Option<u64>, DurabilityError> {
self.counters.add(&self.counters.read_value_calls, 1);
self.inner.read_value(key).await
}
async fn scan_page(
&self,
prefix: &str,
cursor: Option<ScanCursor>,
limit: usize,
) -> Result<ScanPage, DurabilityError> {
self.counters.add(&self.counters.scan_page_calls, 1);
let page = self.inner.scan_page(prefix, cursor, limit).await?;
self.counters
.add(&self.counters.scan_page_rows, widen(page.entries.len()));
Ok(page)
}
async fn flush(&self) -> Result<(), DurabilityError> {
self.counters.add(&self.counters.flushes, 1);
self.counters.flushed();
self.inner.flush().await
}
async fn truncate_before(
&self,
stream_key: &str,
before_seq: u64,
) -> Result<u64, DurabilityError> {
add(&self.counters.truncate_calls, 1);
let deleted = self.inner.truncate_before(stream_key, before_seq).await?;
add(&self.counters.truncate_rows, deleted);
Ok(deleted)
}
async fn scan_routed(
&self,
route_key: &str,
from: &[u8],
to: &[u8],
) -> Result<Vec<(Vec<u8>, Vec<u8>)>, DurabilityError> {
self.inner.scan_routed(route_key, from, to).await
}
async fn remove_routed_range(
&self,
route_key: &str,
from: &[u8],
to: &[u8],
) -> Result<u64, DurabilityError> {
self.inner.remove_routed_range(route_key, from, to).await
}
fn store_counts(&self) -> Option<StoreCountsHandle> {
Some(self.handle())
}
}
#[cfg(test)]
mod tests;