<commit_before>#![cfg_attr(feature = "dev", allow(unstable_features))]
#![cfg_attr(feature = "dev", feature(plugin))]
#![cfg_attr(feature = "dev", plugin(clippy))]

mod logger;
mod toml_file;
mod git;
mod changelog;
mod commit_analyzer;
mod cargo;
mod error;
mod github;
mod config;

extern crate rustc_serialize;
extern crate toml;
extern crate regex;
extern crate semver;
extern crate docopt;
extern crate git2;
extern crate clog;
extern crate hyper;
extern crate hubcaps;
extern crate url;
extern crate travis_after_all;
extern crate env_logger;

use docopt::Docopt;
use commit_analyzer::CommitType;
use config::ConfigBuilder;
use std::process;
use semver::Version;
use std::{env,fs};
use std::path::Path;
use std::error::Error;
use std::thread;
use std::time::Duration;
use url::Url;
use travis_after_all::Build;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const USERAGENT: &'static str = concat!("semantic-rs/", env!("CARGO_PKG_VERSION"));
const USAGE: &'static str = "
semantic.rs 🚀

Usage:
  semantic-rs [options]
  semantic-rs --version

Options:
  -h --help              Show this screen.
  --version              Show version.
  -p PATH, --path=PATH   Specifies the repository path. [default: .]
  -w, --write            Run with writing the changes afterwards.
  -r <r>, --release=<r>  Create release on GitHub and publish on crates.io (only in write mode) [default: yes]
  -b <b>, --branch=<b>   The branch on which releases should happen. [default: master]
";

macro_rules! print_exit {
    ($fmt:expr) => {{
        logger::stderr($fmt);
        process::exit(1);
    }};
    ($fmt:expr, $($arg:tt)*) => {{
        logger::stderr(format!($fmt, $($arg)*));
        process::exit(1);
    }};
}

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_path: String,
    flag_write: bool,
    flag_version: bool,
    flag_release: String,
    flag_branch: String,
}

fn string_to_bool(answer: &str) -> bool {
    match &answer.to_lowercase()[..] {
        "yes" | "true" | "1" => true,
        _ => false
    }
}

fn version_bump(version: &Version, bump: CommitType) -> Option<Version> {
    let mut version = version.clone();
    match bump {
        CommitType::Unknown => return None,
        CommitType::Patch => version.increment_patch(),
        CommitType::Minor => version.increment_minor(),
        CommitType::Major => version.increment_major(),
    }

    Some(version)
}

fn ci_env_set() -> bool {
    env::var("CI").is_ok()
}

fn current_branch(repo: &git2::Repository) -> Option<String> {
    if let Ok(branch) = env::var("TRAVIS_BRANCH") {
        return Some(branch)
    }

    let head = repo.head().expect("No HEAD found for repository");

    if head.is_branch() {
        let short = head.shorthand().expect("No branch name found");
        return Some(short.into());
    }

    None
}

fn is_release_branch(current: &str, release: &str) -> bool {
    if let Ok(pr) = env::var("TRAVIS_PULL_REQUEST") {
        if pr != "false" {
            return false;
        }
    }

    current == release
}

fn user_repo_from_url(url: Url) -> Result<(String, String), String> {
    let path = match url.path() {
        Some(path) => path,
        None => return Err("URL should contain user and repository".into()),
    };

    let user = path[0].clone();
    let repo = match path[1].rfind(".git") {
        None => path[1].clone(),
        Some(suffix_pos) => {
            let valid_pos = path[1].len() - 4;
            if valid_pos == suffix_pos {
                let path = &path[1][0..suffix_pos];
                path.into()
            } else {
                return Err("URL does not point to a git repository".into())
            }
        }
    };

    Ok((user, repo))
}

