<commit_before>#[macro_use] extern crate macro_attr;
#[macro_use] extern crate enum_derive;

use std::io::Write;
use std::env;

mod keybinds;
use keybinds::*;

mod file;
use file::*;

mod terminal;
use terminal::*;

const LINENO_CHARS : i32 = 3;

fn render_file(mut out: &mut Terminal, file: &File) {
    let (screen_width, screen_height) = out.get_size();

    let mut x = 1;
    let mut y = 1;

    for (line_number_maybe, line) in file.wrapped_lines((screen_width - LINENO_CHARS - 1, screen_height - 3)) {
        if let Some(line_number) = line_number_maybe {
            out.goto((x, y));
            out.set_color_fg(Color::Grey);
            write!(out, "{:1$}",
                   line_number + 1, LINENO_CHARS as usize).unwrap();
            out.set_color_fg(Color::Reset);
        }
        x += LINENO_CHARS + 1;
        out.goto((x, y));
        write!(out, "{}", line).unwrap();
        x = 1;
        y += 1;
    }
}

fn render_footer(mut out: &mut Terminal, keys: &KeybindTable) {
    let (screen_width, screen_height) = out.get_size();

    let mut x = 1;
    let mut y = screen_height - 1;
    for (key, action) in keys.entries() {
        out.goto((x, y));
        out.set_color_fg(Color::Black);
        out.set_color_bg(Color::White);
        write!(out, "{}", key).unwrap();
        out.set_color_fg(Color::Reset);
        out.set_color_bg(Color::Reset);
        write!(out, " {}", action).unwrap();
        x += 16;
        if x > screen_width {
            x = 1;
            y -= 1;
        }
    }
}

fn render_status(mut out: &mut Terminal, file: &File) {
    let (width, height) = out.get_size();
    let file_size = (width - LINENO_CHARS - 1, height - 3);
    let x = 1;
    let y = height;
    out.goto((x, y));
    write!(out, "{}", file.debug(file_size)).unwrap();
}

