shep 0.19.0

//! The `start` and `add` verb entry points.
//!
//! The per-target load pipeline: resolve a target, register what the flock
//! does not have, merge declared config into what it does, then respawn.

use std::collections::{BTreeMap, BTreeSet};
use std::io::Read as _;
use std::path::Path;

use shep_client::{Client, START_DEADLINE};
use shep_core::config::{AppConfig, DeclaredApp, ResetDepth};
use shep_core::protocol::{EnvValue, ProcessInfo, Request, Response, SelectorSpec};
use shep_core::selector::ProcessSelector;

use crate::cli::StartArgs;
use crate::commands::lifecycle::configure::{
    apply_declared, apply_interpreters, first_failure, reset_depth, reset_flag,
};
use crate::commands::lifecycle::resolve::{
    TargetError, resolve_target_declared, split_assignments,
};
use crate::commands::lifecycle::respawn::{already_registered, leave_disabled, resume_all};
use crate::commands::lifecycle::selectors::{
    fail_target, flock_matches, flock_now, is_reachable_as_a_name, render_outcome, request_each,
    selector_miss,
};
use crate::commands::lifecycle::staged_start_deadline;
use crate::commands::selector::parse_selector;
use crate::exit::ExitCode;
use crate::output::{FlockRows, Streams};
use crate::remote::{remote_word, ssh_hint};

/// Which of the two verbs that read a Flockfile is running
///
/// They share targets, resolution, merge and refusals; the difference is
/// whether anything is spawned at the end. One code path, so a document
/// cannot register differently depending on which verb read it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Load {
    /// `shep start`: register what the flock does not have, and bring up what
    /// it does.
    Start,
    /// `shep add`: register what the flock does not have, and stop there.
    Add,
}

impl Load {
    /// The verb's own name, for a notice's code and for the `--format json`
    /// envelope's `command`.
    pub(crate) fn verb(self) -> &'static str {
        match self {
            Self::Start => "start",
            Self::Add => "add",
        }
    }
}

/// Registers what `args` resolves to and brings all of it up.
pub async fn start(
    client: &Client,
    streams: &mut Streams<'_>,
    args: &StartArgs,
    discovered: Option<&Path>,
    interpreters: &BTreeMap<String, String>,
) -> ExitCode {
    load(
        client,
        streams,
        args,
        discovered,
        interpreters,
        Load::Start,
        None,
    )
    .await
}

/// [`start`] against the shepherd on `dest`
///
/// Only a sheep or fold the host's flock has can mean the same thing there. A
/// path, `-`, `--flockfile` and a bare `start` would be read on this machine,
/// so each is refused with the `ssh` command that reads it on the host.
pub async fn start_on_host(
    client: &Client,
    streams: &mut Streams<'_>,
    args: &StartArgs,
    dest: &str,
) -> ExitCode {
    load(
        client,
        streams,
        args,
        None,
        &BTreeMap::new(),
        Load::Start,
        Some(dest),
    )
    .await
}

/// Whether `token` would be read on this machine rather than looked up in a
/// flock: `-`, anything under `--flockfile`, or a word no sheep can be named
///
/// Needs no shepherd, so a host link answers it before dialling. A token that
/// is not a selector at all is left to `load_one`, which reports it.
pub(crate) fn reads_this_machine(args: &StartArgs, token: &str) -> bool {
    token == "-"
        || args.flockfile
        || ProcessSelector::parse(token).is_ok_and(|selector| !is_reachable_as_a_name(&selector))
}

/// What the operator typed and the `ssh` line that does the same on `dest`,
/// each carrying every flag, assignment and target of `args`.
fn start_commands(dest: &str, args: &StartArgs) -> (String, String) {
    let mut words = vec!["start".to_string()];
    let valued = [
        ("--name", &args.name),
        ("--fold", &args.fold),
        ("--cwd", &args.cwd),
        ("--interpreter", &args.interpreter),
    ];
    for (flag, value) in valued {
        if let Some(value) = value {
            words.push(flag.to_string());
            words.push(value.clone());
        }
    }
    if args.flockfile {
        words.push("--flockfile".to_string());
    }
    if args.no_log_timestamps {
        words.push("--no-log-timestamps".to_string());
    }
    words.extend(reset_flag(args));
    words.extend(args.targets.iter().cloned());
    let typed: Vec<String> = words.iter().map(|word| remote_word(word)).collect();
    let mut remote = vec!["shep".to_string()];
    remote.extend(words);
    (
        format!("shep --host {dest} {}", typed.join(" ")),
        ssh_hint(dest, &remote),
    )
}

