<commit_before>#![feature(box_patterns)]
#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

mod cli;
mod core;

use rustyline::{Helper, CompletionType, Config, EditMode, Editor};
use rustyline::hint::Hinter;
use rustyline::highlight::Highlighter;
use rustyline::completion::{Completer, FilenameCompleter};
use rustyline::error::ReadlineError;
use std::fs;

mod scheme {
    pub const HTTP: &'static str = "http://";
    pub const TCP: &'static str = "tcp://";

    pub fn is_http(url: &str) -> bool {
        url.starts_with(HTTP)
    }

    pub fn is_tcp(url: &str) -> bool {
        url.starts_with(TCP)
    }
}

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
const PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
const PROMPT: &'static str = ">> ";

#[derive(Default)]
struct Completes {
    file_completer: FilenameCompleter,
}

impl Helper for Completes {}

impl Hinter for Completes {
    fn hint(&self, _line: &str, _pos: usize) -> Option<String> {
        None
    }
}

impl Highlighter for Completes {}

impl Completer for Completes {
    type Candidate = String;
    fn complete(&self, line: &str, _pos: usize) -> rustyline::Result<(usize, Vec<String>)> {
        // TODO Completion for function names
        let cmds = vec![
            command::HELP,
            command::LOAD,
            command::CONNECT,
            command::FNLIST,
            command::ANALYZE,
            command::DOT,
            command::IR,
            command::DECOMPILE,
        ];
        let mut ret: Vec<String> = cmds.into_iter()
            .filter(|s| s.starts_with(line))
            .map(|s| s.to_string())
            .collect();
        match self.file_completer.complete(line, _pos) {
            Ok((n, ss)) => {
                let mut completed_lines = ss.into_iter()
                    .map(|s| format!("{}{}", &line[..n], s.display))
                    .collect();
                ret.append(&mut completed_lines);
            }
            Err(_) => {}
        }

        Ok((0, ret))
    }
}

const SEP: &'static str = "END";
const USAGE: &'static str = "
Usage:
  radeco <bin>
  radeco [--append]
  radeco (--help | --version)

  Options:
  -h, --help        Show this screen.
  --version         Show current version.
  --bin             Load a program.
  --append          Append separator to the end of every output.
";

fn main() {
    let (arg, is_append_mode) = cli::parse_args(USAGE);
    let config = Config::builder()
        .auto_add_history(true)
        .history_ignore_space(true)
        .completion_type(CompletionType::List)
        .edit_mode(EditMode::Emacs)
        .build();
    let mut rl = Editor::with_config(config);
    rl.set_helper(Some(Completes::default()));
    core::PROJ.with(move |proj| {
        *proj.borrow_mut() = arg.and_then(|ref s| {
            if scheme::is_http(s) {
                core::load_proj_http(&s[scheme::HTTP.len()..])
            } else if scheme::is_tcp(s) {
                core::load_proj_tcp(&s[scheme::TCP.len()..])
            } else if is_file(s) {
                Ok(core::load_proj_by_path(s))
            } else {
                Err("Invalid argument")
            }.ok()
        });
    });

    loop {
        let readline = rl.readline(PROMPT);
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2);
                if is_append_mode {
                    println!("{}", SEP);
                }
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const HELP: &'static str = "help";
    pub const LOAD: &'static str = "load";
    pub const CONNECT: &'static str = "connect";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";

    pub fn help() {
        let width = 30;
        println!("{:width$}    Show this help", HELP, width = width);
        println!(
            "{:width$}    Load binary",
            format!("{} path", LOAD),
            width = width
        );
        println!(
            "{:width$}    Connect to radare2 server",
            format!("{} (http|tcp)://<url>", CONNECT),
            width = width
        );
        println!(
            "{:width$}    Show function list",
            FNLIST,
            width = width
        );
        println!(
            "{:width$}    Analyze <func>",
            format!("{} <func>", ANALYZE),
            width = width
        );
        println!(
            "{:width$}    Emit IR of <func>",
            format!("{} <func>", IR),
            width = width
        );
        println!(
            "{:width$}    Emit graph of the IR in Graphviz dot",
            format!("{} <func>", DOT),
            width = width
        );
        println!(
            "{:width$}    Decompile <func>",
            format!("{} <func>", DECOMPILE),
            width = width
        );
    }
}

fn cmd(op1: Option<&str>, op2: Option<&str>) {
    core::PROJ.with(|proj_opt| {
        match (op1, op2) {
            (Some(command::HELP), _) => {
                command::help();
                return;
            }
            (Some(command::LOAD), Some(path)) => {
                if is_file(path) {
                    *proj_opt.borrow_mut() = Some(core::load_proj_by_path(path));
                    return;
                } else {
                    println!("{} is not found.", path);
                    return;
                }
            }
            (Some(command::CONNECT), Some(url)) => {
                let p_opt = if scheme::is_http(&url) {
                    core::load_proj_http(&url[scheme::HTTP.len()..])
                } else if scheme::is_tcp(&url) {
                    core::load_proj_tcp(&url[scheme::TCP.len()..])
                } else {
                    Err("Invalid url")
                };
                match p_opt {
                    Ok(p) => *proj_opt.borrow_mut() = Some(p),
                    Err(msg) => println!("{}", msg),
                }
                return;
            }
            _ => {}
        };
        if proj_opt.borrow().is_none() {
            println!("Load a project first");
            return;
        }
        let mut proj_ = proj_opt.borrow_mut();
        let proj = proj_.as_mut().unwrap();
        match (op1, op2) {
            (Some(command::ANALYZE), Some("*")) => {
                let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                    rmod.functions.values_mut()
                });
                for rfn in rfns {
                    core::analyze(rfn);
                }
            }
            (Some(command::FNLIST), _) => {
                let funcs = core::fn_list(&proj);
                println!("{}", funcs.join("\n"));
            }
            // TODO Show list of dependency information of analyses
            // TODO Add command for individual analyses
            (Some(command::ANALYZE), Some(f)) => {
                if let Some(rfn) = core::get_function_mut(f, proj) {
                    core::analyze(rfn);
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::DOT), Some(f)) => {
                if let Some(rfn) = core::get_function(f, &proj) {
                    println!("{}", core::emit_dot(rfn.ssa()));
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::IR), Some(f)) => {
                if let Some(rfn) = core::get_function(f, &proj) {
                    println!("{}", core::emit_ir(rfn));
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::DECOMPILE), Some(f)) => {
                if let Some(rfn) = core::get_function(f, &proj) {
                    let rmod = proj.iter().map(|i| i.module).next().unwrap();
                    let func_name_map = core::func_names(&rmod);
                    let strings = core::strings(&rmod);
                    match core::decompile(rfn, &func_name_map, &strings) {
                        Ok(res) => println!("{}", res),
                        Err(err) => println!("{}", err),
                    }
                } else {
                    println!("{} is not found", f);
                }
            }
            _ => {
                println!(
                    "Invalid command {} {}",
                    op1.unwrap_or(""),
                    op2.unwrap_or("")
                );
            }
        }
    });
}

fn is_file(path: &str) -> bool {
    fs::metadata(path).map(|f| f.is_file()).unwrap_or(false)
}
<commit_msg>Add -v option<commit_after>#![feature(box_patterns)]
#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

mod cli;
mod core;

use rustyline::{Helper, CompletionType, Config, EditMode, Editor};
use rustyline::hint::Hinter;
use rustyline::highlight::Highlighter;
use rustyline::completion::{Completer, FilenameCompleter};
use rustyline::error::ReadlineError;
use std::fs;

mod scheme {
    pub const HTTP: &'static str = "http://";
    pub const TCP: &'static str = "tcp://";

    pub fn is_http(url: &str) -> bool {
        url.starts_with(HTTP)
    }

    pub fn is_tcp(url: &str) -> bool {
        url.starts_with(TCP)
    }
}

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
const PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
const PROMPT: &'static str = ">> ";

#[derive(Default)]
struct Completes {
    file_completer: FilenameCompleter,
}

impl Helper for Completes {}

impl Hinter for Completes {
    fn hint(&self, _line: &str, _pos: usize) -> Option<String> {
        None
    }
}

impl Highlighter for Completes {}

impl Completer for Completes {
    type Candidate = String;
    fn complete(&self, line: &str, _pos: usize) -> rustyline::Result<(usize, Vec<String>)> {
        // TODO Completion for function names
        let cmds = vec![
            command::HELP,
            command::LOAD,
            command::CONNECT,
            command::FNLIST,
            command::ANALYZE,
            command::DOT,
            command::IR,
            command::DECOMPILE,
        ];
        let mut ret: Vec<String> = cmds.into_iter()
            .filter(|s| s.starts_with(line))
            .map(|s| s.to_string())
            .collect();
        match self.file_completer.complete(line, _pos) {
            Ok((n, ss)) => {
                let mut completed_lines = ss.into_iter()
                    .map(|s| format!("{}{}", &line[..n], s.display))
                    .collect();
                ret.append(&mut completed_lines);
            }
            Err(_) => {}
        }

        Ok((0, ret))
    }
}

const SEP: &'static str = "END";
const USAGE: &'static str = "
Usage:
  radeco <bin>
  radeco [--append]
  radeco (--help | --version)

  Options:
  -h, --help        Show this screen.
  -v, --version     Show current version.
  --bin             Load a program.
  --append          Append separator to the end of every output.
";

