json_job_dispatch/

utils.rs

// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Utilities for director-based tools.
//!
//! Tools written using this crate usually have other tasks related to management of the job files.
//! These functions are meant to be used in the tools so that these tasks are built into the tool
//! rather than managed using external scripts.

use std::fs::{self, File};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::thread;
use std::time::Duration;

use chrono::Utc;
use log::{error, warn};
use lzma::LzmaWriter;
use rand::Rng;
use serde::Serialize;
use serde_json::{json, Value};
use tar::Builder;
use tempfile::TempDir;
use thiserror::Error;

use crate::Outbox;

/// An error creating a utiltiy job.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum JobError {
    /// Failure to create a job.
    #[error("failed to create job {}: {}", filepath.display(), source)]
    Create {
        /// The path to the job.
        filepath: PathBuf,
        /// The source of the error.
        #[source]
        source: io::Error,
    },
    /// Failure to serialize a job.
    #[error("failed to write job file {}: {}", filepath.display(), source)]
    Write {
        /// The path to the job.
        filepath: PathBuf,
        /// The source of the error.
        #[source]
        source: serde_json::Error,
    },
}

impl JobError {
    fn create(filepath: PathBuf, source: io::Error) -> Self {
        JobError::Create {
            filepath,
            source,
        }
    }

    fn write(filepath: PathBuf, source: serde_json::Error) -> Self {
        JobError::Write {
            filepath,
            source,
        }
    }
}

type JobResult<T> = Result<T, JobError>;

/// Write a job object to a new file in a directory.
fn write_job(queue: &Path, data: &Value) -> JobResult<()> {
    let rndpart = rand::rng()
        .sample_iter(&rand::distr::Alphanumeric)
        .map(char::from)
        .take(12)
        .collect::<String>();
    let filename = format!("{}-{}.json", Utc::now().to_rfc3339(), rndpart);
    let job_file = queue.join(filename);

    // Write and close the job file.
    {
        let mut file =
            File::create(&job_file).map_err(|err| JobError::create(job_file.clone(), err))?;
        serde_json::to_writer(&mut file, data)
            .map_err(|err| JobError::write(job_file.clone(), err))?;
    }

    // Wait for the job to have been processed.
    while job_file.exists() {
        thread::sleep(Duration::from_millis(100));
    }

    Ok(())
}

/// Write a job to the given queue.
pub fn drop_job<Q, K, V>(queue: Q, kind: K, data: V) -> JobResult<()>
where
    Q: AsRef<Path>,
    K: AsRef<str>,
    V: Serialize,
{
    let job = json!({
        "kind": kind.as_ref(),
        "data": data,
    });
    write_job(queue.as_ref(), &job)
}

/// Write a restart job to the given queue.
pub fn restart<Q>(queue: Q) -> JobResult<()>
where
    Q: AsRef<Path>,
{
    drop_job(queue, "watchdog:restart", json!({}))
}

/// Write an exit job to the given queue.
pub fn exit<Q>(queue: Q) -> JobResult<()>
where
    Q: AsRef<Path>,
{
    drop_job(queue, "watchdog:exit", json!({}))
}

/// Log that a temporary directory containing jobs moved for archiving but left unarchived due to
/// an error still exists.
fn log_tempdir(tempdir: TempDir) {
    error!(
        "archival failed mid-stream; in-progress archival job files may be found in {}",
        tempdir.into_path().display(),
    );
}

/// An error archiving a queue.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ArchiveQueueError {
    /// Failure to create an output file.
    #[error("failed to create archive file {}: {}", path.display(), source)]
    CreateOutput {
        #[doc(hidden)]
        outbox: Outbox,
        /// The path to the archive file.
        path: PathBuf,
        /// The source of the error.
        #[source]
        source: io::Error,
    },
    /// Failure to create a compressor for the archive.
    #[error("failed to create compressor: {}", source)]
    CreateCompressor {
        #[doc(hidden)]
        outbox: Outbox,
        /// The source of the error.
        #[source]
        source: lzma::LzmaError,
    },
    /// Failure when archiving jobs.
    #[error("failed to archive the {} outbox: {}", outbox, source)]
    Archive {
        #[doc(hidden)]
        outbox: Outbox,
        /// The source of the error.
        #[source]
        source: ArchiveError,
    },
    /// Failure when finishing the archive.
    #[error("failed to finish the {} archive: {}", outbox, source)]
    FinishArchive {
        #[doc(hidden)]
        outbox: Outbox,
        /// The source of the error.
        #[source]
        source: lzma::LzmaError,
    },
    /// Failure to remove an incomplete archive.
    #[error("failed to remove an incomplete {} archive: {}", outbox, source)]
    RemoveIncompleteArchive {
        #[doc(hidden)]
        outbox: Outbox,
        /// The source of the error.
        #[source]
        source: io::Error,
    },
}