fn main() {
    env_logger::init().expect("Can't instantiate env logger");

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let mut config_builder = ConfigBuilder::new();

    if args.flag_version {
        println!("semantic.rs 🚀 -- v{}", VERSION);
        process::exit(0);
    }

    let is_dry_run = if ci_env_set() {
        false
    }
    else {
        !args.flag_write
    };

    let release_mode = string_to_bool(&args.flag_release);

    config_builder.write(args.flag_write);
    config_builder.release(release_mode);
    config_builder.branch(args.flag_branch);

    println!("semantic.rs 🚀");

    logger::stdout("Analyzing your repository");
    let path = Path::new(&args.flag_path);
    let path = fs::canonicalize(path)
        .unwrap_or_else(|_| print_exit!("Path does not exist or a component is not a directory"));
    let repository_path = path.to_str()
        .unwrap_or_else(|| print_exit!("Path is not valid unicode"));

    let repo = match git2::Repository::open(repository_path) {
        Ok(repo) => repo,
        Err(e) => {
            logger::stderr(format!("Could not open the git repository: {:?}", e));
            process::exit(1);
        }
    };

    config_builder.repository_path(repository_path.to_owned());

    // extra scope scope to make sure borrow of `repo` is dropped
    {
        let signature = match git::get_signature(&repo) {
            Ok(sig) => sig,
            Err(e) => {
                logger::stderr(format!("Failed to get the committer's name and email address: {}", e.description()));
                logger::stderr(r"
A release commit needs a committer name and email address.
We tried fetching it from different locations, but couldn't find one.

Committer information is taken from the following environment variables, if set:

GIT_COMMITTER_NAME
GIT_COMMITTER_EMAIL

If none is set the normal git config is tried in the following order:

Local repository config
User config
Global config");
                process::exit(1);
            }
        };

        config_builder.signature(signature.to_owned());
    }

    // In case we are in write-mode AND release mode,
    // we will make sure we got all configuration settings
    if !is_dry_run && release_mode {
        let remote_or_none = repo.find_remote("origin");
        if remote_or_none.is_err() {
            let err = remote_or_none.err().unwrap();
            logger::warn(format!("Could not determine the origin remote url: {:?}", err));
        }
        else  {
            let remote = remote_or_none.ok().unwrap();
            let url = remote.url().expect("Remote URL is not valid UTF-8");
            let remote_url = Url::parse(&url).expect("Remote URL can't be parsed");
            let (user, repo_name) = user_repo_from_url(remote_url)
                .unwrap_or_else(|e| print_exit!("Could not extract user and repository name from URL: {:?}", e));
            config_builder.user(user);
            config_builder.repository_name(repo_name);

            let gh_token = env::var("GH_TOKEN")
                .unwrap_or_else(|err| print_exit!("GH_TOKEN not set: {:?}", err));

            let cargo_token = env::var("CARGO_TOKEN")
                .unwrap_or_else(|err| print_exit!("CARGO_TOKEN not set: {:?}", err));

            config_builder.gh_token(gh_token);
            config_builder.cargo_token(cargo_token);
        }
    }

    config_builder.repository(repo);
    let config = config_builder.build();

    let branch = current_branch(&config.repository)
        .unwrap_or_else(|| print_exit!("Could not determine current branch."));

    if !is_release_branch(&branch, &config.branch) {
        println!("Current branch is '{}', releases are only done from branch '{}'", branch, config.branch);
        println!("No release done from a pull request either.");
        process::exit(0);
    }

    if ci_env_set() {
        let build_run = Build::from_env()
            .unwrap_or_else(|e| print_exit!("CI mode, but can't check other builds. Error: {:?}", e));

        if !build_run.is_leader() {
            println!("Not the build leader. Nothing to do. Bye.");
            process::exit(0);
        }

        println!("I am the build leader. Waiting for other jobs to finish.");
        match build_run.wait_for_others() {
            Ok(()) => println!("Other jobs finished and succeeded. Doing my work now."),
            Err(travis_after_all::Error::FailedBuilds) => {
                print_exit!("Some builds failed. Stopping here.");
            },
            Err(e) => print_exit!("Waiting for other builds failed Reason: {:?}", e),
        }
    }

    let version = toml_file::read_from_file(&config.repository_path)
        .unwrap_or_else(|err| print_exit!("Reading `Cargo.toml` failed: {:?}", err));

    let version = Version::parse(&version).expect("Not a valid version");
    logger::stdout(format!("Current version: {}", version.to_string()));

    logger::stdout("Analyzing commits");

    let bump = git::version_bump_since_latest(&config.repository);
    if is_dry_run {
        logger::stdout(format!("Commits analyzed. Bump would be {:?}", bump));
    }
    else {
        logger::stdout(format!("Commits analyzed. Bump will be {:?}", bump));
    }
    let new_version = match version_bump(&version, bump) {
        Some(new_version) => new_version.to_string(),
        None => {
            logger::stdout("No version bump. Nothing to do.");
            process::exit(0);
        }
    };

    if is_dry_run {
        logger::stdout(format!("New version would be: {}", new_version));
        logger::stdout("Would write the following Changelog:");
        let changelog = match changelog::generate(repository_path, &version.to_string(), &new_version) {
            Ok(log) => log,
            Err(err) => {
                logger::stderr(format!("Generating Changelog failed: {:?}", err));
                process::exit(1);
            }
        };
        logger::stdout("====================================");
        logger::stdout(changelog);
        logger::stdout("====================================");
        logger::stdout("Would create annotated git tag");
    }
    else {
        logger::stdout(format!("New version: {}", new_version));

        toml_file::write_new_version(repository_path, &new_version)
            .unwrap_or_else(|err| print_exit!("Writing `Cargo.toml` failed: {:?}", err));

        logger::stdout("Writing Changelog");
        changelog::write(repository_path, &version.to_string(), &new_version)
            .unwrap_or_else(|err| print_exit!("Writing Changelog failed: {:?}", err));

        if config.release_mode {
            logger::stdout("Updating lockfile");
            if !cargo::update_lockfile(repository_path) {
                print_exit!("`cargo fetch` failed. See above for the cargo error message.");
            }
        }

        logger::stdout("Package crate");
        if !cargo::package(repository_path) {
            print_exit!("`cargo package` failed. See above for the cargo error message.");
        }

        git::commit_files(&config, &new_version)
            .unwrap_or_else(|err| print_exit!("Committing files failed: {:?}", err));

        logger::stdout("Creating annotated git tag");
        let tag_message = changelog::generate(repository_path, &version.to_string(), &new_version)
            .unwrap_or_else(|err| print_exit!("Can't generate changelog: {:?}", err));

        let tag_name = format!("v{}", new_version);
        git::tag(&config, &tag_name, &tag_message)
            .unwrap_or_else(|err| print_exit!("Failed to create git tag: {:?}", err));

        if config.release_mode && config.can_push() {
            logger::stdout("Pushing new commit and tag");
            git::push(&config, &tag_name)
                .unwrap_or_else(|err| print_exit!("Failed to push git: {:?}", err));

            logger::stdout("Waiting a tiny bit, so GitHub can store the git tag");
            thread::sleep(Duration::from_secs(1));

            logger::stdout("Creating GitHub release");
            github::release(&config, &tag_name, &tag_message)
                .unwrap_or_else(|err| print_exit!("Failed to create GitHub release: {:?}", err));

            logger::stdout("Publishing crate on crates.io");
            if !cargo::publish(&config.repository_path, &config.cargo_token.as_ref().unwrap()) {
                print_exit!("Failed to publish on crates.io");
            }

            println!("{} v{} is released. 🚀🚀🚀", config.repository_name.unwrap(), new_version);
        }
    }
}
<commit_msg>feat: provide additional information for the end user<commit_after>#![cfg_attr(feature = "dev", allow(unstable_features))]
#![cfg_attr(feature = "dev", feature(plugin))]
#![cfg_attr(feature = "dev", plugin(clippy))]

mod logger;
mod toml_file;
mod git;
mod changelog;
mod commit_analyzer;
mod cargo;
mod error;
mod github;
mod config;

extern crate rustc_serialize;
extern crate toml;
extern crate regex;
extern crate semver;
extern crate docopt;
extern crate git2;
extern crate clog;
extern crate hyper;
extern crate hubcaps;
extern crate url;
extern crate travis_after_all;
extern crate env_logger;

use docopt::Docopt;
use commit_analyzer::CommitType;
use config::ConfigBuilder;
use std::process;
use semver::Version;
use std::{env,fs};
use std::path::Path;
use std::error::Error;
use std::thread;
use std::time::Duration;
use url::Url;
use travis_after_all::Build;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const USERAGENT: &'static str = concat!("semantic-rs/", env!("CARGO_PKG_VERSION"));
const USAGE: &'static str = "
semantic.rs 🚀

Usage:
  semantic-rs [options]
  semantic-rs --version

Options:
  -h --help              Show this screen.
  --version              Show version.
  -p PATH, --path=PATH   Specifies the repository path. [default: .]
  -w, --write            Run with writing the changes afterwards.
  -r <r>, --release=<r>  Create release on GitHub and publish on crates.io (only in write mode) [default: yes]
  -b <b>, --branch=<b>   The branch on which releases should happen. [default: master]
";

macro_rules! print_exit {
    ($fmt:expr) => {{
        logger::stderr($fmt);
        process::exit(1);
    }};
    ($fmt:expr, $($arg:tt)*) => {{
        logger::stderr(format!($fmt, $($arg)*));
        process::exit(1);
    }};
}

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_path: String,
    flag_write: bool,
    flag_version: bool,
    flag_release: String,
    flag_branch: String,
}

fn string_to_bool(answer: &str) -> bool {
    match &answer.to_lowercase()[..] {
        "yes" | "true" | "1" => true,
        _ => false
    }
}

fn version_bump(version: &Version, bump: CommitType) -> Option<Version> {
    let mut version = version.clone();
    match bump {
        CommitType::Unknown => return None,
        CommitType::Patch => version.increment_patch(),
        CommitType::Minor => version.increment_minor(),
        CommitType::Major => version.increment_major(),
    }

    Some(version)
}

fn ci_env_set() -> bool {
    env::var("CI").is_ok()
}

fn current_branch(repo: &git2::Repository) -> Option<String> {
    if let Ok(branch) = env::var("TRAVIS_BRANCH") {
        return Some(branch)
    }

    let head = repo.head().expect("No HEAD found for repository");

    if head.is_branch() {
        let short = head.shorthand().expect("No branch name found");
        return Some(short.into());
    }

    None
}

fn is_release_branch(current: &str, release: &str) -> bool {
    if let Ok(pr) = env::var("TRAVIS_PULL_REQUEST") {
        if pr != "false" {
            return false;
        }
    }

    current == release
}

fn user_repo_from_url(url: Url) -> Result<(String, String), String> {
    let path = match url.path() {
        Some(path) => path,
        None => return Err("URL should contain user and repository".into()),
    };

    let user = path[0].clone();
    let repo = match path[1].rfind(".git") {
        None => path[1].clone(),
        Some(suffix_pos) => {
            let valid_pos = path[1].len() - 4;
            if valid_pos == suffix_pos {
                let path = &path[1][0..suffix_pos];
                path.into()
            } else {
                return Err("URL does not point to a git repository".into())
            }
        }
    };

    Ok((user, repo))
}