fn main() {
    let (arg, is_append_mode) = cli::parse_args(USAGE);
    let config = Config::builder()
        .auto_add_history(true)
        .history_ignore_space(true)
        .completion_type(CompletionType::List)
        .edit_mode(EditMode::Emacs)
        .build();
    let mut rl = Editor::with_config(config);
    rl.set_helper(Some(Completes::default()));
    core::PROJ.with(move |proj| {
        *proj.borrow_mut() = arg.and_then(|ref s| {
            if scheme::is_http(s) {
                core::load_proj_http(&s[scheme::HTTP.len()..])
            } else if scheme::is_tcp(s) {
                core::load_proj_tcp(&s[scheme::TCP.len()..])
            } else if is_file(s) {
                Ok(core::load_proj_by_path(s))
            } else {
                Err("Invalid argument")
            }.ok()
        });
    });

    loop {
        let readline = rl.readline(PROMPT);
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2);
                if is_append_mode {
                    println!("{}", SEP);
                }
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const HELP: &'static str = "help";
    pub const LOAD: &'static str = "load";
    pub const CONNECT: &'static str = "connect";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";

    pub fn help() {
        let width = 30;
        println!("{:width$}    Show this help", HELP, width = width);
        println!(
            "{:width$}    Load binary",
            format!("{} path", LOAD),
            width = width
        );
        println!(
            "{:width$}    Connect to radare2 server",
            format!("{} (http|tcp)://<url>", CONNECT),
            width = width
        );
        println!(
            "{:width$}    Show function list",
            FNLIST,
            width = width
        );
        println!(
            "{:width$}    Analyze <func>",
            format!("{} <func>", ANALYZE),
            width = width
        );
        println!(
            "{:width$}    Emit IR of <func>",
            format!("{} <func>", IR),
            width = width
        );
        println!(
            "{:width$}    Emit graph of the IR in Graphviz dot",
            format!("{} <func>", DOT),
            width = width
        );
        println!(
            "{:width$}    Decompile <func>",
            format!("{} <func>", DECOMPILE),
            width = width
        );
    }
}

fn cmd(op1: Option<&str>, op2: Option<&str>) {
    core::PROJ.with(|proj_opt| {
        match (op1, op2) {
            (Some(command::HELP), _) => {
                command::help();
                return;
            }
            (Some(command::LOAD), Some(path)) => {
                if is_file(path) {
                    *proj_opt.borrow_mut() = Some(core::load_proj_by_path(path));
                    return;
                } else {
                    println!("{} is not found.", path);
                    return;
                }
            }
            (Some(command::CONNECT), Some(url)) => {
                let p_opt = if scheme::is_http(&url) {
                    core::load_proj_http(&url[scheme::HTTP.len()..])
                } else if scheme::is_tcp(&url) {
                    core::load_proj_tcp(&url[scheme::TCP.len()..])
                } else {
                    Err("Invalid url")
                };
                match p_opt {
                    Ok(p) => *proj_opt.borrow_mut() = Some(p),
                    Err(msg) => println!("{}", msg),
                }
                return;
            }
            _ => {}
        };
        if proj_opt.borrow().is_none() {
            println!("Load a project first");
            return;
        }
        let mut proj_ = proj_opt.borrow_mut();
        let proj = proj_.as_mut().unwrap();
        match (op1, op2) {
            (Some(command::ANALYZE), Some("*")) => {
                let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                    rmod.functions.values_mut()
                });
                for rfn in rfns {
                    core::analyze(rfn);
                }
            }
            (Some(command::FNLIST), _) => {
                let funcs = core::fn_list(&proj);
                println!("{}", funcs.join("\n"));
            }
            // TODO Show list of dependency information of analyses
            // TODO Add command for individual analyses
            (Some(command::ANALYZE), Some(f)) => {
                if let Some(rfn) = core::get_function_mut(f, proj) {
                    core::analyze(rfn);
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::DOT), Some(f)) => {
                if let Some(rfn) = core::get_function(f, &proj) {
                    println!("{}", core::emit_dot(rfn.ssa()));
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::IR), Some(f)) => {
                if let Some(rfn) = core::get_function(f, &proj) {
                    println!("{}", core::emit_ir(rfn));
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::DECOMPILE), Some(f)) => {
                if let Some(rfn) = core::get_function(f, &proj) {
                    let rmod = proj.iter().map(|i| i.module).next().unwrap();
                    let func_name_map = core::func_names(&rmod);
                    let strings = core::strings(&rmod);
                    match core::decompile(rfn, &func_name_map, &strings) {
                        Ok(res) => println!("{}", res),
                        Err(err) => println!("{}", err),
                    }
                } else {
                    println!("{} is not found", f);
                }
            }
            _ => {
                println!(
                    "Invalid command {} {}",
                    op1.unwrap_or(""),
                    op2.unwrap_or("")
                );
            }
        }
    });
}

fn is_file(path: &str) -> bool {
    fs::metadata(path).map(|f| f.is_file()).unwrap_or(false)
}
<|endoftext|>


<commit_before>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::i32;
use std::net::{IpAddr, SocketAddr};
use std::str::FromStr;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};

use engine::Engines;
use futures::{Future, Stream};
use grpcio::{ChannelBuilder, EnvBuilder, Environment, Server as GrpcServer, ServerBuilder};
use kvproto::debugpb_grpc::create_debug;
use kvproto::import_sstpb_grpc::create_import_sst;
use kvproto::tikvpb_grpc::*;
use tokio_threadpool::{Builder as ThreadPoolBuilder, ThreadPool};
use tokio_timer::timer::Handle;

use crate::storage::lock_manager::deadlock::Service as DeadlockService;
use kvproto::deadlock_grpc::create_deadlock;

use crate::coprocessor::Endpoint;
use crate::import::ImportSSTService;
use crate::raftstore::store::SnapManager;
use crate::storage::{Engine, Storage};
use tikv_util::security::SecurityManager;
use tikv_util::timer::GLOBAL_TIMER_HANDLE;
use tikv_util::worker::Worker;
use tikv_util::Either;

use super::load_statistics::*;
use super::raft_client::RaftClient;
use super::resolve::StoreAddrResolver;
use super::service::*;
use super::snap::{Runner as SnapHandler, Task as SnapTask};
use super::transport::{RaftStoreRouter, ServerTransport};
use super::{Config, Result};

const LOAD_STATISTICS_SLOTS: usize = 4;
const LOAD_STATISTICS_INTERVAL: Duration = Duration::from_millis(100);
const MAX_GRPC_RECV_MSG_LEN: i32 = 10 * 1024 * 1024;
pub const GRPC_THREAD_PREFIX: &str = "grpc-server";
pub const STATS_THREAD_PREFIX: &str = "transport-stats";

/// The TiKV server
///
/// It hosts various internal components, including gRPC, the raftstore router
/// and a snapshot worker.
pub struct Server<T: RaftStoreRouter + 'static, S: StoreAddrResolver + 'static> {
    env: Arc<Environment>,
    /// A GrpcServer builder or a GrpcServer.
    ///
    /// If the listening port is configured, the server will be started lazily.
    builder_or_server: Option<Either<ServerBuilder, GrpcServer>>,
    local_addr: SocketAddr,
    // Transport.
    trans: ServerTransport<T, S>,
    raft_router: T,
    // For sending/receiving snapshots.
    snap_mgr: SnapManager,
    snap_worker: Worker<SnapTask>,

    // Currently load statistics is done in the thread.
    stats_pool: Option<ThreadPool>,
    thread_load: Arc<ThreadLoad>,
    timer: Handle,
}

