<commit_before>mod extractor;

extern crate num_cpus;

use clap;
use flate2::write::GzEncoder;
use rayon::prelude::*;
use std::fs;
use std::io::{BufRead, BufWriter};
use std::path::{Path, PathBuf};
use tree_sitter::{Language, Parser, Range};

enum TrapCompression {
    None,
    Gzip,
}

impl TrapCompression {
    fn from_env() -> TrapCompression {
        match std::env::var("CODEQL_RUBY_TRAP_COMPRESSION") {
            Ok(method) => match TrapCompression::from_string(&method) {
                Some(c) => c,
                None => {
                    tracing::error!("Unknown compression method '{}'; using gzip.", &method);
                    TrapCompression::Gzip
                }
            },
            // Default compression method if the env var isn't set:
            Err(_) => TrapCompression::Gzip,
        }
    }

    fn from_string(s: &str) -> Option<TrapCompression> {
        match s.to_lowercase().as_ref() {
            "none" => Some(TrapCompression::None),
            "gzip" => Some(TrapCompression::Gzip),
            _ => None,
        }
    }

    fn extension(&self) -> &str {
        match self {
            TrapCompression::None => ".trap",
            TrapCompression::Gzip => ".trap.gz",
        }
    }
}

/**
 * Gets the number of threads the extractor should use, by reading the
 * CODEQL_THREADS environment variable and using it as described in the
 * extractor spec:
 *
 * "If the number is positive, it indicates the number of threads that should
 * be used. If the number is negative or zero, it should be added to the number
 * of cores available on the machine to determine how many threads to use
 * (minimum of 1). If unspecified, should be considered as set to -1."
 */
fn num_codeql_threads() -> usize {
    let threads_str = std::env::var("CODEQL_THREADS").unwrap_or("-1".to_owned());
    match threads_str.parse::<i32>() {
        Ok(num) if num <= 0 => {
            let reduction = -num as usize;
            std::cmp::max(1, num_cpus::get() - reduction)
        }
        Ok(num) => num as usize,

        Err(_) => {
            tracing::error!(
                "Unable to parse CODEQL_THREADS value '{}'; defaulting to 1 thread.",
                &threads_str
            );
            1
        }
    }
}

fn main() -> std::io::Result<()> {
    tracing_subscriber::fmt()
        .with_target(false)
        .without_time()
        .with_level(true)
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let num_threads = num_codeql_threads();
    tracing::info!(
        "Using {} {}",
        num_threads,
        if num_threads == 1 {
            "thread"
        } else {
            "threads"
        }
    );
    rayon::ThreadPoolBuilder::new()
        .num_threads(num_threads)
        .build_global()
        .unwrap();

    let matches = clap::App::new("Ruby extractor")
        .version("1.0")
        .author("GitHub")
        .about("CodeQL Ruby extractor")
        .args_from_usage(
            "--source-archive-dir=<DIR> 'Sets a custom source archive folder'
                    --output-dir=<DIR>         'Sets a custom trap folder'
                    --file-list=<FILE_LIST>    'A text files containing the paths of the files to extract'",
        )
        .get_matches();
    let src_archive_dir = matches
        .value_of("source-archive-dir")
        .expect("missing --source-archive-dir");
    let src_archive_dir = PathBuf::from(src_archive_dir);

    let trap_dir = matches
        .value_of("output-dir")
        .expect("missing --output-dir");
    let trap_dir = PathBuf::from(trap_dir);
    let trap_compression = TrapCompression::from_env();

    let file_list = matches.value_of("file-list").expect("missing --file-list");
    let file_list = fs::File::open(file_list)?;

    let language = tree_sitter_ruby::language();
    let erb = tree_sitter_embedded_template::language();
    let schema = node_types::read_node_types_str("ruby", tree_sitter_ruby::NODE_TYPES)?;
    let erb_schema =
        node_types::read_node_types_str("erb", tree_sitter_embedded_template::NODE_TYPES)?;
    let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
    let lines = lines?;
    lines
        .par_iter()
        .try_for_each(|line| {
            let path = PathBuf::from(line).canonicalize()?;
            let src_archive_file = path_for(&src_archive_dir, &path, "");
            let mut source = std::fs::read(&path)?;
            let code_ranges;
            let mut trap_writer = extractor::new_trap_writer();
            if path.extension().map_or(false, |x| x == "erb") {
                tracing::info!("scanning: {}", path.display());
                extractor::extract(
                    erb,
                    "erb",
                    &erb_schema,
                    &mut trap_writer,
                    &path,
                    &source,
                    &[],
                )?;

                let (ranges, line_breaks) = scan_erb(erb, &source);
                for i in line_breaks {
                    if i < source.len() {
                        source[i] = b'\n';
                    }
                }
                code_ranges = ranges;
            } else {
                code_ranges = vec![];
            }
            extractor::extract(
                language,
                "ruby",
                &schema,
                &mut trap_writer,
                &path,
                &source,
                &code_ranges,
            )?;
            std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
            std::fs::copy(&path, &src_archive_file)?;
            write_trap(&trap_dir, path, trap_writer, &trap_compression)
        })
        .expect("failed to extract files");

    let path = PathBuf::from("extras");
    let mut trap_writer = extractor::new_trap_writer();
    trap_writer.populate_empty_location();
    write_trap(&trap_dir, path, trap_writer, &trap_compression)
}

fn write_trap(
    trap_dir: &PathBuf,
    path: PathBuf,
    trap_writer: extractor::TrapWriter,
    trap_compression: &TrapCompression,
) -> std::io::Result<()> {
    let trap_file = path_for(&trap_dir, &path, trap_compression.extension());
    std::fs::create_dir_all(&trap_file.parent().unwrap())?;
    let trap_file = std::fs::File::create(&trap_file)?;
    let mut trap_file = BufWriter::new(trap_file);
    match trap_compression {
        TrapCompression::None => trap_writer.output(&mut trap_file),
        TrapCompression::Gzip => {
            let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
            trap_writer.output(&mut compressed_writer)
        }
    }
}

fn scan_erb(erb: Language, source: &std::vec::Vec<u8>) -> (Vec<Range>, Vec<usize>) {
    let mut parser = Parser::new();
    parser.set_language(erb).unwrap();
    let tree = parser.parse(&source, None).expect("Failed to parse file");
    let mut result = Vec::new();
    let mut line_breaks = vec![];

    for n in tree.root_node().children(&mut tree.walk()) {
        let kind = n.kind();
        if kind == "directive" || kind == "output_directive" {
            for c in n.children(&mut tree.walk()) {
                if c.kind() == "code" {
                    let mut range = c.range();
                    if range.end_byte < source.len() {
                        line_breaks.push(range.end_byte);
                        range.end_byte += 1;
                        range.end_point.column += 1;
                    }
                    result.push(range);
                }
            }
        }
    }
    if result.len() == 0 {
        let root = tree.root_node();
        // Add an empty range at the end of the file
        result.push(Range {
            start_byte: root.end_byte(),
            end_byte: root.end_byte(),
            start_point: root.end_position(),
            end_point: root.end_position(),
        });
    }
    (result, line_breaks)
}

fn path_for(dir: &Path, path: &Path, ext: &str) -> PathBuf {
    let mut result = PathBuf::from(dir);
    for component in path.components() {
        match component {
            std::path::Component::Prefix(prefix) => match prefix.kind() {
                std::path::Prefix::Disk(letter) | std::path::Prefix::VerbatimDisk(letter) => {
                    result.push(format!("{}_", letter as char))
                }
                std::path::Prefix::Verbatim(x) | std::path::Prefix::DeviceNS(x) => {
                    result.push(x);
                }
                std::path::Prefix::UNC(server, share)
                | std::path::Prefix::VerbatimUNC(server, share) => {
                    result.push("unc");
                    result.push(server);
                    result.push(share);
                }
            },
            std::path::Component::RootDir => {
                // skip
            }
            std::path::Component::Normal(_) => {
                result.push(component);
            }
            std::path::Component::CurDir => {
                // skip
            }
            std::path::Component::ParentDir => {
                result.pop();
            }
        }
    }
    if let Some(x) = result.extension() {
        let mut new_ext = x.to_os_string();
        new_ext.push(ext);
        result.set_extension(new_ext);
    } else {
        result.set_extension(ext);
    }
    result
}
<commit_msg>Fix trap extension for source files without extensions<commit_after>mod extractor;

extern crate num_cpus;

use clap;
use flate2::write::GzEncoder;
use rayon::prelude::*;
use std::fs;
use std::io::{BufRead, BufWriter};
use std::path::{Path, PathBuf};
use tree_sitter::{Language, Parser, Range};

enum TrapCompression {
    None,
    Gzip,
}

impl TrapCompression {
    fn from_env() -> TrapCompression {
        match std::env::var("CODEQL_RUBY_TRAP_COMPRESSION") {
            Ok(method) => match TrapCompression::from_string(&method) {
                Some(c) => c,
                None => {
                    tracing::error!("Unknown compression method '{}'; using gzip.", &method);
                    TrapCompression::Gzip
                }
            },
            // Default compression method if the env var isn't set:
            Err(_) => TrapCompression::Gzip,
        }
    }

    fn from_string(s: &str) -> Option<TrapCompression> {
        match s.to_lowercase().as_ref() {
            "none" => Some(TrapCompression::None),
            "gzip" => Some(TrapCompression::Gzip),
            _ => None,
        }
    }

    fn extension(&self) -> &str {
        match self {
            TrapCompression::None => "trap",
            TrapCompression::Gzip => "trap.gz",
        }
    }
}

/**
 * Gets the number of threads the extractor should use, by reading the
 * CODEQL_THREADS environment variable and using it as described in the
 * extractor spec:
 *
 * "If the number is positive, it indicates the number of threads that should
 * be used. If the number is negative or zero, it should be added to the number
 * of cores available on the machine to determine how many threads to use
 * (minimum of 1). If unspecified, should be considered as set to -1."
 */
fn num_codeql_threads() -> usize {
    let threads_str = std::env::var("CODEQL_THREADS").unwrap_or("-1".to_owned());
    match threads_str.parse::<i32>() {
        Ok(num) if num <= 0 => {
            let reduction = -num as usize;
            std::cmp::max(1, num_cpus::get() - reduction)
        }
        Ok(num) => num as usize,

        Err(_) => {
            tracing::error!(
                "Unable to parse CODEQL_THREADS value '{}'; defaulting to 1 thread.",
                &threads_str
            );
            1
        }
    }
}

