<commit_before>/*!
`HMAC-SHA-256` `HMAC-SHA-256` is conjectured to meet the standard notion of
unforgeability.
*/
use std::libc::{c_ulonglong, c_int};
use std::vec::raw::{to_ptr, to_mut_ptr};
use randombytes::randombytes_into;
use crypto::verify::verify_32;

#[link(name = "sodium")]
#[link_args = "-lsodium"]
extern {
    fn crypto_auth_hmacsha256(a: *mut u8,
                                 m: *u8,
                                 mlen: c_ulonglong,
                                 k: *u8) -> c_int;
    fn crypto_auth_hmacsha256_verify(a: *u8,
                                        m: *u8,
                                        mlen: c_ulonglong,
                                        k: *u8) -> c_int;
}

pub static KEYBYTES: uint = 32;
pub static TAGBYTES: uint = 32;

/**
 * Authentication `Key`
 *
 * When a `Key` goes out of scope its contents
 * will be zeroed out
 */
pub struct Key([u8, ..KEYBYTES]);

impl Drop for Key {
    fn drop(&mut self) {
        for e in self.mut_iter() { *e = 0 }
    }
}

/**
  * Authentication `Tag`
  *
  * The tag implements the traits `TotalEq` and `Eq` using constant-time
  * comparison functions. See `sodiumoxide::crypto::verify::verify_32`
  */
pub struct Tag([u8, ..TAGBYTES]);
impl TotalEq for Tag {
    fn equals(&self, other: &Tag) -> bool {
        verify_32(&**self, &**other)
    }
}
impl Eq for Tag {
    fn eq(&self, other: &Tag) -> bool {
        verify_32(&**self, &**other)
    }
}

/**
 * `gen_key()` randomly generates a key for authentication
 *
 * THREAD SAFETY: `gen_key()` is thread-safe provided that you have
 * called `sodiumoxide::init()` once before using any other function
 * from sodiumoxide.
 */
pub fn gen_key() -> ~Key {
    let mut key = ~Key([0, ..KEYBYTES]);
    randombytes_into(**key);
    key
}

/**
 * `authenticate()` authenticates a message `m` using a secret key `k`.
 * The function returns an authenticator tag.
 */
#[fixed_stack_segment]
pub fn authenticate(m: &[u8], k: &Key) -> ~Tag {
    unsafe {
        let mut tag = ~Tag([0, ..TAGBYTES]);
        crypto_auth_hmacsha256(to_mut_ptr(**tag),
                               to_ptr(m),
                               m.len() as c_ulonglong,
                               to_ptr(**k));
        tag
    }
}

/**
 * `verify()` returns `true` if `tag` is a correct authenticator of message `m`
 * under a secret key `k`. Otherwise it returns false.
 */
#[fixed_stack_segment]
pub fn verify(tag: &Tag, m: &[u8], k: &Key) -> bool {
    unsafe {
        crypto_auth_hmacsha256_verify(to_ptr(**tag),
                                      to_ptr(m),
                                      m.len() as c_ulonglong,
                                      to_ptr(**k)) == 0
    }
}

#[test]
fn test_auth_verify() {
    use randombytes::randombytes;
    for i in range(0, 256) {
        let k = gen_key();
        let m = randombytes(i as uint);
        let tag = authenticate(m, k);
        assert!(verify(tag, m, k));
    }
}

#[test]
fn test_auth_verify_tamper() {
    use randombytes::randombytes;
    for i in range(0, 32) {
        let k = gen_key();
        let mut m = randombytes(i as uint);
        let mut tag = authenticate(m, k);
        for j in range(0, m.len()) {
            m[j] ^= 0x20;
            assert!(!verify(tag, m, k));
            m[j] ^= 0x20;
        }
        for j in range(0, tag.len()) {
            tag[j] ^= 0x20;
            assert!(!verify(tag, m, k));
            tag[j] ^= 0x20;
        }
    }
}
<commit_msg>added test vector for hmacsha256<commit_after>/*!
`HMAC-SHA-256` `HMAC-SHA-256` is conjectured to meet the standard notion of
unforgeability.
*/
use std::libc::{c_ulonglong, c_int};
use std::vec::raw::{to_ptr, to_mut_ptr};
use randombytes::randombytes_into;
use crypto::verify::verify_32;

#[link(name = "sodium")]
#[link_args = "-lsodium"]
extern {
    fn crypto_auth_hmacsha256(a: *mut u8,
                                 m: *u8,
                                 mlen: c_ulonglong,
                                 k: *u8) -> c_int;
    fn crypto_auth_hmacsha256_verify(a: *u8,
                                        m: *u8,
                                        mlen: c_ulonglong,
                                        k: *u8) -> c_int;
}

pub static KEYBYTES: uint = 32;
pub static TAGBYTES: uint = 32;

/**
 * Authentication `Key`
 *
 * When a `Key` goes out of scope its contents
 * will be zeroed out
 */
pub struct Key([u8, ..KEYBYTES]);

impl Drop for Key {
    fn drop(&mut self) {
        for e in self.mut_iter() { *e = 0 }
    }
}

/**
  * Authentication `Tag`
  *
  * The tag implements the traits `TotalEq` and `Eq` using constant-time
  * comparison functions. See `sodiumoxide::crypto::verify::verify_32`
  */
pub struct Tag([u8, ..TAGBYTES]);
impl TotalEq for Tag {
    fn equals(&self, other: &Tag) -> bool {
        verify_32(&**self, &**other)
    }
}
impl Eq for Tag {
    fn eq(&self, other: &Tag) -> bool {
        verify_32(&**self, &**other)
    }
}

/**
 * `gen_key()` randomly generates a key for authentication
 *
 * THREAD SAFETY: `gen_key()` is thread-safe provided that you have
 * called `sodiumoxide::init()` once before using any other function
 * from sodiumoxide.
 */
pub fn gen_key() -> ~Key {
    let mut key = ~Key([0, ..KEYBYTES]);
    randombytes_into(**key);
    key
}

/**
 * `authenticate()` authenticates a message `m` using a secret key `k`.
 * The function returns an authenticator tag.
 */
#[fixed_stack_segment]
pub fn authenticate(m: &[u8], k: &Key) -> ~Tag {
    unsafe {
        let mut tag = ~Tag([0, ..TAGBYTES]);
        crypto_auth_hmacsha256(to_mut_ptr(**tag),
                               to_ptr(m),
                               m.len() as c_ulonglong,
                               to_ptr(**k));
        tag
    }
}

/**
 * `verify()` returns `true` if `tag` is a correct authenticator of message `m`
 * under a secret key `k`. Otherwise it returns false.
 */
#[fixed_stack_segment]
pub fn verify(tag: &Tag, m: &[u8], k: &Key) -> bool {
    unsafe {
        crypto_auth_hmacsha256_verify(to_ptr(**tag),
                                      to_ptr(m),
                                      m.len() as c_ulonglong,
                                      to_ptr(**k)) == 0
    }
}

#[test]
fn test_auth_verify() {
    use randombytes::randombytes;
    for i in range(0, 256) {
        let k = gen_key();
        let m = randombytes(i as uint);
        let tag = authenticate(m, k);
        assert!(verify(tag, m, k));
    }
}

#[test]
fn test_auth_verify_tamper() {
    use randombytes::randombytes;
    for i in range(0, 32) {
        let k = gen_key();
        let mut m = randombytes(i as uint);
        let mut tag = authenticate(m, k);
        for j in range(0, m.len()) {
            m[j] ^= 0x20;
            assert!(!verify(tag, m, k));
            m[j] ^= 0x20;
        }
        for j in range(0, tag.len()) {
            tag[j] ^= 0x20;
            assert!(!verify(tag, m, k));
            tag[j] ^= 0x20;
        }
    }
}

#[test]
fn test_vector_1() {
    // corresponding to tests/auth2.c from NaCl
    /* "Test Case AUTH256-4 from RFC 4868 */
    let key = Key([0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08
                  ,0x09,0x0a,0x0b,0x0c,0x0d,0x0e,0x0f,0x10
                  ,0x11,0x12,0x13,0x14,0x15,0x16,0x17,0x18
                  ,0x19,0x1a,0x1b,0x1c,0x1d,0x1e,0x1f,0x20]);
    let c = [0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd
            ,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd
            ,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd
            ,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd
            ,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd
            ,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd,0xcd
            ,0xcd,0xcd];
    let a_expected = Tag([0x37,0x2e,0xfc,0xf9,0xb4,0x0b,0x35,0xc2
                         ,0x11,0x5b,0x13,0x46,0x90,0x3d,0x2e,0xf4
                         ,0x2f,0xce,0xd4,0x6f,0x08,0x46,0xe7,0x25
                         ,0x7b,0xb1,0x56,0xd3,0xd7,0xb3,0x0d,0x3f]);
    let a = authenticate(c, &key);
    assert!(*a == a_expected);
}
<|endoftext|>