impl<T: RaftStoreRouter, S: StoreAddrResolver + 'static> Server<T, S> {
    #[allow(clippy::too_many_arguments)]
    pub fn new<E: Engine>(
        cfg: &Arc<Config>,
        security_mgr: &Arc<SecurityManager>,
        storage: Storage<E>,
        cop: Endpoint<E>,
        raft_router: T,
        resolver: S,
        snap_mgr: SnapManager,
        debug_engines: Option<Engines>,
        import_service: Option<ImportSSTService<T>>,
        deadlock_service: Option<DeadlockService>,
    ) -> Result<Self> {
        // A helper thread (or pool) for transport layer.
        let stats_pool = ThreadPoolBuilder::new()
            .pool_size(cfg.stats_concurrency)
            .name_prefix(STATS_THREAD_PREFIX)
            .build();
        let thread_load = Arc::new(ThreadLoad::with_threshold(cfg.heavy_load_threshold));

        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(cfg.grpc_concurrency)
                .name_prefix(thd_name!(GRPC_THREAD_PREFIX))
                .build(),
        );
        let snap_worker = Worker::new("snap-handler");

        let kv_service = KvService::new(
            storage,
            cop,
            raft_router.clone(),
            snap_worker.scheduler(),
            Arc::clone(&thread_load),
        );

        let mut addr = SocketAddr::from_str(&cfg.addr)?;
        let ip = format!("{}", addr.ip());
        let channel_args = ChannelBuilder::new(Arc::clone(&env))
            .stream_initial_window_size(cfg.grpc_stream_initial_window_size.0 as i32)
            .max_concurrent_stream(cfg.grpc_concurrent_stream)
            .max_receive_message_len(MAX_GRPC_RECV_MSG_LEN)
            .max_send_message_len(-1)
            .http2_max_ping_strikes(i32::MAX) // For pings without data from clients.
            .build_args();
        let builder_or_server = {
            let mut sb = ServerBuilder::new(Arc::clone(&env))
                .channel_args(channel_args)
                .register_service(create_tikv(kv_service));
            sb = security_mgr.bind(sb, &ip, addr.port());
            if let Some(engines) = debug_engines {
                let debug_service = DebugService::new(engines, raft_router.clone());
                sb = sb.register_service(create_debug(debug_service));
            }
            if let Some(service) = import_service {
                sb = sb.register_service(create_import_sst(service));
            }
            if let Some(service) = deadlock_service {
                sb = sb.register_service(create_deadlock(service));
            }
            // When port is 0, it has to be binded now to get a valid address, which
            // is then reported to PD before the server is up. 0 is usually used in tests.
            if addr.port() == 0 {
                let server = sb.build()?;
                let (ref host, port) = server.bind_addrs()[0];
                addr = SocketAddr::new(IpAddr::from_str(host)?, port as u16);
                Either::Right(server)
            } else {
                Either::Left(sb)
            }
        };

        info!("listening on addr"; "addr" => addr);

        let raft_client = Arc::new(RwLock::new(RaftClient::new(
            Arc::clone(&env),
            Arc::clone(cfg),
            Arc::clone(security_mgr),
            raft_router.clone(),
            Arc::clone(&thread_load),
            stats_pool.sender().clone(),
        )));

        let trans = ServerTransport::new(
            raft_client,
            snap_worker.scheduler(),
            raft_router.clone(),
            resolver,
        );

        let svr = Server {
            env: Arc::clone(&env),
            builder_or_server: Some(builder_or_server),
            local_addr: addr,
            trans,
            raft_router,
            snap_mgr,
            snap_worker,
            stats_pool: Some(stats_pool),
            thread_load,
            timer: GLOBAL_TIMER_HANDLE.clone(),
        };

        Ok(svr)
    }

    pub fn transport(&self) -> ServerTransport<T, S> {
        self.trans.clone()
    }

    /// Starts the TiKV server.
    pub fn start(&mut self, cfg: Arc<Config>, security_mgr: Arc<SecurityManager>) -> Result<()> {
        let snap_runner = SnapHandler::new(
            Arc::clone(&self.env),
            self.snap_mgr.clone(),
            self.raft_router.clone(),
            security_mgr,
            Arc::clone(&cfg),
        );
        box_try!(self.snap_worker.start(snap_runner));
        let builder_or_server = self.builder_or_server.take().unwrap();
        let mut grpc_server = match builder_or_server {
            Either::Left(builder) => builder.build()?,
            Either::Right(server) => server,
        };
        grpc_server.start();
        self.builder_or_server = Some(Either::Right(grpc_server));

        let mut load_stats = {
            let tl = Arc::clone(&self.thread_load);
            ThreadLoadStatistics::new(LOAD_STATISTICS_SLOTS, GRPC_THREAD_PREFIX, tl)
        };
        self.stats_pool.as_ref().unwrap().spawn(
            self.timer
                .interval(Instant::now(), LOAD_STATISTICS_INTERVAL)
                .map_err(|_| ())
                .for_each(move |i| {
                    load_stats.record(i);
                    Ok(())
                }),
        );

        info!("TiKV is ready to serve");
        Ok(())
    }

    /// Stops the TiKV server.
    pub fn stop(&mut self) -> Result<()> {
        self.snap_worker.stop();
        if let Some(Either::Right(mut server)) = self.builder_or_server.take() {
            server.shutdown();
        }
        if let Some(pool) = self.stats_pool.take() {
            let _ = pool.shutdown_now().wait();
        }
        Ok(())
    }

    // Return listening address, this may only be used for outer test
    // to get the real address because we may use "127.0.0.1:0"
    // in test to avoid port conflict.
    pub fn listening_addr(&self) -> SocketAddr {
        self.local_addr
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::*;
    use std::sync::mpsc::*;
    use std::sync::*;
    use std::time::Duration;

    use super::*;

    use super::super::resolve::{Callback as ResolveCallback, StoreAddrResolver};
    use super::super::transport::RaftStoreRouter;
    use super::super::{Config, Result};
    use crate::coprocessor;
    use crate::raftstore::store::transport::Transport;
    use crate::raftstore::store::*;
    use crate::raftstore::Result as RaftStoreResult;
    use crate::storage::TestStorageBuilder;

    use crate::coprocessor::readpool_impl;
    use kvproto::raft_cmdpb::RaftCmdRequest;
    use kvproto::raft_serverpb::RaftMessage;
    use tikv_util::security::SecurityConfig;

    #[derive(Clone)]
    struct MockResolver {
        quick_fail: Arc<AtomicBool>,
        addr: Arc<Mutex<Option<String>>>,
    }

    impl StoreAddrResolver for MockResolver {
        fn resolve(&self, _: u64, cb: ResolveCallback) -> Result<()> {
            if self.quick_fail.load(Ordering::SeqCst) {
                return Err(box_err!("quick fail"));
            }
            let addr = self.addr.lock().unwrap();
            cb(addr
                .as_ref()
                .map(|s| s.to_owned())
                .ok_or(box_err!("not set")));
            Ok(())
        }
    }

    #[derive(Clone)]
    struct TestRaftStoreRouter {
        tx: Sender<usize>,
        significant_msg_sender: Sender<SignificantMsg>,
    }

    impl RaftStoreRouter for TestRaftStoreRouter {
        fn send_raft_msg(&self, _: RaftMessage) -> RaftStoreResult<()> {
            self.tx.send(1).unwrap();
            Ok(())
        }

        fn send_command(&self, _: RaftCmdRequest, _: Callback) -> RaftStoreResult<()> {
            self.tx.send(1).unwrap();
            Ok(())
        }

        fn significant_send(&self, _: u64, msg: SignificantMsg) -> RaftStoreResult<()> {
            self.significant_msg_sender.send(msg).unwrap();
            Ok(())
        }

        fn casual_send(&self, _: u64, _: CasualMessage) -> RaftStoreResult<()> {
            self.tx.send(1).unwrap();
            Ok(())
        }

        fn broadcast_unreachable(&self, _: u64) {
            let _ = self.tx.send(1);
        }
    }

    fn is_unreachable_to(msg: &SignificantMsg, region_id: u64, to_peer_id: u64) -> bool {
        *msg == SignificantMsg::Unreachable {
            region_id,
            to_peer_id,
        }
    }

    #[test]
    // if this failed, unset the environmental variables 'http_proxy' and 'https_proxy', and retry.
    fn test_peer_resolve() {
        let mut cfg = Config::default();
        cfg.addr = "127.0.0.1:0".to_owned();

        let storage = TestStorageBuilder::new().build().unwrap();

        let (tx, rx) = mpsc::channel();
        let (significant_msg_sender, significant_msg_receiver) = mpsc::channel();
        let router = TestRaftStoreRouter {
            tx,
            significant_msg_sender,
        };

        let quick_fail = Arc::new(AtomicBool::new(false));
        let cfg = Arc::new(cfg);
        let security_mgr = Arc::new(SecurityManager::new(&SecurityConfig::default()).unwrap());

        let cop_read_pool = readpool_impl::build_read_pool_for_test(storage.get_engine());
        let cop = coprocessor::Endpoint::new(&cfg, cop_read_pool);

        let addr = Arc::new(Mutex::new(None));
        let mut server = Server::new(
            &cfg,
            &security_mgr,
            storage,
            cop,
            router,
            MockResolver {
                quick_fail: Arc::clone(&quick_fail),
                addr: Arc::clone(&addr),
            },
            SnapManager::new("", None),
            None,
            None,
            None,
        )
        .unwrap();

        server.start(cfg, security_mgr).unwrap();

        let mut trans = server.transport();
        trans.report_unreachable(RaftMessage::new());
        let mut resp = significant_msg_receiver.try_recv().unwrap();
        assert!(is_unreachable_to(&resp, 0, 0), "{:?}", resp);

        let mut msg = RaftMessage::new();
        msg.set_region_id(1);
        trans.send(msg.clone()).unwrap();
        trans.flush();
        resp = significant_msg_receiver.try_recv().unwrap();
        assert!(is_unreachable_to(&resp, 1, 0), "{:?}", resp);

        *addr.lock().unwrap() = Some(format!("{}", server.listening_addr()));

        trans.send(msg.clone()).unwrap();
        trans.flush();
        assert!(rx.recv_timeout(Duration::from_secs(5)).is_ok());

        msg.mut_to_peer().set_store_id(2);
        msg.set_region_id(2);
        quick_fail.store(true, Ordering::SeqCst);
        trans.send(msg.clone()).unwrap();
        trans.flush();
        resp = significant_msg_receiver.try_recv().unwrap();
        assert!(is_unreachable_to(&resp, 2, 0), "{:?}", resp);
        server.stop().unwrap();
    }
}
<commit_msg>server: remove the limit about received message size (#4992)<commit_after>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::i32;
use std::net::{IpAddr, SocketAddr};
use std::str::FromStr;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};

use engine::Engines;
use futures::{Future, Stream};
use grpcio::{ChannelBuilder, EnvBuilder, Environment, Server as GrpcServer, ServerBuilder};
use kvproto::debugpb_grpc::create_debug;
use kvproto::import_sstpb_grpc::create_import_sst;
use kvproto::tikvpb_grpc::*;
use tokio_threadpool::{Builder as ThreadPoolBuilder, ThreadPool};
use tokio_timer::timer::Handle;

use crate::storage::lock_manager::deadlock::Service as DeadlockService;
use kvproto::deadlock_grpc::create_deadlock;

use crate::coprocessor::Endpoint;
use crate::import::ImportSSTService;
use crate::raftstore::store::SnapManager;
use crate::storage::{Engine, Storage};
use tikv_util::security::SecurityManager;
use tikv_util::timer::GLOBAL_TIMER_HANDLE;
use tikv_util::worker::Worker;
use tikv_util::Either;