fn main() -> std::io::Result<()> {
    tracing_subscriber::fmt()
        .with_target(false)
        .without_time()
        .with_level(true)
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let num_threads = num_codeql_threads();
    tracing::info!(
        "Using {} {}",
        num_threads,
        if num_threads == 1 {
            "thread"
        } else {
            "threads"
        }
    );
    rayon::ThreadPoolBuilder::new()
        .num_threads(num_threads)
        .build_global()
        .unwrap();

    let matches = clap::App::new("Ruby extractor")
        .version("1.0")
        .author("GitHub")
        .about("CodeQL Ruby extractor")
        .args_from_usage(
            "--source-archive-dir=<DIR> 'Sets a custom source archive folder'
                    --output-dir=<DIR>         'Sets a custom trap folder'
                    --file-list=<FILE_LIST>    'A text files containing the paths of the files to extract'",
        )
        .get_matches();
    let src_archive_dir = matches
        .value_of("source-archive-dir")
        .expect("missing --source-archive-dir");
    let src_archive_dir = PathBuf::from(src_archive_dir);

    let trap_dir = matches
        .value_of("output-dir")
        .expect("missing --output-dir");
    let trap_dir = PathBuf::from(trap_dir);
    let trap_compression = TrapCompression::from_env();

    let file_list = matches.value_of("file-list").expect("missing --file-list");
    let file_list = fs::File::open(file_list)?;

    let language = tree_sitter_ruby::language();
    let erb = tree_sitter_embedded_template::language();
    let schema = node_types::read_node_types_str("ruby", tree_sitter_ruby::NODE_TYPES)?;
    let erb_schema =
        node_types::read_node_types_str("erb", tree_sitter_embedded_template::NODE_TYPES)?;
    let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
    let lines = lines?;
    lines
        .par_iter()
        .try_for_each(|line| {
            let path = PathBuf::from(line).canonicalize()?;
            let src_archive_file = path_for(&src_archive_dir, &path, "");
            let mut source = std::fs::read(&path)?;
            let code_ranges;
            let mut trap_writer = extractor::new_trap_writer();
            if path.extension().map_or(false, |x| x == "erb") {
                tracing::info!("scanning: {}", path.display());
                extractor::extract(
                    erb,
                    "erb",
                    &erb_schema,
                    &mut trap_writer,
                    &path,
                    &source,
                    &[],
                )?;

                let (ranges, line_breaks) = scan_erb(erb, &source);
                for i in line_breaks {
                    if i < source.len() {
                        source[i] = b'\n';
                    }
                }
                code_ranges = ranges;
            } else {
                code_ranges = vec![];
            }
            extractor::extract(
                language,
                "ruby",
                &schema,
                &mut trap_writer,
                &path,
                &source,
                &code_ranges,
            )?;
            std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
            std::fs::copy(&path, &src_archive_file)?;
            write_trap(&trap_dir, path, trap_writer, &trap_compression)
        })
        .expect("failed to extract files");

    let path = PathBuf::from("extras");
    let mut trap_writer = extractor::new_trap_writer();
    trap_writer.populate_empty_location();
    write_trap(&trap_dir, path, trap_writer, &trap_compression)
}

fn write_trap(
    trap_dir: &PathBuf,
    path: PathBuf,
    trap_writer: extractor::TrapWriter,
    trap_compression: &TrapCompression,
) -> std::io::Result<()> {
    let trap_file = path_for(&trap_dir, &path, trap_compression.extension());
    std::fs::create_dir_all(&trap_file.parent().unwrap())?;
    let trap_file = std::fs::File::create(&trap_file)?;
    let mut trap_file = BufWriter::new(trap_file);
    match trap_compression {
        TrapCompression::None => trap_writer.output(&mut trap_file),
        TrapCompression::Gzip => {
            let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
            trap_writer.output(&mut compressed_writer)
        }
    }
}

fn scan_erb(erb: Language, source: &std::vec::Vec<u8>) -> (Vec<Range>, Vec<usize>) {
    let mut parser = Parser::new();
    parser.set_language(erb).unwrap();
    let tree = parser.parse(&source, None).expect("Failed to parse file");
    let mut result = Vec::new();
    let mut line_breaks = vec![];

    for n in tree.root_node().children(&mut tree.walk()) {
        let kind = n.kind();
        if kind == "directive" || kind == "output_directive" {
            for c in n.children(&mut tree.walk()) {
                if c.kind() == "code" {
                    let mut range = c.range();
                    if range.end_byte < source.len() {
                        line_breaks.push(range.end_byte);
                        range.end_byte += 1;
                        range.end_point.column += 1;
                    }
                    result.push(range);
                }
            }
        }
    }
    if result.len() == 0 {
        let root = tree.root_node();
        // Add an empty range at the end of the file
        result.push(Range {
            start_byte: root.end_byte(),
            end_byte: root.end_byte(),
            start_point: root.end_position(),
            end_point: root.end_position(),
        });
    }
    (result, line_breaks)
}

fn path_for(dir: &Path, path: &Path, ext: &str) -> PathBuf {
    let mut result = PathBuf::from(dir);
    for component in path.components() {
        match component {
            std::path::Component::Prefix(prefix) => match prefix.kind() {
                std::path::Prefix::Disk(letter) | std::path::Prefix::VerbatimDisk(letter) => {
                    result.push(format!("{}_", letter as char))
                }
                std::path::Prefix::Verbatim(x) | std::path::Prefix::DeviceNS(x) => {
                    result.push(x);
                }
                std::path::Prefix::UNC(server, share)
                | std::path::Prefix::VerbatimUNC(server, share) => {
                    result.push("unc");
                    result.push(server);
                    result.push(share);
                }
            },
            std::path::Component::RootDir => {
                // skip
            }
            std::path::Component::Normal(_) => {
                result.push(component);
            }
            std::path::Component::CurDir => {
                // skip
            }
            std::path::Component::ParentDir => {
                result.pop();
            }
        }
    }
    if let Some(x) = result.extension() {
        let mut new_ext = x.to_os_string();
        new_ext.push(".");
        new_ext.push(ext);
        result.set_extension(new_ext);
    } else {
        result.set_extension(ext);
    }
    result
}
<|endoftext|>


<commit_before>#![feature(semaphore)]
#![feature(plugin)]
#![plugin(clippy)]

extern crate irc;
extern crate rustc_serialize;
extern crate toml;
extern crate byteorder;

mod playpen;
mod cfg;
mod util;
mod eval;

use irc::client::prelude::*;
use std::sync::{Arc, Mutex, Semaphore};
use std::thread;
use std::collections::{VecDeque, HashMap};

use cfg::{EvalbotCfg, LangCfg};
use eval::{script, persistent};

macro_rules! send_msg {
    ($conn:expr, $dest:expr, $line:expr) => {
        if let Err(x) = $conn.send_privmsg($dest, $line) {
            println!("failed to send message: {}", x)
        }
    };
}

macro_rules! send_notice {
    ($conn:expr, $dest:expr, $line:expr) => {
        if let Err(x) = $conn.send_notice($dest, $line) {
            println!("failed to send notice: {}", x)
        }
    };
}

#[derive(Debug)]
struct Req {
    pub is_channel: bool,
    pub sender: String,
    pub target: String,
    pub code: String,
    pub language: Arc<LangCfg>
}

#[derive(Debug)]
enum CommandResult {
    Req(Req),
    Rehash,
    RestartEvaluator(String)
}

#[derive(Debug)]
enum WorkerMessage {
    Req(Req),
    Terminate
}

#[derive(Clone)]
struct State {
    requests: Arc<Mutex<VecDeque<WorkerMessage>>>,
    has_work: Arc<Semaphore>,
    cfg: Arc<EvalbotCfg>,
    languages: Arc<HashMap<String, Arc<LangCfg>>>,
    evaluators: Arc<HashMap<String, eval::persistent::Evaluator>>
}

impl State {
    fn send_message(&self, x: WorkerMessage) {
        self.requests.lock().unwrap().push_back(x); // if mutex is poisoned, just bail
        self.has_work.release();
    }
}

fn evaluate_loop<'a, S, T, U>(conn: Arc<S>, state: State)
    where T: IrcRead,
          U: IrcWrite,
          S: ServerExt<'a, T, U> + Sized {
    let has_work = state.has_work;
    let cfg = state.cfg;
    let requests = state.requests;
    let playpen_args = cfg.playpen_args.iter().map(|s| &**s).collect::<Vec<&str>>();
    loop {
        let work;
        {
            has_work.acquire();
            let mut rvec = requests.lock().unwrap();
            work = rvec.pop_front();
            std::mem::drop(rvec);
        }

        if let Some(WorkerMessage::Req(work)) = work {
            let result = if work.language.persistent {
                state.evaluators
                     .get(&work.language.long_name)
                     .unwrap()
                     .eval(&work.code, cfg.playpen_timeout)
            } else {
                script::eval(&work.code,
                             &*work.language,
                             &playpen_args[..],
                             &cfg.sandbox_dir,
                             cfg.playpen_timeout)
            };

            let (result, err) = match result {
                Ok(x) => (x, false),
                Err(x) => (x, true),
            };

            let dest = if work.is_channel { &work.target } else { &work.sender };

            if result.is_empty() {
                send_msg!(conn, dest, "ok (no output)");
                continue;
            }

            let result = util::wrap_and_trim_output(&result,
                                           if work.is_channel {
                                               cfg.max_channel_line_len
                                           } else {
                                               425
                                           });

            let max_lines = if work.is_channel {
                cfg.max_channel_lines
            } else {
                cfg.max_priv_lines
            };
            let (truncated, result) = util::truncate_output(result, max_lines);
            let iter = result.iter().take(if truncated {
            // truncated: we take up to the 2nd last line, as we add "... (output truncated)"
            // to the last line
                result.len() - 1
            } else {
                result.len()
            });
            let prefix = if work.is_channel && !err {
                &cfg.chan_output_prefix as &str
            } else {
                ""
            };
            for line in iter {
                let line = format!("{}{}", prefix, line);
                send_msg!(conn, dest, &line);
            }
            if truncated {
                send_msg!(conn,
                          dest,
                          &format!("{}{}... (output truncated)", prefix, result.last().unwrap()));
            }
        } else if let Some(WorkerMessage::Terminate) = work {
            println!("got rehashing message; terminating");
            break;
        }
    }
}

