<commit_before>#![crate_name = "uusync"]
/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Alexander Fomin <xander.fomin@ya.ru>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

 /* Last synced with: sync (GNU coreutils) 8.13 */

 #![feature(macro_rules)]

extern crate getopts;
extern crate libc;

use getopts::{optflag, getopts, usage};

#[path = "../common/util.rs"] mod util;

static NAME: &'static str = "sync";
static VERSION: &'static str = "1.0.0";

#[cfg(unix)]
mod platform {
    use super::libc;

    extern {
        fn sync() -> libc::c_void;
    }

    pub unsafe fn do_sync() -> int {
        sync();
        0
    }
}

#[cfg(windows)]
mod platform {
    pub use super::libc;
    use std::{mem, str};
    use std::ptr::null;

    extern "system" {
        fn CreateFileA(lpFileName: *const libc::c_char,
                      dwDesiredAccess: libc::uint32_t,
                      dwShareMode: libc::uint32_t,
                      lpSecurityAttributes: *const libc::c_void, // *LPSECURITY_ATTRIBUTES
                      dwCreationDisposition: libc::uint32_t,
                      dwFlagsAndAttributes: libc::uint32_t,
                      hTemplateFile: *const libc::c_void) -> *const libc::c_void;
        fn GetDriveTypeA(lpRootPathName: *const libc::c_char) -> libc::c_uint;
        fn GetLastError() -> libc::uint32_t;
        fn FindFirstVolumeA(lpszVolumeName: *mut libc::c_char,
                            cchBufferLength: libc::uint32_t) -> *const libc::c_void;
        fn FindNextVolumeA(hFindVolume: *const libc::c_void,
                           lpszVolumeName: *mut libc::c_char,
                           cchBufferLength: libc::uint32_t) -> libc::c_int;
        fn FindVolumeClose(hFindVolume: *const libc::c_void) -> libc::c_int;
        fn FlushFileBuffers(hFile: *const libc::c_void) -> libc::c_int;
    }

    #[allow(unused_unsafe)]
    unsafe fn flush_volume(name: &str) {
        name.to_c_str().with_ref(|name_buffer| {
            if 0x00000003 == GetDriveTypeA(name_buffer) { // DRIVE_FIXED
                let sliced_name = name.slice_to(name.len() - 1); // eliminate trailing backslash
                sliced_name.to_c_str().with_ref(|sliced_name_buffer| {
                    match CreateFileA(sliced_name_buffer,
                                      0xC0000000, // GENERIC_WRITE
                                      0x00000003, // FILE_SHARE_WRITE,
                                      null(),
                                      0x00000003, // OPEN_EXISTING
                                      0,
                                      null()) {
                        _x if _x == -1 as *const libc::c_void => { // INVALID_HANDLE_VALUE
                            crash!(GetLastError(), "failed to create volume handle");
                        }
                        handle @ _ => {
                            if FlushFileBuffers(handle) == 0 {
                                crash!(GetLastError(), "failed to flush file buffer");
                            }
                        }
                    }
                });
            }
        });
    }

    #[allow(unused_unsafe)]
    unsafe fn find_first_volume() -> (String, *const libc::c_void) {
        let mut name: [libc::c_char, ..260] = mem::uninitialized(); // MAX_PATH
        match FindFirstVolumeA(name.as_mut_ptr(),
                               name.len() as libc::uint32_t) {
            _x if _x == -1 as *const libc::c_void => { // INVALID_HANDLE_VALUE
                crash!(GetLastError(), "failed to find first volume");
            }
            handle @ _ => {
                (str::raw::from_c_str(name.as_ptr()), handle)
            }
        }
    }

    #[allow(unused_unsafe)]
    unsafe fn find_all_volumes() -> Vec<String> {
        match find_first_volume() {
            (first_volume, next_volume_handle) => {
                let mut volumes = Vec::from_elem(1, first_volume);
                loop {
                    let mut name: [libc::c_char, ..260] = mem::uninitialized(); // MAX_PATH
                    match FindNextVolumeA(next_volume_handle,
                                          name.as_mut_ptr(),
                                          name.len() as libc::uint32_t) {
                        0 => {
                            match GetLastError() {
                                0x12 => { // ERROR_NO_MORE_FILES
                                    FindVolumeClose(next_volume_handle); // ignore FindVolumeClose() failures
                                    break;
                                }
                                err @ _ => {
                                    crash!(err, "failed to find next volume");
                                }
                            }
                        }
                        _ => {
                            volumes.push(str::raw::from_c_str(name.as_ptr()));
                        }
                    }
                }
                volumes
            }
        }
    }

    pub unsafe fn do_sync() -> int {
        let volumes = find_all_volumes();
        for vol in volumes.iter() {
            flush_volume(vol.as_slice());
        }
        0
    }
}

pub fn uumain(args: Vec<String>) -> int {
    let program = &args[0];

    let options = [
        optflag("h", "help", "display this help and exit"),
        optflag("V", "version", "output version information and exit")
    ];

    let matches = match getopts(args.tail(), options) {
        Ok(m) => { m }
        _ => { help(program.as_slice(), options); return 1 }
    };

    if matches.opt_present("h") {
        help(program.as_slice(), options);
        return 0
    }

    if matches.opt_present("V") {
        version();
        return 0
    }

    uusync();
    0
}

fn version() {
    println!("{} (uutils) {}", NAME, VERSION);
    println!("The MIT License");
    println!("");
    println!("Author -- Alexander Fomin.");
}

fn help(program: &str, options: &[getopts::OptGroup]) {
    println!("Usage: {:s} [OPTION]", program);
    print!("{:s}", usage("Force changed blocks to disk, update the super block.", options));
}

fn uusync() -> int {
    unsafe {
        platform::do_sync()
    }
}
<commit_msg>update deprecated parts of sync<commit_after>#![crate_name = "uusync"]
/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Alexander Fomin <xander.fomin@ya.ru>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

 /* Last synced with: sync (GNU coreutils) 8.13 */

 #![feature(macro_rules)]

extern crate getopts;
extern crate libc;

use getopts::{optflag, getopts, usage};

#[path = "../common/util.rs"] mod util;

static NAME: &'static str = "sync";
static VERSION: &'static str = "1.0.0";

#[cfg(unix)]
mod platform {
    use super::libc;

    extern {
        fn sync() -> libc::c_void;
    }

    pub unsafe fn do_sync() -> int {
        sync();
        0
    }
}

#[cfg(windows)]
mod platform {
    pub use super::libc;
    use std::{mem, string};
    use std::ptr::null;

    extern "system" {
        fn CreateFileA(lpFileName: *const libc::c_char,
                      dwDesiredAccess: libc::uint32_t,
                      dwShareMode: libc::uint32_t,
                      lpSecurityAttributes: *const libc::c_void, // *LPSECURITY_ATTRIBUTES
                      dwCreationDisposition: libc::uint32_t,
                      dwFlagsAndAttributes: libc::uint32_t,
                      hTemplateFile: *const libc::c_void) -> *const libc::c_void;
        fn GetDriveTypeA(lpRootPathName: *const libc::c_char) -> libc::c_uint;
        fn GetLastError() -> libc::uint32_t;
        fn FindFirstVolumeA(lpszVolumeName: *mut libc::c_char,
                            cchBufferLength: libc::uint32_t) -> *const libc::c_void;
        fn FindNextVolumeA(hFindVolume: *const libc::c_void,
                           lpszVolumeName: *mut libc::c_char,
                           cchBufferLength: libc::uint32_t) -> libc::c_int;
        fn FindVolumeClose(hFindVolume: *const libc::c_void) -> libc::c_int;
        fn FlushFileBuffers(hFile: *const libc::c_void) -> libc::c_int;
    }

    #[allow(unused_unsafe)]
    unsafe fn flush_volume(name: &str) {
        let name_buffer = name.to_c_str().as_ptr();
        if 0x00000003 == GetDriveTypeA(name_buffer) { // DRIVE_FIXED
            let sliced_name = name.slice_to(name.len() - 1); // eliminate trailing backslash
            let sliced_name_buffer = sliced_name.to_c_str().as_ptr();
            match CreateFileA(sliced_name_buffer,
                              0xC0000000, // GENERIC_WRITE
                              0x00000003, // FILE_SHARE_WRITE,
                              null(),
                              0x00000003, // OPEN_EXISTING
                              0,
                              null()) {
                _x if _x == -1 as *const libc::c_void => { // INVALID_HANDLE_VALUE
                    crash!(GetLastError(), "failed to create volume handle");
                }
                handle @ _ => {
                    if FlushFileBuffers(handle) == 0 {
                        crash!(GetLastError(), "failed to flush file buffer");
                    }
                }
            }
        }
    }

    #[allow(unused_unsafe)]
    unsafe fn find_first_volume() -> (String, *const libc::c_void) {
        let mut name: [libc::c_char, ..260] = mem::uninitialized(); // MAX_PATH
        match FindFirstVolumeA(name.as_mut_ptr(),
                               name.len() as libc::uint32_t) {
            _x if _x == -1 as *const libc::c_void => { // INVALID_HANDLE_VALUE
                crash!(GetLastError(), "failed to find first volume");
            }
            handle @ _ => {
                (string::raw::from_buf(name.as_ptr() as *const u8), handle)
            }
        }
    }

    #[allow(unused_unsafe)]
    unsafe fn find_all_volumes() -> Vec<String> {
        match find_first_volume() {
            (first_volume, next_volume_handle) => {
                let mut volumes = Vec::from_elem(1, first_volume);
                loop {
                    let mut name: [libc::c_char, ..260] = mem::uninitialized(); // MAX_PATH
                    match FindNextVolumeA(next_volume_handle,
                                          name.as_mut_ptr(),
                                          name.len() as libc::uint32_t) {
                        0 => {
                            match GetLastError() {
                                0x12 => { // ERROR_NO_MORE_FILES
                                    FindVolumeClose(next_volume_handle); // ignore FindVolumeClose() failures
                                    break;
                                }
                                err @ _ => {
                                    crash!(err, "failed to find next volume");
                                }
                            }
                        }
                        _ => {
                            volumes.push(string::raw::from_buf(name.as_ptr() as *const u8));
                        }
                    }
                }
                volumes
            }
        }
    }

