continuum-proxy/src/completion.rs

//! Noticing that a print has finished, and saying so at the machine.
//!
//! Its own task, polling the fleet, rather than a branch inside the uplink's
//! telemetry loop. The uplink's loop only ticks while the cloud connection is
//! up, and a light that only comes on when the internet is working is not
//! much of a notification: the person this is for is standing next to the
//! printer, and the printer is on the same LAN as this process. Nothing about
//! telling them should route through a datacentre.

use std::collections::HashMap;
use std::time::Duration;

use tracing::{info, warn};

use crate::printer::{GenericPrinter, LightState, PrinterHandle, CHAMBER_LIGHT};
use crate::uplink::SharedFleet;

/// How long the lights blink for when a print finishes.
const BLINK_FOR: Duration = Duration::from_secs(10);

/// On for this long, then off for this long. 500ms reads as a deliberate
/// blink from across a room; much faster looks like a fault.
const BLINK_HALF_CYCLE: Duration = Duration::from_millis(500);

/// How often to look for a print that has just finished.
///
/// A notification a few seconds late is fine. What matters more is that this
/// is frequent enough not to miss a state: `FINISH` persists until the next
/// print starts, so there is no race to lose - but the first poll that sees
/// it is the one that fires.
const POLL_INTERVAL: Duration = Duration::from_secs(5);

/// Bambu's word for a print that ran to completion.
///
/// Deliberately only this one. `FAILED` is what a Bambu reports for a print
/// that was cancelled or went wrong - including every `print.stop` this
/// software sends - and celebrating a cancellation would be worse than
/// staying dark.
const FINISHED: &str = "FINISH";

/// Remembers each printer's last known state, so a finish can be told from a
/// printer that has simply been sitting in `FINISH` for an hour.
#[derive(Default)]
pub struct CompletionWatcher {
    last_state: HashMap<String, String>,
}

impl CompletionWatcher {
    pub fn new() -> Self {
        Self::default()
    }

    /// The printers that have *just* finished a print.
    ///
    /// Two rules, both of which exist to avoid a false celebration:
    ///
    /// - It must be a transition. `FINISH` persists until the next print
    ///   starts, so without this every poll for the next hour would fire.
    /// - A printer seen for the first time never fires, even if it is
    ///   already `FINISH`. Otherwise restarting the gateway blinks every
    ///   printer that happens to be sitting in `FINISH`, announcing prints
    ///   that finished yesterday.
    pub fn newly_finished<I, S>(&mut self, states: I) -> Vec<String>
    where
        I: IntoIterator<Item = (String, S)>,
        S: AsRef<str>,
    {
        let mut finished = Vec::new();

        for (printer_id, state) in states {
            let state = state.as_ref();
            let previous = self.last_state.insert(printer_id.clone(), state.to_string());

            if state == FINISHED && previous.is_some_and(|prev| prev != FINISHED) {
                finished.push(printer_id);
            }
        }

        finished
    }
}

/// Watches the fleet and blinks a printer's lights when it finishes a print.
pub async fn run(fleet: SharedFleet) {
    let mut watcher = CompletionWatcher::new();
    let mut ticker = tokio::time::interval(POLL_INTERVAL);

    loop {
        ticker.tick().await;

        let states = {
            let fleet = fleet.lock().await;
            // `print_status` is a cached read, but `refresh_status` below
            // publishes, so this does hold the lock across one MQTT publish
            // per printer. Same shape as the uplink's own `refresh_fleet`,
            // and at a handful of printers it is not worth more machinery;
            // the blink itself is what had to stay out of the lock, because
            // that one takes ten seconds.
            let mut states = Vec::new();
            for printer in fleet.iter() {
                // v1 printers only push a status when asked. Doing it here as
                // well as in the uplink means this keeps working with the
                // cloud connection down, which is the whole point of the
                // separate task.
                if let Err(err) = printer.refresh_status().await {
                    warn!(printer = %printer.base().id, %err, "could not request a status refresh");
                }
                if let Some(state) = printer.status().await.and_then(|s| s.state) {
                    states.push((printer.base().id.clone(), state));
                }
            }
            states
        };

        for printer_id in watcher.newly_finished(states) {
            info!(printer = %printer_id, "print finished; blinking the lights");
            // Spawned, not awaited: this takes ten seconds, and nothing else
            // in the fleet should wait for it.
            tokio::spawn(celebrate(fleet.clone(), printer_id));
        }
    }
}