fn parse_command(message: &str, prefix: &str) -> Option<CommandResult> {
    let tok: Vec<&str> = match message.splitn(2, prefix)
                                      .nth(1)
                                      .map(|x| x.split(' ').collect()) {
        Some(x) => x,
        None => return None,
    };
    if tok.is_empty() {
        return None;
    }
    match tok[0] {
        "rehash" => Some(CommandResult::Rehash),
        "restart" => {
            if tok.len() < 2 {
                None
            } else {
                Some(CommandResult::RestartEvaluator(tok[1].to_owned()))
            }
        }
        _ => None,
    }
}

fn parse_msg(message: &Message, state: &State) -> Option<(CommandResult, String)> {
    let sender = message.get_source_nickname().unwrap_or("");
    if let Ok(Command::PRIVMSG(target, message)) = message.into() {
        let in_channel = target.starts_with('#');
        if !in_channel && message.contains('\x01') {
            // we don't accept code via CTCP
            return None;
        }

        if state.cfg.owners.iter().any(|x| *x == sender) {
            if let Some(cmd) = parse_command(&message, &state.cfg.command_prefix) {
                return Some((cmd, sender.to_owned()));
            }
        }

        let tok: Vec<&str> = message.trim().splitn(2, '>').collect();
        match tok.len() {
            2 => (),
            _ => return None,
        };

        let language = state.languages.get(tok[0]);
        if let Some(language) = language {
            Some((CommandResult::Req(Req {
                is_channel: in_channel,
                sender: sender.to_owned(),
                target: target.to_owned(),
                code: tok[1].to_owned(),
                language: language.clone()
            }),
                  sender.to_owned()))
        } else {
            None
        }
    } else {
        None
    }
}

fn main() {
    let irc_config = match cfg::read("evalbot.irc.toml") {
        Ok(x) => x,
        Err(x) => panic!("could not read irc config; {}", x),
    };
    println!("read irc config: {:?}", irc_config);

    let conn = Arc::new(IrcServer::from_config(irc_config).unwrap());
    let mut rehash_cfg = None;
    loop {
        conn.identify().unwrap();
        'connection: loop {
            let config = {
                if let Some(cfg) = rehash_cfg.take() {
                    Arc::new(cfg)
                } else {
                    Arc::new(match cfg::read::<EvalbotCfg>("evalbot.toml") {
                        Ok(x) => x,
                        Err(x) => panic!("could not read config; {}", x),
                    })
                }
            };
            println!("read config: {:?}", config);
            let evalreqs = Arc::new(Mutex::new(VecDeque::new()));
            let has_work = Arc::new(Semaphore::new(0));

            let mut languages = HashMap::new();
            let mut evaluators = HashMap::new();
            for lang in config.languages.clone() {
                if lang.persistent {
                    let sandbox = config.sandbox_dir.clone();
                    let binary_path = lang.binary_path.clone();
                    let syscalls_path = lang.syscalls_path.clone();
                    let binary_args = lang.binary_args.clone();
                    let playpen_args = config.playpen_args.clone();
                    let childfn = move || {
                        playpen::spawn(&sandbox,
                                       &binary_path,
                                       &syscalls_path,
                                       &playpen_args.iter().map(|s| &**s).collect::<Vec<&str>>()[..],
                                       &binary_args.iter().map(|s| &**s).collect::<Vec<&str>>()[..],
                                       None,
                                       false)
                            .unwrap()
                    };
                    evaluators.insert(lang.long_name.clone(), persistent::new(childfn));
                }
                let lang = Arc::new(lang);
                languages.insert(lang.short_name.clone(), lang.clone());
                languages.insert(lang.long_name.clone(), lang);
            }
            let languages = Arc::new(languages);
            let evaluators = Arc::new(evaluators);

            let state = State {
                cfg: config,
                requests: evalreqs,
                has_work: has_work,
                languages: languages,
                evaluators: evaluators
            };

            for _ in 0..state.cfg.eval_threads {
                let conn = conn.clone();
                let state = state.clone();
                thread::spawn(move || {
                    evaluate_loop(conn, state);
                });
            }

            for maybe_msg in conn.iter() {
                let msg = match maybe_msg {
                    Ok(x) => x,
                    Err(x) => {
                        println!("{}, reconnecting", x);
                        break 'connection;
                    }
                };

                let req = parse_msg(&msg, &state);
                if let Some((x, sender)) = req {
                    println!("{:?}", (&x, &sender));
                    match x {
                        CommandResult::Req(req) => state.send_message(WorkerMessage::Req(req)),
                        CommandResult::Rehash => {
                            for _ in 0..state.cfg.eval_threads {
                                state.send_message(WorkerMessage::Terminate);
                            }
                            for evaluator in state.evaluators.values() {
                                evaluator.terminate();
                            }
                            let maybe_rehash_cfg = cfg::read::<EvalbotCfg>("evalbot.toml");
                            match maybe_rehash_cfg {
                                Ok(cfg) => {
                                    rehash_cfg = Some(cfg);
                                    send_notice!(conn, &sender, "read configuration file OK, rehashing");
                                    continue 'connection;
                                },
                                Err(x) => {
                                    send_notice!(conn, &sender, "error reading configuration, check console");
                                    println!("error while reading config for rehash: {}", x);
                                }
                            };
                        }
                        CommandResult::RestartEvaluator(lang) => {
                            if let Some(lang) = state.languages.get(&lang) {
                                send_notice!(conn,
                                             &sender,
                                             match state.evaluators.get(&lang.long_name) {
                                                 Some(x) => {
                                                     x.restart();
                                                     "restarting evaluator"
                                                 }
                                                 None => "this language's not persistent",
                                             });
                            } else {
                                send_notice!(conn, &sender, "invalid language name");
                            }
                        }
                    };
                }
            }
            break 'connection;
        }
        conn.reconnect().unwrap();
    }
}
<commit_msg>If persistent child spawning fails, keep trying<commit_after>#![feature(semaphore)]
#![feature(plugin)]
#![plugin(clippy)]

extern crate irc;
extern crate rustc_serialize;
extern crate toml;
extern crate byteorder;

mod playpen;
mod cfg;
mod util;
mod eval;

use irc::client::prelude::*;
use std::sync::{Arc, Mutex, Semaphore};
use std::thread;
use std::time::Duration;
use std::collections::{VecDeque, HashMap};

use cfg::{EvalbotCfg, LangCfg};
use eval::{script, persistent};

macro_rules! send_msg {
    ($conn:expr, $dest:expr, $line:expr) => {
        if let Err(x) = $conn.send_privmsg($dest, $line) {
            println!("failed to send message: {}", x)
        }
    };
}

macro_rules! send_notice {
    ($conn:expr, $dest:expr, $line:expr) => {
        if let Err(x) = $conn.send_notice($dest, $line) {
            println!("failed to send notice: {}", x)
        }
    };
}

#[derive(Debug)]
struct Req {
    pub is_channel: bool,
    pub sender: String,
    pub target: String,
    pub code: String,
    pub language: Arc<LangCfg>
}

#[derive(Debug)]
enum CommandResult {
    Req(Req),
    Rehash,
    RestartEvaluator(String)
}

#[derive(Debug)]
enum WorkerMessage {
    Req(Req),
    Terminate
}

#[derive(Clone)]
struct State {
    requests: Arc<Mutex<VecDeque<WorkerMessage>>>,
    has_work: Arc<Semaphore>,
    cfg: Arc<EvalbotCfg>,
    languages: Arc<HashMap<String, Arc<LangCfg>>>,
    evaluators: Arc<HashMap<String, eval::persistent::Evaluator>>
}

impl State {
    fn send_message(&self, x: WorkerMessage) {
        self.requests.lock().unwrap().push_back(x); // if mutex is poisoned, just bail
        self.has_work.release();
    }
}

fn evaluate_loop<'a, S, T, U>(conn: Arc<S>, state: State)
    where T: IrcRead,
          U: IrcWrite,
          S: ServerExt<'a, T, U> + Sized {
    let has_work = state.has_work;
    let cfg = state.cfg;
    let requests = state.requests;
    let playpen_args = cfg.playpen_args.iter().map(|s| &**s).collect::<Vec<&str>>();
    loop {
        let work;
        {
            has_work.acquire();
            let mut rvec = requests.lock().unwrap();
            work = rvec.pop_front();
            std::mem::drop(rvec);
        }

        if let Some(WorkerMessage::Req(work)) = work {
            let result = if work.language.persistent {
                state.evaluators
                     .get(&work.language.long_name)
                     .unwrap()
                     .eval(&work.code, cfg.playpen_timeout)
            } else {
                script::eval(&work.code,
                             &*work.language,
                             &playpen_args[..],
                             &cfg.sandbox_dir,
                             cfg.playpen_timeout)
            };

            let (result, err) = match result {
                Ok(x) => (x, false),
                Err(x) => (x, true),
            };

            let dest = if work.is_channel { &work.target } else { &work.sender };

            if result.is_empty() {
                send_msg!(conn, dest, "ok (no output)");
                continue;
            }

            let result = util::wrap_and_trim_output(&result,
                                           if work.is_channel {
                                               cfg.max_channel_line_len
                                           } else {
                                               425
                                           });

            let max_lines = if work.is_channel {
                cfg.max_channel_lines
            } else {
                cfg.max_priv_lines
            };
            let (truncated, result) = util::truncate_output(result, max_lines);
            let iter = result.iter().take(if truncated {
            // truncated: we take up to the 2nd last line, as we add "... (output truncated)"
            // to the last line
                result.len() - 1
            } else {
                result.len()
            });
            let prefix = if work.is_channel && !err {
                &cfg.chan_output_prefix as &str
            } else {
                ""
            };
            for line in iter {
                let line = format!("{}{}", prefix, line);
                send_msg!(conn, dest, &line);
            }
            if truncated {
                send_msg!(conn,
                          dest,
                          &format!("{}{}... (output truncated)", prefix, result.last().unwrap()));
            }
        } else if let Some(WorkerMessage::Terminate) = work {
            println!("got rehashing message; terminating");
            break;
        }
    }
}

fn parse_command(message: &str, prefix: &str) -> Option<CommandResult> {
    let tok: Vec<&str> = match message.splitn(2, prefix)
                                      .nth(1)
                                      .map(|x| x.split(' ').collect()) {
        Some(x) => x,
        None => return None,
    };
    if tok.is_empty() {
        return None;
    }
    match tok[0] {
        "rehash" => Some(CommandResult::Rehash),
        "restart" => {
            if tok.len() < 2 {
                None
            } else {
                Some(CommandResult::RestartEvaluator(tok[1].to_owned()))
            }
        }
        _ => None,
    }
}

