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.

use crates::chrono::Utc;
use crates::inotify::EventMask;
use crates::itertools::Itertools;
use crates::rand::{self, Rng};
use crates::serde_json::{self, Map, Value};

use error::*;
use handler::{Handler, HandlerResult};
use watcher::Watcher;

use std::borrow::Cow;
use std::collections::hash_map::{Entry, HashMap};
use std::ffi::OsStr;
use std::fmt::{self, Debug};
use std::fs::{self, File};
use std::io::Write;
use std::path::{Path, PathBuf};

/// Dispatch jobs to registered handlers.
///
/// Jobs are sorted into `accept`, `reject`, and `fail` directories based on whether they were
/// accepted by the relevant handler. Once handled, a `.stamp` file containing the timestamp of
/// when the job was completed is created beside the final location. In addition, rejected and
/// failed jobs have a `.reason` file describing what happened.
pub struct Director<'a> {
    /// The path to the directory for accepted jobs.
    accept: PathBuf,
    /// The path to the directory for rejected jobs.
    reject: PathBuf,
    /// The path to the directory for failed jobs.
    fail: PathBuf,

    /// The set of handlers to run.
    handlers: HashMap<String, &'a Handler>,
}

/// The filename of a path for logging purposes.
fn log_file_name<P>(path: &P) -> Cow<str>
    where P: AsRef<Path>,
{
    path.as_ref().file_name().map_or_else(|| Cow::Borrowed("."), OsStr::to_string_lossy)
}

impl<'a> Debug for Director<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "Director {{ accept: {:?}, reject: {:?}, fail: {:?}, handlers: {:?} }}",
               self.accept,
               self.reject,
               self.fail,
               self.handlers.keys().collect::<Vec<_>>())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
/// The result of running jobs.
pub enum RunResult {
    /// The director should continue running.
    Continue,
    /// The director should be restarted from the beginning.
    Restart,
    /// The director should be terminated.
    Done,
}

impl RunResult {
    /// Whether the result indicates that the director should exit.
    fn should_exit(&self) -> bool {
        if let RunResult::Continue = *self {
            false
        } else {
            true
        }
    }

    /// Whether the result indicates that the director is done.
    fn is_done(&self) -> bool {
        if let RunResult::Done = *self {
            true
        } else {
            false
        }
    }
}

impl<'a> Director<'a> {
    /// Creates a new `Director`.
    ///
    /// The `accept` and `reject` directories are created as children of the given directory.
    pub fn new(root: &Path) -> Result<Self> {
        if !root.is_dir() {
            bail!(ErrorKind::NotADirectory(root.to_path_buf()));
        }

        info!(target: "director", "setting up a director in {}", root.display());

        let accept_dir = root.join("accept");
        let reject_dir = root.join("reject");
        let fail_dir = root.join("fail");
        fs::create_dir_all(&accept_dir).chain_err(|| "failed to create the 'accept' directory")?;
        fs::create_dir_all(&reject_dir).chain_err(|| "failed to create the 'reject' directory")?;
        fs::create_dir_all(&fail_dir).chain_err(|| "failed to create the 'fail' directory")?;

        Ok(Director {
            accept: accept_dir,
            reject: reject_dir,
            fail: fail_dir,

            handlers: HashMap::new(),
        })
    }