fn main() {
    let keys = KeybindTable::default();
    let mut term = Terminal::default();
    term.clear();
    term.flush().unwrap();
    render_footer(&mut term, &keys);
    let args: Vec<String> = env::args().collect();
    let filename = args.get(1).cloned().unwrap_or(String::from("README.md"));
    let mut file = File::open(&filename);
    render_file(&mut term, &file);
    render_status(&mut term, &file);
    let (screen_w, screen_h) = term.get_size();
    let file_size = (screen_w - LINENO_CHARS - 1, screen_h - 3);
    term.goto((file.cursor(file_size).x + LINENO_CHARS + 1, file.cursor(file_size).y));
    term.flush().unwrap();
    let mut file_lines = file.lines.len();
    let mut file_wrapped_lines = file.wrapped_lines(file_size).len();
    let mut file_dirty = false;
    let mut screen_dirty = false;
    for evt in term.keys() {
        let (screen_w, screen_h) = term.get_size();
        let file_size = (screen_w - LINENO_CHARS - 1, screen_h - 3);
        match evt {
            Event::Mouse(_) => (),
            Event::Unsupported(_) => (),
            Event::Key(Key::Null) => (),
            Event::Key(Key::Insert) => (),
            Event::Key(Key::F(_)) => (),
            Event::Key(Key::Esc) => (),
            Event::Key(k @ Key::Ctrl(_)) | Event::Key(k @ Key::Alt(_)) => {
                match keys.lookup(k) {
                    Some(Command::Quit) => break,
                    Some(Command::Refresh) => {
                        file.refresh(file_size);
                        screen_dirty = true;
                    },
                    None => (),
                    _ => ()
                }
            },
            Event::Key(Key::Left) => {
                screen_dirty = file.move_cursor_left(file_size);
            },
            Event::Key(Key::Right) => {
                screen_dirty = file.move_cursor_right(file_size);
            },
            Event::Key(Key::Up) => {
                screen_dirty = file.move_cursor_up(file_size);
            },
            Event::Key(Key::Down) => {
                screen_dirty = file.move_cursor_down(file_size);
            },
            Event::Key(Key::Home) => file.move_cursor_home(file_size),
            Event::Key(Key::End) => file.move_cursor_end(file_size),
            Event::Key(Key::PageUp) => {
                file.page_up(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::PageDown) => {
                file.page_down(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Char('\n')) => {
                file.insert_newline(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Delete) => {
                file.delete(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Backspace) => {
                file.backspace(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Char(c)) => {
                file.insert(file_size, c);
                file_dirty = true;
            },
        }
        if screen_dirty {
            file_dirty = true;
        }
        if file_dirty {
            let new_file_lines = file.lines.len();
            let new_file_wrapped_lines = file.wrapped_lines(file_size).len();
            if new_file_lines != file_lines {
                screen_dirty = true;
                file_lines = new_file_lines;
            }
            if new_file_wrapped_lines != file_wrapped_lines {
                screen_dirty = true;
                file_wrapped_lines = new_file_wrapped_lines;
            }
        }
        if screen_dirty {
            term.clear();
            render_footer(&mut term, &keys);
            screen_dirty = false;
        }
        if file_dirty {
            render_file(&mut term, &file);
            file_dirty = false;
        }
        render_status(&mut term, &file);
        term.goto((file.cursor(file_size).x + LINENO_CHARS + 1, file.cursor(file_size).y));
        term.flush().unwrap();
    }
}
<commit_msg>Expand tabs at editor level<commit_after>#[macro_use] extern crate macro_attr;
#[macro_use] extern crate enum_derive;

use std::io::Write;
use std::env;

mod keybinds;
use keybinds::*;

mod file;
use file::*;

mod terminal;
use terminal::*;

const LINENO_CHARS : i32 = 3;

fn render_file(mut out: &mut Terminal, file: &File) {
    let (screen_width, screen_height) = out.get_size();

    let mut x = 1;
    let mut y = 1;

    for (line_number_maybe, line) in file.wrapped_lines((screen_width - LINENO_CHARS - 1, screen_height - 3)) {
        if let Some(line_number) = line_number_maybe {
            out.goto((x, y));
            out.set_color_fg(Color::Grey);
            write!(out, "{:1$}",
                   line_number + 1, LINENO_CHARS as usize).unwrap();
            out.set_color_fg(Color::Reset);
        }
        x += LINENO_CHARS + 1;
        out.goto((x, y));
        write!(out, "{}", line).unwrap();
        x = 1;
        y += 1;
    }
}

fn render_footer(mut out: &mut Terminal, keys: &KeybindTable) {
    let (screen_width, screen_height) = out.get_size();

    let mut x = 1;
    let mut y = screen_height - 1;
    for (key, action) in keys.entries() {
        out.goto((x, y));
        out.set_color_fg(Color::Black);
        out.set_color_bg(Color::White);
        write!(out, "{}", key).unwrap();
        out.set_color_fg(Color::Reset);
        out.set_color_bg(Color::Reset);
        write!(out, " {}", action).unwrap();
        x += 16;
        if x > screen_width {
            x = 1;
            y -= 1;
        }
    }
}

fn render_status(mut out: &mut Terminal, file: &File) {
    let (width, height) = out.get_size();
    let file_size = (width - LINENO_CHARS - 1, height - 3);
    let x = 1;
    let y = height;
    out.goto((x, y));
    write!(out, "{}", file.debug(file_size)).unwrap();
}

fn main() {
    let keys = KeybindTable::default();
    let mut term = Terminal::default();
    term.clear();
    term.flush().unwrap();
    render_footer(&mut term, &keys);
    let args: Vec<String> = env::args().collect();
    let filename = args.get(1).cloned().unwrap_or(String::from("README.md"));
    let mut file = File::open(&filename);
    render_file(&mut term, &file);
    render_status(&mut term, &file);
    let (screen_w, screen_h) = term.get_size();
    let file_size = (screen_w - LINENO_CHARS - 1, screen_h - 3);
    term.goto((file.cursor(file_size).x + LINENO_CHARS + 1, file.cursor(file_size).y));
    term.flush().unwrap();
    let mut file_lines = file.lines.len();
    let mut file_wrapped_lines = file.wrapped_lines(file_size).len();
    let mut file_dirty = false;
    let mut screen_dirty = false;
    for evt in term.keys() {
        let (screen_w, screen_h) = term.get_size();
        let file_size = (screen_w - LINENO_CHARS - 1, screen_h - 3);
        match evt {
            Event::Mouse(_) => (),
            Event::Unsupported(_) => (),
            Event::Key(Key::Null) => (),
            Event::Key(Key::Insert) => (),
            Event::Key(Key::F(_)) => (),
            Event::Key(Key::Esc) => (),
            Event::Key(k @ Key::Ctrl(_)) | Event::Key(k @ Key::Alt(_)) => {
                match keys.lookup(k) {
                    Some(Command::Quit) => break,
                    Some(Command::Refresh) => {
                        file.refresh(file_size);
                        screen_dirty = true;
                    },
                    None => (),
                    _ => ()
                }
            },
            Event::Key(Key::Left) => {
                screen_dirty = file.move_cursor_left(file_size);
            },
            Event::Key(Key::Right) => {
                screen_dirty = file.move_cursor_right(file_size);
            },
            Event::Key(Key::Up) => {
                screen_dirty = file.move_cursor_up(file_size);
            },
            Event::Key(Key::Down) => {
                screen_dirty = file.move_cursor_down(file_size);
            },
            Event::Key(Key::Home) => file.move_cursor_home(file_size),
            Event::Key(Key::End) => file.move_cursor_end(file_size),
            Event::Key(Key::PageUp) => {
                file.page_up(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::PageDown) => {
                file.page_down(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Char('\t')) => {
                for _ in 0..4 {
                    file.insert(file_size, ' ');
                }
                file_dirty = true;
            },
            Event::Key(Key::Char('\n')) => {
                file.insert_newline(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Delete) => {
                file.delete(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Backspace) => {
                file.backspace(file_size);
                screen_dirty = true;
            },
            Event::Key(Key::Char(c)) => {
                file.insert(file_size, c);
                file_dirty = true;
            },
        }
        if screen_dirty {
            file_dirty = true;
        }
        if file_dirty {
            let new_file_lines = file.lines.len();
            let new_file_wrapped_lines = file.wrapped_lines(file_size).len();
            if new_file_lines != file_lines {
                screen_dirty = true;
                file_lines = new_file_lines;
            }
            if new_file_wrapped_lines != file_wrapped_lines {
                screen_dirty = true;
                file_wrapped_lines = new_file_wrapped_lines;
            }
        }
        if screen_dirty {
            term.clear();
            render_footer(&mut term, &keys);
            screen_dirty = false;
        }
        if file_dirty {
            render_file(&mut term, &file);
            file_dirty = false;
        }
        render_status(&mut term, &file);
        term.goto((file.cursor(file_size).x + LINENO_CHARS + 1, file.cursor(file_size).y));
        term.flush().unwrap();
    }
}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::spawn(move || -> Result<()> {
            // Actual data transfer happens here
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        });

        // Actual data transfer happens here
        ::std::io::copy(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn start_serving(&mut self, spec: &str, once: bool);
    fn first_connected_client(&mut self) -> Result<IPeer>;
    fn stop(&mut self);

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

struct TcpServer {

}
impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        unimplemented!()
    }
}

impl Server for TcpServer {
    fn start_serving(&mut self, spec: &str, once: bool) { unimplemented!() }
    fn first_connected_client(&mut self) -> Result<IPeer> { unimplemented!() }
    fn stop(&mut self) { unimplemented!() }
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>    Connect to websocket
  tcp:<socket address>               Connect to TCP
  l-tcp:<socket address>             Listen TCP connections
  -                                  stdin/stdout
  (more to be implemented)

Examples:
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to secure web socket once.
    Like netcat, but for websocket.
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.
  websocat l-tcp::::9559 ws://echo.websocket.org/
    Listen TCPv6 port 9959 on address :: and forward
    all connections to a public loopback websocket

Specify listening part first, unless you want websocat to serve once.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, false);
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.first_connected_client()?;
                    x.stop();
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()?;
        }
    }

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<commit_msg>TCP server<commit_after>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::spawn(move || -> Result<()> {
            // Actual data transfer happens here
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        });

        // Actual data transfer happens here
        ::std::io::copy(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let peer = self.accept_client()?;
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                thread::spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                });
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>    Connect to websocket
  tcp:<socket address>               Connect to TCP
  l-tcp:<socket address>             Listen TCP connections
  -                                  stdin/stdout
  (more to be implemented)

Examples:
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to secure web socket once.
    Like netcat, but for websocket.
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.
  websocat l-tcp::::9559 ws://echo.websocket.org/
    Listen TCPv6 port 9959 on address :: and forward
    all connections to a public loopback websocket

Specify listening part first, unless you want websocat to serve once.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<|endoftext|>


<commit_before><commit_msg>correct some typos in help text<commit_after><|endoftext|>


<commit_before><commit_msg>feat: use job with src/dest<commit_after><|endoftext|>


<commit_before>#![feature(concat_idents,
           custom_attribute,
           inclusive_range_syntax,
           rustc_macro)]
#![recursion_limit = "1024"]
extern crate ansi_term;
extern crate bit_vec;
extern crate crypto;
extern crate difference;
extern crate env_logger;
#[macro_use]
extern crate error_chain;
extern crate getopts;
extern crate libc;
#[macro_use]
extern crate log;
extern crate hyper;
#[macro_use]
extern crate lazy_static;
extern crate nix;
extern crate notify;
#[macro_use]
extern crate quick_error;
extern crate rustc_serialize;
extern crate ssh2;
#[macro_use]
extern crate serde_derive;
extern crate serde;
extern crate serde_json;
extern crate term;
extern crate time;
extern crate toml;
extern crate unqlite;
extern crate url;
extern crate uuid;
extern crate xml;

use getopts::Options;
use std::env;
use std::fs;
use std::os::unix::io::AsRawFd;
use std::sync::Arc;
use nix::sys::signal;
use nix::unistd::{close, dup2, fork, ForkResult, getppid, sleep};
use std::path::Path;
use std::process;
use std::process::Command;

pub mod consts;
use consts::*;

pub mod exec;

pub mod libvirt;

#[macro_use]
pub mod util;
use util::errors::*;
use util::notify::config_hup;

pub mod flota;
use flota::cluster::Cluster;
use flota::config::*;
use flota::config::store::*;
use flota::config::template::Ingredient;
use flota::template::Template;

#[macro_use]
pub mod virt;
use virt::conn::Conn;
use virt::ResourceBlend;
use virt::storage::pool::StoragePool;

pub mod distro;
use distro::Distros;

fn print_usage(opts: Options) {
    let brief = format!("Usage: {}", *PROGNAME);
    print!("{}", opts.usage(&brief));
}

static mut CONFIG_RELOAD: bool = false;

#[allow(unused_variables)]
fn daemonize() -> Result<()> {
    if getppid() == 1 {
        return Err("already daemonized".into());
    }
    match fork().expect("fork failed") {
        ForkResult::Parent { child } => {
            process::exit(0);
        }
        ForkResult::Child => {}
    }
    if unsafe { libc::setsid() } < 0 {
        error!("setsid failed.");
        process::exit(1);
    }
    match (fs::OpenOptions::new().read(true).write(true).open("/dev/null"),
           fs::OpenOptions::new().read(true).write(true).create(true).open(LOGFILE.as_os_str()),
           fs::OpenOptions::new().read(true).write(true).create(true).open(LOGERROR.as_os_str())) {
        (Ok(n), Ok(f), Ok(e)) => {
            let n_raw = n.as_raw_fd();
            let f_raw = f.as_raw_fd();
            let e_raw = e.as_raw_fd();
            dup2(n_raw, libc::STDIN_FILENO).expect("dup2 failed");
            dup2(f_raw, libc::STDOUT_FILENO).expect("dup2 failed");
            dup2(e_raw, libc::STDERR_FILENO).expect("dup2 failed");
            if n_raw > 2 {
                close(n_raw).expect("cannot close fd opened for /dev/null");
            }
        },
        _ => panic!("failed to daemonize. dup2 failed.")
    }
    Ok(())
}

extern "C" fn config_reload(_: i32) {
    unsafe { CONFIG_RELOAD = true };
}

fn verify_env() -> Result<()> {
    // selinux disabled?
    match Command::new("getenforce")
        .output()
        .expect("failed to execute getenforce")
        .stdout {
            ref s if String::from_utf8(s.clone()).unwrap() == "Disabled\n" => {}
            _ => {
                panic!("selinux must be disabled.");
        }
    }

    // data dir exists?
    match DATA_DIR.metadata() {
        Ok(attr) => {
            if ! attr.is_dir() {
                panic!("data dir ({}) does not exists.", DATA_DIR.to_str().unwrap());
            }
        },
        Err(e) => {
            panic!("{}", e);
        }
    }

    // XXX: if dnsmasq running on hosts, check its config and make sure "bind-interfaces" uncommented.
    // XXX: cli utils availability check
    Ok(())
}

fn main() {
    env_logger::init().unwrap();

    // parse options
    let args: Vec<String> = env::args().collect();
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optopt("c", "config", format!(
            "config toml file (DEFAULT: /etc/flota.toml)").as_str(), "FILE");
    opts.optflag("", "clean", "remove all templates and clusters/hosts.");
    opts.optflag("", "reset", "reset all.");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(e) => {
            println!("{}", e.to_string());
            print_usage(opts);
            return;
        }
    };
    if matches.opt_present("h") {
        print_usage(opts);
        return;
    }
    let config_path = match matches.opt_str("c") {
        Some(c) => { c },
        None => { "/etc/flota.toml".to_string() }
    };
    if matches.opt_present("reset") {
        println!("would reset");
        return;
    }
    if matches.opt_present("clean") {
        println!("would clean");
        return;
    }

    // verify environment
    verify_env().unwrap();

    let config_store = unqlite_backed::ConfigStore::new(
        ::consts::CONFIG_HISTORY_DIR.join("hoge").as_path()
    );

    // outermost loop
    'init: loop {
        // read toml
        match Config::from_toml_file(Path::new(&config_path)) {
            Ok(config) => {
                debug!("{:#?}", config);
                config_store.update(&config).unwrap();

                // set up main connection
                let conn = Conn::new(&config.setting.hypervisor);

                // set up default storage pool + default network
                let default_storage_pool = StoragePool::ensure(&conn,
                                                               &config.setting.default_storage_pool_name,
                                                               &config.setting.pool_root)
                    .expect("cannot make sure default storage exists and is active");
                let default_nw_br_ip = &config.setting.default_network.nth_sibling(1);
                let default_network =
                    virt::network::Network::ensure_default(&conn, &default_nw_br_ip, true);
                let mut default_resources = ResourceBlend::new(&conn);
                default_resources.put_network(&default_network);
                default_resources.put_pool(&default_storage_pool);

                if config.setting.daemonized {
                    // if it's already daemonized, returns Err.
                    // in other words, changing the "daemonized" config value
                    // to true is one-way trip.
                    match daemonize() {
                        Ok(_) => {
                            let hup_action = signal::SigAction::new(
                                signal::SigHandler::Handler(config_reload),
                                signal::SaFlags::empty(),
                                signal::SigSet::empty());
                            unsafe {
                                signal::sigaction(signal::SIGHUP, &hup_action)
                                    .expect("sigaction for SIGHUP failed");
                            }
                            let _child = config_hup(Path::new(&config_path))
                                .expect("failed to setup config_hup");
                        },
                        Err(_) => {}
                    }
                }

                // staying in this inner loop
                'cycle: loop {
                    // construct templates.
                    let mut templates = Vec::new();

                    for ref template in &config.templates {
                        let distro = match &template.ingredient {
                            &Ingredient::OffTheShelf {
                                ref distro
                            } => {
                                Distros::search(&distro, &template.arch)
                            },
                            // XXX: linux is not the only choice
                            &Ingredient::Custom {
                                ref iso,
                                ref iso_md5sum,
                                ref vmlinuz,
                                ref initrd,
                            } => {
                                Distros::custom(iso, iso_md5sum, vmlinuz, initrd)
                            }
                        };
                        match Template::new(&default_resources, template, distro) {
                            Ok(t) => {
                                templates.push(Arc::new(t));
                            }
                            Err(e) => {
                                warn!("{}", e);
                                continue;
                            }
                        };
                    }

                    // construct (+ run tests on) clusters.
                    // TODO: safely parallelize
                    for ref cluster in &config.clusters {
                        let _c = match Cluster::new(cluster, &templates) {
                            Ok(c) => c,
                            Err(e) => {
                                error!("{}", e);
                                continue;
                            }
                        };
                    }

                    // if in daemon-mode, sleep five seconds and loop to next.
                    if config.setting.daemonized {
                        sleep(5);
                        if unsafe { CONFIG_RELOAD } {
                            unsafe { CONFIG_RELOAD = false };
                            match Config::from_toml_file(Path::new(&config_path)) {
                                Ok(ref new_config) => {
                                    if let Ok(true) = config_store.update(&new_config) {
                                        break 'cycle;
                                    }
                                },
                                Err(e) => {
                                    error!("{}", e);
                                }
                            }
                        }
                    } else {
                        break 'init;
                    }
                }
            },
            Err(e) => {
                error!("{}", e);
                sleep(5);
            }
        }
    }
}
<commit_msg>add sigterm handling on main loop<commit_after>#![feature(concat_idents,
           custom_attribute,
           inclusive_range_syntax,
           rustc_macro)]