    pub unsafe fn do_sync() -> int {
        let volumes = find_all_volumes();
        for vol in volumes.iter() {
            flush_volume(vol.as_slice());
        }
        0
    }
}

pub fn uumain(args: Vec<String>) -> int {
    let program = &args[0];

    let options = [
        optflag("h", "help", "display this help and exit"),
        optflag("V", "version", "output version information and exit")
    ];

    let matches = match getopts(args.tail(), options) {
        Ok(m) => { m }
        _ => { help(program.as_slice(), options); return 1 }
    };

    if matches.opt_present("h") {
        help(program.as_slice(), options);
        return 0
    }

    if matches.opt_present("V") {
        version();
        return 0
    }

    uusync();
    0
}

fn version() {
    println!("{} (uutils) {}", NAME, VERSION);
    println!("The MIT License");
    println!("");
    println!("Author -- Alexander Fomin.");
}

fn help(program: &str, options: &[getopts::OptGroup]) {
    println!("Usage: {:s} [OPTION]", program);
    print!("{:s}", usage("Force changed blocks to disk, update the super block.", options));
}

fn uusync() -> int {
    unsafe {
        platform::do_sync()
    }
}
<|endoftext|>


<commit_before>/* TOOD: Implement for other kqueue based systems
 */

use {Errno, Result};
#[cfg(not(target_os = "netbsd"))]
use libc::{timespec, time_t, c_int, c_long, intptr_t, uintptr_t};
#[cfg(target_os = "netbsd")]
use libc::{timespec, time_t, c_long, intptr_t, uintptr_t, size_t};
use libc;
use std::os::unix::io::RawFd;
use std::ptr;
use std::mem;

// Redefine kevent in terms of programmer-friendly enums and bitfields.
#[derive(Clone, Copy)]
#[repr(C)]
pub struct KEvent {
    kevent: libc::kevent,
}

#[cfg(any(target_os = "openbsd", target_os = "freebsd",
          target_os = "dragonfly", target_os = "macos"))]
type type_of_udata = *mut ::c_void;
#[cfg(any(target_os = "netbsd"))]
type type_of_udata = intptr_t;

#[cfg(not(target_os = "netbsd"))]
type type_of_event_filter = i16;
#[cfg(not(target_os = "netbsd"))]
#[repr(i16)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum EventFilter {
    EVFILT_AIO = libc::EVFILT_AIO,
    #[cfg(target_os = "dragonfly")]
    EVFILT_EXCEPT = libc::EVFILT_EXCEPT,
    #[cfg(any(target_os = "macos",
              target_os = "dragonfly",
              target_os = "freebsd"))]
    EVFILT_FS = libc::EVFILT_FS,
    #[cfg(target_os = "freebsd")]
    EVFILT_LIO = libc::EVFILT_LIO,
    #[cfg(target_os = "macos")]
    EVFILT_MACHPORT = libc::EVFILT_MACHPORT,
    EVFILT_PROC = libc::EVFILT_PROC,
    EVFILT_READ = libc::EVFILT_READ,
    EVFILT_SIGNAL = libc::EVFILT_SIGNAL,
    EVFILT_TIMER = libc::EVFILT_TIMER,
    #[cfg(any(target_os = "macos",
              target_os = "dragonfly",
              target_os = "freebsd"))]
    EVFILT_USER = libc::EVFILT_USER,
    #[cfg(target_os = "macos")]
    EVFILT_VM = libc::EVFILT_VM,
    EVFILT_VNODE = libc::EVFILT_VNODE,
    EVFILT_WRITE = libc::EVFILT_WRITE,
}

#[cfg(target_os = "netbsd")]
type type_of_event_filter = i32;
#[cfg(target_os = "netbsd")]
#[repr(u32)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum EventFilter {
    EVFILT_READ = libc::EVFILT_READ,
    EVFILT_WRITE = libc::EVFILT_WRITE,
    EVFILT_AIO = libc::EVFILT_AIO,
    EVFILT_VNODE = libc::EVFILT_VNODE,
    EVFILT_PROC = libc::EVFILT_PROC,
    EVFILT_SIGNAL = libc::EVFILT_SIGNAL,
    EVFILT_TIMER = libc::EVFILT_TIMER,
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly"))]
type type_of_event_flag = u16;
#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly"))]
bitflags!(
    flags EventFlag: u16 {
        const EV_ADD       = libc::EV_ADD,
        const EV_CLEAR     = libc::EV_CLEAR,
        const EV_DELETE    = libc::EV_DELETE,
        const EV_DISABLE   = libc::EV_DISABLE,
        const EV_DISPATCH  = libc::EV_DISPATCH,
        #[cfg(target_os = "freebsd")]
        const EV_DROP      = libc::EV_DROP,
        const EV_ENABLE    = libc::EV_ENABLE,
        const EV_EOF       = libc::EV_EOF,
        const EV_ERROR     = libc::EV_ERROR,
        #[cfg(target_os = "macos")]
        const EV_FLAG0     = libc::EV_FLAG0,
        const EV_FLAG1     = libc::EV_FLAG1,
        #[cfg(target_os = "dragonfly")]
        const EV_NODATA    = libc::EV_NODATA,
        const EV_ONESHOT   = libc::EV_ONESHOT,
        #[cfg(target_os = "macos")]
        const EV_OOBAND    = libc::EV_OOBAND,
        #[cfg(target_os = "macos")]
        const EV_POLL      = libc::EV_POLL,
        const EV_RECEIPT   = libc::EV_RECEIPT,
        const EV_SYSFLAGS  = libc::EV_SYSFLAGS,
    }
);

#[cfg(any(target_os = "netbsd", target_os = "openbsd"))]
type type_of_event_flag = u32;
#[cfg(any(target_os = "netbsd", target_os = "openbsd"))]
bitflags!(
    flags EventFlag: u32 {
        const EV_ADD       = libc::EV_ADD,
        const EV_DELETE    = libc::EV_DELETE,
        const EV_ENABLE    = libc::EV_ENABLE,
        const EV_DISABLE   = libc::EV_DISABLE,
        const EV_ONESHOT   = libc::EV_ONESHOT,
        const EV_CLEAR     = libc::EV_CLEAR,
        #[cfg(target_os = "openbsd")]
        const EV_FLAG1     = libc::EV_FLAG1,
        #[cfg(target_os = "netbsd")]
        const EV_RECEIPT   = libc::EV_RECEIPT,
        const EV_DISPATCH  = libc::EV_DISPATCH,
        const EV_SYSFLAGS  = libc::EV_SYSFLAGS,
        const EV_FLAG1     = libc::EV_FLAG1,
        const EV_EOF       = libc::EV_EOF,
        const EV_ERROR     = libc::EV_ERROR,
    }
);

bitflags!(
    flags FilterFlag: u32 {
        #[cfg(target_os = "macos")]
        const NOTE_ABSOLUTE                        = libc::NOTE_ABSOLUTE,
        const NOTE_ATTRIB                          = libc::NOTE_ATTRIB,
        const NOTE_CHILD                           = libc::NOTE_CHILD,
        const NOTE_DELETE                          = libc::NOTE_DELETE,
        #[cfg(target_os = "openbsd")]
        const NOTE_EOF                             = libc::NOTE_EOF,
        const NOTE_EXEC                            = libc::NOTE_EXEC,
        const NOTE_EXIT                            = libc::NOTE_EXIT,
        #[cfg(target_os = "macos")]
        const NOTE_EXIT_REPARENTED                 = libc::NOTE_EXIT_REPARENTED,
        #[cfg(target_os = "macos")]
        const NOTE_EXITSTATUS                      = libc::NOTE_EXITSTATUS,
        const NOTE_EXTEND                          = libc::NOTE_EXTEND,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFAND                           = libc::NOTE_FFAND,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFCOPY                          = libc::NOTE_FFCOPY,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFCTRLMASK                      = libc::NOTE_FFCTRLMASK,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFLAGSMASK                      = libc::NOTE_FFLAGSMASK,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFNOP                           = libc::NOTE_FFNOP,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFOR                            = libc::NOTE_FFOR,
        const NOTE_FORK                            = libc::NOTE_FORK,
        const NOTE_LINK                            = libc::NOTE_LINK,
        const NOTE_LOWAT                           = libc::NOTE_LOWAT,
        #[cfg(target_os = "freebsd")]
        const NOTE_MSECONDS                        = libc::NOTE_MSECONDS,
        #[cfg(target_os = "macos")]
        const NOTE_NONE                            = libc::NOTE_NONE,
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const NOTE_NSECONDS                        = libc::NOTE_NSECONDS,
        #[cfg(target_os = "dragonfly")]
        const NOTE_OOB                             = libc::NOTE_OOB,
        const NOTE_PCTRLMASK                       = libc::NOTE_PCTRLMASK,
        const NOTE_PDATAMASK                       = libc::NOTE_PDATAMASK,
        #[cfg(target_os = "macos")]
        const NOTE_REAP                            = libc::NOTE_REAP,
        const NOTE_RENAME                          = libc::NOTE_RENAME,
        const NOTE_REVOKE                          = libc::NOTE_REVOKE,
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const NOTE_SECONDS                         = libc::NOTE_SECONDS,
        #[cfg(target_os = "macos")]
        const NOTE_SIGNAL                          = libc::NOTE_SIGNAL,
        const NOTE_TRACK                           = libc::NOTE_TRACK,
        const NOTE_TRACKERR                        = libc::NOTE_TRACKERR,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_TRIGGER                         = libc::NOTE_TRIGGER,
        #[cfg(target_os = "openbsd")]
        const NOTE_TRUNCATE                        = libc::NOTE_TRUNCATE,
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const NOTE_USECONDS                        = libc::NOTE_USECONDS,
        #[cfg(target_os = "macos")]
        const NOTE_VM_ERROR                        = libc::NOTE_VM_ERROR,
        #[cfg(target_os = "macos")]
        const NOTE_VM_PRESSURE                     = libc::NOTE_VM_PRESSURE,
        #[cfg(target_os = "macos")]
        const NOTE_VM_PRESSURE_SUDDEN_TERMINATE    = libc::NOTE_VM_PRESSURE_SUDDEN_TERMINATE,
        #[cfg(target_os = "macos")]
        const NOTE_VM_PRESSURE_TERMINATE           = libc::NOTE_VM_PRESSURE_TERMINATE,
        const NOTE_WRITE                           = libc::NOTE_WRITE,
    }
);

