<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86_64` architecture-specific implementation.
// pub mod cpu;
pub mod drivers;
pub mod interrupts;
// pub mod memory;

#[path = "../x86_all/bda.rs"] pub mod bda;
#[path = "../x86_all/multiboot2.rs"] pub mod multiboot2;

pub const ARCH_BITS: u8 = 64;

use memory::PAddr;

/// Entry point for architecture-specific kernel init
///
/// This expects to be passed the address of a valid
/// Multiboot 2 info struct. It's the bootloader's responsibility to ensure
/// that this is passed in the correct register as expected by the calling
/// convention (`edi` on x86). If this isn't there, you can expect to have a
/// bad problem and not go to space today.
#[no_mangle]
pub extern "C" fn arch_init(multiboot_addr: PAddr) {
    use cpu::{control_regs, msr};
    use elf;
    use memory::{PAddr, Page, PhysicalPage};
    use params::InitParams;
    use ::kernel_init;

    kinfoln!(dots: " . ", "Beginning `arch_init()` for x86_64");

    ::io::term::CONSOLE.lock().clear();
    ::logger::initialize()
        .expect("Could not initialize logger!");

    // -- Unpack multiboot tag -----------------------------------------------
    // try to interpret the structure at the multiboot address as a multiboot
    // info struct. if it's invalid, fail.
    let boot_info
        = unsafe { multiboot2::Info::from(multiboot_addr)
                    .expect("Could not unpack multiboot2 information!") };

    // Extract the memory map tag from the multiboot info
    let mmap_tag
        = boot_info.mem_map()
                   .expect("Memory map tag required!");

    kinfoln!(dots: " . ", "Detected memory areas:");
    for a in mmap_tag.areas() {
        kinfoln!( dots: " . . ", "start: {:#08x}, end: {:#08x}"
                , a.base, a.length );
        // TODO: add these to a list of memory areas?
        //       - eliza, 1/23/2017
    }

    // Extract ELF sections tag from the multiboot info
    let elf_sections_tag
        = boot_info.elf_sections()
                   .expect("ELF sections tag required!");

    kinfoln!(dots: " . ", "Detecting kernel ELF sections:");

    // Extract kernel ELF sections from  multiboot info
    let mut n_elf_sections = 0;
    let kernel_begin
        = elf_sections_tag.sections()
            .inspect(|s| {
                kinfoln!( dots: " . . "
                        , "address: {:#08x}, size: {:#08x}, flags: {:#08x}"
                        , s.addr()
                        , s.length()
                        , s.flags() );
                n_elf_sections += 1;
                })
            // the ELF section with the lowest address is the kernel start
            // section
            .min_by_key(elf::Section::addr)
            .expect("Could not find kernel start section!\
                    \nSomething is deeply wrong.");

    let kernel_end
        = elf_sections_tag.sections()
            .max_by_key(elf::Section::addr)
            .expect("Could not find kernel end section!\
                    \nSomething is deeply wrong.");

    kinfoln!( dots: " . ", "Detected {} kernel ELF sections.", n_elf_sections);
    kinfoln!( dots: " . . ", "Kernel begins at {:#p} and ends at {:#p}."
            , kernel_begin.addr(), kernel_end.addr() );

    let multiboot_end = multiboot_addr + boot_info.length as u64;

    kinfoln!( dots: " . . ", "Multiboot info begins at {:#x} and ends at {:#x}."
            , multiboot_addr, multiboot_end);

    let params = InitParams { multiboot_start: Some(multiboot_addr)
                            , multiboot_end: Some(multiboot_end)
                            , ..Default::default()
                            };

     // -- enable flags needed for paging ------------------------------------
     unsafe {
         control_regs::cr0::enable_write_protect(true);
         kinfoln!(dots: " . ", "Page write protect ENABED" );

         msr::enable_nxe();
         kinfoln!(dots: " . ", "Page no execute bit ENABLED");
     }

    kinfoln!(dots: " . ", "Transferring to `kernel_init()`.");
    kernel_init(params);
}
<commit_msg>[kernel] remove unused imports<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86_64` architecture-specific implementation.
// pub mod cpu;
pub mod drivers;
pub mod interrupts;
// pub mod memory;

#[path = "../x86_all/bda.rs"] pub mod bda;
#[path = "../x86_all/multiboot2.rs"] pub mod multiboot2;

pub const ARCH_BITS: u8 = 64;

use memory::PAddr;

/// Entry point for architecture-specific kernel init
///
/// This expects to be passed the address of a valid
/// Multiboot 2 info struct. It's the bootloader's responsibility to ensure
/// that this is passed in the correct register as expected by the calling
/// convention (`edi` on x86). If this isn't there, you can expect to have a
/// bad problem and not go to space today.
#[no_mangle]
pub extern "C" fn arch_init(multiboot_addr: PAddr) {
    use cpu::{control_regs, msr};
    use elf;
    use params::InitParams;

    kinfoln!(dots: " . ", "Beginning `arch_init()` for x86_64");

    ::io::term::CONSOLE.lock().clear();
    ::logger::initialize()
        .expect("Could not initialize logger!");

    // -- Unpack multiboot tag -----------------------------------------------
    // try to interpret the structure at the multiboot address as a multiboot
    // info struct. if it's invalid, fail.
    let boot_info
        = unsafe { multiboot2::Info::from(multiboot_addr)
                    .expect("Could not unpack multiboot2 information!") };

    // Extract the memory map tag from the multiboot info
    let mmap_tag
        = boot_info.mem_map()
                   .expect("Memory map tag required!");

    kinfoln!(dots: " . ", "Detected memory areas:");
    for a in mmap_tag.areas() {
        kinfoln!( dots: " . . ", "start: {:#08x}, end: {:#08x}"
                , a.base, a.length );
        // TODO: add these to a list of memory areas?
        //       - eliza, 1/23/2017
    }

    // Extract ELF sections tag from the multiboot info
    let elf_sections_tag
        = boot_info.elf_sections()
                   .expect("ELF sections tag required!");

    kinfoln!(dots: " . ", "Detecting kernel ELF sections:");

    // Extract kernel ELF sections from  multiboot info
    let mut n_elf_sections = 0;
    let kernel_begin
        = elf_sections_tag.sections()
            .inspect(|s| {
                kinfoln!( dots: " . . "
                        , "address: {:#08x}, size: {:#08x}, flags: {:#08x}"
                        , s.addr()
                        , s.length()
                        , s.flags() );
                n_elf_sections += 1;
                })
            // the ELF section with the lowest address is the kernel start
            // section
            .min_by_key(elf::Section::addr)
            .expect("Could not find kernel start section!\
                    \nSomething is deeply wrong.");

    let kernel_end
        = elf_sections_tag.sections()
            .max_by_key(elf::Section::addr)
            .expect("Could not find kernel end section!\
                    \nSomething is deeply wrong.");

    kinfoln!( dots: " . ", "Detected {} kernel ELF sections.", n_elf_sections);
    kinfoln!( dots: " . . ", "Kernel begins at {:#p} and ends at {:#p}."
            , kernel_begin.addr(), kernel_end.addr() );

    let multiboot_end = multiboot_addr + boot_info.length as u64;

    kinfoln!( dots: " . . ", "Multiboot info begins at {:#x} and ends at {:#x}."
            , multiboot_addr, multiboot_end);

    let params = InitParams { multiboot_start: Some(multiboot_addr)
                            , multiboot_end: Some(multiboot_end)
                            , ..Default::default()
                            };

     // -- enable flags needed for paging ------------------------------------
     unsafe {
         control_regs::cr0::enable_write_protect(true);
         kinfoln!(dots: " . ", "Page write protect ENABED" );

         msr::enable_nxe();
         kinfoln!(dots: " . ", "Page no execute bit ENABLED");
     }

    kinfoln!(dots: " . ", "Transferring to `kernel_init()`.");
    ::kernel_init(params);
}
<|endoftext|>


<commit_before>extern crate gio;
extern crate glib;
extern crate gtk;

use std::cell::RefCell;
use std::rc::Rc;

use gio::{ApplicationExt, ApplicationExtManual};
use gtk::{
    Application,
    ApplicationWindow,
    Button,
    ButtonExt,
    ContainerExt,
    GtkWindowExt,
    WidgetExt,
};
use glib::clone;

struct State {
    started: bool,
    count: i32,
}

impl State {
    fn new() -> Self {
        Self {
            started: false,
            count: 0,
        }
    }
}

fn main() {
    let application = Application::new(
        Some("com.github.gtk-rs.examples.basic"),
        Default::default(),
    ).expect("failed to initialize GTK application");

    let state = Rc::new(RefCell::new(State::new()));

    application.connect_activate(clone!(@weak state => move |app| {
        state.borrow_mut().started = true;

        let window = ApplicationWindow::new(app);
        window.set_title("First GTK+ Program");
        window.set_default_size(350, 70);

        let button = Button::new_with_label("Click me!");
        button.connect_clicked(clone!(state => move |_| {
            let mut state = state.borrow_mut();
            println!("Clicked (started: {}): {}!", state.started, state.count);
            state.count += 1;
        }));
        window.add(&button);

        window.show_all();
    }));

    application.run(&[]);
}
<commit_msg>Fix example<commit_after>extern crate gio;
extern crate glib;
extern crate gtk;

use std::cell::RefCell;
use std::rc::Rc;

use gio::prelude::*;
use gtk::{
    Application,
    ApplicationWindow,
    Button,
    prelude::*,
};
use glib::clone;

#[derive(Default)]
struct State {
    started: bool,
    count: i32,
}

impl State {
    fn new() -> Self {
        Self {
            started: false,
            count: 0,
        }
    }
}

