use thiserror::Error;
use crate::compat::Compatibility;
use crate::stager::{
Activator, HealthCheck, Outcome, ReleaseId, SignatureVerifier, StageError, Stager, UpdateSource,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Checkpoint {
last_committed_turn: String,
paused_approvals: usize,
}
impl Checkpoint {
#[must_use]
pub fn new(last_committed_turn: impl Into<String>, paused_approvals: usize) -> Self {
Self {
last_committed_turn: last_committed_turn.into(),
paused_approvals,
}
}
#[must_use]
pub fn last_committed_turn(&self) -> &str {
&self.last_committed_turn
}
#[must_use]
pub const fn paused_approvals(&self) -> usize {
self.paused_approvals
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResumeReport {
resumed_from: String,
redriven_approvals: usize,
}
impl ResumeReport {
#[must_use]
pub fn new(resumed_from: impl Into<String>, redriven_approvals: usize) -> Self {
Self {
resumed_from: resumed_from.into(),
redriven_approvals,
}
}
#[must_use]
pub fn resumed_from(&self) -> &str {
&self.resumed_from
}
#[must_use]
pub const fn redriven_approvals(&self) -> usize {
self.redriven_approvals
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WarmSwapReport {
pub swap: Outcome,
pub checkpoint: Checkpoint,
pub resume: ResumeReport,
}
#[derive(Debug, Error)]
pub enum WarmSwapError {
#[error("refused: only a warm binary swap takes the click-to-restart path ({})", verdict_label(.0))]
NotWarm(Compatibility),
#[error("could not stop admitting new turns: {0}")]
Admission(String),
#[error("draining to a committed checkpoint failed: {0}")]
Drain(String),
#[error("lease {op} failed: {reason}")]
Lease {
op: &'static str,
reason: String,
},
#[error(transparent)]
Stage(#[from] StageError),
#[error("resuming from the last committed turn failed: {0}")]
Resume(String),
#[error(
"resume restarted from turn {resumed_from:?} but drained to {expected:?} — \
a committed turn would be dropped or replayed"
)]
ResumeMismatch {
resumed_from: String,
expected: String,
},
}
const fn verdict_label(verdict: &Compatibility) -> &'static str {
match verdict {
Compatibility::Warm => "warm",
Compatibility::Hot => "a config change that reloads with no restart",
Compatibility::Cold => "a format change that needs a coordinated redeploy",
Compatibility::Incompatible(_) => "built for a different runtime",
}
}
pub fn ensure_warm(verdict: &Compatibility) -> Result<(), WarmSwapError> {
if *verdict == Compatibility::Warm {
Ok(())
} else {
Err(WarmSwapError::NotWarm(verdict.clone()))
}
}
pub trait AdmissionGate {
fn close(&self) -> Result<(), WarmSwapError>;
}
impl<F> AdmissionGate for F
where
F: Fn() -> Result<(), WarmSwapError>,
{
fn close(&self) -> Result<(), WarmSwapError> {
self()
}
}
pub trait Drain {
fn drain(&self) -> Result<Checkpoint, WarmSwapError>;
}
impl<F> Drain for F
where
F: Fn() -> Result<Checkpoint, WarmSwapError>,
{
fn drain(&self) -> Result<Checkpoint, WarmSwapError> {
self()
}
}
pub trait LeaseControl {
fn release(&self) -> Result<(), WarmSwapError>;
fn reacquire(&self) -> Result<(), WarmSwapError>;
}
pub trait Resume {
fn resume(&self, checkpoint: &Checkpoint) -> Result<ResumeReport, WarmSwapError>;
}
impl<F> Resume for F
where
F: Fn(&Checkpoint) -> Result<ResumeReport, WarmSwapError>,
{
fn resume(&self, checkpoint: &Checkpoint) -> Result<ResumeReport, WarmSwapError> {
self(checkpoint)
}
}
pub struct WarmSwap<G, D, L, R> {
gate: G,
drain: D,
lease: L,
resume: R,
}
impl<G, D, L, R> WarmSwap<G, D, L, R>
where
G: AdmissionGate,
D: Drain,
L: LeaseControl,
R: Resume,
{
pub const fn new(gate: G, drain: D, lease: L, resume: R) -> Self {
Self {
gate,
drain,
lease,
resume,
}
}
pub fn run<S, V, A, H>(
&self,
verdict: &Compatibility,
stager: &mut Stager<S, V, A, H>,
release: &ReleaseId,
) -> Result<WarmSwapReport, WarmSwapError>
where
S: UpdateSource,
V: SignatureVerifier,
A: Activator,
H: HealthCheck,
{
ensure_warm(verdict)?;
self.gate.close()?;
let checkpoint = self.drain.drain()?;
self.lease
.release()
.map_err(|e| relabel_lease(e, "release"))?;
let swap = stager.stage_and_apply(release);
self.lease
.reacquire()
.map_err(|e| relabel_lease(e, "re-acquire"))?;
let swap = swap?;
let resume = self.resume.resume(&checkpoint)?;
if resume.resumed_from() != checkpoint.last_committed_turn() {
return Err(WarmSwapError::ResumeMismatch {
resumed_from: resume.resumed_from().to_owned(),
expected: checkpoint.last_committed_turn().to_owned(),
});
}
Ok(WarmSwapReport {
swap,
checkpoint,
resume,
})
}
}
fn relabel_lease(err: WarmSwapError, op: &'static str) -> WarmSwapError {
match err {
WarmSwapError::Lease { reason, .. } => WarmSwapError::Lease { op, reason },
other => other,
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use std::cell::RefCell;
use std::path::PathBuf;
use std::rc::Rc;
use super::*;
use crate::compat::{Incompatibility, RuntimeTarget, StagedBundle};
use crate::stager::{Health, StagedArtifact};
type Trace = Rc<RefCell<Vec<String>>>;
fn trace() -> Trace {
Rc::new(RefCell::new(Vec::new()))
}
fn log(trace: &Trace, step: impl Into<String>) {
trace.borrow_mut().push(step.into());
}
#[derive(Clone)]
struct RecordingGate {
trace: Trace,
open: Rc<RefCell<bool>>,
}
impl RecordingGate {
fn new(trace: &Trace) -> Self {
Self {
trace: Rc::clone(trace),
open: Rc::new(RefCell::new(true)),
}
}
fn admits(&self) -> bool {
*self.open.borrow()
}
}
impl AdmissionGate for RecordingGate {
fn close(&self) -> Result<(), WarmSwapError> {
*self.open.borrow_mut() = false;
log(&self.trace, "close");
Ok(())
}
}
#[derive(Clone)]
struct RecordingDrain {
trace: Trace,
gate: RecordingGate,
log: Vec<(&'static str, bool)>,
paused_approvals: usize,
}
impl RecordingDrain {
fn new(
trace: &Trace,
gate: &RecordingGate,
model_log: Vec<(&'static str, bool)>,
paused_approvals: usize,
) -> Self {
Self {
trace: Rc::clone(trace),
gate: gate.clone(),
log: model_log,
paused_approvals,
}
}
}
impl Drain for RecordingDrain {
fn drain(&self) -> Result<Checkpoint, WarmSwapError> {
assert!(
!self.gate.admits(),
"drain began while new turns were still admitted"
);
log(&self.trace, "drain");
let last_committed = self
.log
.iter()
.rev()
.find_map(|(id, committed)| committed.then_some(*id))
.expect("model log has at least one committed turn");
Ok(Checkpoint::new(last_committed, self.paused_approvals))
}
}
#[derive(Clone)]
struct RecordingLease {
trace: Trace,
held: Rc<RefCell<bool>>,
fail_release: bool,
fail_reacquire: bool,
}
impl RecordingLease {
fn new(trace: &Trace) -> Self {
Self {
trace: Rc::clone(trace),
held: Rc::new(RefCell::new(true)),
fail_release: false,
fail_reacquire: false,
}
}
fn failing_reacquire(trace: &Trace) -> Self {
Self {
fail_reacquire: true,
..Self::new(trace)
}
}
fn is_held(&self) -> bool {
*self.held.borrow()
}
}
impl LeaseControl for RecordingLease {
fn release(&self) -> Result<(), WarmSwapError> {
if self.fail_release {
return Err(WarmSwapError::Lease {
op: "release",
reason: "api timeout".to_owned(),
});
}
assert!(self.is_held(), "released a lease that was not held");
*self.held.borrow_mut() = false;
log(&self.trace, "lease.release");
Ok(())
}
fn reacquire(&self) -> Result<(), WarmSwapError> {
if self.fail_reacquire {
return Err(WarmSwapError::Lease {
op: "re-acquire",
reason: "held by another".to_owned(),
});
}
assert!(
!self.is_held(),
"re-acquired a lease that was never released — split-brain writer"
);
*self.held.borrow_mut() = true;
log(&self.trace, "lease.reacquire");
Ok(())
}
}
#[derive(Clone)]
struct RecordingResume {
trace: Trace,
gate: RecordingGate,
lease: RecordingLease,
drift_to: Option<&'static str>,
}
impl RecordingResume {
fn new(trace: &Trace, gate: &RecordingGate, lease: &RecordingLease) -> Self {
Self {
trace: Rc::clone(trace),
gate: gate.clone(),
lease: lease.clone(),
drift_to: None,
}
}
fn drifting(
trace: &Trace,
gate: &RecordingGate,
lease: &RecordingLease,
to: &'static str,
) -> Self {
Self {
drift_to: Some(to),
..Self::new(trace, gate, lease)
}
}
}
impl Resume for RecordingResume {
fn resume(&self, checkpoint: &Checkpoint) -> Result<ResumeReport, WarmSwapError> {
assert!(self.lease.is_held(), "resumed without a re-acquired lease");
log(
&self.trace,
format!("resume@{}", checkpoint.last_committed_turn()),
);
*self.gate.open.borrow_mut() = true;
let from = self.drift_to.unwrap_or(checkpoint.last_committed_turn());
Ok(ResumeReport::new(from, checkpoint.paused_approvals()))
}
}
#[derive(Clone)]
struct SwapActivator {
trace: Trace,
}
impl SwapActivator {
fn new(trace: &Trace) -> Self {
Self {
trace: Rc::clone(trace),
}
}
}
impl Activator for SwapActivator {
fn activate(&self, release: &ReleaseId) -> Result<(), StageError> {
log(&self.trace, format!("swap:{release}"));
Ok(())
}
}
fn artifact_for(release: &ReleaseId) -> StagedArtifact {
StagedArtifact {
release: release.clone(),
staged_path: PathBuf::from(format!("/var/lib/polychrome/staged/{release}")),
bundle: StagedBundle::new(
RuntimeTarget::new(3, 7, "polychrome.uno/v1"),
format!("catalog-{release}"),
),
signed_bytes: format!("bytes-of-{release}").into_bytes(),
signature: vec![0xAB; 4],
signer_public_key: vec![0xCD; 4],
}
}
fn ok_source(release: &ReleaseId) -> Result<StagedArtifact, StageError> {
Ok(artifact_for(release))
}
#[test]
fn ensure_warm_admits_only_warm() {
assert!(ensure_warm(&Compatibility::Warm).is_ok());
for verdict in [
Compatibility::Hot,
Compatibility::Cold,
Compatibility::Incompatible(Incompatibility::Wire),
] {
let err = ensure_warm(&verdict).unwrap_err();
assert!(
matches!(err, WarmSwapError::NotWarm(v) if v == verdict),
"non-warm verdict must be refused: {verdict:?}",
);
}
}
#[test]
fn a_non_warm_release_touches_no_seam() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(&trace, &gate, vec![("t1", true)], 0);
let lease = RecordingLease::new(&trace);
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate.clone(), drain, lease.clone(), resume);
let mut stager = Stager::new(
|_: &ReleaseId| -> Result<StagedArtifact, StageError> {
panic!("download must not run for a non-warm release")
},
|_: &StagedArtifact| panic!("verify must not run for a non-warm release"),
SwapActivator::new(&trace),
|| panic!("health check must not run for a non-warm release"),
ReleaseId::new("v1"),
);
let err = swap
.run(&Compatibility::Hot, &mut stager, &ReleaseId::new("v2"))
.unwrap_err();
assert!(matches!(err, WarmSwapError::NotWarm(Compatibility::Hot)));
assert!(
gate.admits(),
"admission must stay open for a refused release"
);
assert!(
lease.is_held(),
"the lease must stay held for a refused release"
);
assert!(
trace.borrow().is_empty(),
"no seam runs for a refused release"
);
}
#[test]
fn warm_swap_walks_close_drain_release_swap_reacquire_resume_in_order() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(
&trace,
&gate,
vec![("t1", true), ("t2-interrupted", false)],
2,
);
let lease = RecordingLease::new(&trace);
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate.clone(), drain, lease.clone(), resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
SwapActivator::new(&trace),
|| Health::Healthy,
ReleaseId::new("v1"),
);
let report = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap();
assert_eq!(
*trace.borrow(),
vec![
"close".to_owned(),
"drain".to_owned(),
"lease.release".to_owned(),
"swap:v2".to_owned(),
"lease.reacquire".to_owned(),
"resume@t1".to_owned(),
],
);
assert_eq!(
report.swap,
Outcome::Committed {
version: ReleaseId::new("v2"),
}
);
assert_eq!(report.checkpoint.last_committed_turn(), "t1");
assert_eq!(report.resume.resumed_from(), "t1");
assert_eq!(report.checkpoint.paused_approvals(), 2);
assert_eq!(report.resume.redriven_approvals(), 2);
assert!(lease.is_held());
assert!(gate.admits());
}
#[test]
fn no_new_turn_is_admitted_once_the_drain_begins() {
let trace = trace();
let gate = RecordingGate::new(&trace);
assert!(gate.admits(), "admission starts open");
gate.close().unwrap();
assert!(
!gate.admits(),
"a turn arriving after admission closes is refused"
);
}
#[test]
fn a_paused_approval_survives_the_warm_restart() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(&trace, &gate, vec![("t7", true)], 3);
let lease = RecordingLease::new(&trace);
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate, drain, lease, resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
SwapActivator::new(&trace),
|| Health::Healthy,
ReleaseId::new("v1"),
);
let report = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap();
assert_eq!(report.checkpoint.paused_approvals(), 3);
assert_eq!(report.resume.redriven_approvals(), 3);
assert_eq!(report.resume.resumed_from(), "t7");
}
#[test]
fn a_rolled_back_swap_still_reacquires_the_lease_and_resumes() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(&trace, &gate, vec![("t1", true)], 1);
let lease = RecordingLease::new(&trace);
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate.clone(), drain, lease.clone(), resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
SwapActivator::new(&trace),
|| Health::Unhealthy("readiness probe timed out".to_owned()),
ReleaseId::new("v1"),
);
let report = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap();
assert_eq!(
report.swap,
Outcome::RolledBack {
stayed_on: ReleaseId::new("v1"),
reason: "readiness probe timed out".to_owned(),
}
);
assert!(
lease.is_held(),
"lease re-acquired even on a rolled-back swap"
);
assert!(gate.admits(), "admission re-opened on the previous binary");
assert_eq!(report.resume.resumed_from(), "t1");
assert_eq!(
*trace.borrow(),
vec![
"close".to_owned(),
"drain".to_owned(),
"lease.release".to_owned(),
"swap:v2".to_owned(),
"swap:v1".to_owned(),
"lease.reacquire".to_owned(),
"resume@t1".to_owned(),
],
);
}
#[test]
fn an_unverified_binary_is_refused_and_the_lease_is_reacquired() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(&trace, &gate, vec![("t1", true)], 0);
let lease = RecordingLease::new(&trace);
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate, drain, lease.clone(), resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| false,
SwapActivator::new(&trace),
|| panic!("health check must not run for an unverified binary"),
ReleaseId::new("v1"),
);
let err = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap_err();
assert!(matches!(err, WarmSwapError::Stage(StageError::Unverified)));
assert!(
lease.is_held(),
"the lease is re-acquired after a refused swap"
);
assert_eq!(
*trace.borrow(),
vec![
"close".to_owned(),
"drain".to_owned(),
"lease.release".to_owned(),
"lease.reacquire".to_owned(),
],
);
assert!(!trace.borrow().iter().any(|s| s.starts_with("resume")));
}
#[test]
fn resume_from_the_wrong_turn_is_a_mismatch_error() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(&trace, &gate, vec![("t1", true)], 0);
let lease = RecordingLease::new(&trace);
let resume = RecordingResume::drifting(&trace, &gate, &lease, "t0-stale");
let swap = WarmSwap::new(gate, drain, lease, resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
SwapActivator::new(&trace),
|| Health::Healthy,
ReleaseId::new("v1"),
);
let err = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap_err();
assert!(matches!(
err,
WarmSwapError::ResumeMismatch { resumed_from, expected }
if resumed_from == "t0-stale" && expected == "t1"
));
}
#[test]
fn a_failed_drain_never_releases_the_lease_or_swaps() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let lease = RecordingLease::new(&trace);
let failing_drain = {
let trace = Rc::clone(&trace);
move || -> Result<Checkpoint, WarmSwapError> {
log(&trace, "drain");
Err(WarmSwapError::Drain("journal sync failed".to_owned()))
}
};
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate.clone(), failing_drain, lease.clone(), resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| panic!("verify must not run when the drain fails"),
SwapActivator::new(&trace),
|| panic!("health check must not run when the drain fails"),
ReleaseId::new("v1"),
);
let err = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap_err();
assert!(matches!(err, WarmSwapError::Drain(_)));
assert!(lease.is_held(), "a failed drain keeps the lease held");
assert_eq!(
*trace.borrow(),
vec!["close".to_owned(), "drain".to_owned()]
);
}
#[test]
fn a_failed_reacquire_surfaces_as_a_lease_error() {
let trace = trace();
let gate = RecordingGate::new(&trace);
let drain = RecordingDrain::new(&trace, &gate, vec![("t1", true)], 0);
let lease = RecordingLease::failing_reacquire(&trace);
let resume = RecordingResume::new(&trace, &gate, &lease);
let swap = WarmSwap::new(gate, drain, lease, resume);
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
SwapActivator::new(&trace),
|| Health::Healthy,
ReleaseId::new("v1"),
);
let err = swap
.run(&Compatibility::Warm, &mut stager, &ReleaseId::new("v2"))
.unwrap_err();
assert!(matches!(
err,
WarmSwapError::Lease {
op: "re-acquire",
..
}
));
}
#[test]
fn not_warm_error_reads_plainly() {
let msg = WarmSwapError::NotWarm(Compatibility::Hot).to_string();
assert_eq!(
msg,
"refused: only a warm binary swap takes the click-to-restart path \
(a config change that reloads with no restart)",
);
for banned in ["sorry", "please", "unfortunately"] {
assert!(!msg.to_lowercase().contains(banned));
}
}
}