use super::load_statistics::*;
use super::raft_client::RaftClient;
use super::resolve::StoreAddrResolver;
use super::service::*;
use super::snap::{Runner as SnapHandler, Task as SnapTask};
use super::transport::{RaftStoreRouter, ServerTransport};
use super::{Config, Result};

const LOAD_STATISTICS_SLOTS: usize = 4;
const LOAD_STATISTICS_INTERVAL: Duration = Duration::from_millis(100);
pub const GRPC_THREAD_PREFIX: &str = "grpc-server";
pub const STATS_THREAD_PREFIX: &str = "transport-stats";

/// The TiKV server
///
/// It hosts various internal components, including gRPC, the raftstore router
/// and a snapshot worker.
pub struct Server<T: RaftStoreRouter + 'static, S: StoreAddrResolver + 'static> {
    env: Arc<Environment>,
    /// A GrpcServer builder or a GrpcServer.
    ///
    /// If the listening port is configured, the server will be started lazily.
    builder_or_server: Option<Either<ServerBuilder, GrpcServer>>,
    local_addr: SocketAddr,
    // Transport.
    trans: ServerTransport<T, S>,
    raft_router: T,
    // For sending/receiving snapshots.
    snap_mgr: SnapManager,
    snap_worker: Worker<SnapTask>,

    // Currently load statistics is done in the thread.
    stats_pool: Option<ThreadPool>,
    thread_load: Arc<ThreadLoad>,
    timer: Handle,
}