pub fn kqueue() -> Result<RawFd> {
    let res = unsafe { libc::kqueue() };

    Errno::result(res)
}

impl KEvent {
    pub fn new(ident: uintptr_t, filter: EventFilter, flags: EventFlag,
               fflags:FilterFlag, data: intptr_t, udata: uintptr_t) -> KEvent {
        KEvent { kevent: libc::kevent {
            ident: ident,
            filter: filter as type_of_event_filter,
            flags: flags.bits(),
            fflags: fflags.bits(),
            data: data,
            udata: udata as type_of_udata
        } }
    }

    pub fn ident(&self) -> uintptr_t {
        self.kevent.ident
    }

    pub fn filter(&self) -> EventFilter {
        unsafe { mem::transmute(self.kevent.filter as type_of_event_filter) }
    }

    pub fn flags(&self) -> EventFlag {
        EventFlag::from_bits(self.kevent.flags).unwrap()
    }

    pub fn fflags(&self) -> FilterFlag {
        FilterFlag::from_bits(self.kevent.fflags).unwrap()
    }

    pub fn data(&self) -> intptr_t {
        self.kevent.data
    }

    pub fn udata(&self) -> uintptr_t {
        self.kevent.udata as uintptr_t
    }
}

pub fn kevent(kq: RawFd,
              changelist: &[KEvent],
              eventlist: &mut [KEvent],
              timeout_ms: usize) -> Result<usize> {

    // Convert ms to timespec
    let timeout = timespec {
        tv_sec: (timeout_ms / 1000) as time_t,
        tv_nsec: ((timeout_ms % 1000) * 1_000_000) as c_long
    };

    kevent_ts(kq, changelist, eventlist, Some(timeout))
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "openbsd"))]
type type_of_nchanges = c_int;
#[cfg(target_os = "netbsd")]
type type_of_nchanges = size_t;

pub fn kevent_ts(kq: RawFd,
              changelist: &[KEvent],
              eventlist: &mut [KEvent],
              timeout_opt: Option<timespec>) -> Result<usize> {

    let res = unsafe {
        libc::kevent(
            kq,
            changelist.as_ptr() as *const libc::kevent,
            changelist.len() as type_of_nchanges,
            eventlist.as_mut_ptr() as *mut libc::kevent,
            eventlist.len() as type_of_nchanges,
            if let Some(ref timeout) = timeout_opt {timeout as *const timespec} else {ptr::null()})
    };

    Errno::result(res).map(|r| r as usize)
}

#[inline]
pub fn ev_set(ev: &mut KEvent,
              ident: usize,
              filter: EventFilter,
              flags: EventFlag,
              fflags: FilterFlag,
              udata: uintptr_t) {

    ev.kevent.ident  = ident as uintptr_t;
    ev.kevent.filter = filter as type_of_event_filter;
    ev.kevent.flags  = flags.bits();
    ev.kevent.fflags = fflags.bits();
    ev.kevent.data   = 0;
    ev.kevent.udata  = udata as type_of_udata;
}

#[test]
fn test_struct_kevent() {
    let udata : uintptr_t = 12345;

    let expected = libc::kevent{ident: 0xdeadbeef,
                                filter: libc::EVFILT_READ,
                                flags: libc::EV_DISPATCH | libc::EV_ADD,
                                fflags: libc::NOTE_CHILD | libc::NOTE_EXIT,
                                data: 0x1337,
                                udata: udata as type_of_udata};
    let actual = KEvent::new(0xdeadbeef,
                             EventFilter::EVFILT_READ,
                             EV_DISPATCH | EV_ADD,
                             NOTE_CHILD | NOTE_EXIT,
                             0x1337,
                             udata);
    assert!(expected.ident == actual.ident());
    assert!(expected.filter == actual.filter() as type_of_event_filter);
    assert!(expected.flags == actual.flags().bits());
    assert!(expected.fflags == actual.fflags().bits());
    assert!(expected.data == actual.data());
    assert!(expected.udata == actual.udata() as type_of_udata);
    assert!(mem::size_of::<libc::kevent>() == mem::size_of::<KEvent>());
}
<commit_msg>Use libc_bitflags! for EventFlag<commit_after>/* TOOD: Implement for other kqueue based systems
 */

use {Errno, Result};
#[cfg(not(target_os = "netbsd"))]
use libc::{timespec, time_t, c_int, c_long, intptr_t, uintptr_t};
#[cfg(target_os = "netbsd")]
use libc::{timespec, time_t, c_long, intptr_t, uintptr_t, size_t};
use libc;
use std::os::unix::io::RawFd;
use std::ptr;
use std::mem;

// Redefine kevent in terms of programmer-friendly enums and bitfields.
#[derive(Clone, Copy)]
#[repr(C)]
pub struct KEvent {
    kevent: libc::kevent,
}

#[cfg(any(target_os = "openbsd", target_os = "freebsd",
          target_os = "dragonfly", target_os = "macos"))]
type type_of_udata = *mut ::c_void;
#[cfg(any(target_os = "netbsd"))]
type type_of_udata = intptr_t;

#[cfg(not(target_os = "netbsd"))]
type type_of_event_filter = i16;
#[cfg(not(target_os = "netbsd"))]
#[repr(i16)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum EventFilter {
    EVFILT_AIO = libc::EVFILT_AIO,
    #[cfg(target_os = "dragonfly")]
    EVFILT_EXCEPT = libc::EVFILT_EXCEPT,
    #[cfg(any(target_os = "macos",
              target_os = "dragonfly",
              target_os = "freebsd"))]
    EVFILT_FS = libc::EVFILT_FS,
    #[cfg(target_os = "freebsd")]
    EVFILT_LIO = libc::EVFILT_LIO,
    #[cfg(target_os = "macos")]
    EVFILT_MACHPORT = libc::EVFILT_MACHPORT,
    EVFILT_PROC = libc::EVFILT_PROC,
    EVFILT_READ = libc::EVFILT_READ,
    EVFILT_SIGNAL = libc::EVFILT_SIGNAL,
    EVFILT_TIMER = libc::EVFILT_TIMER,
    #[cfg(any(target_os = "macos",
              target_os = "dragonfly",
              target_os = "freebsd"))]
    EVFILT_USER = libc::EVFILT_USER,
    #[cfg(target_os = "macos")]
    EVFILT_VM = libc::EVFILT_VM,
    EVFILT_VNODE = libc::EVFILT_VNODE,
    EVFILT_WRITE = libc::EVFILT_WRITE,
}

#[cfg(target_os = "netbsd")]
type type_of_event_filter = i32;
#[cfg(target_os = "netbsd")]
#[repr(u32)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum EventFilter {
    EVFILT_READ = libc::EVFILT_READ,
    EVFILT_WRITE = libc::EVFILT_WRITE,
    EVFILT_AIO = libc::EVFILT_AIO,
    EVFILT_VNODE = libc::EVFILT_VNODE,
    EVFILT_PROC = libc::EVFILT_PROC,
    EVFILT_SIGNAL = libc::EVFILT_SIGNAL,
    EVFILT_TIMER = libc::EVFILT_TIMER,
}

#[cfg(any(target_os = "macos", target_os = "freebsd", target_os = "dragonfly"))]
pub type type_of_event_flag = u16;
#[cfg(any(target_os = "netbsd", target_os = "openbsd"))]
pub type type_of_event_flag = u32;
libc_bitflags!{
    flags EventFlag: type_of_event_flag {
        EV_ADD,
        EV_CLEAR,
        EV_DELETE,
        EV_DISABLE,
        EV_DISPATCH,
        #[cfg(target_os = "freebsd")]
        EV_DROP,
        EV_ENABLE,
        EV_EOF,
        EV_ERROR,
        #[cfg(target_os = "macos")]
        EV_FLAG0,
        EV_FLAG1,
        #[cfg(target_os = "dragonfly")]
        EV_NODATA,
        EV_ONESHOT,
        #[cfg(target_os = "macos")]
        EV_OOBAND,
        #[cfg(target_os = "macos")]
        EV_POLL,
        #[cfg(not(target_os = "openbsd"))]
        EV_RECEIPT,
        EV_SYSFLAGS,
    }
}