<commit_before>//! Defines the metrics system.
//! The main design goals of this system are:
//! - Use lockless operations, preferably ones that don't require anything other than
//!   simple reads/writes being atomic.
//! - Exploit interior mutability and atomics being Sync to allow all methods (including the ones
//!   which are effectively mutable) to be callable on a global non-mut static.
//! - Rely on Serde to provide the actual serialization for logging the metrics.
//! - Since all metrics start at 0, we implement the Default trait via derive for all of them,
//!   to avoid having to initialize everything by hand.

//! Moreover, the value of a metric is currently NOT reset to 0 each time it's being logged. The
//! current approach is to store two values (current and previous) and compute the delta between
//! them each time we do a flush (i.e by serialization). There are a number of advantages
//! to this approach, including:
//! - We don't have to introduce an additional write (to reset the value) from the thread which
//!   does to actual logging, so less synchronization effort is required.
//! - We don't have to worry at all that much about losing some data if logging fails for a while
//!   (this could be a concern, I guess).
//!
//! If if turns out this approach is not really what we want, it's pretty easy to resort to
//! something else, while working behind the same interface.

use std::sync::atomic::{AtomicUsize, Ordering};

use chrono;
use serde::{Serialize, Serializer};

const SYSCALL_MAX: usize = 350;

/// Used for defining new types of metrics that can be either incremented with an unit
/// or an arbitrary amount of units.
// This trait helps with writing less code. It has to be in scope (via an use directive) in order
// for its methods to be available to call on structs that implement it.
pub trait Metric {
    fn add(&self, value: usize);
    fn inc(&self) {
        self.add(1);
    }
    fn count(&self) -> usize;
}

// A simple metric is one meant to be incremented from a single thread, so it can use simple
// loads + stores. Loads are currently Relaxed everywhere, because we don't do anything besides
// logging the retrieved value (their outcome os not used to modify some memory location in a
// potentially inconsistent manner). There's no way currently to make sure a SimpleMetric is only
// incremented by a single thread, this has to be enforced via judicious use (although, every
// non-vCPU related metric is associated with a particular thread, so it shouldn't be that easy
// to misuse SimpleMetric fields).
#[derive(Default)]
pub struct SimpleMetric(AtomicUsize);

impl Metric for SimpleMetric {
    fn add(&self, value: usize) {
        let ref count = self.0;
        count.store(count.load(Ordering::Relaxed) + value, Ordering::Relaxed);
    }

    fn count(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
}

impl Serialize for SimpleMetric {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        // There's no serializer.serialize_usize() for some reason :(
        serializer.serialize_u64(self.0.load(Ordering::Relaxed) as u64)
    }
}

// A shared metric is one which is expected to be incremented from more than one thread, so more
// synchronization is necessary. It's currently used for vCPU metrics. An alternative here would be
// to have one instance of every metric for each thread (like a per-thread SimpleMetric), and to
// aggregate them when logging. However this probably overkill unless we have a lot of vCPUs
// incrementing metrics very often. Still, it's there if we ever need it :-s
#[derive(Default)]
// We will be keeping two values for each metric for being able to reset
// counters on each metric.
// 1st member - current value being updated
// 2nd member - old value that gets the current value whenever metrics is flushed to disk
pub struct SharedMetric(AtomicUsize, AtomicUsize);

impl Metric for SharedMetric {
    // While the order specified for this operation is still Relaxed, the actual instruction will
    // be an asm "LOCK; something" and thus atomic across multiple threads, simply because of the
    // fetch_and_add (as opposed to "store(load() + 1)") implementation for atomics.
    // TODO: would a stronger ordering make a difference here?
    fn add(&self, value: usize) {
        self.0.fetch_add(value, Ordering::Relaxed);
    }

    fn count(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
}

impl Serialize for SharedMetric {
    /// Reset counters of each metrics. Here we suppose that Serialize's goal is to help with the
    /// flushing of metrics.
    /// !!! Any print of the metrics will also reset them. Use with caution !!!
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        // There's no serializer.serialize_usize() for some reason :(
        let snapshot = self.0.load(Ordering::Relaxed);
        let res = serializer.serialize_u64(snapshot as u64 - self.1.load(Ordering::Relaxed) as u64);

        if res.is_ok() {
            self.1.store(snapshot, Ordering::Relaxed);
        }
        res
    }
}

// The following structs are used to define a certain organization for the set of metrics we
// are interested in. Whenever the name of a field differs from its ideal textual representation
// in the serialized form, we can use the #[serde(rename = "name")] attribute to, well, rename it.

// Metrics related to the internal api server
#[derive(Default, Serialize)]
pub struct ApiServerMetrics {
    pub async_missed_actions_count: SharedMetric,
    pub async_outcome_fails: SharedMetric,
    pub async_vmm_send_timeout_count: SharedMetric,
    pub instance_info_fails: SharedMetric,
    pub sync_outcome_fails: SharedMetric,
    pub sync_vmm_send_timeout_count: SharedMetric,
}

// Metrics on GET Api Requests
#[derive(Default, Serialize)]
pub struct GetRequestsMetrics {
    pub action_info_count: SharedMetric,
    pub actions_count: SharedMetric,
    pub actions_fails: SharedMetric,
    pub instance_info_count: SharedMetric,
    pub machine_cfg_count: SharedMetric,
    pub machine_cfg_fails: SharedMetric,
}

// Metrics on PUT Api Requests
#[derive(Default, Serialize)]
pub struct PutRequestsMetrics {
    pub actions_count: SharedMetric,
    pub actions_fails: SharedMetric,
    pub boot_source_count: SharedMetric,
    pub boot_source_fails: SharedMetric,
    pub drive_fails: SharedMetric,
    pub drive_count: SharedMetric,
    pub logger_count: SharedMetric,
    pub logger_fails: SharedMetric,
    pub machine_cfg_count: SharedMetric,
    pub machine_cfg_fails: SharedMetric,
    pub network_count: SharedMetric,
    pub network_fails: SharedMetric,
}