impl<T: RaftStoreRouter, S: StoreAddrResolver + 'static> Server<T, S> {
    #[allow(clippy::too_many_arguments)]
    pub fn new<E: Engine>(
        cfg: &Arc<Config>,
        security_mgr: &Arc<SecurityManager>,
        storage: Storage<E>,
        cop: Endpoint<E>,
        raft_router: T,
        resolver: S,
        snap_mgr: SnapManager,
        debug_engines: Option<Engines>,
        import_service: Option<ImportSSTService<T>>,
        deadlock_service: Option<DeadlockService>,
    ) -> Result<Self> {
        // A helper thread (or pool) for transport layer.
        let stats_pool = ThreadPoolBuilder::new()
            .pool_size(cfg.stats_concurrency)
            .name_prefix(STATS_THREAD_PREFIX)
            .build();
        let thread_load = Arc::new(ThreadLoad::with_threshold(cfg.heavy_load_threshold));

        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(cfg.grpc_concurrency)
                .name_prefix(thd_name!(GRPC_THREAD_PREFIX))
                .build(),
        );
        let snap_worker = Worker::new("snap-handler");

        let kv_service = KvService::new(
            storage,
            cop,
            raft_router.clone(),
            snap_worker.scheduler(),
            Arc::clone(&thread_load),
        );

        let mut addr = SocketAddr::from_str(&cfg.addr)?;
        let ip = format!("{}", addr.ip());
        let channel_args = ChannelBuilder::new(Arc::clone(&env))
            .stream_initial_window_size(cfg.grpc_stream_initial_window_size.0 as i32)
            .max_concurrent_stream(cfg.grpc_concurrent_stream)
            .max_receive_message_len(-1)
            .max_send_message_len(-1)
            .http2_max_ping_strikes(i32::MAX) // For pings without data from clients.
            .build_args();
        let builder_or_server = {
            let mut sb = ServerBuilder::new(Arc::clone(&env))
                .channel_args(channel_args)
                .register_service(create_tikv(kv_service));
            sb = security_mgr.bind(sb, &ip, addr.port());
            if let Some(engines) = debug_engines {
                let debug_service = DebugService::new(engines, raft_router.clone());
                sb = sb.register_service(create_debug(debug_service));
            }
            if let Some(service) = import_service {
                sb = sb.register_service(create_import_sst(service));
            }
            if let Some(service) = deadlock_service {
                sb = sb.register_service(create_deadlock(service));
            }
            // When port is 0, it has to be binded now to get a valid address, which
            // is then reported to PD before the server is up. 0 is usually used in tests.
            if addr.port() == 0 {
                let server = sb.build()?;
                let (ref host, port) = server.bind_addrs()[0];
                addr = SocketAddr::new(IpAddr::from_str(host)?, port as u16);
                Either::Right(server)
            } else {
                Either::Left(sb)
            }
        };

        info!("listening on addr"; "addr" => addr);

        let raft_client = Arc::new(RwLock::new(RaftClient::new(
            Arc::clone(&env),
            Arc::clone(cfg),
            Arc::clone(security_mgr),
            raft_router.clone(),
            Arc::clone(&thread_load),
            stats_pool.sender().clone(),
        )));

        let trans = ServerTransport::new(
            raft_client,
            snap_worker.scheduler(),
            raft_router.clone(),
            resolver,
        );

        let svr = Server {
            env: Arc::clone(&env),
            builder_or_server: Some(builder_or_server),
            local_addr: addr,
            trans,
            raft_router,
            snap_mgr,
            snap_worker,
            stats_pool: Some(stats_pool),
            thread_load,
            timer: GLOBAL_TIMER_HANDLE.clone(),
        };

        Ok(svr)
    }

    pub fn transport(&self) -> ServerTransport<T, S> {
        self.trans.clone()
    }

    /// Starts the TiKV server.
    pub fn start(&mut self, cfg: Arc<Config>, security_mgr: Arc<SecurityManager>) -> Result<()> {
        let snap_runner = SnapHandler::new(
            Arc::clone(&self.env),
            self.snap_mgr.clone(),
            self.raft_router.clone(),
            security_mgr,
            Arc::clone(&cfg),
        );
        box_try!(self.snap_worker.start(snap_runner));
        let builder_or_server = self.builder_or_server.take().unwrap();
        let mut grpc_server = match builder_or_server {
            Either::Left(builder) => builder.build()?,
            Either::Right(server) => server,
        };
        grpc_server.start();
        self.builder_or_server = Some(Either::Right(grpc_server));

        let mut load_stats = {
            let tl = Arc::clone(&self.thread_load);
            ThreadLoadStatistics::new(LOAD_STATISTICS_SLOTS, GRPC_THREAD_PREFIX, tl)
        };
        self.stats_pool.as_ref().unwrap().spawn(
            self.timer
                .interval(Instant::now(), LOAD_STATISTICS_INTERVAL)
                .map_err(|_| ())
                .for_each(move |i| {
                    load_stats.record(i);
                    Ok(())
                }),
        );

        info!("TiKV is ready to serve");
        Ok(())
    }

    /// Stops the TiKV server.
    pub fn stop(&mut self) -> Result<()> {
        self.snap_worker.stop();
        if let Some(Either::Right(mut server)) = self.builder_or_server.take() {
            server.shutdown();
        }
        if let Some(pool) = self.stats_pool.take() {
            let _ = pool.shutdown_now().wait();
        }
        Ok(())
    }

    // Return listening address, this may only be used for outer test
    // to get the real address because we may use "127.0.0.1:0"
    // in test to avoid port conflict.
    pub fn listening_addr(&self) -> SocketAddr {
        self.local_addr
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::*;
    use std::sync::mpsc::*;
    use std::sync::*;
    use std::time::Duration;

    use super::*;

    use super::super::resolve::{Callback as ResolveCallback, StoreAddrResolver};
    use super::super::transport::RaftStoreRouter;
    use super::super::{Config, Result};
    use crate::coprocessor;
    use crate::raftstore::store::transport::Transport;
    use crate::raftstore::store::*;
    use crate::raftstore::Result as RaftStoreResult;
    use crate::storage::TestStorageBuilder;

    use crate::coprocessor::readpool_impl;
    use kvproto::raft_cmdpb::RaftCmdRequest;
    use kvproto::raft_serverpb::RaftMessage;
    use tikv_util::security::SecurityConfig;

    #[derive(Clone)]
    struct MockResolver {
        quick_fail: Arc<AtomicBool>,
        addr: Arc<Mutex<Option<String>>>,
    }

    impl StoreAddrResolver for MockResolver {
        fn resolve(&self, _: u64, cb: ResolveCallback) -> Result<()> {
            if self.quick_fail.load(Ordering::SeqCst) {
                return Err(box_err!("quick fail"));
            }
            let addr = self.addr.lock().unwrap();
            cb(addr
                .as_ref()
                .map(|s| s.to_owned())
                .ok_or(box_err!("not set")));
            Ok(())
        }
    }

    #[derive(Clone)]
    struct TestRaftStoreRouter {
        tx: Sender<usize>,
        significant_msg_sender: Sender<SignificantMsg>,
    }

    impl RaftStoreRouter for TestRaftStoreRouter {
        fn send_raft_msg(&self, _: RaftMessage) -> RaftStoreResult<()> {
            self.tx.send(1).unwrap();
            Ok(())
        }

        fn send_command(&self, _: RaftCmdRequest, _: Callback) -> RaftStoreResult<()> {
            self.tx.send(1).unwrap();
            Ok(())
        }

        fn significant_send(&self, _: u64, msg: SignificantMsg) -> RaftStoreResult<()> {
            self.significant_msg_sender.send(msg).unwrap();
            Ok(())
        }

        fn casual_send(&self, _: u64, _: CasualMessage) -> RaftStoreResult<()> {
            self.tx.send(1).unwrap();
            Ok(())
        }

        fn broadcast_unreachable(&self, _: u64) {
            let _ = self.tx.send(1);
        }
    }

    fn is_unreachable_to(msg: &SignificantMsg, region_id: u64, to_peer_id: u64) -> bool {
        *msg == SignificantMsg::Unreachable {
            region_id,
            to_peer_id,
        }
    }

    #[test]
    // if this failed, unset the environmental variables 'http_proxy' and 'https_proxy', and retry.
    fn test_peer_resolve() {
        let mut cfg = Config::default();
        cfg.addr = "127.0.0.1:0".to_owned();

        let storage = TestStorageBuilder::new().build().unwrap();

        let (tx, rx) = mpsc::channel();
        let (significant_msg_sender, significant_msg_receiver) = mpsc::channel();
        let router = TestRaftStoreRouter {
            tx,
            significant_msg_sender,
        };

        let quick_fail = Arc::new(AtomicBool::new(false));
        let cfg = Arc::new(cfg);
        let security_mgr = Arc::new(SecurityManager::new(&SecurityConfig::default()).unwrap());

        let cop_read_pool = readpool_impl::build_read_pool_for_test(storage.get_engine());
        let cop = coprocessor::Endpoint::new(&cfg, cop_read_pool);

        let addr = Arc::new(Mutex::new(None));
        let mut server = Server::new(
            &cfg,
            &security_mgr,
            storage,
            cop,
            router,
            MockResolver {
                quick_fail: Arc::clone(&quick_fail),
                addr: Arc::clone(&addr),
            },
            SnapManager::new("", None),
            None,
            None,
            None,
        )
        .unwrap();

        server.start(cfg, security_mgr).unwrap();

        let mut trans = server.transport();
        trans.report_unreachable(RaftMessage::new());
        let mut resp = significant_msg_receiver.try_recv().unwrap();
        assert!(is_unreachable_to(&resp, 0, 0), "{:?}", resp);

        let mut msg = RaftMessage::new();
        msg.set_region_id(1);
        trans.send(msg.clone()).unwrap();
        trans.flush();
        resp = significant_msg_receiver.try_recv().unwrap();
        assert!(is_unreachable_to(&resp, 1, 0), "{:?}", resp);

        *addr.lock().unwrap() = Some(format!("{}", server.listening_addr()));

        trans.send(msg.clone()).unwrap();
        trans.flush();
        assert!(rx.recv_timeout(Duration::from_secs(5)).is_ok());

        msg.mut_to_peer().set_store_id(2);
        msg.set_region_id(2);
        quick_fail.store(true, Ordering::SeqCst);
        trans.send(msg.clone()).unwrap();
        trans.flush();
        resp = significant_msg_receiver.try_recv().unwrap();
        assert!(is_unreachable_to(&resp, 2, 0), "{:?}", resp);
        server.stop().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::io;
use std::mem;
use std::num;
use std::fs;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::io::Read;
use std::io::Write;
use std::io::BufRead;
use std::usize;

struct CooccurRec {
    word1: u32,
    word2: u32,
    val: f32,
}

fn get_filename(file_head: &str, n: i32) -> String {
    return format!("{}_{:>04}.bin", file_head, n);
}

fn get_word(fin: &io::Read) -> Option<String> {
    let mut word: Vec<u8> = vec![];
    let mut byte: [u8; 1];
    while fin.read_exact(&mut byte).is_ok() {
        match byte[0] {
            b' ' | b'\t' | b'\n' => { break; }
            c => { word.push(c); }
        }
    }
    if word.len() > 0 { String::from_utf8(word).ok() } else { None }
}

fn get_cooccurrence(verbose: i32, symmetric: bool, window_size: i32, max_product: usize, overflow_length: usize,
                    vocab_file: &str, file_head: &str) -> i32 {
    macro_rules! progress {
        ($v:expr, $fmt:expr) => (
            if verbose > $v { write!(io::stderr(), $fmt).unwrap(); }
        );
        ($v:expr, $fmt:expr, $($arg:tt)*) => (
            if verbose > $v { write!(io::stderr(), $fmt, $($arg)*).unwrap(); }
        );
    }
    macro_rules! log {
        ($v:expr, $fmt:expr) => (
            if verbose > $v { writeln!(io::stderr(), $fmt).unwrap(); }
        );
        ($v:expr, $fmt:expr, $($arg:tt)*) => (
            if verbose > $v { writeln!(io::stderr(), $fmt, $($arg)*).unwrap(); }
        );
    }
    macro_rules! min {
        ($a: expr, $b: expr) => { if $a <= $b { $a } else { $b } };
    }
    macro_rules! max {
        ($a: expr, $b: expr) => { if $a >= $b { $a } else { $b } };
    }

    log!(-1, "COUNTING COOCCURRENCES");
    log!(0, "window size: {}", window_size);
    log!(0, "context: {}", if symmetric { "symmetric" } else { "asymmetric" });
    log!(1, "max product: {}", max_product);
    log!(1, "overflow length: {}", overflow_length);

    let mut vocab_hash: HashMap<String, i64> = HashMap::new();
    {
        log!(1, "Reading vocab from file \"{}\"...", vocab_file);
        let file = io::BufReader::new(fs::File::open(vocab_file).expect("Unable to open vocab file."));
        let mut vocab_rank = 1i64;
        for line in file.lines() {
            let line = line.expect("vocab file read error.");
            let vec: Vec<&str> = line.split_whitespace().collect();
            let word = vec[0];
            //let num = vec[1].parse::<u32>().expect("Parse error.");

            vocab_hash.insert(word.to_string(), vocab_rank);
            vocab_rank += 1;
        }
        log!(1, "loaded {} words.\nBuilding lookup table...", vocab_hash.len());
    }
    let mut table: Vec<Vec<f32>> = Vec::with_capacity(vocab_hash.len());
    {
        for a in 0..vocab_hash.len() {
            table.push(vec![0.0_f32 ; min!(max_product / a, vocab_hash.len())]);
        }
    }
    {
        let mut fcounter = 0i32;
        let mut wcounter = 0i64;
        let mut cr: Vec<CooccurRec> = Vec::with_capacity(overflow_length + 1);
        let mut history = vec![0i64 ; window_size];
        let stdin = io::BufReader::new(io::stdin());
        'outer: loop {  // for each file.
            progress!(1, "Processing token: 0");
            let mut file = fs::File::open(get_filename(file_head, fcounter)).expect("File open error.");
            loop {  // for each line.
                let mut n_words = 0;
                while let Some(word) = get_word(&stdin) {
                    wcounter += 1;
                    if counter % 100000 == 0 { progress!(1, "\x1b[19G{}", wcounter); }
                    if let Some(w2) = vocab_hash.get(word) {
                        for k in max!(wcounter - window_size, 0) .. (wcounter - 1).rev() {
                            let w1 = history[k % window_size];
                            if w1 < max_product / w2 {  // Product is small enough to store in a full array
                                table[w1 - 1][w2 - 2] += (1.0 / (wcounter - k) as f64) as f32;
                                if symmetric { table[w2 - 1][w1 - 2] += (1.0 / (wcounter - k) as f64) as f32; }
                            }
                            else {  // Product is too big, data is likely to be sparse. Store these entries in a temporary buffer to be sorted, mered (accumulated), and written to file when it gets full.
                                cr.push(CooccurRec { word1: w1, word2: w2, val: get_val(?)});
                                if symmetric { cr.push(CooccurRec { word1: w2, word2: w1, val: get_val(?)}); }
                            }
                            history[wcounter % window_size] = w2;
                            // TODO: write here...
                        }
                    }
                }
                if n_words == 0 { break; }
            }
            // TODO:
            // sort cr
            // write cr
        }
        // TODO: write out ...
    }
    0
}

fn main() {
    let mut verbose = 1i32;
    let mut window_size = 15i32;
    let mut symmetric = true;
    let mut memory_limit = 3.0;
    let mut vocab_file = "vocab.txt".to_string();
    let mut file_head = "overflow".to_string();
    let mut max_product: Option<usize> = None;
    let mut overflow_length: Option<usize> = None;
    {
        let args: Vec<String> = env::args().collect();

        if args.len() == 1 {
            println!("Tool to calculate word-word cooccurrence statistics
This program is transportation of GloVe cooccur.
Original program (written in C) author is Jeffrey Pennington (jpennin@stanford.edu)
Transpoted to rust by Takumi Kato (takumi.kt@gmail.com)

Usage options:
\t-verbose <int>
\t\tSet verbosity: 0, 1, or 2 (default)
\t-symmetric <int>
\t\tIf <int> = 0, only use left context; if <int> = 1 (default), use left and right
\t-window-size <int>
\t\tNumber of context words to the left (and to the right, if symmetric = 1); default 15
\t-vocab-file <file>
\t\tFile containing vocabulary (truncated unigram counts, produced by 'vocab_count'); default vocab.txt
\t-memory <float>
\t\tSoft limit for memory consumption, in GB -- based on simple heuristic, so not extremely accurate; default 4.0
\t-max-product <int>
\t\tLimit the size of dense cooccurrence array by specifying the max product <int> of the frequency counts of the two cooccurring words.
\t\tThis value overrides that which is automatically produced by '-memory'. Typically only needs adjustment for use with very large corpora.
\toverflow-file <file>
\t\tFilename, excluding extension, for temporary file; default overflow

Example usage:
./cooccur -verbose 2 -symmetric 0 -window-size 10 -vocab-file vocab.txt -memory 8.0 -overflow-file tempoverflow < corpus.txt > cooccurrences.bin

");
            std::process::exit(0);
        }
        let mut is_skip = true;
        let mut it = args.into_iter().peekable();
        loop {
            let arg = match it.next() {
                Some(s) => s,
                None => break,
            };
            if is_skip {
                is_skip = false;
                continue
            }
            match arg.as_ref() {
                "-verbose" => {
                    verbose = it.peek().and_then(|v| v.parse().ok())
                        .and_then(|x| if 0 <= x && x <= 2 { Some(x) } else { None })
                        .expect("-verbose <int>, verbosity: 0, 1, or 2 (default)");
                    is_skip = true;
                },
                "-symmetric" => {
                    symmetric = it.peek().and_then(|v| v.parse().ok())
                        .and_then(|x| if 0 <= x && x <= 1 { Some(x) } else { None })
                        .expect("symmetric <int>, 0 (only use left context) or 1 (use left and right; default)") == 1;
                    is_skip = true;
                },
                "-window-size" => {
                    window_size = it.peek().and_then(|v| v.parse().ok()).expect("-window-size <int>");
                    is_skip = true;
                },
                "-memory" => {
                    memory_limit = it.peek().and_then(|v| v.parse().ok()).expect("-memory <float>");
                    is_skip = true;
                },
                "-vocab-file" => {
                    vocab_file = it.peek().expect("-vocab-file <file>").to_string();
                    is_skip = true;
                },
                "-overflow" => {
                    file_head = it.peek().expect("-overflow <file>").to_string();
                    is_skip = true;
                },
                "-max-product" => {
                    max_product = Some(it.peek().and_then(|v| v.parse().ok()).expect("-max-product <int>"));
                    is_skip = true;
                },
                "-overflow-length" => {
                    overflow_length = Some(it.peek().and_then(|v| v.parse().ok()).expect("-overflow-length <int>"));
                    is_skip = true;
                },
                &_ => {},
            }
            let rlimit = 0.85 * memory_limit * 1073741824_f64 / mem::size_of::<f64>() as f64;
            let mut n = 1e5_f64;
            while (rlimit - n * (n.ln() + 0.1544313298)).abs() > 1e-3 { n = rlimit / (n.ln() + 0.1544313298) }
            max_product = max_product.or_else(|| Some(n as usize));
            overflow_length = overflow_length.or_else(|| Some((rlimit / 6.0) as usize));
        }
        //std::process::exit(get_counts(verbose, max_vocab, min_count));
    }
}
<commit_msg>Update cooccur (still not working again)<commit_after>use std::{env, io, mem, num, i32, usize, fs, slice, string};
use std::io::{Read, Write, BufRead, Error, ErrorKind};
use std::cmp::Ordering;
use std::collections::{HashMap,VecDeque};

#[derive(Copy, Clone)]
struct CooccurRec {
    word1: u32,
    word2: u32,
    val: f32,
}

static mut verbose: i32 = 2;

macro_rules! progress {
    ($v:expr, $fmt:expr) => (
        if verbose > $v { write!(io::stderr(), $fmt).unwrap(); }
    );
    ($v:expr, $fmt:expr, $($arg:tt)*) => (
        if verbose > $v { write!(io::stderr(), $fmt, $($arg)*).unwrap(); }
    );
}
macro_rules! log {
    ($v:expr, $fmt:expr) => (
        if verbose > $v { writeln!(io::stderr(), $fmt).unwrap(); }
    );
    ($v:expr, $fmt:expr, $($arg:tt)*) => (
        if verbose > $v { writeln!(io::stderr(), $fmt, $($arg)*).unwrap(); }
    );
}

enum GetWord {
    Word(String),
    EndLine,
    EndOfFile,
    IOError(io::Error),
    Utf8Error(std::string::FromUtf8Error),
}

fn get_word(fin: &io::Read) -> GetWord {
    static mut next_endline: bool = false;
    if next_endline {
        next_endline = false;
        return GetWord::EndLine
    }

    let mut eof: bool = false;
    let mut word: Vec<u8> = vec![];
    let mut byte: [u8; 1];
    loop {
        match fin.read_exact(&mut byte) {
            Ok(_) => match byte[0] {
                b' ' | b'\t' => { break; },
                b'\n' => {
                    next_endline = true;
                    break;
                },
                c => { word.push(c); }
            },
            Err(e) => if e.kind() == ErrorKind::UnexpectedEof {
                eof = true;
                break;
            } else {
                return GetWord::IOError(e);
            }
        }
    }
    if word.len() > 0 {
        match String::from_utf8(word) {
            Ok(x) => GetWord::Word(x),
            Err(e) => GetWord::Utf8Error(e),
        }
    } else if eof {
        GetWord::EndOfFile
    } else {
        get_word(fin)
    }
}

fn merge_files(file_head: &str, num: i32) {
    progress!(1, "Merging cooccurrence files: processed 0 lines.");
    let pq = VecDeque::<(CooccurRec, i32)>::new();
    let fid = Vec::<fs::File>::new();
    for i in 0..num {
        let mut data: CooccurRec;
        let f = fs::File::open(format!("{}_{:>4}.bin", file_head, i)).unwrap();
        f.read_exact(unsafe {
            slice::from_raw_parts_mut((&mut data as *mut CooccurRec) as *mut u8,
            mem::size_of::<CooccurRec>())}).unwrap();
        pq.push_back((data, i));
        fid.push(f);
    }
    // Pop op node, save it in old to see if the next entry is a duplicate
    let mut size = num as usize;
    let (mut old, mut i) = match pq.pop_front().unwrap();

}

fn get_cooccurrence(symmetric: bool, window_size: usize,
                    max_product: usize, overflow_length: usize,
                    vocab_file: &str, file_head: &str) -> i32 {
    macro_rules! min {
        ($a: expr, $b: expr) => { if $a <= $b { $a } else { $b } };
    }
    macro_rules! max {
        ($a: expr, $b: expr) => { if $a >= $b { $a } else { $b } };
    }

    log!(-1, "COUNTING COOCCURRENCES");
    log!(0, "window size: {}", window_size);
    log!(0, "context: {}", if symmetric { "symmetric" } else { "asymmetric" });
    log!(1, "max product: {}", max_product);
    log!(1, "overflow length: {}", overflow_length);

    let mut vocab_hash: HashMap<String, i64> = HashMap::new();
    log!(1, "Reading vocab from file \"{}\"...", vocab_file);
    {
        let file = io::BufReader::new(fs::File::open(vocab_file)
                                      .expect("Unable to open vocab file."));
        for (i, line) in file.lines().enumerate() {
            let line = line.expect("vocab file read error.");
            let vec: Vec<&str> = line.split_whitespace().collect();
            let word = vec[0];
            //let num = vec[1].parse::<u32>().expect("Parse error.");

            vocab_hash.insert(word.to_string(), i as i64);
        }
        log!(1, "loaded {} words.", vocab_hash.len());
    }
    let mut table: Vec<Vec<f32>> = Vec::with_capacity(vocab_hash.len());
    {
        let mut size = 1;
        table.push(vec![0.0_f32]);
        for a in 1..vocab_hash.len() {
            size += min!(max_product / a, vocab_hash.len());
            table.push(vec![0.0_f32 ; size]);
        }
    }
    // for each token in input stream, calculate a weighted cooccurrence sum within window_size.
    let value = |j, k| 1.0f32 / (j - k) as f32;
    let write_chunk = |cr: &Vec<CooccurRec>, n| {
        let mut f = fs::File::create(format!("{}_{:>04}.bin", file_head, n)).unwrap();
        let mut old = cr[0];

        for x in &cr[1..] {
            if x.word1 == old.word1 && x.word2 == old.word2 {
                old.val += x.val;
                continue;
            }
            f.write(unsafe {
                slice::from_raw_parts((&old as *const CooccurRec) as *const u8,
                mem::size_of::<CooccurRec>())}).unwrap();
            old = *x;
        }
        f.write(unsafe {
            slice::from_raw_parts((&old as *const CooccurRec) as *const u8,
            mem::size_of::<CooccurRec>())}).unwrap();
    };

    let mut fcounter = 1i32;
    let mut cr: Vec<CooccurRec> = Vec::with_capacity(overflow_length + 1);
    let mut history = vec![0i64 ; window_size as usize];
    let mut n_words = 0usize;
    let mut j = 0i64;
    let stdin = io::BufReader::new(io::stdin());
    progress!(1, "Processing token: 0");
    loop {
        let word = match get_word(&stdin) {
            GetWord::Word(w) => w,
            GetWord::EndLine => { j = 0; continue },
            GetWord::EndOfFile => break,
            GetWord::Utf8Error(_) => continue,
            GetWord::IOError(e) => panic!(format!("{:?}", e)),
        };
        j += 1;
        n_words += 1;
        if n_words % 100000 == 0 { progress!(1, "\x1b[19G{}", j); }
        if let Some(&w2) = vocab_hash.get(&word) {
            for k in (max!(j - window_size as i64, 0) .. j - 1).rev() {
                let w1 = history[k as usize % window_size];
                if (w1 as usize) < max_product / (w2 as usize) {  // Product is small enough to store in a full array
                    table[(w1 - 1) as usize][(w2 - 2) as usize] += value(j, k);
                    if symmetric { table[(w2 - 1) as usize][(w1 - 2) as usize] += value(j, k); }
                }
                else {  // Product is too big, data is likely to be sparse. Store these entries in a temporary buffer to be sorted, merged (accumulated), and written to file when it gets full.
                    cr.push(CooccurRec { word1: w1 as u32, word2: w2 as u32, val: value(j, k)});
                    if symmetric { cr.push(CooccurRec { word1: w2 as u32, word2: w1 as u32, val: value(j, k)}); }
                    if cr.len() >= overflow_length - window_size {
                        cr.sort_by(|lhs, rhs| match lhs.word1.cmp(&rhs.word1) {
                            Ordering::Equal => lhs.word2.cmp(&rhs.word2),
                            x => x,
                        });
                        write_chunk(&cr, fcounter);
                        fcounter += 1;
                        cr.clear();
                    }
                }
                history[j as usize % window_size] = w2;
                j += 1;
            }
        }
    }
    log!(1, "\x1b[0GProcessed {} tokens.", n_words);
    cr.sort_by(|lhs, rhs| match lhs.word1.cmp(&rhs.word1) {
        Ordering::Equal => lhs.word2.cmp(&rhs.word2),
        x => x,
    });
    write_chunk(&cr, fcounter);
    log!(1, "Writing cooccurrences to disk");
    let j = i32::MAX;
    let file = fs::File::create(format!("{}_0000.bin", file_head)).unwrap();
    for (x, &v) in table.iter().enumerate() {
        if ((0.75f64 * (vocab_hash.len() / x) as f64).ln() as i32) < j {
            progress!(1, ".")
        }
        for (y, &r) in v.iter().enumerate() {
            if r != 0f32 {
                file.write(unsafe {mem::transmute::<u32, &[u8; 4]>(x as u32)});
                file.write(unsafe {mem::transmute::<u32, &[u8; 4]>(y as u32)});
                file.write(unsafe {mem::transmute::<f32, &[u8; 4]>(r)});
            }
        }
    }
    log!(1, "{} files in total.", fcounter);
    merge_files(file_head, fcounter + 1);
    0
}

fn main() {
    let mut window_size = 15usize;
    let mut symmetric = true;
    let mut memory_limit = 3.0;
    let mut vocab_file = "vocab.txt".to_string();
    let mut file_head = "overflow".to_string();
    let mut max_product: Option<usize> = None;
    let mut overflow_length: Option<usize> = None;
    {
        let args: Vec<String> = env::args().collect();

        if args.len() == 1 {
            println!("Tool to calculate word-word cooccurrence statistics
This program is transportation of GloVe cooccur.
Original program (written in C) author is Jeffrey Pennington (jpennin@stanford.edu)
Transpoted to rust by Takumi Kato (takumi.kt@gmail.com)

Usage options:
\t-verbose <int>
\t\tSet verbosity: 0, 1, or 2 (default)
\t-symmetric <int>
\t\tIf <int> = 0, only use left context; if <int> = 1 (default), use left and right
\t-window-size <int>
\t\tNumber of context words to the left (and to the right, if symmetric = 1); default 15
\t-vocab-file <file>
\t\tFile containing vocabulary (truncated unigram counts, produced by 'vocab_count'); default vocab.txt
\t-memory <float>
\t\tSoft limit for memory consumption, in GB -- based on simple heuristic, so not extremely accurate; default 4.0
\t-max-product <int>
\t\tLimit the size of dense cooccurrence array by specifying the max product <int> of the frequency counts of the two cooccurring words.
\t\tThis value overrides that which is automatically produced by '-memory'. Typically only needs adjustment for use with very large corpora.
\toverflow-file <file>
\t\tFilename, excluding extension, for temporary file; default overflow

Example usage:
./cooccur -verbose 2 -symmetric 0 -window-size 10 -vocab-file vocab.txt -memory 8.0 -overflow-file tempoverflow < corpus.txt > cooccurrences.bin

");
            std::process::exit(0);
        }
        let mut is_skip = true;
        let mut it = args.into_iter().peekable();
        loop {
            let arg = match it.next() {
                Some(s) => s,
                None => break,
            };
            if is_skip {
                is_skip = false;
                continue
            }
            match arg.as_ref() {
                "-verbose" => {
                    verbose = it.peek().and_then(|v| v.parse().ok())
                        .and_then(|x| if 0 <= x && x <= 2 { Some(x) } else { None })
                        .expect("-verbose <int>, verbosity: 0, 1, or 2 (default)");
                    is_skip = true;
                },
                "-symmetric" => {
                    symmetric = it.peek().and_then(|v| v.parse().ok())
                        .and_then(|x| if 0 <= x && x <= 1 { Some(x) } else { None })
                        .expect("symmetric <int>, 0 (only use left context) or 1 (use left and right; default)") == 1;
                    is_skip = true;
                },
                "-window-size" => {
                    window_size = it.peek().and_then(|v| v.parse().ok()).expect("-window-size <int>");
                    is_skip = true;
                },
                "-memory" => {
                    memory_limit = it.peek().and_then(|v| v.parse().ok()).expect("-memory <float>");
                    is_skip = true;
                },
                "-vocab-file" => {
                    vocab_file = it.peek().expect("-vocab-file <file>").to_string();
                    is_skip = true;
                },
                "-overflow" => {
                    file_head = it.peek().expect("-overflow <file>").to_string();
                    is_skip = true;
                },
                "-max-product" => {
                    max_product = Some(it.peek().and_then(|v| v.parse().ok()).expect("-max-product <int>"));
                    is_skip = true;
                },
                "-overflow-length" => {
                    overflow_length = Some(it.peek().and_then(|v| v.parse().ok()).expect("-overflow-length <int>"));
                    is_skip = true;
                },
                &_ => {},
            }
            let rlimit = 0.85 * memory_limit * 1073741824_f64 / mem::size_of::<f64>() as f64;
            let mut n = 1e5_f64;
            while (rlimit - n * (n.ln() + 0.1544313298)).abs() > 1e-3 { n = rlimit / (n.ln() + 0.1544313298) }
            max_product = max_product.or_else(|| Some(n as usize));
            overflow_length = overflow_length.or_else(|| Some((rlimit / 6.0) as usize));
        }
        //std::process::exit(get_counts(verbose, max_vocab, min_count));
    }
}
<|endoftext|>


<commit_before>use crate::discord_client::CommandHandler;
use rusqlite::Connection;
use rusqlite::NO_PARAMS;

trait TagRepository {
    fn create_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>>;
    fn read_tag(&self, tag_name: &str) -> Result<String, Box<std::error::Error>>;
    fn update_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>>;
    fn delete_tag(&self, tag_name: &str) -> Result<(), Box<std::error::Error>>;
}

pub struct TagCommand {
    tag_repository: Box<TagRepository>,
}
pub struct SqliteTagRepository {
    db_connection: Connection,
}

impl TagCommand {
    fn new() -> Result<TagCommand, Box<std::error::Error>> {
        let tag_repository = Box::new(SqliteTagRepository::new()?);

        Ok(TagCommand { tag_repository })
    }
}

impl CommandHandler for TagCommand {
    fn process_command(
        &self,
        command: &str,
        send_message_callback: &Fn(&str) -> (),
    ) -> Result<(), Box<std::error::Error>> {
        if command.starts_with("ntag") {

        } else if command.starts_with("atags") {

        } else {

        }

        Ok(())
    }
}

impl SqliteTagRepository {
    fn new() -> Result<SqliteTagRepository, Box<std::error::Error>> {
        let db_connection = Connection::open("tags.db")?;

        db_connection.execute(
            "CREATE TABLE IF NOT EXISTS tags (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL UNIQUE,
            body TEXT NOT NULL
        )",
            NO_PARAMS,
        )?;

        Ok(SqliteTagRepository { db_connection })
    }
}

impl TagRepository for SqliteTagRepository {
    fn create_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>> {
        self.db_connection.execute(
            "INSERT INTO tags
            (name, body) VALUES (?1, ?1)",
            &[tag_name, tag_body],
        )?;
        Ok(())
    }

    fn read_tag(&self, tag_name: &str) -> Result<String, Box<std::error::Error>> {
        Ok(self
            .db_connection
            .prepare(
                "SELECT body
                FROM tags
                WHERE name = (?1)",
            )?
            .query_row(&[tag_name], |row| Ok(row.get(0)?))?)
    }

    fn update_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>> {
        self.db_connection.execute(
            "UPDATE tags
                SET body = (?2)
                WHERE name = (?1)",
            &[tag_name, tag_body],
        )?;

        Ok(())
    }

    fn delete_tag(&self, tag_name: &str) -> Result<(), Box<std::error::Error>> {
        self.db_connection.execute(
            "DELETE FROM tags
                WHERE name = (?1)",
            &[tag_name],
        )?;

        Ok(())
    }
}
<commit_msg>Add code to deal with parsing for creating a new tag<commit_after>use crate::discord_client::CommandHandler;
use rusqlite::Connection;
use rusqlite::NO_PARAMS;

trait TagRepository {
    fn create_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>>;
    fn read_tag(&self, tag_name: &str) -> Result<String, Box<std::error::Error>>;
    fn update_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>>;
    fn delete_tag(&self, tag_name: &str) -> Result<(), Box<std::error::Error>>;
}

pub struct TagCommand {
    tag_repository: Box<TagRepository>,
}
pub struct SqliteTagRepository {
    db_connection: Connection,
}

struct Tag {
    name: String,
    body: String,
}

fn parse_ntag(command: &str) -> Option<Tag> {
    let mut result: Option<Tag> = None;

    let mut broken_down_command = command.split_whitespace().skip(1);
    let name = broken_down_command.next();
    let body: Vec<&str> = broken_down_command.collect();
    let body: String = body.join(" ");
    let body: Option<String> = if body.is_empty() { None } else { Some(body) };

    if let Some(name) = name {
        if let Some(body) = body {
            result = Some(Tag {
                name: String::from(name),
                body,
            });
        }
    }

    result
}

impl TagCommand {
    fn new() -> Result<TagCommand, Box<std::error::Error>> {
        let tag_repository = Box::new(SqliteTagRepository::new()?);

        Ok(TagCommand { tag_repository })
    }
}

impl CommandHandler for TagCommand {
    fn process_command(
        &self,
        command: &str,
        send_message_callback: &Fn(&str) -> (),
    ) -> Result<(), Box<std::error::Error>> {
        if command.starts_with("ntag") {
            match parse_ntag(command) {
                Some(new_tag) => {
                    self.tag_repository.create_tag(&new_tag.name, &new_tag.body);
                    send_message_callback(&format!(
                        "Created new tag {} with body {}",
                        &new_tag.name, &new_tag.body
                    ));
                }
                None => send_message_callback("Syntax error, could not create tag"),
            }
        } else if command.starts_with("atags") {

        } else {

        }

        Ok(())
    }
}

impl SqliteTagRepository {
    fn new() -> Result<SqliteTagRepository, Box<std::error::Error>> {
        let db_connection = Connection::open("tags.db")?;

        db_connection.execute(
            "CREATE TABLE IF NOT EXISTS tags (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL UNIQUE,
            body TEXT NOT NULL
        )",
            NO_PARAMS,
        )?;

        Ok(SqliteTagRepository { db_connection })
    }
}

impl TagRepository for SqliteTagRepository {
    fn create_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>> {
        self.db_connection.execute(
            "INSERT INTO tags
            (name, body) VALUES (?1, ?1)",
            &[tag_name, tag_body],
        )?;
        Ok(())
    }

    fn read_tag(&self, tag_name: &str) -> Result<String, Box<std::error::Error>> {
        Ok(self
            .db_connection
            .prepare(
                "SELECT body
                FROM tags
                WHERE name = (?1)",
            )?
            .query_row(&[tag_name], |row| Ok(row.get(0)?))?)
    }

    fn update_tag(&self, tag_name: &str, tag_body: &str) -> Result<(), Box<std::error::Error>> {
        self.db_connection.execute(
            "UPDATE tags
                SET body = (?2)
                WHERE name = (?1)",
            &[tag_name, tag_body],
        )?;

        Ok(())
    }

    fn delete_tag(&self, tag_name: &str) -> Result<(), Box<std::error::Error>> {
        self.db_connection.execute(
            "DELETE FROM tags
                WHERE name = (?1)",
            &[tag_name],
        )?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_ntag() {
        let result = parse_ntag("ntag tag_name tag_body").unwrap();
        assert_eq!("tag_name", result.name);
        assert_eq!("tag_body", result.body);

        let result =
            parse_ntag("ntag tag_name long tag body with whitespace and 日本語").unwrap();
        assert_eq!("tag_name", result.name);
        assert_eq!("long tag body with whitespace and 日本語", result.body);
    }
}
<|endoftext|>


<commit_before>pub use crate::wait::NotifyReady;
use crate::{wait::Wait, Error, Result};
use std::fmt;

pub struct Confirmation<T, I = ()> {
    kind: ConfirmationKind<T, I>,
}

impl<T, I> Confirmation<T, I> {
    pub(crate) fn new(wait: Wait<T>) -> Self {
        Self {
            kind: ConfirmationKind::Wait(wait),
        }
    }

    pub(crate) fn new_error(error: Error) -> Self {
        let (wait, wait_handle) = Wait::new();
        wait_handle.error(error);
        Self::new(wait)
    }

    // FIXME: remove
    pub(crate) fn into_error(self) -> Result<()> {
        self.try_wait().transpose().map(|_| ())
    }

    pub fn subscribe(&self, task: Box<dyn NotifyReady + Send>) {
        match &self.kind {
            ConfirmationKind::Wait(wait) => wait.subscribe(task),
            ConfirmationKind::Map(wait, _) => wait.subscribe(task),
        }
    }

    pub fn has_subscriber(&self) -> bool {
        match &self.kind {
            ConfirmationKind::Wait(wait) => wait.has_subscriber(),
            ConfirmationKind::Map(wait, _) => wait.has_subscriber(),
        }
    }

    pub fn try_wait(&self) -> Option<Result<T>> {
        match &self.kind {
            ConfirmationKind::Wait(wait) => wait.try_wait(),
            ConfirmationKind::Map(wait, f) => wait.try_wait().map(|res| res.map(f)),
        }
    }

    pub fn wait(self) -> Result<T> {
        match self.kind {
            ConfirmationKind::Wait(wait) => wait.wait(),
            ConfirmationKind::Map(wait, f) => wait.wait().map(f),
        }
    }
}

impl<T> Confirmation<T> {
    pub(crate) fn map<M>(self, f: Box<dyn Fn(T) -> M + Send + 'static>) -> Confirmation<M, T> {
        Confirmation {
            kind: ConfirmationKind::Map(Box::new(self), f),
        }
    }
}

enum ConfirmationKind<T, I> {
    Wait(Wait<T>),
    Map(Box<Confirmation<I>>, Box<dyn Fn(I) -> T + Send + 'static>),
}

impl<T> From<Result<Wait<T>>> for Confirmation<T> {
    fn from(res: Result<Wait<T>>) -> Self {
        match res {
            Ok(wait) => Confirmation::new(wait),
            Err(err) => Confirmation::new_error(err),
        }
    }
}

impl<T> fmt::Debug for Confirmation<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Confirmation")
    }
}

#[cfg(feature = "futures")]
pub(crate) mod futures {
    use super::*;

    use std::{
        future::Future,
        pin::Pin,
        task::{Context, Poll, Waker},
    };

    impl<T, I> Future for Confirmation<T, I> {
        type Output = Result<T>;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if !self.has_subscriber() {
                self.subscribe(Box::new(Watcher(cx.waker().clone())));
            }
            self.try_wait().map(Poll::Ready).unwrap_or(Poll::Pending)
        }
    }

    pub(crate) struct Watcher(pub(crate) Waker);

    impl NotifyReady for Watcher {
        fn notify(&self) {
            self.0.wake_by_ref();
        }
    }
}
<commit_msg>Warn on unused confirmation<commit_after>pub use crate::wait::NotifyReady;
use crate::{wait::Wait, Error, Result};
use std::fmt;

