sausage/cli/src/upgrades.rs

//! Grace periods, service restarts, and retirement of replaced versions (PROTOCOL.md section 3.7).
use crate::{Host, services::Service};
use serde_json::json;
use std::{collections::BTreeMap, time::Duration};
use tokio::sync::watch;

pub(crate) struct OldProcess {
    pid: String,
    version: String,
    ended: watch::Receiver<bool>,
}
pub(crate) struct Upgrade {
    plugin: String,
    previous: String,
    grace: Duration,
    processes: Vec<OldProcess>,
    services: Vec<(Service, Option<watch::Receiver<bool>>)>,
}
impl Host {
    /// Snapshots old admissions while activation holds the admission write guard.
    pub(crate) fn snapshot_upgrade(
        &self,
        plugin: &str,
        previous: String,
        current: &str,
        grace: f64,
    ) -> Upgrade {
        let mut state = self.0.state.lock().unwrap();
        let processes = state
            .processes
            .iter()
            .filter(|(_, process)| process.info.plugin == plugin && process.info.version != current)
            .map(|(pid, process)| OldProcess {
                pid: pid.clone(),
                version: process.info.version.clone(),
                ended: process.ended.subscribe(),
            })
            .collect();
        let mut services = Vec::new();
        let mut retained = Vec::new();
        for service in std::mem::take(&mut state.services) {
            if service
                .method
                .split_once('.')
                .is_some_and(|(name, _)| name == plugin)
            {
                let ended = state
                    .processes
                    .values()
                    .find(|process| process.id.as_ref() == Some(&service.request))
                    .map(|process| process.ended.subscribe());
                services.push((service, ended));
            } else {
                retained.push(service);
            }
        }
        state.services = retained;
        Upgrade {
            plugin: plugin.into(),
            previous,
            grace: Duration::from_secs_f64(grace),
            processes,
            services,
        }
    }
    /// Lets old calls finish, stops them after grace, and starts each lifecycle hook once.
    pub(crate) fn begin_upgrade(&self, upgrade: Upgrade) {
        let host = self.clone();
        tokio::spawn(async move {
            host.0.journal.write(
                "core",
                &upgrade.plugin,
                "core",
                &format!("upgrade grace {} seconds", upgrade.grace.as_secs_f64()),
            );
            let mut versions: BTreeMap<String, Vec<watch::Receiver<bool>>> = BTreeMap::new();
            versions.entry(upgrade.previous).or_default();
            for process in upgrade.processes {
                versions
                    .entry(process.version)
                    .or_default()
                    .push(process.ended.clone());
                let host = host.clone();
                let grace = upgrade.grace;
                tokio::spawn(async move {
                    let mut ended = process.ended;
                    tokio::select! {
                        _ = ended.wait_for(|ended|*ended) => (),
                        _ = tokio::time::sleep(grace) => { host.begin_stop(&process.pid); }
                    }
                });
            }
            for (service, ended) in upgrade.services {
                let host = host.clone();
                tokio::spawn(async move {
                    if let Some(mut ended) = ended {
                        let _ = ended.wait_for(|ended| *ended).await;
                    }
                    host.start_service(service.method, service.params).await;
                });
            }
            for (version, processes) in versions {
                let host = host.clone();
                let plugin = upgrade.plugin.clone();
                let retirement = (plugin.clone(), version.clone());
                if !host
                    .0
                    .retirements
                    .lock()
                    .unwrap()
                    .insert(retirement.clone())
                {
                    continue;
                }
                tokio::spawn(async move {
                    for mut ended in processes {
                        let _ = ended.wait_for(|ended| *ended).await;
                    }
                    if host
                        .0
                        .graph
                        .lock()
                        .unwrap()
                        .plugins
                        .get(&plugin)
                        .is_some_and(|current| current.hash == version)
                    {
                        host.0.retirements.lock().unwrap().remove(&retirement);
                        return;
                    }
                    host.0.journal.write(
                        "core",
                        &plugin,
                        "core",
                        &format!("retired start {version}"),
                    );
                    let result = host
                        .route(
                            "core".into(),
                            Some(json!(uuid::Uuid::now_v7().to_string())),
                            format!("{plugin}.rpc.retired"),
                            json!({"previous":version}),
                            true,
                        )
                        .await;
                    host.0.journal.write(
                        "core",
                        &plugin,
                        "core",
                        &format!("retired result {result:?}"),
                    );
                    host.0.retirements.lock().unwrap().remove(&retirement);
                });
            }
        });
    }
}