rushdb_storage/
usage.rs

1//! Per-tenant logical byte accounting.
2//!
3//! Every backend write already reads the value it replaces (for MVCC history),
4//! so the change in a tenant's stored `key + value` bytes is known at apply
5//! time at no extra read. A backend folds each op of a batch into
6//! [`UsageBytes`] and writes one [`usage_bytes_key`] row per touched tenant in
7//! the same backend transaction, so the figure is exact at every commit index
8//! and replay (WAL, replicas) reproduces it.
9//!
10//! The figure is *logical*: it sums the row bytes the engine stores, not the
11//! file pages that hold them (redb pages run partly full, so the file is
12//! larger). The usage row itself is excluded, which keeps it a fixed point.
13
14use std::collections::BTreeMap;
15
16use rushdb_types::ids::TenantId;
17use rushdb_types::keys::{parse_tenant, usage_bytes_key, F_USAGE_BYTES, TENANT_PREFIX_LEN};
18
19/// Whether `key` is a tenant's usage-bytes row.
20#[must_use]
21pub fn is_usage_bytes_key(key: &[u8]) -> bool {
22    key.len() == TENANT_PREFIX_LEN + 1
23        && parse_tenant(key).is_some()
24        && key[TENANT_PREFIX_LEN] == F_USAGE_BYTES
25}
26
27/// Byte deltas of one batch, per tenant.
28#[derive(Default)]
29pub(crate) struct UsageBytes {
30    deltas: BTreeMap<TenantId, i64>,
31}
32
33impl UsageBytes {
34    /// Record that `key` went from a value of `old` bytes to one of `new`
35    /// bytes (`None` = absent).
36    pub(crate) fn record(&mut self, key: &[u8], old: Option<usize>, new: Option<usize>) {
37        let Some(tenant) = parse_tenant(key) else {
38            return;
39        };
40        if is_usage_bytes_key(key) {
41            return;
42        }
43        let size = |value: Option<usize>| {
44            value.map_or(0, |len| i64::try_from(key.len() + len).unwrap_or(i64::MAX))
45        };
46        let delta = size(new) - size(old);
47        if delta != 0 {
48            *self.deltas.entry(tenant).or_default() += delta;
49        }
50    }
51
52    /// The usage rows to write: for each tenant, its new total from `current`
53    /// (the row's value before the batch). A total of zero deletes the row.
54    ///
55    /// # Errors
56    /// Returns a message when a stored total is malformed or the delta would
57    /// take it below zero, which means the accounting has diverged.
58    pub(crate) fn rows(
59        self,
60        mut current: impl FnMut(&[u8]) -> Result<Option<Vec<u8>>, String>,
61    ) -> Result<Vec<(Vec<u8>, Option<Vec<u8>>)>, String> {
62        let mut out = Vec::with_capacity(self.deltas.len());
63        for (tenant, delta) in self.deltas {
64            if delta == 0 {
65                continue;
66            }
67            let key = usage_bytes_key(tenant);
68            let total = rushdb_types::keys::decode_counter(current(&key)?.as_deref())
69                .ok_or_else(|| format!("tenant {tenant}: malformed usage-bytes row"))?;
70            let next = total
71                .checked_add_signed(delta)
72                .ok_or_else(|| format!("tenant {tenant}: usage bytes would underflow"))?;
73            out.push((key, (next > 0).then(|| next.to_be_bytes().to_vec())));
74        }
75        Ok(out)
76    }
77}
rushdb_storage/usage.rs

rushdb_storage/
usage.rs
