Skip to main content

automation_structures/compositions/
rate_limit.rs

1// Executable RateLimit contract. The TLA+ carrier models a per-window grant
2// bound with count, window_start, and a runtime-provided clock. It has two
3// actions:
4//
5//   TryAcquire — a three-branch atomic step:
6//     (1) window expired (clock - window_start >= WindowDuration):
7//         window_start' = clock and count' = 1;
8//     (2) else, headroom (count < MaxPerWindow): count' = count + 1;
9//     (3) else: UNCHANGED (the acquire is rejected).
10//   Tick — guard clock < MaxClock; clock' = clock + 1.
11//
12// Maintained predicates:
13//
14//   TypeInvariant       == count ∈ 0..MaxPerWindow /\ window_start ∈ Nat
15//                          /\ clock ∈ Nat
16//   WindowCountBound     == count <= MaxPerWindow
17//   WindowStartNotFuture == window_start <= clock
18//
19// The natural-number state is represented by u64. MaxPerWindow >= 1 is a
20// constructor precondition because the rollover grants one request. TryAcquire
21// returns whether the request was granted; Tick's guard is a method precondition.
22//
23// Evidence boundary: this is a sequential witness. Each merged TLA+ action is
24// one &mut self method and Rust's exclusive mutable borrow is the atomic
25// boundary. A concurrent realization must supply equivalent exclusion.
26
27use vstd::prelude::*;
28
29#[allow(unused_imports)]
30use crate::connectives::{counter, cursor};
31
32verus! {
33
34/// A per-window operation bound: at most `max_per_window` acquires per
35/// `window_duration` clock units, the window re-anchored at `window_start`.
36pub struct RateLimit {
37    /// Budget component for operations admitted in the current window.
38    pub budget: crate::primitives::budget::Budget,
39    /// WindowDuration (constant): the window length in clock units.
40    pub window_duration: u64,
41    /// MaxClock (constant): the model's clock bound (Tick's guard).
42    pub max_clock: u64,
43    /// window_start ∈ Nat: when the current window was anchored.
44    pub window_start: u64,
45    /// clock ∈ Nat: the runtime-given clock.
46    pub clock: u64,
47}
48
49impl RateLimit {
50    // ── Specifications ──────────────────────────────────────────────────
51
52    /// TLA+ TypeInvariant's range clause (count ∈ 0..MaxPerWindow; the Nat
53    /// typings are carried by u64), plus the constants clause MaxPerWindow >= 1
54    /// that the rollover branch's count' = 1 needs (see the header note).
55    pub open spec fn type_invariant(&self) -> bool {
56        &&& self.budget.capacity >= 1
57        &&& self.budget.safety_invariant()
58        &&& self.budget.reserved == 0
59        &&& self.budget.pending_eviction == 0
60    }
61
62    /// TLA+ `WindowCountBound == count <= MaxPerWindow` (the same bound
63    /// TypeInvariant's range clause states, kept under its own name for
64    /// fidelity to the .cfg's invariant list).
65    pub open spec fn window_count_bound(&self) -> bool {
66        self.budget.allocated <= self.budget.capacity
67    }
68
69    /// TLA+ `WindowStartNotFuture == window_start <= clock`.
70    pub open spec fn window_start_not_future(&self) -> bool {
71        cursor::cursor_admitted(self.window_start as nat, self.clock as nat)
72    }
73
74    /// TryAcquire's branch-1 condition: `clock - window_start >= WindowDuration`
75    /// (stated over int so the subtraction is exact; WindowStartNotFuture keeps it
76    /// non-negative at every reachable state).
77    pub open spec fn window_expired(&self) -> bool {
78        self.clock as int - self.window_start as int >= self.window_duration as int
79    }
80
81    // ── Init (TLA+ Init) ────────────────────────────────────────────────
82
83    /// Construct the initial state: count = 0, window_start = 0, clock = 0.
84    /// Realises the TLA+ `Init` predicate and establishes all three invariants.
85    pub fn new(max_per_window: u64, window_duration: u64, max_clock: u64) -> (r: RateLimit)
86        requires
87            max_per_window >= 1,   // constants clause (header note)
88        ensures
89            r.budget.capacity == max_per_window,
90            r.window_duration == window_duration,
91            r.max_clock == max_clock,
92            r.budget.allocated == 0,
93            r.window_start == 0,
94            r.clock == 0,
95            r.type_invariant(),
96            r.window_count_bound(),
97            r.window_start_not_future(),
98    {
99        RateLimit {
100            budget: crate::primitives::budget::Budget::new(max_per_window),
101            window_duration,
102            max_clock,
103            window_start: 0,
104            clock: 0,
105        }
106    }
107
108    // ── TryAcquire (TLA+ TryAcquire) ────────────────────────────────────
109
110    /// Try to acquire one operation. The whole three-branch TLA+ IF is this
111    /// one method: on an expired window, re-anchor AND grant in the same step
112    /// (window_start' = clock, count' = 1); on headroom, grant (count' + 1);
113    /// otherwise reject (UNCHANGED). Returns whether the acquire was granted.
114    pub fn try_acquire(&mut self) -> (acquired: bool)
115        requires
116            old(self).type_invariant(),
117            old(self).window_start_not_future(),
118        ensures
119            final(self).budget.capacity == old(self).budget.capacity,
120            final(self).window_duration == old(self).window_duration,
121            final(self).max_clock == old(self).max_clock,
122            final(self).clock == old(self).clock,   // TryAcquire never moves the clock
123            counter::stutter(old(self).clock as int, final(self).clock as int),
124            cursor::cursor_admitted(
125                old(self).window_start as nat,
126                final(self).window_start as nat,
127            ),
128            acquired == (old(self).window_expired()
129                || old(self).budget.allocated < old(self).budget.capacity),
130            // Branch 1: rollover — re-anchor and grant, fused.
131            old(self).window_expired() ==> {
132                &&& final(self).window_start == old(self).clock
133                &&& final(self).budget.allocated == 1
134            },
135            // Branch 2: grant within the window.
136            (!old(self).window_expired()
137                && old(self).budget.allocated < old(self).budget.capacity) ==> {
138                &&& final(self).window_start == old(self).window_start
139                &&& final(self).budget.allocated == old(self).budget.allocated + 1
140            },
141            // Branch 3: reject — UNCHANGED vars.
142            (!old(self).window_expired()
143                && old(self).budget.allocated >= old(self).budget.capacity) ==> {
144                &&& final(self).window_start == old(self).window_start
145                &&& final(self).budget.allocated == old(self).budget.allocated
146            },
147            final(self).type_invariant(),
148            final(self).window_count_bound(),
149            final(self).window_start_not_future(),
150    {
151        // WindowStartNotFuture makes the subtraction safe.
152        let elapsed = self.clock - self.window_start;
153        if elapsed >= self.window_duration {
154            // Rollover: re-anchor the window at the current clock and grant
155            // the acquire, in one step. count' = 1 <= max_per_window by the
156            // constants clause; window_start' = clock keeps WindowStartNotFuture.
157            self.window_start = self.clock;
158            let allocated = self.budget.allocated;
159            self.budget.release(allocated);
160            let _accepted = self.budget.try_allocate(1);
161            assert(_accepted);
162            true
163        } else {
164            self.budget.try_allocate(1)
165        }
166    }
167
168    // ── Tick (TLA+ Tick) ────────────────────────────────────────────────
169
170    /// Advance the runtime-given clock by one. Realises the TLA+ `Tick`
171    /// action: its guard (clock < MaxClock) is a `requires`, so the action is
172    /// callable exactly when the TLA+ action is enabled.
173    pub fn tick(&mut self)
174        requires
175            old(self).type_invariant(),
176            old(self).window_start_not_future(),
177            old(self).clock < old(self).max_clock,   // Tick's guard
178        ensures
179            final(self).budget.capacity == old(self).budget.capacity,
180            final(self).window_duration == old(self).window_duration,
181            final(self).max_clock == old(self).max_clock,
182            final(self).budget.allocated == old(self).budget.allocated,
183            final(self).window_start == old(self).window_start,
184            final(self).clock == old(self).clock + 1,
185            counter::increment(old(self).clock as int, final(self).clock as int),
186            final(self).type_invariant(),
187            final(self).window_count_bound(),
188            final(self).window_start_not_future(),
189    {
190        self.clock = self.clock + 1;
191    }
192}
193
194}