fn main() {
    env_logger::init().expect("Can't instantiate env logger");

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let mut config_builder = ConfigBuilder::new();

    if args.flag_version {
        println!("semantic.rs 🚀 -- v{}", VERSION);
        process::exit(0);
    }

    let is_dry_run = if ci_env_set() {
        false
    }
    else {
        !args.flag_write
    };

    let release_mode = string_to_bool(&args.flag_release);

    config_builder.write(args.flag_write);
    config_builder.release(release_mode);
    config_builder.branch(args.flag_branch);

    println!("semantic.rs 🚀");

    logger::stdout("Analyzing your repository");
    let path = Path::new(&args.flag_path);
    let path = fs::canonicalize(path)
        .unwrap_or_else(|_| print_exit!("Path does not exist or a component is not a directory"));
    let repository_path = path.to_str()
        .unwrap_or_else(|| print_exit!("Path is not valid unicode"));

    let repo = match git2::Repository::open(repository_path) {
        Ok(repo) => repo,
        Err(e) => {
            logger::stderr(format!("Could not open the git repository: {:?}", e));
            process::exit(1);
        }
    };

    config_builder.repository_path(repository_path.to_owned());

    // extra scope scope to make sure borrow of `repo` is dropped
    {
        let signature = match git::get_signature(&repo) {
            Ok(sig) => sig,
            Err(e) => {
                logger::stderr(format!("Failed to get the committer's name and email address: {}", e.description()));
                logger::stderr(r"
A release commit needs a committer name and email address.
We tried fetching it from different locations, but couldn't find one.

Committer information is taken from the following environment variables, if set:

GIT_COMMITTER_NAME
GIT_COMMITTER_EMAIL

If none is set the normal git config is tried in the following order:

Local repository config
User config
Global config");
                process::exit(1);
            }
        };

        config_builder.signature(signature.to_owned());
    }

    // In case we are in write-mode AND release mode,
    // we will make sure we got all configuration settings
    if !is_dry_run && release_mode {
        let remote_or_none = repo.find_remote("origin");
        if remote_or_none.is_err() {
            let err = remote_or_none.err().unwrap();
            logger::warn(format!("Could not determine the origin remote url: {:?}", err));
            logger::warn("semantic-rs therefor does not push anything to a remote repository does not create a release on GitHub");
        }
        else  {
            let remote = remote_or_none.ok().unwrap();
            let url = remote.url().expect("Remote URL is not valid UTF-8");
            let remote_url = Url::parse(&url).expect("Remote URL can't be parsed");
            let (user, repo_name) = user_repo_from_url(remote_url)
                .unwrap_or_else(|e| print_exit!("Could not extract user and repository name from URL: {:?}", e));
            config_builder.user(user);
            config_builder.repository_name(repo_name);

            let gh_token = env::var("GH_TOKEN")
                .unwrap_or_else(|err| print_exit!("GH_TOKEN not set: {:?}", err));

            let cargo_token = env::var("CARGO_TOKEN")
                .unwrap_or_else(|err| print_exit!("CARGO_TOKEN not set: {:?}", err));

            config_builder.gh_token(gh_token);
            config_builder.cargo_token(cargo_token);
        }
    }

    config_builder.repository(repo);
    let config = config_builder.build();

    let branch = current_branch(&config.repository)
        .unwrap_or_else(|| print_exit!("Could not determine current branch."));

    if !is_release_branch(&branch, &config.branch) {
        println!("Current branch is '{}', releases are only done from branch '{}'", branch, config.branch);
        println!("No release done from a pull request either.");
        process::exit(0);
    }

    if ci_env_set() {
        let build_run = Build::from_env()
            .unwrap_or_else(|e| print_exit!("CI mode, but can't check other builds. Error: {:?}", e));

        if !build_run.is_leader() {
            println!("Not the build leader. Nothing to do. Bye.");
            process::exit(0);
        }

        println!("I am the build leader. Waiting for other jobs to finish.");
        match build_run.wait_for_others() {
            Ok(()) => println!("Other jobs finished and succeeded. Doing my work now."),
            Err(travis_after_all::Error::FailedBuilds) => {
                print_exit!("Some builds failed. Stopping here.");
            },
            Err(e) => print_exit!("Waiting for other builds failed Reason: {:?}", e),
        }
    }

    let version = toml_file::read_from_file(&config.repository_path)
        .unwrap_or_else(|err| print_exit!("Reading `Cargo.toml` failed: {:?}", err));

    let version = Version::parse(&version).expect("Not a valid version");
    logger::stdout(format!("Current version: {}", version.to_string()));

    logger::stdout("Analyzing commits");

    let bump = git::version_bump_since_latest(&config.repository);
    if is_dry_run {
        logger::stdout(format!("Commits analyzed. Bump would be {:?}", bump));
    }
    else {
        logger::stdout(format!("Commits analyzed. Bump will be {:?}", bump));
    }
    let new_version = match version_bump(&version, bump) {
        Some(new_version) => new_version.to_string(),
        None => {
            logger::stdout("No version bump. Nothing to do.");
            process::exit(0);
        }
    };

    if is_dry_run {
        logger::stdout(format!("New version would be: {}", new_version));
        logger::stdout("Would write the following Changelog:");
        let changelog = match changelog::generate(repository_path, &version.to_string(), &new_version) {
            Ok(log) => log,
            Err(err) => {
                logger::stderr(format!("Generating Changelog failed: {:?}", err));
                process::exit(1);
            }
        };
        logger::stdout("====================================");
        logger::stdout(changelog);
        logger::stdout("====================================");
        logger::stdout("Would create annotated git tag");
    }
    else {
        logger::stdout(format!("New version: {}", new_version));

        toml_file::write_new_version(repository_path, &new_version)
            .unwrap_or_else(|err| print_exit!("Writing `Cargo.toml` failed: {:?}", err));

        logger::stdout("Writing Changelog");
        changelog::write(repository_path, &version.to_string(), &new_version)
            .unwrap_or_else(|err| print_exit!("Writing Changelog failed: {:?}", err));

        if config.release_mode {
            logger::stdout("Updating lockfile");
            if !cargo::update_lockfile(repository_path) {
                print_exit!("`cargo fetch` failed. See above for the cargo error message.");
            }
        }

        logger::stdout("Package crate");
        if !cargo::package(repository_path) {
            print_exit!("`cargo package` failed. See above for the cargo error message.");
        }

        git::commit_files(&config, &new_version)
            .unwrap_or_else(|err| print_exit!("Committing files failed: {:?}", err));

        logger::stdout("Creating annotated git tag");
        let tag_message = changelog::generate(repository_path, &version.to_string(), &new_version)
            .unwrap_or_else(|err| print_exit!("Can't generate changelog: {:?}", err));

        let tag_name = format!("v{}", new_version);
        git::tag(&config, &tag_name, &tag_message)
            .unwrap_or_else(|err| print_exit!("Failed to create git tag: {:?}", err));

        if config.release_mode && config.can_push() {
            logger::stdout("Pushing new commit and tag");
            git::push(&config, &tag_name)
                .unwrap_or_else(|err| print_exit!("Failed to push git: {:?}", err));

            logger::stdout("Waiting a tiny bit, so GitHub can store the git tag");
            thread::sleep(Duration::from_secs(1));

            logger::stdout("Creating GitHub release");
            github::release(&config, &tag_name, &tag_message)
                .unwrap_or_else(|err| print_exit!("Failed to create GitHub release: {:?}", err));

            logger::stdout("Publishing crate on crates.io");
            if !cargo::publish(&config.repository_path, &config.cargo_token.as_ref().unwrap()) {
                print_exit!("Failed to publish on crates.io");
            }

            println!("{} v{} is released. 🚀🚀🚀", config.repository_name.unwrap(), new_version);
        }
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::collections::HashSet;
use std::net::SocketAddr;

// MIO
use mio::tcp::TcpListener;
use mio::{Token, EventLoop, Handler, ReadHint};

// Data structures.
use store::Store;
use replica::Replica;
use state_machine::StateMachine;

// Cap'n Proto
use capnp::serialize_packed;
use capnp::{MessageBuilder, MessageReader, ReaderOptions, MallocMessageBuilder};
use messages_capnp::{
    rpc_request,
    rpc_response,
    client_request,
    client_response,
    request_vote_response,
    append_entries_response,
    append_entries_request,
};

// MIO Tokens
const SOCKET:  Token = Token(0);
const TIMEOUT: Token = Token(1);

/// The Raft Distributed Consensus Algorithm requires two RPC calls to be available:
///
///   * `append_entries` which is used as both a heartbeat (with no payload) and the primary
///     interface for requests.
///   * `request_vote` which is used by candidates during campaigns to obtain a vote.
///
/// A `RaftNode` acts as a replicated state machine. The server's role in the cluster depends on it's
/// own status. It will maintain both volatile state (which can be safely lost) and persistent
/// state (which must be carefully stored and kept safe).
///
/// Currently, the `RaftNode` API is not well defined. **We are looking for feedback and suggestions.**
pub struct RaftNode<S, M> where S: Store, M: StateMachine {
    replica: Replica<S, M>,
    // Channels and Sockets
    listener: TcpListener, // TODO: Just use streams??
}

type Reactor<S, M> = EventLoop<RaftNode<S, M>>;

/// The implementation of the RaftNode. In most use cases, creating a `RaftNode` should just be
/// done via `::new()`.
impl<S, M> RaftNode<S, M> where S: Store, M: StateMachine {

    /// Creates a new Raft node with the cluster members specified.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address of the new node.
    /// * `peers` - The address of all peers in the Raft cluster.
    /// * `store` - The persitent log store.
    /// * `state_machine` - The client state machine to which client commands will be applied.
    pub fn spawn(addr: SocketAddr,
                 peers: HashSet<SocketAddr>,
                 store: S,
                 state_machine: M) {
        // Create an event loop
        let mut event_loop = Reactor::<S, M>::new().unwrap();
        // Setup the socket, make it not block.
        let listener = TcpListener::bind(&addr).unwrap();
        event_loop.register(&listener, SOCKET).unwrap();
        event_loop.timeout_ms(TIMEOUT, 250).unwrap();
        let replica = Replica::new(addr, peers, store, state_machine);
        // Fire up the thread.
        thread::Builder::new().name(format!("RaftNode {}", addr)).spawn(move || {
            let mut raft_node = RaftNode {
                listener: listener,
                replica: replica,
            };
            event_loop.run(&mut raft_node).unwrap();
        }).unwrap();
    }
}

impl<S, M> Handler for RaftNode<S, M> where S: Store, M: StateMachine {

    type Message = ();
    type Timeout = Token;

    /// A registered IoHandle has available data to read
    fn readable(&mut self, _reactor: &mut Reactor<S, M>, token: Token, hint: ReadHint) {
        let mut message = MallocMessageBuilder::new_default();
        // TODO: Determine the stream we got a message on?
        if token == SOCKET && hint == ReadHint::data() {
            let (mut stream, from) = self.listener.accept().unwrap();
            let message_reader = serialize_packed::new_reader_unbuffered(&stream, ReaderOptions::new())
                .unwrap();
            if let Ok(request) = message_reader.get_root::<rpc_request::Reader>() {
                // We will be responding.
                match request.which().unwrap() {
                    // TODO: Move these into replica?
                    rpc_request::Which::AppendEntries(Ok(call)) => {
                        let res = message.init_root::<append_entries_response::Builder>();
                        self.replica.append_entries_request(from, call, res);
                    },
                    rpc_request::Which::RequestVote(Ok(call)) => {
                        let res = message.init_root::<request_vote_response::Builder>();
                        self.replica.request_vote_request(from, call, res);
                    },
                    _ => unimplemented!(),
                }
                serialize_packed::write_packed_message_unbuffered(&mut stream, &mut message).unwrap();
            } else if let Ok(response) = message_reader.get_root::<rpc_response::Reader>() {
                // We won't be responding. This is already a response.
                match response.which().unwrap() {
                    rpc_response::Which::AppendEntries(Ok(call)) => {
                        self.replica.append_entries_response(from, call);
                    },
                    rpc_response::Which::RequestVote(Ok(call)) => {
                        let res = message.init_root::<append_entries_request::Builder>();
                        self.replica.request_vote_response(from, call, res);
                        // TODO: send the AppendEntries requests if necessary
                        // TODO: Can we just reset the timer?
                        unimplemented!();
                    },
                    _ => unimplemented!(),
                }
            } else if let Ok(client_req) = message_reader.get_root::<client_request::Reader>() {
                let mut should_die = false;
                // We will be responding.
                match client_req.which().unwrap() {
                    client_request::Which::Append(Ok(call)) => {
                        let mut res = message.init_root::<client_response::Builder>();
                        self.replica.client_append(from, call, res)
                    },
                    client_request::Which::Die(Ok(call)) => {
                        should_die = true;
                        let mut res = message.init_root::<client_response::Builder>();
                        res.set_success(());
                        debug!("Got a Die request from Client({}). Reason: {}", from, call);
                    },
                    client_request::Which::LeaderRefresh(()) => {
                        let mut res = message.init_root::<client_response::Builder>();
                        self.replica.client_leader_refresh(from, res)
                    },
                    _ => unimplemented!(),
                }
                serialize_packed::write_packed_message_unbuffered(&mut stream, &mut message);
                // Do this here so that we can send the response.
                if should_die {
                    panic!("Got a Die request.");
                }
            } else {
                // It's something we don't understand.
                unimplemented!();
            }
        } else {
            // TODO Log this?
            unimplemented!();
        }
    }

    /// A registered timer has expired
    fn timeout(&mut self, reactor: &mut Reactor<S, M>, token: Token) {
        let mut message = MallocMessageBuilder::new_default();
        let request = message.init_root::<rpc_request::Builder>();
        let (timeout, _send_message) = self.replica.timeout(request.init_request_vote());
        reactor.timeout_ms(TIMEOUT, timeout).unwrap();
        // TODO: send messages if necessary
    }
}
<commit_msg>Per-connection events. Capn Proto impl is makeshift.<commit_after>use std::thread;
use std::collections::HashSet;
use std::net::SocketAddr;
use std::ops::Deref;
// MIO
use mio::tcp::{listen, TcpListener, TcpStream};
use mio::util::Slab;
use mio::Socket;
use mio::buf::{ByteBuf, MutByteBuf, SliceBuf};
use mio::{Interest, PollOpt, NonBlock, Token, EventLoop, Handler, ReadHint};

// Data structures.
use store::Store;
use replica::Replica;
use state_machine::StateMachine;

// Cap'n Proto
use capnp::serialize_packed;
use capnp::{MessageBuilder, MessageReader, ReaderOptions, MallocMessageBuilder};
use messages_capnp::{
    rpc_request,
    rpc_response,
    client_request,
    client_response,
    request_vote_response,
    append_entries_response,
    append_entries_request,
};
use super::RaftError;

// MIO Tokens
const TIMEOUT: Token = Token(0);
const LISTENER:  Token = Token(1);

/// The Raft Distributed Consensus Algorithm requires two RPC calls to be available:
///
///   * `append_entries` which is used as both a heartbeat (with no payload) and the primary
///     interface for requests.
///   * `request_vote` which is used by candidates during campaigns to obtain a vote.
///
/// A `RaftNode` acts as a replicated state machine. The server's role in the cluster depends on it's
/// own status. It will maintain both volatile state (which can be safely lost) and persistent
/// state (which must be carefully stored and kept safe).
///
/// Currently, the `RaftNode` API is not well defined. **We are looking for feedback and suggestions.**
pub struct RaftNode<S, M> where S: Store, M: StateMachine {
    replica: Replica<S, M>,
    // Channels and Sockets
    listener: NonBlock<TcpListener>,
    connections: Slab<RaftConnection>,
}

type Reactor<S, M> = EventLoop<RaftNode<S, M>>;

/// The implementation of the RaftNode. In most use cases, creating a `RaftNode` should just be
/// done via `::new()`.
impl<S, M> RaftNode<S, M> where S: Store, M: StateMachine {

    /// Creates a new Raft node with the cluster members specified.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address of the new node.
    /// * `peers` - The address of all peers in the Raft cluster.
    /// * `store` - The persitent log store.
    /// * `state_machine` - The client state machine to which client commands will be applied.
    pub fn spawn(addr: SocketAddr,
                 peers: HashSet<SocketAddr>,
                 store: S,
                 state_machine: M) {
        // Create an event loop
        let mut event_loop = Reactor::<S, M>::new().unwrap();
        // Setup the socket, make it not block.
        let listener = listen(&addr).unwrap();
        listener.set_reuseaddr(true);
        event_loop.register(&listener, LISTENER).unwrap();
        event_loop.timeout_ms(TIMEOUT, 250).unwrap();
        let replica = Replica::new(addr, peers, store, state_machine);
        // Fire up the thread.
        thread::Builder::new().name(format!("RaftNode {}", addr)).spawn(move || {
            let mut raft_node = RaftNode {
                listener: listener,
                replica: replica,
                connections: Slab::new_starting_at(Token(2), 128),
            };
            event_loop.run(&mut raft_node).unwrap();
        }).unwrap();
    }
}

impl<S, M> Handler for RaftNode<S, M> where S: Store, M: StateMachine {

    type Message = ();
    type Timeout = Token;

    /// A registered IoHandle has available writing space.
    fn writable(&mut self, reactor: &mut Reactor<S, M>, token: Token) {
        match token {
            TIMEOUT => unreachable!(),
            LISTENER => unreachable!(),
            tok => {
                self.connections[tok].writable(reactor, &mut self.replica);
            }
        }
    }

    /// A registered IoHandle has available data to read
    fn readable(&mut self, reactor: &mut Reactor<S, M>, token: Token, hint: ReadHint) {
        match token {
            TIMEOUT => unreachable!(),
            LISTENER => {
                let stream = self.listener.accept().unwrap().unwrap(); // Result<Option<_>,_>
                let conn = RaftConnection::new(stream);
                let tok = self.connections.insert(conn)
                    .ok().expect("Could not add connection to slab.");

                // Register the connection
                self.connections[tok].token = tok;
                reactor.register_opt(&self.connections[tok].stream, tok, Interest::readable(), PollOpt::edge() | PollOpt::oneshot())
                    .ok().expect("Could not register socket with event loop.");
            },
            tok => {
                self.connections[tok].readable(reactor, &mut self.replica);
            }
        }
    }

    /// A registered timer has expired
    fn timeout(&mut self, reactor: &mut Reactor<S, M>, token: Token) {
        let mut message = MallocMessageBuilder::new_default();
        let request = message.init_root::<rpc_request::Builder>();
        let (timeout, _send_message) = self.replica.timeout(request.init_request_vote());
        reactor.timeout_ms(TIMEOUT, timeout).unwrap();
        // TODO: send messages if necessary
    }
}

struct RaftConnection {
    stream: NonBlock<TcpStream>,
    token: Token,
    interest: Interest,
    current_read: Option<MutByteBuf>,
    current_write: Option<ByteBuf>,
    next_write: Option<ByteBuf>,
}

impl RaftConnection {
    fn new(sock: NonBlock<TcpStream>) -> RaftConnection {
        RaftConnection {
            stream: sock,
            token: Token(-1),
            interest: Interest::hup(),
            current_read: None,
            current_write: None, // TODO: Size??
            next_write: None,
        }
    }

    fn writable<S, M>(&mut self, event_loop: &mut EventLoop<RaftNode<S, M>>, replica: &mut Replica<S,M>)
                      -> Result<(), RaftError>
    where S: Store, M: StateMachine {
        unimplemented!();
    }

    fn readable<S, M>(&mut self, event_loop: &mut EventLoop<RaftNode<S, M>>, replica: &mut Replica<S,M>)
                      -> Result<(), RaftError>
    where S: Store, M: StateMachine {
        let mut message = MallocMessageBuilder::new_default();
        let from = self.stream.peer_addr().unwrap();
        let message_reader = serialize_packed::new_reader_unbuffered(&mut self.stream.deref(), ReaderOptions::new())
            .unwrap();
        if let Ok(request) = message_reader.get_root::<rpc_request::Reader>() {
            // We will be responding.
            match request.which().unwrap() {
                // TODO: Move these into replica?
                rpc_request::Which::AppendEntries(Ok(call)) => {
                    let res = message.init_root::<append_entries_response::Builder>();
                    replica.append_entries_request(from, call, res);
                },
                rpc_request::Which::RequestVote(Ok(call)) => {
                    let res = message.init_root::<request_vote_response::Builder>();
                    replica.request_vote_request(from, call, res);
                },
                _ => unimplemented!(),
            }
            serialize_packed::write_packed_message_unbuffered(&mut self.stream.deref(), &mut message).unwrap();
            Ok(())
        } else if let Ok(response) = message_reader.get_root::<rpc_response::Reader>() {
            // We won't be responding. This is already a response.
            match response.which().unwrap() {
                rpc_response::Which::AppendEntries(Ok(call)) => {
                    replica.append_entries_response(from, call);
                    Ok(())
                },
                rpc_response::Which::RequestVote(Ok(call)) => {
                    let res = message.init_root::<append_entries_request::Builder>();
                    replica.request_vote_response(from, call, res);
                    // TODO: send the AppendEntries requests if necessary
                    // TODO: Can we just reset the timer?
                    unimplemented!();
                },
                _ => unimplemented!(),
            }
        } else if let Ok(client_req) = message_reader.get_root::<client_request::Reader>() {
            let mut should_die = false;
            // We will be responding.
            match client_req.which().unwrap() {
                client_request::Which::Append(Ok(call)) => {
                    let mut res = message.init_root::<client_response::Builder>();
                    replica.client_append(from, call, res)
                },
                client_request::Which::Die(Ok(call)) => {
                    should_die = true;
                    let mut res = message.init_root::<client_response::Builder>();
                    res.set_success(());
                    debug!("Got a Die request from Client({}). Reason: {}", from, call);
                },
                client_request::Which::LeaderRefresh(()) => {
                    let mut res = message.init_root::<client_response::Builder>();
                    replica.client_leader_refresh(from, res)
                },
                _ => unimplemented!(),
            }
            serialize_packed::write_packed_message_unbuffered(&mut self.stream.deref(), &mut message);
            // Do this here so that we can send the response.
            if should_die {
                panic!("Got a Die request.");
            }
            Ok(())
        } else {
            // It's something we don't understand.
            unimplemented!();
        }
    }
}

// Old
// let mut message = MallocMessageBuilder::new_default();
// // TODO: Determine the stream we got a message on?
// if token == SOCKET && hint == ReadHint::data() {
//
// } else {
//     // TODO Log this?
//     unimplemented!();
// }
<|endoftext|>


<commit_before>//! Tool to remove PCR duplicates from UMI-tagged reads.
//!
//! This tool takes two FASTQ files (forward and reverse)
//! and returns two FASTQ files in which all PCR duplicates
//! have been merged into a consensus read.
//! Duplicates are identified by a Unique Molecular Identifier (UMI).
//!
//! ## Requirements:
//!
//!  - starcode
//!
//!
//! ## Usage:
//!
//! ```bash
//! $ rbt call-consensus-reads \
//!   <Path to FASTQ file with forward reads> \
//!   <Path to FASTQ file with reverse reads> \
//!   <Path for output forward FASTQ file> \
//!   <Path for output reverse FASTQ file> \
//!   -l <Length of UMI sequence> \
//!   -D <Maximum distance between sequences in a cluster> \  # See step 1 below
//!   -d <Maximum distance between UMIs in a cluster> \  # See step 2 below
//!   --umi-on-reverse  # if the UMIs are part of the reverse reads
//! ```
//!
//! ## Assumptions:
//!
//!  - Reads are of equal length
//!  - UMI is the prefix of the reads
//!
//! ## Workflow:
//!
//! The main steps are:
//!
//! 1. Preparation
//!    1. Remove UMI sequence from read (and save it for later use).
//!    2. Concatenate forward and reverse sequence.
//!        ```text
//!        Forward Read: [================]
//!        Reverse Read: [(UMI)-----------]
//!        Sequence for clustering in step 3: [================-----------]
//!        ```
//!
//! 2. Cluster all reads by their UMIs using starcode.
//!    Each cluster generated in this step contains reads with similar UMIs.
//!    However, all PCR duplicates of a read are within one cluster, since they
//!    share a UMI sequence.
//!    The size of these clusters highly depends on the length of the used UMI.
//!
//! 2. For each cluster from step two:
//!    1. Cluster reads by their concatenated sequences (without UMI) using starcode.
//!    2. Each new cluster contains reads that have a similar UMI (from step 2)
//!       as well as similar sequences. Consequently, these sets of reads are
//!       likely to be PCR duplicates of each other.
//!
//! 3. For each cluster from step three: Compute a consensus sequence.
//!
//!    At each position in the read, all bases and quality values are used
//!    to compute the base with Maximum a-posteriori probability (MAP).
//!
//!      1. For one position, compute the likelihood for the four alleles
//!         A, C, G, and T, incorporating the number of bases as well as
//!         their quality values.
//!      2. Choose the allele with the largest likelihood for the consensus read.
//!      3. Compute the quality value of the consensus read from the maximum posterior
//!         probability used to select the allele.
//!
//! 4. Write consensus reads to output file.
//!
//!
//!
// Since this is a binary crate, documentation needs to be compiled with this 'ancient incantation':
// https://github.com/rust-lang/cargo/issues/1865#issuecomment-394179125
mod calc_consensus;
mod pipeline;

use bio::io::fastq;
use flate2::bufread::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use pipeline::{CallConsensusReads, CallNonOverlappingConsensusRead, CallOverlappingConsensusRead};
use std::error::Error;
use std::fs;
use std::io::BufReader;
use std::str;

/// Build readers for the given input and output FASTQ files and pass them to
/// `call_consensus_reads`.
///
/// The type of the readers (writers) depends on the file ending.
/// If the input file names end with '.gz' a gzipped reader (writer) is used.
pub fn call_consensus_reads_from_paths(
    fq1: &str,
    fq2: &str,
    fq1_out: &str,
    fq2_out: &str,
    fq3_out: Option<&str>,
    umi_len: usize,
    seq_dist: usize,
    umi_dist: usize,
    reverse_umi: bool,
    insert_size: Option<usize>,
    std_dev: Option<usize>,
) -> Result<(), Box<dyn Error>> {
    match fq3_out {
        None => {
            eprintln!("Reading input files:\n    {}\n    {}", fq1, fq2);
            eprintln!("Writing output to:\n    {}\n    {}", fq1_out, fq2_out);
            match (fq1.ends_with(".gz"), fq2.ends_with(".gz"), fq1_out.ends_with(".gz"), fq2_out.ends_with(".gz")) {
                (false, false, false, false) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, false, false) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                (false, false, true, true) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, true, true) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                _ => panic!("Invalid combination of files. Each pair of files (input and output) need to be both gzipped or both not zipped.")
            }
        }
        Some(fq3_path) => {
            eprintln!("Reading input files:\n    {}\n    {}", fq1, fq2);
            eprintln!(
                "Writing output to:\n    {}\n    {}\n    {}",
                fq1_out, fq2_out, fq3_path
            );
            match (fq1.ends_with(".gz"), fq2.ends_with(".gz"), fq1_out.ends_with(".gz"), fq2_out.ends_with(".gz"), fq3_path.ends_with(".gz")) {
                (false, false, false, false, false) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    &mut fastq::Writer::to_file(fq3_path)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, false, false, false) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    &mut fastq::Writer::to_file(fq3_path)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                (false, false, true, true, true) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq3_path)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, true, true, true) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(GzDecoder::new)?),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq3_path)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                _ => panic!("Invalid combination of files. Each pair of files (input and output) need to be both gzipped or both not zipped.")
            }
        }
    }
}
<commit_msg>Replaced GzDecoder with MultiGzDecoder to add support for BGZF zipped files.<commit_after>//! Tool to remove PCR duplicates from UMI-tagged reads.
//!
//! This tool takes two FASTQ files (forward and reverse)
//! and returns two FASTQ files in which all PCR duplicates
//! have been merged into a consensus read.
//! Duplicates are identified by a Unique Molecular Identifier (UMI).
//!
//! ## Requirements:
//!
//!  - starcode
//!
//!
//! ## Usage:
//!
//! ```bash
//! $ rbt call-consensus-reads \
//!   <Path to FASTQ file with forward reads> \
//!   <Path to FASTQ file with reverse reads> \
//!   <Path for output forward FASTQ file> \
//!   <Path for output reverse FASTQ file> \
//!   -l <Length of UMI sequence> \
//!   -D <Maximum distance between sequences in a cluster> \  # See step 1 below
//!   -d <Maximum distance between UMIs in a cluster> \  # See step 2 below
//!   --umi-on-reverse  # if the UMIs are part of the reverse reads
//! ```
//!
//! ## Assumptions:
//!
//!  - Reads are of equal length
//!  - UMI is the prefix of the reads
//!
//! ## Workflow:
//!
//! The main steps are:
//!
//! 1. Preparation
//!    1. Remove UMI sequence from read (and save it for later use).
//!    2. Concatenate forward and reverse sequence.
//!        ```text
//!        Forward Read: [================]
//!        Reverse Read: [(UMI)-----------]
//!        Sequence for clustering in step 3: [================-----------]
//!        ```
//!
//! 2. Cluster all reads by their UMIs using starcode.
//!    Each cluster generated in this step contains reads with similar UMIs.
//!    However, all PCR duplicates of a read are within one cluster, since they
//!    share a UMI sequence.
//!    The size of these clusters highly depends on the length of the used UMI.
//!
//! 2. For each cluster from step two:
//!    1. Cluster reads by their concatenated sequences (without UMI) using starcode.
//!    2. Each new cluster contains reads that have a similar UMI (from step 2)
//!       as well as similar sequences. Consequently, these sets of reads are
//!       likely to be PCR duplicates of each other.
//!
//! 3. For each cluster from step three: Compute a consensus sequence.
//!
//!    At each position in the read, all bases and quality values are used
//!    to compute the base with Maximum a-posteriori probability (MAP).
//!
//!      1. For one position, compute the likelihood for the four alleles
//!         A, C, G, and T, incorporating the number of bases as well as
//!         their quality values.
//!      2. Choose the allele with the largest likelihood for the consensus read.
//!      3. Compute the quality value of the consensus read from the maximum posterior
//!         probability used to select the allele.
//!
//! 4. Write consensus reads to output file.
//!
//!
//!
// Since this is a binary crate, documentation needs to be compiled with this 'ancient incantation':
// https://github.com/rust-lang/cargo/issues/1865#issuecomment-394179125
mod calc_consensus;
mod pipeline;

