filespooler 1.2.4

/*! Tools executing jobs */

/*
    Copyright (C) 2022 John Goerzen <jgoerzen@complete.org>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

use crate::cmd::cmd_queue::QueueOptsWithDecoder;
use crate::exec;
use crate::header::{CommandTypeV1, FSMetaV1};
use crate::jobfile;
use crate::jobqueue::*;
use crate::seqfile;
use crate::seqfile::SeqFile;
use crate::util::get_unbuffered_stdin;
use anyhow::{anyhow, bail};
use clap::Args;
use std::collections::BTreeMap;
use std::ffi::OsStr;
use std::ffi::OsString;
use std::fmt;
use std::fmt::Debug;
use std::fs::{remove_file, File};
use std::os::unix::ffi::OsStrExt;
use std::process::{ExitStatus, Stdio};
use std::str::FromStr;
use tracing::*;

/// Option for what to do if an error is encountered while processing a queue
#[derive(Debug, Clone)]
pub enum OnError {
    /// Delete the file and mark the job as done; do not return an error code.  If --never-delete
    /// is given, this becomes the same as Leave.
    Delete,

    /// Leave the file and mark the job as done; do not return an error code.
    Leave,

    /// Abort processing with an error code and do not mark the job as done.  It can be retried by a subsequent queue processing command.
    Retry,
}

// consider enum-utils fromstr derivation or others at
// https://crates.io/search?q=enum_utils
impl FromStr for OnError {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, anyhow::Error> {
        match s.trim().to_lowercase().as_str() {
            "delete" => Ok(OnError::Delete),
            "leave" => Ok(OnError::Leave),
            "retry" => Ok(OnError::Retry),
            _ => Err(anyhow!(
                "on-error must be 'delete', 'leave', or 'retry'"
            )),
        }
    }
}

impl fmt::Display for OnError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self, f)
    }
}

/// Option for what to do with output
#[derive(Debug, Clone)]
pub enum OutputTo {
    /// Don't capture; use the calling environment's stdout/stderr
    PassBoth,

    /// Write both stdout and stderr to a file
    SaveBoth,
}

// consider enum-utils fromstr derivation or others at
// https://crates.io/search?q=enum_utils
impl FromStr for OutputTo {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, anyhow::Error> {
        match s.trim().to_lowercase().as_str() {
            "passboth" => Ok(OutputTo::PassBoth),
            "saveboth" => Ok(OutputTo::SaveBoth),
            _ => Err(anyhow!("output-to must be 'passboth' or 'saveboth'")),
        }
    }
}

impl fmt::Display for OutputTo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self, f)
    }
}

#[derive(Debug, Args)]
pub struct ProcessOpts {
    /// Allow the parameters from the job to be appended to the command's
    /// command line (not doing so by default)
    #[arg(long)]
    pub allow_job_params: bool,

    /// Ignore the payload; do not pass it to the command.
    #[arg(long)]
    pub ignore_payload: bool,

    /// A timeout, in seconds, for the command.
    #[arg(long, value_name = "SECONDS")]
    pub timeout: Option<u64>,

    /// The command to run
    #[arg(value_name = "COMMAND")]
    pub command: OsString,

    /// Optional parameters to pass to COMMAND, before those contained
    /// in the request file.
    #[arg(last = true)]
    pub params: Vec<OsString>,
}

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum OrderBy {
    /// Order by sequence number (default)
    Sequence,

    /// Order by timestamp
    Timestamp,
}

// consider enum-utils fromstr derivation or others at
// https://crates.io/search?q=enum_utils
impl FromStr for OrderBy {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, anyhow::Error> {
        match s.trim().to_lowercase().as_str() {
            "sequence" => Ok(OrderBy::Sequence),
            "timestamp" => Ok(OrderBy::Timestamp),
            _ => Err(anyhow!("order-by must be 'sequence' or 'timestamp'")),
        }
    }
}

impl fmt::Display for OrderBy {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self, f)
    }
}
#[derive(Debug, Args)]
pub struct QueueProcessOpts {
    #[clap(flatten)]
    pub qdopts: QueueOptsWithDecoder,

    /// What to do with an error.  Can be Delete, Leave, or Retry.
    #[arg(long, value_name="ONERROR", default_value_t = OnError::Retry)]
    pub on_error: OnError,

    /// In what order to process jobs.  Can be Sequence or Timestamp.
    #[arg(long, value_name="ORDERING", default_value_t = OrderBy::Sequence)]
    pub order_by: OrderBy,

    /// Never delete processed job files in any circumstance
    #[arg(long)]
    pub never_delete: bool,

    /// Maximum number of jobs to process
    #[arg(long, value_name = "JOBS", short = 'n')]
    pub maxjobs: Option<u64>,

    /// What to do with stdout and stderr of the process.  Can be PassBoth or SaveBoth.
    #[arg(long, value_name="DEST", default_value_t = OutputTo::PassBoth)]
    pub output_to: OutputTo,

    #[clap(flatten)]
    pub processopts: ProcessOpts,
}

/// The result type from running a command.
#[derive(Debug, Eq, PartialEq)]
pub enum CmdResult {
    Success,
    Failure,
    ExitResult(ExitStatus),
}

impl CmdResult {
    fn success(&self) -> bool {
        match self {
            CmdResult::Success => true,
            CmdResult::Failure => false,
            CmdResult::ExitResult(es) => es.success(),
        }
    }
}

/// Process some input to execute, optionally writing the output to a given File.
/// If the input is None, use stdin.
pub fn cmd_exec_generic(
    opts: &ProcessOpts,
    mut input: Option<DecoderHandle>,
    inputfilename: Option<&OsString>,
    queuedir: Option<&OsString>,
    output: Option<File>,
) -> Result<CmdResult, anyhow::Error> {
    let mut stdindrop = None;
    let meta = match &mut input {
        Some(x) => x.with_read(|y| jobfile::read_jobfile_header(y))?,
        None => {
            // Standard stdin is buffered; bleh.  Have to deal with that here.
            let mut stdin = get_unbuffered_stdin();
            let ret = jobfile::read_jobfile_header(&mut stdin)?;
            stdindrop = Some(stdin);
            ret
        }
    };
    let input = if opts.ignore_payload {
        // Just to be really clear, we will not be processing input
        // in this case!
        std::mem::drop(input);
        None
    } else {
        input
    };

    let res = match meta.cmd_type {
        CommandTypeV1::NOP => Ok(CmdResult::Success),
        CommandTypeV1::Fail => Ok(CmdResult::Failure),
        CommandTypeV1::Command(ref copts) => {
            let mut params = opts.params.clone();
            if opts.allow_job_params {
                let jobparams: Vec<OsString> = copts
                    .params
                    .iter()
                    .map(|x| OsStr::from_bytes(x).to_os_string())
                    .collect();

                params.extend(jobparams);
            } else if !copts.params.is_empty() {
                warn!("Parameters given in job ignored because --allow-job-params not given");
            }

            debug!(
                "Preparing to execute job {} from {:?}",
                meta.seq, &inputfilename
            );
            let env = meta.get_envvars(inputfilename, queuedir);

            if let Some(outfile) = output {
                Ok(CmdResult::ExitResult(exec::exec_job(
                    &opts.command,
                    &params,
                    env,
                    input,
                    Stdio::from(outfile.try_clone()?),
                    Stdio::from(outfile),
                    opts.timeout,
                )?))
            } else {
                Ok(CmdResult::ExitResult(exec::exec_job(
                    &opts.command,
                    &params,
                    env,
                    input,
                    Stdio::inherit(),
                    Stdio::inherit(),
                    opts.timeout,
                )?))
            }
        }
    };

    // Explicitly hold it until here.  Not strictly necessary; just for clarity.
    std::mem::drop(stdindrop);
    res
}

/// Execute a job using a job file piped in on stdin.  This is just a very
/// thin wrapper around [`cmd_exec_generic`].
pub fn cmd_execstdin(opts: &ProcessOpts) -> Result<(), anyhow::Error> {
    let res = cmd_exec_generic(opts, None, None, None, None)?;

    if res.success() {
        Ok(())
    } else {
        bail!("Command exited with non-success status {:?}", res);
    }
}

pub trait SeqIncrement {
    /// Increments an embedded seqf if we are a sequence-order iterator.
    fn increment_if_seq_order(&mut self) -> Result<(), anyhow::Error>;
}

/// An iterator over the jobs in sequence order.  An external command
/// must be used to increment the seqfile.
pub struct JobsBySeqIter<'a> {
    seqf: SeqFile<'a>,
    map: BTreeMap<u64, (OsString, FSMetaV1)>,
}

impl<'a> Iterator for JobsBySeqIter<'a> {
    type Item = (OsString, FSMetaV1);

    fn next(&mut self) -> Option<Self::Item> {
        self.map.remove(&self.seqf.get_next())
    }
}

/// An iterator over the jobs in timestamp order.
pub struct JobsByTimestampIter<'a> {
    #[allow(dead_code)]
    seqf: SeqFile<'a>,

    entries: Box<dyn Iterator<Item = (OsString, FSMetaV1)>>,
}

impl<'a> JobsByTimestampIter<'a> {
    pub fn new(seqf: SeqFile<'a>, mut entries: Vec<(OsString, FSMetaV1)>) -> Self {
        entries.sort_by(|(_, ma), (_, mb)| ma.ctime.cmp(&mb.ctime));
        JobsByTimestampIter {
            seqf,
            entries: Box::new(entries.into_iter()),
        }
    }
}

impl<'a> Iterator for JobsByTimestampIter<'a> {
    type Item = (OsString, FSMetaV1);

    fn next(&mut self) -> Option<Self::Item> {
        self.entries.next()
    }
}

impl<'a> SeqIncrement for JobsBySeqIter<'a> {
    fn increment_if_seq_order(&mut self) -> Result<(), anyhow::Error> {
        self.seqf.increment()?;
        Ok(())
    }
}

impl<'a> SeqIncrement for JobsByTimestampIter<'a> {
    fn increment_if_seq_order(&mut self) -> Result<(), anyhow::Error> {
        Ok(())
    }
}

pub trait IncrementJobsIter: Iterator<Item = (OsString, FSMetaV1)> + SeqIncrement {}

impl<'a> IncrementJobsIter for JobsByTimestampIter<'a> {}
impl<'a> IncrementJobsIter for JobsBySeqIter<'a> {}

/// The main worker of the program; it executes a queue.
pub fn cmd_execqueue(opts: &QueueProcessOpts) -> Result<(), anyhow::Error> {
    let seqfn = get_seqfile(&opts.qdopts.qopts.queuedir);
    let mut lock = seqfile::prepare_seqfile_lock(&seqfn, false)?;

    let seqf = SeqFile::open(&seqfn, &mut lock)?;

    let mut iter: Box<dyn IncrementJobsIter> = match opts.order_by {
        OrderBy::Sequence => {
            let qmap = scanqueue_map(&opts.qdopts.qopts.queuedir, &opts.qdopts.decoder)?;
            Box::new(JobsBySeqIter { seqf, map: qmap })
        }
        OrderBy::Timestamp => {
            let entries: Vec<(OsString, FSMetaV1)> =
                scanqueue(&opts.qdopts.qopts.queuedir, &opts.qdopts.decoder)?
                    .flatten()
                    .collect();
            Box::new(JobsByTimestampIter::new(seqf, entries))
        }
    };

    let mut jobsprocessed = 0;
    let maxjobs = opts.maxjobs.unwrap_or(u64::MAX);
    while let Some((filename, _meta)) = iter.next() {
        if jobsprocessed >= maxjobs {
            debug!(
                "Stopping processing as requested maximum jobs processed {} has been hit",
                maxjobs
            );
            return Ok(());
        }
        jobsprocessed += 1;

        let mut inputinfo =
            queue_openjob(&opts.qdopts.qopts.queuedir, &filename, &opts.qdopts.decoder)?;
        let inhandle = inputinfo.reader.take().unwrap();
        let fullfilename = queue_genfilename(&opts.qdopts.qopts.queuedir, &filename);
        let res = match opts.output_to {
            OutputTo::PassBoth => cmd_exec_generic(
                &opts.processopts,
                Some(inhandle),
                Some(&filename.clone()),
                Some(&opts.qdopts.qopts.queuedir),
                None,
            )?,
            OutputTo::SaveBoth => {
                let mut savefilename = OsString::from(&fullfilename);
                savefilename.push(".out");
                trace!(
                    "Writing stdout and stderr to {:?} per request",
                    savefilename
                );
                let savefile = File::create(savefilename)?;
                cmd_exec_generic(
                    &opts.processopts,
                    Some(inhandle),
                    Some(&filename.clone()),
                    Some(&opts.qdopts.qopts.queuedir),
                    Some(savefile),
                )?
            }
        };

        match (res.success(), &opts.on_error) {
            (true, _) | (false, OnError::Delete) => {
                iter.increment_if_seq_order()?;
                if !opts.never_delete {
                    // Now we can make sure to drop the decoder before deleting.  Probably
                    // harmless not to on POSIX - early versions didn't with no problems -
                    // but it feels like good form to do so.
                    std::mem::drop(inputinfo);

                    debug!("Deleting {:?}", fullfilename);
                    remove_file(fullfilename)?;
                }
            }
            (false, OnError::Retry) => bail!(
                "Aborting processing due to exit status {:?} from command",
                res
            ),
            (false, OnError::Leave) => {
                iter.increment_if_seq_order()?;
            }
        }
    }
    match opts.order_by {
        OrderBy::Sequence => debug!(
            "Sequence processing: Stopping processing as there is no job with the next sequence ID to process"),
        OrderBy::Timestamp => debug!(
            "Timestamp processing: Stopping processing as all jobs in queue have been processed"),
    };

    Ok(())
}