fn main() {
    let application = Application::new(
        Some("com.github.gtk-rs.examples.basic"),
        Default::default(),
    ).expect("failed to initialize GTK application");

    let state = Rc::new(RefCell::new(State::new()));

    {
        let state2 = Rc::new(RefCell::new(State::new()));

        application.connect_activate(clone!(@weak state, state2 => move |app| {
            state.borrow_mut().started = true;

            let window = ApplicationWindow::new(app);
            window.set_title("First GTK+ Program");
            window.set_default_size(350, 70);

            let button = Button::new_with_label("Click me!");
            button.connect_clicked(clone!(@weak state, state2 => move |_| {
                let mut state = state.borrow_mut();
                let mut state2 = state2.borrow_mut();
                println!("Clicked (started: {}): {} - {}!", state.started, state.count, state2.count);
                state.count += 1;
                state2.count += 1;
            }));
            window.add(&button);

            window.show_all();
        }));
    }

    application.run(&[]);
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use futures::AsyncSink;
use futures::future::{FutureResult, ok, err};
use futures::sink::{Sink};
use minihttp::websocket;
use minihttp::websocket::{Error};
use minihttp::websocket::Frame::{self, Text, Binary, Ping, Pong};
use tokio_core::reactor::Handle;
use rustc_serialize::json::Json;

use runtime::Runtime;
use config::chat::Chat;
use config::SessionPool;
use chat::{Cid, ConnectionSender};
use chat::cid::serialize_cid;
use chat::message::{decode_message, get_active, Meta, Args, Kwargs};
use chat::processor::ProcessorPool;
use chat::backend::CallCodec;

pub struct Dispatcher {
    pub cid: Cid,
    pub auth: Arc<String>,
    pub runtime: Arc<Runtime>,
    pub settings: Arc<Chat>,
    pub pool_settings: Arc<SessionPool>,
    pub processor: ProcessorPool,
    pub handle: Handle, // Does it belong here?
    pub channel: ConnectionSender,
}

impl websocket::Dispatcher for Dispatcher {
    type Future = FutureResult<(), Error>;
    fn frame(&mut self, frame: &Frame) -> FutureResult<(), Error> {
        match *frame {
            Text(data) => match decode_message(data) {
                Ok((name, mut meta, args, kwargs)) => {
                    if let Some(duration) = get_active(&meta) {
                        // TODO(tailhook) update activity
                    }
                    meta.insert("connection_id".to_string(),
                        Json::String(serialize_cid(&self.cid)));
                    self.method_call(name, meta, args, kwargs);
                    ok(()) // no backpressure, yet
                }
                Err(e) => {
                    debug!("Message error: {}", e);
                    // TODO(tailhook) better error
                    err(Error::Closed)
                }
            },
            Binary(_) => {
                debug!("Binary messages are not supported yet");
                // TODO(tailhook) better error
                err(Error::Closed)
            }
            Ping(_)|Pong(_) => unreachable!(),
        }
    }
}

impl Dispatcher {
    fn method_call(&self, name: String, meta: Meta, args: Args, kw: Kwargs) {
        let dest = self.settings.message_handlers.resolve(&name);
        let mut path = name.replace("/", ".");
        if dest.path == "/" {
            path.insert(0, '/');
        } else {
            path = dest.path.clone() + "/" + &path;
        };
        let mut up = self.runtime.http_pools.upstream(&dest.upstream);
        let codec = Box::new(CallCodec::new(
            self.auth.clone(),
            path, meta, args, kw,
            self.channel.clone()));
        match up.get_mut().get_mut() {
            Some(pool) => {
                match pool.start_send(codec) {
                    Ok(AsyncSink::NotReady(codec)) => {
                        // codec.into_inner().send(Err(Status::ServiceUnavailable))
                        unimplemented!();
                    }
                    Ok(AsyncSink::Ready) => {
                        debug!("Sent /tangle/authorize_connection to proxy");
                    }
                    Err(e) => {
                        error!("Error sending to pool {:?}: {}", dest.upstream, e);
                        // TODO(tailhook) ensure that sender is closed
                    }
                }
            }
            None => {
                error!("No such destination {:?}", dest.upstream);
                // TODO(tailhook) return error to user
                // TODO(tailhook) deregister connection in pool
                // codec.into_inner().send(Err(Status::NotFound))
                //
                unimplemented!();
            }
        }
    }
}
<commit_msg>Fix typo in dispatcher. Now multi-user-chat finally works<commit_after>use std::sync::Arc;

use futures::AsyncSink;
use futures::future::{FutureResult, ok, err};
use futures::sink::{Sink};
use minihttp::websocket;
use minihttp::websocket::{Error};
use minihttp::websocket::Frame::{self, Text, Binary, Ping, Pong};
use tokio_core::reactor::Handle;
use rustc_serialize::json::Json;

use runtime::Runtime;
use config::chat::Chat;
use config::SessionPool;
use chat::{Cid, ConnectionSender};
use chat::cid::serialize_cid;
use chat::message::{decode_message, get_active, Meta, Args, Kwargs};
use chat::processor::ProcessorPool;
use chat::backend::CallCodec;

pub struct Dispatcher {
    pub cid: Cid,
    pub auth: Arc<String>,
    pub runtime: Arc<Runtime>,
    pub settings: Arc<Chat>,
    pub pool_settings: Arc<SessionPool>,
    pub processor: ProcessorPool,
    pub handle: Handle, // Does it belong here?
    pub channel: ConnectionSender,
}

impl websocket::Dispatcher for Dispatcher {
    type Future = FutureResult<(), Error>;
    fn frame(&mut self, frame: &Frame) -> FutureResult<(), Error> {
        match *frame {
            Text(data) => match decode_message(data) {
                Ok((name, mut meta, args, kwargs)) => {
                    if let Some(duration) = get_active(&meta) {
                        // TODO(tailhook) update activity
                    }
                    meta.insert("connection_id".to_string(),
                        Json::String(serialize_cid(&self.cid)));
                    self.method_call(name, meta, args, kwargs);
                    ok(()) // no backpressure, yet
                }
                Err(e) => {
                    debug!("Message error: {}", e);
                    // TODO(tailhook) better error
                    err(Error::Closed)
                }
            },
            Binary(_) => {
                debug!("Binary messages are not supported yet");
                // TODO(tailhook) better error
                err(Error::Closed)
            }
            Ping(_)|Pong(_) => unreachable!(),
        }
    }
}

impl Dispatcher {
    fn method_call(&self, name: String, meta: Meta, args: Args, kw: Kwargs) {
        let dest = self.settings.message_handlers.resolve(&name);
        let mut path = name.replace(".", "/");
        if dest.path == "/" {
            path.insert(0, '/');
        } else {
            path = dest.path.clone() + "/" + &path;
        };
        let mut up = self.runtime.http_pools.upstream(&dest.upstream);
        let codec = Box::new(CallCodec::new(
            self.auth.clone(),
            path, meta, args, kw,
            self.channel.clone()));
        match up.get_mut().get_mut() {
            Some(pool) => {
                match pool.start_send(codec) {
                    Ok(AsyncSink::NotReady(codec)) => {
                        // codec.into_inner().send(Err(Status::ServiceUnavailable))
                        unimplemented!();
                    }
                    Ok(AsyncSink::Ready) => {
                        debug!("Sent /tangle/authorize_connection to proxy");
                    }
                    Err(e) => {
                        error!("Error sending to pool {:?}: {}", dest.upstream, e);
                        // TODO(tailhook) ensure that sender is closed
                    }
                }
            }
            None => {
                error!("No such destination {:?}", dest.upstream);
                // TODO(tailhook) return error to user
                // TODO(tailhook) deregister connection in pool
                // codec.into_inner().send(Err(Status::NotFound))
                //
                unimplemented!();
            }
        }
    }
}
<|endoftext|>


<commit_before>//! On windows we cannot get read (or write) events for sockets in the
//! IOCP model; we only can run asynchronous reads!
//!
//! So we need to use select() to poll for read, and run it in a
//! separate thread.  To cancel that select() we'd need another socket
//! to select() for, and there is no socketpair() - we could use a
//! loopback TCP connection, but a firewall might block it.
//!
//! Instead we use a small (1 second) timeout for the select; it is only
//! used to terminate the thread anyway.
//!
//! This of course wastes one thread per fd we want to watch; a bigger
//! solution would reuse the same backend thread over and over, but then
//! we'd have to try the loopback TCP connection to wake it and fall
//! back to a smaller timeout.

use futures::{
	sink::Wait,
	sync::mpsc as futures_mpsc,
	Async,
	Poll,
	Sink,
	Stream,
};
use log::debug;
use std::{
	cell::UnsafeCell,
	io,
	os::raw::c_int,
	sync::mpsc as std_mpsc,
	thread,
	time::Duration,
};
use winapi::um::winsock2;

use self::fd_set::FdSet;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum PollRequest {
	Poll,
	Close,
}

struct SelectFdRead {
	fd: c_int,
	read_fds: FdSet,
}
impl SelectFdRead {
	pub fn new(fd: c_int) -> Self {
		SelectFdRead { fd, read_fds: FdSet::new() }
	}

	pub fn select(&mut self, timeout: Option<Duration>) -> bool {
		use std::ptr::null_mut;
		let mut timeout = timeout.map(|timeout| winsock2::timeval {
			tv_sec: timeout.as_secs() as libc::c_long,
			tv_usec: (timeout.subsec_nanos() / 1000) as libc::c_long,
		});
		self.read_fds.set(self.fd);
		unsafe {
			winsock2::select(
				self.fd + 1,
				self.read_fds.inner(),
				null_mut(),
				null_mut(),
				timeout.as_mut().map(|x| x as *mut _).unwrap_or(null_mut()),
			);
		}
		self.read_fds.is_set(self.fd)
	}
}

struct Inner {
	/// file descriptor to watch read events for
	fd: c_int,
	/// background select thread
	_thread: thread::JoinHandle<()>,
	/// either the select thread is running a Poll request or we manually
	/// sent a response through `send_response`
	pending_request: bool,
	/// send poll or close request to select thread
	send_request: std_mpsc::SyncSender<PollRequest>,
	/// when clear_read_ready() is called we use this to trigger a response if
	/// we already know the read event is pending
	send_response: Wait<futures_mpsc::Sender<()>>,
	/// a response means a read event is pending
	recv_response: futures_mpsc::Receiver<()>,
}

impl Inner {
	fn poll_read_ready(&mut self) -> Poll<(), io::Error> {
		debug!("poll read");
		if !self.pending_request {
			let mut read_fds = SelectFdRead::new(self.fd);
			if read_fds.select(Some(Duration::from_millis(0))) {
				debug!("poll read: local ready");
				return Ok(Async::Ready(()));
			} else {
				debug!("poll read: not ready, start thread");
				self.send_request
					.send(PollRequest::Poll)
					.expect("select thread terminated");
				self.pending_request = true;
			}
		}

		match self.recv_response.poll().unwrap() {
			Async::Ready(None) => unreachable!(),
			Async::Ready(Some(())) => {
				debug!("poll read: thread ready");
				self.pending_request = false;
				Ok(Async::Ready(()))
			},
			Async::NotReady => {
				debug!("poll read: thread not ready");
				Ok(Async::NotReady)
			},
		}
	}

	fn clear_read_ready(&mut self) -> io::Result<()> {
		// we need to get Async::NotReady from recv_response.poll
		match self.recv_response.poll().unwrap() {
			Async::Ready(None) => unreachable!(),
			Async::Ready(Some(())) => {
				// was ready. damn...
				assert!(self.pending_request);
				// try again - can't be ready again
				match self.recv_response.poll().unwrap() {
					Async::Ready(None) => unreachable!(),
					Async::Ready(Some(())) => unreachable!(),
					Async::NotReady => (),
				}
				// now send a response - it was ready after all
				self.send_response.send(()).unwrap();
			},
			Async::NotReady => {
				// yay!
				//
				// now we need something to trigger a response
				let mut read_fds = SelectFdRead::new(self.fd);
				self.pending_request = true;
				if read_fds.select(Some(Duration::from_millis(0))) {
					// ready, send a response
					self.send_response.send(()).unwrap();
				} else {
					debug!("poll need read: not ready, start thread");
					self.send_request
						.send(PollRequest::Poll)
						.expect("select thread terminated");
				}
			},
		}
		Ok(())
	}
}

pub fn is_readable(fd: c_int) -> io::Result<bool> {
	let mut read_fds = SelectFdRead::new(fd);
	Ok(read_fds.select(Some(Duration::from_millis(0))))
}

pub struct PollReadFd(UnsafeCell<Inner>);
impl PollReadFd {
	/// does not take overship of fd
	pub fn new(fd: c_int) -> io::Result<Self> {
		// buffer one request for "Close"
		let (send_request, recv_request) = std_mpsc::sync_channel(1);
		// buffer one notification
		let (send_response, recv_response) = futures_mpsc::channel(1);
		let outer_send_response = send_response.clone().wait();

		let thread = thread::spawn(move || {
			let mut read_fds = SelectFdRead::new(fd);
			let mut send_response = send_response.wait();
			loop {
				debug!("[select thread] waiting for request");
				match recv_request.recv() {
					Ok(PollRequest::Poll) => (),
					Ok(PollRequest::Close) => return,
					Err(_) => return,
				}
				debug!("[select thread] start polling");

				while !read_fds.select(Some(Duration::from_millis(1000))) {
					match recv_request.try_recv() {
						Ok(PollRequest::Poll) => unreachable!(),
						Ok(PollRequest::Close) => return,
						Err(_) => (), // back to select()
					}
				}

				debug!("[select thread] read event");

				if send_response.send(()).is_err() {
					return;
				}
			}
		});

		Ok(PollReadFd(UnsafeCell::new(Inner {
			fd,
			_thread: thread,
			pending_request: false,
			send_request,
			send_response: outer_send_response,
			recv_response,
		})))
	}

	fn inner(&self) -> &mut Inner {
		unsafe { &mut *self.0.get() }
	}

	pub fn poll_read_ready(&self) -> Poll<(), io::Error> {
		self.inner().poll_read_ready()
	}

	pub fn clear_read_ready(&self) -> io::Result<()> {
		self.inner().clear_read_ready()
	}
}

impl Drop for PollReadFd {
	fn drop(&mut self) {
		let _ = self.inner().send_request.send(PollRequest::Close);
	}
}

mod fd_set {
	use libc::{c_int, c_uint};
	use winapi::um::winsock2::{
		fd_set as RawFdSet,
		FD_SETSIZE,
		SOCKET,
	};

	pub(super) struct FdSet(RawFdSet);

	impl FdSet {
		pub fn new() -> Self {
			FdSet(RawFdSet {
				fd_count: 0,
				fd_array: [0; FD_SETSIZE],
			})
		}

		pub fn inner(&mut self) -> *mut RawFdSet {
			&mut self.0
		}

		pub fn set(&mut self, fd: c_int) {
			if self.is_set(fd) {
				return;
			}
			let count = self.0.fd_count as usize;
			if count < FD_SETSIZE {
				self.0.fd_array[count] = fd as c_uint as SOCKET;
				self.0.fd_count += 1;
			}
		}

		pub fn is_set(&self, fd: c_int) -> bool {
			let fd = fd as c_uint as SOCKET;
			let count = self.0.fd_count as usize;
			self.0.fd_array[..count].iter().any(|i| *i == fd)
		}
	}
}
<commit_msg>windows fd_set: use MaybeUninit array to avoid zeroing unused entries<commit_after>//! On windows we cannot get read (or write) events for sockets in the
//! IOCP model; we only can run asynchronous reads!
//!
//! So we need to use select() to poll for read, and run it in a
//! separate thread.  To cancel that select() we'd need another socket
//! to select() for, and there is no socketpair() - we could use a
//! loopback TCP connection, but a firewall might block it.
//!
//! Instead we use a small (1 second) timeout for the select; it is only
//! used to terminate the thread anyway.
//!
//! This of course wastes one thread per fd we want to watch; a bigger
//! solution would reuse the same backend thread over and over, but then
//! we'd have to try the loopback TCP connection to wake it and fall
//! back to a smaller timeout.

use futures::{
	sink::Wait,
	sync::mpsc as futures_mpsc,
	Async,
	Poll,
	Sink,
	Stream,
};
use log::debug;
use std::{
	cell::UnsafeCell,
	io,
	os::raw::c_int,
	sync::mpsc as std_mpsc,
	thread,
	time::Duration,
};
use winapi::um::winsock2;

use self::fd_set::FdSet;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum PollRequest {
	Poll,
	Close,
}

struct SelectFdRead {
	fd: c_int,
	read_fds: FdSet,
}
impl SelectFdRead {
	pub fn new(fd: c_int) -> Self {
		SelectFdRead { fd, read_fds: FdSet::new() }
	}

	pub fn select(&mut self, timeout: Option<Duration>) -> bool {
		use std::ptr::null_mut;
		let mut timeout = timeout.map(|timeout| winsock2::timeval {
			tv_sec: timeout.as_secs() as libc::c_long,
			tv_usec: (timeout.subsec_nanos() / 1000) as libc::c_long,
		});
		self.read_fds.set(self.fd);
		unsafe {
			winsock2::select(
				self.fd + 1,
				self.read_fds.inner(),
				null_mut(),
				null_mut(),
				timeout.as_mut().map(|x| x as *mut _).unwrap_or(null_mut()),
			);
		}
		self.read_fds.is_set(self.fd)
	}
}

struct Inner {
	/// file descriptor to watch read events for
	fd: c_int,
	/// background select thread
	_thread: thread::JoinHandle<()>,
	/// either the select thread is running a Poll request or we manually
	/// sent a response through `send_response`
	pending_request: bool,
	/// send poll or close request to select thread
	send_request: std_mpsc::SyncSender<PollRequest>,
	/// when clear_read_ready() is called we use this to trigger a response if
	/// we already know the read event is pending
	send_response: Wait<futures_mpsc::Sender<()>>,
	/// a response means a read event is pending
	recv_response: futures_mpsc::Receiver<()>,
}

impl Inner {
	fn poll_read_ready(&mut self) -> Poll<(), io::Error> {
		debug!("poll read");
		if !self.pending_request {
			let mut read_fds = SelectFdRead::new(self.fd);
			if read_fds.select(Some(Duration::from_millis(0))) {
				debug!("poll read: local ready");
				return Ok(Async::Ready(()));
			} else {
				debug!("poll read: not ready, start thread");
				self.send_request
					.send(PollRequest::Poll)
					.expect("select thread terminated");
				self.pending_request = true;
			}
		}

		match self.recv_response.poll().unwrap() {
			Async::Ready(None) => unreachable!(),
			Async::Ready(Some(())) => {
				debug!("poll read: thread ready");
				self.pending_request = false;
				Ok(Async::Ready(()))
			},
			Async::NotReady => {
				debug!("poll read: thread not ready");
				Ok(Async::NotReady)
			},
		}
	}

	fn clear_read_ready(&mut self) -> io::Result<()> {
		// we need to get Async::NotReady from recv_response.poll
		match self.recv_response.poll().unwrap() {
			Async::Ready(None) => unreachable!(),
			Async::Ready(Some(())) => {
				// was ready. damn...
				assert!(self.pending_request);
				// try again - can't be ready again
				match self.recv_response.poll().unwrap() {
					Async::Ready(None) => unreachable!(),
					Async::Ready(Some(())) => unreachable!(),
					Async::NotReady => (),
				}
				// now send a response - it was ready after all
				self.send_response.send(()).unwrap();
			},
			Async::NotReady => {
				// yay!
				//
				// now we need something to trigger a response
				let mut read_fds = SelectFdRead::new(self.fd);
				self.pending_request = true;
				if read_fds.select(Some(Duration::from_millis(0))) {
					// ready, send a response
					self.send_response.send(()).unwrap();
				} else {
					debug!("poll need read: not ready, start thread");
					self.send_request
						.send(PollRequest::Poll)
						.expect("select thread terminated");
				}
			},
		}
		Ok(())
	}
}

pub fn is_readable(fd: c_int) -> io::Result<bool> {
	let mut read_fds = SelectFdRead::new(fd);
	Ok(read_fds.select(Some(Duration::from_millis(0))))
}

pub struct PollReadFd(UnsafeCell<Inner>);
impl PollReadFd {
	/// does not take overship of fd
	pub fn new(fd: c_int) -> io::Result<Self> {
		// buffer one request for "Close"
		let (send_request, recv_request) = std_mpsc::sync_channel(1);
		// buffer one notification
		let (send_response, recv_response) = futures_mpsc::channel(1);
		let outer_send_response = send_response.clone().wait();

		let thread = thread::spawn(move || {
			let mut read_fds = SelectFdRead::new(fd);
			let mut send_response = send_response.wait();
			loop {
				debug!("[select thread] waiting for request");
				match recv_request.recv() {
					Ok(PollRequest::Poll) => (),
					Ok(PollRequest::Close) => return,
					Err(_) => return,
				}
				debug!("[select thread] start polling");

				while !read_fds.select(Some(Duration::from_millis(1000))) {
					match recv_request.try_recv() {
						Ok(PollRequest::Poll) => unreachable!(),
						Ok(PollRequest::Close) => return,
						Err(_) => (), // back to select()
					}
				}

				debug!("[select thread] read event");

				if send_response.send(()).is_err() {
					return;
				}
			}
		});

		Ok(PollReadFd(UnsafeCell::new(Inner {
			fd,
			_thread: thread,
			pending_request: false,
			send_request,
			send_response: outer_send_response,
			recv_response,
		})))
	}

	fn inner(&self) -> &mut Inner {
		unsafe { &mut *self.0.get() }
	}

	pub fn poll_read_ready(&self) -> Poll<(), io::Error> {
		self.inner().poll_read_ready()
	}

	pub fn clear_read_ready(&self) -> io::Result<()> {
		self.inner().clear_read_ready()
	}
}

impl Drop for PollReadFd {
	fn drop(&mut self) {
		let _ = self.inner().send_request.send(PollRequest::Close);
	}
}

mod fd_set {
	use libc::{c_int, c_uint};
	use winapi::um::winsock2::{
		u_int,
		fd_set,
		FD_SETSIZE,
		SOCKET,
	};
	use std::mem::MaybeUninit;
	use std::ptr;

	/// Layout compatible struct of `fd_set`, but it holds maybe uninitialized `fd_array`.
	///
	/// # Safety
	/// The first `fd_count` slots of `fd_array` must be initialized,
	/// and the rest may be uninitialized.
	#[repr(C)]
	pub(super) struct FdSet {
		fd_count: u_int,
		fd_array: [MaybeUninit<SOCKET>; FD_SETSIZE],
	}

	impl FdSet {
		pub fn new() -> Self {
			FdSet {
				fd_count: 0,
				// safe according to:
				// https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#initializing-an-array-element-by-element
				fd_array: unsafe { MaybeUninit::uninit().assume_init() },
			}
		}

		pub fn inner(&mut self) -> *mut fd_set {
			self as *mut FdSet as *mut _
		}

		pub fn set(&mut self, fd: c_int) {
			if self.is_set(fd) {
				return;
			}
			let count = self.fd_count as usize;
			if count < FD_SETSIZE {
				let fd = fd as c_uint as SOCKET;
				// This is safe because this slot is uninitialized.
				unsafe { ptr::write(self.fd_array[count].as_mut_ptr(), fd) };
				self.fd_count += 1;
			}
		}

		pub fn is_set(&self, fd: c_int) -> bool {
			let fd = fd as c_uint as SOCKET;
			let count = self.fd_count as usize;
			self.fd_array[..count].iter().any(|slot| {
				// This is safe because it's reading from first `fd_count` slots.
				fd == unsafe { ptr::read(slot.as_ptr()) }
			})
		}
	}

	#[cfg(test)]
	mod tests {
		use super::*;
		use std::mem::{needs_drop, transmute};

		// Check that `FdSet` is layout compatible with `fd_set`.
		#[test]
		fn fd_set_layout_compatible() {
			let mut fd_set = FdSet::new();
			(0..FD_SETSIZE as c_int).for_each(|i| fd_set.set(i));
			let fd_set: fd_set = unsafe { transmute(fd_set) };
			assert_eq!(fd_set.fd_count, FD_SETSIZE as u32);
			for i in 0..FD_SETSIZE as usize {
				assert_eq!(fd_set.fd_array[i], i as SOCKET);
			}
		}

		// Check that `SOCKET` doesn't need to be dropped, so that we don't need to
		// implement `Drop` for `FdSet`.
		#[test]
		fn socket_not_drop() {
			assert!(!needs_drop::<SOCKET>());
		}
	}
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate std;
extern crate sync;

use super::liblua;
use super::Lua;
use super::Pushable;

extern fn wrapper1(lua: *mut liblua::lua_State) -> libc::c_int {
    unsafe {
        let argumentsCount = liblua::lua_gettop(lua);

        let data = liblua::lua_touserdata(lua, liblua::lua_upvalueindex(1));
        let wrapper2ptr = liblua::lua_touserdata(lua, liblua::lua_upvalueindex(2));
        let wrapper2: fn(*mut liblua::lua_State, libc::c_int, *mut libc::c_void)->libc::c_int = std::mem::transmute(wrapper2ptr);

        wrapper2(lua, argumentsCount, data)
    }
}

/*fn wrapper2proc<Ret: Pushable>(lua: *mut liblua::lua_State, argumentsCount: libc::c_int, p: &proc()->Ret) -> libc::c_int {
    let ret = p();
    ret.push_to_lua(Lua{lua:lua});
    1
}*/

fn wrapper2fn<Ret: Pushable>(lua: *mut liblua::lua_State, argumentsCount: libc::c_int, p: fn()->Ret) -> libc::c_int {
    let ret = p();
    ret.push_to_lua(&mut Lua{lua:lua});
    1
}

/*impl<Ret: Pushable> Pushable for proc()->Ret {
    fn push_to_lua(&self, lua: &Lua) {
        unsafe {
            // pushing the std::raw::Procedure as a userdata
            let rawProc: &std::raw::Procedure = std::mem::transmute(self);
            let mut userDataPtr = liblua::lua_newuserdata(lua.lua, std::mem::size_of_val(rawProc) as libc::size_t);
            let userData: &mut std::raw::Procedure = std::mem::transmute(userDataPtr);
            userData.code = rawProc.code;
            userData.env = rawProc.env;

            // pushing wrapper2 as a lightuserdata
            let wrapper2: &fn(*mut liblua::lua_State, libc::c_int, &proc()->Ret)->libc::c_int = &wrapper2proc;
            liblua::lua_pushlightuserdata(lua.lua, std::mem::transmute(wrapper2));

            // pushing wrapper1 as a closure
            liblua::lua_pushcclosure(lua.lua, std::mem::transmute(wrapper1), 2);
        }
    }
}*/

impl<Ret: Pushable> Pushable for fn()->Ret {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe {
            // pushing the function pointer as a lightuserdata
            liblua::lua_pushlightuserdata(lua.lua, std::mem::transmute(*self));

            // pushing wrapper2 as a lightuserdata
            let wrapper2: fn(*mut liblua::lua_State, libc::c_int, fn()->Ret)->libc::c_int = wrapper2fn;
            liblua::lua_pushlightuserdata(lua.lua, std::mem::transmute(wrapper2));

            // pushing wrapper1 as a closure
            liblua::lua_pushcclosure(lua.lua, std::mem::transmute(wrapper1), 2);
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn simple_function() {
        let mut lua = super::super::Lua::new();

        fn ret5() -> int { 5 };
        lua.set("ret5", ret5).unwrap();

        let val: int = lua.execute("return ret5()").unwrap();
        assert_eq!(val, 5);
    }
}
<commit_msg>Partial implementation of pushable for static closures<commit_after>extern crate libc;
extern crate std;
extern crate sync;

use super::liblua;
use super::Lua;
use super::Pushable;

extern fn wrapper1(lua: *mut liblua::lua_State) -> libc::c_int {
    unsafe {
        let argumentsCount = liblua::lua_gettop(lua);

        let data = liblua::lua_touserdata(lua, liblua::lua_upvalueindex(1));
        let wrapper2ptr = liblua::lua_touserdata(lua, liblua::lua_upvalueindex(2));
        let wrapper2: fn(*mut liblua::lua_State, libc::c_int, *mut libc::c_void)->libc::c_int = std::mem::transmute(wrapper2ptr);

        wrapper2(lua, argumentsCount, data)
    }
}

fn wrapper2closure<Ret: Pushable>(lua: *mut liblua::lua_State, argumentsCount: libc::c_int, p: &||->Ret) -> libc::c_int {
    unimplemented!()
    /*let ret = p();
    ret.push_to_lua(&mut Lua{lua:lua});
    1*/
}

fn wrapper2fn<Ret: Pushable>(lua: *mut liblua::lua_State, argumentsCount: libc::c_int, p: fn()->Ret) -> libc::c_int {
    let ret = p();
    ret.push_to_lua(&mut Lua{lua:lua});
    1
}

impl<Ret: Pushable> Pushable for ||:'static->Ret {
    fn push_to_lua(&self, lua: &mut Lua) {
        // pushing the std::raw::Closure as a userdata
        let rawClosure: &std::raw::Closure = unsafe { std::mem::transmute(self) };
        let mut userDataPtr = unsafe { liblua::lua_newuserdata(lua.lua, std::mem::size_of_val(rawClosure) as libc::size_t) };
        let userData: &mut std::raw::Closure = unsafe { std::mem::transmute(userDataPtr) };
        userData.code = rawClosure.code;
        userData.env = rawClosure.env;

        // pushing wrapper2 as a lightuserdata
        let wrapper2: &fn(*mut liblua::lua_State, libc::c_int, &||->Ret)->libc::c_int = &wrapper2closure;
        unsafe { liblua::lua_pushlightuserdata(lua.lua, std::mem::transmute(wrapper2)) };

        // pushing wrapper1 as a closure
        unsafe { liblua::lua_pushcclosure(lua.lua, std::mem::transmute(wrapper1), 2) };
    }
}

impl<Ret: Pushable> Pushable for fn()->Ret {
    fn push_to_lua(&self, lua: &mut Lua) {
        // pushing the function pointer as a lightuserdata
        unsafe { liblua::lua_pushlightuserdata(lua.lua, std::mem::transmute(*self)) };

        // pushing wrapper2 as a lightuserdata
        let wrapper2: fn(*mut liblua::lua_State, libc::c_int, fn()->Ret)->libc::c_int = wrapper2fn;
        unsafe { liblua::lua_pushlightuserdata(lua.lua, std::mem::transmute(wrapper2)) };

        // pushing wrapper1 as a closure
        unsafe { liblua::lua_pushcclosure(lua.lua, std::mem::transmute(wrapper1), 2) };
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn simple_function() {
        let mut lua = super::super::Lua::new();

        fn ret5() -> int { 5 };
        lua.set("ret5", ret5).unwrap();

        let val: int = lua.execute("return ret5()").unwrap();
        assert_eq!(val, 5);
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::VecDeque,
    fmt,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc,
        Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        loadbalancing::server::LoadBalancer,
        socks5::Address,
        tcprelay::client::ServerClient as TcpServerClient,
        udprelay::client::ServerClient as UdpServerClient,
    },
};

use log::{debug, info};
use tokio::{
    self,
    io::{AsyncReadExt, AsyncWriteExt},
    sync::Barrier,
    time,
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

#[derive(Debug, Copy, Clone)]
enum Score {
    Latency(u64),
    Errored,
}

struct ServerLatencyInner {
    latency_queue: VecDeque<Score>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: Score) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            // Never checked, assume it is the worst of all
            return 2 * 1000;
        }

        // 1. Mid Latency
        // 2. Proportion of Errors
        let mut vec_lat = Vec::with_capacity(self.latency_queue.len());
        let mut acc_err = 0;
        for lat in &self.latency_queue {
            match lat {
                Score::Latency(l) => vec_lat.push(l),
                Score::Errored => acc_err += 1,
            }
        }

        let max_lat = DEFAULT_CHECK_TIMEOUT_SEC * 1000;

        // Find the mid of latencies
        let mid_lat = if vec_lat.is_empty() {
            // The whole array are errors
            max_lat
        } else {
            vec_lat.sort();
            let mid = vec_lat.len() / 2;
            if vec_lat.len() % 2 == 0 {
                (vec_lat[mid] + vec_lat[mid - 1]) / 2
            } else {
                *vec_lat[mid]
            }
        };

        // Score = norm_lat + prop_err
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        let norm_lat = mid_lat as f64 / max_lat as f64;
        let prop_err = acc_err as f64 / self.latency_queue.len() as f64;

        ((norm_lat + prop_err) * 1000.0) as u64
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: Score) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 2; // Latency shouldn't greater than 2 secs, that's too long

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    async fn new(context: SharedContext, server_type: ServerType) -> Inner {
        let config = context.config();
        assert!(!config.server.is_empty());

        // Load balancer is only required in multi-server configuration
        let check_required = config.server.len() > 1;

        // Wait for all ping tasks to be started
        let barrier = Arc::new(Barrier::new(config.server.len() + 1));

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();
            let barrier = barrier.clone();

            tokio::spawn(
                // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                async move {
                    let start_time = Instant::now();

                    debug!(
                        "{:?} server {} latency ping task initializing",
                        server_type,
                        sc.config.addr()
                    );

                    // Quickly collect some latency data
                    //
                    // Maximum wait duration: DEFAULT_CHECK_TIMEOUT_SEC
                    while Instant::now() - start_time < Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC) {
                        let mut succ = false;

                        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
                            Ok(d) => {
                                succ = true;
                                latency.push(Score::Latency(d))
                            }
                            Err(..) => {
                                latency.push(Score::Errored) // Penalty
                            }
                        };
                        debug!(
                            "updated remote {:?} server {} (score: {})",
                            server_type,
                            sc.config.addr(),
                            score
                        );
                        sc.score.store(score, Ordering::Release);

                        if succ {
                            break;
                        }
                    }

                    // Wait until all the other tasks are finished initializing
                    barrier.wait().await;
                    drop(barrier);

                    debug!(
                        "{:?} server {} latency ping task started",
                        server_type,
                        sc.config.addr()
                    );

                    while context.server_running() {
                        // First round may be failed, plugins are started asynchronously

                        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
                            Ok(d) => latency.push(Score::Latency(d)),
                            Err(..) => latency.push(Score::Errored), // Penalty
                        };
                        debug!(
                            "updated remote {:?} server {} (score: {})",
                            server_type,
                            sc.config.addr(),
                            score
                        );
                        sc.score.store(score, Ordering::Release);

                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }

                    debug!(
                        "{:?} server {} latency ping task stopped",
                        server_type,
                        sc.config.addr()
                    );
                },
            );
        }

        barrier.wait().await;

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    async fn check_request_tcp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static GET_BODY: &[u8] =
            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);

        let TcpServerClient { mut stream } = TcpServerClient::connect(context, &addr, sc).await?;
        stream.write_all(GET_BODY).await?;
        stream.flush().await?;
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await?;

        Ok(())
    }

    async fn check_request_udp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static DNS_QUERY: &[u8] = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01";

        let addr = Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53));

        let mut client = UdpServerClient::new(sc).await?;
        client.send_to(context, &addr, DNS_QUERY).await?;
        let _ = client.recv_from(context).await?;

        Ok(())
    }

    async fn check_request(sc: Arc<ServerConfig>, context: &Context, server_type: ServerType) -> io::Result<()> {
        match server_type {
            ServerType::Tcp => Inner::check_request_tcp(sc, context).await,
            ServerType::Udp => Inner::check_request_udp(sc, context).await,
        }
    }

    async fn check_delay(sc: &Server, context: &Context, server_type: ServerType) -> io::Result<u64> {
        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
        let res = time::timeout(timeout, Inner::check_request(sc.config.clone(), context, server_type)).await;

        let elapsed = Instant::now() - start;
        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
        match res {
            Ok(Ok(..)) => {
                // Got the result ... record its time
                debug!(
                    "checked remote {:?} server {} latency with {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );
                Ok(elapsed)
            }
            Ok(Err(err)) => {
                debug!(
                    "failed to check {:?} server {}, error: {}",
                    server_type,
                    sc.config.addr(),
                    err
                );

                // NOTE: connection / handshake error, server is down
                Err(err)
            }
            Err(..) => {
                // Timeout
                debug!(
                    "checked remote {:?} server {} latency timeout, elapsed {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );

                // NOTE: timeout is still available, but server is too slow
                Ok(elapsed)
            }
        }
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ServerType {
    Tcp,
    Udp,
}

/// Load balancer based on pinging latencies of all servers
#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    /// Create a PingBalancer
    pub async fn new(context: SharedContext, server_type: ServerType) -> PingBalancer {
        // Wait until all tasks are started
        let inner = Arc::new(Inner::new(context.clone(), server_type).await);

        if inner.servers.len() > 1 {
            let barrier = Arc::new(Barrier::new(2));

            let cloned_inner = inner.clone();
            let cloned_barrier = barrier.clone();
            tokio::spawn(async move {
                let inner = cloned_inner;
                assert!(!inner.servers.is_empty());

                let mut opt_barrier = Some(cloned_barrier);

                while context.server_running() {
                    // Choose the best one
                    let mut svr_idx = 0;

                    for (idx, svr) in inner.servers.iter().enumerate() {
                        let choosen_svr = &inner.servers[svr_idx];
                        if svr.score() < choosen_svr.score() {
                            svr_idx = idx;
                        }
                    }

                    let choosen_svr = &inner.servers[svr_idx];
                    if choosen_svr.score() != inner.servers[inner.best_idx()].score() {
                        if opt_barrier.is_none() {
                            info!(
                                "switched {:?} server from {} (score: {}) to {} (score: {})",
                                server_type,
                                inner.best_server().config.addr(),
                                inner.best_server().score(),
                                choosen_svr.config.addr(),
                                choosen_svr.score(),
                            );
                        }

                        inner.set_best_idx(svr_idx);
                    }

                    if let Some(barrier) = opt_barrier.take() {
                        barrier.wait().await;
                        debug!("ping {:?} server choosing task started", server_type);
                    }

                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                }
            });

            barrier.wait().await;
        }

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<commit_msg>Fixed bug, single server doesnt need to wait<commit_after>use std::{
    collections::VecDeque,
    fmt,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc,
        Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        loadbalancing::server::LoadBalancer,
        socks5::Address,
        tcprelay::client::ServerClient as TcpServerClient,
        udprelay::client::ServerClient as UdpServerClient,
    },
};

use log::{debug, info};
use tokio::{
    self,
    io::{AsyncReadExt, AsyncWriteExt},
    sync::Barrier,
    time,
};

struct Server {
    config: Arc<ServerConfig>,
    score: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            score: AtomicU64::new(0),
        }
    }

    fn score(&self) -> u64 {
        self.score.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 37;

#[derive(Debug, Copy, Clone)]
enum Score {
    Latency(u64),
    Errored,
}

struct ServerLatencyInner {
    latency_queue: VecDeque<Score>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: Score) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.score()
    }

    fn score(&self) -> u64 {
        if self.latency_queue.is_empty() {
            // Never checked, assume it is the worst of all
            return 2 * 1000;
        }

        // 1. Mid Latency
        // 2. Proportion of Errors
        let mut vec_lat = Vec::with_capacity(self.latency_queue.len());
        let mut acc_err = 0;
        for lat in &self.latency_queue {
            match lat {
                Score::Latency(l) => vec_lat.push(l),
                Score::Errored => acc_err += 1,
            }
        }

        let max_lat = DEFAULT_CHECK_TIMEOUT_SEC * 1000;

        // Find the mid of latencies
        let mid_lat = if vec_lat.is_empty() {
            // The whole array are errors
            max_lat
        } else {
            vec_lat.sort();
            let mid = vec_lat.len() / 2;
            if vec_lat.len() % 2 == 0 {
                (vec_lat[mid] + vec_lat[mid - 1]) / 2
            } else {
                *vec_lat[mid]
            }
        };

        // Score = norm_lat + prop_err
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        let norm_lat = mid_lat as f64 / max_lat as f64;
        let prop_err = acc_err as f64 / self.latency_queue.len() as f64;

        ((norm_lat + prop_err) * 1000.0) as u64
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: Score) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 2; // Latency shouldn't greater than 2 secs, that's too long

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    async fn new(context: SharedContext, server_type: ServerType) -> Inner {
        let config = context.config();
        assert!(!config.server.is_empty());

        // Load balancer is only required in multi-server configuration
        let check_required = config.server.len() > 1;

        // Wait for all ping tasks to be started
        let barrier = Arc::new(Barrier::new(config.server.len() + 1));

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();
            let barrier = barrier.clone();

            tokio::spawn(
                // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                async move {
                    let start_time = Instant::now();

                    debug!(
                        "{:?} server {} latency ping task initializing",
                        server_type,
                        sc.config.addr()
                    );

                    // Quickly collect some latency data
                    //
                    // Maximum wait duration: DEFAULT_CHECK_TIMEOUT_SEC
                    while Instant::now() - start_time < Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC) {
                        let mut succ = false;

                        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
                            Ok(d) => {
                                succ = true;
                                latency.push(Score::Latency(d))
                            }
                            Err(..) => {
                                latency.push(Score::Errored) // Penalty
                            }
                        };
                        debug!(
                            "updated remote {:?} server {} (score: {})",
                            server_type,
                            sc.config.addr(),
                            score
                        );
                        sc.score.store(score, Ordering::Release);

                        if succ {
                            break;
                        }
                    }

                    // Wait until all the other tasks are finished initializing
                    barrier.wait().await;
                    drop(barrier);

                    debug!(
                        "{:?} server {} latency ping task started",
                        server_type,
                        sc.config.addr()
                    );

                    while context.server_running() {
                        // First round may be failed, plugins are started asynchronously

                        let score = match Inner::check_delay(&*sc, &*context, server_type).await {
                            Ok(d) => latency.push(Score::Latency(d)),
                            Err(..) => latency.push(Score::Errored), // Penalty
                        };
                        debug!(
                            "updated remote {:?} server {} (score: {})",
                            server_type,
                            sc.config.addr(),
                            score
                        );
                        sc.score.store(score, Ordering::Release);

                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }

                    debug!(
                        "{:?} server {} latency ping task stopped",
                        server_type,
                        sc.config.addr()
                    );
                },
            );
        }

        if check_required {
            barrier.wait().await;
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    async fn check_request_tcp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static GET_BODY: &[u8] =
            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);

        let TcpServerClient { mut stream } = TcpServerClient::connect(context, &addr, sc).await?;
        stream.write_all(GET_BODY).await?;
        stream.flush().await?;
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await?;

        Ok(())
    }

    async fn check_request_udp(sc: Arc<ServerConfig>, context: &Context) -> io::Result<()> {
        static DNS_QUERY: &[u8] = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01";

        let addr = Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53));

        let mut client = UdpServerClient::new(sc).await?;
        client.send_to(context, &addr, DNS_QUERY).await?;
        let _ = client.recv_from(context).await?;

        Ok(())
    }

    async fn check_request(sc: Arc<ServerConfig>, context: &Context, server_type: ServerType) -> io::Result<()> {
        match server_type {
            ServerType::Tcp => Inner::check_request_tcp(sc, context).await,
            ServerType::Udp => Inner::check_request_udp(sc, context).await,
        }
    }

    async fn check_delay(sc: &Server, context: &Context, server_type: ServerType) -> io::Result<u64> {
        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
        let res = time::timeout(timeout, Inner::check_request(sc.config.clone(), context, server_type)).await;

        let elapsed = Instant::now() - start;
        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
        match res {
            Ok(Ok(..)) => {
                // Got the result ... record its time
                debug!(
                    "checked remote {:?} server {} latency with {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );
                Ok(elapsed)
            }
            Ok(Err(err)) => {
                debug!(
                    "failed to check {:?} server {}, error: {}",
                    server_type,
                    sc.config.addr(),
                    err
                );

                // NOTE: connection / handshake error, server is down
                Err(err)
            }
            Err(..) => {
                // Timeout
                debug!(
                    "checked remote {:?} server {} latency timeout, elapsed {} ms",
                    server_type,
                    sc.config.addr(),
                    elapsed
                );

                // NOTE: timeout is still available, but server is too slow
                Ok(elapsed)
            }
        }
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ServerType {
    Tcp,
    Udp,
}

/// Load balancer based on pinging latencies of all servers
#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    /// Create a PingBalancer
    pub async fn new(context: SharedContext, server_type: ServerType) -> PingBalancer {
        // Wait until all tasks are started
        let inner = Arc::new(Inner::new(context.clone(), server_type).await);

        if inner.servers.len() > 1 {
            let barrier = Arc::new(Barrier::new(2));

            let cloned_inner = inner.clone();
            let cloned_barrier = barrier.clone();
            tokio::spawn(async move {
                let inner = cloned_inner;
                assert!(!inner.servers.is_empty());

                let mut opt_barrier = Some(cloned_barrier);

                while context.server_running() {
                    // Choose the best one
                    let mut svr_idx = 0;

                    for (idx, svr) in inner.servers.iter().enumerate() {
                        let choosen_svr = &inner.servers[svr_idx];
                        if svr.score() < choosen_svr.score() {
                            svr_idx = idx;
                        }
                    }

                    let choosen_svr = &inner.servers[svr_idx];
                    if choosen_svr.score() != inner.servers[inner.best_idx()].score() {
                        if opt_barrier.is_none() {
                            info!(
                                "switched {:?} server from {} (score: {}) to {} (score: {})",
                                server_type,
                                inner.best_server().config.addr(),
                                inner.best_server().score(),
                                choosen_svr.config.addr(),
                                choosen_svr.score(),
                            );
                        }

                        inner.set_best_idx(svr_idx);
                    }

                    if let Some(barrier) = opt_barrier.take() {
                        barrier.wait().await;
                        debug!("ping {:?} server choosing task started", server_type);
                    }

                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                }
            });

            barrier.wait().await;
        }

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<|endoftext|>


<commit_before>//! Unix terminal formatter and drain for slog-rs
//!
//! ```
//! #[macro_use]
//! extern crate slog;
//! extern crate slog_term;
//!
//! use slog::*;
//!
//! fn main() {
//!     let root = Logger::root(slog_term::streamer().build().fuse(), o!("build-id" => "8dfljdf"));
//! }
//! ```
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate slog_stream;
extern crate isatty;
extern crate ansi_term;
extern crate chrono;

use std::{io, fmt, sync};

use ansi_term::Colour;
use isatty::{stderr_isatty, stdout_isatty};

use slog::Record;
use slog::ser;
use slog::{Level, OwnedKeyValueList};
use slog_stream::Format as StreamFormat;
use slog_stream::{Decorator, RecordDecorator, Streamer, AsyncStreamer};

/// Formatting mode
enum FormatMode {
    /// Compact logging format
    Compact,
    /// Full logging format
    Full,
}

/// Full formatting with optional color support
pub struct Format<D: Decorator> {
    mode: FormatMode,
    decorator: D,
    history: sync::Mutex<Vec<usize>>,
}

impl<D: Decorator> Format<D> {
    /// New Format format that prints using color
    fn new(mode: FormatMode, d: D) -> Self {
        Format {
            decorator: d,
            mode: mode,
            history: sync::Mutex::new(vec![]),
        }
    }

    fn print_msg_header(&self,
                        io: &mut io::Write,
                        rd: &D::RecordDecorator,
                        info: &Record)
                        -> io::Result<()> {
        let ts = chrono::Local::now();
        try!(rd.fmt_timestamp(io, format_args!("{} ", ts.format("%b %d %H:%M:%S%.3f"))));
        try!(rd.fmt_level(io, format_args!("{} ", info.level().as_short_str())));

        try!(rd.fmt_msg(io, format_args!("{}", info.msg())));
        Ok(())
    }
    fn format_full(&self,
                   io: &mut io::Write,
                   info: &Record,
                   logger_values: &OwnedKeyValueList)
                   -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);

        try!(self.print_msg_header(io, &r_decorator, info));
        let mut serializer = Serializer::new(io, r_decorator);

        for &(ref k, ref v) in logger_values.iter() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }

        for &(k, v) in info.values() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }
        let (mut io, _decorator_r) = serializer.finish();

        let _ = try!(write!(io, "\n"));

        Ok(())
    }

    fn format_compact(&self,
                      io: &mut io::Write,
                      info: &Record,
                      logger_values: &OwnedKeyValueList)
                      -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);
        let mut ser = Serializer::new(io, r_decorator);

        let indent = try!(self.format_recurse(&mut ser, info, logger_values));

        try!(self.print_indent(&mut ser.io, indent));

        try!(self.print_msg_header(&mut ser.io, &ser.decorator, info));

        for &(k, v) in info.values() {
            try!(ser.print_comma());
            try!(v.serialize(info, k, &mut ser));
        }
        try!(write!(&mut ser.io, "\n"));

        Ok(())
    }

    fn print_indent<W: io::Write>(&self, io: &mut W, indent: usize) -> io::Result<()> {
        for _ in 0..indent {
            try!(write!(io, "  "));
        }
        Ok(())
    }

    // record in the history, and check if should print
    // given set of values
    fn should_print(&self, address: usize, indent: usize) -> bool {
        let mut history = self.history.lock().unwrap();
        if history.len() <= indent {
            debug_assert_eq!(history.len(), indent);
            history.push(address);
            true
        } else {
            let should = history[indent] != address;
            history[indent] = address;
            should
        }
    }

    fn format_recurse<W: io::Write>(&self,
                                    ser: &mut Serializer<W, D::RecordDecorator>,
                                    info: &slog::Record,
                                    logger_values: &slog::OwnedKeyValueList)
                                    -> io::Result<usize> {
        let mut indent = if logger_values.parent().is_none() {
            0
        } else {
            try!(self.format_recurse(ser, info, logger_values.parent().as_ref().unwrap()))
        };

        if !logger_values.values().is_empty() {
            if self.should_print(logger_values.values() as *const _ as usize, indent) {
                try!(self.print_indent(&mut ser.io, indent));
                let mut clean = true;
                for &(ref k, ref v) in logger_values.values() {
                    if !clean {
                        try!(ser.print_comma());
                    }
                    try!(v.serialize(info, k, ser));
                    clean = false;
                }
                try!(write!(&mut ser.io, "\n"));
            }
            indent += 1
        }

        Ok(indent)
    }
}