use bio::io::fastq;
use flate2::bufread::MultiGzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use pipeline::{CallConsensusReads, CallNonOverlappingConsensusRead, CallOverlappingConsensusRead};
use std::error::Error;
use std::fs;
use std::io::BufReader;
use std::str;

/// Build readers for the given input and output FASTQ files and pass them to
/// `call_consensus_reads`.
///
/// The type of the readers (writers) depends on the file ending.
/// If the input file names end with '.gz' a gzipped reader (writer) is used.
pub fn call_consensus_reads_from_paths(
    fq1: &str,
    fq2: &str,
    fq1_out: &str,
    fq2_out: &str,
    fq3_out: Option<&str>,
    umi_len: usize,
    seq_dist: usize,
    umi_dist: usize,
    reverse_umi: bool,
    insert_size: Option<usize>,
    std_dev: Option<usize>,
) -> Result<(), Box<dyn Error>> {
    match fq3_out {
        None => {
            eprintln!("Reading input files:\n    {}\n    {}", fq1, fq2);
            eprintln!("Writing output to:\n    {}\n    {}", fq1_out, fq2_out);
            match (fq1.ends_with(".gz"), fq2.ends_with(".gz"), fq1_out.ends_with(".gz"), fq2_out.ends_with(".gz")) {
                (false, false, false, false) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, false, false) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                (false, false, true, true) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, true, true) => CallNonOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    reverse_umi,
                ).call_consensus_reads(),
                _ => panic!("Invalid combination of files. Each pair of files (input and output) need to be both gzipped or both not zipped.")
            }
        }
        Some(fq3_path) => {
            eprintln!("Reading input files:\n    {}\n    {}", fq1, fq2);
            eprintln!(
                "Writing output to:\n    {}\n    {}\n    {}",
                fq1_out, fq2_out, fq3_path
            );
            match (fq1.ends_with(".gz"), fq2.ends_with(".gz"), fq1_out.ends_with(".gz"), fq2_out.ends_with(".gz"), fq3_path.ends_with(".gz")) {
                (false, false, false, false, false) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    &mut fastq::Writer::to_file(fq3_path)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, false, false, false) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Writer::to_file(fq1_out)?,
                    &mut fastq::Writer::to_file(fq2_out)?,
                    &mut fastq::Writer::to_file(fq3_path)?,
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                (false, false, true, true, true) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::from_file(fq1)?,
                    &mut fastq::Reader::from_file(fq2)?,
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq3_path)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                (true, true, true, true, true) => CallOverlappingConsensusRead::new(
                    &mut fastq::Reader::new(fs::File::open(fq1).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Reader::new(fs::File::open(fq2).map(BufReader::new).map(MultiGzDecoder::new)?),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq1_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq2_out)?, Compression::default())),
                    &mut fastq::Writer::new(GzEncoder::new(fs::File::create(fq3_path)?, Compression::default())),
                    umi_len,
                    seq_dist,
                    umi_dist,
                    insert_size.unwrap(),
                    std_dev.unwrap(),
                    reverse_umi,
                ).call_consensus_reads(),
                _ => panic!("Invalid combination of files. Each pair of files (input and output) need to be both gzipped or both not zipped.")
            }
        }
    }
}
<|endoftext|>