bitflags!(
    flags FilterFlag: u32 {
        #[cfg(target_os = "macos")]
        const NOTE_ABSOLUTE                        = libc::NOTE_ABSOLUTE,
        const NOTE_ATTRIB                          = libc::NOTE_ATTRIB,
        const NOTE_CHILD                           = libc::NOTE_CHILD,
        const NOTE_DELETE                          = libc::NOTE_DELETE,
        #[cfg(target_os = "openbsd")]
        const NOTE_EOF                             = libc::NOTE_EOF,
        const NOTE_EXEC                            = libc::NOTE_EXEC,
        const NOTE_EXIT                            = libc::NOTE_EXIT,
        #[cfg(target_os = "macos")]
        const NOTE_EXIT_REPARENTED                 = libc::NOTE_EXIT_REPARENTED,
        #[cfg(target_os = "macos")]
        const NOTE_EXITSTATUS                      = libc::NOTE_EXITSTATUS,
        const NOTE_EXTEND                          = libc::NOTE_EXTEND,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFAND                           = libc::NOTE_FFAND,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFCOPY                          = libc::NOTE_FFCOPY,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFCTRLMASK                      = libc::NOTE_FFCTRLMASK,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFLAGSMASK                      = libc::NOTE_FFLAGSMASK,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFNOP                           = libc::NOTE_FFNOP,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_FFOR                            = libc::NOTE_FFOR,
        const NOTE_FORK                            = libc::NOTE_FORK,
        const NOTE_LINK                            = libc::NOTE_LINK,
        const NOTE_LOWAT                           = libc::NOTE_LOWAT,
        #[cfg(target_os = "freebsd")]
        const NOTE_MSECONDS                        = libc::NOTE_MSECONDS,
        #[cfg(target_os = "macos")]
        const NOTE_NONE                            = libc::NOTE_NONE,
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const NOTE_NSECONDS                        = libc::NOTE_NSECONDS,
        #[cfg(target_os = "dragonfly")]
        const NOTE_OOB                             = libc::NOTE_OOB,
        const NOTE_PCTRLMASK                       = libc::NOTE_PCTRLMASK,
        const NOTE_PDATAMASK                       = libc::NOTE_PDATAMASK,
        #[cfg(target_os = "macos")]
        const NOTE_REAP                            = libc::NOTE_REAP,
        const NOTE_RENAME                          = libc::NOTE_RENAME,
        const NOTE_REVOKE                          = libc::NOTE_REVOKE,
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const NOTE_SECONDS                         = libc::NOTE_SECONDS,
        #[cfg(target_os = "macos")]
        const NOTE_SIGNAL                          = libc::NOTE_SIGNAL,
        const NOTE_TRACK                           = libc::NOTE_TRACK,
        const NOTE_TRACKERR                        = libc::NOTE_TRACKERR,
        #[cfg(any(target_os = "macos",
                  target_os = "freebsd",
                  target_os = "dragonfly"))]
        const NOTE_TRIGGER                         = libc::NOTE_TRIGGER,
        #[cfg(target_os = "openbsd")]
        const NOTE_TRUNCATE                        = libc::NOTE_TRUNCATE,
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const NOTE_USECONDS                        = libc::NOTE_USECONDS,
        #[cfg(target_os = "macos")]
        const NOTE_VM_ERROR                        = libc::NOTE_VM_ERROR,
        #[cfg(target_os = "macos")]
        const NOTE_VM_PRESSURE                     = libc::NOTE_VM_PRESSURE,
        #[cfg(target_os = "macos")]
        const NOTE_VM_PRESSURE_SUDDEN_TERMINATE    = libc::NOTE_VM_PRESSURE_SUDDEN_TERMINATE,
        #[cfg(target_os = "macos")]
        const NOTE_VM_PRESSURE_TERMINATE           = libc::NOTE_VM_PRESSURE_TERMINATE,
        const NOTE_WRITE                           = libc::NOTE_WRITE,
    }
);

pub fn kqueue() -> Result<RawFd> {
    let res = unsafe { libc::kqueue() };

    Errno::result(res)
}

impl KEvent {
    pub fn new(ident: uintptr_t, filter: EventFilter, flags: EventFlag,
               fflags:FilterFlag, data: intptr_t, udata: uintptr_t) -> KEvent {
        KEvent { kevent: libc::kevent {
            ident: ident,
            filter: filter as type_of_event_filter,
            flags: flags.bits(),
            fflags: fflags.bits(),
            data: data,
            udata: udata as type_of_udata
        } }
    }

    pub fn ident(&self) -> uintptr_t {
        self.kevent.ident
    }

    pub fn filter(&self) -> EventFilter {
        unsafe { mem::transmute(self.kevent.filter as type_of_event_filter) }
    }

    pub fn flags(&self) -> EventFlag {
        EventFlag::from_bits(self.kevent.flags).unwrap()
    }

    pub fn fflags(&self) -> FilterFlag {
        FilterFlag::from_bits(self.kevent.fflags).unwrap()
    }

    pub fn data(&self) -> intptr_t {
        self.kevent.data
    }

    pub fn udata(&self) -> uintptr_t {
        self.kevent.udata as uintptr_t
    }
}

pub fn kevent(kq: RawFd,
              changelist: &[KEvent],
              eventlist: &mut [KEvent],
              timeout_ms: usize) -> Result<usize> {

    // Convert ms to timespec
    let timeout = timespec {
        tv_sec: (timeout_ms / 1000) as time_t,
        tv_nsec: ((timeout_ms % 1000) * 1_000_000) as c_long
    };

    kevent_ts(kq, changelist, eventlist, Some(timeout))
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "openbsd"))]
type type_of_nchanges = c_int;
#[cfg(target_os = "netbsd")]
type type_of_nchanges = size_t;

pub fn kevent_ts(kq: RawFd,
              changelist: &[KEvent],
              eventlist: &mut [KEvent],
              timeout_opt: Option<timespec>) -> Result<usize> {

    let res = unsafe {
        libc::kevent(
            kq,
            changelist.as_ptr() as *const libc::kevent,
            changelist.len() as type_of_nchanges,
            eventlist.as_mut_ptr() as *mut libc::kevent,
            eventlist.len() as type_of_nchanges,
            if let Some(ref timeout) = timeout_opt {timeout as *const timespec} else {ptr::null()})
    };

    Errno::result(res).map(|r| r as usize)
}

#[inline]
pub fn ev_set(ev: &mut KEvent,
              ident: usize,
              filter: EventFilter,
              flags: EventFlag,
              fflags: FilterFlag,
              udata: uintptr_t) {

    ev.kevent.ident  = ident as uintptr_t;
    ev.kevent.filter = filter as type_of_event_filter;
    ev.kevent.flags  = flags.bits();
    ev.kevent.fflags = fflags.bits();
    ev.kevent.data   = 0;
    ev.kevent.udata  = udata as type_of_udata;
}

#[test]
fn test_struct_kevent() {
    let udata : uintptr_t = 12345;

    let expected = libc::kevent{ident: 0xdeadbeef,
                                filter: libc::EVFILT_READ,
                                flags: libc::EV_DISPATCH | libc::EV_ADD,
                                fflags: libc::NOTE_CHILD | libc::NOTE_EXIT,
                                data: 0x1337,
                                udata: udata as type_of_udata};
    let actual = KEvent::new(0xdeadbeef,
                             EventFilter::EVFILT_READ,
                             EV_DISPATCH | EV_ADD,
                             NOTE_CHILD | NOTE_EXIT,
                             0x1337,
                             udata);
    assert!(expected.ident == actual.ident());
    assert!(expected.filter == actual.filter() as type_of_event_filter);
    assert!(expected.flags == actual.flags().bits());
    assert!(expected.fflags == actual.fflags().bits());
    assert!(expected.data == actual.data());
    assert!(expected.udata == actual.udata() as type_of_udata);
    assert!(mem::size_of::<libc::kevent>() == mem::size_of::<KEvent>());
}
<|endoftext|>


<commit_before>pub struct Node<V> {
    c: char,
    left:  Option<Box<Node<V>>>,
    mid:   Option<Box<Node<V>>>,
    right: Option<Box<Node<V>>>,
    val: Option<V>,
}

impl<V> Node<V> {
    fn new(c: char) -> Node<V> {
        Node {
            c: c,
            left: None,
            mid: None,
            right: None,
            val: None
        }
    }

    fn put(mut x: Option<Box<Node<V>>>, key: &str, val: Option<V>, d: usize) -> (Option<Box<Node<V>>>, Option<V>) {
        let replaced;
        let c = key.char_at(d);
        if x.is_none() {
            if val.is_none() {  // no need to call put further
                return (x, None);
            }
            x = Some(Box::new(Node::new(c)));
        }
        let xc = x.as_ref().unwrap().c;
        if c < xc {
            let (left, repl) = Node::put(x.as_mut().unwrap().left.take(), key, val, d);
            x.as_mut().map(|n| n.left = left);
            replaced = repl;
        } else if c > xc {
            let (right, repl) = Node::put(x.as_mut().unwrap().right.take(), key, val, d);
            x.as_mut().map(|n| n.right = right);
            replaced = repl;
        } else if d < key.len()-1 {
            let (mid, repl) = Node::put(x.as_mut().unwrap().mid.take(), key, val, d+1);
            x.as_mut().map(|n| n.mid = mid);
            replaced = repl;
        } else {
            replaced = x.as_mut().unwrap().val.take();
            x.as_mut().map(|n| n.val = val);
        }
        (x, replaced)
    }

    fn get<'a>(x: Option<&'a Box<Node<V>>>, key: &str, d: usize) -> Option<&'a Box<Node<V>>> {
        if x.is_none() {
            return None;
        }
        let c = key.char_at(d);
        let xc = x.unwrap().c;
        if c < xc {
            Node::get(x.unwrap().left.as_ref(), key, d)
        } else if c > xc {
            Node::get(x.unwrap().right.as_ref(), key, d)
        } else if d < key.len()-1 {
            Node::get(x.unwrap().mid.as_ref(), key, d+1)
        } else {
            x
        }
    }
}

/// Symbol table with string keys, implemented using a ternary search trie (TST).
pub struct TernarySearchTrie<V> {
    root: Option<Box<Node<V>>>,
    n: usize
}

impl<V> TernarySearchTrie<V>  {
    pub fn new() -> TernarySearchTrie<V> {
        TernarySearchTrie { root: None, n: 0 }
    }

    pub fn put(&mut self, key: &str, val: V) {
        let (root, replaced) = Node::put(self.root.take(), key, Some(val), 0);
        self.root = root;
        // replace old val? or insert new?
        if replaced.is_none() {
            self.n += 1;
        }
    }

    pub fn get(&self, key: &str) -> Option<&V> {
        assert!(key.len() > 0, "key must have length >= 1");
        Node::get(self.root.as_ref(), key, 0).map_or(None, |n| n.val.as_ref())
    }

    pub fn delete(&mut self, key: &str) {
        let (root, replaced) = Node::put(self.root.take(), key, None, 0);
        self.root = root;
        // deleted?
        if replaced.is_some() {
            self.n -= 1;
        }
    }

    pub fn size(&self) -> usize {
        self.n
    }

    pub fn is_empty(&self) -> bool {
        self.n == 0
    }

    pub fn contains(&self, key: &str) -> bool {
        self.get(key).is_some()
    }
}

#[test]
fn test_tst() {
    let mut t = TernarySearchTrie::new();
    assert_eq!(t.size(), 0);
    t.put("name", "Andelf");
    assert_eq!(t.size(), 1);
    t.put("name", "Fledna");
    assert_eq!(t.size(), 1);
    t.put("language", "Rust");
    assert_eq!(t.size(), 2);

    assert_eq!(t.get("name"), Some(&"Fledna"));
    assert_eq!(t.get("whatever"), None);

    t.delete("name");
    assert_eq!(t.size(), 1);
    assert_eq!(t.get("name"), None);
}
<commit_msg>tries/tst: implement keys()/keys_with_prefix() for TernarySearchTrie<commit_after>use adivon::queue::Queue;