fn severity_to_color(lvl: Level) -> u8 {
    match lvl {
        Level::Critical => 1,
        Level::Error => 9,
        Level::Warning => 3,
        Level::Info => 2,
        Level::Debug => 6,
        Level::Trace => 4,
    }
}

/// Record decorator (color) for terminal output
pub struct ColorDecorator {
    use_color: bool,
}

impl ColorDecorator {
    /// New decorator that does color records
    pub fn new_colored() -> Self {
        ColorDecorator { use_color: true }
    }
    /// New decorator that does not color records
    pub fn new_plain() -> Self {
        ColorDecorator { use_color: true }
    }
}

/// Particular record decorator (color) for terminal output
pub struct ColorRecordDecorator {
    level_color: Option<Colour>,
    key_style: Option<ansi_term::Style>,
}

impl Decorator for ColorDecorator {
    type RecordDecorator = ColorRecordDecorator;

    fn decorate(&self, record: &Record) -> ColorRecordDecorator {
        if self.use_color {
            ColorRecordDecorator {
                level_color: Some(Colour::Fixed(severity_to_color(record.level()))),
                key_style: Some(ansi_term::Style::new().bold()),
            }
        } else {
            ColorRecordDecorator {
                level_color: None,
                key_style: None,
            }
        }
    }
}