<commit_before>use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use std::io::{stderr, Read, Write, BufRead, BufReader};
use std::sync::mpsc::SyncSender;
use super::Splitter;

pub struct NulSplitter {
    tx: SyncSender<Vec<u8>>,
    decoder: Box<Decoder>,
    encoder: Box<Encoder>
}

impl<T: Read> Splitter<T> for NulSplitter {
    fn run(&self, buf_reader: BufReader<T>) {
        for line in buf_reader.split(0) {
            let line = match line {
                Err(_) => {
                    let _ = writeln!(stderr(), "Invalid UTF-8 input");
                    continue;
                }
                Ok(line) => match String::from_utf8(line) {
                    Err(_) => {
                        let _ = writeln!(stderr(), "Invalid UTF-8 input");
                        continue;
                    }
                    Ok(line) => line
                }
            };
            if let Err(e) = handle_line(&line, &self.tx, &self.decoder, &self.encoder) {
                let _ = writeln!(stderr(), "{}: [{}]", e, line.trim());
            }
        }
    }
}

impl NulSplitter {
    pub fn new(tx: SyncSender<Vec<u8>>, decoder: Box<Decoder>, encoder: Box<Encoder>) -> NulSplitter {
        NulSplitter {
            tx: tx,
            decoder: decoder,
            encoder: encoder
        }
    }
}