#[must_use = "Confirmation should be used or you can miss errors"]
pub struct Confirmation<T, I = ()> {
    kind: ConfirmationKind<T, I>,
}

impl<T, I> Confirmation<T, I> {
    pub(crate) fn new(wait: Wait<T>) -> Self {
        Self {
            kind: ConfirmationKind::Wait(wait),
        }
    }

    pub(crate) fn new_error(error: Error) -> Self {
        let (wait, wait_handle) = Wait::new();
        wait_handle.error(error);
        Self::new(wait)
    }

    // FIXME: remove
    pub(crate) fn into_error(self) -> Result<()> {
        self.try_wait().transpose().map(|_| ())
    }

    pub fn subscribe(&self, task: Box<dyn NotifyReady + Send>) {
        match &self.kind {
            ConfirmationKind::Wait(wait) => wait.subscribe(task),
            ConfirmationKind::Map(wait, _) => wait.subscribe(task),
        }
    }

    pub fn has_subscriber(&self) -> bool {
        match &self.kind {
            ConfirmationKind::Wait(wait) => wait.has_subscriber(),
            ConfirmationKind::Map(wait, _) => wait.has_subscriber(),
        }
    }

    pub fn try_wait(&self) -> Option<Result<T>> {
        match &self.kind {
            ConfirmationKind::Wait(wait) => wait.try_wait(),
            ConfirmationKind::Map(wait, f) => wait.try_wait().map(|res| res.map(f)),
        }
    }