/// The refusal for a target that would be read on this machine
///
/// `token` is `None` for a bare `start`, which reads the Flockfile in this
/// directory.
///
pub(crate) fn refuse_local_read(
    streams: &mut Streams<'_>,
    dest: &str,
    args: &StartArgs,
    token: Option<&str>,
) -> ExitCode {
    let read = match token {
        Some("-") => "stdin".to_string(),
        Some(token) => token.to_string(),
        None => "the Flockfile in this directory".to_string(),
    };
    let (typed, ssh) = start_commands(dest, args);
    let message =
        format!("`{typed}` would read {read} on this machine. To start it there, run: {ssh}");
    streams.fail(ExitCode::Usage, &message)
}

/// Registers what `args` resolves to and starts none of it
///
/// The app lands registered and stopped, so a template shipping `env = {
/// DB_HOST = "" }` can be filled in before it spawns.
///
/// An app the flock already has is merged into and left as it is, running
/// or not, so re-running this cannot stop a service.
pub async fn add(
    client: &Client,
    streams: &mut Streams<'_>,
    args: &StartArgs,
    discovered: Option<&Path>,
    interpreters: &BTreeMap<String, String>,
) -> ExitCode {
    load(
        client,
        streams,
        args,
        discovered,
        interpreters,
        Load::Add,
        None,
    )
    .await
}

/// The body [`start`] and [`add`] share; see [`Load`] for what they do not
async fn load(
    client: &Client,
    streams: &mut Streams<'_>,
    args: &StartArgs,
    discovered: Option<&Path>,
    interpreters: &BTreeMap<String, String>,
    mode: Load,
    host: Option<&str>,
) -> ExitCode {
    let (assignments, targets) = split_assignments(&args.targets);
    // An environment belongs to one sheep for the reason a name does, and
    // the shell has no analogue to borrow: one assignment prefix there
    // introduces one command.
    if !assignments.is_empty() && targets.len() > 1 {
        let message = "an assignment takes one target: an environment belongs to one sheep";
        return streams.fail(ExitCode::Usage, message);
    }
    // A word holding whitespace is one a shell never hands over: the
    // operator quoted the assignment and its target together, which is the
    // one spelling that would need shep to split and requote for itself.
    if !assignments.is_empty() && targets.is_empty() {
        let quoted = args
            .targets
            .iter()
            .find(|word| word.split_whitespace().count() > 1);
        let message = match quoted {
            Some(word) => format!(
                "an assignment needs a target; `{word}` is one word, so drop the quotes and \
                 let the shell split it"
            ),
            None => {
                "an assignment needs a target: a script path, or a sheep the flock has".to_string()
            }
        };
        return streams.fail(ExitCode::Usage, &message);
    }
    // `--name` renames the sheep a target becomes, and a name is unique to
    // one sheep, so it cannot mean anything across several targets.
    if args.name.is_some() && targets.len() > 1 {
        let message = "--name takes one target: a name belongs to one sheep";
        return streams.fail(ExitCode::Usage, message);
    }
    if targets.is_empty() {
        let mut started = Vec::new();
        let code = load_one(
            client,
            streams,
            args,
            None,
            discovered,
            interpreters,
            &mut started,
            mode,
            &assignments,
            host,
        )
        .await;
        // Printed whenever the verb succeeded, not only when it touched a
        // row. A failing verb leaves stdout empty, crate-wide.
        if code == ExitCode::Success {
            let wrote = render_outcome(client, streams, mode.verb(), FlockRows(started)).await;
            if wrote != ExitCode::Success {
                return wrote;
            }
        }
        return code;
    }
    // In turn, not atomically: if the second target fails the first is
    // already up. The exit code is the first failure.
    let mut failure: Option<ExitCode> = None;
    let mut started = Vec::new();
    for target in targets {
        let code = load_one(
            client,
            streams,
            args,
            Some(target),
            discovered,
            interpreters,
            &mut started,
            mode,
            &assignments,
            host,
        )
        .await;
        if code != ExitCode::Success {
            failure = failure.or(Some(code));
        }
    }
    // One table for the whole invocation, keyed on the outcome alone and
    // never on `started` being non-empty: every row is attempted, so a
    // partly-failed fold ends non-empty, and a table there would sit beside
    // an error envelope under `--format json`.
    if failure.is_none() {
        let wrote = render_outcome(client, streams, mode.verb(), FlockRows(started)).await;
        if wrote != ExitCode::Success {
            return wrote;
        }
    }
    failure.unwrap_or(ExitCode::Success)
}