fn handle_line(line: &String, tx: &SyncSender<Vec<u8>>, decoder: &Box<Decoder>, encoder: &Box<Encoder>) -> Result<(), &'static str> {
    let decoded = try!(decoder.decode(&line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
<commit_msg>Zero-allocations NUL splitter<commit_after>use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use std::io::{stderr, Read, Write, BufRead, BufReader};
use std::str;
use std::sync::mpsc::SyncSender;
use super::Splitter;

pub struct NulSplitter {
    tx: SyncSender<Vec<u8>>,
    decoder: Box<Decoder>,
    encoder: Box<Encoder>
}

impl<T: Read> Splitter<T> for NulSplitter {
    fn run(&self, buf_reader: BufReader<T>) {
        for line in buf_reader.split(0) {
            let line = match line {
                Err(_) => {
                    let _ = writeln!(stderr(), "EOF?");
                    continue;
                }
                Ok(line) => line
            };
            let line = match str::from_utf8(&line) {
                Err(_) => {
                    let _ = writeln!(stderr(), "Invalid UTF-8 input");
                    continue;
                }
                Ok(line) => line
            };
            if let Err(e) = handle_line(line, &self.tx, &self.decoder, &self.encoder) {
                let _ = writeln!(stderr(), "{}: [{}]", e, line.trim());
            }
        }
    }
}

impl NulSplitter {
    pub fn new(tx: SyncSender<Vec<u8>>, decoder: Box<Decoder>, encoder: Box<Encoder>) -> NulSplitter {
        NulSplitter {
            tx: tx,
            decoder: decoder,
            encoder: encoder
        }
    }
}

fn handle_line(line: &str, tx: &SyncSender<Vec<u8>>, decoder: &Box<Decoder>, encoder: &Box<Encoder>) -> Result<(), &'static str> {
    let decoded = try!(decoder.decode(&line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! Simple file-locking apis for each OS.
//!
//! This is not meant to be in the standard library, it does nothing with
//! green/native threading. This is just a bare-bones enough solution for
//! librustdoc, it is not production quality at all.

#![allow(non_camel_case_types)]
use std::path::Path;

pub use self::imp::Lock;

#[cfg(unix)]
mod imp {
    use std::ffi::{CString, OsStr};
    use std::os::unix::prelude::*;
    use std::path::Path;
    use std::io;
    use libc;

    #[cfg(target_os = "linux")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 0;
        pub const F_WRLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_SETLK: libc::c_int = 6;
        pub const F_SETLKW: libc::c_int = 7;
    }

    #[cfg(target_os = "freebsd")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_WRLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 12;
        pub const F_SETLKW: libc::c_int = 13;
    }

    #[cfg(any(target_os = "dragonfly",
              target_os = "bitrig",
              target_os = "netbsd",
              target_os = "openbsd"))]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_WRLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 8;
        pub const F_SETLKW: libc::c_int = 9;
    }

    #[cfg(target_os = "haiku")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_UNLCK: libc::c_short = 0x0200;
        pub const F_WRLCK: libc::c_short = 0x0400;
        pub const F_SETLK: libc::c_int = 0x0080;
        pub const F_SETLKW: libc::c_int = 0x0100;
    }

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_WRLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 8;
        pub const F_SETLKW: libc::c_int = 9;
    }

    #[cfg(target_os = "solaris")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_sysid: libc::c_int,
            pub l_pid: libc::pid_t,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_WRLCK: libc::c_short = 2;
        pub const F_UNLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 6;
        pub const F_SETLKW: libc::c_int = 7;
    }

    #[derive(Debug)]
    pub struct Lock {
        fd: libc::c_int,
    }

    impl Lock {
        pub fn new(p: &Path,
                   wait: bool,
                   create: bool,
                   exclusive: bool)
                   -> io::Result<Lock> {
            let os: &OsStr = p.as_ref();
            let buf = CString::new(os.as_bytes()).unwrap();
            let open_flags = if create {
                libc::O_RDWR | libc::O_CREAT
            } else {
                libc::O_RDWR
            };

            let fd = unsafe {
                libc::open(buf.as_ptr(), open_flags,
                           libc::S_IRWXU as libc::c_int)
            };

            if fd < 0 {
                return Err(io::Error::last_os_error());
            }

            let lock_type = if exclusive {
                os::F_WRLCK
            } else {
                os::F_RDLCK
            };

            let flock = os::flock {
                l_start: 0,
                l_len: 0,
                l_pid: 0,
                l_whence: libc::SEEK_SET as libc::c_short,
                l_type: lock_type,
                l_sysid: 0,
            };
            let cmd = if wait { os::F_SETLKW } else { os::F_SETLK };
            let ret = unsafe {
                libc::fcntl(fd, cmd, &flock)
            };
            if ret == -1 {
                let err = io::Error::last_os_error();
                unsafe { libc::close(fd); }
                Err(err)
            } else {
                Ok(Lock { fd: fd })
            }
        }
    }

    impl Drop for Lock {
        fn drop(&mut self) {
            let flock = os::flock {
                l_start: 0,
                l_len: 0,
                l_pid: 0,
                l_whence: libc::SEEK_SET as libc::c_short,
                l_type: os::F_UNLCK,
                l_sysid: 0,
            };
            unsafe {
                libc::fcntl(self.fd, os::F_SETLK, &flock);
                libc::close(self.fd);
            }
        }
    }
}