#![recursion_limit = "1024"]
extern crate ansi_term;
extern crate bit_vec;
extern crate crypto;
extern crate difference;
extern crate env_logger;
#[macro_use]
extern crate error_chain;
extern crate getopts;
extern crate libc;
#[macro_use]
extern crate log;
extern crate hyper;
#[macro_use]
extern crate lazy_static;
extern crate nix;
extern crate notify;
#[macro_use]
extern crate quick_error;
extern crate rustc_serialize;
extern crate ssh2;
#[macro_use]
extern crate serde_derive;
extern crate serde;
extern crate serde_json;
extern crate term;
extern crate time;
extern crate toml;
extern crate unqlite;
extern crate url;
extern crate uuid;
extern crate xml;

use getopts::Options;
use std::env;
use std::fs;
use std::os::unix::io::AsRawFd;
use std::sync::{Arc, Mutex};
use nix::sys::signal;
use nix::unistd::{close, dup2, fork, ForkResult, getppid, sleep};
use std::path::Path;
use std::process;
use std::process::Command;

pub mod consts;
use consts::*;

pub mod exec;

pub mod libvirt;

#[macro_use]
pub mod util;
use util::errors::*;
use util::notify::config_hup;

pub mod flota;
use flota::cluster::Cluster;
use flota::config::*;
use flota::config::store::*;
use flota::config::template::Ingredient;
use flota::template::Template;