/// Sets each assignment on `name` as an operator override
///
/// One request per key, since [`Request::SetSheepEnv`] carries one. The
/// first failure stops the run and comes back unphrased: a value the sheep
/// is already running with is a notice, and one it has never seen is a
/// refusal, and only the caller knows which it has.
///
/// # Errors
/// The first key the shepherd did not record, and why.
async fn set_assignments(
    client: &Client,
    name: &str,
    assignments: &BTreeMap<String, String>,
) -> Result<(), String> {
    for (key, value) in assignments {
        let body = Request::SetSheepEnv {
            name: name.to_string(),
            key: key.clone(),
            value: Some(EnvValue::from(value.clone())),
        };
        match client.request(body).await {
            Ok(Response::SheepEnvSet { .. }) => {}
            Ok(_unrecognised) => {
                return Err(format!(
                    "{name}: the shepherd answered {key} with something this client does not \
                     understand"
                ));
            }
            Err(err) => return Err(format!("{name}: {key} was not recorded ({err})")),
        }
    }
    Ok(())
}

/// Writes `--no-log-timestamps` onto each sheep the flock already has
///
/// The flag is a value the operator typed, so it travels as an operator
/// override, not as a declared key: `merge_declared` at
/// [`ResetDepth::None`] takes a declared key only when nobody has established it.
///
/// It parks: the shepherd's `pending` answer is reported.
async fn set_log_timestamps_off(
    client: &Client,
    streams: &mut Streams<'_>,
    mode: Load,
    names: &[&str],
) -> ExitCode {
    for name in names {
        let body = Request::SetSheepField {
            name: name.to_string(),
            key: "log_timestamps".to_string(),
            value: serde_json::Value::Bool(false),
        };
        match client.request(body).await {
            // A write that landed and a write that landed parked are not the
            // same run, and only the shepherd knows which it answered. The
            // clause is the one `apply_declared` prints for a parked field.
            Ok(Response::SheepFieldSet {
                pending, warning, ..
            }) => {
                if let Some(warning) = warning {
                    streams.aside("warning", &warning);
                }
                if pending {
                    streams.aside(
                        mode.verb(),
                        &format!(
                            "{name}: log_timestamps waiting on the next spawn \
                             (`shep reload {name}` promotes them)"
                        ),
                    );
                }
            }
            // `Response` is `#[non_exhaustive]`, so an answer this client
            // does not recognise is a daemon-side fault rather than a bad
            // value.
            Ok(_other) => {
                let message = format!(
                    "{name}: the shepherd answered log_timestamps with a response this client \
                     does not understand, so the stamp stays on"
                );
                return streams.fail(ExitCode::Internal, &message);
            }
            // A refusal keeps the code its own `RpcErrorCode` maps to, the
            // way a refused field in an `ApplyConfig` report does.
            Err(err) => {
                let message = format!(
                    "{name}: log_timestamps could not be turned off ({err}), so the stamp stays \
                     on"
                );
                return streams.fail(ExitCode::from(&err), &message);
            }
        }
    }
    ExitCode::Success
}