impl ArchiveQueueError {
    /// The outbox being archived when the error occurred.
    pub fn outbox(&self) -> Outbox {
        match self {
            ArchiveQueueError::CreateOutput {
                outbox, ..
            }
            | ArchiveQueueError::CreateCompressor {
                outbox, ..
            }
            | ArchiveQueueError::Archive {
                outbox, ..
            }
            | ArchiveQueueError::FinishArchive {
                outbox, ..
            }
            | ArchiveQueueError::RemoveIncompleteArchive {
                outbox, ..
            } => *outbox,
        }
    }

    fn create_output(outbox: Outbox, path: PathBuf, source: io::Error) -> Self {
        ArchiveQueueError::CreateOutput {
            outbox,
            path,
            source,
        }
    }

    fn create_compressor(outbox: Outbox, source: lzma::LzmaError) -> Self {
        ArchiveQueueError::CreateCompressor {
            outbox,
            source,
        }
    }

    fn archive(outbox: Outbox, source: ArchiveError) -> Self {
        ArchiveQueueError::Archive {
            outbox,
            source,
        }
    }

    fn finish_archive(tempdir: TempDir, outbox: Outbox, source: lzma::LzmaError) -> Self {
        log_tempdir(tempdir);
        ArchiveQueueError::FinishArchive {
            outbox,
            source,
        }
    }

    fn remove_incomplete_archive(outbox: Outbox, source: io::Error) -> Self {
        ArchiveQueueError::RemoveIncompleteArchive {
            outbox,
            source,
        }
    }
}

/// The LZMA compression level to use for archiving.
const LZMA_COMPRESSION: u32 = 6;

/// Archive the jobs in the given queue into a tarball in the output directory.
///
/// Each subdirectory, `accept`, `fail`, and `reject` will be archived separately.
pub fn archive_queue<Q, O>(queue: Q, output: O) -> Result<(), ArchiveQueueError>
where
    Q: AsRef<Path>,
    O: AsRef<Path>,
{
    archive_queue_impl(queue.as_ref(), output.as_ref())
}

fn archive_queue_impl(queue: &Path, output: &Path) -> Result<(), ArchiveQueueError> {
    for outbox in [Outbox::Accept, Outbox::Fail, Outbox::Reject]
        .iter()
        .cloned()
    {
        let (filename, file) = archive_file(output, outbox)?;
        let opt_writer = archive_directory(queue, output, outbox, file)
            .map_err(|err| ArchiveQueueError::archive(outbox, err))?;
        if let Some((tempdir, writer)) = opt_writer {
            writer
                .finish()
                .map_err(|err| ArchiveQueueError::finish_archive(tempdir, outbox, err))?;
        } else {
            fs::remove_file(filename)
                .map_err(|err| ArchiveQueueError::remove_incomplete_archive(outbox, err))?;
        }
    }

    Ok(())
}

/// An error when writing to an archive.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ArchiveError {
    /// Failure to create a temporary directory for working.
    #[error("failed to create temporary directory: {}", source)]
    CreateTempdir {
        /// The source of the error.
        #[source]
        source: io::Error,
    },
    /// Failure to add the root directory to the archive.
    #[error("failed to add directory: {}", source)]
    AddDirectory {
        /// The source of the error.
        #[source]
        source: io::Error,
    },
    /// Failure to read the directory for processed job files.
    #[error("failed to read directory: {}", source)]
    ReadDirectory {
        /// The source of the error.
        #[source]
        source: io::Error,
    },
    /// Failure to append a job to the archive.
    #[error("failed to append job: {}", source)]
    AppendJob {
        /// The file that could not be appended.
        path: PathBuf,
        /// The source of the error.
        #[source]
        source: io::Error,
    },
    /// Failure to finish the tar stream.
    #[error("failed to finish tar stream: {}", source)]
    FinishTar {
        /// The source of the error.
        #[source]
        source: io::Error,
    },
}