pub struct Node<V> {
    c: char,
    left:  Option<Box<Node<V>>>,
    mid:   Option<Box<Node<V>>>,
    right: Option<Box<Node<V>>>,
    val: Option<V>,
}

impl<V> Node<V> {
    fn new(c: char) -> Node<V> {
        Node {
            c: c,
            left: None,
            mid: None,
            right: None,
            val: None
        }
    }

    fn put(mut x: Option<Box<Node<V>>>, key: &str, val: Option<V>, d: usize) -> (Option<Box<Node<V>>>, Option<V>) {
        let replaced;
        let c = key.char_at(d);
        if x.is_none() {
            if val.is_none() {  // no need to call put further
                return (x, None);
            }
            x = Some(Box::new(Node::new(c)));
        }
        let xc = x.as_ref().unwrap().c;
        if c < xc {
            let (left, repl) = Node::put(x.as_mut().unwrap().left.take(), key, val, d);
            x.as_mut().map(|n| n.left = left);
            replaced = repl;
        } else if c > xc {
            let (right, repl) = Node::put(x.as_mut().unwrap().right.take(), key, val, d);
            x.as_mut().map(|n| n.right = right);
            replaced = repl;
        } else if d < key.len()-1 {
            let (mid, repl) = Node::put(x.as_mut().unwrap().mid.take(), key, val, d+1);
            x.as_mut().map(|n| n.mid = mid);
            replaced = repl;
        } else {
            replaced = x.as_mut().unwrap().val.take();
            x.as_mut().map(|n| n.val = val);
        }
        (x, replaced)
    }

    fn get<'a>(x: Option<&'a Box<Node<V>>>, key: &str, d: usize) -> Option<&'a Box<Node<V>>> {
        if x.is_none() {
            return None;
        }
        let c = key.char_at(d);
        let xc = x.unwrap().c;
        if c < xc {
            Node::get(x.unwrap().left.as_ref(), key, d)
        } else if c > xc {
            Node::get(x.unwrap().right.as_ref(), key, d)
        } else if d < key.len()-1 {
            Node::get(x.unwrap().mid.as_ref(), key, d+1)
        } else {
            x
        }
    }

    fn collect(x: Option<&Box<Node<V>>>, mut prefix: String, queue: &mut Queue<String>) {
        if x.is_none() {
            return;
        }
        Node::collect(x.unwrap().left.as_ref(), prefix.clone(), queue);
        let xc = x.unwrap().c;
        prefix.push(xc);
        if x.unwrap().val.is_some() {
            queue.enqueue(prefix.clone());
        }
        Node::collect(x.unwrap().mid.as_ref(), prefix.clone(), queue);
        prefix.pop();
        Node::collect(x.unwrap().right.as_ref(), prefix, queue);
    }
}

/// Symbol table with string keys, implemented using a ternary search trie (TST).
pub struct TernarySearchTrie<V> {
    root: Option<Box<Node<V>>>,
    n: usize
}

impl<V> TernarySearchTrie<V>  {
    pub fn new() -> TernarySearchTrie<V> {
        TernarySearchTrie { root: None, n: 0 }
    }

    pub fn put(&mut self, key: &str, val: V) {
        let (root, replaced) = Node::put(self.root.take(), key, Some(val), 0);
        self.root = root;
        // replace old val? or insert new?
        if replaced.is_none() {
            self.n += 1;
        }
    }

    pub fn get(&self, key: &str) -> Option<&V> {
        assert!(key.len() > 0, "key must have length >= 1");
        Node::get(self.root.as_ref(), key, 0).map_or(None, |n| n.val.as_ref())
    }

    pub fn delete(&mut self, key: &str) {
        let (root, replaced) = Node::put(self.root.take(), key, None, 0);
        self.root = root;
        // deleted?
        if replaced.is_some() {
            self.n -= 1;
        }
    }

    pub fn size(&self) -> usize {
        self.n
    }

    pub fn is_empty(&self) -> bool {
        self.n == 0
    }

    pub fn contains(&self, key: &str) -> bool {
        self.get(key).is_some()
    }

    pub fn keys_with_prefix(&self, prefix: &str) -> Vec<String> {
        let mut queue = Queue::new();
        let x = Node::get(self.root.as_ref(), prefix, 0);
        if x.is_some() {
            if x.unwrap().val.is_some() {
                queue.enqueue(prefix.into());
            }
            Node::collect(x.unwrap().mid.as_ref(), prefix.into(), &mut queue);
        }
        queue.into_iter().collect()
    }

    pub fn keys(&self) -> Vec<String> {
        let mut queue = Queue::new();
        Node::collect(self.root.as_ref(), "".into(), &mut queue);
        queue.into_iter().collect()
    }
}

#[test]
fn test_tst() {
    let mut t = TernarySearchTrie::new();
    assert_eq!(t.size(), 0);
    t.put("name", "Andelf");
    assert_eq!(t.size(), 1);
    assert!(t.contains("name"));
    t.put("name", "Fledna");
    assert_eq!(t.size(), 1);
    t.put("language", "Rust");
    assert_eq!(t.size(), 2);

    assert_eq!(t.get("name"), Some(&"Fledna"));
    assert_eq!(t.get("whatever"), None);

    t.delete("name");
    assert_eq!(t.size(), 1);
    assert_eq!(t.get("name"), None);

    t.put("name", "Lednaf");
    assert!(t.keys().contains(&"name".into()));
    assert!(t.keys().contains(&"language".into()));
    assert!(t.keys_with_prefix("lang").contains(&"language".into()));
}
<|endoftext|>


<commit_before>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

/*!

This module provides an implementation of mpsc channel based on
crossbeam_channel. Comparing to the crossbeam_channel, this implementation
supports closed detection and try operations.

*/

#![cfg_attr(feature = "cargo-clippy", allow(cast_ptr_alignment))]

use crossbeam_channel;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{RecvError, RecvTimeoutError, SendError, TryRecvError, TrySendError};
use std::sync::Arc;
use std::time::Duration;

struct State {
    sender_cnt: AtomicIsize,
    receiver_cnt: AtomicIsize,
}

impl State {
    fn new() -> State {
        State {
            sender_cnt: AtomicIsize::new(1),
            receiver_cnt: AtomicIsize::new(1),
        }
    }

    #[inline]
    fn is_receiver_closed(&self) -> bool {
        self.receiver_cnt.load(Ordering::Acquire) == 0
    }

    #[inline]
    fn is_sender_closed(&self) -> bool {
        self.sender_cnt.load(Ordering::Acquire) == 0
    }
}

pub struct Sender<T> {
    sender: crossbeam_channel::Sender<T>,
    state: Arc<State>,
}

impl<T> Clone for Sender<T> {
    #[inline]
    fn clone(&self) -> Sender<T> {
        self.state.sender_cnt.fetch_add(1, Ordering::AcqRel);
        Sender {
            sender: self.sender.clone(),
            state: self.state.clone(),
        }
    }
}

impl<T> Drop for Sender<T> {
    #[inline]
    fn drop(&mut self) {
        self.state.sender_cnt.fetch_add(-1, Ordering::AcqRel);
    }
}

pub struct Receiver<T> {
    receiver: crossbeam_channel::Receiver<T>,
    state: Arc<State>,
}

impl<T> Sender<T> {
    #[inline]
    pub fn send(&self, t: T) -> Result<(), SendError<T>> {
        if !self.state.is_receiver_closed() {
            self.sender.send(t);
            return Ok(());
        }
        Err(SendError(t))
    }

    #[inline]
    pub fn try_send(&self, t: T) -> Result<(), TrySendError<T>> {
        if !self.state.is_receiver_closed() {
            crossbeam_channel::select! {
                send(self.sender, t) => Ok(()),
                default => Err(TrySendError::Full(t)),
            }
        } else {
            Err(TrySendError::Disconnected(t))
        }
    }
}

impl<T> Receiver<T> {
    #[inline]
    pub fn recv(&self) -> Result<T, RecvError> {
        match self.receiver.recv() {
            Some(t) => Ok(t),
            None => Err(RecvError),
        }
    }

    #[inline]
    pub fn try_recv(&self) -> Result<T, TryRecvError> {
        match self.receiver.try_recv() {
            Some(t) => Ok(t),
            None => {
                if !self.state.is_sender_closed() {
                    return Err(TryRecvError::Empty);
                }
                self.receiver.try_recv().ok_or(TryRecvError::Disconnected)
            }
        }
    }

    #[inline]
    pub fn recv_timeout(&self, timeout: Duration) -> Result<T, RecvTimeoutError> {
        crossbeam_channel::select! {
            recv(self.receiver, task) => match task {
                Some(t) => Ok(t),
                None => Err(RecvTimeoutError::Disconnected),
            }
            recv(crossbeam_channel::after(timeout)) => Err(RecvTimeoutError::Timeout),
        }
    }
}

impl<T> Drop for Receiver<T> {
    #[inline]
    fn drop(&mut self) {
        self.state.receiver_cnt.fetch_add(-1, Ordering::AcqRel);
    }
}

#[inline]
pub fn unbounded<T>() -> (Sender<T>, Receiver<T>) {
    let state = Arc::new(State::new());
    let (sender, receiver) = crossbeam_channel::unbounded();
    (
        Sender {
            sender,
            state: state.clone(),
        },
        Receiver { receiver, state },
    )
}

#[inline]
pub fn bounded<T>(cap: usize) -> (Sender<T>, Receiver<T>) {
    let state = Arc::new(State::new());
    let (sender, receiver) = crossbeam_channel::bounded(cap);
    (
        Sender {
            sender,
            state: state.clone(),
        },
        Receiver { receiver, state },
    )
}

const CHECK_INTERVAL: usize = 8;

/// A sender of channel that limits the maximun pending messages count loosely.
pub struct LooseBoundedSender<T> {
    sender: Sender<T>,
    tried_cnt: usize,
    limit: usize,
}