// Metrics on PATCH Api Requests
#[derive(Default, Serialize)]
pub struct PatchRequestsMetrics {
    pub drive_fails: SharedMetric,
    pub drive_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct BlockDeviceMetrics {
    pub activate_fails: SharedMetric,
    pub cfg_fails: SharedMetric,
    pub event_fails: SharedMetric,
    pub execute_fails: SharedMetric,
    pub invalid_reqs_count: SharedMetric,
    pub flush_count: SharedMetric,
    pub queue_event_count: SharedMetric,
    pub rate_limiter_event_count: SharedMetric,
    pub update_count: SharedMetric,
    pub update_fails: SharedMetric,
    pub read_count: SharedMetric,
    pub write_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct I8042DeviceMetrics {
    pub error_count: SharedMetric,
    pub missed_read_count: SharedMetric,
    pub missed_write_count: SharedMetric,
    pub read_count: SharedMetric,
    pub reset_count: SharedMetric,
    pub write_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct LoggerSystemMetrics {
    pub missed_metrics_count: SharedMetric,
    pub metrics_fails: SharedMetric,
    pub missed_log_count: SharedMetric,
    pub log_fails: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct MmdsMetrics {
    pub rx_accepted: SharedMetric,
    pub rx_accepted_err: SharedMetric,
    pub rx_bad_eth: SharedMetric,
    pub tx_bytes: SharedMetric,
    pub tx_errors: SharedMetric,
    pub tx_frames: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct NetDeviceMetrics {
    pub activate_fails: SharedMetric,
    pub cfg_fails: SharedMetric,
    pub event_fails: SharedMetric,
    pub rx_queue_event_count: SharedMetric,
    pub rx_event_rate_limiter_count: SharedMetric,
    pub rx_tap_event_count: SharedMetric,
    pub rx_bytes_count: SharedMetric,
    pub rx_packets_count: SharedMetric,
    pub rx_fails: SharedMetric,
    pub tx_queue_event_count: SharedMetric,
    pub tx_rate_limiter_event_count: SharedMetric,
    pub tx_bytes_count: SharedMetric,
    pub tx_packets_count: SharedMetric,
    pub tx_fails: SharedMetric,
}

#[derive(Serialize)]
pub struct SeccompMetrics {
    pub num_faults: SharedMetric,
    pub bad_syscalls: Vec<SharedMetric>,
}

impl Default for SeccompMetrics {
    fn default() -> SeccompMetrics {
        let mut def_syscalls = vec![];
        for _syscall in 0..SYSCALL_MAX {
            def_syscalls.push(SharedMetric::default());
        }
        SeccompMetrics {
            num_faults: SharedMetric::default(),
            bad_syscalls: def_syscalls,
        }
    }
}

#[derive(Default, Serialize)]
pub struct SerialDeviceMetrics {
    pub error_count: SharedMetric,
    pub flush_count: SharedMetric,
    pub missed_read_count: SharedMetric,
    pub missed_write_count: SharedMetric,
    pub read_count: SharedMetric,
    pub write_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct VcpuMetrics {
    pub eagain: SharedMetric,
    pub eintr: SharedMetric,
    pub exit_io_in: SharedMetric,
    pub exit_io_out: SharedMetric,
    pub exit_mmio_read: SharedMetric,
    pub exit_mmio_write: SharedMetric,
    pub failures: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct VmmMetrics {
    pub device_events: SharedMetric,
    pub panic_count: SharedMetric,
}

// The sole purpose of this struct is to produce an UTC timestamp when an instance is serialized.
#[derive(Default)]
struct SerializeToUtcTimestampMs;

impl Serialize for SerializeToUtcTimestampMs {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_i64(chrono::Utc::now().timestamp_millis())
    }
}

#[derive(Default, Serialize)]
pub struct FirecrackerMetrics {
    utc_timestamp_ms: SerializeToUtcTimestampMs,
    pub api_server: ApiServerMetrics,
    pub block: BlockDeviceMetrics,
    pub get_api_requests: GetRequestsMetrics,
    pub i8042: I8042DeviceMetrics,
    pub logger: LoggerSystemMetrics,
    pub mmds: MmdsMetrics,
    pub net: NetDeviceMetrics,
    pub patch_api_requests: PatchRequestsMetrics,
    pub put_api_requests: PutRequestsMetrics,
    pub seccomp: SeccompMetrics,
    pub vcpu: VcpuMetrics,
    pub vmm: VmmMetrics,
    pub uart: SerialDeviceMetrics,
}

lazy_static! {
    /// Static instance used for handling metrics.
    ///
    pub static ref METRICS: FirecrackerMetrics = FirecrackerMetrics::default();
}

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

    use std::sync::Arc;
    use std::thread;

    #[test]
    fn test_metric() {
        let m1 = SimpleMetric::default();

        m1.inc();
        m1.inc();
        m1.add(5);
        m1.inc();

        assert_eq!(m1.count(), 8);

        let m2 = Arc::new(SharedMetric::default());

        // We're going to create a number of threads that will attempt to increase this metric
        // in parallel. If everything goes fine we still can't be sure the synchronization works,
        // but it something fails, then we definitely have a problem :-s

        const NUM_THREADS_TO_SPAWN: usize = 4;
        const NUM_INCREMENTS_PER_THREAD: usize = 100000;
        const M2_INITIAL_COUNT: usize = 123;

        m2.add(M2_INITIAL_COUNT);

        let mut v = Vec::with_capacity(NUM_THREADS_TO_SPAWN);

        for _ in 0..NUM_THREADS_TO_SPAWN {
            let r = m2.clone();
            v.push(thread::spawn(move || {
                for _ in 0..NUM_INCREMENTS_PER_THREAD {
                    r.inc();
                }
            }));
        }

        for handle in v {
            handle.join().unwrap();
        }

        assert_eq!(
            m2.count(),
            M2_INITIAL_COUNT + NUM_THREADS_TO_SPAWN * NUM_INCREMENTS_PER_THREAD
        );
    }

    #[test]
    fn test_serialize() {
        let s = serde_json::to_string(&FirecrackerMetrics::default());
        assert!(s.is_ok());
    }
}
<commit_msg>Unused metrics removed<commit_after>//! Defines the metrics system.
//! The main design goals of this system are:
//! - Use lockless operations, preferably ones that don't require anything other than
//!   simple reads/writes being atomic.
//! - Exploit interior mutability and atomics being Sync to allow all methods (including the ones
//!   which are effectively mutable) to be callable on a global non-mut static.
//! - Rely on Serde to provide the actual serialization for logging the metrics.
//! - Since all metrics start at 0, we implement the Default trait via derive for all of them,
//!   to avoid having to initialize everything by hand.

//! Moreover, the value of a metric is currently NOT reset to 0 each time it's being logged. The
//! current approach is to store two values (current and previous) and compute the delta between
//! them each time we do a flush (i.e by serialization). There are a number of advantages
//! to this approach, including:
//! - We don't have to introduce an additional write (to reset the value) from the thread which
//!   does to actual logging, so less synchronization effort is required.
//! - We don't have to worry at all that much about losing some data if logging fails for a while
//!   (this could be a concern, I guess).
//!
//! If if turns out this approach is not really what we want, it's pretty easy to resort to
//! something else, while working behind the same interface.

use std::sync::atomic::{AtomicUsize, Ordering};

use chrono;
use serde::{Serialize, Serializer};

const SYSCALL_MAX: usize = 350;

/// Used for defining new types of metrics that can be either incremented with an unit
/// or an arbitrary amount of units.
// This trait helps with writing less code. It has to be in scope (via an use directive) in order
// for its methods to be available to call on structs that implement it.
pub trait Metric {
    fn add(&self, value: usize);
    fn inc(&self) {
        self.add(1);
    }
    fn count(&self) -> usize;
}

// A simple metric is one meant to be incremented from a single thread, so it can use simple
// loads + stores. Loads are currently Relaxed everywhere, because we don't do anything besides
// logging the retrieved value (their outcome os not used to modify some memory location in a
// potentially inconsistent manner). There's no way currently to make sure a SimpleMetric is only
// incremented by a single thread, this has to be enforced via judicious use (although, every
// non-vCPU related metric is associated with a particular thread, so it shouldn't be that easy
// to misuse SimpleMetric fields).
#[derive(Default)]
pub struct SimpleMetric(AtomicUsize);

impl Metric for SimpleMetric {
    fn add(&self, value: usize) {
        let ref count = self.0;
        count.store(count.load(Ordering::Relaxed) + value, Ordering::Relaxed);
    }

    fn count(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
}

impl Serialize for SimpleMetric {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        // There's no serializer.serialize_usize() for some reason :(
        serializer.serialize_u64(self.0.load(Ordering::Relaxed) as u64)
    }
}

// A shared metric is one which is expected to be incremented from more than one thread, so more
// synchronization is necessary. It's currently used for vCPU metrics. An alternative here would be
// to have one instance of every metric for each thread (like a per-thread SimpleMetric), and to
// aggregate them when logging. However this probably overkill unless we have a lot of vCPUs
// incrementing metrics very often. Still, it's there if we ever need it :-s
#[derive(Default)]
// We will be keeping two values for each metric for being able to reset
// counters on each metric.
// 1st member - current value being updated
// 2nd member - old value that gets the current value whenever metrics is flushed to disk
pub struct SharedMetric(AtomicUsize, AtomicUsize);

impl Metric for SharedMetric {
    // While the order specified for this operation is still Relaxed, the actual instruction will
    // be an asm "LOCK; something" and thus atomic across multiple threads, simply because of the
    // fetch_and_add (as opposed to "store(load() + 1)") implementation for atomics.
    // TODO: would a stronger ordering make a difference here?
    fn add(&self, value: usize) {
        self.0.fetch_add(value, Ordering::Relaxed);
    }

    fn count(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
}

impl Serialize for SharedMetric {
    /// Reset counters of each metrics. Here we suppose that Serialize's goal is to help with the
    /// flushing of metrics.
    /// !!! Any print of the metrics will also reset them. Use with caution !!!
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        // There's no serializer.serialize_usize() for some reason :(
        let snapshot = self.0.load(Ordering::Relaxed);
        let res = serializer.serialize_u64(snapshot as u64 - self.1.load(Ordering::Relaxed) as u64);

        if res.is_ok() {
            self.1.store(snapshot, Ordering::Relaxed);
        }
        res
    }
}

// The following structs are used to define a certain organization for the set of metrics we
// are interested in. Whenever the name of a field differs from its ideal textual representation
// in the serialized form, we can use the #[serde(rename = "name")] attribute to, well, rename it.

// Metrics related to the internal api server
#[derive(Default, Serialize)]
pub struct ApiServerMetrics {
    pub instance_info_fails: SharedMetric,
    pub sync_outcome_fails: SharedMetric,
    pub sync_vmm_send_timeout_count: SharedMetric,
}

// Metrics on GET Api Requests
#[derive(Default, Serialize)]
pub struct GetRequestsMetrics {
    pub instance_info_count: SharedMetric,
    pub machine_cfg_count: SharedMetric,
    pub machine_cfg_fails: SharedMetric,
}

// Metrics on PUT Api Requests
#[derive(Default, Serialize)]
pub struct PutRequestsMetrics {
    pub actions_count: SharedMetric,
    pub actions_fails: SharedMetric,
    pub boot_source_count: SharedMetric,
    pub boot_source_fails: SharedMetric,
    pub drive_fails: SharedMetric,
    pub drive_count: SharedMetric,
    pub logger_count: SharedMetric,
    pub logger_fails: SharedMetric,
    pub machine_cfg_count: SharedMetric,
    pub machine_cfg_fails: SharedMetric,
    pub network_count: SharedMetric,
    pub network_fails: SharedMetric,
}

// Metrics on PATCH Api Requests
#[derive(Default, Serialize)]
pub struct PatchRequestsMetrics {
    pub drive_fails: SharedMetric,
    pub drive_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct BlockDeviceMetrics {
    pub activate_fails: SharedMetric,
    pub cfg_fails: SharedMetric,
    pub event_fails: SharedMetric,
    pub execute_fails: SharedMetric,
    pub invalid_reqs_count: SharedMetric,
    pub flush_count: SharedMetric,
    pub queue_event_count: SharedMetric,
    pub rate_limiter_event_count: SharedMetric,
    pub update_count: SharedMetric,
    pub update_fails: SharedMetric,
    pub read_count: SharedMetric,
    pub write_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct I8042DeviceMetrics {
    pub error_count: SharedMetric,
    pub missed_read_count: SharedMetric,
    pub missed_write_count: SharedMetric,
    pub read_count: SharedMetric,
    pub reset_count: SharedMetric,
    pub write_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct LoggerSystemMetrics {
    pub missed_metrics_count: SharedMetric,
    pub metrics_fails: SharedMetric,
    pub missed_log_count: SharedMetric,
    pub log_fails: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct MmdsMetrics {
    pub rx_accepted: SharedMetric,
    pub rx_accepted_err: SharedMetric,
    pub rx_bad_eth: SharedMetric,
    pub tx_bytes: SharedMetric,
    pub tx_errors: SharedMetric,
    pub tx_frames: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct NetDeviceMetrics {
    pub activate_fails: SharedMetric,
    pub cfg_fails: SharedMetric,
    pub event_fails: SharedMetric,
    pub rx_queue_event_count: SharedMetric,
    pub rx_event_rate_limiter_count: SharedMetric,
    pub rx_tap_event_count: SharedMetric,
    pub rx_bytes_count: SharedMetric,
    pub rx_packets_count: SharedMetric,
    pub rx_fails: SharedMetric,
    pub tx_queue_event_count: SharedMetric,
    pub tx_rate_limiter_event_count: SharedMetric,
    pub tx_bytes_count: SharedMetric,
    pub tx_packets_count: SharedMetric,
    pub tx_fails: SharedMetric,
}

#[derive(Serialize)]
pub struct SeccompMetrics {
    pub num_faults: SharedMetric,
    pub bad_syscalls: Vec<SharedMetric>,
}

impl Default for SeccompMetrics {
    fn default() -> SeccompMetrics {
        let mut def_syscalls = vec![];
        for _syscall in 0..SYSCALL_MAX {
            def_syscalls.push(SharedMetric::default());
        }
        SeccompMetrics {
            num_faults: SharedMetric::default(),
            bad_syscalls: def_syscalls,
        }
    }
}

#[derive(Default, Serialize)]
pub struct SerialDeviceMetrics {
    pub error_count: SharedMetric,
    pub flush_count: SharedMetric,
    pub missed_read_count: SharedMetric,
    pub missed_write_count: SharedMetric,
    pub read_count: SharedMetric,
    pub write_count: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct VcpuMetrics {
    pub eagain: SharedMetric,
    pub eintr: SharedMetric,
    pub exit_io_in: SharedMetric,
    pub exit_io_out: SharedMetric,
    pub exit_mmio_read: SharedMetric,
    pub exit_mmio_write: SharedMetric,
    pub failures: SharedMetric,
}

#[derive(Default, Serialize)]
pub struct VmmMetrics {
    pub device_events: SharedMetric,
    pub panic_count: SharedMetric,
}

// The sole purpose of this struct is to produce an UTC timestamp when an instance is serialized.
#[derive(Default)]
struct SerializeToUtcTimestampMs;

impl Serialize for SerializeToUtcTimestampMs {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_i64(chrono::Utc::now().timestamp_millis())
    }
}

#[derive(Default, Serialize)]
pub struct FirecrackerMetrics {
    utc_timestamp_ms: SerializeToUtcTimestampMs,
    pub api_server: ApiServerMetrics,
    pub block: BlockDeviceMetrics,
    pub get_api_requests: GetRequestsMetrics,
    pub i8042: I8042DeviceMetrics,
    pub logger: LoggerSystemMetrics,
    pub mmds: MmdsMetrics,
    pub net: NetDeviceMetrics,
    pub patch_api_requests: PatchRequestsMetrics,
    pub put_api_requests: PutRequestsMetrics,
    pub seccomp: SeccompMetrics,
    pub vcpu: VcpuMetrics,
    pub vmm: VmmMetrics,
    pub uart: SerialDeviceMetrics,
}

lazy_static! {
    /// Static instance used for handling metrics.
    ///
    pub static ref METRICS: FirecrackerMetrics = FirecrackerMetrics::default();
}

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

    use std::sync::Arc;
    use std::thread;

    #[test]
    fn test_metric() {
        let m1 = SimpleMetric::default();

        m1.inc();
        m1.inc();
        m1.add(5);
        m1.inc();

        assert_eq!(m1.count(), 8);

        let m2 = Arc::new(SharedMetric::default());

        // We're going to create a number of threads that will attempt to increase this metric
        // in parallel. If everything goes fine we still can't be sure the synchronization works,
        // but it something fails, then we definitely have a problem :-s

        const NUM_THREADS_TO_SPAWN: usize = 4;
        const NUM_INCREMENTS_PER_THREAD: usize = 100000;
        const M2_INITIAL_COUNT: usize = 123;

        m2.add(M2_INITIAL_COUNT);

        let mut v = Vec::with_capacity(NUM_THREADS_TO_SPAWN);

        for _ in 0..NUM_THREADS_TO_SPAWN {
            let r = m2.clone();
            v.push(thread::spawn(move || {
                for _ in 0..NUM_INCREMENTS_PER_THREAD {
                    r.inc();
                }
            }));
        }

        for handle in v {
            handle.join().unwrap();
        }

        assert_eq!(
            m2.count(),
            M2_INITIAL_COUNT + NUM_THREADS_TO_SPAWN * NUM_INCREMENTS_PER_THREAD
        );
    }

    #[test]
    fn test_serialize() {
        let s = serde_json::to_string(&FirecrackerMetrics::default());
        assert!(s.is_ok());
    }
}
<|endoftext|>


<commit_before>use crate::grammar::ast::{eq::AstEq, Expression, Parens};
use crate::grammar::model::{Fragment, HasFragment};
use crate::grammar::parsers::expression::ToExpression;
use crate::grammar::parsers::with_input;
use nom::character::complete::{char as ch, multispace0};
use nom::combinator::map;
use nom::sequence::delimited;
use nom::IResult;

impl<'s> Parens<'s> {
    fn inner(frag: Fragment<'s>) -> IResult<Fragment<'s>, Expression<'s>> {
        delimited(
            multispace0,
            delimited(ch('('), Expression::parse, ch(')')),
            multispace0,
        )(frag)
    }

    /// Parse parentheses and the expression between them in source code. Will
    /// ignore any whitespace before and after.
    pub fn parse(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        map(with_input(Self::inner), |(consumed, expr)| Parens {
            frag: consumed,
            inner: Box::new(expr),
        })(input)
    }
}

impl<'s> HasFragment<'s> for Parens<'s> {
    fn get_fragment(&self) -> Fragment<'s> {
        self.frag
    }
}

impl<'s> ToExpression<'s> for Parens<'s> {
    fn create_expr(self) -> Expression<'s> {
        Expression::Parens(self)
    }
}

impl<'s> AstEq for Parens<'s> {
    #[inline]
    fn ast_eq(fst: &Self, snd: &Self) -> bool {
        AstEq::ast_eq(&*fst.inner, &*snd.inner)
    }
}
<commit_msg>parens<commit_after>use crate::grammar::ast::{eq::AstEq, Expression, Parens};
use crate::grammar::model::{Fragment, HasFragment};
use crate::grammar::parsers::expression::ToExpression;
use crate::grammar::parsers::whitespace::token_delimiter;
use crate::grammar::parsers::with_input;
use nom::character::complete::{char as ch, multispace0};
use nom::combinator::map;
use nom::sequence::{delimited, terminated};
use nom::IResult;

impl<'s> Parens<'s> {
    fn inner(frag: Fragment<'s>) -> IResult<Fragment<'s>, Expression<'s>> {
        delimited(
            token_delimiter,
            delimited(
                terminated(ch('('), token_delimiter),
                Expression::parse,
                terminated(ch(')'), token_delimiter),
            ),
            token_delimiter,
        )(frag)
    }

    /// Parse parentheses and the expression between them in source code. Will
    /// ignore any whitespace before and after.
    pub fn parse(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        map(with_input(Self::inner), |(consumed, expr)| Parens {
            frag: consumed,
            inner: Box::new(expr),
        })(input)
    }
}

impl<'s> HasFragment<'s> for Parens<'s> {
    fn get_fragment(&self) -> Fragment<'s> {
        self.frag
    }
}

impl<'s> ToExpression<'s> for Parens<'s> {
    fn create_expr(self) -> Expression<'s> {
        Expression::Parens(self)
    }
}

impl<'s> AstEq for Parens<'s> {
    #[inline]
    fn ast_eq(fst: &Self, snd: &Self) -> bool {
        AstEq::ast_eq(&*fst.inner, &*snd.inner)
    }
}
<|endoftext|>


<commit_before>use crate::prelude::*;

impl Uuid {
    /// Creates a random UUID.
    ///
    /// This uses the [`getrandom`] crate to utilise the operating system's RNG
    /// as the source of random numbers. If you'd like to use a custom generator,
    /// don't use this method: use the `rand::Rand trait`'s `rand()` method instead.
    ///
    /// Note that usage of this method requires the `v4` feature of this crate
    /// to be enabled.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// use uuid::Uuid;
    ///
    /// let uuid = Uuid::new_v4();
    /// ```
    ///
    /// [`getrandom`]: https://crates.io/crates/getrandom
    /// [`rand`]: https://crates.io/crates/rand
    // TODO: change signature to support uuid's Error.
    pub fn new_v4() -> Result<Uuid, getrandom::Error> {
        let mut bytes = [0u8; 16];
        getrandom::getrandom(&mut bytes)?;

        let uuid = crate::builder::Builder::from_bytes(bytes)
            .set_variant(Variant::RFC4122)
            .set_version(Version::Random)
            .build();
        Ok(uuid)
    }
}

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