#[macro_use]
pub mod virt;
use virt::conn::Conn;
use virt::ResourceBlend;
use virt::storage::pool::StoragePool;

pub mod distro;
use distro::Distros;

fn print_usage(opts: Options) {
    let brief = format!("Usage: {}", *PROGNAME);
    print!("{}", opts.usage(&brief));
}

static mut CONFIG_RELOAD: bool = false;
static mut SIGTERM_RECVED: bool = false;
lazy_static! {
    static ref WAIT_FOR: Mutex<Vec<i32>> = Mutex::new(vec![]);
}

#[allow(unused_variables)]
fn daemonize() -> Result<()> {
    if getppid() == 1 {
        return Err("already daemonized".into());
    }
    match fork().expect("fork failed") {
        ForkResult::Parent { child } => {
            process::exit(0);
        }
        ForkResult::Child => {}
    }
    if unsafe { libc::setsid() } < 0 {
        error!("setsid failed.");
        process::exit(1);
    }
    match (fs::OpenOptions::new().read(true).write(true).open("/dev/null"),
           fs::OpenOptions::new().read(true).write(true).create(true).open(LOGFILE.as_os_str()),
           fs::OpenOptions::new().read(true).write(true).create(true).open(LOGERROR.as_os_str())) {
        (Ok(n), Ok(f), Ok(e)) => {
            let n_raw = n.as_raw_fd();
            let f_raw = f.as_raw_fd();
            let e_raw = e.as_raw_fd();
            dup2(n_raw, libc::STDIN_FILENO).expect("dup2 failed");
            dup2(f_raw, libc::STDOUT_FILENO).expect("dup2 failed");
            dup2(e_raw, libc::STDERR_FILENO).expect("dup2 failed");
            if n_raw > 2 {
                close(n_raw).expect("cannot close fd opened for /dev/null");
            }
        },
        _ => panic!("failed to daemonize. dup2 failed.")
    }
    Ok(())
}

extern "C" fn config_reload(_: i32) {
    unsafe { CONFIG_RELOAD = true };
}

extern "C" fn sigterm_received(_: i32) {
    unsafe { SIGTERM_RECVED = true };
}

fn verify_env() -> Result<()> {
    // selinux disabled?
    match Command::new("getenforce")
        .output()
        .expect("failed to execute getenforce")
        .stdout {
            ref s if String::from_utf8(s.clone()).unwrap() == "Disabled\n" => {}
            _ => {
                panic!("selinux must be disabled.");
        }
    }

    // data dir exists?
    match DATA_DIR.metadata() {
        Ok(attr) => {
            if ! attr.is_dir() {
                panic!("data dir ({}) does not exists.", DATA_DIR.to_str().unwrap());
            }
        },
        Err(e) => {
            panic!("{}", e);
        }
    }

    // XXX: if dnsmasq running on hosts, check its config and make sure "bind-interfaces" uncommented.
    // XXX: cli utils availability check
    Ok(())
}

fn main() {
    env_logger::init().unwrap();

    // parse options
    let args: Vec<String> = env::args().collect();
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.optopt("c", "config", format!(
            "config toml file (DEFAULT: /etc/flota.toml)").as_str(), "FILE");
    opts.optflag("", "clean", "remove all templates and clusters/hosts.");
    opts.optflag("", "reset", "reset all.");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(e) => {
            println!("{}", e.to_string());
            print_usage(opts);
            return;
        }
    };
    if matches.opt_present("h") {
        print_usage(opts);
        return;
    }
    let config_path = match matches.opt_str("c") {
        Some(c) => { c },
        None => { "/etc/flota.toml".to_string() }
    };
    if matches.opt_present("reset") {
        println!("would reset");
        return;
    }
    if matches.opt_present("clean") {
        println!("would clean");
        return;
    }

    // verify environment
    verify_env().unwrap();

    let config_store = unqlite_backed::ConfigStore::new(
        ::consts::CONFIG_HISTORY_DIR.join("hoge").as_path()
    );

    // outermost loop
    'init: loop {
        if unsafe { SIGTERM_RECVED } { break }
        // read toml
        match Config::from_toml_file(Path::new(&config_path)) {
            Ok(config) => {
                debug!("{:#?}", config);
                config_store.update(&config).unwrap();

                // set up main connection
                let conn = Conn::new(&config.setting.hypervisor);

                // set up default storage pool + default network
                let default_storage_pool = StoragePool::ensure(&conn,
                                                               &config.setting.default_storage_pool_name,
                                                               &config.setting.pool_root)
                    .expect("cannot make sure default storage exists and is active");
                let default_nw_br_ip = &config.setting.default_network.nth_sibling(1);
                let default_network =
                    virt::network::Network::ensure_default(&conn, &default_nw_br_ip, true);
                let mut default_resources = ResourceBlend::new(&conn);
                default_resources.put_network(&default_network);
                default_resources.put_pool(&default_storage_pool);

                if config.setting.daemonized {
                    // if it's already daemonized, returns Err.
                    // in other words, changing the "daemonized" config value
                    // to true is one-way trip.
                    match daemonize() {
                        Ok(_) => {
                            let hup_action = signal::SigAction::new(
                                signal::SigHandler::Handler(config_reload),
                                signal::SaFlags::empty(),
                                signal::SigSet::empty());
                            let term_action = signal::SigAction::new(
                                signal::SigHandler::Handler(sigterm_received),
                                signal::SaFlags::empty(),
                                signal::SigSet::empty());
                            unsafe {
                                signal::sigaction(signal::SIGHUP, &hup_action)
                                    .expect("sigaction for SIGHUP failed");
                                signal::sigaction(signal::SIGTERM, &term_action)
                                    .expect("sigaction for SIGTERM failed");
                                WAIT_FOR.lock().unwrap().push(config_hup(Path::new(&config_path))
                                    .expect("failed to setup config_hup"));
                            }
                        },
                        Err(_) => {}
                    }
                }

                // staying in this inner loop
                'cycle: loop {
                    // construct templates.
                    let mut templates = Vec::new();

                    for ref template in &config.templates {
                        let distro = match &template.ingredient {
                            &Ingredient::OffTheShelf {
                                ref distro
                            } => {
                                Distros::search(&distro, &template.arch)
                            },
                            // XXX: linux is not the only choice
                            &Ingredient::Custom {
                                ref iso,
                                ref iso_md5sum,
                                ref vmlinuz,
                                ref initrd,
                            } => {
                                Distros::custom(iso, iso_md5sum, vmlinuz, initrd)
                            }
                        };
                        match Template::new(&default_resources, template, distro) {
                            Ok(t) => {
                                templates.push(Arc::new(t));
                            }
                            Err(e) => {
                                warn!("{}", e);
                                continue;
                            }
                        };
                    }

                    // construct (+ run tests on) clusters.
                    // TODO: safely parallelize
                    for ref cluster in &config.clusters {
                        let _c = match Cluster::new(cluster, &templates) {
                            Ok(c) => c,
                            Err(e) => {
                                error!("{}", e);
                                continue;
                            }
                        };
                    }

                    if ! config.setting.daemonized ||
                       unsafe { SIGTERM_RECVED } { break 'init }

                    sleep(5);
                    if unsafe { CONFIG_RELOAD } {
                        unsafe { CONFIG_RELOAD = false };
                        match Config::from_toml_file(Path::new(&config_path)) {
                            Ok(ref new_config) => {
                                if let Ok(true) = config_store.update(&new_config) {
                                    break 'cycle;
                                }
                            },
                            Err(e) => {
                                error!("{}", e);
                            }
                        }
                    }
                }
            },
            Err(e) => {
                error!("{}", e);
                sleep(5);
            }
        }
    }
    for child in WAIT_FOR.lock().unwrap().iter() {
        signal::kill(*child, libc::SIGKILL)
            .expect(format!("failed to send SIGKILL to {}", child).as_str());
    }
}
<|endoftext|>