#[cfg(windows)]
#[allow(bad_style)]
mod imp {
    use std::io;
    use std::mem;
    use std::os::windows::prelude::*;
    use std::os::windows::raw::HANDLE;
    use std::path::Path;
    use std::fs::{File, OpenOptions};
    use std::os::raw::{c_ulong, c_ulonglong, c_int};

    type DWORD = c_ulong;
    type BOOL = c_int;
    type ULONG_PTR = c_ulonglong;

    type LPOVERLAPPED = *mut OVERLAPPED;
    const LOCKFILE_EXCLUSIVE_LOCK: DWORD = 0x00000002;
    const LOCKFILE_FAIL_IMMEDIATELY: DWORD = 0x00000001;

    const FILE_SHARE_DELETE: DWORD = 0x4;
    const FILE_SHARE_READ: DWORD = 0x1;
    const FILE_SHARE_WRITE: DWORD = 0x2;

    #[repr(C)]
    struct OVERLAPPED {
        Internal: ULONG_PTR,
        InternalHigh: ULONG_PTR,
        Offset: DWORD,
        OffsetHigh: DWORD,
        hEvent: HANDLE,
    }

    extern "system" {
        fn LockFileEx(hFile: HANDLE,
                      dwFlags: DWORD,
                      dwReserved: DWORD,
                      nNumberOfBytesToLockLow: DWORD,
                      nNumberOfBytesToLockHigh: DWORD,
                      lpOverlapped: LPOVERLAPPED) -> BOOL;
    }

    #[derive(Debug)]
    pub struct Lock {
        _file: File,
    }

    impl Lock {
        pub fn new(p: &Path,
                   wait: bool,
                   create: bool,
                   exclusive: bool)
                   -> io::Result<Lock> {
            assert!(p.parent().unwrap().exists(),
                "Parent directory of lock-file must exist: {}",
                p.display());

            let share_mode = FILE_SHARE_DELETE | FILE_SHARE_READ | FILE_SHARE_WRITE;

            let mut open_options = OpenOptions::new();
            open_options.read(true)
                        .share_mode(share_mode);

            if create {
                open_options.create(true)
                            .write(true);
            }

            debug!("Attempting to open lock file `{}`", p.display());
            let file = match open_options.open(p) {
                Ok(file) => {
                    debug!("Lock file opened successfully");
                    file
                }
                Err(err) => {
                    debug!("Error opening lock file: {}", err);
                    return Err(err)
                }
            };

            let ret = unsafe {
                let mut overlapped: OVERLAPPED = mem::zeroed();

                let mut dwFlags = 0;
                if !wait {
                    dwFlags |= LOCKFILE_FAIL_IMMEDIATELY;
                }

                if exclusive {
                    dwFlags |= LOCKFILE_EXCLUSIVE_LOCK;
                }

                debug!("Attempting to acquire lock on lock file `{}`",
                       p.display());
                LockFileEx(file.as_raw_handle(),
                           dwFlags,
                           0,
                           0xFFFF_FFFF,
                           0xFFFF_FFFF,
                           &mut overlapped)
            };
            if ret == 0 {
                let err = io::Error::last_os_error();
                debug!("Failed acquiring file lock: {}", err);
                Err(err)
            } else {
                debug!("Successfully acquired lock.");
                Ok(Lock { _file: file })
            }
        }
    }

