krabka_verified/
quota.rs

1//! Kafka quota-entity precedence and quota bucket debt accounting.
2
3use {
4    crate::throttle::{
5        AvailableTokens, BurstCapacity, RefillTokens, RequestedTokens, plan_consume,
6    },
7    creusot_std::prelude::*,
8};
9
10mod refill;
11pub use refill::quota_refill;
12
13mod debt_cap;
14pub use debt_cap::quota_debt_cap;
15
16mod request;
17pub use request::quota_whole_request;
18
19/// Credit micro-tokens to a quota bucket, repaying debt before filling its
20/// balance.
21///
22/// Both refill and refund use this transition; rate and clock conversion and
23/// atomic state publication remain with the caller.
24#[ensures(result.0@ <= burst@)]
25#[ensures(result.1@ == debt@ - if debt@ <= credit@ { debt@ } else { credit@ })]
26#[ensures(result.0@ == crate::throttle::capped(available@,
27    credit@ - if debt@ <= credit@ { debt@ } else { credit@ }, burst@))]
28#[ensures((available@ == 0 || debt@ == 0) ==> (result.0@ == 0 || result.1@ == 0))]
29#[must_use]
30pub fn quota_credit(available: u64, debt: u64, credit: u64, burst: u64) -> (u64, u64) {
31    let repaid = credit.min(debt);
32    let (_, new_available) = plan_consume(
33        AvailableTokens(available),
34        RefillTokens(credit - repaid),
35        BurstCapacity(burst),
36        RequestedTokens(0),
37    );
38    (new_available.0, debt - repaid)
39}
40
41/// Charge the effective micro-token request in full, retaining its unpaid
42/// portion as debt up to `debt_cap`.
43///
44/// A smaller cap intentionally forgets debt; `u64::MAX` is the unbounded mode's
45/// representability limit.
46#[ensures(result.0@ == if requested@ <= crate::throttle::capped(available@, 0, burst@) {
47    crate::throttle::capped(available@, 0, burst@) - requested@
48} else { 0 })]
49#[ensures(result.1@ == crate::throttle::capped(debt@,
50    requested@ - (crate::throttle::capped(available@, 0, burst@) - result.0@), debt_cap@))]
51#[ensures(result.0@ <= burst@ && result.1@ <= debt_cap@)]
52#[ensures((available@ == 0 || debt@ == 0) ==> (result.0@ == 0 || result.1@ == 0))]
53#[must_use]
54pub fn quota_charge(
55    available: u64,
56    debt: u64,
57    requested: u64,
58    burst: u64,
59    debt_cap: u64,
60) -> (u64, u64) {
61    let (grant, new_available) = plan_consume(
62        AvailableTokens(available),
63        RefillTokens(0),
64        BurstCapacity(burst),
65        RequestedTokens(requested),
66    );
67    let new_debt = debt.saturating_add(requested - grant.0).min(debt_cap);
68    (new_available.0, new_debt)
69}
70
71/// Selected user/client quota candidate, in Kafka's precedence order.
72///
73/// The order is `DefaultQuotaCallback`'s in Kafka's `ClientQuotaManager`:
74/// every `user=U` level ranks above every `user=<default>` level, and both
75/// rank above the `client-id`-only levels.
76#[krabka_macros::model_type]
77pub enum UserClientQuotaPrecedence {
78    /// 1. `user=U, client-id=C`
79    ExactPair,
80    /// 2. `user=U, client-id=<default>`
81    ExactUserDefaultClient,
82    /// 3. `user=U`
83    ExactUser,
84    /// 4. `user=<default>, client-id=C`
85    DefaultUserExactClient,
86    /// 5. `user=<default>, client-id=<default>`
87    DefaultPair,
88    /// 6. `user=<default>`
89    DefaultUser,
90    /// 7. `client-id=C`
91    ExactClient,
92    /// 8. `client-id=<default>`
93    DefaultClient,
94    None,
95}
96
97/// Whether one canonical quota candidate exists in the metadata image.
98#[krabka_macros::model_type]
99pub enum QuotaCandidatePresence {
100    Absent,
101    Present,
102}
103
104/// Presence of each canonical user/client quota candidate, in Kafka's
105/// precedence order.
106#[krabka_macros::model_type]
107pub struct UserClientQuotaFacts {
108    pub exact_pair: QuotaCandidatePresence,
109    pub exact_user_default_client: QuotaCandidatePresence,
110    pub exact_user: QuotaCandidatePresence,
111    pub default_user_exact_client: QuotaCandidatePresence,
112    pub default_pair: QuotaCandidatePresence,
113    pub default_user: QuotaCandidatePresence,
114    pub exact_client: QuotaCandidatePresence,
115    pub default_client: QuotaCandidatePresence,
116}
117
118// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
119#[cfg(creusot)]
120#[cfg_attr(test, mutants::skip)]
121#[logic(open)]
122pub fn exact_user_quota_absent(facts: UserClientQuotaFacts) -> bool {
123    pearlite! { facts.exact_pair == QuotaCandidatePresence::Absent
124    && facts.exact_user_default_client == QuotaCandidatePresence::Absent
125    && facts.exact_user == QuotaCandidatePresence::Absent }
126}
127
128// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
129#[cfg(creusot)]
130#[cfg_attr(test, mutants::skip)]
131#[logic(open)]
132fn all_user_quota_absent(facts: UserClientQuotaFacts) -> bool {
133    pearlite! { exact_user_quota_absent(facts)
134    && facts.default_user_exact_client == QuotaCandidatePresence::Absent
135    && facts.default_pair == QuotaCandidatePresence::Absent
136    && facts.default_user == QuotaCandidatePresence::Absent }
137}
138
139/// Select Kafka's first present user/client quota candidate.
140#[ensures((result == UserClientQuotaPrecedence::ExactPair)
141    == (facts.exact_pair == QuotaCandidatePresence::Present))]
142#[ensures((result == UserClientQuotaPrecedence::ExactUserDefaultClient)
143    == (facts.exact_pair == QuotaCandidatePresence::Absent
144        && facts.exact_user_default_client == QuotaCandidatePresence::Present))]
145#[ensures((result == UserClientQuotaPrecedence::ExactUser)
146    == (facts.exact_pair == QuotaCandidatePresence::Absent
147        && facts.exact_user_default_client == QuotaCandidatePresence::Absent
148        && facts.exact_user == QuotaCandidatePresence::Present))]
149#[ensures((result == UserClientQuotaPrecedence::DefaultUserExactClient)
150    == (exact_user_quota_absent(facts)
151        && facts.default_user_exact_client == QuotaCandidatePresence::Present))]
152#[ensures((result == UserClientQuotaPrecedence::DefaultPair)
153    == (exact_user_quota_absent(facts)
154        && facts.default_user_exact_client == QuotaCandidatePresence::Absent
155        && facts.default_pair == QuotaCandidatePresence::Present))]
156#[ensures((result == UserClientQuotaPrecedence::DefaultUser)
157    == (exact_user_quota_absent(facts)
158        && facts.default_user_exact_client == QuotaCandidatePresence::Absent
159        && facts.default_pair == QuotaCandidatePresence::Absent
160        && facts.default_user == QuotaCandidatePresence::Present))]
161#[ensures((result == UserClientQuotaPrecedence::ExactClient)
162    == (all_user_quota_absent(facts)
163        && facts.exact_client == QuotaCandidatePresence::Present))]
164#[ensures((result == UserClientQuotaPrecedence::DefaultClient)
165    == (all_user_quota_absent(facts)
166        && facts.exact_client == QuotaCandidatePresence::Absent
167        && facts.default_client == QuotaCandidatePresence::Present))]
168#[ensures((result == UserClientQuotaPrecedence::None) == !(
169    facts.exact_pair == QuotaCandidatePresence::Present
170        || facts.exact_user_default_client == QuotaCandidatePresence::Present
171        || facts.exact_user == QuotaCandidatePresence::Present
172        || facts.default_user_exact_client == QuotaCandidatePresence::Present
173        || facts.default_pair == QuotaCandidatePresence::Present
174        || facts.default_user == QuotaCandidatePresence::Present
175        || facts.exact_client == QuotaCandidatePresence::Present
176        || facts.default_client == QuotaCandidatePresence::Present))]
177#[must_use]
178pub fn user_client_quota_precedence(facts: UserClientQuotaFacts) -> UserClientQuotaPrecedence {
179    if matches!(facts.exact_pair, QuotaCandidatePresence::Present) {
180        UserClientQuotaPrecedence::ExactPair
181    } else if matches!(
182        facts.exact_user_default_client,
183        QuotaCandidatePresence::Present
184    ) {
185        UserClientQuotaPrecedence::ExactUserDefaultClient
186    } else if matches!(facts.exact_user, QuotaCandidatePresence::Present) {
187        UserClientQuotaPrecedence::ExactUser
188    } else if matches!(
189        facts.default_user_exact_client,
190        QuotaCandidatePresence::Present
191    ) {
192        UserClientQuotaPrecedence::DefaultUserExactClient
193    } else if matches!(facts.default_pair, QuotaCandidatePresence::Present) {
194        UserClientQuotaPrecedence::DefaultPair
195    } else if matches!(facts.default_user, QuotaCandidatePresence::Present) {
196        UserClientQuotaPrecedence::DefaultUser
197    } else if matches!(facts.exact_client, QuotaCandidatePresence::Present) {
198        UserClientQuotaPrecedence::ExactClient
199    } else if matches!(facts.default_client, QuotaCandidatePresence::Present) {
200        UserClientQuotaPrecedence::DefaultClient
201    } else {
202        UserClientQuotaPrecedence::None
203    }
204}
205
206/// Selected IP quota candidate.
207#[krabka_macros::model_type]
208pub enum IpQuotaPrecedence {
209    Exact,
210    Default,
211    None,
212}
213
214/// Select an exact IP quota before the IP default.
215#[ensures((result == IpQuotaPrecedence::Exact) == exact)]
216#[ensures((result == IpQuotaPrecedence::Default) == (!exact && default))]
217#[ensures((result == IpQuotaPrecedence::None) == (!exact && !default))]
218#[must_use]
219pub fn ip_quota_precedence(exact: bool, default: bool) -> IpQuotaPrecedence {
220    if exact {
221        IpQuotaPrecedence::Exact
222    } else if default {
223        IpQuotaPrecedence::Default
224    } else {
225        IpQuotaPrecedence::None
226    }
227}
228
229// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
230#[cfg(creusot)]
231#[cfg_attr(test, mutants::skip)]
232#[logic(open)]
233/// Spendable credit is bounded by the burst and is absent while debt remains.
234pub(crate) fn balance_within_burst(available: Int, debt: Int, burst: Int) -> bool {
235    pearlite! { available <= burst && (available == 0 || debt == 0) }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::{
241        IpQuotaPrecedence, QuotaCandidatePresence, UserClientQuotaFacts, UserClientQuotaPrecedence,
242        ip_quota_precedence, user_client_quota_precedence,
243    };
244
245    #[test]
246    fn selectors_choose_the_first_present_candidate() {
247        for mask in 0_u16..256 {
248            let present = |index| mask & (1_u16 << index) != 0_u16;
249            let candidate = |index| {
250                if present(index) {
251                    QuotaCandidatePresence::Present
252                } else {
253                    QuotaCandidatePresence::Absent
254                }
255            };
256            let got = user_client_quota_precedence(UserClientQuotaFacts {
257                exact_pair: candidate(0),
258                exact_user_default_client: candidate(1),
259                exact_user: candidate(2),
260                default_user_exact_client: candidate(3),
261                default_pair: candidate(4),
262                default_user: candidate(5),
263                exact_client: candidate(6),
264                default_client: candidate(7),
265            });
266            let expected = match (0..8).find(|index| present(*index)) {
267                Some(0) => UserClientQuotaPrecedence::ExactPair,
268                Some(1) => UserClientQuotaPrecedence::ExactUserDefaultClient,
269                Some(2) => UserClientQuotaPrecedence::ExactUser,
270                Some(3) => UserClientQuotaPrecedence::DefaultUserExactClient,
271                Some(4) => UserClientQuotaPrecedence::DefaultPair,
272                Some(5) => UserClientQuotaPrecedence::DefaultUser,
273                Some(6) => UserClientQuotaPrecedence::ExactClient,
274                Some(7) => UserClientQuotaPrecedence::DefaultClient,
275                Some(_) | None => UserClientQuotaPrecedence::None,
276            };
277            assert2::check!(got == expected, "mask {mask:#010b}");
278        }
279        assert2::check!(ip_quota_precedence(true, true) == IpQuotaPrecedence::Exact);
280        assert2::check!(ip_quota_precedence(false, true) == IpQuotaPrecedence::Default);
281        assert2::check!(ip_quota_precedence(false, false) == IpQuotaPrecedence::None);
282    }
283}
krabka_verified/quota.rs

krabka_verified/
quota.rs