<commit_before><commit_msg>Add comments<commit_after><|endoftext|>


<commit_before><commit_msg>Parser work<commit_after><|endoftext|>


<commit_before><commit_msg>Fix soundness issue with badly-behaved comparators.<commit_after><|endoftext|>


<commit_before><commit_msg>stop using the super::* glob import<commit_after><|endoftext|>


<commit_before><commit_msg>Use `NonZeroU32::new_unchecked` to convert wasi error (#233)<commit_after><|endoftext|>


<commit_before>//! Implements a command for sending events to Sentry.
use std::borrow::Cow;
use std::env;

use clap::{App, Arg, ArgMatches};
use failure::{err_msg, Error};
use itertools::Itertools;
use sentry::protocol::{Event, Level, LogEntry, User};
use serde_json::Value;
use username::get_user_name;

use crate::config::Config;
use crate::utils::event::{attach_logfile, get_sdk_info, with_sentry_client};
use crate::utils::releases::detect_release_name;

pub fn make_app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
    app.about("Send a manual event to Sentry.")
        .long_about(
            "Send a manual event to Sentry.{n}{n}\
             NOTE: This command will validate input parameters and attempt to send an event to \
             Sentry. Due to network errors, rate limits or sampling the event is not guaranteed to \
             actually arrive. Check debug output for transmission errors by passing --log-level=\
             debug or setting `SENTRY_LOG_LEVEL=debug`.",
        )
        .arg(
            Arg::with_name("level")
                .value_name("LEVEL")
                .long("level")
                .short("l")
                .help("Optional event severity/log level. [defaults to 'error']"),
        )
        .arg(
            Arg::with_name("release")
                .value_name("RELEASE")
                .long("release")
                .short("r")
                .help("Optional identifier of the release."),
        )
        .arg(
            Arg::with_name("dist")
                .value_name("DISTRIBUTION")
                .long("dist")
                .short("d")
                .help("Set the distribution."),
        )
        .arg(
            Arg::with_name("environment")
                .value_name("ENVIRONMENT")
                .long("env")
                .short("E")
                .help("Send with a specific environment."),
        )
        .arg(
            Arg::with_name("no_environ")
                .long("no-environ")
                .help("Do not send environment variables along"),
        )
        .arg(
            Arg::with_name("message")
                .value_name("MESSAGE")
                .long("message")
                .short("m")
                .multiple(true)
                .number_of_values(1)
                .help("The event message."),
        )
        .arg(
            Arg::with_name("message_args")
                .value_name("MESSAGE_ARG")
                .long("message-arg")
                .short("a")
                .multiple(true)
                .number_of_values(1)
                .help("Arguments for the event message."),
        )
        .arg(
            Arg::with_name("platform")
                .value_name("PLATFORM")
                .long("platform")
                .short("p")
                .help("Override the default 'other' platform specifier."),
        )
        .arg(
            Arg::with_name("tags")
                .value_name("KEY:VALUE")
                .long("tag")
                .short("t")
                .multiple(true)
                .number_of_values(1)
                .help("Add a tag (key:value) to the event."),
        )
        .arg(
            Arg::with_name("extra")
                .value_name("KEY:VALUE")
                .long("extra")
                .short("e")
                .multiple(true)
                .number_of_values(1)
                .help("Add extra information (key:value) to the event."),
        )
        .arg(
            Arg::with_name("user_data")
                .value_name("KEY:VALUE")
                .long("user")
                .short("u")
                .multiple(true)
                .number_of_values(1)
                .help(
                    "Add user information (key:value) to the event. \
                     [eg: id:42, username:foo]",
                ),
        )
        .arg(
            Arg::with_name("fingerprint")
                .value_name("FINGERPRINT")
                .long("fingerprint")
                .short("f")
                .multiple(true)
                .number_of_values(1)
                .help("Change the fingerprint of the event."),
        )
        .arg(
            Arg::with_name("logfile")
                .value_name("PATH")
                .long("logfile")
                .help("Send a logfile as breadcrumbs with the event (last 100 records)"),
        )
        .arg(
            Arg::with_name("with_categories")
                .long("with-categories")
                .help("Parses off a leading category for breadcrumbs from the logfile")
                .long_help(
                    "When logfile is provided, this flag will try to assign correct level \
                    to extracted log breadcrumbs. It uses standard log format of \"category: message\". \
                    eg. \"INFO: Something broke\" will be parsed as a breadcrumb \
                    \"{\"level\": \"info\", \"message\": \"Something broke\"}\"")
        )
}