impl RecordDecorator for ColorRecordDecorator {
    fn fmt_level(&self, io: &mut io::Write, args: fmt::Arguments) -> io::Result<()> {
        if let Some(level_color) = self.level_color {
            write!(io, "{}", level_color.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_msg(&self, io: &mut io::Write, args: fmt::Arguments) -> io::Result<()> {
        if let Some(key_style) = self.key_style {
            write!(io, "{}", key_style.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_key(&self, io: &mut io::Write, args: fmt::Arguments) -> io::Result<()> {
        self.fmt_msg(io, args)
    }
}

struct Serializer<W, D: RecordDecorator> {
    io: W,
    decorator: D,
}

impl<W: io::Write, D: RecordDecorator> Serializer<W, D> {
    fn new(io: W, d: D) -> Self {
        Serializer {
            io: io,
            decorator: d,
        }
    }

    fn print_comma(&mut self) -> io::Result<()> {
        try!(self.decorator.fmt_separator(&mut self.io, format_args!(", ")));
        Ok(())
    }

    fn finish(self) -> (W, D) {
        (self.io, self.decorator)
    }
}

macro_rules! s(
    ($s:expr, $k:expr, $v:expr) => {
        try!($s.decorator.fmt_key(&mut $s.io, format_args!("{}", $k)));
        try!($s.decorator.fmt_separator(&mut $s.io, format_args!(": ")));
        try!($s.decorator.fmt_value(&mut $s.io, format_args!("{}", $v)));
    };
);

impl<W: io::Write, D: RecordDecorator> slog::ser::Serializer for Serializer<W, D> {
    fn emit_none(&mut self, key: &str) -> ser::Result {
        s!(self, key, "None");
        Ok(())
    }
    fn emit_unit(&mut self, key: &str) -> ser::Result {
        s!(self, key, "()");
        Ok(())
    }

    fn emit_bool(&mut self, key: &str, val: bool) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

    fn emit_char(&mut self, key: &str, val: char) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

    fn emit_usize(&mut self, key: &str, val: usize) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_isize(&mut self, key: &str, val: isize) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

    fn emit_u8(&mut self, key: &str, val: u8) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i8(&mut self, key: &str, val: i8) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u16(&mut self, key: &str, val: u16) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i16(&mut self, key: &str, val: i16) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u32(&mut self, key: &str, val: u32) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i32(&mut self, key: &str, val: i32) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f32(&mut self, key: &str, val: f32) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u64(&mut self, key: &str, val: u64) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i64(&mut self, key: &str, val: i64) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f64(&mut self, key: &str, val: f64) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_str(&mut self, key: &str, val: &str) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_arguments(&mut self, key: &str, val: &fmt::Arguments) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

}

impl<D: Decorator + Send + Sync> StreamFormat for Format<D> {
    fn format(&self,
              io: &mut io::Write,
              info: &Record,
              logger_values: &OwnedKeyValueList)
              -> io::Result<()> {
        match self.mode {
            FormatMode::Compact => self.format_compact(io, info, logger_values),
            FormatMode::Full => self.format_full(io, info, logger_values),
        }
    }
}

/// Streamer builder
pub struct StreamerBuilder {
    color: Option<bool>, // None = auto
    stdout: bool,
    async: bool,
    mode: FormatMode,
}

impl StreamerBuilder {
    /// New `StreamerBuilder`
    pub fn new() -> Self {
        StreamerBuilder {
            color: None,
            stdout: true,
            async: false,
            mode: FormatMode::Full,
        }
    }

    /// Force colored output
    pub fn color(mut self) -> Self {
        self.color = Some(true);
        self
    }

    /// Force plain output
    pub fn plain(mut self) -> Self {
        self.color = Some(false);
        self
    }

    /// Auto detect color (default)
    pub fn auto_color(mut self) -> Self {
        self.color = None;
        self
    }

    /// Output to stderr
    pub fn stderr(mut self) -> Self {
        self.stdout = false;
        self
    }

    /// Output to stdout (default)
    pub fn stdout(mut self) -> Self {
        self.stdout = true;
        self
    }

    /// Output using full mode
    pub fn full(mut self) -> Self {
        self.mode = FormatMode::Full;
        self
    }

    /// Output using compact mode (default)
    pub fn compact(mut self) -> Self {
        self.mode = FormatMode::Compact;
        self
    }

    /// Use asynchronous streamer
    pub fn async(mut self) -> Self {
        self.async = true;
        self
    }

    /// Use synchronous streamer (default)
    pub fn sync(mut self) -> Self {
        self.async = false;
        self
    }

    /// Build the streamer
    pub fn build(self) -> Box<slog::Drain<Error=io::Error>> {
        let color = self.color.unwrap_or(if self.stdout {
            stdout_isatty()
        } else {
            stderr_isatty()
        });

        let format = Format::new(self.mode, ColorDecorator { use_color: color });
        let io = if self.stdout {
            Box::new(io::stdout()) as Box<io::Write + Send>
        } else {
            Box::new(io::stderr()) as Box<io::Write + Send>
        };

        if self.async {
            Box::new(AsyncStreamer::new(io, format))
        } else {
            Box::new(Streamer::new(io, format))
        }
    }
}

/// Build `slog_stream::Streamer`/`slog_stream::AsyncStreamer` that
/// will output logging records to stderr/stderr.
pub fn streamer() -> StreamerBuilder {
    StreamerBuilder::new()
}
<commit_msg>feat(term): Add option for custom timestamp<commit_after>//! Unix terminal formatter and drain for slog-rs
//!
//! ```
//! #[macro_use]
//! extern crate slog;
//! extern crate slog_term;
//!
//! use slog::*;
//!
//! fn main() {
//!     let root = Logger::root(slog_term::streamer().build().fuse(), o!("build-id" => "8dfljdf"));
//! }
//! ```
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate slog_stream;
extern crate isatty;
extern crate ansi_term;
extern crate chrono;

use std::{io, fmt, sync};

use ansi_term::Colour;
use isatty::{stderr_isatty, stdout_isatty};

use slog::Record;
use slog::ser;
use slog::{Level, OwnedKeyValueList};
use slog_stream::Format as StreamFormat;
use slog_stream::{Decorator, RecordDecorator, Streamer, AsyncStreamer};

/// Formatting mode
enum FormatMode {
    /// Compact logging format
    Compact,
    /// Full logging format
    Full,
}

/// Full formatting with optional color support
pub struct Format<D: Decorator> {
    mode: FormatMode,
    decorator: D,
    history: sync::Mutex<Vec<usize>>,
    fn_timestamp: Box<Fn() -> String + Send + Sync>,
}

impl<D: Decorator> Format<D> {
    /// New Format format that prints using color
    fn new(mode: FormatMode, d: D, fn_timestamp: Box<Fn() -> String + Send + Sync>) -> Self {
        Format {
            decorator: d,
            mode: mode,
            history: sync::Mutex::new(vec![]),
            fn_timestamp: fn_timestamp,
        }
    }

    fn print_msg_header(&self,
                        io: &mut io::Write,
                        rd: &D::RecordDecorator,
                        info: &Record)
                        -> io::Result<()> {
        try!(rd.fmt_timestamp(io, format_args!("{} ", (self.fn_timestamp)())));
        try!(rd.fmt_level(io, format_args!("{} ", info.level().as_short_str())));

        try!(rd.fmt_msg(io, format_args!("{}", info.msg())));
        Ok(())
    }
    fn format_full(&self,
                   io: &mut io::Write,
                   info: &Record,
                   logger_values: &OwnedKeyValueList)
                   -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);

        try!(self.print_msg_header(io, &r_decorator, info));
        let mut serializer = Serializer::new(io, r_decorator);

        for &(ref k, ref v) in logger_values.iter() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }

        for &(k, v) in info.values() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }
        let (mut io, _decorator_r) = serializer.finish();

        let _ = try!(write!(io, "\n"));

        Ok(())
    }

    fn format_compact(&self,
                      io: &mut io::Write,
                      info: &Record,
                      logger_values: &OwnedKeyValueList)
                      -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);
        let mut ser = Serializer::new(io, r_decorator);

        let indent = try!(self.format_recurse(&mut ser, info, logger_values));

        try!(self.print_indent(&mut ser.io, indent));

        try!(self.print_msg_header(&mut ser.io, &ser.decorator, info));

        for &(k, v) in info.values() {
            try!(ser.print_comma());
            try!(v.serialize(info, k, &mut ser));
        }
        try!(write!(&mut ser.io, "\n"));

        Ok(())
    }