impl ArchiveError {
    fn create_tempdir(source: io::Error) -> Self {
        ArchiveError::CreateTempdir {
            source,
        }
    }

    fn add_directory(source: io::Error) -> Self {
        ArchiveError::AddDirectory {
            source,
        }
    }

    fn read_directory(source: io::Error) -> Self {
        ArchiveError::ReadDirectory {
            source,
        }
    }

    fn append_job(tempdir: TempDir, path: PathBuf, source: io::Error) -> Self {
        log_tempdir(tempdir);
        ArchiveError::AppendJob {
            path,
            source,
        }
    }

    fn finish_tar(tempdir: TempDir, source: io::Error) -> Self {
        log_tempdir(tempdir);
        ArchiveError::FinishTar {
            source,
        }
    }
}

/// Create an archive file stream in the given path for the `outbox` files.
fn archive_file(
    path: &Path,
    outbox: Outbox,
) -> Result<(PathBuf, LzmaWriter<File>), ArchiveQueueError> {
    let now = Utc::now();
    let filepath = path.join(format!("{}-{}.tar.xz", now.to_rfc3339(), outbox));
    let file = File::create(&filepath)
        .map_err(|err| ArchiveQueueError::create_output(outbox, filepath.clone(), err))?;
    let writer = LzmaWriter::new_compressor(file, LZMA_COMPRESSION)
        .map_err(|err| ArchiveQueueError::create_compressor(outbox, err))?;

    Ok((filepath, writer))
}

