Skip to main content

krabka_verified/
reassignment.rs

1//! Partition-reassignment transition decisions.
2
3use creusot_std::prelude::*;
4
5/// The next safe phase of one in-flight partition reassignment.
6#[krabka_macros::model_type]
7pub enum ReassignmentAction {
8    WaitForReplication,
9    WaitForLeader,
10    Handoff(usize),
11    Complete,
12}
13
14/// Pointwise membership in a reassignment's union, additions, and removals.
15#[krabka_macros::model_type]
16pub struct ReassignmentSetMembership {
17    pub in_union: bool,
18    pub adding: bool,
19    pub removing: bool,
20}
21
22/// Classify one unique replica against the current and requested assignments.
23#[ensures(result.in_union == (in_current || in_target))]
24#[ensures(result.adding == (in_target && !in_current))]
25#[ensures(result.removing == (in_current && !in_target))]
26#[ensures(!(result.adding && result.removing))]
27#[must_use]
28pub fn reassignment_set_membership(in_current: bool, in_target: bool) -> ReassignmentSetMembership {
29    ReassignmentSetMembership {
30        in_union: in_current || in_target,
31        adding: in_target && !in_current,
32        removing: in_current && !in_target,
33    }
34}
35
36/// Final mutation admission for either a reassignment start or cancel.
37#[ensures(result == (
38    readiness.0
39        && readiness.1
40        && if mode.0 {
41            mode.1 && mode.2
42        } else {
43            start.0 && start.1
44        }
45))]
46#[must_use]
47pub fn reassignment_plan_admission(
48    mode: (bool, bool, bool),
49    start: (bool, bool),
50    readiness: (bool, bool),
51) -> bool {
52    let (is_cancel, cancel_approved, cancel_in_progress) = mode;
53    let (target_nonempty_unique_registered, rf_policy_satisfied) = start;
54    let (leader_eligible, epoch_available) = readiness;
55    leader_eligible
56        && epoch_available
57        && if is_cancel {
58            cancel_approved && cancel_in_progress
59        } else {
60            target_nonempty_unique_registered && rf_policy_satisfied
61        }
62}
63
64/// Wait until the reassignment may complete, hand leadership to the first
65/// eligible target, or complete while the existing leader remains in the target
66/// assignment.
67///
68/// `completion_ready` is Kafka's `PartitionReassignmentReplicas
69/// .maybeCompleteReassignment` predicate: every adding replica is in the ISR,
70/// and a replication factor decrease does not shrink the ISR. The host
71/// evaluates it, so a handoff, like a completion, never runs ahead of it.
72#[ensures(!completion_ready ==> result == ReassignmentAction::WaitForReplication)]
73#[ensures(completion_ready && !leader_removed ==> result == ReassignmentAction::Complete)]
74#[ensures(match result {
75    ReassignmentAction::WaitForLeader => completion_ready
76        && leader_removed
77        && (forall<i: Int> 0 <= i && i < [email protected]()
78            ==> !eligible_handoffs@[i]),
79    ReassignmentAction::Handoff(index) => completion_ready
80        && leader_removed
81        && index@ < [email protected]()
82        && eligible_handoffs@[index@]
83        && (forall<i: Int> 0 <= i && i < index@ ==> !eligible_handoffs@[i]),
84    _ => true,
85})]
86#[must_use]
87pub fn reassignment_action(
88    completion_ready: bool,
89    leader_removed: bool,
90    eligible_handoffs: &[bool],
91) -> ReassignmentAction {
92    if !completion_ready {
93        return ReassignmentAction::WaitForReplication;
94    }
95    if !leader_removed {
96        return ReassignmentAction::Complete;
97    }
98    let mut i = 0usize;
99    #[invariant(i@ <= [email protected]())]
100    #[invariant(forall<j: Int> 0 <= j && j < i@ ==> !eligible_handoffs@[j])]
101    #[variant([email protected]() - i@)]
102    while i < eligible_handoffs.len() {
103        if matches!(eligible_handoffs.get(i), Some(true)) {
104            return ReassignmentAction::Handoff(i);
105        }
106        i += 1;
107    }
108    ReassignmentAction::WaitForLeader
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn reassignment_phase_table_is_fail_closed() {
117        use ReassignmentAction::{Complete, Handoff, WaitForLeader, WaitForReplication};
118
119        assert2::assert!(reassignment_action(false, false, &[true]) == WaitForReplication);
120        assert2::assert!(reassignment_action(false, true, &[true]) == WaitForReplication);
121        assert2::assert!(reassignment_action(true, false, &[]) == Complete);
122        assert2::assert!(reassignment_action(true, true, &[false, true, true]) == Handoff(1));
123        assert2::assert!(reassignment_action(true, true, &[false]) == WaitForLeader);
124    }
125
126    #[yare::parameterized(
127        neither = {
128            false,
129            false,
130            ReassignmentSetMembership { in_union: false, adding: false, removing: false }
131        },
132        current_only = {
133            true,
134            false,
135            ReassignmentSetMembership { in_union: true, adding: false, removing: true }
136        },
137        target_only = {
138            false,
139            true,
140            ReassignmentSetMembership { in_union: true, adding: true, removing: false }
141        },
142        both = {
143            true,
144            true,
145            ReassignmentSetMembership { in_union: true, adding: false, removing: false }
146        },
147    )]
148    fn planning_kernels_cover_set_algebra_and_fences(
149        in_current: bool,
150        in_target: bool,
151        expected: ReassignmentSetMembership,
152    ) {
153        assert2::assert!(reassignment_set_membership(in_current, in_target) == expected);
154    }
155
156    #[yare::parameterized(
157        start_admitted = { (false, false, false), (true, true), (true, true), true },
158        cancel_admitted = { (true, true, true), (false, false), (true, true), true },
159        missing_target = { (false, false, false), (false, true), (true, true), false },
160        invalid_replication_factor = { (false, false, false), (true, false), (true, true), false },
161        cancel_not_approved = { (true, false, true), (false, false), (true, true), false },
162        cancel_not_in_progress = { (true, true, false), (false, false), (true, true), false },
163        ineligible_leader = { (false, false, false), (true, true), (false, true), false },
164        unavailable_epoch = { (false, false, false), (true, true), (true, false), false },
165    )]
166    fn planning_admission_requires_every_fence(
167        mode: (bool, bool, bool),
168        start: (bool, bool),
169        readiness: (bool, bool),
170        expected: bool,
171    ) {
172        assert2::assert!(reassignment_plan_admission(mode, start, readiness) == expected);
173    }
174}