pub fn execute<'a>(matches: &ArgMatches<'a>) -> Result<(), Error> {
    let config = Config::current();
    let mut event = Event::default();

    event.sdk = Some(get_sdk_info());

    event.level = matches
        .value_of("level")
        .and_then(|l| l.parse().ok())
        .unwrap_or(Level::Error);

    if let Some(release) = matches.value_of("release") {
        event.release = Some(release.to_string().into());
    } else {
        event.release = detect_release_name().ok().map(Cow::from);
    }

    event.dist = matches.value_of("dist").map(|x| x.to_string().into());
    event.platform = matches
        .value_of("platform")
        .unwrap_or("other")
        .to_string()
        .into();
    event.environment = matches
        .value_of("environment")
        .map(|x| x.to_string().into());

    if let Some(mut lines) = matches.values_of("message") {
        event.logentry = Some(LogEntry {
            message: lines.join("\n"),
            params: matches
                .values_of("message_args")
                .map(|args| args.map(|x| x.into()).collect())
                .unwrap_or_default(),
        });
    }

    for tag in matches.values_of("tags").unwrap_or_default() {
        let mut split = tag.splitn(2, ':');
        let key = split.next().ok_or_else(|| err_msg("missing tag key"))?;
        let value = split.next().ok_or_else(|| err_msg("missing tag value"))?;
        event.tags.insert(key.into(), value.into());
    }

    if !matches.is_present("no_environ") {
        event.extra.insert(
            "environ".into(),
            Value::Object(env::vars().map(|(k, v)| (k, Value::String(v))).collect()),
        );
    }

    for pair in matches.values_of("extra").unwrap_or_default() {
        let mut split = pair.splitn(2, ':');
        let key = split.next().ok_or_else(|| err_msg("missing extra key"))?;
        let value = split.next().ok_or_else(|| err_msg("missing extra value"))?;
        event.extra.insert(key.into(), Value::String(value.into()));
    }

    if let Some(user_data) = matches.values_of("user_data") {
        let mut user = User::default();
        for pair in user_data {
            let mut split = pair.splitn(2, ':');
            let key = split.next().ok_or_else(|| err_msg("missing user key"))?;
            let value = split.next().ok_or_else(|| err_msg("missing user value"))?;

            match key {
                "id" => user.id = Some(value.into()),
                "email" => user.email = Some(value.into()),
                "ip_address" => user.ip_address = Some(value.parse()?),
                "username" => user.username = Some(value.into()),
                _ => {
                    user.other.insert(key.into(), value.into());
                }
            };
        }

        user.ip_address.get_or_insert(Default::default());
        event.user = Some(user);
    } else {
        event.user = get_user_name().ok().map(|n| User {
            username: Some(n),
            ip_address: Some(Default::default()),
            ..Default::default()
        });
    }

    if let Some(fingerprint) = matches.values_of("fingerprint") {
        event.fingerprint = fingerprint
            .map(|x| x.to_string().into())
            .collect::<Vec<_>>()
            .into();
    }

    if let Some(logfile) = matches.value_of("logfile") {
        attach_logfile(&mut event, logfile, matches.is_present("with_categories"))?;
    }

    let id = with_sentry_client(config.get_dsn()?, |c| c.capture_event(event, None));
    println!("{}", id);

    Ok(())
}
<commit_msg>feat: Allow for passing custom timestamp with send-event (#826)<commit_after>//! Implements a command for sending events to Sentry.
use std::borrow::Cow;
use std::env;

use clap::{App, Arg, ArgMatches};
use failure::{err_msg, Error};
use itertools::Itertools;
use sentry::protocol::{Event, Level, LogEntry, User};
use serde_json::Value;
use username::get_user_name;

use crate::config::Config;
use crate::utils::args::{get_timestamp, validate_timestamp};
use crate::utils::event::{attach_logfile, get_sdk_info, with_sentry_client};
use crate::utils::releases::detect_release_name;

pub fn make_app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
    app.about("Send a manual event to Sentry.")
        .long_about(
            "Send a manual event to Sentry.{n}{n}\
             NOTE: This command will validate input parameters and attempt to send an event to \
             Sentry. Due to network errors, rate limits or sampling the event is not guaranteed to \
             actually arrive. Check debug output for transmission errors by passing --log-level=\
             debug or setting `SENTRY_LOG_LEVEL=debug`.",
        )
        .arg(
            Arg::with_name("level")
                .value_name("LEVEL")
                .long("level")
                .short("l")
                .help("Optional event severity/log level. [defaults to 'error']"),
        )
        .arg(Arg::with_name("timestamp")
                 .long("timestamp")
                 .validator(validate_timestamp)
                 .value_name("TIMESTAMP")
                 .help("Optional event timestamp in one of supported formats: unix timestamp, RFC2822 or RFC3339."))
        .arg(
            Arg::with_name("release")
                .value_name("RELEASE")
                .long("release")
                .short("r")
                .help("Optional identifier of the release."),
        )
        .arg(
            Arg::with_name("dist")
                .value_name("DISTRIBUTION")
                .long("dist")
                .short("d")
                .help("Set the distribution."),
        )
        .arg(
            Arg::with_name("environment")
                .value_name("ENVIRONMENT")
                .long("env")
                .short("E")
                .help("Send with a specific environment."),
        )
        .arg(
            Arg::with_name("no_environ")
                .long("no-environ")
                .help("Do not send environment variables along"),
        )
        .arg(
            Arg::with_name("message")
                .value_name("MESSAGE")
                .long("message")
                .short("m")
                .multiple(true)
                .number_of_values(1)
                .help("The event message."),
        )
        .arg(
            Arg::with_name("message_args")
                .value_name("MESSAGE_ARG")
                .long("message-arg")
                .short("a")
                .multiple(true)
                .number_of_values(1)
                .help("Arguments for the event message."),
        )
        .arg(
            Arg::with_name("platform")
                .value_name("PLATFORM")
                .long("platform")
                .short("p")
                .help("Override the default 'other' platform specifier."),
        )
        .arg(
            Arg::with_name("tags")
                .value_name("KEY:VALUE")
                .long("tag")
                .short("t")
                .multiple(true)
                .number_of_values(1)
                .help("Add a tag (key:value) to the event."),
        )
        .arg(
            Arg::with_name("extra")
                .value_name("KEY:VALUE")
                .long("extra")
                .short("e")
                .multiple(true)
                .number_of_values(1)
                .help("Add extra information (key:value) to the event."),
        )
        .arg(
            Arg::with_name("user_data")
                .value_name("KEY:VALUE")
                .long("user")
                .short("u")
                .multiple(true)
                .number_of_values(1)
                .help(
                    "Add user information (key:value) to the event. \
                     [eg: id:42, username:foo]",
                ),
        )
        .arg(
            Arg::with_name("fingerprint")
                .value_name("FINGERPRINT")
                .long("fingerprint")
                .short("f")
                .multiple(true)
                .number_of_values(1)
                .help("Change the fingerprint of the event."),
        )
        .arg(
            Arg::with_name("logfile")
                .value_name("PATH")
                .long("logfile")
                .help("Send a logfile as breadcrumbs with the event (last 100 records)"),
        )
        .arg(
            Arg::with_name("with_categories")
                .long("with-categories")
                .help("Parses off a leading category for breadcrumbs from the logfile")
                .long_help(
                    "When logfile is provided, this flag will try to assign correct level \
                    to extracted log breadcrumbs. It uses standard log format of \"category: message\". \
                    eg. \"INFO: Something broke\" will be parsed as a breadcrumb \
                    \"{\"level\": \"info\", \"message\": \"Something broke\"}\"")
        )
}