fn parse_msg(message: &Message, state: &State) -> Option<(CommandResult, String)> {
    let sender = message.get_source_nickname().unwrap_or("");
    if let Ok(Command::PRIVMSG(target, message)) = message.into() {
        let in_channel = target.starts_with('#');
        if !in_channel && message.contains('\x01') {
            // we don't accept code via CTCP
            return None;
        }

        if state.cfg.owners.iter().any(|x| *x == sender) {
            if let Some(cmd) = parse_command(&message, &state.cfg.command_prefix) {
                return Some((cmd, sender.to_owned()));
            }
        }

        let tok: Vec<&str> = message.trim().splitn(2, '>').collect();
        match tok.len() {
            2 => (),
            _ => return None,
        };

        let language = state.languages.get(tok[0]);
        if let Some(language) = language {
            Some((CommandResult::Req(Req {
                is_channel: in_channel,
                sender: sender.to_owned(),
                target: target.to_owned(),
                code: tok[1].to_owned(),
                language: language.clone()
            }),
                  sender.to_owned()))
        } else {
            None
        }
    } else {
        None
    }
}

fn main() {
    let irc_config = match cfg::read("evalbot.irc.toml") {
        Ok(x) => x,
        Err(x) => panic!("could not read irc config; {}", x),
    };
    println!("read irc config: {:?}", irc_config);

    let conn = Arc::new(IrcServer::from_config(irc_config).unwrap());
    let mut rehash_cfg = None;
    loop {
        conn.identify().unwrap();
        'connection: loop {
            let config = {
                if let Some(cfg) = rehash_cfg.take() {
                    Arc::new(cfg)
                } else {
                    Arc::new(match cfg::read::<EvalbotCfg>("evalbot.toml") {
                        Ok(x) => x,
                        Err(x) => panic!("could not read config; {}", x),
                    })
                }
            };
            println!("read config: {:?}", config);
            let evalreqs = Arc::new(Mutex::new(VecDeque::new()));
            let has_work = Arc::new(Semaphore::new(0));

            let mut languages = HashMap::new();
            let mut evaluators = HashMap::new();
            for lang in config.languages.clone() {
                if lang.persistent {
                    let sandbox = config.sandbox_dir.clone();
                    let binary_path = lang.binary_path.clone();
                    let syscalls_path = lang.syscalls_path.clone();
                    let binary_args = lang.binary_args.clone();
                    let playpen_args = config.playpen_args.clone();
                    let childfn = move || {
                        loop {
                            if let Ok(x) = playpen::spawn(&sandbox,
                                                          &binary_path,
                                                          &syscalls_path,
                                                          &playpen_args.iter().map(|s| &**s).collect::<Vec<&str>>()[..],
                                                          &binary_args.iter().map(|s| &**s).collect::<Vec<&str>>()[..],
                                                          None,
                                                          false) {
                                return x;
                            } else {
                                thread::sleep(Duration::new(1, 0));
                            }
                        }
                    };
                    evaluators.insert(lang.long_name.clone(), persistent::new(childfn));
                }
                let lang = Arc::new(lang);
                languages.insert(lang.short_name.clone(), lang.clone());
                languages.insert(lang.long_name.clone(), lang);
            }
            let languages = Arc::new(languages);
            let evaluators = Arc::new(evaluators);

            let state = State {
                cfg: config,
                requests: evalreqs,
                has_work: has_work,
                languages: languages,
                evaluators: evaluators
            };

            for _ in 0..state.cfg.eval_threads {
                let conn = conn.clone();
                let state = state.clone();
                thread::spawn(move || {
                    evaluate_loop(conn, state);
                });
            }

            for maybe_msg in conn.iter() {
                let msg = match maybe_msg {
                    Ok(x) => x,
                    Err(x) => {
                        println!("{}, reconnecting", x);
                        break 'connection;
                    }
                };

                let req = parse_msg(&msg, &state);
                if let Some((x, sender)) = req {
                    println!("{:?}", (&x, &sender));
                    match x {
                        CommandResult::Req(req) => state.send_message(WorkerMessage::Req(req)),
                        CommandResult::Rehash => {
                            for _ in 0..state.cfg.eval_threads {
                                state.send_message(WorkerMessage::Terminate);
                            }
                            for evaluator in state.evaluators.values() {
                                evaluator.terminate();
                            }
                            let maybe_rehash_cfg = cfg::read::<EvalbotCfg>("evalbot.toml");
                            match maybe_rehash_cfg {
                                Ok(cfg) => {
                                    rehash_cfg = Some(cfg);
                                    send_notice!(conn, &sender, "read configuration file OK, rehashing");
                                    continue 'connection;
                                },
                                Err(x) => {
                                    send_notice!(conn, &sender, "error reading configuration, check console");
                                    println!("error while reading config for rehash: {}", x);
                                }
                            };
                        }
                        CommandResult::RestartEvaluator(lang) => {
                            if let Some(lang) = state.languages.get(&lang) {
                                send_notice!(conn,
                                             &sender,
                                             match state.evaluators.get(&lang.long_name) {
                                                 Some(x) => {
                                                     x.restart();
                                                     "restarting evaluator"
                                                 }
                                                 None => "this language's not persistent",
                                             });
                            } else {
                                send_notice!(conn, &sender, "invalid language name");
                            }
                        }
                    };
                }
            }
            break 'connection;
        }
        conn.reconnect().unwrap();
    }
}
<|endoftext|>


<commit_before>#![crate_id = "gossip"]
#![crate_type = "lib"]
#![feature(globs, macro_rules)]
#![allow(unused_variable,unused_imports,unused_must_use)]

/*!
Gossip protocol engine written in Rust.

WIP
*/

extern crate collections;
extern crate uuid;
extern crate rand;

pub mod server;
pub mod state;
pub mod util;
pub mod message;
pub mod broadcast;
pub mod protocol;
pub mod transport;
pub mod tcp;
pub mod connection;<commit_msg>Don't export the util module<commit_after>#![crate_id = "gossip"]
#![crate_type = "lib"]
#![feature(globs, macro_rules)]
#![allow(unused_variable,unused_imports,unused_must_use)]

/*!
Gossip protocol engine written in Rust.

WIP
*/

extern crate collections;
extern crate uuid;
extern crate rand;

pub mod server;
pub mod state;
mod util;
pub mod message;
pub mod broadcast;
pub mod protocol;
pub mod transport;
pub mod tcp;
pub mod connection;<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use std::fmt;

#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
pub enum EdgeKind {
    Uses,
    Owns,
}

impl EdgeKind {
    pub fn display_name(&self) -> &'static str {
        match self {
            Self::Uses => "uses",
            Self::Owns => "owns",
        }
    }
}

impl fmt::Display for EdgeKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::Uses => "Uses",
            Self::Owns => "Owns",
        };
        write!(f, "{}", name)
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Edge {
    pub kind: EdgeKind,
}
<commit_msg>Add methods for merging `Edge`s<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use std::fmt;

#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
pub enum EdgeKind {
    Uses,
    Owns,
}

impl EdgeKind {
    pub fn display_name(&self) -> &'static str {
        match self {
            Self::Uses => "uses",
            Self::Owns => "owns",
        }
    }
}

impl fmt::Display for EdgeKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::Uses => "Uses",
            Self::Owns => "Owns",
        };
        write!(f, "{}", name)
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Edge {
    pub kind: EdgeKind,
}

impl Edge {
    pub fn merge_with(&mut self, other: &Edge) {
        match (self.kind, other.kind) {
            (EdgeKind::Uses, EdgeKind::Uses) => {
                self.kind = EdgeKind::Uses;
            }
            (EdgeKind::Uses, EdgeKind::Owns)
            | (EdgeKind::Owns, EdgeKind::Uses)
            | (EdgeKind::Owns, EdgeKind::Owns) => {
                self.kind = EdgeKind::Owns;
            }
        }
    }

    pub fn merged_with(&self, other: &Edge) -> Self {
        let mut clone = self.clone();
        clone.merge_with(other);
        clone
    }
}
<|endoftext|>


<commit_before>// Copyright 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.

use std::sync::Mutex;

fn par_for<I, F>(iter: I, f: F)
    where I: Iterator,
          I::Item: Send,
          F: Fn(I::Item) + Sync
{
    for item in iter {
        f(item)
    }
}

fn sum(x: &[i32]) {
    let sum_lengths = Mutex::new(0);
    par_for(x.windows(4), |x| {
        *sum_lengths.lock().unwrap() += x.len()
    });

    assert_eq!(*sum_lengths.lock().unwrap(), (x.len() - 3) * 4);
}

fn main() {
    let mut elements = [0; 20];

    // iterators over references into this stack frame
    par_for(elements.iter_mut().enumerate(), |(i, x)| {
        *x = i as i32
    });

    sum(&elements)
}
<commit_msg>no mir for mutex::lock in windows-gnu<commit_after>// Copyright 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.

//ignore-windows-gnu

use std::sync::Mutex;

fn par_for<I, F>(iter: I, f: F)
    where I: Iterator,
          I::Item: Send,
          F: Fn(I::Item) + Sync
{
    for item in iter {
        f(item)
    }
}

fn sum(x: &[i32]) {
    let sum_lengths = Mutex::new(0);
    par_for(x.windows(4), |x| {
        *sum_lengths.lock().unwrap() += x.len()
    });

    assert_eq!(*sum_lengths.lock().unwrap(), (x.len() - 3) * 4);
}

fn main() {
    let mut elements = [0; 20];

    // iterators over references into this stack frame
    par_for(elements.iter_mut().enumerate(), |(i, x)| {
        *x = i as i32
    });

    sum(&elements)
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::marker;

use crate::core::Message;

use crate::descriptor::DescriptorProto;
use crate::descriptor::FileDescriptorProto;

use crate::reflect::accessor::FieldAccessor;
use crate::reflect::find_message_or_enum::find_message_or_enum;
use crate::reflect::find_message_or_enum::MessageOrEnum;
use crate::reflect::reflect_deep_eq::ReflectDeepEq;
use crate::reflect::FieldDescriptor;

trait MessageFactory: Send + Sync + 'static {
    fn new_instance(&self) -> Box<dyn Message>;
    fn default_instance(&self) -> &dyn Message;
    fn clone(&self, message: &dyn Message) -> Box<dyn Message>;
    fn eq(&self, a: &dyn Message, b: &dyn Message) -> bool;
}