/// The miss for a name the host's flock does not have
///
/// Hints at the `ssh` command too, since the token may have been a script
/// that lives on the host.
fn unknown_on_host(
    streams: &mut Streams<'_>,
    dest: &str,
    args: &StartArgs,
    token: &str,
    selector: &ProcessSelector,
    flock: &[ProcessInfo],
) -> ExitCode {
    let miss = selector_miss(token, selector, flock)
        .unwrap_or_else(|| format!("{token} is not a sheep or a fold"));
    let (_, ssh) = start_commands(dest, args);
    let message = format!("{miss} on {dest}. If {token} is a script there, run: {ssh}");
    streams.fail(ExitCode::NotFound, &message)
}

/// One target's worth of [`load`]
///
/// `interpreters` is read once by the caller rather than per target. `host` is
/// the destination whose shepherd `client` speaks to, and with one set the
/// function returns before anything is read on this machine.
#[allow(clippy::too_many_arguments)]
async fn load_one(
    client: &Client,
    streams: &mut Streams<'_>,
    args: &StartArgs,
    target: Option<&str>,
    discovered: Option<&Path>,
    interpreters: &BTreeMap<String, String>,
    started: &mut Vec<shep_core::protocol::ProcessInfo>,
    mode: Load,
    assignments: &BTreeMap<String, String>,
    host: Option<&str>,
) -> ExitCode {
    if let Some(dest) = host
        && (target.is_none() || target.is_some_and(|token| reads_this_machine(args, token)))
    {
        return refuse_local_read(streams, dest, args, target);
    }
    // Everything from here to `resolve_target` is the precedence in
    // `StartArgs::targets`' own help: a sheep by id or name, then a fold,
    // then a Flockfile, then a path. Each tier claims the token only if it
    // resolves there. `-` and `--flockfile` skip the flock entirely.
    let mut listing: Option<Vec<ProcessInfo>> = None;
    let mut missed: Option<String> = None;
    if let Some(token) = target
        && token != "-"
        && !args.flockfile
    {
        // Parsed client-side, as every selector-taking verb does, so a
        // malformed one is a local usage error rather than a round trip.
        let selector = match parse_selector(streams, token) {
            Ok(selector) => selector,
            Err(code) => return code,
        };
        if is_reachable_as_a_name(&selector) {
            // Over a host a miss is final, so a listing that failed must not
            // pass for an empty flock.
            let flock = if host.is_some() {
                match client.list_flock().await {
                    Ok(flock) => flock,
                    Err(err) => return streams.fail(ExitCode::from(&err), &err.to_string()),
                }
            } else {
                flock_now(client).await
            };
            let matched = flock_matches(&selector, &flock);
            if !matched.is_empty() {
                // Returns rather than falling through: a token that resolved
                // to a sheep reads no file, so `shep start web` cannot apply
                // whatever Flockfile sits in the operator's directory. A
                // reset flag is refused: there is no template to reset to.
                if let Some(flag) = reset_flag(args) {
                    let message = format!(
                        "{flag} needs a Flockfile to reset to; {token} names a sheep, not a file"
                    );
                    return streams.fail(ExitCode::Usage, &message);
                }
                // Instances of one name are one sheep; two names are two, and
                // an assignment names neither of them. Recorded before the
                // resume below, which is the spawn that promotes it.
                if !assignments.is_empty() {
                    let names: BTreeSet<&str> =
                        matched.iter().map(|info| info.name.as_str()).collect();
                    if names.len() > 1 {
                        let message = format!(
                            "an assignment takes a script path or one sheep; {token} names several"
                        );
                        return streams.fail(ExitCode::Usage, &message);
                    }
                    for name in &names {
                        if let Err(reason) = set_assignments(client, name, assignments).await {
                            return streams.fail(ExitCode::Failure, &reason);
                        }
                    }
                }
                // The flag reaches a sheep the token named the way it reaches
                // one a file resolved: written before the resume below, which
                // is the spawn that promotes a parked field. Names rather
                // than rows: instances of one name are one sheep.
                if args.no_log_timestamps {
                    let names: Vec<&str> = matched
                        .iter()
                        .map(|info| info.name.as_str())
                        .collect::<BTreeSet<&str>>()
                        .into_iter()
                        .collect();
                    let stamp = set_log_timestamps_off(client, streams, mode, &names).await;
                    if stamp != ExitCode::Success {
                        return stamp;
                    }
                }
                return match mode {
                    Load::Start => {
                        resume_all(client, streams, Some(token), &matched, started).await
                    }
                    // The sheep is registered, which is all `add` was asked
                    // for. Said anyway, so a printed table is explained.
                    Load::Add => already_registered(streams, &matched),
                };
            }
            if let Some(dest) = host {
                return unknown_on_host(streams, dest, args, token, &selector, &flock);
            }
            // Held for the failure path below rather than reported here: the
            // token may still name a Flockfile or a path.
            missed = selector_miss(token, &selector, &flock);
            listing = Some(flock);
        }
    }

    let discovered = discovered.map(|p| p.to_string_lossy().into_owned());
    let target: &str = match (target, discovered.as_deref()) {
        (Some(target), _) => target,
        (None, Some(found)) => found,
        (None, None) => {
            let message = "no target and no Flockfile in this directory";
            return streams.fail(ExitCode::Usage, message);
        }
    };

    let stdin = if target == "-" {
        let mut buf = Vec::new();
        if let Err(source) = std::io::stdin().lock().read_to_end(&mut buf) {
            return fail_target(streams, &TargetError::Stdin(source));
        }
        buf
    } else {
        Vec::new()
    };

    let mut apps =
        match resolve_target_declared(target, args.name.as_deref(), &stdin, args.flockfile) {
            Ok(apps) => apps,
            // The token named nothing anywhere. One written unmistakably as
            // a selector is reported as one, at exit 3, matching every other
            // verb's answer for a selector that matched nothing.
            Err(TargetError::Unresolvable { .. }) if missed.is_some() => {
                let message = missed.unwrap_or_default();
                return streams.fail(ExitCode::NotFound, &message);
            }
            Err(err) => return fail_target(streams, &err),
        };

    // The other target that supplies no template: a bare script path gets
    // an empty declared set, so a reset would exit 0 having reset nothing.
    // Tested on the declared set, not the token's shape, since that is what
    // a reset acts on. Every Flockfile app declares at least `script`.
    if let Some(flag) = reset_flag(args)
        && apps.iter().all(|app| app.declared.is_empty())
    {
        let message =
            format!("{flag} needs a Flockfile to reset to; {target} is a script, not a Flockfile");
        return streams.fail(ExitCode::Usage, &message);
    }

    // A file may declare several apps and an assignment names none of them,
    // so nothing says which one the operator meant. Tested on the declared
    // set for the reason the reset above is: it is what a file supplies.
    if !assignments.is_empty() && apps.iter().any(|app| !app.declared.is_empty()) {
        let message = format!(
            "an assignment takes a script path or one sheep; {target} is a Flockfile, which \
             may declare several"
        );
        return streams.fail(ExitCode::Usage, &message);
    }

    if let Some(fold) = &args.fold {
        for app in &mut apps {
            app.config.fold = Some(fold.clone());
        }
    }
    // After the per-app defaults above, so an explicit flag wins over both
    // the script form's default and a Flockfile's own value.
    if let Some(cwd) = &args.cwd {
        for app in &mut apps {
            app.config.cwd = Some(cwd.clone());
        }
    }
    if args.no_log_timestamps {
        for app in &mut apps {
            app.config.log_timestamps = false;
        }
    }
    // After the file's own values, since an assignment is something an
    // operator typed for this one sheep.
    for app in &mut apps {
        for (key, value) in assignments {
            app.config.env.insert(key.clone(), value.clone());
        }
    }
    apply_interpreters(&mut apps, interpreters, args.interpreter.as_deref());
    // The bare-name rule again, now the target is a set of named apps: a
    // name the flock already has is that sheep. Without it `shep start
    // ./thing` spawns a second copy of a one-instance app.
    let flock = match listing {
        Some(flock) => flock,
        None => flock_now(client).await,
    };
    // Every row the name has, not the first: respawns go out per row, so one
    // row standing in for a clustered app leaves the other instances down.
    let mut resumed: Vec<(DeclaredApp, Vec<ProcessInfo>)> = Vec::new();
    let mut fresh = Vec::new();
    for app in apps {
        let rows: Vec<ProcessInfo> = flock
            .iter()
            .filter(|info| info.name == app.config.name)
            .cloned()
            .collect();
        if rows.is_empty() {
            fresh.push(app);
        } else {
            resumed.push((app, rows));
        }
    }
    // Before the resumes below, since a `NeedsRespawn` field parks for the
    // next spawn and that is the resume immediately below. The code is
    // carried rather than returned: a refused field must not stop the flock
    // coming back up. The merge runs under `add` too, only the resume does not.
    let declared: Vec<DeclaredApp> = resumed.iter().map(|(app, _)| app.clone()).collect();
    let applied = apply_declared(client, streams, &declared, reset_depth(args), mode).await;
    // The flag is not in `declared`, so `apply_declared` cannot carry it.
    // It is written before the resume for the reason the resume follows the
    // merge: a parked field is promoted by the next spawn, and that spawn is
    // this one. `add` runs the write too; only the resume is `start`'s.
    if args.no_log_timestamps {
        let names: Vec<&str> = resumed
            .iter()
            .map(|(app, _)| app.config.name.as_str())
            .collect();
        let stamp = set_log_timestamps_off(client, streams, mode, &names).await;
        if stamp != ExitCode::Success {
            return first_failure(applied, stamp);
        }
    }
    if !resumed.is_empty() && mode == Load::Start {
        // `None`, not the operator's token: this arm reached the flock
        // through a Flockfile or a path, so there is no selector to quote.
        let existing: Vec<ProcessInfo> = resumed
            .iter()
            .flat_map(|(_, rows)| rows.iter().cloned())
            .collect();
        let existing = leave_disabled(streams, existing);
        let code = resume_all(client, streams, None, &existing, started).await;
        if code != ExitCode::Success {
            return first_failure(applied, code);
        }
    }
    if fresh.is_empty() {
        return applied;
    }
    let apps: Vec<AppConfig> = fresh.iter().map(|app| app.config.clone()).collect();

    let (procs, failure) = request_each(
        client,
        streams, // Both requests carry the apps, so this "selector" is a
        // placeholder the body closure ignores.
        &[SelectorSpec::All],
        Some(match mode {
            // `add` registers and starts nothing, so it runs no stages and
            // needs no more than the batch budget. `START_DEADLINE`: a
            // single-threaded actor behind a batch of cold spawns outruns the
            // client's default 5s on its own.
            Load::Add => START_DEADLINE,
            Load::Start => staged_start_deadline(&apps),
        }),
        |_| match mode {
            Load::Start => Request::Start { apps: apps.clone() },
            Load::Add => Request::Add { apps: apps.clone() },
        },
        |response| match response {
            Response::Started(procs) | Response::Added(procs) => Some(procs),
            _ => None,
        },
    )
    .await;
    // Establishes what a fresh app's file declared: `Start` and `Add` write
    // nothing to the override store. `ResetDepth::None` whatever flag was
    // passed, since a reset would drop the record this call writes. Only
    // apps the daemon registered: a failed spawn is not in `procs`.
    let registered: BTreeSet<&str> = procs.iter().map(|info| info.name.as_str()).collect();
    let established: Vec<DeclaredApp> = fresh
        .into_iter()
        .filter(|app| registered.contains(app.config.name.as_str()))
        .collect();
    let recorded = apply_declared(client, streams, &established, ResetDepth::None, mode).await;
    // `Start` and `Add` write nothing to the override store, so a value only
    // the config carries is one the next Flockfile load replaces. The sheep
    // is up with it either way, which is what makes this a notice.
    for name in &registered {
        if let Err(reason) = set_assignments(client, name, assignments).await {
            let message = format!(
                "{reason}; it is set for this spawn, but a Flockfile load could replace it"
            );
            streams.aside(mode.verb(), &message);
        }
    }
    started.extend(procs);
    first_failure(
        applied,
        first_failure(failure.unwrap_or(ExitCode::Success), recorded),
    )
}