/// Archive a directory into an output stream.
fn archive_directory<O>(
    path: &Path,
    workdir: &Path,
    outbox: Outbox,
    output: O,
) -> Result<Option<(TempDir, O)>, ArchiveError>
where
    O: Write,
{
    let tempdir = TempDir::new_in(workdir).map_err(ArchiveError::create_tempdir)?;
    let outbox_name = outbox.name();
    let outbox_path = PathBuf::from(outbox_name);
    let mut archive = Builder::new(output);
    archive
        .append_dir(outbox_name, &tempdir)
        .map_err(ArchiveError::add_directory)?;

    // Any errors after this point will need to perform `tempdir.into_path()` to prevent it from
    // removing any job files which have been staged there in preparation for removal because the
    // archive will not be successful at this point.

    let entries = fs::read_dir(path.join(outbox_name)).map_err(ArchiveError::read_directory)?;
    let mut is_empty = true;
    for entry in entries {
        let entry = match entry {
            Ok(entry) => entry,
            Err(err) => {
                warn!("failed to read directory entry; skipping: {:?}", err);
                continue;
            },
        };
        let path = entry.path();
        let file_name = entry.file_name();
        if let Err(err) = archive.append_path_with_name(&path, outbox_path.join(&file_name)) {
            // The state of `archive` is unknown if an error occurs when adding data to it; it
            // cannot be trusted anymore, so bail out.
            //
            // See https://github.com/alexcrichton/tar-rs/issues/213
            return Err(ArchiveError::append_job(tempdir, path, err));
        }
        // The file has been added to the archive; move it to the staging directory.
        let target_path = tempdir.path().join(&file_name);
        match fs::rename(&path, &target_path) {
            Ok(()) => is_empty = false,
            Err(err) => {
                warn!(
                    "failed to rename {} to {}: {:?}",
                    path.display(),
                    target_path.display(),
                    err,
                )
            },
        }
    }

    if is_empty {
        // All good to let `tempdir` be destructed here; we didn't move anything to staging
        // directory anyways.
        return Ok(None);
    }

    // Both branches here need to pass on `tempdir`. The `Ok` path might still have an error before
    // everything is good to go.
    match archive.into_inner() {
        Ok(output) => Ok(Some((tempdir, output))),
        Err(err) => Err(ArchiveError::finish_tar(tempdir, err)),
    }
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::sync::{Arc, Mutex};
    use std::thread;
    use std::time::Duration;

    use serde_json::json;

    use crate::{test, utils};
    use crate::{Director, DirectorWatchdog, HandlerCore, RunResult};

    #[test]
    fn test_log_tempdir_leaves_directory() {
        let tempdir = test::test_workspace_dir();
        let path = tempdir.path().to_path_buf();

        super::log_tempdir(tempdir);
        assert!(path.exists(), "{} should still exist", path.display());
        fs::remove_dir_all(path).unwrap();
    }

    #[test]
    fn test_write_job_waits_for_consumption() {
        let tempdir = test::test_workspace_dir();
        let mut director = Director::new(tempdir.path()).unwrap();
        let mutex = Arc::new(Mutex::new(false));
        let signal = test::FlaggingWatcher::new(mutex.clone());
        signal.add_to_director(&mut director).unwrap();
        DirectorWatchdog.add_to_director(&mut director).unwrap();

        let thread_mutex = mutex.clone();
        let job_path = tempdir.path().to_path_buf();
        let write_thread = thread::spawn(move || {
            let data = json!({
                "kind": "flagging:set",
                "data": {},
            });
            utils::write_job(&job_path, &data).unwrap();
            assert!(
                *thread_mutex.lock().unwrap(),
                "`write_job` returned before the handler set the flag",
            );

            utils::exit(job_path).unwrap();
        });

        // Wait for a file to appear.
        loop {
            let entries = fs::read_dir(tempdir.path()).unwrap();
            let nfiles = entries
                .filter(|entry| entry.as_ref().unwrap().file_type().unwrap().is_file())
                .count();
            if nfiles == 1 {
                break;
            }
        }
        // Let the other thread make progress.
        thread::sleep(Duration::from_millis(100));

        let res = director.watch_directory(tempdir.path()).unwrap();
        assert_eq!(res, RunResult::Done);

        write_thread.join().unwrap();
    }

    #[test]
    fn test_archive_queue() {
        let tempdir = test::test_workspace_dir();
        let archive = tempdir.path().join("archive");
        let output = tempdir.path().join("output");
        fs::create_dir_all(&archive).unwrap();
        let _director = Director::new(&archive).unwrap(); // Create structure.
        fs::create_dir_all(&output).unwrap();

        let accept = archive.join("accept");
        let reject = archive.join("reject");
        let fail = archive.join("fail");

        assert_eq!(
            accept.read_dir().unwrap().count(),
            0,
            "`accept` is not an empty directory",
        );
        assert_eq!(
            reject.read_dir().unwrap().count(),
            0,
            "`reject` is not an empty directory",
        );
        assert_eq!(
            fail.read_dir().unwrap().count(),
            0,
            "`fail` is not an empty directory",
        );
        assert_eq!(
            output.read_dir().unwrap().count(),
            0,
            "`output` is not an empty directory",
        );
        utils::archive_queue(&archive, &output).unwrap();
        assert_eq!(
            output.read_dir().unwrap().count(),
            0,
            "archived empty queues",
        );

        fs::write(accept.join("accept.json"), b"{}").unwrap();
        fs::write(reject.join("reject.json"), b"{}").unwrap();
        fs::write(fail.join("fail.json"), b"{}").unwrap();
        utils::archive_queue(archive, &output).unwrap();
        assert_eq!(
            output.read_dir().unwrap().count(),
            3,
            "not all archives archived",
        );
    }

    #[test]
    fn test_drop_restart() {
        let tempdir = test::test_workspace_dir();
        let mut director = Director::new(tempdir.path()).unwrap();
        DirectorWatchdog.add_to_director(&mut director).unwrap();

        let job_path = tempdir.path().to_path_buf();
        let write_thread = thread::spawn(move || {
            utils::restart(job_path).unwrap();
        });

        let res = director.watch_directory(tempdir.path()).unwrap();
        assert_eq!(res, RunResult::Restart);

        write_thread.join().unwrap();
    }
}