    #[test]
    fn test_new() {
        let uuid = Uuid::new_v4().unwrap();

        assert_eq!(uuid.get_version(), Some(Version::Random));
        assert_eq!(uuid.get_variant(), Some(Variant::RFC4122));
    }

    #[test]
    fn test_get_version() {
        let uuid = Uuid::new_v4().unwrap();

        assert_eq!(uuid.get_version(), Some(Version::Random));
        assert_eq!(uuid.get_version_num(), 4)
    }
}
<commit_msg>Clarify documentation around the relation to the rand crate<commit_after>use crate::prelude::*;

impl Uuid {
    /// Creates a random UUID.
    ///
    /// This uses the [`getrandom`] crate to utilise the operating system's RNG
    /// as the source of random numbers. If you'd like to use a custom generator,
    /// don't use this method: generate random bytes using your custom generator
    /// and pass them to the [`uuid::Builder::from_bytes`][from_bytes] function
    /// instead.
    ///
    /// Note that usage of this method requires the `v4` feature of this crate
    /// to be enabled.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// use uuid::Uuid;
    ///
    /// let uuid = Uuid::new_v4();
    /// ```
    ///
    /// [`getrandom`]: https://crates.io/crates/getrandom
    /// [from_bytes]: struct.Builder.html#method.from_bytes
    // TODO: change signature to support uuid's Error.
    pub fn new_v4() -> Result<Uuid, getrandom::Error> {
        let mut bytes = [0u8; 16];
        getrandom::getrandom(&mut bytes)?;

        let uuid = crate::builder::Builder::from_bytes(bytes)
            .set_variant(Variant::RFC4122)
            .set_version(Version::Random)
            .build();
        Ok(uuid)
    }
}

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

    #[test]
    fn test_new() {
        let uuid = Uuid::new_v4().unwrap();

        assert_eq!(uuid.get_version(), Some(Version::Random));
        assert_eq!(uuid.get_variant(), Some(Variant::RFC4122));
    }

    #[test]
    fn test_get_version() {
        let uuid = Uuid::new_v4().unwrap();

        assert_eq!(uuid.get_version(), Some(Version::Random));
        assert_eq!(uuid.get_version_num(), 4)
    }
}
<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    conn.connect();

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();
    t.poll();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      debug!("Heartbeat");
      self.start_send(Frame::Heartbeat(0));
      self.poll_complete();
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      self.poll_complete();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      self.upstream.start_send(f);
      self.poll_complete();
    }
    //self.poll_complete();
  }

  pub fn handle_frames(&mut self) {
    while let Ok(Async::Ready(Some(_))) = self.poll() {}
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat();
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<commit_msg>futures: split and use send_frame<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    conn.connect();

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();
    t.poll();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      trace!("Sending heartbeat");
      if let Err(e) = self.send_frame(Frame::Heartbeat(0)) {
        debug!("Failed to send heartbeat: {:?}", e);
      }
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      self.send_frame(f);
    }
    self.poll_complete();
  }

  fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  pub fn handle_frames(&mut self) {
    while let Ok(Async::Ready(Some(_))) = self.poll() {}
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat();
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<|endoftext|>


<commit_before>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#![feature(proc_macro)]

extern crate gtk;
#[macro_use]
extern crate relm;
extern crate relm_attributes;
#[macro_use]
extern crate relm_derive;
#[macro_use]
extern crate relm_test;

use gtk::{
    ButtonExt,
    Inhibit,
    LabelExt,
    OrientableExt,
    WidgetExt,
};
use gtk::Orientation::Vertical;
use relm::Widget;
use relm_attributes::widget;

use self::Msg::*;

// Define the structure of the model.
pub struct Model {
    counter: i32,
}

// The messages that can be sent to the update function.
#[derive(Msg)]
pub enum Msg {
    #[cfg(test)] Test,
    Decrement,
    Increment,
    Quit,
}

#[widget]
impl Widget for Win {
    // The initial model.
    fn model() -> Model {
        Model {
            counter: 0,
        }
    }

    // Update the model according to the message received.
    fn update(&mut self, event: Msg) {
        match event {
            #[cfg(test)] Test => (),
            Decrement => self.model.counter -= 1,
            Increment => self.model.counter += 1,
            Quit => gtk::main_quit(),
        }
    }

    view! {
        gtk::Window {
            gtk::Box {
                // Set the orientation property of the Box.
                orientation: Vertical,
                // Create a Button inside the Box.
                #[name="inc_button"]
                gtk::Button {
                    // Send the message Increment when the button is clicked.
                    clicked => Increment,
                    // TODO: check if using two events of the same name work.
                    label: "+",
                },
                #[name="label"]
                gtk::Label {
                    // Bind the text property of the label to the counter attribute of the model.
                    text: &self.model.counter.to_string(),
                },
                #[name="dec_button"]
                gtk::Button {
                    clicked => Decrement,
                    label: "-",
                },
            },
            delete_event(_, _) => (Quit, Inhibit(false)),
        }
    }
}

fn main() {
    Win::run(()).unwrap();
}

#[cfg(test)]
mod tests {
    use gtk::LabelExt;

    use relm;
    use relm_test::click;

    use Win;

    #[test]
    fn label_change() {
        let (_component, widgets) = relm::init_test::<Win>(()).unwrap();
        let inc_button = &widgets.inc_button;
        let dec_button = &widgets.dec_button;
        let label = &widgets.label;

        assert_label!(inc_button, "+");
        assert_label!(dec_button, "-");

        assert_text!(label, 0);
        click(inc_button);
        assert_text!(label, 1);
        click(inc_button);
        assert_text!(label, 2);
        click(inc_button);
        assert_text!(label, 3);
        click(inc_button);
        assert_text!(label, 4);

        click(dec_button);
        assert_text!(label, 3);
        click(dec_button);
        assert_text!(label, 2);
        click(dec_button);
        assert_text!(label, 1);
        click(dec_button);
        assert_text!(label, 0);
        click(dec_button);
        assert_text!(label, -1);
    }
}
<commit_msg>Add missing import<commit_after>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#![feature(proc_macro)]

extern crate gtk;
#[macro_use]
extern crate relm;
extern crate relm_attributes;
#[macro_use]
extern crate relm_derive;
#[macro_use]
extern crate relm_test;

use gtk::{
    ButtonExt,
    Inhibit,
    LabelExt,
    OrientableExt,
    WidgetExt,
};
use gtk::Orientation::Vertical;
use relm::Widget;
use relm_attributes::widget;

use self::Msg::*;

// Define the structure of the model.
pub struct Model {
    counter: i32,
}

// The messages that can be sent to the update function.
#[derive(Msg)]
pub enum Msg {
    #[cfg(test)] Test,
    Decrement,
    Increment,
    Quit,
}

#[widget]
impl Widget for Win {
    // The initial model.
    fn model() -> Model {
        Model {
            counter: 0,
        }
    }

    // Update the model according to the message received.
    fn update(&mut self, event: Msg) {
        match event {
            #[cfg(test)] Test => (),
            Decrement => self.model.counter -= 1,
            Increment => self.model.counter += 1,
            Quit => gtk::main_quit(),
        }
    }

    view! {
        gtk::Window {
            gtk::Box {
                // Set the orientation property of the Box.
                orientation: Vertical,
                // Create a Button inside the Box.
                #[name="inc_button"]
                gtk::Button {
                    // Send the message Increment when the button is clicked.
                    clicked => Increment,
                    // TODO: check if using two events of the same name work.
                    label: "+",
                },
                #[name="label"]
                gtk::Label {
                    // Bind the text property of the label to the counter attribute of the model.
                    text: &self.model.counter.to_string(),
                },
                #[name="dec_button"]
                gtk::Button {
                    clicked => Decrement,
                    label: "-",
                },
            },
            delete_event(_, _) => (Quit, Inhibit(false)),
        }
    }
}

fn main() {
    Win::run(()).unwrap();
}

#[cfg(test)]
mod tests {
    use gtk::{ButtonExt, LabelExt};

    use relm;
    use relm_test::click;

    use Win;

    #[test]
    fn label_change() {
        let (_component, widgets) = relm::init_test::<Win>(()).unwrap();
        let inc_button = &widgets.inc_button;
        let dec_button = &widgets.dec_button;
        let label = &widgets.label;

        assert_label!(inc_button, "+");
        assert_label!(dec_button, "-");

        assert_text!(label, 0);
        click(inc_button);
        assert_text!(label, 1);
        click(inc_button);
        assert_text!(label, 2);
        click(inc_button);
        assert_text!(label, 3);
        click(inc_button);
        assert_text!(label, 4);

        click(dec_button);
        assert_text!(label, 3);
        click(dec_button);
        assert_text!(label, 2);
        click(dec_button);
        assert_text!(label, 1);
        click(dec_button);
        assert_text!(label, 0);
        click(dec_button);
        assert_text!(label, -1);
    }
}
<|endoftext|>


<commit_before>use petgraph::Graph;
use petgraph::visit::Bfs;
use petgraph::prelude::NodeIndex;
use rand::{Rng, SeedableRng, StdRng};

/// Basic Markov chain for generating random strings which are similar to the
/// trained data provided.
pub struct MarkovChain {
    graph: Graph<String, f64>,
    start: NodeIndex,
    end: NodeIndex,
    strings: Vec<String>,
    rng_generator: StdRng,
}

impl MarkovChain {
    /// Creates a new MarkovChain and training it with the provided strings,
    /// also configures the random generator with the given string.
    /// Providing the same training data and seed will produce the same set
    /// of names in the same order.
    pub fn new(starting_strings: &Vec<String>, seed: u32) -> MarkovChain {
        let max_len = starting_strings
            .iter()
            .clone()
            .max_by_key(|s| s.len())
            .unwrap()
            .len();
        let mut graph = Graph::<String, f64>::new();
        let start = graph.add_node(String::from("Start"));
        let end = graph.add_node(String::from("End"));

        let alphabet = vec![
            "A",
            "B",
            "C",
            "D",
            "E",
            "F",
            "G",
            "H",
            "I",
            "J",
            "K",
            "L",
            "M",
            "N",
            "O",
            "P",
            "Q",
            "R",
            "S",
            "T",
            "U",
            "V",
            "W",
            "X",
            "Y",
            "Z",
        ];

        // Temporary container of all nodes in previous layer
        let mut prev_layer = Vec::<NodeIndex>::new();
        prev_layer.push(start);

        // Initialize graph with layers, nodes, and edges.
        for _ in 0..max_len {
            let mut current_layer = vec![];

            for chr in alphabet.clone() {
                let current_node = graph.add_node(String::from(chr));
                current_layer.push(current_node);

                // Add edges to all nodes in previous layer
                for prev_node in &prev_layer {
                    graph.add_edge(prev_node.clone(), current_node, 0.0);
                }

                // Also include 'end' node so that words can be shorter
                graph.add_edge(current_node, end, 0.0);
            }
            prev_layer = current_layer;
        }

        // Create and initialize random generator using seed
        let new_seed: &[_] = &[seed.clone() as usize];
        let rng: StdRng = SeedableRng::from_seed(new_seed);

        let mut markov = MarkovChain {
            graph: graph,
            start: start,
            end: end,
            strings: starting_strings.clone(),
            rng_generator: rng,
        };

        // Train the model
        for string in starting_strings {
            markov.train(string);
        }

        // Normalize all edges such that the sum of all outgoing edges from a
        // node is 1
        markov.normalize();
        markov
    }

    /// Train the model with the given input string.
    /// Traverses the graph, going from start to end through the nodes with
    /// resulting in the given string.
    /// Increments all edges it walks over by one.
    fn train(&mut self, input: &String) {
        let mut chars = input.chars();
        let mut current_node = self.start;

        fn increment_edge(
            first_node: &NodeIndex,
            second_node: &NodeIndex,
            graph: &mut Graph<String, f64>,
        ) {
            let edge = graph.find_edge(*first_node, *second_node).unwrap();
            let current_edge_weight = graph.edge_weight(edge).unwrap().clone();
            graph.update_edge(*first_node, *second_node, current_edge_weight + 1.0);
        }

        while let Some(current_char) = chars.next() {
            let next_node = self.graph
                .neighbors(current_node)
                .find(|node| {
                    self.graph.node_weight(*node) == Some(&current_char.to_string())
                })
                .unwrap();

            increment_edge(&current_node, &next_node, &mut self.graph);
            current_node = next_node;
        }

        increment_edge(&current_node, &self.end, &mut self.graph);
    }

    /// Normalizes all edges for all nodes such that the sum of the edge weights
    /// for all edges from a node is 1. I.e the edge weights will be normalized
    /// to a value in [0,1] reflecting the probability of that edge being used
    /// in the training data.
    fn normalize(&mut self) {
        let mut bfs = Bfs::new(&self.graph, self.start);
        while let Some(node) = bfs.next(&self.graph) {
            let edge_sum = match self.graph.edges(node).fold(0.0, |sum, edge| {
                sum + *edge.weight() as f64
            }) {
                x if x > 0.0 => x,
                _ => 1.0,
            };

            // Normalize edges for this node
            let mut edges = self.graph.neighbors(node).detach();
            while let Some(edge) = edges.next_edge(&self.graph) {
                self.graph[edge] /= edge_sum;
            }
        }
    }

    /// Generate a new string from the model, the new string cannot be longer
    /// than any string in the training data.
    /// Can currently generate duplicate strings, also may produce string which
    /// exist in the training data.
    pub fn generate(&mut self) -> String {
        let mut final_string = String::new();
        let mut current_node = self.start;

        // Traverse until we hit end
        while current_node != self.end {
            let mut sample = self.rng_generator.gen::<f64>();

            // Go through all edges, when an edge weight exceeds the random
            // value, go through that edge
            let mut edges = self.graph.neighbors(current_node).detach();
            while let Some(edge) = edges.next_edge(&self.graph) {
                sample -= self.graph[edge];

                // Go through this edge
                if sample <= 0.0 {
                    // Add current node value to the final string
                    final_string += self.graph.node_weight(current_node).unwrap();
                    if let Some(nodes) = self.graph.edge_endpoints(edge) {
                        current_node = nodes.1;
                        break;
                    }
                }
            }
        }
        // TODO: Ensure that all generated strings are unique, and not part of
        // training data
        final_string.drain(0..self.graph.node_weight(self.start).unwrap().len());
        final_string
    }
}

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

    const DEFAULT_SEED: u32 = 42;

    #[test]
    fn test_new_markov() {
        let starting_strings = vec![String::from("HELLO")];
        let markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);
        assert_eq!(markov.graph.node_count(), 5 * 26 + 2);
        assert_eq!(markov.graph.edge_count(), 26 + 4 * 26 * 27 + 26);
    }

    #[test]
    fn test_train() {
        let training_str = String::from("ABC");
        let starting_strings = vec![training_str.clone()];
        let markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);

        assert_eq!(
            markov.graph.edge_references().fold(0.0, |sum, edge| {
                sum + *edge.weight() as f64
            }),
            training_str.len() as f64 + 1.0
        );
    }

    #[test]
    fn test_normalize() {
        let starting_strings = vec![
            String::from("ABC"),
            String::from("ACC"),
            String::from("ACD"),
        ];
        let mut markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);
        for string in &starting_strings {
            markov.train(string);
        }

        markov.normalize();
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == markov.start && edge.target() == NodeIndex::new(2) &&
                edge.weight == 1.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2) && edge.target() == NodeIndex::new(2 + 26 + 1) &&
                edge.weight == 1.0 / 3.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2) && edge.target() == NodeIndex::new(2 + 26 + 2) &&
                edge.weight == 2.0 / 3.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 26 + 1) &&
                edge.target() == NodeIndex::new(2 + 2 * 26 + 2) && edge.weight == 1.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 26 + 2) &&
                edge.target() == NodeIndex::new(2 + 2 * 26 + 2) &&
                edge.weight == 1.0 / 2.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 26 + 2) &&
                edge.target() == NodeIndex::new(2 + 2 * 26 + 3) &&
                edge.weight == 1.0 / 2.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 2 * 26 + 3) && edge.target() == markov.end &&
                edge.weight == 1.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 2 * 26 + 2) && edge.target() == markov.end &&
                edge.weight == 1.0
        }));
    }

    #[test]
    fn test_generate() {
        let starting_strings = vec![
            String::from("ABCD"),
            String::from("ACCC"),
            String::from("ACDC"),
        ];
        let mut markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);

        assert_eq!(markov.generate(), String::from("ABCC"));
        assert_eq!(markov.generate(), String::from("ACDC"));
        assert_eq!(markov.generate(), String::from("ACCC"));
    }
}
<commit_msg>Make markov chain char based instead of string<commit_after>use petgraph::Graph;
use petgraph::visit::Bfs;
use petgraph::prelude::NodeIndex;
use rand::{Rng, SeedableRng, StdRng};

