use crate::error::ErrorClass;
use std::fmt;
pub mod driver;
pub const SPLIT_ID_MAX_LEN: usize = 128;
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SplitId(String);
impl SplitId {
pub fn new(id: impl Into<String>) -> Result<SplitId, CoordinationError> {
let id = id.into();
if id.is_empty() || id.len() > SPLIT_ID_MAX_LEN {
return Err(CoordinationError::new(
CoordinationErrorKind::Fatal,
format!(
"split id must be 1..={SPLIT_ID_MAX_LEN} bytes, got {} ({id:?})",
id.len()
),
));
}
if let Some(bad) = id
.chars()
.find(|c| !(c.is_ascii_alphanumeric() || *c == '_' || *c == '-'))
{
return Err(CoordinationError::new(
CoordinationErrorKind::Fatal,
format!("split id may only contain [A-Za-z0-9_-], got {bad:?} in {id:?}"),
));
}
Ok(SplitId(id))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for SplitId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct LeaseEpoch(pub u64);
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct SplitSpec {
pub id: SplitId,
pub descriptor: Vec<u8>,
pub weight: u64,
}
impl SplitSpec {
#[must_use]
pub fn new(id: SplitId, descriptor: Vec<u8>) -> SplitSpec {
SplitSpec {
id,
descriptor,
weight: 1,
}
}
#[must_use]
pub fn with_weight(mut self, weight: u64) -> SplitSpec {
self.weight = weight;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct PlannedSplit {
pub spec: SplitSpec,
pub seed: Option<SplitProgress>,
}
impl PlannedSplit {
#[must_use]
pub fn new(spec: SplitSpec) -> PlannedSplit {
PlannedSplit { spec, seed: None }
}
#[must_use]
pub fn with_seed(mut self, seed: SplitProgress) -> PlannedSplit {
self.seed = Some(seed);
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PlanFinality {
Open,
Final,
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct SplitPlan {
pub splits: Vec<PlannedSplit>,
pub finality: PlanFinality,
pub planner_state: Option<Vec<u8>>,
}
impl SplitPlan {
#[must_use]
pub fn new(splits: Vec<PlannedSplit>, finality: PlanFinality) -> SplitPlan {
SplitPlan {
splits,
finality,
planner_state: None,
}
}
#[must_use]
pub fn with_planner_state(mut self, state: Vec<u8>) -> SplitPlan {
self.planner_state = Some(state);
self
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct PlanContext<'a> {
pub planner_state: Option<&'a [u8]>,
pub generation: u64,
}
impl<'a> PlanContext<'a> {
#[must_use]
pub fn new(planner_state: Option<&'a [u8]>, generation: u64) -> PlanContext<'a> {
PlanContext {
planner_state,
generation,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct SplitProgress {
pub watermark: i64,
pub state: Vec<u8>,
pub completed: bool,
}
impl SplitProgress {
#[must_use]
pub fn new(watermark: i64, state: Vec<u8>) -> SplitProgress {
SplitProgress {
watermark,
state,
completed: false,
}
}
#[must_use]
pub fn completed(watermark: i64, state: Vec<u8>) -> SplitProgress {
SplitProgress {
watermark,
state,
completed: true,
}
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum CoordinationEvent {
Gained {
split: SplitSpec,
epoch: LeaseEpoch,
progress: Option<SplitProgress>,
},
Lost {
split: SplitId,
},
RevokeRequested {
split: SplitId,
},
Quarantined {
split: SplitId,
attempts: u32,
},
AllComplete,
Stalled {
completed: u64,
quarantined: u64,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum CoordinationErrorKind {
Fenced,
Retryable,
Fatal,
}
#[derive(Debug, thiserror::Error)]
#[error("coordination error ({kind:?}): {reason}")]
#[non_exhaustive]
pub struct CoordinationError {
pub kind: CoordinationErrorKind,
pub reason: String,
}
impl CoordinationError {
pub fn new(kind: CoordinationErrorKind, reason: impl Into<String>) -> CoordinationError {
CoordinationError {
kind,
reason: reason.into(),
}
}
#[must_use]
pub fn class(&self) -> ErrorClass {
match self.kind {
CoordinationErrorKind::Retryable => ErrorClass::Retryable,
CoordinationErrorKind::Fenced | CoordinationErrorKind::Fatal => ErrorClass::Fatal,
}
}
}
pub trait SplitPlanner: Send {
fn fingerprint(&self) -> String;
fn plan(&mut self, ctx: PlanContext<'_>) -> Result<SplitPlan, CoordinationError>;
}
pub trait SplitCoordinator: Send {
fn start(&mut self, planner: Box<dyn SplitPlanner>) -> Result<(), CoordinationError>;
fn set_waker(&mut self, waker: ControlWaker);
fn poll(&mut self) -> Result<Vec<CoordinationEvent>, CoordinationError>;
fn commit(
&mut self,
split: &SplitId,
progress: &SplitProgress,
) -> Result<(), CoordinationError>;
fn fail(&mut self, split: &SplitId, reason: &str) -> Result<(), CoordinationError>;
fn release(&mut self, splits: &[SplitId]) -> Result<(), CoordinationError>;
fn release_drained(&mut self, splits: &[SplitId]) -> Result<(), CoordinationError> {
self.release(splits)
}
fn decline_revoke(&mut self, _split: &SplitId) -> Result<(), CoordinationError> {
Ok(())
}
}
#[derive(Clone, Debug)]
pub struct ControlWaker(crossbeam_channel::Sender<()>);
impl ControlWaker {
#[must_use]
pub fn inert() -> ControlWaker {
let (tx, rx) = crossbeam_channel::bounded(1);
drop(rx);
ControlWaker(tx)
}
pub fn wake(&self) {
let _ = self.0.try_send(());
}
}
pub(crate) fn control_channel() -> (ControlWaker, crossbeam_channel::Receiver<()>) {
let (tx, rx) = crossbeam_channel::bounded(1);
(ControlWaker(tx), rx)
}
#[cfg(test)]
mod tests {
use super::*;
struct NoopCoordinator;
impl SplitCoordinator for NoopCoordinator {
fn start(&mut self, _planner: Box<dyn SplitPlanner>) -> Result<(), CoordinationError> {
Ok(())
}
fn set_waker(&mut self, _waker: ControlWaker) {}
fn poll(&mut self) -> Result<Vec<CoordinationEvent>, CoordinationError> {
Ok(vec![])
}
fn commit(
&mut self,
split: &SplitId,
_progress: &SplitProgress,
) -> Result<(), CoordinationError> {
Err(CoordinationError::new(
CoordinationErrorKind::Fenced,
format!("split {split} is owned by a peer at epoch 2"),
))
}
fn fail(&mut self, _split: &SplitId, _reason: &str) -> Result<(), CoordinationError> {
Ok(())
}
fn release(&mut self, _splits: &[SplitId]) -> Result<(), CoordinationError> {
Ok(())
}
}
struct NoopPlanner;
impl SplitPlanner for NoopPlanner {
fn fingerprint(&self) -> String {
"noop:v1".into()
}
fn plan(&mut self, ctx: PlanContext<'_>) -> Result<SplitPlan, CoordinationError> {
assert!(ctx.planner_state.is_none());
Ok(SplitPlan::new(vec![], PlanFinality::Final))
}
}
#[test]
fn split_coordinator_is_object_safe() {
let mut c: Box<dyn SplitCoordinator> = Box::new(NoopCoordinator);
c.start(Box::new(NoopPlanner)).unwrap();
assert!(c.poll().unwrap().is_empty());
c.release(&[SplitId::new("s-0").unwrap()]).unwrap();
c.release_drained(&[SplitId::new("s-0").unwrap()]).unwrap();
}
#[test]
fn split_ids_validate_charset_and_length() {
assert_eq!(
SplitId::new("rows-000000-000125").unwrap().as_str(),
"rows-000000-000125"
);
assert_eq!(SplitId::new("A_z9").unwrap().to_string(), "A_z9");
for bad in [
"",
"a.b",
"a b",
"a/b",
"å",
&"x".repeat(SPLIT_ID_MAX_LEN + 1),
] {
let err = SplitId::new(bad).unwrap_err();
assert_eq!(err.kind, CoordinationErrorKind::Fatal, "{bad:?}");
}
assert!(SplitId::new("x".repeat(SPLIT_ID_MAX_LEN)).is_ok());
}
#[test]
fn error_kinds_map_to_the_framework_taxonomy() {
let fenced = CoordinationError::new(CoordinationErrorKind::Fenced, "seized");
assert_eq!(fenced.class(), ErrorClass::Fatal, "unintercepted backstop");
assert_eq!(
CoordinationError::new(CoordinationErrorKind::Retryable, "timeout").class(),
ErrorClass::Retryable
);
assert_eq!(
CoordinationError::new(CoordinationErrorKind::Fatal, "no CAS").class(),
ErrorClass::Fatal
);
assert!(fenced.to_string().contains("seized"));
}
#[test]
fn builders_cover_the_non_exhaustive_structs() {
let spec = SplitSpec::new(SplitId::new("s").unwrap(), b"d".to_vec());
assert_eq!(spec.weight, 1, "default weight");
assert_eq!(spec.clone().with_weight(64 << 20).weight, 64 << 20);
let planned = PlannedSplit::new(spec.clone());
assert!(planned.seed.is_none());
let seeded = planned.with_seed(SplitProgress::new(7, b"state".to_vec()));
assert_eq!(seeded.seed.as_ref().unwrap().watermark, 7);
let plan = SplitPlan::new(vec![seeded], PlanFinality::Open);
assert!(plan.planner_state.is_none());
assert_eq!(
plan.with_planner_state(b"cursor".to_vec())
.planner_state
.as_deref(),
Some(b"cursor".as_slice())
);
let running = SplitProgress::new(7, vec![]);
assert!(!running.completed);
assert!(SplitProgress::completed(7, vec![]).completed);
let ctx = PlanContext::new(Some(b"cursor"), 3);
assert_eq!(ctx.generation, 3);
}
#[test]
fn lease_epochs_order_across_owners() {
assert!(LeaseEpoch(2) > LeaseEpoch(1));
}
}