use std::collections::{BTreeMap, BTreeSet};
use semver::Version;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BundleMeta {
pub version: Version,
pub archive_sha256: String,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub struct State {
pub committed: Option<u64>,
pub last_good: Option<u64>,
pub staged: Option<u64>,
pub booting: Option<u64>,
pub failed: BTreeSet<String>,
pub max_version_seen: Option<Version>,
pub versions: BTreeMap<u64, BundleMeta>,
}
impl State {
fn referenced(&self) -> BTreeSet<u64> {
[self.committed, self.last_good, self.staged, self.booting]
.into_iter()
.flatten()
.collect()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Active {
Embedded,
Ota(u64),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Effect {
DeleteBundle(u64),
}
#[derive(Debug)]
pub struct BootOutcome {
pub state: State,
pub active: Active,
pub rolled_back: Option<u64>,
pub effects: Vec<Effect>,
}
pub fn resolve_boot(mut state: State, embedded_version: &Version, present: &BTreeSet<u64>) -> BootOutcome {
let rolled_back = state.booting.take();
if let Some(seq) = rolled_back {
if let Some(meta) = state.versions.get(&seq) {
state.failed.insert(meta.archive_sha256.clone());
}
}
state.max_version_seen = Some(match state.max_version_seen.take() {
Some(seen) => seen.max(embedded_version.clone()),
None => embedded_version.clone(),
});
let ota_beats_embedded = |seq: u64, state: &State| -> bool {
present.contains(&seq)
&& state
.versions
.get(&seq)
.is_some_and(|meta| meta.version > *embedded_version)
};
if let Some(seq) = state.committed {
if !ota_beats_embedded(seq, &state) {
state.committed = None;
}
}
if let Some(seq) = state.last_good {
if !ota_beats_embedded(seq, &state) {
state.last_good = None;
}
}
if let Some(seq) = state.staged.take() {
let arm = present.contains(&seq)
&& state.versions.get(&seq).is_some_and(|meta| {
!state.failed.contains(&meta.archive_sha256)
&& state
.max_version_seen
.as_ref()
.map_or(true, |seen| meta.version > *seen)
});
if arm {
state.booting = Some(seq);
}
}
let active = match state.booting.or(state.committed) {
Some(seq) => Active::Ota(seq),
None => Active::Embedded,
};
let referenced = state.referenced();
let version_seqs: BTreeSet<u64> = state.versions.keys().copied().collect();
let mut effects = Vec::new();
for seq in present.union(&version_seqs).copied() {
if !referenced.contains(&seq) {
if present.contains(&seq) {
effects.push(Effect::DeleteBundle(seq));
}
state.versions.remove(&seq);
}
}
BootOutcome {
state,
active,
rolled_back,
effects,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AckOutcome {
Committed(u64),
AlreadyCommitted(u64),
EmbeddedNoop,
Stale(u64),
}
pub fn ack(mut state: State, booted: Active) -> (State, AckOutcome, Vec<Effect>) {
let seq = match booted {
Active::Embedded => return (state, AckOutcome::EmbeddedNoop, Vec::new()),
Active::Ota(seq) => seq,
};
if state.booting != Some(seq) {
let outcome = if state.committed == Some(seq) {
AckOutcome::AlreadyCommitted(seq)
} else {
AckOutcome::Stale(seq)
};
return (state, outcome, Vec::new());
}
state.booting = None;
let previous = state.committed.replace(seq);
state.last_good = Some(seq);
if let Some(meta) = state.versions.get(&seq) {
state.max_version_seen = Some(match state.max_version_seen.take() {
Some(seen) => seen.max(meta.version.clone()),
None => meta.version.clone(),
});
}
let mut effects = Vec::new();
if let Some(prev) = previous {
if !state.referenced().contains(&prev) {
state.versions.remove(&prev);
effects.push(Effect::DeleteBundle(prev));
}
}
(state, AckOutcome::Committed(seq), effects)
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum StageError {
#[error("bundle version {version} is not newer than the watermark {watermark}")]
VersionNotNewer { version: Version, watermark: Version },
#[error("archive sha256 {0} previously failed a trial boot and is blacklisted")]
HashBlacklisted(String),
#[error("seq {0} is already referenced by the state")]
SeqInUse(u64),
}
pub fn stage(
mut state: State,
seq: u64,
meta: BundleMeta,
) -> (State, Result<Vec<Effect>, StageError>) {
if let Some(watermark) = &state.max_version_seen {
if meta.version <= *watermark {
let err = StageError::VersionNotNewer {
version: meta.version,
watermark: watermark.clone(),
};
return (state, Err(err));
}
}
if state.failed.contains(&meta.archive_sha256) {
let err = StageError::HashBlacklisted(meta.archive_sha256);
return (state, Err(err));
}
if state.referenced().contains(&seq) || state.versions.contains_key(&seq) {
return (state, Err(StageError::SeqInUse(seq)));
}
let mut effects = Vec::new();
if let Some(prev) = state.staged.replace(seq) {
state.versions.remove(&prev);
effects.push(Effect::DeleteBundle(prev));
}
state.versions.insert(seq, meta);
(state, Ok(effects))
}
pub fn reset(state: State, embedded_version: &Version, active: Active) -> (State, Vec<Effect>) {
let keep = match active {
Active::Embedded => None,
Active::Ota(seq) => Some(seq),
};
let effects = state
.versions
.keys()
.copied()
.filter(|seq| Some(*seq) != keep)
.map(Effect::DeleteBundle)
.collect();
let fresh = State {
max_version_seen: Some(embedded_version.clone()),
..State::default()
};
(fresh, effects)
}
#[cfg(test)]
mod tests;