/// Basic Markov chain for generating random strings which are similar to the
/// trained data provided.
pub struct MarkovChain {
    graph: Graph<char, f64>,
    start: NodeIndex,
    end: NodeIndex,
    strings: Vec<String>,
    rng_generator: StdRng,
}

impl MarkovChain {
    /// Creates a new MarkovChain and training it with the provided strings,
    /// also configures the random generator with the given string.
    /// Providing the same training data and seed will produce the same set
    /// of names in the same order.
    pub fn new(starting_strings: &Vec<String>, seed: u32) -> MarkovChain {
        let max_len = starting_strings
            .iter()
            .clone()
            .max_by_key(|s| s.len())
            .unwrap()
            .len();
        let mut graph = Graph::<char, f64>::new();
        let start = graph.add_node('<');
        let end = graph.add_node('>');

        let alphabet: Vec<char> = String::from("ABCDEFGHIJKLMNOPQRSTUVWXYZ").chars().collect();

        // Temporary container of all nodes in previous layer
        let mut prev_layer = Vec::<NodeIndex>::new();
        prev_layer.push(start);

        // Initialize graph with layers, nodes, and edges.
        for _ in 0..max_len {
            let mut current_layer = vec![];

            for chr in alphabet.clone() {
                let current_node = graph.add_node(chr);
                current_layer.push(current_node);

                // Add edges to all nodes in previous layer
                for prev_node in &prev_layer {
                    graph.add_edge(prev_node.clone(), current_node, 0.0);
                }

                // Also include 'end' node so that words can be shorter
                graph.add_edge(current_node, end, 0.0);
            }
            prev_layer = current_layer;
        }

        // Create and initialize random generator using seed
        let new_seed: &[_] = &[seed.clone() as usize];
        let rng: StdRng = SeedableRng::from_seed(new_seed);

        let mut markov = MarkovChain {
            graph: graph,
            start: start,
            end: end,
            strings: starting_strings.clone(),
            rng_generator: rng,
        };

        // Train the model
        for string in starting_strings {
            markov.train(string);
        }

        // Normalize all edges such that the sum of all outgoing edges from a
        // node is 1
        markov.normalize();
        markov
    }

