use std::path::PathBuf;
use thiserror::Error;
use crate::compat::StagedBundle;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ReleaseId(String);
impl ReleaseId {
#[must_use]
pub fn new(id: impl Into<String>) -> Self {
Self(id.into())
}
#[must_use]
pub const fn as_str(&self) -> &str {
self.0.as_str()
}
}
impl std::fmt::Display for ReleaseId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Debug, Clone)]
pub struct StagedArtifact {
pub release: ReleaseId,
pub staged_path: PathBuf,
pub bundle: StagedBundle,
pub signed_bytes: Vec<u8>,
pub signature: Vec<u8>,
pub signer_public_key: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Health {
Healthy,
Unhealthy(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Outcome {
Committed {
version: ReleaseId,
},
RolledBack {
stayed_on: ReleaseId,
reason: String,
},
}
impl Outcome {
#[must_use]
pub fn summary(&self) -> String {
match self {
Self::Committed { version } => format!("Now running {version}."),
Self::RolledBack { stayed_on, reason } => {
format!(
"Stayed on {stayed_on} — the new version failed its health check ({reason})."
)
}
}
}
}
#[derive(Debug, Error)]
pub enum StageError {
#[error("staging the update failed: {0}")]
Download(String),
#[error("the update's signature did not verify against a trusted key")]
Unverified,
#[error("activating the staged update failed: {0}")]
Activate(String),
}
pub trait UpdateSource {
fn fetch(&self, release: &ReleaseId) -> Result<StagedArtifact, StageError>;
}
impl<F> UpdateSource for F
where
F: Fn(&ReleaseId) -> Result<StagedArtifact, StageError>,
{
fn fetch(&self, release: &ReleaseId) -> Result<StagedArtifact, StageError> {
self(release)
}
}
pub trait SignatureVerifier {
fn verify(&self, artifact: &StagedArtifact) -> bool;
}
impl<F> SignatureVerifier for F
where
F: Fn(&StagedArtifact) -> bool,
{
fn verify(&self, artifact: &StagedArtifact) -> bool {
self(artifact)
}
}
pub trait Activator {
fn activate(&self, release: &ReleaseId) -> Result<(), StageError>;
}
impl<F> Activator for F
where
F: Fn(&ReleaseId) -> Result<(), StageError>,
{
fn activate(&self, release: &ReleaseId) -> Result<(), StageError> {
self(release)
}
}
pub trait HealthCheck {
fn check(&self) -> Health;
}
impl<F> HealthCheck for F
where
F: Fn() -> Health,
{
fn check(&self) -> Health {
self()
}
}
pub struct Stager<S, V, A, H> {
source: S,
verifier: V,
activator: A,
health: H,
live: ReleaseId,
previous: Option<ReleaseId>,
}
impl<S, V, A, H> Stager<S, V, A, H>
where
S: UpdateSource,
V: SignatureVerifier,
A: Activator,
H: HealthCheck,
{
pub const fn new(
source: S,
verifier: V,
activator: A,
health: H,
initial_live: ReleaseId,
) -> Self {
Self {
source,
verifier,
activator,
health,
live: initial_live,
previous: None,
}
}
#[must_use]
pub const fn live(&self) -> &ReleaseId {
&self.live
}
#[must_use]
pub const fn previous(&self) -> Option<&ReleaseId> {
self.previous.as_ref()
}
pub fn stage_and_apply(&mut self, release: &ReleaseId) -> Result<Outcome, StageError> {
let artifact = self.source.fetch(release)?;
if !self.verifier.verify(&artifact) {
return Err(StageError::Unverified);
}
let rollback_to = self.live.clone();
self.activator.activate(&artifact.release)?;
match self.health.check() {
Health::Healthy => {
self.previous = Some(rollback_to);
self.live = artifact.release.clone();
Ok(Outcome::Committed {
version: artifact.release,
})
}
Health::Unhealthy(reason) => {
self.activator.activate(&rollback_to)?;
Ok(Outcome::RolledBack {
stayed_on: rollback_to,
reason,
})
}
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use std::cell::RefCell;
use std::rc::Rc;
use super::*;
use crate::compat::RuntimeTarget;
#[derive(Clone)]
struct RecordingActivator {
flips: Rc<RefCell<Vec<ReleaseId>>>,
fail_on: Option<ReleaseId>,
}
impl RecordingActivator {
fn new() -> Self {
Self {
flips: Rc::new(RefCell::new(Vec::new())),
fail_on: None,
}
}
fn failing_on(release: ReleaseId) -> Self {
Self {
flips: Rc::new(RefCell::new(Vec::new())),
fail_on: Some(release),
}
}
fn flips(&self) -> Vec<ReleaseId> {
self.flips.borrow().clone()
}
}
impl Activator for RecordingActivator {
fn activate(&self, release: &ReleaseId) -> Result<(), StageError> {
if self.fail_on.as_ref() == Some(release) {
return Err(StageError::Activate(format!("cannot flip to {release}")));
}
self.flips.borrow_mut().push(release.clone());
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.dev/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 verified_healthy_update_commits_and_keeps_previous_staged() {
let activator = RecordingActivator::new();
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
activator.clone(),
|| Health::Healthy,
ReleaseId::new("v1"),
);
let outcome = stager.stage_and_apply(&ReleaseId::new("v2")).unwrap();
assert_eq!(
outcome,
Outcome::Committed {
version: ReleaseId::new("v2")
}
);
assert_eq!(stager.live(), &ReleaseId::new("v2"));
assert_eq!(stager.previous(), Some(&ReleaseId::new("v1")));
assert_eq!(activator.flips(), vec![ReleaseId::new("v2")]);
}
#[test]
fn unverified_bundle_is_never_applied() {
let activator = RecordingActivator::new();
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| false,
activator.clone(),
|| panic!("health check must not run for an unverified bundle"),
ReleaseId::new("v1"),
);
let err = stager.stage_and_apply(&ReleaseId::new("v2")).unwrap_err();
assert!(matches!(err, StageError::Unverified));
assert!(activator.flips().is_empty());
assert_eq!(stager.live(), &ReleaseId::new("v1"));
assert_eq!(stager.previous(), None);
}
#[test]
fn failed_health_check_auto_rolls_back_to_previous() {
let activator = RecordingActivator::new();
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
activator.clone(),
|| Health::Unhealthy("readiness probe timed out".to_owned()),
ReleaseId::new("v1"),
);
let outcome = stager.stage_and_apply(&ReleaseId::new("v2")).unwrap();
assert_eq!(
outcome,
Outcome::RolledBack {
stayed_on: ReleaseId::new("v1"),
reason: "readiness probe timed out".to_owned(),
}
);
assert_eq!(stager.live(), &ReleaseId::new("v1"));
assert_eq!(stager.previous(), None);
assert_eq!(
activator.flips(),
vec![ReleaseId::new("v2"), ReleaseId::new("v1")]
);
}
#[test]
fn rollback_outcome_surfaces_an_honest_summary() {
let outcome = Outcome::RolledBack {
stayed_on: ReleaseId::new("v1"),
reason: "readiness probe timed out".to_owned(),
};
let summary = outcome.summary();
assert_eq!(
summary,
"Stayed on v1 — the new version failed its health check (readiness probe timed out).",
);
for banned in ["sorry", "please", "unfortunately"] {
assert!(
!summary.to_lowercase().contains(banned),
"rollback summary must not contain {banned:?}",
);
}
}
#[test]
fn committed_outcome_summary_names_the_new_version() {
let summary = Outcome::Committed {
version: ReleaseId::new("v2"),
}
.summary();
assert_eq!(summary, "Now running v2.");
}
#[test]
fn download_failure_applies_nothing() {
let activator = RecordingActivator::new();
let mut stager = Stager::new(
|_: &ReleaseId| Err(StageError::Download("side location is full".to_owned())),
|_: &StagedArtifact| panic!("verify must not run when the download fails"),
activator.clone(),
|| panic!("health check must not run when the download fails"),
ReleaseId::new("v1"),
);
let err = stager.stage_and_apply(&ReleaseId::new("v2")).unwrap_err();
assert!(matches!(err, StageError::Download(_)));
assert!(activator.flips().is_empty());
assert_eq!(stager.live(), &ReleaseId::new("v1"));
}
#[test]
fn apply_flip_failure_leaves_live_untouched() {
let activator = RecordingActivator::failing_on(ReleaseId::new("v2"));
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
activator.clone(),
|| Health::Healthy,
ReleaseId::new("v1"),
);
let err = stager.stage_and_apply(&ReleaseId::new("v2")).unwrap_err();
assert!(matches!(err, StageError::Activate(_)));
assert!(activator.flips().is_empty());
assert_eq!(stager.live(), &ReleaseId::new("v1"));
assert_eq!(stager.previous(), None);
}
#[test]
fn a_second_update_restages_the_prior_live_as_previous() {
let activator = RecordingActivator::new();
let mut stager = Stager::new(
ok_source,
|_: &StagedArtifact| true,
activator.clone(),
|| Health::Healthy,
ReleaseId::new("v1"),
);
stager.stage_and_apply(&ReleaseId::new("v2")).unwrap();
stager.stage_and_apply(&ReleaseId::new("v3")).unwrap();
assert_eq!(stager.live(), &ReleaseId::new("v3"));
assert_eq!(stager.previous(), Some(&ReleaseId::new("v2")));
assert_eq!(
activator.flips(),
vec![ReleaseId::new("v2"), ReleaseId::new("v3")]
);
}
#[test]
fn staged_artifact_carries_its_classified_bundle() {
let artifact = artifact_for(&ReleaseId::new("v2"));
assert!(
artifact
.staged_path
.starts_with("/var/lib/polychrome/staged")
);
assert_eq!(artifact.bundle.catalog_hash, "catalog-v2");
}
}