/// Blinks one printer's chamber light for `BLINK_FOR`, then leaves it off.
///
/// The light is left off whatever happens, including when a blink fails
/// part-way: a printer abandoned mid-blink with its light on looks like a
/// machine still working.
/// Public so `examples/bambu_print_probe.rs` can run the real blink against a
/// real printer. Twenty light commands in ten seconds is a rate a printer can
/// refuse, and a blink reimplemented in the probe to match would not have
/// told us whether this one works.
pub async fn celebrate(fleet: SharedFleet, printer_id: String) {
    let deadline = tokio::time::Instant::now() + BLINK_FOR;
    let mut state = LightState::On;

    while tokio::time::Instant::now() < deadline {
        if let Err(err) = set_light(&fleet, &printer_id, state).await {
            // Stop at the first failure rather than logging twenty times,
            // but still fall through to the final Off below.
            warn!(printer = %printer_id, %err, "stopped blinking after a failed light command");
            break;
        }

        state = match state {
            LightState::On => LightState::Off,
            LightState::Off => LightState::On,
        };
        tokio::time::sleep(BLINK_HALF_CYCLE).await;
    }

    if let Err(err) = set_light(&fleet, &printer_id, LightState::Off).await {
        warn!(printer = %printer_id, %err, "could not turn the light back off");
    }
}

/// One light command, holding the fleet lock only for as long as it takes.
///
/// Locking per command rather than once around the whole blink: the fleet
/// mutex is also what telemetry and cloud command dispatch need, and holding
/// it for ten seconds would stall both for every printer on the farm.
async fn set_light(
    fleet: &SharedFleet,
    printer_id: &str,
    state: LightState,
) -> Result<(), String> {
    let fleet = fleet.lock().await;
    let printer = fleet
        .iter()
        .find(|p| p.base().id == printer_id)
        .ok_or_else(|| format!("no printer with id '{printer_id}' in this gateway's fleet"))?;

    // Bambu-only for now. A Prusa has no addressable light and Klipper's
    // depends entirely on the user's own config, so there is nothing
    // vendor-neutral to put on `GenericPrinter` yet - and inventing a
    // trait method that three of four implementations refuse is worse than
    // matching here.
    match printer {
        PrinterHandle::Bambu(bambu) => {
            bambu.set_light(CHAMBER_LIGHT, state).await.map_err(|err| err.to_string())
        }
        other => Err(format!(
            "{} has no light this software knows how to control",
            other.display_name()
        )),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn states(pairs: &[(&str, &str)]) -> Vec<(String, String)> {
        pairs.iter().map(|(id, s)| (id.to_string(), s.to_string())).collect()
    }

    /// The regression that matters most: a gateway restarting while a
    /// printer sits in FINISH must not announce a print that finished
    /// yesterday.
    #[test]
    fn a_printer_first_seen_already_finished_does_not_fire() {
        let mut watcher = CompletionWatcher::new();
        assert!(watcher.newly_finished(states(&[("p1p", "FINISH")])).is_empty());
    }

    #[test]
    fn a_transition_into_finish_fires_once_and_only_once() {
        let mut watcher = CompletionWatcher::new();
        watcher.newly_finished(states(&[("p1p", "RUNNING")]));

        assert_eq!(watcher.newly_finished(states(&[("p1p", "FINISH")])), vec!["p1p"]);
        // FINISH persists until the next print starts; without the
        // transition check this would fire every poll for an hour.
        assert!(watcher.newly_finished(states(&[("p1p", "FINISH")])).is_empty());
    }

    /// A Bambu reports an operator cancellation as FAILED - including every
    /// `print.stop` this software sends. Celebrating that would be worse
    /// than staying dark.
    #[test]
    fn a_cancelled_print_never_fires() {
        let mut watcher = CompletionWatcher::new();
        watcher.newly_finished(states(&[("p1p", "RUNNING")]));
        assert!(watcher.newly_finished(states(&[("p1p", "FAILED")])).is_empty());
    }

    #[test]
    fn a_second_print_fires_again() {
        let mut watcher = CompletionWatcher::new();
        watcher.newly_finished(states(&[("p1p", "RUNNING")]));
        watcher.newly_finished(states(&[("p1p", "FINISH")]));
        watcher.newly_finished(states(&[("p1p", "RUNNING")]));
        assert_eq!(watcher.newly_finished(states(&[("p1p", "FINISH")])), vec!["p1p"]);
    }

    #[test]
    fn printers_are_tracked_independently() {
        let mut watcher = CompletionWatcher::new();
        watcher.newly_finished(states(&[("p1p", "RUNNING"), ("h2c", "RUNNING")]));

        let finished = watcher.newly_finished(states(&[("p1p", "FINISH"), ("h2c", "RUNNING")]));
        assert_eq!(finished, vec!["p1p"]);
    }

    /// Ten seconds of 500ms half-cycles is 20 commands, ending on Off. Worth
    /// pinning: the ask was "blink for ten seconds then turn the lights off",
    /// and a half-cycle that does not divide the duration would end mid-blink.
    #[test]
    fn the_blink_divides_evenly_into_the_duration() {
        assert_eq!(BLINK_FOR.as_millis() % BLINK_HALF_CYCLE.as_millis(), 0);
        assert_eq!(BLINK_FOR.as_millis() / BLINK_HALF_CYCLE.as_millis(), 20);
    }
}