futu_auth/limits/runtime/

batch.rs

//! Split from runtime.rs: batch (atomic multi-leg daily receipt) admission.
//!
//! `BatchLimitGuard` commit / receipt release and `RuntimeCounters::check_batch_limits`.

use super::*;
use chrono::{DateTime, Utc};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;

impl BatchLimitGuard<'_> {
    /// Commit all currency buckets exactly once for this caller+intent.
    pub fn commit_daily_once(self) -> Result<(), LimitOutcome> {
        if self.pending_daily.is_empty() || self.already_committed {
            return Ok(());
        }
        let mut receipts = self.counters.daily_receipts.lock();
        // Sample the date only after receipt serialization. This gives
        // midnight commits a total order: an old preflight can never write the
        // prior date after a newer-date commit, even if no D+1 request raced it.
        let commit_day = Utc::now().date_naive();
        self.commit_daily_once_locked(&mut receipts, commit_day)
    }

    fn commit_daily_once_locked(
        &self,
        receipts: &mut HashMap<String, DailyReceipt>,
        commit_day: NaiveDate,
    ) -> Result<(), LimitOutcome> {
        if commit_day < self.today {
            return Err(stale_day_outcome(commit_day, self.today));
        }
        receipts.retain(|_, receipt| receipt.day >= commit_day);
        if let Some(existing) = receipts.get(&self.receipt_key) {
            return if existing.fingerprint == self.fingerprint {
                Ok(())
            } else {
                Err(LimitOutcome::ValueReject(
                    "daily quota receipt was reused with different order facts".into(),
                ))
            };
        }
        let counter = self
            .counters
            .counters
            .entry(self.key_id.clone())
            .or_insert_with(|| DailyCounter::new(commit_day));
        counter
            .try_add_batch(&self.pending_daily, self.daily_cap, commit_day)
            .map_err(batch_daily_error)?;
        receipts.insert(
            self.receipt_key.clone(),
            DailyReceipt {
                fingerprint: self.fingerprint,
                day: commit_day,
                deltas: self.pending_daily.clone(),
            },
        );
        Ok(())
    }

    #[must_use]
    pub fn is_already_committed(&self) -> bool {
        self.already_committed
    }

    #[cfg(test)]
    pub(in crate::limits) fn commit_daily_once_on(
        self,
        commit_day: NaiveDate,
    ) -> Result<(), LimitOutcome> {
        if self.pending_daily.is_empty() || self.already_committed {
            return Ok(());
        }
        let mut receipts = self.counters.daily_receipts.lock();
        self.commit_daily_once_locked(&mut receipts, commit_day)
    }
}

fn batch_fingerprint(contexts: &[CheckCtx]) -> [u8; 32] {
    fn text(digest: &mut Sha256, value: Option<&str>) {
        match value {
            Some(value) => {
                digest.update([1]);
                digest.update((value.len() as u64).to_be_bytes());
                digest.update(value.as_bytes());
            }
            None => digest.update([0]),
        }
    }
    let mut digest = Sha256::new();
    digest.update((contexts.len() as u64).to_be_bytes());
    for context in contexts {
        text(&mut digest, Some(&context.market));
        text(&mut digest, Some(&context.symbol));
        match context.order_value {
            Some(value) => {
                digest.update([1]);
                digest.update(value.to_bits().to_be_bytes());
            }
            None => digest.update([0]),
        }
        text(&mut digest, context.trd_side.as_deref());
        match context.acc_id {
            Some(value) => {
                digest.update([1]);
                digest.update(value.to_be_bytes());
            }
            None => digest.update([0]),
        }
        digest.update([u8::from(context.mutation_no_exposure)]);
        text(&mut digest, context.currency.as_deref());
    }
    digest.finalize().into()
}