struct MessageFactoryImpl<M>(marker::PhantomData<M>);

impl<M> MessageFactory for MessageFactoryImpl<M>
where
    M: 'static + Message + Default + Clone + PartialEq,
{
    fn new_instance(&self) -> Box<dyn Message> {
        let m: M = Default::default();
        Box::new(m)
    }

    fn default_instance(&self) -> &dyn Message {
        M::default_instance() as &dyn Message
    }

    fn clone(&self, message: &dyn Message) -> Box<dyn Message> {
        let m: &M = message.downcast_ref().expect("wrong message type");
        Box::new(m.clone())
    }

    fn eq(&self, a: &dyn Message, b: &dyn Message) -> bool {
        let a: &M = a.downcast_ref().expect("wrong message type");
        let b: &M = b.downcast_ref().expect("wrong message type");
        a == b
    }
}

/// Dynamic representation of message type.
///
/// Used for reflection.
pub struct MessageDescriptor {
    full_name: String,
    file_descriptor_proto: &'static FileDescriptorProto,
    proto: &'static DescriptorProto,
    factory: &'static dyn MessageFactory,
    fields: Vec<FieldDescriptor>,

    index_by_name: HashMap<String, usize>,
    index_by_name_or_json_name: HashMap<String, usize>,
    index_by_number: HashMap<u32, usize>,
}

impl MessageDescriptor {
    /// Get underlying `DescriptorProto` object.
    pub fn get_proto(&self) -> &DescriptorProto {
        self.proto
    }

    /// Get a message descriptor for given message type
    pub fn for_type<M: Message>() -> &'static MessageDescriptor {
        M::descriptor_static()
    }

    fn compute_full_name(package: &str, path_to_package: &str, proto: &DescriptorProto) -> String {
        let mut full_name = package.to_owned();
        if path_to_package.len() != 0 {
            if full_name.len() != 0 {
                full_name.push('.');
            }
            full_name.push_str(path_to_package);
        }
        if full_name.len() != 0 {
            full_name.push('.');
        }
        full_name.push_str(proto.get_name());
        full_name
    }

    // Non-generic part of `new` is a separate function
    // to reduce code bloat from multiple instantiations.
    fn new_non_generic(
        protobuf_name_to_package: &'static str,
        fields: Vec<FieldAccessor>,
        file_descriptor_proto: &'static FileDescriptorProto,
        factory: &'static dyn MessageFactory,
    ) -> MessageDescriptor {
        let (path_to_package, proto) =
            match find_message_or_enum(file_descriptor_proto, protobuf_name_to_package) {
                (path_to_package, MessageOrEnum::Message(m)) => (path_to_package, m),
                (_, MessageOrEnum::Enum(_)) => panic!("not a message"),
            };

        let mut field_proto_by_name = HashMap::new();
        for field_proto in &proto.field {
            field_proto_by_name.insert(field_proto.get_name(), field_proto);
        }

        let mut index_by_name = HashMap::new();
        let mut index_by_name_or_json_name = HashMap::new();
        let mut index_by_number = HashMap::new();

        let fields: Vec<_> = fields
            .into_iter()
            .map(|f| {
                let proto = *field_proto_by_name.get(f.name).unwrap();
                FieldDescriptor::new(f, proto)
            })
            .collect();

        for (i, f) in fields.iter().enumerate() {
            assert!(index_by_number
                .insert(f.proto().get_number() as u32, i)
                .is_none());
            assert!(index_by_name
                .insert(f.proto().get_name().to_owned(), i)
                .is_none());
            assert!(index_by_name_or_json_name
                .insert(f.proto().get_name().to_owned(), i)
                .is_none());

            let json_name = f.json_name().to_owned();

            if json_name != f.proto().get_name() {
                assert!(index_by_name_or_json_name.insert(json_name, i).is_none());
            }
        }
        MessageDescriptor {
            full_name: MessageDescriptor::compute_full_name(
                file_descriptor_proto.get_package(),
                &path_to_package,
                &proto,
            ),
            proto,
            factory,
            fields,
            index_by_name,
            index_by_name_or_json_name,
            index_by_number,
            file_descriptor_proto,
        }
    }

    /// Construct a new message descriptor.
    ///
    /// This operation is called from generated code and rarely
    /// need to be called directly.
    ///
    /// This function is not a part of public API.
    #[doc(hidden)]
    pub fn new<M: 'static + Message + Default + Clone + PartialEq>(
        protobuf_name_to_package: &'static str,
        fields: Vec<FieldAccessor>,
        file_descriptor_proto: &'static FileDescriptorProto,
    ) -> MessageDescriptor {
        let factory = &MessageFactoryImpl(marker::PhantomData::<M>);
        MessageDescriptor::new_non_generic(
            protobuf_name_to_package,
            fields,
            file_descriptor_proto,
            factory,
        )
    }

    /// `FileDescriptorProto` containg this message type
    pub fn file_descriptor_proto(&self) -> &FileDescriptorProto {
        self.file_descriptor_proto
    }

    /// New empty message
    pub fn new_instance(&self) -> Box<dyn Message> {
        self.factory.new_instance()
    }

    /// Shared immutable empty message
    pub fn default_instance(&self) -> &dyn Message {
        self.factory.default_instance()
    }

    /// Clone a message
    pub fn clone(&self, message: &dyn Message) -> Box<dyn Message> {
        self.factory.clone(message)
    }

    /// Check if two messages equal.
    ///
    /// # Panics
    ///
    /// Is any message has different type than this descriptor.
    pub fn eq(&self, a: &dyn Message, b: &dyn Message) -> bool {
        self.factory.eq(a, b)
    }

    /// Similar to `eq`, but considers `NaN` values equal.
    ///
    /// # Panics
    ///
    /// Is any message has different type than this descriptor.
    pub fn deep_eq(&self, a: &dyn Message, b: &dyn Message) -> bool {
        // Explicitly force panic even if field list is empty
        assert_eq!(
            self as *const MessageDescriptor,
            a.descriptor() as *const MessageDescriptor
        );
        assert_eq!(
            self as *const MessageDescriptor,
            b.descriptor() as *const MessageDescriptor
        );

        for field in self.fields() {
            let af = field.get_reflect(a);
            let bf = field.get_reflect(b);
            if !af.reflect_deep_eq(&bf) {
                return false;
            }
        }
        true
    }

    /// Message name as given in `.proto` file
    pub fn name(&self) -> &'static str {
        self.proto.get_name()
    }

    /// Fully qualified protobuf message name
    pub fn full_name(&self) -> &str {
        &self.full_name[..]
    }

    /// Message field descriptors.
    pub fn fields(&self) -> &[FieldDescriptor] {
        &self.fields
    }

    /// Find message field by protobuf field name
    ///
    /// Note: protobuf field name might be different for Rust field name.
    pub fn field_by_name(&self, name: &str) -> Option<&FieldDescriptor> {
        let &index = self.index_by_name.get(name)?;
        Some(&self.fields[index])
    }

    /// Find message field by field name or field JSON name
    pub fn field_by_name_or_json_name<'a>(&'a self, name: &str) -> Option<&'a FieldDescriptor> {
        let &index = self.index_by_name_or_json_name.get(name)?;
        Some(&self.fields[index])
    }

    /// Find message field by field name
    pub fn field_by_number<'a>(&'a self, number: u32) -> Option<&'a FieldDescriptor> {
        let &index = self.index_by_number.get(&number)?;
        Some(&self.fields[index])
    }

    /// This operation is not needed here
    // TODO: remove it
    pub fn cast<M: 'static>(&self, message: Box<dyn Message>) -> Result<M, Box<dyn Message>> {
        message.downcast_box::<M>().map(|m| *m)
    }
}

/// Identity comparison: message descriptor are equal if their addresses are equal
impl PartialEq for MessageDescriptor {
    fn eq(&self, other: &MessageDescriptor) -> bool {
        self as *const MessageDescriptor == other as *const MessageDescriptor
    }
}
<commit_msg>Fix field_by_name function signature<commit_after>use std::collections::HashMap;
use std::marker;

use crate::core::Message;

use crate::descriptor::DescriptorProto;
use crate::descriptor::FileDescriptorProto;

use crate::reflect::accessor::FieldAccessor;
use crate::reflect::find_message_or_enum::find_message_or_enum;
use crate::reflect::find_message_or_enum::MessageOrEnum;
use crate::reflect::reflect_deep_eq::ReflectDeepEq;
use crate::reflect::FieldDescriptor;

trait MessageFactory: Send + Sync + 'static {
    fn new_instance(&self) -> Box<dyn Message>;
    fn default_instance(&self) -> &dyn Message;
    fn clone(&self, message: &dyn Message) -> Box<dyn Message>;
    fn eq(&self, a: &dyn Message, b: &dyn Message) -> bool;
}

struct MessageFactoryImpl<M>(marker::PhantomData<M>);

impl<M> MessageFactory for MessageFactoryImpl<M>
where
    M: 'static + Message + Default + Clone + PartialEq,
{
    fn new_instance(&self) -> Box<dyn Message> {
        let m: M = Default::default();
        Box::new(m)
    }

    fn default_instance(&self) -> &dyn Message {
        M::default_instance() as &dyn Message
    }

    fn clone(&self, message: &dyn Message) -> Box<dyn Message> {
        let m: &M = message.downcast_ref().expect("wrong message type");
        Box::new(m.clone())
    }

    fn eq(&self, a: &dyn Message, b: &dyn Message) -> bool {
        let a: &M = a.downcast_ref().expect("wrong message type");
        let b: &M = b.downcast_ref().expect("wrong message type");
        a == b
    }
}

/// Dynamic representation of message type.
///
/// Used for reflection.
pub struct MessageDescriptor {
    full_name: String,
    file_descriptor_proto: &'static FileDescriptorProto,
    proto: &'static DescriptorProto,
    factory: &'static dyn MessageFactory,
    fields: Vec<FieldDescriptor>,

    index_by_name: HashMap<String, usize>,
    index_by_name_or_json_name: HashMap<String, usize>,
    index_by_number: HashMap<u32, usize>,
}

impl MessageDescriptor {
    /// Get underlying `DescriptorProto` object.
    pub fn get_proto(&self) -> &DescriptorProto {
        self.proto
    }