impl<T> LooseBoundedSender<T> {
    /// Send a message regardless its capacity limit.
    #[inline]
    pub fn force_send(&mut self, t: T) -> Result<(), SendError<T>> {
        self.tried_cnt += 1;
        self.sender.send(t)
    }

    #[inline]
    pub fn try_send(&mut self, t: T) -> Result<(), TrySendError<T>> {
        if self.tried_cnt < CHECK_INTERVAL {
            self.force_send(t)
                .map_err(|SendError(t)| TrySendError::Disconnected(t))
        } else if self.sender.sender.len() < self.limit {
            self.tried_cnt = 0;
            self.force_send(t)
                .map_err(|SendError(t)| TrySendError::Disconnected(t))
        } else {
            Err(TrySendError::Full(t))
        }
    }
}

impl<T> Clone for LooseBoundedSender<T> {
    #[inline]
    fn clone(&self) -> LooseBoundedSender<T> {
        LooseBoundedSender {
            sender: self.sender.clone(),
            tried_cnt: self.tried_cnt,
            limit: self.limit,
        }
    }
}

pub fn loose_bounded<T>(cap: usize) -> (LooseBoundedSender<T>, Receiver<T>) {
    let (sender, receiver) = unbounded();
    (
        LooseBoundedSender {
            sender,
            tried_cnt: 0,
            limit: cap,
        },
        receiver,
    )
}

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

    #[test]
    fn test_bounded() {
        let (tx, rx) = super::bounded::<u64>(10);
        tx.try_send(1).unwrap();
        for i in 2..11 {
            tx.clone().send(i).unwrap();
        }
        assert_eq!(tx.try_send(11), Err(TrySendError::Full(11)));

        assert_eq!(rx.try_recv(), Ok(1));
        for i in 2..11 {
            assert_eq!(rx.recv(), Ok(i));
        }
        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        let timer = Instant::now();
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Timeout)
        );
        assert!((i64::from(timer.elapsed().subsec_nanos()) - 100_000_000).abs() < 1_000_000);

        drop(rx);
        assert_eq!(tx.send(2), Err(SendError(2)));
        assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));

        let (tx, rx) = super::bounded::<u64>(10);
        tx.send(2).unwrap();
        tx.send(3).unwrap();
        drop(tx);
        assert_eq!(rx.try_recv(), Ok(2));
        assert_eq!(rx.recv(), Ok(3));
        assert_eq!(rx.recv(), Err(RecvError));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Disconnected)
        );

        let (tx1, rx1) = super::bounded::<u64>(10);
        let (tx2, rx2) = super::bounded::<u64>(0);
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            tx1.send(10).unwrap();
            thread::sleep(Duration::from_millis(100));
            assert_eq!(rx2.recv(), Ok(2));
        });
        let timer = Instant::now();
        assert_eq!(rx1.recv(), Ok(10));
        assert!((i64::from(timer.elapsed().subsec_nanos()) - 100_000_000).abs() < 1_000_000);
        let timer = Instant::now();
        tx2.send(2).unwrap();
        assert!(timer.elapsed() > Duration::from_millis(50));
    }

    #[test]
    fn test_unbounded() {
        let (tx, rx) = super::unbounded::<u64>();
        tx.try_send(1).unwrap();
        tx.send(2).unwrap();

        assert_eq!(rx.try_recv(), Ok(1));
        assert_eq!(rx.recv(), Ok(2));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        let timer = Instant::now();
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Timeout)
        );
        assert!((i64::from(timer.elapsed().subsec_nanos()) - 100_000_000).abs() < 1_000_000);

        drop(rx);
        assert_eq!(tx.send(2), Err(SendError(2)));
        assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));

        let (tx, rx) = super::unbounded::<u64>();
        drop(tx);
        assert_eq!(rx.recv(), Err(RecvError));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Disconnected)
        );

        let (tx, rx) = super::unbounded::<u64>();
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            tx.send(10).unwrap();
        });
        let timer = Instant::now();
        assert_eq!(rx.recv(), Ok(10));
        assert!((i64::from(timer.elapsed().subsec_nanos()) - 100_000_000).abs() < 1_000_000);
    }

    #[test]
    fn test_loose() {
        let (mut tx, rx) = super::loose_bounded(10);
        tx.try_send(1).unwrap();
        for i in 2..11 {
            tx.clone().try_send(i).unwrap();
        }
        for i in 1..super::CHECK_INTERVAL {
            tx.force_send(i).unwrap();
        }
        assert_eq!(tx.try_send(4), Err(TrySendError::Full(4)));
        tx.force_send(5).unwrap();
        assert_eq!(tx.try_send(6), Err(TrySendError::Full(6)));

        assert_eq!(rx.try_recv(), Ok(1));
        for i in 2..11 {
            assert_eq!(rx.recv(), Ok(i));
        }
        for i in 1..super::CHECK_INTERVAL {
            assert_eq!(rx.try_recv(), Ok(i));
        }
        assert_eq!(rx.try_recv(), Ok(5));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        let timer = Instant::now();
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Timeout)
        );
        assert!((i64::from(timer.elapsed().subsec_nanos()) - 100_000_000).abs() < 1_000_000);

        tx.force_send(1).unwrap();
        drop(rx);
        assert_eq!(tx.force_send(2), Err(SendError(2)));
        assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));
        for _ in 0..super::CHECK_INTERVAL {
            assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));
        }

        let (mut tx, rx) = super::loose_bounded(10);
        tx.try_send(2).unwrap();
        drop(tx);
        assert_eq!(rx.recv(), Ok(2));
        assert_eq!(rx.recv(), Err(RecvError));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Disconnected)
        );

        let (mut tx, rx) = super::loose_bounded(10);
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            tx.try_send(10).unwrap();
        });
        let timer = Instant::now();
        assert_eq!(rx.recv(), Ok(10));
        assert!((i64::from(timer.elapsed().subsec_nanos()) - 100_000_000).abs() < 1_000_000);
    }
}
<commit_msg>relax timeout check (#3312)<commit_after>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

/*!

This module provides an implementation of mpsc channel based on
crossbeam_channel. Comparing to the crossbeam_channel, this implementation
supports closed detection and try operations.

*/

#![cfg_attr(feature = "cargo-clippy", allow(cast_ptr_alignment))]

use crossbeam_channel;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{RecvError, RecvTimeoutError, SendError, TryRecvError, TrySendError};
use std::sync::Arc;
use std::time::Duration;

struct State {
    sender_cnt: AtomicIsize,
    receiver_cnt: AtomicIsize,
}

impl State {
    fn new() -> State {
        State {
            sender_cnt: AtomicIsize::new(1),
            receiver_cnt: AtomicIsize::new(1),
        }
    }

    #[inline]
    fn is_receiver_closed(&self) -> bool {
        self.receiver_cnt.load(Ordering::Acquire) == 0
    }

    #[inline]
    fn is_sender_closed(&self) -> bool {
        self.sender_cnt.load(Ordering::Acquire) == 0
    }
}

pub struct Sender<T> {
    sender: crossbeam_channel::Sender<T>,
    state: Arc<State>,
}

impl<T> Clone for Sender<T> {
    #[inline]
    fn clone(&self) -> Sender<T> {
        self.state.sender_cnt.fetch_add(1, Ordering::AcqRel);
        Sender {
            sender: self.sender.clone(),
            state: self.state.clone(),
        }
    }
}

impl<T> Drop for Sender<T> {
    #[inline]
    fn drop(&mut self) {
        self.state.sender_cnt.fetch_add(-1, Ordering::AcqRel);
    }
}

pub struct Receiver<T> {
    receiver: crossbeam_channel::Receiver<T>,
    state: Arc<State>,
}

impl<T> Sender<T> {
    #[inline]
    pub fn send(&self, t: T) -> Result<(), SendError<T>> {
        if !self.state.is_receiver_closed() {
            self.sender.send(t);
            return Ok(());
        }
        Err(SendError(t))
    }

    #[inline]
    pub fn try_send(&self, t: T) -> Result<(), TrySendError<T>> {
        if !self.state.is_receiver_closed() {
            crossbeam_channel::select! {
                send(self.sender, t) => Ok(()),
                default => Err(TrySendError::Full(t)),
            }
        } else {
            Err(TrySendError::Disconnected(t))
        }
    }
}

impl<T> Receiver<T> {
    #[inline]
    pub fn recv(&self) -> Result<T, RecvError> {
        match self.receiver.recv() {
            Some(t) => Ok(t),
            None => Err(RecvError),
        }
    }

    #[inline]
    pub fn try_recv(&self) -> Result<T, TryRecvError> {
        match self.receiver.try_recv() {
            Some(t) => Ok(t),
            None => {
                if !self.state.is_sender_closed() {
                    return Err(TryRecvError::Empty);
                }
                self.receiver.try_recv().ok_or(TryRecvError::Disconnected)
            }
        }
    }

    #[inline]
    pub fn recv_timeout(&self, timeout: Duration) -> Result<T, RecvTimeoutError> {
        crossbeam_channel::select! {
            recv(self.receiver, task) => match task {
                Some(t) => Ok(t),
                None => Err(RecvTimeoutError::Disconnected),
            }
            recv(crossbeam_channel::after(timeout)) => Err(RecvTimeoutError::Timeout),
        }
    }
}

impl<T> Drop for Receiver<T> {
    #[inline]
    fn drop(&mut self) {
        self.state.receiver_cnt.fetch_add(-1, Ordering::AcqRel);
    }
}

#[inline]
pub fn unbounded<T>() -> (Sender<T>, Receiver<T>) {
    let state = Arc::new(State::new());
    let (sender, receiver) = crossbeam_channel::unbounded();
    (
        Sender {
            sender,
            state: state.clone(),
        },
        Receiver { receiver, state },
    )
}

#[inline]
pub fn bounded<T>(cap: usize) -> (Sender<T>, Receiver<T>) {
    let state = Arc::new(State::new());
    let (sender, receiver) = crossbeam_channel::bounded(cap);
    (
        Sender {
            sender,
            state: state.clone(),
        },
        Receiver { receiver, state },
    )
}

const CHECK_INTERVAL: usize = 8;

/// A sender of channel that limits the maximun pending messages count loosely.
pub struct LooseBoundedSender<T> {
    sender: Sender<T>,
    tried_cnt: usize,
    limit: usize,
}

impl<T> LooseBoundedSender<T> {
    /// Send a message regardless its capacity limit.
    #[inline]
    pub fn force_send(&mut self, t: T) -> Result<(), SendError<T>> {
        self.tried_cnt += 1;
        self.sender.send(t)
    }

    #[inline]
    pub fn try_send(&mut self, t: T) -> Result<(), TrySendError<T>> {
        if self.tried_cnt < CHECK_INTERVAL {
            self.force_send(t)
                .map_err(|SendError(t)| TrySendError::Disconnected(t))
        } else if self.sender.sender.len() < self.limit {
            self.tried_cnt = 0;
            self.force_send(t)
                .map_err(|SendError(t)| TrySendError::Disconnected(t))
        } else {
            Err(TrySendError::Full(t))
        }
    }
}

impl<T> Clone for LooseBoundedSender<T> {
    #[inline]
    fn clone(&self) -> LooseBoundedSender<T> {
        LooseBoundedSender {
            sender: self.sender.clone(),
            tried_cnt: self.tried_cnt,
            limit: self.limit,
        }
    }
}

pub fn loose_bounded<T>(cap: usize) -> (LooseBoundedSender<T>, Receiver<T>) {
    let (sender, receiver) = unbounded();
    (
        LooseBoundedSender {
            sender,
            tried_cnt: 0,
            limit: cap,
        },
        receiver,
    )
}

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

    #[test]
    fn test_bounded() {
        let (tx, rx) = super::bounded::<u64>(10);
        tx.try_send(1).unwrap();
        for i in 2..11 {
            tx.clone().send(i).unwrap();
        }
        assert_eq!(tx.try_send(11), Err(TrySendError::Full(11)));

        assert_eq!(rx.try_recv(), Ok(1));
        for i in 2..11 {
            assert_eq!(rx.recv(), Ok(i));
        }
        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        let timer = Instant::now();
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Timeout)
        );
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(100), "{:?}", elapsed);

        drop(rx);
        assert_eq!(tx.send(2), Err(SendError(2)));
        assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));

        let (tx, rx) = super::bounded::<u64>(10);
        tx.send(2).unwrap();
        tx.send(3).unwrap();
        drop(tx);
        assert_eq!(rx.try_recv(), Ok(2));
        assert_eq!(rx.recv(), Ok(3));
        assert_eq!(rx.recv(), Err(RecvError));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Disconnected)
        );

        let (tx1, rx1) = super::bounded::<u64>(10);
        let (tx2, rx2) = super::bounded::<u64>(0);
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            tx1.send(10).unwrap();
            thread::sleep(Duration::from_millis(100));
            assert_eq!(rx2.recv(), Ok(2));
        });
        let timer = Instant::now();
        assert_eq!(rx1.recv(), Ok(10));
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(100), "{:?}", elapsed);
        let timer = Instant::now();
        tx2.send(2).unwrap();
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(50), "{:?}", elapsed);
    }

    #[test]
    fn test_unbounded() {
        let (tx, rx) = super::unbounded::<u64>();
        tx.try_send(1).unwrap();
        tx.send(2).unwrap();

        assert_eq!(rx.try_recv(), Ok(1));
        assert_eq!(rx.recv(), Ok(2));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        let timer = Instant::now();
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Timeout)
        );
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(100), "{:?}", elapsed);

        drop(rx);
        assert_eq!(tx.send(2), Err(SendError(2)));
        assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));

        let (tx, rx) = super::unbounded::<u64>();
        drop(tx);
        assert_eq!(rx.recv(), Err(RecvError));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Disconnected)
        );

        let (tx, rx) = super::unbounded::<u64>();
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            tx.send(10).unwrap();
        });
        let timer = Instant::now();
        assert_eq!(rx.recv(), Ok(10));
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(100), "{:?}", elapsed);
    }

    #[test]
    fn test_loose() {
        let (mut tx, rx) = super::loose_bounded(10);
        tx.try_send(1).unwrap();
        for i in 2..11 {
            tx.clone().try_send(i).unwrap();
        }
        for i in 1..super::CHECK_INTERVAL {
            tx.force_send(i).unwrap();
        }
        assert_eq!(tx.try_send(4), Err(TrySendError::Full(4)));
        tx.force_send(5).unwrap();
        assert_eq!(tx.try_send(6), Err(TrySendError::Full(6)));

        assert_eq!(rx.try_recv(), Ok(1));
        for i in 2..11 {
            assert_eq!(rx.recv(), Ok(i));
        }
        for i in 1..super::CHECK_INTERVAL {
            assert_eq!(rx.try_recv(), Ok(i));
        }
        assert_eq!(rx.try_recv(), Ok(5));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        let timer = Instant::now();
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Timeout)
        );
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(100), "{:?}", elapsed);

        tx.force_send(1).unwrap();
        drop(rx);
        assert_eq!(tx.force_send(2), Err(SendError(2)));
        assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));
        for _ in 0..super::CHECK_INTERVAL {
            assert_eq!(tx.try_send(2), Err(TrySendError::Disconnected(2)));
        }

        let (mut tx, rx) = super::loose_bounded(10);
        tx.try_send(2).unwrap();
        drop(tx);
        assert_eq!(rx.recv(), Ok(2));
        assert_eq!(rx.recv(), Err(RecvError));
        assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
        assert_eq!(
            rx.recv_timeout(Duration::from_millis(100)),
            Err(RecvTimeoutError::Disconnected)
        );

        let (mut tx, rx) = super::loose_bounded(10);
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            tx.try_send(10).unwrap();
        });
        let timer = Instant::now();
        assert_eq!(rx.recv(), Ok(10));
        let elapsed = timer.elapsed();
        assert!(elapsed >= Duration::from_millis(100), "{:?}", elapsed);
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::io::{Read, Write};

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

