1use std::collections::BTreeSet;
4
5pub const TWO_PHASE_MAX_GROUPS: usize = 3;
7pub const TWO_PHASE_MAX_STEPS: usize = 16;
9pub const TWO_PHASE_MAX_PAYLOAD: usize = 64 * 1024;
11pub const TWO_PHASE_DEFAULT_PREPARE_TIMEOUT_MS: u64 = 300_000;
13
14#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct TwoPhaseStep {
17 pub key: Vec<u8>,
19 pub command: Vec<u8>,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct TwoPhasePlan {
26 pub tx_id: Vec<u8>,
28 pub steps: Vec<TwoPhaseStep>,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
34pub enum TwoPhasePlanError {
35 EmptyTxId,
37 EmptyPlan,
39 TooManySteps,
41 PayloadTooLarge {
43 step: usize,
45 },
46 UnroutableKey {
48 step: usize,
50 },
51 TooManyGroups {
53 groups: usize,
55 },
56}
57
58impl std::fmt::Display for TwoPhasePlanError {
59 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60 match self {
61 Self::EmptyTxId => f.write_str("transaction id must not be empty"),
62 Self::EmptyPlan => f.write_str("plan must contain at least one step"),
63 Self::TooManySteps => write!(f, "plan exceeds {TWO_PHASE_MAX_STEPS} steps"),
64 Self::PayloadTooLarge { step } => {
65 write!(
66 f,
67 "step {step} payload exceeds {TWO_PHASE_MAX_PAYLOAD} bytes"
68 )
69 }
70 Self::UnroutableKey { step } => write!(f, "step {step} key is not routable"),
71 Self::TooManyGroups { groups } => write!(
72 f,
73 "plan spans {groups} groups; maximum is {TWO_PHASE_MAX_GROUPS}"
74 ),
75 }
76 }
77}
78
79impl std::error::Error for TwoPhasePlanError {}
80
81pub fn validate_two_phase_plan(
88 plan: &TwoPhasePlan,
89 group_for_key: impl Fn(&[u8]) -> Option<u32>,
90) -> Result<(), TwoPhasePlanError> {
91 if plan.tx_id.is_empty() {
92 return Err(TwoPhasePlanError::EmptyTxId);
93 }
94 if plan.steps.is_empty() {
95 return Err(TwoPhasePlanError::EmptyPlan);
96 }
97 if plan.steps.len() > TWO_PHASE_MAX_STEPS {
98 return Err(TwoPhasePlanError::TooManySteps);
99 }
100
101 let mut groups = BTreeSet::new();
102 for (step, item) in plan.steps.iter().enumerate() {
103 if item.command.len() > TWO_PHASE_MAX_PAYLOAD {
104 return Err(TwoPhasePlanError::PayloadTooLarge { step });
105 }
106 let Some(group) = group_for_key(&item.key) else {
107 return Err(TwoPhasePlanError::UnroutableKey { step });
108 };
109 groups.insert(group);
110 }
111 if groups.len() > TWO_PHASE_MAX_GROUPS {
112 return Err(TwoPhasePlanError::TooManyGroups {
113 groups: groups.len(),
114 });
115 }
116 Ok(())
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122
123 #[test]
124 fn accepts_two_group_plan() {
125 let plan = TwoPhasePlan {
126 tx_id: b"tx".to_vec(),
127 steps: vec![
128 TwoPhaseStep {
129 key: b"a".to_vec(),
130 command: vec![1],
131 },
132 TwoPhaseStep {
133 key: b"b".to_vec(),
134 command: vec![2],
135 },
136 ],
137 };
138 validate_two_phase_plan(&plan, |key| Some(u32::from(key != b"a"))).expect("valid");
139 }
140
141 #[test]
142 fn rejects_four_groups() {
143 let plan = TwoPhasePlan {
144 tx_id: b"tx".to_vec(),
145 steps: (0..4)
146 .map(|i| TwoPhaseStep {
147 key: vec![u8::try_from(i).expect("test key fits u8")],
148 command: vec![1],
149 })
150 .collect(),
151 };
152 assert!(matches!(
153 validate_two_phase_plan(&plan, |key| Some(u32::from(key[0]))),
154 Err(TwoPhasePlanError::TooManyGroups { groups: 4 })
155 ));
156 }
157}