    /// Get a message descriptor for given message type
    pub fn for_type<M: Message>() -> &'static MessageDescriptor {
        M::descriptor_static()
    }

    fn compute_full_name(package: &str, path_to_package: &str, proto: &DescriptorProto) -> String {
        let mut full_name = package.to_owned();
        if path_to_package.len() != 0 {
            if full_name.len() != 0 {
                full_name.push('.');
            }
            full_name.push_str(path_to_package);
        }
        if full_name.len() != 0 {
            full_name.push('.');
        }
        full_name.push_str(proto.get_name());
        full_name
    }

    // Non-generic part of `new` is a separate function
    // to reduce code bloat from multiple instantiations.
    fn new_non_generic(
        protobuf_name_to_package: &'static str,
        fields: Vec<FieldAccessor>,
        file_descriptor_proto: &'static FileDescriptorProto,
        factory: &'static dyn MessageFactory,
    ) -> MessageDescriptor {
        let (path_to_package, proto) =
            match find_message_or_enum(file_descriptor_proto, protobuf_name_to_package) {
                (path_to_package, MessageOrEnum::Message(m)) => (path_to_package, m),
                (_, MessageOrEnum::Enum(_)) => panic!("not a message"),
            };

        let mut field_proto_by_name = HashMap::new();
        for field_proto in &proto.field {
            field_proto_by_name.insert(field_proto.get_name(), field_proto);
        }

        let mut index_by_name = HashMap::new();
        let mut index_by_name_or_json_name = HashMap::new();
        let mut index_by_number = HashMap::new();

        let fields: Vec<_> = fields
            .into_iter()
            .map(|f| {
                let proto = *field_proto_by_name.get(f.name).unwrap();
                FieldDescriptor::new(f, proto)
            })
            .collect();

        for (i, f) in fields.iter().enumerate() {
            assert!(index_by_number
                .insert(f.proto().get_number() as u32, i)
                .is_none());
            assert!(index_by_name
                .insert(f.proto().get_name().to_owned(), i)
                .is_none());
            assert!(index_by_name_or_json_name
                .insert(f.proto().get_name().to_owned(), i)
                .is_none());

            let json_name = f.json_name().to_owned();

            if json_name != f.proto().get_name() {
                assert!(index_by_name_or_json_name.insert(json_name, i).is_none());
            }
        }
        MessageDescriptor {
            full_name: MessageDescriptor::compute_full_name(
                file_descriptor_proto.get_package(),
                &path_to_package,
                &proto,
            ),
            proto,
            factory,
            fields,
            index_by_name,
            index_by_name_or_json_name,
            index_by_number,
            file_descriptor_proto,
        }
    }

    /// Construct a new message descriptor.
    ///
    /// This operation is called from generated code and rarely
    /// need to be called directly.
    ///
    /// This function is not a part of public API.
    #[doc(hidden)]
    pub fn new<M: 'static + Message + Default + Clone + PartialEq>(
        protobuf_name_to_package: &'static str,
        fields: Vec<FieldAccessor>,
        file_descriptor_proto: &'static FileDescriptorProto,
    ) -> MessageDescriptor {
        let factory = &MessageFactoryImpl(marker::PhantomData::<M>);
        MessageDescriptor::new_non_generic(
            protobuf_name_to_package,
            fields,
            file_descriptor_proto,
            factory,
        )
    }

    /// `FileDescriptorProto` containg this message type
    pub fn file_descriptor_proto(&self) -> &FileDescriptorProto {
        self.file_descriptor_proto
    }

    /// New empty message
    pub fn new_instance(&self) -> Box<dyn Message> {
        self.factory.new_instance()
    }

    /// Shared immutable empty message
    pub fn default_instance(&self) -> &dyn Message {
        self.factory.default_instance()
    }

    /// Clone a message
    pub fn clone(&self, message: &dyn Message) -> Box<dyn Message> {
        self.factory.clone(message)
    }

    /// Check if two messages equal.
    ///
    /// # Panics
    ///
    /// Is any message has different type than this descriptor.
    pub fn eq(&self, a: &dyn Message, b: &dyn Message) -> bool {
        self.factory.eq(a, b)
    }

    /// Similar to `eq`, but considers `NaN` values equal.
    ///
    /// # Panics
    ///
    /// Is any message has different type than this descriptor.
    pub fn deep_eq(&self, a: &dyn Message, b: &dyn Message) -> bool {
        // Explicitly force panic even if field list is empty
        assert_eq!(
            self as *const MessageDescriptor,
            a.descriptor() as *const MessageDescriptor
        );
        assert_eq!(
            self as *const MessageDescriptor,
            b.descriptor() as *const MessageDescriptor
        );

        for field in self.fields() {
            let af = field.get_reflect(a);
            let bf = field.get_reflect(b);
            if !af.reflect_deep_eq(&bf) {
                return false;
            }
        }
        true
    }

    /// Message name as given in `.proto` file
    pub fn name(&self) -> &'static str {
        self.proto.get_name()
    }

    /// Fully qualified protobuf message name
    pub fn full_name(&self) -> &str {
        &self.full_name[..]
    }

    /// Message field descriptors.
    pub fn fields(&self) -> &[FieldDescriptor] {
        &self.fields
    }

    /// Find message field by protobuf field name
    ///
    /// Note: protobuf field name might be different for Rust field name.
    pub fn field_by_name<'a>(&'a self, name: &str) -> Option<&'a FieldDescriptor> {
        let &index = self.index_by_name.get(name)?;
        Some(&self.fields[index])
    }

    /// Find message field by field name or field JSON name
    pub fn field_by_name_or_json_name<'a>(&'a self, name: &str) -> Option<&'a FieldDescriptor> {
        let &index = self.index_by_name_or_json_name.get(name)?;
        Some(&self.fields[index])
    }

    /// Find message field by field name
    pub fn field_by_number<'a>(&'a self, number: u32) -> Option<&'a FieldDescriptor> {
        let &index = self.index_by_number.get(&number)?;
        Some(&self.fields[index])
    }

    /// This operation is not needed here
    // TODO: remove it
    pub fn cast<M: 'static>(&self, message: Box<dyn Message>) -> Result<M, Box<dyn Message>> {
        message.downcast_box::<M>().map(|m| *m)
    }
}

/// Identity comparison: message descriptor are equal if their addresses are equal
impl PartialEq for MessageDescriptor {
    fn eq(&self, other: &MessageDescriptor) -> bool {
        self as *const MessageDescriptor == other as *const MessageDescriptor
    }
}
<|endoftext|>


<commit_before>//! This example demonstrates using Raft to implement a replicated hashmap over `n` servers and
//! interact with them over `m` clients.
//!
//! Comments below will aim to be tailored towards Raft and it's usage. If you have any questions,
//! Please, just open an issue.
//!
//! TODO: For the sake of simplicity of this example, we don't implement a `Log` and just use a
//! simple testing one. We should improve this in the future.

// In order to use Serde we need to enable these nightly features.
#![feature(plugin)]
#![feature(custom_derive)]
#![plugin(serde_macros)]

extern crate raft; // <--- Kind of a big deal
extern crate env_logger;
extern crate docopt;
extern crate serde;
extern crate rustc_serialize;

use std::net::SocketAddr;
use std::str::FromStr;
use std::io::{Error, Result};
use std::collections::HashMap;
use std::fmt::Debug;
use std::hash::Hash;

use serde::json::{self, Value};
use serde::{Serialize, Deserialize};
use docopt::Docopt;

// Raft's major components. See comments in code on usage and things.
use raft::{
    Server,
    Client,
    state_machine,
    persistent_log,
    ServerId,
};
// A payload datatype. We're just using a simple enum. You can use whatever.
use Message::*;

// Using docopt we define the overall usage of the application.
static USAGE: &'static str = "
A replicated mutable value. Operations on the register have serializable
consistency, but no durability (once all register servers are terminated the
value is lost).

Each register server holds a replica of the register, and coordinates with its
peers to update the register's value according to client commands. The register
is available for reading and writing only if a majority of register servers are
available.

Commands:

  get     Returns the current value of the key.

  put     Sets the current value of the key, and returns the previous
          value.

  cas     (compare and set) Conditionally sets the value of the key if the
          current value matches an expected value, returning true if the
          key was set.

  server  Starts a key server. Servers must be provided a unique ID and
          address (ip:port) at startup, along with the ID and address of all
          peer servers.register

Usage:
  hashmap get <key> (<node-address>)...
  hashmap put <key> <new-value> (<node-address>)...
  hashmap cas <key> <expected-value> <new-value> (<node-address>)...
  hashmap server <id> [<node-id> <node-address>]...
  hashmap (-h | --help)

Options:
  -h --help   Show a help message.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    cmd_server: bool,
    cmd_get: bool,
    cmd_put: bool,
    cmd_cas: bool,

    // When creating a server you will necessarily need some sort of unique ID for it as well
    // as a list of peers. In this example we just accept them straight from args. You might
    // find it best to use a `toml` or `yaml` or `json` file.
    arg_id: Option<u64>,
    arg_node_id: Vec<u64>,
    arg_node_address: Vec<String>,

    // In this example keys and values are associated. In your application you can model your data
    // however you please.
    arg_key: String,
    arg_new_value: String,
    arg_expected_value: String,
}

/// This is the defined message type for this example. For the sake of simplicity we don't go very
/// far with this. In a "real" application you may want to more distinctly distinguish between
/// data meant for `.query()` and data meant for `.propose()`.
#[derive(Serialize, Deserialize)]
pub enum Message {
    Get(String),
    Put(String, Value),
    Cas(String, Value, Value),
}

/// Just a plain old boring "parse args and dispatch" call.
fn main() {
    let _ = env_logger::init();
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());
    if args.cmd_server {
        server(&args);
    } else if args.cmd_get {
        get(&args);
    } else if args.cmd_put {
        put(&args);
    } else if args.cmd_cas {
        cas(&args);
    }
}

/// A simple convenience method since this is an example and it should exit if given invalid params.
fn parse_addr(addr: &str) -> SocketAddr {
    SocketAddr::from_str(addr)
               .ok()
               .expect(&format!("unable to parse socket address: {}", addr))
}

/// Creates a Raft server using the specified ID from the list of nodes.
fn server(args: &Args) {
    // Creating a raft server requires several things:

    // A persistent log implementation, which manages the persistent, replicated log...
    let persistent_log = persistent_log::MemLog::new();

    // A state machine which replicated state. This state should be the same on all nodes.
    let state_machine = HashmapStateMachine::new();

    // As well as a unique server id.
    let id = ServerId::from(args.arg_id.unwrap());

    // ...  And a list of peers.
    let mut peers = args.arg_node_id
                    .iter()
                    .zip(args.arg_node_address.iter())
                    .map(|(&id, addr)| (ServerId::from(id), parse_addr(&addr)))
                    .collect::<HashMap<_,_>>();

    // The Raft Server will panic if it's ID is inside of it's peer set. Don't do that.
    // Instead, take it out and use it!
    let addr = peers.remove(&id).unwrap();

    // Using all of the above components.
    // You probably shouldn't `.unwrap()` in production code unless you're totally sure it works
    // 100% of the time, all the time.
    Server::run(id, addr, peers, persistent_log, state_machine).unwrap();
}