    fn print_indent<W: io::Write>(&self, io: &mut W, indent: usize) -> io::Result<()> {
        for _ in 0..indent {
            try!(write!(io, "  "));
        }
        Ok(())
    }

    // record in the history, and check if should print
    // given set of values
    fn should_print(&self, address: usize, indent: usize) -> bool {
        let mut history = self.history.lock().unwrap();
        if history.len() <= indent {
            debug_assert_eq!(history.len(), indent);
            history.push(address);
            true
        } else {
            let should = history[indent] != address;
            history[indent] = address;
            should
        }
    }

    fn format_recurse<W: io::Write>(&self,
                                    ser: &mut Serializer<W, D::RecordDecorator>,
                                    info: &slog::Record,
                                    logger_values: &slog::OwnedKeyValueList)
                                    -> io::Result<usize> {
        let mut indent = if logger_values.parent().is_none() {
            0
        } else {
            try!(self.format_recurse(ser, info, logger_values.parent().as_ref().unwrap()))
        };

        if !logger_values.values().is_empty() {
            if self.should_print(logger_values.values() as *const _ as usize, indent) {
                try!(self.print_indent(&mut ser.io, indent));
                let mut clean = true;
                for &(ref k, ref v) in logger_values.values() {
                    if !clean {
                        try!(ser.print_comma());
                    }
                    try!(v.serialize(info, k, ser));
                    clean = false;
                }
                try!(write!(&mut ser.io, "\n"));
            }
            indent += 1
        }

        Ok(indent)
    }
}