    pub fn wait(self) -> Result<T> {
        match self.kind {
            ConfirmationKind::Wait(wait) => wait.wait(),
            ConfirmationKind::Map(wait, f) => wait.wait().map(f),
        }
    }
}

impl<T> Confirmation<T> {
    pub(crate) fn map<M>(self, f: Box<dyn Fn(T) -> M + Send + 'static>) -> Confirmation<M, T> {
        Confirmation {
            kind: ConfirmationKind::Map(Box::new(self), f),
        }
    }
}

enum ConfirmationKind<T, I> {
    Wait(Wait<T>),
    Map(Box<Confirmation<I>>, Box<dyn Fn(I) -> T + Send + 'static>),
}

impl<T> From<Result<Wait<T>>> for Confirmation<T> {
    fn from(res: Result<Wait<T>>) -> Self {
        match res {
            Ok(wait) => Confirmation::new(wait),
            Err(err) => Confirmation::new_error(err),
        }
    }
}

impl<T> fmt::Debug for Confirmation<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Confirmation")
    }
}

#[cfg(feature = "futures")]
pub(crate) mod futures {
    use super::*;

    use std::{
        future::Future,
        pin::Pin,
        task::{Context, Poll, Waker},
    };

    impl<T, I> Future for Confirmation<T, I> {
        type Output = Result<T>;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if !self.has_subscriber() {
                self.subscribe(Box::new(Watcher(cx.waker().clone())));
            }
            self.try_wait().map(Poll::Ready).unwrap_or(Poll::Pending)
        }
    }

    pub(crate) struct Watcher(pub(crate) Waker);

    impl NotifyReady for Watcher {
        fn notify(&self) {
            self.0.wake_by_ref();
        }
    }
}
<|endoftext|>