/// Gets a value for a given key from the provided Raft cluster.
fn get(args: &Args) {
    // Clients necessarily need to now the valid set of nodes which they can talk to.
    // This is both so they can try to talk to all the nodes if some are failing, and so that it
    // can verify that it's not being lead astray somehow in redirections on leadership changes.
    let cluster = args.arg_node_address.iter()
        .map(|v| parse_addr(&v))
        .collect();

    // Clients and be stored and reused, or used once and discarded.
    // There is very small overhead in connecting a new client to a cluster as it must discover and
    // identify itself to the leader.
    let mut client = Client::new(cluster);

    // In this example `serde::json` is used to serialize and deserialize messages.
    // Since Raft accepts `[u8]` the way you structure your data, the serialization method you
    // choose, and how you interpret that data is entirely up to you.
    let payload = json::to_string(&Message::Get(args.arg_key.clone())).unwrap();

    // A query executes **immutably** on the leader of the cluster and does not pass through the
    // persistent log. This is intended for querying the current state of the state machine.
    let response = client.query(payload.as_bytes()).unwrap();

    // A response will block until it's query is complete. This is intended and expected behavior
    // based on the papers specifications.
    println!("{}", String::from_utf8(response).unwrap())
}

/// Sets a value for a given key in the provided Raft cluster.
fn put(args: &Args) {
    // Same as above.
    let cluster = args.arg_node_address.iter()
        .map(|v| parse_addr(&v))
        .collect();

    let mut client = Client::new(cluster);

    let new_value = json::to_value(&args.arg_new_value);
    let payload = json::to_string(&Message::Put(args.arg_key.clone(), new_value)).unwrap();

    // A propose will go through the persistent log and mutably modify the state machine in some
    // way. This is **much** slower than `.query()`.
    let response = client.propose(payload.as_bytes()).unwrap();

    // A response will block until it's proposal is complete. This is intended and expected behavior
    // based on the papers specifications.
    println!("{}", String::from_utf8(response).unwrap())
}

fn cas(_args: &Args) {
    panic!("unimplemented: waiting on changes to the Raft Client and StateMachine APIs");
}

/// A state machine that holds a hashmap.
#[derive(Debug)]
pub struct HashmapStateMachine {
    map: HashMap<String, Value>,
}

/// Implement anything you want... A `new()` is generally a great idea.
impl HashmapStateMachine {
    pub fn new() -> HashmapStateMachine {
        HashmapStateMachine {
            map: HashMap::new(),
        }
    }
}

/// Implementing `state_machine::StateMachine` allows your application specific state machine to be
/// used in Raft. Feel encouraged to base yours of one of ours in these examples.
impl state_machine::StateMachine for HashmapStateMachine {

    type Error = Error;

    /// `apply()` is called on when a client's `.propose()` is commited and reaches the state
    /// machine. At this point it is durable and is going to be applied on at least half the nodes
    /// within the next couple round trips.
    fn apply(&mut self, new_value: &[u8]) -> Result<Vec<u8>> {
        // Deserialize
        let string = String::from_utf8_lossy(new_value);
        let message = json::from_str(&string).unwrap();

        // Handle
        let response = match message {
            Get(key) => {
                let old_value = &self.map.get(&key).map(|v| v.clone());
                json::to_string(old_value)
            },
            Put(key, value) => {
                let old_value = &self.map.insert(key, value);
                json::to_string(old_value)
            },
            Cas(key, old_check, new) => {
                if *self.map.get(&key).unwrap() == old_check {
                    let _ = self.map.insert(key, new);
                    json::to_string(&true)
                } else {
                    json::to_string(&false)
                }
            },
        };

        // Respond.
        Ok(response.unwrap().into_bytes())
    }

    /// `query()` is called on when a client's `.query()` is recieved. It does not go through the
    /// persistent log, it does not mutate the state of the state machine, and it is intended to be
    /// fast.
    fn query(&self, query: &[u8]) -> Result<Vec<u8>> {
        // Deserialize
        let string = String::from_utf8_lossy(query);
        let message = json::from_str(&string).unwrap();

        // Handle
        let response = match message {
            Get(key) => {
                let old_value = &self.map.get(&key).map(|v| v.clone());
                json::to_string(old_value)
            },
            _ => panic!("Can't do mutating requests in query"),
        };

        // Respond.
        Ok(response.unwrap().into_bytes())
    }

    fn snapshot(&self) -> Result<Vec<u8>> {
        Ok(json::to_string(&self.map)
            .unwrap()
            .into_bytes())
    }

    fn restore_snapshot(&mut self, snapshot_value: Vec<u8>) -> Result<()> {
        self.map = json::from_str(&String::from_utf8_lossy(&snapshot_value)).unwrap();
        Ok(())
    }
}
<commit_msg>Remove some unnecessary<commit_after>//! This example demonstrates using Raft to implement a replicated hashmap over `n` servers and
//! interact with them over `m` clients.
//!
//! Comments below will aim to be tailored towards Raft and it's usage. If you have any questions,
//! Please, just open an issue.
//!
//! TODO: For the sake of simplicity of this example, we don't implement a `Log` and just use a
//! simple testing one. We should improve this in the future.

// In order to use Serde we need to enable these nightly features.
#![feature(plugin)]
#![feature(custom_derive)]
#![plugin(serde_macros)]

extern crate raft; // <--- Kind of a big deal
extern crate env_logger;
extern crate docopt;
extern crate serde;
extern crate rustc_serialize;

use std::net::SocketAddr;
use std::str::FromStr;
use std::io::{Error, Result};
use std::collections::HashMap;

use serde::json::{self, Value};
use docopt::Docopt;

// Raft's major components. See comments in code on usage and things.
use raft::{
    Server,
    Client,
    state_machine,
    persistent_log,
    ServerId,
};
// A payload datatype. We're just using a simple enum. You can use whatever.
use Message::*;

// Using docopt we define the overall usage of the application.
static USAGE: &'static str = "
A replicated mutable value. Operations on the register have serializable
consistency, but no durability (once all register servers are terminated the
value is lost).

Each register server holds a replica of the register, and coordinates with its
peers to update the register's value according to client commands. The register
is available for reading and writing only if a majority of register servers are
available.

Commands:

  get     Returns the current value of the key.

  put     Sets the current value of the key, and returns the previous
          value.

  cas     (compare and set) Conditionally sets the value of the key if the
          current value matches an expected value, returning true if the
          key was set.

  server  Starts a key server. Servers must be provided a unique ID and
          address (ip:port) at startup, along with the ID and address of all
          peer servers.register

Usage:
  hashmap get <key> (<node-address>)...
  hashmap put <key> <new-value> (<node-address>)...
  hashmap cas <key> <expected-value> <new-value> (<node-address>)...
  hashmap server <id> [<node-id> <node-address>]...
  hashmap (-h | --help)

Options:
  -h --help   Show a help message.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    cmd_server: bool,
    cmd_get: bool,
    cmd_put: bool,
    cmd_cas: bool,

    // When creating a server you will necessarily need some sort of unique ID for it as well
    // as a list of peers. In this example we just accept them straight from args. You might
    // find it best to use a `toml` or `yaml` or `json` file.
    arg_id: Option<u64>,
    arg_node_id: Vec<u64>,
    arg_node_address: Vec<String>,

    // In this example keys and values are associated. In your application you can model your data
    // however you please.
    arg_key: String,
    arg_new_value: String,
    arg_expected_value: String,
}

/// This is the defined message type for this example. For the sake of simplicity we don't go very
/// far with this. In a "real" application you may want to more distinctly distinguish between
/// data meant for `.query()` and data meant for `.propose()`.
#[derive(Serialize, Deserialize)]
pub enum Message {
    Get(String),
    Put(String, Value),
    Cas(String, Value, Value),
}

/// Just a plain old boring "parse args and dispatch" call.
fn main() {
    let _ = env_logger::init();
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());
    if args.cmd_server {
        server(&args);
    } else if args.cmd_get {
        get(&args);
    } else if args.cmd_put {
        put(&args);
    } else if args.cmd_cas {
        cas(&args);
    }
}

/// A simple convenience method since this is an example and it should exit if given invalid params.
fn parse_addr(addr: &str) -> SocketAddr {
    SocketAddr::from_str(addr)
               .ok()
               .expect(&format!("unable to parse socket address: {}", addr))
}

/// Creates a Raft server using the specified ID from the list of nodes.
fn server(args: &Args) {
    // Creating a raft server requires several things:

    // A persistent log implementation, which manages the persistent, replicated log...
    let persistent_log = persistent_log::MemLog::new();

    // A state machine which replicated state. This state should be the same on all nodes.
    let state_machine = HashmapStateMachine::new();

    // As well as a unique server id.
    let id = ServerId::from(args.arg_id.unwrap());

    // ...  And a list of peers.
    let mut peers = args.arg_node_id
                    .iter()
                    .zip(args.arg_node_address.iter())
                    .map(|(&id, addr)| (ServerId::from(id), parse_addr(&addr)))
                    .collect::<HashMap<_,_>>();

    // The Raft Server will panic if it's ID is inside of it's peer set. Don't do that.
    // Instead, take it out and use it!
    let addr = peers.remove(&id).unwrap();

    // Using all of the above components.
    // You probably shouldn't `.unwrap()` in production code unless you're totally sure it works
    // 100% of the time, all the time.
    Server::run(id, addr, peers, persistent_log, state_machine).unwrap();
}