fn severity_to_color(lvl: Level) -> u8 {
    match lvl {
        Level::Critical => 1,
        Level::Error => 9,
        Level::Warning => 3,
        Level::Info => 2,
        Level::Debug => 6,
        Level::Trace => 4,
    }
}

/// Record decorator (color) for terminal output
pub struct ColorDecorator {
    use_color: bool,
}

impl ColorDecorator {
    /// New decorator that does color records
    pub fn new_colored() -> Self {
        ColorDecorator { use_color: true }
    }
    /// New decorator that does not color records
    pub fn new_plain() -> Self {
        ColorDecorator { use_color: true }
    }
}

/// Particular record decorator (color) for terminal output
pub struct ColorRecordDecorator {
    level_color: Option<Colour>,
    key_style: Option<ansi_term::Style>,
}

impl Decorator for ColorDecorator {
    type RecordDecorator = ColorRecordDecorator;

    fn decorate(&self, record: &Record) -> ColorRecordDecorator {
        if self.use_color {
            ColorRecordDecorator {
                level_color: Some(Colour::Fixed(severity_to_color(record.level()))),
                key_style: Some(ansi_term::Style::new().bold()),
            }
        } else {
            ColorRecordDecorator {
                level_color: None,
                key_style: None,
            }
        }
    }
}