pub fn execute<'a>(matches: &ArgMatches<'a>) -> Result<(), Error> {
    let config = Config::current();
    let mut event = Event::default();

    event.sdk = Some(get_sdk_info());

    event.level = matches
        .value_of("level")
        .and_then(|l| l.parse().ok())
        .unwrap_or(Level::Error);

    if let Some(timestamp) = matches.value_of("timestamp") {
        event.timestamp = get_timestamp(timestamp)?;
    }

    if let Some(release) = matches.value_of("release") {
        event.release = Some(release.to_string().into());
    } else {
        event.release = detect_release_name().ok().map(Cow::from);
    }

    event.dist = matches.value_of("dist").map(|x| x.to_string().into());
    event.platform = matches
        .value_of("platform")
        .unwrap_or("other")
        .to_string()
        .into();
    event.environment = matches
        .value_of("environment")
        .map(|x| x.to_string().into());

    if let Some(mut lines) = matches.values_of("message") {
        event.logentry = Some(LogEntry {
            message: lines.join("\n"),
            params: matches
                .values_of("message_args")
                .map(|args| args.map(|x| x.into()).collect())
                .unwrap_or_default(),
        });
    }

    for tag in matches.values_of("tags").unwrap_or_default() {
        let mut split = tag.splitn(2, ':');
        let key = split.next().ok_or_else(|| err_msg("missing tag key"))?;
        let value = split.next().ok_or_else(|| err_msg("missing tag value"))?;
        event.tags.insert(key.into(), value.into());
    }

    if !matches.is_present("no_environ") {
        event.extra.insert(
            "environ".into(),
            Value::Object(env::vars().map(|(k, v)| (k, Value::String(v))).collect()),
        );
    }

    for pair in matches.values_of("extra").unwrap_or_default() {
        let mut split = pair.splitn(2, ':');
        let key = split.next().ok_or_else(|| err_msg("missing extra key"))?;
        let value = split.next().ok_or_else(|| err_msg("missing extra value"))?;
        event.extra.insert(key.into(), Value::String(value.into()));
    }

    if let Some(user_data) = matches.values_of("user_data") {
        let mut user = User::default();
        for pair in user_data {
            let mut split = pair.splitn(2, ':');
            let key = split.next().ok_or_else(|| err_msg("missing user key"))?;
            let value = split.next().ok_or_else(|| err_msg("missing user value"))?;

            match key {
                "id" => user.id = Some(value.into()),
                "email" => user.email = Some(value.into()),
                "ip_address" => user.ip_address = Some(value.parse()?),
                "username" => user.username = Some(value.into()),
                _ => {
                    user.other.insert(key.into(), value.into());
                }
            };
        }

        user.ip_address.get_or_insert(Default::default());
        event.user = Some(user);
    } else {
        event.user = get_user_name().ok().map(|n| User {
            username: Some(n),
            ip_address: Some(Default::default()),
            ..Default::default()
        });
    }

    if let Some(fingerprint) = matches.values_of("fingerprint") {
        event.fingerprint = fingerprint
            .map(|x| x.to_string().into())
            .collect::<Vec<_>>()
            .into();
    }

    if let Some(logfile) = matches.value_of("logfile") {
        attach_logfile(&mut event, logfile, matches.is_present("with_categories"))?;
    }

    let id = with_sentry_client(config.get_dsn()?, |c| c.capture_event(event, None));
    println!("{}", id);

    Ok(())
}
<|endoftext|>


<commit_before>use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex};
use std::usize;

use sleep::Sleep;

/// We define various kinds of latches, which are all a primitive signaling
/// mechanism. A latch starts as false. Eventually someone calls `set()` and
/// it becomes true. You can test if it has been set by calling `probe()`.
///
/// Some kinds of latches, but not all, support a `wait()` operation
/// that will wait until the latch is set, blocking efficiently. That
/// is not part of the trait since it is not possibly to do with all
/// latches.
///
/// The intention is that `set()` is called once, but `probe()` may be
/// called any number of times. Once `probe()` returns true, the memory
/// effects that occurred before `set()` become visible.
///
/// It'd probably be better to refactor the API into two paired types,
/// but that's a bit of work, and this is not a public API.
///
/// ## Memory ordering
///
/// Latches need to guarantee two things:
///
/// - Once `probe()` returns true, all memory effects from the `set()`
///   are visible (in other words, the set should synchronize-with
///   the probe).
/// - Once `set()` occurs, the next `probe()` *will* observe it.  This
///   typically requires a seq-cst ordering. See [the "tickle-then-get-sleepy" scenario in the sleep
///   README](/src/sleep/README.md#tickle-then-get-sleepy) for details.
pub(super) trait Latch: LatchProbe {
    /// Set the latch, signalling others.
    fn set(&self);
}

pub(super) trait LatchProbe {
    /// Test if the latch is set.
    fn probe(&self) -> bool;
}

/// Spin latches are the simplest, most efficient kind, but they do
/// not support a `wait()` operation. They just have a boolean flag
/// that becomes true when `set()` is called.
pub(super) struct SpinLatch {
    b: AtomicBool,
}

impl SpinLatch {
    #[inline]
    pub(super) fn new() -> SpinLatch {
        SpinLatch {
            b: AtomicBool::new(false),
        }
    }
}

impl LatchProbe for SpinLatch {
    #[inline]
    fn probe(&self) -> bool {
        self.b.load(Ordering::SeqCst)
    }
}

impl Latch for SpinLatch {
    #[inline]
    fn set(&self) {
        self.b.store(true, Ordering::SeqCst);
    }
}

/// A Latch starts as false and eventually becomes true. You can block
/// until it becomes true.
pub(super) struct LockLatch {
    m: Mutex<bool>,
    v: Condvar,
}

impl LockLatch {
    #[inline]
    pub(super) fn new() -> LockLatch {
        LockLatch {
            m: Mutex::new(false),
            v: Condvar::new(),
        }
    }

    /// Resets this lock latch so it can be reused again.
    pub(super) fn reset(&self) {
        *self.m.lock().unwrap() = false;
    }

    /// Block until latch is set.
    pub(super) fn wait(&self) {
        let mut guard = self.m.lock().unwrap();
        while !*guard {
            guard = self.v.wait(guard).unwrap();
        }
    }
}

impl LatchProbe for LockLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Not particularly efficient, but we don't really use this operation
        let guard = self.m.lock().unwrap();
        *guard
    }
}

impl Latch for LockLatch {
    #[inline]
    fn set(&self) {
        let mut guard = self.m.lock().unwrap();
        *guard = true;
        self.v.notify_all();
    }
}

/// Counting latches are used to implement scopes. They track a
/// counter. Unlike other latches, calling `set()` does not
/// necessarily make the latch be considered `set()`; instead, it just
/// decrements the counter. The latch is only "set" (in the sense that
/// `probe()` returns true) once the counter reaches zero.
#[derive(Debug)]
pub(super) struct CountLatch {
    counter: AtomicUsize,
}