/// Gets a value for a given key from the provided Raft cluster.
fn get(args: &Args) {
    // Clients necessarily need to now the valid set of nodes which they can talk to.
    // This is both so they can try to talk to all the nodes if some are failing, and so that it
    // can verify that it's not being lead astray somehow in redirections on leadership changes.
    let cluster = args.arg_node_address.iter()
        .map(|v| parse_addr(&v))
        .collect();

    // Clients and be stored and reused, or used once and discarded.
    // There is very small overhead in connecting a new client to a cluster as it must discover and
    // identify itself to the leader.
    let mut client = Client::new(cluster);

    // In this example `serde::json` is used to serialize and deserialize messages.
    // Since Raft accepts `[u8]` the way you structure your data, the serialization method you
    // choose, and how you interpret that data is entirely up to you.
    let payload = json::to_string(&Message::Get(args.arg_key.clone())).unwrap();

    // A query executes **immutably** on the leader of the cluster and does not pass through the
    // persistent log. This is intended for querying the current state of the state machine.
    let response = client.query(payload.as_bytes()).unwrap();

    // A response will block until it's query is complete. This is intended and expected behavior
    // based on the papers specifications.
    println!("{}", String::from_utf8(response).unwrap())
}

/// Sets a value for a given key in the provided Raft cluster.
fn put(args: &Args) {
    // Same as above.
    let cluster = args.arg_node_address.iter()
        .map(|v| parse_addr(&v))
        .collect();

    let mut client = Client::new(cluster);

    let new_value = json::to_value(&args.arg_new_value);
    let payload = json::to_string(&Message::Put(args.arg_key.clone(), new_value)).unwrap();

    // A propose will go through the persistent log and mutably modify the state machine in some
    // way. This is **much** slower than `.query()`.
    let response = client.propose(payload.as_bytes()).unwrap();

    // A response will block until it's proposal is complete. This is intended and expected behavior
    // based on the papers specifications.
    println!("{}", String::from_utf8(response).unwrap())
}

fn cas(_args: &Args) {
    panic!("unimplemented: waiting on changes to the Raft Client and StateMachine APIs");
}

/// A state machine that holds a hashmap.
#[derive(Debug)]
pub struct HashmapStateMachine {
    map: HashMap<String, Value>,
}

/// Implement anything you want... A `new()` is generally a great idea.
impl HashmapStateMachine {
    pub fn new() -> HashmapStateMachine {
        HashmapStateMachine {
            map: HashMap::new(),
        }
    }
}

/// Implementing `state_machine::StateMachine` allows your application specific state machine to be
/// used in Raft. Feel encouraged to base yours of one of ours in these examples.
impl state_machine::StateMachine for HashmapStateMachine {

    type Error = Error;

    /// `apply()` is called on when a client's `.propose()` is commited and reaches the state
    /// machine. At this point it is durable and is going to be applied on at least half the nodes
    /// within the next couple round trips.
    fn apply(&mut self, new_value: &[u8]) -> Result<Vec<u8>> {
        // Deserialize
        let string = String::from_utf8_lossy(new_value);
        let message = json::from_str(&string).unwrap();

        // Handle
        let response = match message {
            Get(key) => {
                let old_value = &self.map.get(&key).map(|v| v.clone());
                json::to_string(old_value)
            },
            Put(key, value) => {
                let old_value = &self.map.insert(key, value);
                json::to_string(old_value)
            },
            Cas(key, old_check, new) => {
                if *self.map.get(&key).unwrap() == old_check {
                    let _ = self.map.insert(key, new);
                    json::to_string(&true)
                } else {
                    json::to_string(&false)
                }
            },
        };

        // Respond.
        Ok(response.unwrap().into_bytes())
    }

    /// `query()` is called on when a client's `.query()` is recieved. It does not go through the
    /// persistent log, it does not mutate the state of the state machine, and it is intended to be
    /// fast.
    fn query(&self, query: &[u8]) -> Result<Vec<u8>> {
        // Deserialize
        let string = String::from_utf8_lossy(query);
        let message = json::from_str(&string).unwrap();

        // Handle
        let response = match message {
            Get(key) => {
                let old_value = &self.map.get(&key).map(|v| v.clone());
                json::to_string(old_value)
            },
            _ => panic!("Can't do mutating requests in query"),
        };

        // Respond.
        Ok(response.unwrap().into_bytes())
    }

    fn snapshot(&self) -> Result<Vec<u8>> {
        Ok(json::to_string(&self.map)
            .unwrap()
            .into_bytes())
    }

    fn restore_snapshot(&mut self, snapshot_value: Vec<u8>) -> Result<()> {
        self.map = json::from_str(&String::from_utf8_lossy(&snapshot_value)).unwrap();
        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate rustbox;

pub struct UIBuffer {
    width: uint,
    height: uint,
    cells: Vec<Vec<Cell>>
}

impl UIBuffer {
    pub fn new(width: uint, height: uint) -> UIBuffer {
        let cells = Cell::create_grid(width, height, ' ');

        UIBuffer {
            width: width,
            height: height,
            cells: cells,
        }
    }

    pub fn draw_range(&self, start: uint, stop: uint) {
        let rows = self.cells.slice(start, stop);
        for row in rows.iter() {
            for cell in row.iter() {
                rustbox::print_char(cell.x, cell.y, rustbox::Style::Normal, cell.fg, cell.bg, cell.ch);
            }
        }
    }

    /// Recreated the entire grid, will cells containing `ch`.
    pub fn fill(&mut self, ch: char) {
        self.cells = Cell::create_grid(self.width, self.height, ch);
    }

    /// Update the `ch` attribute of an individual cell
    pub fn update_cell_content(&mut self, x: uint, y: uint, ch: char) {
        self.cells[y][x].ch = ch
    }

    /// Update the `ch`, `fg`, and `bg` attributes of an indivudual cell
    pub fn update_cell(&mut self, x: uint, y: uint, ch: char, fg: rustbox::Color, bg: rustbox::Color) {
        let cell = &mut self.cells[y][x];
        cell.ch = ch;
        cell.fg = fg;
        cell.bg = bg;
    }
}

pub struct Cell {
    pub bg: rustbox::Color,
    pub fg: rustbox::Color,
    pub ch: char,
    pub x: uint,
    pub y: uint,
}

impl Cell {
    pub fn new() -> Cell {
        Cell {
            bg: rustbox::Color::Default,
            fg: rustbox::Color::White,
            ch: ' ',
            x: 0,
            y: 0,
        }
    }

    pub fn create_grid(width: uint, height: uint, ch: char) -> Vec<Vec<Cell>> {
        let mut rows = Vec::new();
        for voffset in range(0, height) {
            let mut cells = Vec::new();
            for boffset in range(0, width) {
                let mut cell = Cell::new();
                cell.x = boffset;
                cell.y = voffset;
                cell.ch = ch;
                cells.push(cell);
            }
            rows.push(cells);
        }

        return rows
    }
}

<commit_msg>Added tests for uibuf module<commit_after>extern crate rustbox;

pub struct UIBuffer {
    width: uint,
    height: uint,
    rows: Vec<Vec<Cell>>
}

impl UIBuffer {
    pub fn new(width: uint, height: uint) -> UIBuffer {
        let rows = Cell::create_grid(width, height, ' ');

        UIBuffer {
            width: width,
            height: height,
            rows: rows,
        }
    }

    pub fn draw_range(&self, start: uint, stop: uint) {
        let rows = self.rows.slice(start, stop);
        for row in rows.iter() {
            for cell in row.iter() {
                rustbox::print_char(cell.x, cell.y, rustbox::Style::Normal, cell.fg, cell.bg, cell.ch);
            }
        }
    }

    /// Recreated the entire grid, will cells containing `ch`.
    pub fn fill(&mut self, ch: char) {
        self.rows = Cell::create_grid(self.width, self.height, ch);
    }

    /// Update the `ch` attribute of an individual cell
    pub fn update_cell_content(&mut self, cell_num: uint, row_num: uint, ch: char) {
        self.rows[row_num][cell_num].ch = ch
    }

    /// Update the `ch`, `fg`, and `bg` attributes of an indivudual cell
    pub fn update_cell(&mut self, cell_num: uint, row_num: uint, ch: char, fg: rustbox::Color, bg: rustbox::Color) {
        let cell = &mut self.rows[row_num][cell_num];
        cell.ch = ch;
        cell.fg = fg;
        cell.bg = bg;
    }
}

pub struct Cell {
    pub bg: rustbox::Color,
    pub fg: rustbox::Color,
    pub ch: char,
    pub x: uint,
    pub y: uint,
}

impl Cell {
    pub fn new() -> Cell {
        Cell {
            bg: rustbox::Color::Default,
            fg: rustbox::Color::White,
            ch: ' ',
            x: 0,
            y: 0,
        }
    }

    pub fn create_grid(width: uint, height: uint, ch: char) -> Vec<Vec<Cell>> {
        let mut rows = Vec::new();
        for voffset in range(0, height) {
            let mut cells = Vec::new();
            for boffset in range(0, width) {
                let mut cell = Cell::new();
                cell.x = boffset;
                cell.y = voffset;
                cell.ch = ch;
                cells.push(cell);
            }
            rows.push(cells);
        }

        return rows
    }
}

#[cfg(test)]
mod tests {

    extern crate rustbox;

    use uibuf::UIBuffer;

    fn setup_uibuf() -> UIBuffer {
        UIBuffer::new(50, 50)
    }

    #[test]
    fn grid_is_created_correctly() {
        let uibuf = setup_uibuf();
        assert_eq!(uibuf.height, 50);
        assert_eq!(uibuf.width, 50);
        assert_eq!(uibuf.rows.len(), 50);
        assert_eq!(uibuf.rows[0][5].ch, ' ');
    }

    #[test]
    fn fill_updates_contents_of_all_cells() {
        let mut uibuf = setup_uibuf();
        uibuf.fill('x');

        // check some random cells
        assert_eq!(uibuf.rows[20][5].ch, 'x');
        assert_eq!(uibuf.rows[0][30].ch, 'x');
    }

    #[test]
    fn update_cell_content_updates_a_single_cell() {
        let mut uibuf = setup_uibuf();
        // update cell 10 in row 32
        uibuf.update_cell_content(10, 32, 'y');

        assert_eq!(uibuf.rows[32][10].ch, 'y');
    }

    #[test]
    fn update_cell_updates_all_attrs_of_cell() {
        let mut uibuf = setup_uibuf();
        let cell_num = 10u;
        let row_num = 0u;
        let ch = 'q';
        let fg = rustbox::Color::Red;
        let bg = rustbox::Color::Blue;

        uibuf.update_cell(cell_num, row_num, ch, fg, bg);

        let cell = &uibuf.rows[row_num][cell_num];
        assert_eq!(cell.ch, ch);

        // FIXME(greg): this fails cos '==' is not implemented for rustbox::Color.
        // assert_eq!(cell.fg, fg);
        // assert_eq!(cell.bg, bg);
    }

}
<|endoftext|>