    // Note that we don't need a Drop impl on the Windows: The file is unlocked
    // automatically when it's closed.
}

impl imp::Lock {
    pub fn panicking_new(p: &Path,
                         wait: bool,
                         create: bool,
                         exclusive: bool)
                         -> Lock {
        Lock::new(p, wait, create, exclusive).unwrap_or_else(|err| {
            panic!("could not lock `{}`: {}", p.display(), err);
        })
    }
}
<commit_msg>Rollup merge of #39724 - malbarbo:android, r=alexcrichton<commit_after>// Copyright 2014-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! Simple file-locking apis for each OS.
//!
//! This is not meant to be in the standard library, it does nothing with
//! green/native threading. This is just a bare-bones enough solution for
//! librustdoc, it is not production quality at all.

#![allow(non_camel_case_types)]
use std::path::Path;

pub use self::imp::Lock;

#[cfg(unix)]
mod imp {
    use std::ffi::{CString, OsStr};
    use std::os::unix::prelude::*;
    use std::path::Path;
    use std::io;
    use libc;

    #[cfg(any(target_os = "linux", target_os = "android"))]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 0;
        pub const F_WRLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_SETLK: libc::c_int = 6;
        pub const F_SETLKW: libc::c_int = 7;
    }

    #[cfg(target_os = "freebsd")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_WRLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 12;
        pub const F_SETLKW: libc::c_int = 13;
    }

    #[cfg(any(target_os = "dragonfly",
              target_os = "bitrig",
              target_os = "netbsd",
              target_os = "openbsd"))]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_WRLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 8;
        pub const F_SETLKW: libc::c_int = 9;
    }

    #[cfg(target_os = "haiku")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_UNLCK: libc::c_short = 0x0200;
        pub const F_WRLCK: libc::c_short = 0x0400;
        pub const F_SETLK: libc::c_int = 0x0080;
        pub const F_SETLKW: libc::c_int = 0x0100;
    }

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_pid: libc::pid_t,
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,

            // not actually here, but brings in line with freebsd
            pub l_sysid: libc::c_int,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_UNLCK: libc::c_short = 2;
        pub const F_WRLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 8;
        pub const F_SETLKW: libc::c_int = 9;
    }

    #[cfg(target_os = "solaris")]
    mod os {
        use libc;

        #[repr(C)]
        pub struct flock {
            pub l_type: libc::c_short,
            pub l_whence: libc::c_short,
            pub l_start: libc::off_t,
            pub l_len: libc::off_t,
            pub l_sysid: libc::c_int,
            pub l_pid: libc::pid_t,
        }

        pub const F_RDLCK: libc::c_short = 1;
        pub const F_WRLCK: libc::c_short = 2;
        pub const F_UNLCK: libc::c_short = 3;
        pub const F_SETLK: libc::c_int = 6;
        pub const F_SETLKW: libc::c_int = 7;
    }

    #[derive(Debug)]
    pub struct Lock {
        fd: libc::c_int,
    }

    impl Lock {
        pub fn new(p: &Path,
                   wait: bool,
                   create: bool,
                   exclusive: bool)
                   -> io::Result<Lock> {
            let os: &OsStr = p.as_ref();
            let buf = CString::new(os.as_bytes()).unwrap();
            let open_flags = if create {
                libc::O_RDWR | libc::O_CREAT
            } else {
                libc::O_RDWR
            };

            let fd = unsafe {
                libc::open(buf.as_ptr(), open_flags,
                           libc::S_IRWXU as libc::c_int)
            };

            if fd < 0 {
                return Err(io::Error::last_os_error());
            }

            let lock_type = if exclusive {
                os::F_WRLCK
            } else {
                os::F_RDLCK
            };

            let flock = os::flock {
                l_start: 0,
                l_len: 0,
                l_pid: 0,
                l_whence: libc::SEEK_SET as libc::c_short,
                l_type: lock_type,
                l_sysid: 0,
            };
            let cmd = if wait { os::F_SETLKW } else { os::F_SETLK };
            let ret = unsafe {
                libc::fcntl(fd, cmd, &flock)
            };
            if ret == -1 {
                let err = io::Error::last_os_error();
                unsafe { libc::close(fd); }
                Err(err)
            } else {
                Ok(Lock { fd: fd })
            }
        }
    }

    impl Drop for Lock {
        fn drop(&mut self) {
            let flock = os::flock {
                l_start: 0,
                l_len: 0,
                l_pid: 0,
                l_whence: libc::SEEK_SET as libc::c_short,
                l_type: os::F_UNLCK,
                l_sysid: 0,
            };
            unsafe {
                libc::fcntl(self.fd, os::F_SETLK, &flock);
                libc::close(self.fd);
            }
        }
    }
}

#[cfg(windows)]
#[allow(bad_style)]
mod imp {
    use std::io;
    use std::mem;
    use std::os::windows::prelude::*;
    use std::os::windows::raw::HANDLE;
    use std::path::Path;
    use std::fs::{File, OpenOptions};
    use std::os::raw::{c_ulong, c_ulonglong, c_int};

    type DWORD = c_ulong;
    type BOOL = c_int;
    type ULONG_PTR = c_ulonglong;

    type LPOVERLAPPED = *mut OVERLAPPED;
    const LOCKFILE_EXCLUSIVE_LOCK: DWORD = 0x00000002;
    const LOCKFILE_FAIL_IMMEDIATELY: DWORD = 0x00000001;

    const FILE_SHARE_DELETE: DWORD = 0x4;
    const FILE_SHARE_READ: DWORD = 0x1;
    const FILE_SHARE_WRITE: DWORD = 0x2;

    #[repr(C)]
    struct OVERLAPPED {
        Internal: ULONG_PTR,
        InternalHigh: ULONG_PTR,
        Offset: DWORD,
        OffsetHigh: DWORD,
        hEvent: HANDLE,
    }

    extern "system" {
        fn LockFileEx(hFile: HANDLE,
                      dwFlags: DWORD,
                      dwReserved: DWORD,
                      nNumberOfBytesToLockLow: DWORD,
                      nNumberOfBytesToLockHigh: DWORD,
                      lpOverlapped: LPOVERLAPPED) -> BOOL;
    }

    #[derive(Debug)]
    pub struct Lock {
        _file: File,
    }

    impl Lock {
        pub fn new(p: &Path,
                   wait: bool,
                   create: bool,
                   exclusive: bool)
                   -> io::Result<Lock> {
            assert!(p.parent().unwrap().exists(),
                "Parent directory of lock-file must exist: {}",
                p.display());

            let share_mode = FILE_SHARE_DELETE | FILE_SHARE_READ | FILE_SHARE_WRITE;

            let mut open_options = OpenOptions::new();
            open_options.read(true)
                        .share_mode(share_mode);

            if create {
                open_options.create(true)
                            .write(true);
            }

            debug!("Attempting to open lock file `{}`", p.display());
            let file = match open_options.open(p) {
                Ok(file) => {
                    debug!("Lock file opened successfully");
                    file
                }
                Err(err) => {
                    debug!("Error opening lock file: {}", err);
                    return Err(err)
                }
            };

            let ret = unsafe {
                let mut overlapped: OVERLAPPED = mem::zeroed();

                let mut dwFlags = 0;
                if !wait {
                    dwFlags |= LOCKFILE_FAIL_IMMEDIATELY;
                }

                if exclusive {
                    dwFlags |= LOCKFILE_EXCLUSIVE_LOCK;
                }

                debug!("Attempting to acquire lock on lock file `{}`",
                       p.display());
                LockFileEx(file.as_raw_handle(),
                           dwFlags,
                           0,
                           0xFFFF_FFFF,
                           0xFFFF_FFFF,
                           &mut overlapped)
            };
            if ret == 0 {
                let err = io::Error::last_os_error();
                debug!("Failed acquiring file lock: {}", err);
                Err(err)
            } else {
                debug!("Successfully acquired lock.");
                Ok(Lock { _file: file })
            }
        }
    }

    // Note that we don't need a Drop impl on the Windows: The file is unlocked
    // automatically when it's closed.
}

impl imp::Lock {
    pub fn panicking_new(p: &Path,
                         wait: bool,
                         create: bool,
                         exclusive: bool)
                         -> Lock {
        Lock::new(p, wait, create, exclusive).unwrap_or_else(|err| {
            panic!("could not lock `{}`: {}", p.display(), err);
        })
    }
}
<|endoftext|>