krabka_verified/uniform_assignor/
select_least_loaded.rs

1use {
2    super::{HomogeneousMemberQuota, SubscriberLoad, UniformQuotaSplit},
3    creusot_std::prelude::*,
4};
5
6/// Split `total_partitions` over `member_count` members, as
7/// `UniformHomogeneousAssignmentBuilder.build` does.
8///
9/// The contract is the Euclidean division property, which names the floor
10/// and the remainder uniquely: the quotas add up to the total, and fewer
11/// members than there are members get the extra partition.
12#[requires(member_count@ > 0)]
13#[ensures(result.minimum_quota@ * member_count@ + result.extra_quotas@ == total_partitions@)]
14#[ensures(result.extra_quotas@ < member_count@)]
15#[must_use]
16pub fn uniform_quota_split(total_partitions: usize, member_count: usize) -> UniformQuotaSplit {
17    UniformQuotaSplit {
18        minimum_quota: total_partitions / member_count,
19        extra_quotas: total_partitions % member_count,
20    }
21}
22
23/// Kafka's rule for the tentative extra partition in
24/// `UniformHomogeneousAssignmentBuilder.maybeRevokePartitions`.
25///
26/// While extra partitions are left to hand out, a member keeps the extra slot
27/// when it already owns more than the minimum. Otherwise it gives the slot up
28/// for a later member, unless the members that are left, this one included,
29/// are no more than the extra partitions that are left. Then it must take it.
30// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
31#[cfg(creusot)]
32#[cfg_attr(test, mutants::skip)]
33#[logic]
34pub fn keeps_extra_quota(
35    minimum_quota: Int,
36    extras_left: Int,
37    owned: Int,
38    members_left: Int,
39) -> bool {
40    pearlite! { extras_left > 0 && (owned > minimum_quota || members_left <= extras_left) }
41}
42
43/// The extra partitions that are still unclaimed when the member at `index`
44/// is visited. Members are visited in order, and each member that keeps the
45/// extra slot claims one.
46// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
47#[cfg(creusot)]
48#[cfg_attr(test, mutants::skip)]
49#[logic]
50#[variant(index)]
51pub fn extras_left_before(
52    minimum_quota: Int,
53    extra_quotas: Int,
54    owned: Seq<usize>,
55    index: Int,
56) -> Int {
57    pearlite! {
58        if index <= 0 {
59            extra_quotas
60        } else {
61            extras_left_before(minimum_quota, extra_quotas, owned, index - 1)
62                - (if keeps_extra_quota(
63                    minimum_quota,
64                    extras_left_before(minimum_quota, extra_quotas, owned, index - 1),
65                    owned[index - 1]@,
66                    owned.len() - (index - 1),
67                ) { 1 } else { 0 })
68        }
69    }
70}
71
72/// The partition target for a member: the minimum, plus one when it keeps
73/// the extra slot.
74// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
75#[cfg(creusot)]
76#[cfg_attr(test, mutants::skip)]
77#[logic]
78pub fn member_target(minimum_quota: Int, quota: HomogeneousMemberQuota) -> Int {
79    pearlite! { minimum_quota + (if quota.takes_extra { 1 } else { 0 }) }
80}
81
82/// Every member's quota in the homogeneous builder, in member order.
83///
84/// `owned[i]` is how many partitions of the subscribed topics member `i`
85/// holds in its current target assignment. The contract states Kafka's
86/// balance-then-stickiness rule:
87///
88/// - Balance: each member's target is `minimum_quota` or `minimum_quota + 1`,
89///   and exactly `extra_quotas` members get the larger target, because no
90///   extra partition is left unclaimed after the last member.
91/// - Stickiness: each member keeps as much of its current ownership as its
92///   target allows, and receives only the difference.
93/// - The extra slots go by `keeps_extra_quota`, in member order.
94#[requires(extra_quotas@ <= [email protected]())]
95#[requires(minimum_quota@ < usize::MAX@)]
96#[ensures([email protected]() == [email protected]())]
97#[ensures(forall<i: Int> 0 <= i && i < [email protected]() ==> result@[i].takes_extra == keeps_extra_quota(
98    minimum_quota@,
99    extras_left_before(minimum_quota@, extra_quotas@, owned@, i),
100    owned@[i]@,
101    [email protected]() - i,
102))]
103#[ensures(forall<i: Int> 0 <= i && i < [email protected]()
104    ==> result@[i].retain@ + result@[i].fill@ == member_target(minimum_quota@, result@[i]))]
105#[ensures(forall<i: Int> 0 <= i && i < [email protected]()
106    ==> result@[i].retain@ == (if owned@[i]@ < member_target(minimum_quota@, result@[i]) {
107        owned@[i]@
108    } else {
109        member_target(minimum_quota@, result@[i])
110    }))]
111#[ensures(extras_left_before(minimum_quota@, extra_quotas@, owned@, [email protected]()) == 0)]
112#[must_use]
113pub fn homogeneous_member_quotas(
114    minimum_quota: usize,
115    extra_quotas: usize,
116    owned: &[usize],
117) -> Vec<HomogeneousMemberQuota> {
118    let mut quotas: Vec<HomogeneousMemberQuota> = Vec::new();
119    let mut extras_left = extra_quotas;
120    let mut i = 0usize;
121    #[invariant(i@ <= [email protected]())]
122    #[invariant([email protected]() == i@)]
123    #[invariant(extras_left@ == extras_left_before(minimum_quota@, extra_quotas@, owned@, i@))]
124    #[invariant(extras_left@ <= [email protected]() - i@)]
125    #[invariant(forall<k: Int> 0 <= k && k < i@ ==> quotas@[k].takes_extra == keeps_extra_quota(
126        minimum_quota@,
127        extras_left_before(minimum_quota@, extra_quotas@, owned@, k),
128        owned@[k]@,
129        [email protected]() - k,
130    ))]
131    #[invariant(forall<k: Int> 0 <= k && k < i@
132        ==> quotas@[k].retain@ + quotas@[k].fill@ == member_target(minimum_quota@, quotas@[k]))]
133    #[invariant(forall<k: Int> 0 <= k && k < i@
134        ==> quotas@[k].retain@ == (if owned@[k]@ < member_target(minimum_quota@, quotas@[k]) {
135            owned@[k]@
136        } else {
137            member_target(minimum_quota@, quotas@[k])
138        }))]
139    #[variant([email protected]() - i@)]
140    while i < owned.len() {
141        let Some(&owned_now) = owned.get(i) else {
142            break;
143        };
144        let members_left = owned.len() - i;
145        let takes_extra =
146            extras_left > 0 && (owned_now > minimum_quota || members_left <= extras_left);
147        let target = if takes_extra {
148            minimum_quota + 1
149        } else {
150            minimum_quota
151        };
152        let retain = if owned_now < target {
153            owned_now
154        } else {
155            target
156        };
157        quotas.push(HomogeneousMemberQuota {
158            takes_extra,
159            retain,
160            fill: target - retain,
161        });
162        if takes_extra {
163            extras_left -= 1;
164        }
165        i += 1;
166    }
167    quotas
168}
169
170/// `a` goes before `b` in the heterogeneous builder's least-loaded order: the
171/// smaller current load first, then the smaller load at the start of the
172/// topic. The candidate position breaks the remaining ties.
173///
174/// `MemberAssignmentBalancer.nextLeastLoadedMember` sorts the subscribers once
175/// per topic by `(load, member index)` and then fills them level by level in
176/// that order. Among the members at the lowest current level, the next one it
177/// picks is therefore the first by starting load, then by member index.
178// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
179#[cfg(creusot)]
180#[cfg_attr(test, mutants::skip)]
181#[logic]
182pub fn load_precedes(a: SubscriberLoad, b: SubscriberLoad) -> bool {
183    pearlite! {
184        a.assigned@ < b.assigned@
185            || (a.assigned@ == b.assigned@ && a.assigned_at_topic_start@ < b.assigned_at_topic_start@)
186    }
187}
188
189/// Choose the subscriber that receives the next unassigned partition of a
190/// topic in `UniformHeterogeneousAssignmentBuilder.assignRemainingPartitions`.
191///
192/// The host passes the topic's subscribers in ascending member order. The
193/// result is `None` exactly for an empty slice. Otherwise no candidate
194/// precedes the chosen one, and the chosen one strictly precedes every
195/// earlier candidate, so the first of the least-loaded candidates wins.
196#[ensures((result == None) == ([email protected]() == 0))]
197#[ensures(forall<chosen: usize> result == Some(chosen) ==> chosen@ < [email protected]()
198    && (forall<j: Int> 0 <= j && j < [email protected]()
199        ==> !load_precedes(candidates@[j], candidates@[chosen@]))
200    && (forall<j: Int> 0 <= j && j < chosen@
201        ==> load_precedes(candidates@[chosen@], candidates@[j])))]
202#[must_use]
203pub fn select_least_loaded(candidates: &[SubscriberLoad]) -> Option<usize> {
204    let count = candidates.len();
205    if count == 0 {
206        return None;
207    }
208    let mut best = 0usize;
209    let mut i = 1usize;
210    #[invariant(1 <= i@ && i@ <= [email protected]())]
211    #[invariant(best@ < i@)]
212    #[invariant(forall<j: Int> 0 <= j && j < i@
213        ==> !load_precedes(candidates@[j], candidates@[best@]))]
214    #[invariant(forall<j: Int> 0 <= j && j < best@
215        ==> load_precedes(candidates@[best@], candidates@[j]))]
216    #[variant([email protected]() - i@)]
217    while i < count {
218        if let (Some(candidate), Some(current)) = (candidates.get(i), candidates.get(best))
219            && (candidate.assigned < current.assigned
220                || (candidate.assigned == current.assigned
221                    && candidate.assigned_at_topic_start < current.assigned_at_topic_start))
222        {
223            best = i;
224        }
225        i += 1;
226    }
227    Some(best)
228}
krabka_verified/uniform_assignor/select_least_loaded.rs

krabka_verified/uniform_assignor/
select_least_loaded.rs