use super::{ReadError, WriteError};

/// A Fleetspeak client communication message.
///
/// This structure represents incoming or outgoing message objects delivered by
/// Fleetspeak. This is a simplified version of the underlying Protocol Buffers
/// message that exposes too much irrelevant fields and makes the protocol easy
/// to misuse.
pub struct Message<M: protobuf::Message> {
    /// A name of the server-side service that sent or should receive the data.
    pub service: String,
    /// An optional message type that can be used by the server-side service.
    pub kind: Option<String>,
    /// A message to sent to the specified service.
    pub data: M,
}

/// Executes the handshake procedure.
///
/// The handshake procedure consists of writing and reading magic numbers from
/// the connection buffers. This validates that the communication between the
/// Fleetspeak client and the service daemon is working as expected.
///
/// All Fleetspeak connection buffers are required to perform the handshake
/// before they became usable for sending and receiving messages.
pub fn handshake<R, W>(input: &mut R, output: &mut W) -> std::io::Result<()>
where
    R: Read,
    W: Write,
{
    write_magic(output)?;
    output.flush()?;
    read_magic(input)?;

    Ok(())
}

/// Sends a heartbeat information through this connection.
///
/// All client services should heartbeat from time to time. Otherwise, from
/// the Fleetspeak perspective, the service is unresponsive and should be
/// restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(String::from("Heartbeat"));
    msg.mut_destination().set_service_name(String::from("system"));

    write_raw(output, msg)
}

/// Sends the startup information through this connection.
///
/// All clients are required to send this information on startup. If the
/// client does not receive this information quickly enough, the service
/// will be killed.
///
/// The `version` string should contain a self-reported version of the
/// service. This data is used primarily for statistics.
pub fn startup<W>(output: &mut W, version: &str) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    let mut data = fleetspeak_proto::channel::StartupData::new();
    data.set_pid(i64::from(std::process::id()));
    data.set_version(String::from(version));

    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(String::from("StartupData"));
    msg.mut_destination().set_service_name(String::from("system"));
    msg.mut_data().set_type_url(type_url(&data));
    msg.mut_data().set_value(data.write_to_bytes()?);

    write_raw(output, msg)
}

/// Sends the message to the Fleetspeak server through the output buffer.
///
/// The message is sent to the server-side `service` and tagged with the
/// `kind` type. Note that this message type is rather irrelevant for
/// Fleetspeak and it is up to the service what to do with this information.
pub fn send<W, M>(output: &mut W, message: Message<M>) -> Result<(), WriteError>
where
    W: Write,
    M: protobuf::Message,
{
    let mut proto = fleetspeak_proto::common::Message::new();
    proto.set_message_type(message.kind.unwrap_or_else(String::new));
    proto.mut_destination().set_service_name(message.service);
    proto.mut_data().set_type_url(type_url(&message.data));
    proto.mut_data().set_value(message.data.write_to_bytes()?);

    write_raw(output, proto)
}

/// Receives the message from the Fleetspeak server through the input buffer.
///
/// This function will block until there is a message to be read in the
/// input. Errors are reported in case of any I/O failure or if the read
/// message was malformed (e.g. it cannot be parsed to the expected type).
pub fn receive<R, M>(input: &mut R) -> Result<Message<M>, ReadError>
where
    R: Read,
    M: protobuf::Message,
{
    let mut msg = read_raw(input)?;

    // While missing source address might not be consider a critical error
    // in most cases, for our own sanity we just disregard such messages.
    // Allowing such behaviour might indicate a more severe problem with
    // Fleetspeak and ignoring it simply masks the issue. This might be
    // reconsidered in the future.
    let service = if msg.has_source() {
        msg.take_source().take_service_name()
    } else {
        return Err(ReadError::malformed("missing source address"))
    };

    // It is not clear what is the best approach here. If there is no data,
    // should we error-out or return a default value? For the time being we
    // stick to the default approach, but if this proves to be not working
    // well in practice, it might be reconsidered.
    let data = if msg.has_data() {
        msg.take_data()
    } else {
        log::warn!(target: "fleetspeak", "empty message from '{}'", service);
        Default::default()
    };

    Ok(Message {
        service: service,
        kind: Some(msg.message_type),
        data: protobuf::Message::parse_from_bytes(&data.value[..])?,
    })
}