impl CountLatch {
    #[inline]
    pub(super) fn new() -> CountLatch {
        CountLatch {
            counter: AtomicUsize::new(1),
        }
    }

    #[inline]
    pub(super) fn increment(&self) {
        debug_assert!(!self.probe());
        self.counter.fetch_add(1, Ordering::Relaxed);
    }
}

impl LatchProbe for CountLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Need to acquire any memory reads before latch was set:
        self.counter.load(Ordering::SeqCst) == 0
    }
}

impl Latch for CountLatch {
    /// Set the latch to true, releasing all threads who are waiting.
    #[inline]
    fn set(&self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}

/// A tickling latch wraps another latch type, and will also awaken a thread
/// pool when it is set.  This is useful for jobs injected between thread pools,
/// so the source pool can continue processing its own work while waiting.
pub(super) struct TickleLatch<'a, L: Latch> {
    inner: L,
    sleep: &'a Sleep,
}

impl<'a, L: Latch> TickleLatch<'a, L> {
    #[inline]
    pub(super) fn new(latch: L, sleep: &'a Sleep) -> Self {
        TickleLatch {
            inner: latch,
            sleep,
        }
    }
}

impl<'a, L: Latch> LatchProbe for TickleLatch<'a, L> {
    #[inline]
    fn probe(&self) -> bool {
        self.inner.probe()
    }
}

impl<'a, L: Latch> Latch for TickleLatch<'a, L> {
    #[inline]
    fn set(&self) {
        self.inner.set();
        self.sleep.tickle(usize::MAX);
    }
}

impl<'a, L> LatchProbe for &'a L where L: LatchProbe {
    fn probe(&self) -> bool {
        L::probe(&self)
    }
}

impl<'a, L> Latch for &'a L where L: Latch {
    fn set(&self) {
        L::set(&self);
    }
}
<commit_msg>We don't need these.<commit_after>use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex};
use std::usize;

use sleep::Sleep;

/// We define various kinds of latches, which are all a primitive signaling
/// mechanism. A latch starts as false. Eventually someone calls `set()` and
/// it becomes true. You can test if it has been set by calling `probe()`.
///
/// Some kinds of latches, but not all, support a `wait()` operation
/// that will wait until the latch is set, blocking efficiently. That
/// is not part of the trait since it is not possibly to do with all
/// latches.
///
/// The intention is that `set()` is called once, but `probe()` may be
/// called any number of times. Once `probe()` returns true, the memory
/// effects that occurred before `set()` become visible.
///
/// It'd probably be better to refactor the API into two paired types,
/// but that's a bit of work, and this is not a public API.
///
/// ## Memory ordering
///
/// Latches need to guarantee two things:
///
/// - Once `probe()` returns true, all memory effects from the `set()`
///   are visible (in other words, the set should synchronize-with
///   the probe).
/// - Once `set()` occurs, the next `probe()` *will* observe it.  This
///   typically requires a seq-cst ordering. See [the "tickle-then-get-sleepy" scenario in the sleep
///   README](/src/sleep/README.md#tickle-then-get-sleepy) for details.
pub(super) trait Latch: LatchProbe {
    /// Set the latch, signalling others.
    fn set(&self);
}

pub(super) trait LatchProbe {
    /// Test if the latch is set.
    fn probe(&self) -> bool;
}

/// Spin latches are the simplest, most efficient kind, but they do
/// not support a `wait()` operation. They just have a boolean flag
/// that becomes true when `set()` is called.
pub(super) struct SpinLatch {
    b: AtomicBool,
}

impl SpinLatch {
    #[inline]
    pub(super) fn new() -> SpinLatch {
        SpinLatch {
            b: AtomicBool::new(false),
        }
    }
}

impl LatchProbe for SpinLatch {
    #[inline]
    fn probe(&self) -> bool {
        self.b.load(Ordering::SeqCst)
    }
}

impl Latch for SpinLatch {
    #[inline]
    fn set(&self) {
        self.b.store(true, Ordering::SeqCst);
    }
}

/// A Latch starts as false and eventually becomes true. You can block
/// until it becomes true.
pub(super) struct LockLatch {
    m: Mutex<bool>,
    v: Condvar,
}

impl LockLatch {
    #[inline]
    pub(super) fn new() -> LockLatch {
        LockLatch {
            m: Mutex::new(false),
            v: Condvar::new(),
        }
    }

    /// Resets this lock latch so it can be reused again.
    pub(super) fn reset(&self) {
        *self.m.lock().unwrap() = false;
    }

    /// Block until latch is set.
    pub(super) fn wait(&self) {
        let mut guard = self.m.lock().unwrap();
        while !*guard {
            guard = self.v.wait(guard).unwrap();
        }
    }
}

impl LatchProbe for LockLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Not particularly efficient, but we don't really use this operation
        let guard = self.m.lock().unwrap();
        *guard
    }
}

impl Latch for LockLatch {
    #[inline]
    fn set(&self) {
        let mut guard = self.m.lock().unwrap();
        *guard = true;
        self.v.notify_all();
    }
}

/// Counting latches are used to implement scopes. They track a
/// counter. Unlike other latches, calling `set()` does not
/// necessarily make the latch be considered `set()`; instead, it just
/// decrements the counter. The latch is only "set" (in the sense that
/// `probe()` returns true) once the counter reaches zero.
#[derive(Debug)]
pub(super) struct CountLatch {
    counter: AtomicUsize,
}

impl CountLatch {
    #[inline]
    pub(super) fn new() -> CountLatch {
        CountLatch {
            counter: AtomicUsize::new(1),
        }
    }

    #[inline]
    pub(super) fn increment(&self) {
        debug_assert!(!self.probe());
        self.counter.fetch_add(1, Ordering::Relaxed);
    }
}

impl LatchProbe for CountLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Need to acquire any memory reads before latch was set:
        self.counter.load(Ordering::SeqCst) == 0
    }
}

impl Latch for CountLatch {
    /// Set the latch to true, releasing all threads who are waiting.
    #[inline]
    fn set(&self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}

/// A tickling latch wraps another latch type, and will also awaken a thread
/// pool when it is set.  This is useful for jobs injected between thread pools,
/// so the source pool can continue processing its own work while waiting.
pub(super) struct TickleLatch<'a, L: Latch> {
    inner: L,
    sleep: &'a Sleep,
}

impl<'a, L: Latch> TickleLatch<'a, L> {
    #[inline]
    pub(super) fn new(latch: L, sleep: &'a Sleep) -> Self {
        TickleLatch {
            inner: latch,
            sleep,
        }
    }
}

impl<'a, L: Latch> LatchProbe for TickleLatch<'a, L> {
    #[inline]
    fn probe(&self) -> bool {
        self.inner.probe()
    }
}

impl<'a, L: Latch> Latch for TickleLatch<'a, L> {
    #[inline]
    fn set(&self) {
        self.inner.set();
        self.sleep.tickle(usize::MAX);
    }
}

impl<'a, L> LatchProbe for &'a L where L: LatchProbe {
    fn probe(&self) -> bool {
        L::probe(self)
    }
}

impl<'a, L> Latch for &'a L where L: Latch {
    fn set(&self) {
        L::set(self);
    }
}
<|endoftext|>