impl RuntimeCounters {
    /// Validate a complete group write and return one atomic daily receipt.
    /// `receipt_id` must be the durable caller+intent identity shared by every
    /// transport surface and the gateway writer.
    pub fn check_batch_limits<'a>(
        &'a self,
        key_id: &str,
        limits: &(impl LimitPolicy + ?Sized),
        contexts: &[CheckCtx],
        now: DateTime<Utc>,
        commit_rate: bool,
        receipt_id: &str,
    ) -> Result<BatchLimitGuard<'a>, LimitOutcome> {
        if receipt_id.is_empty() {
            return Err(LimitOutcome::ValueReject(
                "batch daily admission requires a durable receipt id".into(),
            ));
        }
        for (index, context) in contexts.iter().enumerate() {
            let mut non_committing = context.clone();
            non_committing.mutation_no_exposure = true;
            let _ = self.check_limits(
                key_id,
                limits,
                &non_committing,
                now,
                commit_rate && index == 0,
            )?;
        }

        let mut aggregate = BTreeMap::<Option<String>, f64>::new();
        if limits.max_daily_value().is_some() {
            for context in contexts {
                if context.mutation_no_exposure {
                    continue;
                }
                if let Some(value) = context.order_value {
                    let currency = context.daily_currency();
                    let next = aggregate.get(&currency).copied().unwrap_or(0.0) + value;
                    validate_order_value(next).map_err(|reason| {
                        LimitOutcome::ValueReject(format!(
                            "order value invalid ({reason:?}) in batch aggregate"
                        ))
                    })?;
                    aggregate.insert(currency, next);
                }
            }
        }
        let pending_daily = aggregate.into_iter().collect::<Vec<_>>();
        let fingerprint = batch_fingerprint(contexts);
        let receipt_key = format!("{}\0{}", key_id, receipt_id);
        let already_committed = {
            let mut receipts = self.daily_receipts.lock();
            receipts.retain(|_, receipt| receipt.day >= now.date_naive());
            match receipts.get(&receipt_key) {
                Some(existing) if existing.fingerprint == fingerprint => true,
                Some(_) => {
                    return Err(LimitOutcome::ValueReject(
                        "daily quota receipt was reused with different order facts".into(),
                    ));
                }
                None => false,
            }
        };
        if !already_committed && !pending_daily.is_empty() {
            let counter = self
                .counters
                .entry(key_id.to_string())
                .or_insert_with(|| DailyCounter::new(now.date_naive()));
            counter
                .peek_add_batch(&pending_daily, limits.max_daily_value(), now.date_naive())
                .map_err(batch_daily_error)?;
        }
        Ok(BatchLimitGuard {
            counters: self,
            key_id: key_id.to_string(),
            receipt_key,
            fingerprint,
            pending_daily,
            daily_cap: limits.max_daily_value(),
            today: now.date_naive(),
            already_committed,
        })
    }

    /// Release an authoritative explicit-reject reservation exactly once.
    /// Unknown/timeout/malformed outcomes must not call this method.
    pub fn release_batch_receipt(
        &self,
        key_id: &str,
        receipt_id: &str,
        contexts: &[CheckCtx],
    ) -> Result<bool, LimitOutcome> {
        let fingerprint = batch_fingerprint(contexts);
        self.release_batch_receipt_internal(key_id, receipt_id, Some(&fingerprint))
    }

    /// Release the exact durable caller+intent receipt after an authoritative
    /// asynchronous failure. The original request contexts are already bound
    /// into the immutable receipt fingerprint and deltas at commit time.
    pub fn release_batch_receipt_by_id(
        &self,
        key_id: &str,
        receipt_id: &str,
    ) -> Result<bool, LimitOutcome> {
        self.release_batch_receipt_internal(key_id, receipt_id, None)
    }

    fn release_batch_receipt_internal(
        &self,
        key_id: &str,
        receipt_id: &str,
        expected_fingerprint: Option<&[u8; 32]>,
    ) -> Result<bool, LimitOutcome> {
        let receipt_key = format!("{}\0{}", key_id, receipt_id);
        let mut receipts = self.daily_receipts.lock();
        let Some(receipt) = receipts.get(&receipt_key).cloned() else {
            return Ok(false);
        };
        if expected_fingerprint.is_some_and(|fingerprint| receipt.fingerprint != *fingerprint) {
            return Err(LimitOutcome::ValueReject(
                "daily quota receipt release used different order facts".into(),
            ));
        }
        let counter = self.counters.get(key_id).ok_or_else(|| {
            LimitOutcome::ValueReject("daily quota receipt has no counter state".into())
        })?;
        counter
            .try_subtract_batch(&receipt.deltas, receipt.day)
            .map_err(batch_daily_error)?;
        receipts.remove(&receipt_key);
        Ok(true)
    }
}