    /// Train the model with the given input string.
    /// Traverses the graph, going from start to end through the nodes with
    /// resulting in the given string.
    /// Increments all edges it walks over by one.
    fn train(&mut self, input: &String) {
        let mut chars = input.chars();
        let mut current_node = self.start;

        fn increment_edge(
            first_node: &NodeIndex,
            second_node: &NodeIndex,
            graph: &mut Graph<char, f64>,
        ) {
            let edge = graph.find_edge(*first_node, *second_node).unwrap();
            let current_edge_weight = graph.edge_weight(edge).unwrap().clone();
            graph.update_edge(*first_node, *second_node, current_edge_weight + 1.0);
        }

        while let Some(current_char) = chars.next() {
            let next_node = self.graph
                .neighbors(current_node)
                .find(|node| self.graph.node_weight(*node) == Some(&current_char))
                .unwrap();

            increment_edge(&current_node, &next_node, &mut self.graph);
            current_node = next_node;
        }

        increment_edge(&current_node, &self.end, &mut self.graph);
    }

    /// Normalizes all edges for all nodes such that the sum of the edge weights
    /// for all edges from a node is 1. I.e the edge weights will be normalized
    /// to a value in [0,1] reflecting the probability of that edge being used
    /// in the training data.
    fn normalize(&mut self) {
        let mut bfs = Bfs::new(&self.graph, self.start);
        while let Some(node) = bfs.next(&self.graph) {
            let edge_sum = match self.graph.edges(node).fold(0.0, |sum, edge| {
                sum + *edge.weight() as f64
            }) {
                x if x > 0.0 => x,
                _ => 1.0,
            };

            // Normalize edges for this node
            let mut edges = self.graph.neighbors(node).detach();
            while let Some(edge) = edges.next_edge(&self.graph) {
                self.graph[edge] /= edge_sum;
            }
        }
    }