    /// Add a handler for jobs labeled as `kind`.
    pub fn add_handler<K>(&mut self, kind: K, handler: &'a Handler) -> Result<()>
        where K: ToString,
    {
        match self.handlers.entry(kind.to_string()) {
            Entry::Occupied(o) => bail!(ErrorKind::DuplicateHandler(o.key().clone())),
            Entry::Vacant(v) => {
                debug!(target: "director", "adding handler '{}'", v.key());

                v.insert(handler);
                Ok(())
            },
        }
    }

    /// Move the job file into the appropriate directory.
    fn tag(&self, target_dir: &Path, file: &Path) -> Result<PathBuf> {
        let mut target_path = target_dir.to_path_buf();
        target_path.push(file.file_name().expect("expected the input file to have a file name"));

        // Write the stamp file.
        let mut stamp_file = File::create(target_path.with_extension("stamp")).chain_err(|| {
            format!("failed to create a stamp file for {}",
                    log_file_name(&target_path))
        })?;
        let time = Utc::now();
        write!(stamp_file, "{}\n", time.to_string()).chain_err(|| {
            format!("failed to write the stamp file for {}",
                    log_file_name(&target_path))
        })?;

        // Rename the file into the target path.
        fs::rename(file, &target_path)
            .chain_err(|| format!("failed to move to {}", target_path.display()))?;

        Ok(target_path)
    }

    /// Tag a job file into a directory with a given reason.
    fn tag_with_reason<R>(&self, target_dir: &Path, file: &Path, reason: R) -> Result<()>
        where R: fmt::Display,
    {
        let target_file = self.tag(target_dir, file)?;
        self.write_reason(&target_file, reason)
    }

    /// Write a reason file.
    fn write_reason<R>(&self, target_file: &Path, reason: R) -> Result<()>
        where R: fmt::Display,
    {
        let mut reason_file = File::create(target_file.with_extension("reason")).chain_err(|| {
            format!("failed to create the reason file for {}",
                    log_file_name(&target_file))
        })?;
        write!(reason_file, "{}\n", reason).chain_err(|| {
            format!("failed to write the reason file for {}",
                    log_file_name(&target_file))
        })?;

        Ok(())
    }

    /// Handle a single file.
    ///
    /// Returns `true` if the loop should terminate.
    pub fn run_one(&self, path: &Path) -> RunResult {
        match self.dispatch(path) {
            Ok(res) => res,
            Err(err) => {
                error!("error when handling {}: {}", log_file_name(&path), err);
                RunResult::Done
            },
        }
    }

    /// Handle a batch of files at once.
    ///
    /// Returns `true` if the loop should terminate.
    pub fn run(&self, paths: Vec<PathBuf>) -> RunResult {
        let mut result = RunResult::Continue;

        for path in paths {
            let one_result = self.run_one(&path);
            if one_result > result {
                result = one_result;
            }

            if result.is_done() {
                break;
            }
        }

        result
    }

    /// Watch a directory for new job files and act upon them.
    ///
    /// Upon startup, this searches the directory for existing files. These files are processed in
    /// lexigraphic order.
    pub fn watch_directory(&self, path: &Path) -> Result<RunResult> {
        let mut watcher = {
            let paths = path.read_dir()
                .chain_err(|| format!("failed to list the queue directory {}", path.display()))?;
            let watcher = Watcher::new(path).chain_err(|| {
                format!("failed to start the watcher in the queue directory {}",
                        path.display())
            })?;

            let path_entries = paths.filter_map(|e| e.ok())
                .filter_map(|e| {
                    if e.path().is_dir() {
                        None
                    } else {
                        Some(e.path())
                    }
                })
                .sorted();
            let result = self.run(path_entries);
            if result.should_exit() {
                return Ok(result);
            }

            watcher
        };

        let loop_result;

        loop {
            let events = watcher.events()
                .chain_err(|| format!("failed to watch the queue directory {}", path.display()))?;

            let path_entries = events.into_iter()
                .filter_map(|event| {
                    if event.mask.contains(EventMask::ISDIR) {
                        None
                    } else {
                        event.name.map(|name| path.join(name))
                    }
                })
                .sorted();
            let result = self.run(path_entries);
            if result.should_exit() {
                loop_result = result;
                break;
            }
        }

        Ok(loop_result)
    }

    /// Takes the file and directs it to the right location.
    fn dispatch(&self, file: &Path) -> Result<RunResult> {
        trace!(target: "director", "handling file {}", log_file_name(&file));

        let ext = file.extension()
            .map(OsStr::to_string_lossy);
        let ext_str = ext.as_ref()
            .map(AsRef::as_ref);

        if let Some("json") = ext_str {
        } else {
            return Ok(RunResult::Continue);
        }

        let event_file = File::open(file)
            .chain_err(|| format!("failed to open an input file: {}", file.display()))?;
        let mut payload: Value = if let Ok(payload) = serde_json::from_reader(event_file) {
            payload
        } else {
            return Ok(RunResult::Continue);
        };

        if !payload.is_object() {
            return Ok(RunResult::Continue);
        }

        let ret = match self.handle(&payload)? {
            HandlerResult::Accept => {
                debug!("accepted {}", log_file_name(&file));
                self.tag(&self.accept, file)?;
                RunResult::Continue
            },
            HandlerResult::Defer(ref reason) => {
                debug!("deferring {}: {}", log_file_name(&file), reason);
                let rndpart = rand::thread_rng()
                    .gen_ascii_chars()
                    .take(12)
                    .collect::<String>();
                let defer_job_file =
                    file.with_file_name(format!("{}-{}.json", Utc::now().to_rfc3339(), rndpart));
                Self::add_reason_to_object(&mut payload, file.to_string_lossy(), &reason);
                let mut defer_file = File::create(&defer_job_file).chain_err(|| {
                    format!("failed to create the job for the deferred job {} as {}",
                            log_file_name(&file),
                            log_file_name(&defer_job_file))
                })?;
                serde_json::to_writer(&mut defer_file, &payload).chain_err(|| {
                    format!("failed to write the deferred job to {}",
                            log_file_name(&defer_job_file))
                })?;
                self.tag_with_reason(&self.reject, file, reason)?;
                RunResult::Continue
            },
            HandlerResult::Reject(ref reason) => {
                debug!("rejecting {}: {}", file.display(), reason);
                self.tag_with_reason(&self.reject, file, reason)?;
                RunResult::Continue
            },
            HandlerResult::Fail(ref reason) => {
                debug!("failed {}: {:?}", file.display(), reason);
                self.tag_with_reason(&self.fail, file, format!("{:?}", reason))?;
                RunResult::Continue
            },
            HandlerResult::Restart => {
                info!(target: "director", "restarting via {}", log_file_name(&file));

                self.tag(&self.accept, file)?;
                RunResult::Restart
            },
            HandlerResult::Done => {
                info!(target: "director", "completed via {}", log_file_name(&file));

                self.tag(&self.accept, file)?;
                RunResult::Done
            },
        };

        trace!(target: "director", "handled file {}", log_file_name(&file));

        Ok(ret)
    }

    /// Add a reason to a JSON job object.
    fn add_reason_to_object<N, R>(object: &mut Value, name: N, reason: R)
        where N: ToString,
              R: ToString,
    {
        let retry_map = object.as_object_mut()
            .expect("expected an object; internal logic failure")
            .entry("retry")
            .or_insert_with(|| Value::Object(Map::new()));

        if !retry_map.is_object() {
            *retry_map = Value::Object(Map::new());
        }

        retry_map.as_object_mut()
            .expect("expected an object; internal logic failure") // Verified above.
            .insert(name.to_string(), Value::String(reason.to_string()));
    }

    /// Dispatch the contents of the file to the right handler.
    fn handle(&self, payload: &Value) -> Result<HandlerResult> {
        let kind = match payload.pointer("/kind") {
            Some(&Value::String(ref kind)) => kind,
            Some(_) => return Ok(HandlerResult::Reject("'kind' is not a string".to_string())),
            None => return Ok(HandlerResult::Reject("no 'kind'".to_string())),
        };
        let data = match payload.pointer("/data") {
            Some(data) => data,
            None => return Ok(HandlerResult::Reject("no 'data'".to_string())),
        };
        let retry_reasons = match payload.pointer("/retry") {
            Some(&Value::Object(ref reasons)) => {
                let retry_reasons = reasons.iter()
                    .map(|(_, reason)| {
                        reason.as_str()
                            .map(ToString::to_string)
                            .ok_or_else(|| {
                                HandlerResult::Reject("retry reason is not a string".to_string())
                            })
                    })
                    .collect::<::std::result::Result<Vec<_>, HandlerResult>>();

                match retry_reasons {
                    Ok(reasons) => reasons,
                    Err(reject) => return Ok(reject),
                }
            },
            Some(_) => return Ok(HandlerResult::Reject("'retry' is not an object".to_string())),
            None => vec![],
        };

        match self.handlers.get(kind) {
            Some(handler) => handler.handle_retry(kind, data, retry_reasons),
            None => Ok(HandlerResult::Reject(format!("no handler for kind '{}'", kind))),
        }
    }
}