/// Writes a raw Fleetspeak message to the output buffer.
///
/// This method does not perform any validation of the message being emitted
/// and assumes that all the required fields are present.
///
/// Note that this call will fail only if the message cannot be written to
/// the output or cannot be properly encoded but will succeed even if the
/// message is not what the server expects.
fn write_raw<W>(output: &mut W, msg: fleetspeak_proto::common::Message) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    output.write_u32::<LittleEndian>(msg.compute_size())?;
    msg.write_to_writer(output)?;
    write_magic(output)?;
    output.flush()?;

    Ok(())
}

/// Reads a raw Fleetspeeak message from the input buffer.
///
/// This function will block until there is a message to be read from the
/// input. It will fail in case of any I/O error or if the message cannot
/// be parsed as a Fleetspeak message.
fn read_raw<R>(input: &mut R) -> Result<fleetspeak_proto::common::Message, ReadError>
where
    R: Read,
{
    let len = input.read_u32::<LittleEndian>()? as usize;
    let mut buf = vec!(0; len);

    input.read_exact(&mut buf[..])?;
    read_magic(input)?;

    Ok(protobuf::Message::parse_from_bytes(&buf[..])?)
}

/// Writes the Fleetspeak magic to the output buffer.
fn write_magic<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    output.write_u32::<LittleEndian>(MAGIC)?;

    Ok(())
}

/// Reads the Fleetspeak magic from the input buffer.
fn read_magic<R>(input: &mut R) -> Result<(), ReadError>
where
    R: Read,
{
    let magic = input.read_u32::<LittleEndian>()?;
    if magic != MAGIC {
        return Err(ReadError::Magic(magic));
    }

    Ok(())
}

const MAGIC: u32 = 0xf1ee1001;

// Computes a type URL of the given Protocol Buffers message.
//
// This function should probably be part of the `protobuf` package but for some
// reason it is not and we have to implement it ourselves.
fn type_url<M: protobuf::Message>(message: &M) -> String {
    format!("{}/{}", TYPE_URL_PREFIX, message.descriptor().full_name())
}

const TYPE_URL_PREFIX: &'static str = "type.googleapis.com";

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

    #[test]
    fn handshake_good_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(MAGIC).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_ok());

        let mut cur = Cursor::new(&mut buf_out[..]);
        assert_eq!(cur.read_u32::<LittleEndian>().unwrap(), MAGIC);
    }

    #[test]
    fn handshake_bad_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(0xf1ee1337).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_err());
    }

    #[test]
    fn type_url_startup_data() {
        assert_eq! {
            type_url(&fleetspeak_proto::channel::StartupData::new()),
            "type.googleapis.com/fleetspeak.channel.StartupData"
        };
    }
}
<commit_msg>Turn comment into a doc.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::io::{Read, Write};

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

use super::{ReadError, WriteError};

/// A Fleetspeak client communication message.
///
/// This structure represents incoming or outgoing message objects delivered by
/// Fleetspeak. This is a simplified version of the underlying Protocol Buffers
/// message that exposes too much irrelevant fields and makes the protocol easy
/// to misuse.
pub struct Message<M: protobuf::Message> {
    /// A name of the server-side service that sent or should receive the data.
    pub service: String,
    /// An optional message type that can be used by the server-side service.
    pub kind: Option<String>,
    /// A message to sent to the specified service.
    pub data: M,
}

/// Executes the handshake procedure.
///
/// The handshake procedure consists of writing and reading magic numbers from
/// the connection buffers. This validates that the communication between the
/// Fleetspeak client and the service daemon is working as expected.
///
/// All Fleetspeak connection buffers are required to perform the handshake
/// before they became usable for sending and receiving messages.
pub fn handshake<R, W>(input: &mut R, output: &mut W) -> std::io::Result<()>
where
    R: Read,
    W: Write,
{
    write_magic(output)?;
    output.flush()?;
    read_magic(input)?;

    Ok(())
}

/// Sends a heartbeat information through this connection.
///
/// All client services should heartbeat from time to time. Otherwise, from
/// the Fleetspeak perspective, the service is unresponsive and should be
/// restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(String::from("Heartbeat"));
    msg.mut_destination().set_service_name(String::from("system"));

    write_raw(output, msg)
}

/// Sends the startup information through this connection.
///
/// All clients are required to send this information on startup. If the
/// client does not receive this information quickly enough, the service
/// will be killed.
///
/// The `version` string should contain a self-reported version of the
/// service. This data is used primarily for statistics.
pub fn startup<W>(output: &mut W, version: &str) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    let mut data = fleetspeak_proto::channel::StartupData::new();
    data.set_pid(i64::from(std::process::id()));
    data.set_version(String::from(version));

    let mut msg = fleetspeak_proto::common::Message::new();
    msg.set_message_type(String::from("StartupData"));
    msg.mut_destination().set_service_name(String::from("system"));
    msg.mut_data().set_type_url(type_url(&data));
    msg.mut_data().set_value(data.write_to_bytes()?);

    write_raw(output, msg)
}

/// Sends the message to the Fleetspeak server through the output buffer.
///
/// The message is sent to the server-side `service` and tagged with the
/// `kind` type. Note that this message type is rather irrelevant for
/// Fleetspeak and it is up to the service what to do with this information.
pub fn send<W, M>(output: &mut W, message: Message<M>) -> Result<(), WriteError>
where
    W: Write,
    M: protobuf::Message,
{
    let mut proto = fleetspeak_proto::common::Message::new();
    proto.set_message_type(message.kind.unwrap_or_else(String::new));
    proto.mut_destination().set_service_name(message.service);
    proto.mut_data().set_type_url(type_url(&message.data));
    proto.mut_data().set_value(message.data.write_to_bytes()?);

    write_raw(output, proto)
}

/// Receives the message from the Fleetspeak server through the input buffer.
///
/// This function will block until there is a message to be read in the
/// input. Errors are reported in case of any I/O failure or if the read
/// message was malformed (e.g. it cannot be parsed to the expected type).
pub fn receive<R, M>(input: &mut R) -> Result<Message<M>, ReadError>
where
    R: Read,
    M: protobuf::Message,
{
    let mut msg = read_raw(input)?;

    // While missing source address might not be consider a critical error
    // in most cases, for our own sanity we just disregard such messages.
    // Allowing such behaviour might indicate a more severe problem with
    // Fleetspeak and ignoring it simply masks the issue. This might be
    // reconsidered in the future.
    let service = if msg.has_source() {
        msg.take_source().take_service_name()
    } else {
        return Err(ReadError::malformed("missing source address"))
    };

    // It is not clear what is the best approach here. If there is no data,
    // should we error-out or return a default value? For the time being we
    // stick to the default approach, but if this proves to be not working
    // well in practice, it might be reconsidered.
    let data = if msg.has_data() {
        msg.take_data()
    } else {
        log::warn!(target: "fleetspeak", "empty message from '{}'", service);
        Default::default()
    };

    Ok(Message {
        service: service,
        kind: Some(msg.message_type),
        data: protobuf::Message::parse_from_bytes(&data.value[..])?,
    })
}

/// Writes a raw Fleetspeak message to the output buffer.
///
/// This method does not perform any validation of the message being emitted
/// and assumes that all the required fields are present.
///
/// Note that this call will fail only if the message cannot be written to
/// the output or cannot be properly encoded but will succeed even if the
/// message is not what the server expects.
fn write_raw<W>(output: &mut W, msg: fleetspeak_proto::common::Message) -> Result<(), WriteError>
where
    W: Write,
{
    use protobuf::Message as _;

    output.write_u32::<LittleEndian>(msg.compute_size())?;
    msg.write_to_writer(output)?;
    write_magic(output)?;
    output.flush()?;

    Ok(())
}

/// Reads a raw Fleetspeeak message from the input buffer.
///
/// This function will block until there is a message to be read from the
/// input. It will fail in case of any I/O error or if the message cannot
/// be parsed as a Fleetspeak message.
fn read_raw<R>(input: &mut R) -> Result<fleetspeak_proto::common::Message, ReadError>
where
    R: Read,
{
    let len = input.read_u32::<LittleEndian>()? as usize;
    let mut buf = vec!(0; len);

    input.read_exact(&mut buf[..])?;
    read_magic(input)?;

    Ok(protobuf::Message::parse_from_bytes(&buf[..])?)
}

/// Writes the Fleetspeak magic to the output buffer.
fn write_magic<W>(output: &mut W) -> Result<(), WriteError>
where
    W: Write,
{
    output.write_u32::<LittleEndian>(MAGIC)?;

    Ok(())
}

/// Reads the Fleetspeak magic from the input buffer.
fn read_magic<R>(input: &mut R) -> Result<(), ReadError>
where
    R: Read,
{
    let magic = input.read_u32::<LittleEndian>()?;
    if magic != MAGIC {
        return Err(ReadError::Magic(magic));
    }

    Ok(())
}

const MAGIC: u32 = 0xf1ee1001;

/// Computes a type URL of the given Protocol Buffers message.
///
/// This function should probably be part of the `protobuf` package but for some
/// reason it is not and we have to implement it ourselves.
fn type_url<M: protobuf::Message>(message: &M) -> String {
    format!("{}/{}", TYPE_URL_PREFIX, message.descriptor().full_name())
}

const TYPE_URL_PREFIX: &'static str = "type.googleapis.com";

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

    #[test]
    fn handshake_good_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(MAGIC).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_ok());

        let mut cur = Cursor::new(&mut buf_out[..]);
        assert_eq!(cur.read_u32::<LittleEndian>().unwrap(), MAGIC);
    }

    #[test]
    fn handshake_bad_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(0xf1ee1337).is_ok());

        let mut cur_in = Cursor::new(&mut buf_in[..]);
        let mut cur_out = Cursor::new(&mut buf_out[..]);
        assert!(handshake(&mut cur_in, &mut cur_out).is_err());
    }

    #[test]
    fn type_url_startup_data() {
        assert_eq! {
            type_url(&fleetspeak_proto::channel::StartupData::new()),
            "type.googleapis.com/fleetspeak.channel.StartupData"
        };
    }
}
<|endoftext|>