    /// Generate a new string from the model, the new string cannot be longer
    /// than any string in the training data.
    /// Can currently generate duplicate strings, also may produce string which
    /// exist in the training data.
    pub fn generate(&mut self) -> String {
        let mut final_string = String::new();
        let mut current_node = self.start;

        // Traverse until we hit end
        while current_node != self.end {
            let mut sample = self.rng_generator.gen::<f64>();

            // Go through all edges, when an edge weight exceeds the random
            // value, go through that edge
            let mut edges = self.graph.neighbors(current_node).detach();
            while let Some(edge) = edges.next_edge(&self.graph) {
                sample -= self.graph[edge];

                // Go through this edge
                if sample <= 0.0 {
                    // Add current node value to the final string
                    final_string.push(*self.graph.node_weight(current_node).unwrap());
                    if let Some(nodes) = self.graph.edge_endpoints(edge) {
                        current_node = nodes.1;
                        break;
                    }
                }
            }
        }
        // TODO: Ensure that all generated strings are unique, and not part of
        // training data
        final_string.drain(0..1);
        final_string
    }
}

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

    const DEFAULT_SEED: u32 = 42;

    #[test]
    fn test_new_markov() {
        let starting_strings = vec![String::from("HELLO")];
        let markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);
        assert_eq!(markov.graph.node_count(), 5 * 26 + 2);
        assert_eq!(markov.graph.edge_count(), 26 + 4 * 26 * 27 + 26);
    }

    #[test]
    fn test_train() {
        let training_str = String::from("ABC");
        let starting_strings = vec![training_str.clone()];
        let markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);

        assert_eq!(
            markov.graph.edge_references().fold(0.0, |sum, edge| {
                sum + *edge.weight() as f64
            }),
            training_str.len() as f64 + 1.0
        );
    }

    #[test]
    fn test_normalize() {
        let starting_strings = vec![
            String::from("ABC"),
            String::from("ACC"),
            String::from("ACD"),
        ];
        let mut markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);
        for string in &starting_strings {
            markov.train(string);
        }

        markov.normalize();
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == markov.start && edge.target() == NodeIndex::new(2) &&
                edge.weight == 1.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2) && edge.target() == NodeIndex::new(2 + 26 + 1) &&
                edge.weight == 1.0 / 3.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2) && edge.target() == NodeIndex::new(2 + 26 + 2) &&
                edge.weight == 2.0 / 3.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 26 + 1) &&
                edge.target() == NodeIndex::new(2 + 2 * 26 + 2) && edge.weight == 1.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 26 + 2) &&
                edge.target() == NodeIndex::new(2 + 2 * 26 + 2) &&
                edge.weight == 1.0 / 2.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 26 + 2) &&
                edge.target() == NodeIndex::new(2 + 2 * 26 + 3) &&
                edge.weight == 1.0 / 2.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 2 * 26 + 3) && edge.target() == markov.end &&
                edge.weight == 1.0
        }));
        assert!(markov.graph.raw_edges().iter().any(|ref edge| {
            edge.source() == NodeIndex::new(2 + 2 * 26 + 2) && edge.target() == markov.end &&
                edge.weight == 1.0
        }));
    }

    #[test]
    fn test_generate() {
        let starting_strings = vec![
            String::from("ABCD"),
            String::from("ACCC"),
            String::from("ACDC"),
        ];
        let mut markov = MarkovChain::new(&starting_strings, DEFAULT_SEED);

        assert_eq!(markov.generate(), String::from("ABCC"));
        assert_eq!(markov.generate(), String::from("ACDC"));
        assert_eq!(markov.generate(), String::from("ACCC"));
    }
}
<|endoftext|>


<commit_before>extern mod extra;

use std::io;
use extra::net;
use extra::uv;

struct Client {
    writer: @io::WriterUtil,
    reader: @io::ReaderUtil
}

impl Client {
    fn flush(&self) -> Result<~str, ~str> {
        self.writer.write_str("flush_all\r\n");
        let result_str = self.reader.read_line();
        if result_str == ~"OK\r" {  // XXX: read_line dose not remove \r, is it a bug?
            return Ok(result_str);
        }
        else {
            return Err(result_str);
        }
    }

