json-job-dispatch 1.2.0

// Copyright 2016 Kitware, Inc.
//
// 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 crates::chrono::Utc;
use crates::lzma::LzmaWriter;
use crates::rand::{self, Rng};
use crates::serde::Serialize;
use crates::serde_json::{self, Value};
use crates::tar::Builder;
use crates::tempdir::TempDir;

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

error_chain! { }

/// Write a job object to a new file in a directory.
fn write_job(queue: &Path, data: &Value) -> Result<()> {
    let rndpart = rand::thread_rng()
        .gen_ascii_chars()
        .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).chain_err(|| "failed to create the job file")?;
        serde_json::to_writer(&mut file, data)
            .chain_err(|| format!("failed to write the job to {}", filename))?;
    }

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

    loop {
        // Wait for the job to have been processed.
        if !job_file.exists() {
            break;
        }

        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) -> Result<()>
    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) -> Result<()>
    where Q: AsRef<Path>,
{
    drop_job(queue, "watchdog:restart", json!({}))
}

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

/// 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<()>
    where Q: AsRef<Path>,
          O: AsRef<Path>,
{
    for result in &["accept", "fail", "reject"] {
        let (filename, file) = archive_file(output.as_ref(), result)?;
        let opt_writer = archive_directory(queue.as_ref(), output.as_ref(), result, file)?;
        if let Some(mut writer) = opt_writer {
            writer.finish()
                .chain_err(|| format!("failed to finish archive stream for {}", result))?;
        } else {
            fs::remove_file(&filename)
                .chain_err(|| format!("failed to delete file {}", filename.display()))?;
        }
    }

    Ok(())
}

/// Create an archive file stream in the given path for the `result` files.
fn archive_file(path: &Path, result: &str) -> Result<(PathBuf, LzmaWriter<File>)> {
    let now = Utc::now();
    let filepath = path.join(format!("{}-{}.tar.xz", now.to_rfc3339(), result));
    let file = File::create(&filepath).chain_err(|| {
        format!("failed to create output file {}",
                filepath.display())
    })?;
    let writer = LzmaWriter::new_compressor(file, LZMA_COMPRESSION)
        .chain_err(|| "failed to construct LZMA writer")?;

    Ok((filepath, writer))
}

/// Archive a directory into an output stream.
fn archive_directory<O>(path: &Path, workdir: &Path, subdir: &str, output: O) -> Result<Option<O>>
    where O: Write,
{
    let tempdir = TempDir::new_in(workdir, "archive-jobs")
        .chain_err(|| "failed to create temporary directory")?;
    let targetdir = tempdir.path().join(subdir);
    fs::create_dir_all(&targetdir).chain_err(|| "failed to create target directory")?;
    let mut archive = Builder::new(output);
    archive.append_dir(subdir, &targetdir)
        .chain_err(|| format!("failed to append directory to {}", subdir))?;
    let entries = fs::read_dir(path.join(subdir)).chain_err(|| "failed to read input directory")?;
    let mut is_empty = true;
    for entry in entries {
        is_empty = false;
        let entry = entry.chain_err(|| "failed to read entry")?;
        let path = entry.path();
        let file_name = path.file_name().expect("expected the path to have a filename");
        let target_path = targetdir.join(&file_name);
        fs::rename(&path, &target_path)
            .chain_err(|| format!("failed to move input file {}", path.display()))?;
        let mut added_file = File::open(&target_path)
            .chain_err(|| format!("failed to open job file {}", target_path.display()))?;
        archive.append_file(format!("{}/{}", subdir, file_name.to_string_lossy()),
                            &mut added_file)
            .chain_err(|| {
                format!("failed to append {} to the archive",
                        target_path.display())
            })?;
    }

    if is_empty {
        return Ok(None);
    }

    Ok(Some(archive.into_inner()
        .chain_err(|| format!("failed to finish TAR stream for {}", subdir))?))
}