impl RecordDecorator for ColorRecordDecorator {
    fn fmt_level(&self, io: &mut io::Write, args: fmt::Arguments) -> io::Result<()> {
        if let Some(level_color) = self.level_color {
            write!(io, "{}", level_color.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_msg(&self, io: &mut io::Write, args: fmt::Arguments) -> io::Result<()> {
        if let Some(key_style) = self.key_style {
            write!(io, "{}", key_style.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_key(&self, io: &mut io::Write, args: fmt::Arguments) -> io::Result<()> {
        self.fmt_msg(io, args)
    }
}

struct Serializer<W, D: RecordDecorator> {
    io: W,
    decorator: D,
}

impl<W: io::Write, D: RecordDecorator> Serializer<W, D> {
    fn new(io: W, d: D) -> Self {
        Serializer {
            io: io,
            decorator: d,
        }
    }

    fn print_comma(&mut self) -> io::Result<()> {
        try!(self.decorator.fmt_separator(&mut self.io, format_args!(", ")));
        Ok(())
    }

    fn finish(self) -> (W, D) {
        (self.io, self.decorator)
    }
}

macro_rules! s(
    ($s:expr, $k:expr, $v:expr) => {
        try!($s.decorator.fmt_key(&mut $s.io, format_args!("{}", $k)));
        try!($s.decorator.fmt_separator(&mut $s.io, format_args!(": ")));
        try!($s.decorator.fmt_value(&mut $s.io, format_args!("{}", $v)));
    };
);

impl<W: io::Write, D: RecordDecorator> slog::ser::Serializer for Serializer<W, D> {
    fn emit_none(&mut self, key: &str) -> ser::Result {
        s!(self, key, "None");
        Ok(())
    }
    fn emit_unit(&mut self, key: &str) -> ser::Result {
        s!(self, key, "()");
        Ok(())
    }

    fn emit_bool(&mut self, key: &str, val: bool) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

    fn emit_char(&mut self, key: &str, val: char) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

    fn emit_usize(&mut self, key: &str, val: usize) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_isize(&mut self, key: &str, val: isize) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

    fn emit_u8(&mut self, key: &str, val: u8) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i8(&mut self, key: &str, val: i8) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u16(&mut self, key: &str, val: u16) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i16(&mut self, key: &str, val: i16) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u32(&mut self, key: &str, val: u32) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i32(&mut self, key: &str, val: i32) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f32(&mut self, key: &str, val: f32) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u64(&mut self, key: &str, val: u64) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i64(&mut self, key: &str, val: i64) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f64(&mut self, key: &str, val: f64) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_str(&mut self, key: &str, val: &str) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }
    fn emit_arguments(&mut self, key: &str, val: &fmt::Arguments) -> ser::Result {
        s!(self, key, val);
        Ok(())
    }

}

impl<D: Decorator + Send + Sync> StreamFormat for Format<D> {
    fn format(&self,
              io: &mut io::Write,
              info: &Record,
              logger_values: &OwnedKeyValueList)
              -> io::Result<()> {
        match self.mode {
            FormatMode::Compact => self.format_compact(io, info, logger_values),
            FormatMode::Full => self.format_full(io, info, logger_values),
        }
    }
}

const TIMESTAMP_FORMAT: &'static str = "%b %d %H:%M:%S%.3f";

fn timestamp_local() -> String {
    chrono::Local::now().format(TIMESTAMP_FORMAT).to_string()
}

fn timestamp_utc() -> String {
    chrono::UTC::now().format(TIMESTAMP_FORMAT).to_string()
}

/// Streamer builder
pub struct StreamerBuilder {
    color: Option<bool>, // None = auto
    stdout: bool,
    async: bool,
    mode: FormatMode,
    fn_timestamp: Box<Fn() -> String + Send + Sync>,
}

impl StreamerBuilder {
    /// New `StreamerBuilder`
    pub fn new() -> Self {
        StreamerBuilder {
            color: None,
            stdout: true,
            async: false,
            mode: FormatMode::Full,
            fn_timestamp: Box::new(timestamp_local),
        }
    }

    /// Force colored output
    pub fn color(mut self) -> Self {
        self.color = Some(true);
        self
    }

    /// Force plain output
    pub fn plain(mut self) -> Self {
        self.color = Some(false);
        self
    }

    /// Auto detect color (default)
    pub fn auto_color(mut self) -> Self {
        self.color = None;
        self
    }

    /// Output to stderr
    pub fn stderr(mut self) -> Self {
        self.stdout = false;
        self
    }

    /// Output to stdout (default)
    pub fn stdout(mut self) -> Self {
        self.stdout = true;
        self
    }

    /// Output using full mode
    pub fn full(mut self) -> Self {
        self.mode = FormatMode::Full;
        self
    }

    /// Output using compact mode (default)
    pub fn compact(mut self) -> Self {
        self.mode = FormatMode::Compact;
        self
    }

    /// Use asynchronous streamer
    pub fn async(mut self) -> Self {
        self.async = true;
        self
    }

    /// Use synchronous streamer (default)
    pub fn sync(mut self) -> Self {
        self.async = false;
        self
    }

    /// Use the UTC time zone for the timestamp
    pub fn use_utc_timestamp(mut self) -> Self {
        self.fn_timestamp = Box::new(timestamp_utc);
        self
    }

    /// Use the local time zone for the timestamp (default)
    pub fn use_local_timestamp(mut self) -> Self {
        self.fn_timestamp = Box::new(timestamp_local);
        self
    }

    /// Provide a custom function to generate the timestamp
    pub fn use_custom_timestamp(mut self, f: Box<Fn() -> String + Send + Sync>) -> Self  {
        self.fn_timestamp = f;
        self
    }

    /// Build the streamer
    pub fn build(self) -> Box<slog::Drain<Error=io::Error>> {
        let color = self.color.unwrap_or(if self.stdout {
            stdout_isatty()
        } else {
            stderr_isatty()
        });

        let format = Format::new(
            self.mode,
            ColorDecorator { use_color: color },
            self.fn_timestamp
        );

        let io = if self.stdout {
            Box::new(io::stdout()) as Box<io::Write + Send>
        } else {
            Box::new(io::stderr()) as Box<io::Write + Send>
        };

        if self.async {
            Box::new(AsyncStreamer::new(io, format))
        } else {
            Box::new(Streamer::new(io, format))
        }
    }
}

/// Build `slog_stream::Streamer`/`slog_stream::AsyncStreamer` that
/// will output logging records to stderr/stderr.
pub fn streamer() -> StreamerBuilder {
    StreamerBuilder::new()
}
<|endoftext|>