    fn _store(&self, action: ~str, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self.writer.write_str(fmt!("%s %s 0 %u %u\r\n", action, key, exp_time, value.len()));
        self.writer.write_str(fmt!("%s\r\n", value));
        let result_str = self.reader.read_line();
        if result_str == ~"STORED\r" {
            return Ok(result_str);
        }
        else {
            return Err(result_str);
        }
    }

    fn set(&self, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self._store(~"set", key, value, exp_time)
    }

    fn add(&self, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self._store(~"add", key, value, exp_time)
    }

    fn replace(&self, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self._store(~"replace", key, value, exp_time)
    }
}

fn main() {
    let client = connect(~"127.0.0.1", 11211).get();

    client.flush();

    client.set(~"foo", ~"bar", 0);
    client.add(~"foo", ~"bar", 0);
    client.replace(~"foo", ~"bar", 0);

}

fn connect(addr: ~str, port: uint) -> Result<Client, net::tcp::TcpConnectErrData> {
    match net::tcp::connect(net::ip::v4::parse_addr(addr), port, &uv::global_loop::get()) {
        Err(err) => {
            //match err {
            //    net::tcp::GenericConnectErr(name,msg) => io::println(fmt!("Connection error %s: %s", name, msg)),
            //    net::tcp::ConnectionRefused => io::println("Connection refused")
            //}
            return Err(err);
        },
        Ok(socket) => {
            let socket_buf = net::tcp::socket_buf(socket);
            let writer = @socket_buf as @io::WriterUtil;
            let reader = @socket_buf as @io::ReaderUtil;
            let client = Client{writer: writer, reader: reader};
            return Ok(client);
        }
    }

}
<commit_msg>implement other command<commit_after>extern mod extra;

use std::io;
use extra::net;
use extra::uv;

struct Client {
    writer: @io::WriterUtil,
    reader: @io::ReaderUtil
}

impl Client {
    fn flush(&self) -> Result<~str, ~str> {
        self.writer.write_str("flush_all\r\n");
        let result_str = self.reader.read_line();
        if result_str == ~"OK\r" {  // XXX: read_line dose not remove \r, is it a bug?
            return Ok(result_str);
        }
        else {
            return Err(result_str);
        }
    }

    fn delete(&self, key: ~str, time: uint) -> Result<~str, ~str> {
        self.writer.write_str(fmt!("delete %s %u\r\n", key, time));
        let result_str = self.reader.read_line();
        if result_str == ~ "DELETED\r" {
            return Ok(result_str);
        }
        else {
            return Err(result_str);
        }
    }

    fn get(&self, key: ~str) -> Result<~str, ~str> {
        self.writer.write_str(fmt!("get %s\r\n", key));
        let result_str = self.reader.read_line();
        if result_str == ~"END\r" {
            return Err(~"NOT_FOUND");
        }
        // TODO:
        let result_str = self.reader.read_line();
        return Ok(result_str);

    }

    fn incr(&self, key: ~str, value: int) -> Result<~str, ~str> {
        self.writer.write_str(fmt!("incr %s %i\r\n", key, value));
        let result_str = self.reader.read_line();
        if result_str == ~"NOT_FOUND\r" {
            return Err(result_str);
        }
        else {
            return Ok(result_str);
        }
    }

    fn _store(&self, action: ~str, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self.writer.write_str(fmt!("%s %s 0 %u %u\r\n", action, key, exp_time, value.len()));
        self.writer.write_str(fmt!("%s\r\n", value));
        let result_str = self.reader.read_line();
        if result_str == ~"STORED\r" {
            return Ok(result_str);
        }
        else {
            return Err(result_str);
        }
    }

    fn set(&self, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self._store(~"set", key, value, exp_time)
    }

    fn add(&self, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self._store(~"add", key, value, exp_time)
    }

    fn replace(&self, key: ~str, value: ~str, exp_time: uint) -> Result<~str, ~str> {
        self._store(~"replace", key, value, exp_time)
    }
}

fn main() {
    let client = connect(~"127.0.0.1", 11211).get();

    client.flush();
    client.set(~"foo", ~"bar", 0);
    client.add(~"foo", ~"bar", 0);
    client.get(~"foo").get();
    client.replace(~"foo", ~"0", 0);
    client.incr(~"foo", 1);
    client.delete(~"foo", 0);

}

fn connect(addr: ~str, port: uint) -> Result<Client, net::tcp::TcpConnectErrData> {
    match net::tcp::connect(net::ip::v4::parse_addr(addr), port, &uv::global_loop::get()) {
        Err(err) => {
            //match err {
            //    net::tcp::GenericConnectErr(name,msg) => io::println(fmt!("Connection error %s: %s", name, msg)),
            //    net::tcp::ConnectionRefused => io::println("Connection refused")
            //}
            return Err(err);
        },
        Ok(socket) => {
            let socket_buf = net::tcp::socket_buf(socket);
            let writer = @socket_buf as @io::WriterUtil;
            let reader = @socket_buf as @io::ReaderUtil;
            let client = Client{writer: writer, reader: reader};
            return Ok(client);
        }
    }

}
<|endoftext|>


<commit_before>#![allow(dead_code)]

extern crate jam;
extern crate alto;
extern crate cgmath;

use alto::Alto;
// use jam::audio::*;
use jam::audio::{Listener, SoundEvent};
use jam::audio::engine::{process};
use jam::audio::engine::SoundEngineUpdate::*;
use jam::HashMap;
use jam::Vec3f;
use cgmath::Zero;

#[cfg(target_os = "windows")]
const OPENAL_PATH: &'static str = "./native/windows/OpenAL64.dll";
#[cfg(target_os = "mac_os")]
const OPENAL_PATH: &'static str = "./native/mac/openal.dylib";

fn main() {
    let alto = Alto::load(OPENAL_PATH).unwrap();
    let dev = alto.open(None).unwrap();
    let ctx = dev.new_context(None).unwrap();
    let mut cb = jam::audio::context::create_sound_context(&ctx, "resources/sound", "ogg");
    cb.create_sources(32, 4).unwrap();

    let sounds = vec![("teleport".into(), 1.0), ("water".into(), 1.0)];
    process(&mut cb, Preload(sounds)).unwrap();

    let listener = Listener::default();

    let sound_event = SoundEvent {
        name: "teleport".into(),
        position: Vec3f::zero(),
        gain: 1.0,
        pitch: 1.5,
        attenuation:1.0,
    };
    let sound_event_b = SoundEvent {
        name: "water".into(),
        position: Vec3f::zero(),
        gain: 1.0,
        pitch: 1.0,
        attenuation:1.0,
    };

    process(&mut cb, Render { master_gain: 1.0, sounds:vec![sound_event, sound_event_b], persistent_sounds:HashMap::default(), listener: listener }).unwrap();

    std::thread::sleep(std::time::Duration::new(2, 0));

    process(&mut cb, Clear).unwrap();

    std::thread::sleep(std::time::Duration::new(1, 0));
}<commit_msg>Fixed MacOS library loading<commit_after>#![allow(dead_code)]

extern crate jam;
extern crate alto;
extern crate cgmath;

use alto::Alto;
// use jam::audio::*;
use jam::audio::{Listener, SoundEvent};
use jam::audio::engine::{process};
use jam::audio::engine::SoundEngineUpdate::*;
use jam::HashMap;
use jam::Vec3f;
use cgmath::Zero;

#[cfg(target_os = "windows")]
const OPENAL_PATH: &'static str = "./native/windows/OpenAL64.dll";
#[cfg(target_os = "macos")]
const OPENAL_PATH: &'static str = "./native/mac/openal.dylib";

fn main() {
    let alto = Alto::load(OPENAL_PATH).unwrap();
    let dev = alto.open(None).unwrap();
    let ctx = dev.new_context(None).unwrap();
    let mut cb = jam::audio::context::create_sound_context(&ctx, "resources/sound", "ogg");
    cb.create_sources(32, 4).unwrap();

    let sounds = vec![("teleport".into(), 1.0), ("water".into(), 1.0)];
    process(&mut cb, Preload(sounds)).unwrap();

    let listener = Listener::default();

    let sound_event = SoundEvent {
        name: "teleport".into(),
        position: Vec3f::zero(),
        gain: 1.0,
        pitch: 1.5,
        attenuation:1.0,
    };
    let sound_event_b = SoundEvent {
        name: "water".into(),
        position: Vec3f::zero(),
        gain: 1.0,
        pitch: 1.0,
        attenuation:1.0,
    };

    process(&mut cb, Render { master_gain: 1.0, sounds:vec![sound_event, sound_event_b], persistent_sounds:HashMap::default(), listener: listener }).unwrap();

    std::thread::sleep(std::time::Duration::new(2, 0));

    process(&mut cb, Clear).unwrap();

    std::thread::sleep(std::time::Duration::new(1, 0));
}<|endoftext|>