<commit_before>use crate::{internal_prelude::*, Duration};
use core::{
    cmp::{Ord, Ordering, PartialEq, PartialOrd},
    ops::{Add, AddAssign, Sub, SubAssign},
    time::Duration as StdDuration,
};
use std::time::Instant as StdInstant;

/// A measurement of a monotonically non-decreasing clock. Opaque and useful
/// only with [`Duration`].
///
/// Instants are always guaranteed to be no less than any previously measured
/// instant when created, and are often useful for tasks such as measuring
/// benchmarks or timing how long an operation takes.
///
/// Note, however, that instants are not guaranteed to be **steady**. In other
/// words, each tick of the underlying clock may not be the same length (e.g.
/// some seconds may be longer than others). An instant may jump forwards or
/// experience time dilation (slow down or speed up), but it will never go
/// backwards.
///
/// Instants are opaque types that can only be compared to one another. There is
/// no method to get "the number of seconds" from an instant. Instead, it only
/// allows measuring the duration between two instants (or comparing two
/// instants).
///
/// This implementation allows for operations with signed [`Duration`]s, but is
/// otherwise identical to [`std::time::Instant`].
#[cfg_attr(docs, doc(cfg(feature = "std")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Instant {
    /// Inner representation, using `std::time::Instant`.
    inner: StdInstant,
}

impl Instant {
    /// Returns an `Instant` corresponding to "now".
    ///
    /// ```rust,no_run
    /// # use time::Instant;
    /// println!("{:?}", Instant::now());
    /// ```
    #[inline(always)]
    pub fn now() -> Self {
        Self {
            inner: StdInstant::now(),
        }
    }

    /// Returns the amount of time elapsed since this instant was created. The
    /// duration will always be nonnegative if the instant is not synthetically
    /// created.
    ///
    /// ```rust
    /// # use time::{Instant, prelude::*};
    /// # use std::thread;
    /// let instant = Instant::now();
    /// thread::sleep(100.std_milliseconds());
    /// assert!(instant.elapsed() >= 100.milliseconds());
    /// ```
    #[inline(always)]
    pub fn elapsed(self) -> Duration {
        Self::now() - self
    }

    /// Returns `Some(t)` where `t` is the time `self + duration` if `t` can be
    /// represented as `Instant` (which means it's inside the bounds of the
    /// underlying data structure), `None` otherwise.
    ///
    /// ```rust
    /// # use time::{Instant, prelude::*};
    /// let now = Instant::now();
    /// assert_eq!(
    ///     now.checked_add(5.seconds()),
    ///     Some(now + 5.seconds())
    /// );
    /// assert_eq!(
    ///     now.checked_add((-5).seconds()),
    ///     Some(now + (-5).seconds())
    /// );
    /// ```
    ///
    /// This function is only present when using rustc >= 1.34.0.
    #[inline]
    #[cfg(__time_02_instant_checked_ops)]
    pub fn checked_add(self, duration: Duration) -> Option<Self> {
        if duration.is_zero() {
            Some(self)
        } else if duration.is_positive() {
            self.inner.checked_add(duration.abs_std()).map(From::from)
        } else {
            // duration.is_negative()
            self.inner.checked_sub(duration.abs_std()).map(From::from)
        }
    }

    /// Returns `Some(t)` where `t` is the time `self - duration` if `t` can be
    /// represented as `Instant` (which means it's inside the bounds of the
    /// underlying data structure), `None` otherwise.
    ///
    /// ```rust
    /// # use time::{Instant, prelude::*};
    /// let now = Instant::now();
    /// assert_eq!(
    ///     now.checked_sub(5.seconds()),
    ///     Some(now - 5.seconds())
    /// );
    /// assert_eq!(
    ///     now.checked_sub((-5).seconds()),
    ///     Some(now - (-5).seconds())
    /// );
    /// ```
    ///
    /// This function is only present when using rustc >= 1.34.0.
    #[inline(always)]
    #[cfg(__time_02_instant_checked_ops)]
    pub fn checked_sub(self, duration: Duration) -> Option<Self> {
        self.checked_add(-duration)
    }
}

#[allow(clippy::missing_docs_in_private_items)]
impl Instant {
    #[inline(always)]
    #[cfg(feature = "deprecated")]
    #[cfg_attr(tarpaulin, skip)]
    #[deprecated(since = "0.2.0", note = "Use `rhs - lhs`")]
    pub fn to(&self, later: Self) -> Duration {
        later - *self
    }
}

impl From<StdInstant> for Instant {
    #[inline(always)]
    fn from(instant: StdInstant) -> Self {
        Self { inner: instant }
    }
}

impl From<Instant> for StdInstant {
    #[inline(always)]
    fn from(instant: Instant) -> Self {
        instant.inner
    }
}

impl Sub for Instant {
    type Output = Duration;

    #[inline(always)]
    fn sub(self, other: Self) -> Self::Output {
        match self.inner.cmp(&other.inner) {
            Ordering::Equal => Duration::zero(),
            Ordering::Greater => (self.inner - other.inner)
                .try_into()
                .expect("overflow converting `std::time::Duration` to `time::Duration`"),
            Ordering::Less => -Duration::try_from(other.inner - self.inner)
                .expect("overflow converting `std::time::Duration` to `time::Duration`"),
        }
    }
}

impl Sub<StdInstant> for Instant {
    type Output = Duration;

    #[inline(always)]
    fn sub(self, other: StdInstant) -> Self::Output {
        self - Self::from(other)
    }
}

impl Sub<Instant> for StdInstant {
    type Output = Duration;

    #[inline(always)]
    fn sub(self, other: Instant) -> Self::Output {
        Instant::from(self) - other
    }
}

impl Add<Duration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn add(self, duration: Duration) -> Self::Output {
        if duration.is_positive() {
            (self.inner + duration.abs_std()).into()
        } else if duration.is_negative() {
            (self.inner - duration.abs_std()).into()
        } else {
            self
        }
    }
}

impl Add<Duration> for StdInstant {
    type Output = Self;

    #[inline(always)]
    fn add(self, duration: Duration) -> Self::Output {
        (Instant::from(self) + duration).into()
    }
}

impl Add<StdDuration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn add(self, duration: StdDuration) -> Self::Output {
        Self {
            inner: self.inner + duration,
        }
    }
}

impl AddAssign<Duration> for Instant {
    #[inline(always)]
    fn add_assign(&mut self, duration: Duration) {
        *self = *self + duration;
    }
}

impl AddAssign<Duration> for StdInstant {
    #[inline(always)]
    fn add_assign(&mut self, duration: Duration) {
        *self = *self + duration;
    }
}

impl AddAssign<StdDuration> for Instant {
    #[inline(always)]
    fn add_assign(&mut self, duration: StdDuration) {
        *self = *self + duration;
    }
}

impl Sub<Duration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn sub(self, duration: Duration) -> Self::Output {
        self + -duration
    }
}

impl Sub<Duration> for StdInstant {
    type Output = Self;

    #[inline(always)]
    fn sub(self, duration: Duration) -> Self::Output {
        (Instant::from(self) - duration).into()
    }
}

impl Sub<StdDuration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn sub(self, duration: StdDuration) -> Self::Output {
        Self {
            inner: self.inner - duration,
        }
    }
}

impl SubAssign<Duration> for Instant {
    #[inline(always)]
    fn sub_assign(&mut self, duration: Duration) {
        *self = *self - duration;
    }
}

impl SubAssign<Duration> for StdInstant {
    #[inline(always)]
    fn sub_assign(&mut self, duration: Duration) {
        *self = *self - duration;
    }
}

impl SubAssign<StdDuration> for Instant {
    #[inline(always)]
    fn sub_assign(&mut self, duration: StdDuration) {
        *self = *self - duration;
    }
}

impl PartialEq<StdInstant> for Instant {
    #[inline(always)]
    fn eq(&self, rhs: &StdInstant) -> bool {
        self.inner.eq(rhs)
    }
}

impl PartialEq<Instant> for StdInstant {
    #[inline(always)]
    fn eq(&self, rhs: &Instant) -> bool {
        self.eq(&rhs.inner)
    }
}

impl PartialOrd<StdInstant> for Instant {
    #[inline(always)]
    fn partial_cmp(&self, rhs: &StdInstant) -> Option<Ordering> {
        self.inner.partial_cmp(rhs)
    }
}

impl PartialOrd<Instant> for StdInstant {
    #[inline(always)]
    fn partial_cmp(&self, rhs: &Instant) -> Option<Ordering> {
        self.partial_cmp(&rhs.inner)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::prelude::*;
    use std::thread;

    #[test]
    fn elapsed() {
        let instant = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(instant.elapsed() >= 100.milliseconds());
    }

    #[test]
    #[cfg(__time_02_instant_checked_ops)]
    fn checked_add() {
        let now = Instant::now();
        assert_eq!(now.checked_add(5.seconds()), Some(now + 5.seconds()));
        assert_eq!(now.checked_add((-5).seconds()), Some(now + (-5).seconds()));
    }

    #[test]
    #[cfg(__time_02_instant_checked_ops)]
    fn checked_sub() {
        let now = Instant::now();
        assert_eq!(now.checked_sub(5.seconds()), Some(now - 5.seconds()));
        assert_eq!(now.checked_sub((-5).seconds()), Some(now - (-5).seconds()));
    }

    #[test]
    fn from_std() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time);
        assert_eq!(now_time, now_std);
    }

    #[test]
    fn to_std() {
        let now_std = StdInstant::now();
        let now_time = Instant::from(now_std);
        assert_eq!(now_time, now_std);
    }

    #[test]
    fn sub() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(Instant::now() - start >= 100.milliseconds());
    }

    #[test]
    fn sub_std() {
        let start = StdInstant::now();
        thread::sleep(100.std_milliseconds());
        assert!(Instant::now() - start >= 100.milliseconds());
    }

    #[test]
    fn std_sub() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(StdInstant::now() - start >= 100.milliseconds());
    }

    #[test]
    fn add_duration() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(start + 100.milliseconds() <= Instant::now());
    }

    #[test]
    fn std_add_duration() {
        let start = StdInstant::now();
        thread::sleep(100.std_milliseconds());
        assert!(start + 100.milliseconds() <= StdInstant::now());
    }

    #[test]
    fn add_std_duration() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(start + 100.std_milliseconds() <= Instant::now());
    }

    #[test]
    fn add_assign_duration() {
        let mut start = Instant::now();
        thread::sleep(100.std_milliseconds());
        start += 100.milliseconds();
        assert!(start <= Instant::now());
    }

    #[test]
    fn std_add_assign_duration() {
        let mut start = StdInstant::now();
        thread::sleep(100.std_milliseconds());
        start += 100.milliseconds();
        assert!(start <= StdInstant::now());
    }

    #[test]
    fn add_assign_std_duration() {
        let mut start = Instant::now();
        thread::sleep(100.std_milliseconds());
        start += 100.std_milliseconds();
        assert!(start <= Instant::now());
    }

    #[test]
    fn sub_duration() {
        let instant = Instant::now();
        assert!(instant - 100.milliseconds() <= Instant::now());
    }

    #[test]
    fn std_sub_duration() {
        let instant = StdInstant::now();
        assert!(instant - 100.milliseconds() <= StdInstant::now());
    }

    #[test]
    fn sub_std_duration() {
        let instant = Instant::now();
        assert!(instant - 100.std_milliseconds() <= Instant::now());
    }

    #[test]
    fn sub_assign_duration() {
        let mut instant = Instant::now();
        instant -= 100.milliseconds();
        assert!(instant <= Instant::now());
    }

    #[test]
    fn std_sub_assign_duration() {
        let mut instant = StdInstant::now();
        instant -= 100.milliseconds();
        assert!(instant <= StdInstant::now());
    }

    #[test]
    fn sub_assign_std_duration() {
        let mut instant = Instant::now();
        instant -= 100.std_milliseconds();
        assert!(instant <= Instant::now());
    }

    #[test]
    fn eq_std() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time);
        assert_eq!(now_time, now_std);
    }

    #[test]
    fn std_eq() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time);
        assert_eq!(now_std, now_time);
    }

    #[test]
    fn ord_std() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) + 1.seconds();
        assert!(now_time < now_std);

        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) - 1.seconds();
        assert!(now_time > now_std);
    }

    #[test]
    fn std_ord() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) + 1.seconds();
        assert!(now_std > now_time);

        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) - 1.seconds();
        assert!(now_std < now_time);
    }

    #[test]
    fn sub_regression() {
        let now = Instant::now();
        let future = now + Duration::seconds(5);
        let past = now - Duration::seconds(5);

        assert_eq!(future - now, Duration::seconds(5));
        assert_eq!(now - past, Duration::seconds(5));
        assert_eq!(future - past, Duration::seconds(10));

        assert_eq!(now - future, Duration::seconds(-5));
        assert_eq!(past - now, Duration::seconds(-5));
        assert_eq!(past - future, Duration::seconds(-10));
    }
}
<commit_msg>Silence lint in testing<commit_after>use crate::{internal_prelude::*, Duration};
use core::{
    cmp::{Ord, Ordering, PartialEq, PartialOrd},
    ops::{Add, AddAssign, Sub, SubAssign},
    time::Duration as StdDuration,
};
use std::time::Instant as StdInstant;

/// A measurement of a monotonically non-decreasing clock. Opaque and useful
/// only with [`Duration`].
///
/// Instants are always guaranteed to be no less than any previously measured
/// instant when created, and are often useful for tasks such as measuring
/// benchmarks or timing how long an operation takes.
///
/// Note, however, that instants are not guaranteed to be **steady**. In other
/// words, each tick of the underlying clock may not be the same length (e.g.
/// some seconds may be longer than others). An instant may jump forwards or
/// experience time dilation (slow down or speed up), but it will never go
/// backwards.
///
/// Instants are opaque types that can only be compared to one another. There is
/// no method to get "the number of seconds" from an instant. Instead, it only
/// allows measuring the duration between two instants (or comparing two
/// instants).
///
/// This implementation allows for operations with signed [`Duration`]s, but is
/// otherwise identical to [`std::time::Instant`].
#[cfg_attr(docs, doc(cfg(feature = "std")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Instant {
    /// Inner representation, using `std::time::Instant`.
    inner: StdInstant,
}

impl Instant {
    /// Returns an `Instant` corresponding to "now".
    ///
    /// ```rust,no_run
    /// # use time::Instant;
    /// println!("{:?}", Instant::now());
    /// ```
    #[inline(always)]
    pub fn now() -> Self {
        Self {
            inner: StdInstant::now(),
        }
    }

    /// Returns the amount of time elapsed since this instant was created. The
    /// duration will always be nonnegative if the instant is not synthetically
    /// created.
    ///
    /// ```rust
    /// # use time::{Instant, prelude::*};
    /// # use std::thread;
    /// let instant = Instant::now();
    /// thread::sleep(100.std_milliseconds());
    /// assert!(instant.elapsed() >= 100.milliseconds());
    /// ```
    #[inline(always)]
    pub fn elapsed(self) -> Duration {
        Self::now() - self
    }

    /// Returns `Some(t)` where `t` is the time `self + duration` if `t` can be
    /// represented as `Instant` (which means it's inside the bounds of the
    /// underlying data structure), `None` otherwise.
    ///
    /// ```rust
    /// # use time::{Instant, prelude::*};
    /// let now = Instant::now();
    /// assert_eq!(
    ///     now.checked_add(5.seconds()),
    ///     Some(now + 5.seconds())
    /// );
    /// assert_eq!(
    ///     now.checked_add((-5).seconds()),
    ///     Some(now + (-5).seconds())
    /// );
    /// ```
    ///
    /// This function is only present when using rustc >= 1.34.0.
    #[inline]
    #[cfg(__time_02_instant_checked_ops)]
    pub fn checked_add(self, duration: Duration) -> Option<Self> {
        if duration.is_zero() {
            Some(self)
        } else if duration.is_positive() {
            self.inner.checked_add(duration.abs_std()).map(From::from)
        } else {
            // duration.is_negative()
            self.inner.checked_sub(duration.abs_std()).map(From::from)
        }
    }

    /// Returns `Some(t)` where `t` is the time `self - duration` if `t` can be
    /// represented as `Instant` (which means it's inside the bounds of the
    /// underlying data structure), `None` otherwise.
    ///
    /// ```rust
    /// # use time::{Instant, prelude::*};
    /// let now = Instant::now();
    /// assert_eq!(
    ///     now.checked_sub(5.seconds()),
    ///     Some(now - 5.seconds())
    /// );
    /// assert_eq!(
    ///     now.checked_sub((-5).seconds()),
    ///     Some(now - (-5).seconds())
    /// );
    /// ```
    ///
    /// This function is only present when using rustc >= 1.34.0.
    #[inline(always)]
    #[cfg(__time_02_instant_checked_ops)]
    pub fn checked_sub(self, duration: Duration) -> Option<Self> {
        self.checked_add(-duration)
    }
}

#[allow(clippy::missing_docs_in_private_items)]
impl Instant {
    #[inline(always)]
    #[cfg(feature = "deprecated")]
    #[cfg_attr(tarpaulin, skip)]
    #[deprecated(since = "0.2.0", note = "Use `rhs - lhs`")]
    pub fn to(&self, later: Self) -> Duration {
        later - *self
    }
}

impl From<StdInstant> for Instant {
    #[inline(always)]
    fn from(instant: StdInstant) -> Self {
        Self { inner: instant }
    }
}

impl From<Instant> for StdInstant {
    #[inline(always)]
    fn from(instant: Instant) -> Self {
        instant.inner
    }
}

impl Sub for Instant {
    type Output = Duration;

    #[inline(always)]
    fn sub(self, other: Self) -> Self::Output {
        match self.inner.cmp(&other.inner) {
            Ordering::Equal => Duration::zero(),
            Ordering::Greater => (self.inner - other.inner)
                .try_into()
                .expect("overflow converting `std::time::Duration` to `time::Duration`"),
            Ordering::Less => -Duration::try_from(other.inner - self.inner)
                .expect("overflow converting `std::time::Duration` to `time::Duration`"),
        }
    }
}

impl Sub<StdInstant> for Instant {
    type Output = Duration;

    #[inline(always)]
    fn sub(self, other: StdInstant) -> Self::Output {
        self - Self::from(other)
    }
}

impl Sub<Instant> for StdInstant {
    type Output = Duration;

    #[inline(always)]
    fn sub(self, other: Instant) -> Self::Output {
        Instant::from(self) - other
    }
}

impl Add<Duration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn add(self, duration: Duration) -> Self::Output {
        if duration.is_positive() {
            (self.inner + duration.abs_std()).into()
        } else if duration.is_negative() {
            (self.inner - duration.abs_std()).into()
        } else {
            self
        }
    }
}

impl Add<Duration> for StdInstant {
    type Output = Self;

    #[inline(always)]
    fn add(self, duration: Duration) -> Self::Output {
        (Instant::from(self) + duration).into()
    }
}

impl Add<StdDuration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn add(self, duration: StdDuration) -> Self::Output {
        Self {
            inner: self.inner + duration,
        }
    }
}

impl AddAssign<Duration> for Instant {
    #[inline(always)]
    fn add_assign(&mut self, duration: Duration) {
        *self = *self + duration;
    }
}

impl AddAssign<Duration> for StdInstant {
    #[inline(always)]
    fn add_assign(&mut self, duration: Duration) {
        *self = *self + duration;
    }
}

impl AddAssign<StdDuration> for Instant {
    #[inline(always)]
    fn add_assign(&mut self, duration: StdDuration) {
        *self = *self + duration;
    }
}

impl Sub<Duration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn sub(self, duration: Duration) -> Self::Output {
        self + -duration
    }
}

impl Sub<Duration> for StdInstant {
    type Output = Self;

    #[inline(always)]
    fn sub(self, duration: Duration) -> Self::Output {
        (Instant::from(self) - duration).into()
    }
}

impl Sub<StdDuration> for Instant {
    type Output = Self;

    #[inline(always)]
    fn sub(self, duration: StdDuration) -> Self::Output {
        Self {
            inner: self.inner - duration,
        }
    }
}

impl SubAssign<Duration> for Instant {
    #[inline(always)]
    fn sub_assign(&mut self, duration: Duration) {
        *self = *self - duration;
    }
}

impl SubAssign<Duration> for StdInstant {
    #[inline(always)]
    fn sub_assign(&mut self, duration: Duration) {
        *self = *self - duration;
    }
}

impl SubAssign<StdDuration> for Instant {
    #[inline(always)]
    fn sub_assign(&mut self, duration: StdDuration) {
        *self = *self - duration;
    }
}

impl PartialEq<StdInstant> for Instant {
    #[inline(always)]
    fn eq(&self, rhs: &StdInstant) -> bool {
        self.inner.eq(rhs)
    }
}

impl PartialEq<Instant> for StdInstant {
    #[inline(always)]
    fn eq(&self, rhs: &Instant) -> bool {
        self.eq(&rhs.inner)
    }
}

impl PartialOrd<StdInstant> for Instant {
    #[inline(always)]
    fn partial_cmp(&self, rhs: &StdInstant) -> Option<Ordering> {
        self.inner.partial_cmp(rhs)
    }
}

impl PartialOrd<Instant> for StdInstant {
    #[inline(always)]
    fn partial_cmp(&self, rhs: &Instant) -> Option<Ordering> {
        self.partial_cmp(&rhs.inner)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    #[allow(unused_imports)]
    use crate::prelude::*;
    use std::thread;

    #[test]
    fn elapsed() {
        let instant = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(instant.elapsed() >= 100.milliseconds());
    }

    #[test]
    #[cfg(__time_02_instant_checked_ops)]
    fn checked_add() {
        let now = Instant::now();
        assert_eq!(now.checked_add(5.seconds()), Some(now + 5.seconds()));
        assert_eq!(now.checked_add((-5).seconds()), Some(now + (-5).seconds()));
    }

    #[test]
    #[cfg(__time_02_instant_checked_ops)]
    fn checked_sub() {
        let now = Instant::now();
        assert_eq!(now.checked_sub(5.seconds()), Some(now - 5.seconds()));
        assert_eq!(now.checked_sub((-5).seconds()), Some(now - (-5).seconds()));
    }

    #[test]
    fn from_std() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time);
        assert_eq!(now_time, now_std);
    }

    #[test]
    fn to_std() {
        let now_std = StdInstant::now();
        let now_time = Instant::from(now_std);
        assert_eq!(now_time, now_std);
    }

    #[test]
    fn sub() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(Instant::now() - start >= 100.milliseconds());
    }

    #[test]
    fn sub_std() {
        let start = StdInstant::now();
        thread::sleep(100.std_milliseconds());
        assert!(Instant::now() - start >= 100.milliseconds());
    }

    #[test]
    fn std_sub() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(StdInstant::now() - start >= 100.milliseconds());
    }

    #[test]
    fn add_duration() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(start + 100.milliseconds() <= Instant::now());
    }

    #[test]
    fn std_add_duration() {
        let start = StdInstant::now();
        thread::sleep(100.std_milliseconds());
        assert!(start + 100.milliseconds() <= StdInstant::now());
    }

    #[test]
    fn add_std_duration() {
        let start = Instant::now();
        thread::sleep(100.std_milliseconds());
        assert!(start + 100.std_milliseconds() <= Instant::now());
    }

    #[test]
    fn add_assign_duration() {
        let mut start = Instant::now();
        thread::sleep(100.std_milliseconds());
        start += 100.milliseconds();
        assert!(start <= Instant::now());
    }

    #[test]
    fn std_add_assign_duration() {
        let mut start = StdInstant::now();
        thread::sleep(100.std_milliseconds());
        start += 100.milliseconds();
        assert!(start <= StdInstant::now());
    }

    #[test]
    fn add_assign_std_duration() {
        let mut start = Instant::now();
        thread::sleep(100.std_milliseconds());
        start += 100.std_milliseconds();
        assert!(start <= Instant::now());
    }

    #[test]
    fn sub_duration() {
        let instant = Instant::now();
        assert!(instant - 100.milliseconds() <= Instant::now());
    }

    #[test]
    fn std_sub_duration() {
        let instant = StdInstant::now();
        assert!(instant - 100.milliseconds() <= StdInstant::now());
    }

    #[test]
    fn sub_std_duration() {
        let instant = Instant::now();
        assert!(instant - 100.std_milliseconds() <= Instant::now());
    }

    #[test]
    fn sub_assign_duration() {
        let mut instant = Instant::now();
        instant -= 100.milliseconds();
        assert!(instant <= Instant::now());
    }

    #[test]
    fn std_sub_assign_duration() {
        let mut instant = StdInstant::now();
        instant -= 100.milliseconds();
        assert!(instant <= StdInstant::now());
    }

    #[test]
    fn sub_assign_std_duration() {
        let mut instant = Instant::now();
        instant -= 100.std_milliseconds();
        assert!(instant <= Instant::now());
    }

    #[test]
    fn eq_std() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time);
        assert_eq!(now_time, now_std);
    }

    #[test]
    fn std_eq() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time);
        assert_eq!(now_std, now_time);
    }

    #[test]
    fn ord_std() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) + 1.seconds();
        assert!(now_time < now_std);

        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) - 1.seconds();
        assert!(now_time > now_std);
    }

    #[test]
    fn std_ord() {
        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) + 1.seconds();
        assert!(now_std > now_time);

        let now_time = Instant::now();
        let now_std = StdInstant::from(now_time) - 1.seconds();
        assert!(now_std < now_time);
    }

    #[test]
    fn sub_regression() {
        let now = Instant::now();
        let future = now + Duration::seconds(5);
        let past = now - Duration::seconds(5);

        assert_eq!(future - now, Duration::seconds(5));
        assert_eq!(now - past, Duration::seconds(5));
        assert_eq!(future - past, Duration::seconds(10));

        assert_eq!(now - future, Duration::seconds(-5));
        assert_eq!(past - now, Duration::seconds(-5));
        assert_eq!(past - future, Duration::seconds(-10));
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serialized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn cancel_serialized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        let res = self.do_run();
        self.internal_rpc.poll(&self.channels).and(res)
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.should_heartbeat() {
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.channels.set_connection_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!("We've written {} but didn't expect to write anything", written);
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.send_continue()
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.channels.set_connection_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.channels.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.channels.set_connection_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>properly handle mio 'error' readiness events<commit_after>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serialized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn cancel_serialized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        let res = self.do_run();
        self.internal_rpc.poll(&self.channels).and(res)
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.should_heartbeat() {
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!("We've written {} but didn't expect to write anything", written);
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.send_continue()
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.channels.set_connection_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.channels.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.channels.set_connection_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Linus Färnstrand <faern@faern.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets on a fake network managed
//! by in memory FIFO queues. Useful for writing tests.

use std::io;
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time;

use datalink::{self, EthernetDataLinkChannelIterator, EthernetDataLinkReceiver,
               EthernetDataLinkSender, NetworkInterface};
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use util::MacAddr;

/// Configuration for the dummy datalink backend
#[derive(Debug)]
pub struct Config {
    /// The fake network will pull packets (or errors) form this `Receiver`
    pub in_packets: Receiver<io::Result<Box<[u8]>>>,

    /// All packets sent to this fake network will end up on this `Sender`
    pub out_packets: Sender<Vec<u8>>,
}

impl<'a> From<&'a datalink::Config> for Config {
    /// This conversion will not allow injecting and reading packets from the dummy network.
    /// To do that please create your own `dummy::Config`
    fn from(_config: &datalink::Config) -> Config {
        Config {
            in_packets: mpsc::channel().1,
            out_packets: mpsc::channel().0,
        }
    }
}

impl Default for Config {
    /// This conversion will not allow injecting and reading packets from the dummy network.
    /// To do that please create your own `dummy::Config`
    fn default() -> Config {
        Config {
            in_packets: mpsc::channel().1,
            out_packets: mpsc::channel().0,
        }
    }
}

/// Create a data link channel backed by FIFO queues. Useful for debugging and testing.
pub fn channel(_: &NetworkInterface, config: Config) -> io::Result<datalink::Channel> {
    let sender = Box::new(MockEthernetDataLinkSender { out_packets: config.out_packets_tx });
    let receiver =
        Box::new(MockEthernetDataLinkReceiver { in_packets: Some(config.in_packets_rx) });

    Ok(datalink::Channel::Ethernet(sender, receiver))
}

struct MockEthernetDataLinkSender {
    out_packets: Sender<Vec<u8>>,
}

impl EthernetDataLinkSender for MockEthernetDataLinkSender {
    fn build_and_send(&mut self,
                      _num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
        -> Option<io::Result<()>> {
        let mut buffer = vec![0; packet_size];
        {
            let pkg = MutableEthernetPacket::new(&mut buffer[..]).unwrap();
            func(pkg);
        }
        // Send the data to the queue. Don't care if it's closed
        self.out_packets.send(buffer).unwrap_or(());
        Some(Ok(()))
    }

    fn send_to(&mut self,
               _packet: &EthernetPacket,
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        panic!("Not implemented in mock");
    }
}

struct MockEthernetDataLinkReceiver {
    in_packets: Option<Receiver<io::Result<Box<[u8]>>>>,
}

impl EthernetDataLinkReceiver for MockEthernetDataLinkReceiver {
    fn iter<'a>(&'a mut self) -> Box<EthernetDataLinkChannelIterator + 'a> {
        Box::new(MockEthernetDataLinkChannelIterator {
            in_packets: self.in_packets.take().expect("Only one receiver allowed"),
            used_packets: vec![],
        })
    }
}

struct MockEthernetDataLinkChannelIterator {
    in_packets: Receiver<io::Result<Box<[u8]>>>,
    used_packets: Vec<Box<[u8]>>,
}

impl<'a> EthernetDataLinkChannelIterator<'a> for MockEthernetDataLinkChannelIterator {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        match self.in_packets.recv() {
            Ok(result) => {
                match result {
                    Ok(buffer) => {
                        self.used_packets.push(buffer);
                        let buffer_ref = &*self.used_packets[self.used_packets.len() - 1];
                        let packet = EthernetPacket::new(buffer_ref).unwrap();
                        Ok(packet)
                    }
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // When we run out of test packets we sleep forever.
                // Simulating no more packets arrive
                loop {
                    thread::sleep(time::Duration::new(1, 0));
                }
            }
        }
    }
}

/// Get three fake interfaces generated with `dummy_interface`
pub fn interfaces() -> Vec<NetworkInterface> {
    (0..3).map(|i| dummy_interface(i)).collect()
}

/// Generates a fake `NetworkInterface`.
/// The name of the interface will be `ethX` where X is the integer `i`.
/// The index will be `i`.
/// The MAC will be `01:02:03:04:05:i`.
pub fn dummy_interface(i: u8) -> NetworkInterface {
    NetworkInterface {
        name: format!("eth{}", i),
        index: i as u32,
        mac: Some(MacAddr::new(1, 2, 3, 4, 5, i)),
        ips: None,
        flags: 0,
    }
}
<commit_msg>Improve dummy::Config with inject/read handles<commit_after>// Copyright (c) 2016 Linus Färnstrand <faern@faern.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for sending and receiving data link layer packets on a fake network managed
//! by in memory FIFO queues. Useful for writing tests.

use std::io;
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time;

use datalink::{self, EthernetDataLinkChannelIterator, EthernetDataLinkReceiver,
               EthernetDataLinkSender, NetworkInterface};
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use util::MacAddr;

/// Configuration for the dummy datalink backend
#[derive(Debug)]
pub struct Config {
    /// The fake network will pull packets (or errors) form this `Receiver`
    pub in_packets_rx: Receiver<io::Result<Box<[u8]>>>,

    in_packets_tx: Option<Sender<io::Result<Box<[u8]>>>>,

    /// All packets sent to this fake network will end up on this `Sender`
    pub out_packets_tx: Sender<Vec<u8>>,

    out_packets_rx: Option<Receiver<Vec<u8>>>,
}

impl Config {
    /// Get the `Sender` handle that can inject packets in the fake network
    pub fn inject_handle(&mut self) -> Option<Sender<io::Result<Box<[u8]>>>> {
        self.in_packets_tx.take()
    }

    /// Get the `Receiver` handle where packets sent to the fake network can be read
    pub fn read_handle(&mut self) -> Option<Receiver<Vec<u8>>> {
        self.out_packets_rx.take()
    }
}

impl<'a> From<&'a datalink::Config> for Config {
    /// This conversion will not allow injecting and reading packets from the dummy network.
    /// To do that please create your own `dummy::Config`
    fn from(_config: &datalink::Config) -> Config {
        Config::default()
    }
}

impl Default for Config {
    /// This conversion will not allow injecting and reading packets from the dummy network.
    /// To do that please create your own `dummy::Config`
    fn default() -> Config {
        let (in_tx, in_rx) = mpsc::channel();
        let (out_tx, out_rx) = mpsc::channel();
        Config {
            in_packets_rx: in_rx,
            in_packets_tx: Some(in_tx),
            out_packets_tx: out_tx,
            out_packets_rx: Some(out_rx),
        }
    }
}

/// Create a data link channel backed by FIFO queues. Useful for debugging and testing.
pub fn channel(_: &NetworkInterface, config: Config) -> io::Result<datalink::Channel> {
    let sender = Box::new(MockEthernetDataLinkSender { out_packets: config.out_packets_tx });
    let receiver =
        Box::new(MockEthernetDataLinkReceiver { in_packets: Some(config.in_packets_rx) });

    Ok(datalink::Channel::Ethernet(sender, receiver))
}

struct MockEthernetDataLinkSender {
    out_packets: Sender<Vec<u8>>,
}

impl EthernetDataLinkSender for MockEthernetDataLinkSender {
    fn build_and_send(&mut self,
                      _num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
        -> Option<io::Result<()>> {
        let mut buffer = vec![0; packet_size];
        {
            let pkg = MutableEthernetPacket::new(&mut buffer[..]).unwrap();
            func(pkg);
        }
        // Send the data to the queue. Don't care if it's closed
        self.out_packets.send(buffer).unwrap_or(());
        Some(Ok(()))
    }

    fn send_to(&mut self,
               _packet: &EthernetPacket,
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        panic!("Not implemented in mock");
    }
}

struct MockEthernetDataLinkReceiver {
    in_packets: Option<Receiver<io::Result<Box<[u8]>>>>,
}

impl EthernetDataLinkReceiver for MockEthernetDataLinkReceiver {
    fn iter<'a>(&'a mut self) -> Box<EthernetDataLinkChannelIterator + 'a> {
        Box::new(MockEthernetDataLinkChannelIterator {
            in_packets: self.in_packets.take().expect("Only one receiver allowed"),
            used_packets: vec![],
        })
    }
}

struct MockEthernetDataLinkChannelIterator {
    in_packets: Receiver<io::Result<Box<[u8]>>>,
    used_packets: Vec<Box<[u8]>>,
}

impl<'a> EthernetDataLinkChannelIterator<'a> for MockEthernetDataLinkChannelIterator {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        match self.in_packets.recv() {
            Ok(result) => {
                match result {
                    Ok(buffer) => {
                        self.used_packets.push(buffer);
                        let buffer_ref = &*self.used_packets[self.used_packets.len() - 1];
                        let packet = EthernetPacket::new(buffer_ref).unwrap();
                        Ok(packet)
                    }
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // When we run out of test packets we sleep forever.
                // Simulating no more packets arrive
                loop {
                    thread::sleep(time::Duration::new(1, 0));
                }
            }
        }
    }
}

/// Get three fake interfaces generated with `dummy_interface`
pub fn interfaces() -> Vec<NetworkInterface> {
    (0..3).map(|i| dummy_interface(i)).collect()
}

/// Generates a fake `NetworkInterface`.
/// The name of the interface will be `ethX` where X is the integer `i`.
/// The index will be `i`.
/// The MAC will be `01:02:03:04:05:i`.
pub fn dummy_interface(i: u8) -> NetworkInterface {
    NetworkInterface {
        name: format!("eth{}", i),
        index: i as u32,
        mac: Some(MacAddr::new(1, 2, 3, 4, 5, i)),
        ips: None,
        flags: 0,
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Manages jailing and connecting virtio devices to the system bus.

use std::fmt;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::{Arc, Mutex};

use io_jail::Minijail;
use kvm::IoeventAddress;
use sys_util::GuestMemory;
use sys_util;

use hw;
use kernel_cmdline;
use vm_control::VmRequest;

/// Errors for device manager.
#[derive(Debug)]
pub enum Error {
    /// Could not create the mmio device to wrap a VirtioDevice.
    CreateMmioDevice(sys_util::Error),
    /// Failed to clone a queue's ioeventfd.
    CloneIoeventFd(sys_util::Error),
    /// Failed to clone the mmio irqfd.
    CloneIrqFd(sys_util::Error),
    /// Failed to initialize proxy device for jailed device.
    ProxyDeviceCreation(io::Error),
    /// Appending to kernel command line failed.
    Cmdline(kernel_cmdline::Error),
    /// No more IRQs are available.
    IrqsExhausted,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::CreateMmioDevice(ref e) => write!(f, "failed to create mmio device: {:?}", e),
            &Error::CloneIoeventFd(ref e) => write!(f, "failed to clone ioeventfd: {:?}", e),
            &Error::CloneIrqFd(ref e) => write!(f, "failed to clone irqfd: {:?}", e),
            &Error::ProxyDeviceCreation(ref e) => write!(f, "failed to create proxy device: {}", e),
            &Error::Cmdline(ref e) => {
                write!(f, "unable to add device to kernel command line: {}", e)
            }
            &Error::IrqsExhausted => {
                write!(f, "no more IRQs are available")
            }
        }
    }
}

type Result<T> = ::std::result::Result<T, Error>;

const MAX_IRQ: u32 = 15;

/// Manages the complexities of adding a device.
pub struct DeviceManager {
    pub bus: hw::Bus,
    pub vm_requests: Vec<VmRequest>,
    guest_mem: GuestMemory,
    mmio_len: u64,
    mmio_base: u64,
    irq: u32,
}

impl DeviceManager {
    /// Create a new DeviceManager.
    pub fn new(guest_mem: GuestMemory, mmio_len: u64, mmio_base: u64, irq_base: u32) -> DeviceManager {
        DeviceManager {
            guest_mem: guest_mem,
            vm_requests: Vec::new(),
            mmio_len: mmio_len,
            mmio_base: mmio_base,
            irq: irq_base,
            bus: hw::Bus::new(),
        }
    }

    /// Register a device to be used via MMIO transport.
    pub fn register_mmio(&mut self,
                         device: Box<hw::virtio::VirtioDevice>,
                         jail: Option<Minijail>,
                         cmdline: &mut kernel_cmdline::Cmdline)
                         -> Result<()> {
        if self.irq > MAX_IRQ {
            return Err(Error::IrqsExhausted);
        }

        // List of FDs to keep open in the child after it forks.
        let mut keep_fds: Vec<RawFd> = device.keep_fds();

        let mmio_device = hw::virtio::MmioDevice::new(self.guest_mem.clone(), device)
            .map_err(Error::CreateMmioDevice)?;
        for (i, queue_evt) in mmio_device.queue_evts().iter().enumerate() {
            let io_addr = IoeventAddress::Mmio(self.mmio_base +
                                               hw::virtio::NOITFY_REG_OFFSET as u64);
            self.vm_requests.push(VmRequest::RegisterIoevent(queue_evt
                                                                 .try_clone()
                                                                 .map_err(Error::CloneIoeventFd)?,
                                                             io_addr,
                                                             i as u32));
            keep_fds.push(queue_evt.as_raw_fd());
        }

        if let Some(interrupt_evt) = mmio_device.interrupt_evt() {
            self.vm_requests.push(VmRequest::RegisterIrqfd(interrupt_evt
                                                               .try_clone()
                                                               .map_err(Error::CloneIrqFd)?,
                                                           self.irq));
            keep_fds.push(interrupt_evt.as_raw_fd());
        }

        if let Some(jail) = jail {
            let proxy_dev = hw::ProxyDevice::new(mmio_device, move |keep_pipe| {
                keep_fds.push(keep_pipe.as_raw_fd());
                // Need to panic here as there isn't a way to recover from a
                // partly-jailed process.
                unsafe {
                    // This is OK as we have whitelisted all the FDs we need open.
                    jail.enter(Some(&keep_fds)).unwrap();
                }
            }).map_err(|e| Error::ProxyDeviceCreation(e))?;
            self.bus
                .insert(Arc::new(Mutex::new(proxy_dev)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        } else {
            self.bus
                .insert(Arc::new(Mutex::new(mmio_device)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        }

        cmdline
            .insert("virtio_mmio.device",
                    &format!("4K@0x{:08x}:{}", self.mmio_base, self.irq))
            .map_err(Error::Cmdline)?;
        self.mmio_base += self.mmio_len;
        self.irq += 1;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;
    use std::os::unix::io::RawFd;
    use sys_util::{EventFd, GuestAddress, GuestMemory};
    use device_manager;
    use kernel_cmdline;
    use hw;

    const QUEUE_SIZES: &'static [u16] = &[64];

    #[allow(dead_code)]
    struct DummyDevice {
        dummy: u32,
    }

    impl hw::virtio::VirtioDevice for DummyDevice {
        fn keep_fds(&self) -> Vec<RawFd> {
            Vec::new()
        }

        fn device_type(&self) -> u32 {
            0
        }

        fn queue_max_sizes(&self) -> &[u16] {
            QUEUE_SIZES
        }

        #[allow(unused_variables)]
        #[allow(unused_mut)]
        fn activate(&mut self,
                    mem: GuestMemory,
                    interrupt_evt: EventFd,
                    status: Arc<AtomicUsize>,
                    queues: Vec<hw::virtio::Queue>,
                    mut queue_evts: Vec<EventFd>) {
        }
    }

    #[test]
    fn register_device() {
        let start_addr1 = GuestAddress(0x0);
        let start_addr2 = GuestAddress(0x1000);
        let guest_mem = GuestMemory::new(&vec![(start_addr1, 0x1000), (start_addr2, 0x1000)])
            .unwrap();
        let mut device_manager =
            device_manager::DeviceManager::new(guest_mem, 0x1000, 0xd0000000, 5);

        let mut cmdline = kernel_cmdline::Cmdline::new(4096);
        let dummy_box = Box::new(DummyDevice { dummy: 0 });
        device_manager
            .register_mmio(dummy_box, None, &mut cmdline)
            .unwrap();
    }
}
<commit_msg>crosvm: jailed child processes inherit stderr<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Manages jailing and connecting virtio devices to the system bus.

use std::fmt;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::{Arc, Mutex};

use libc::STDERR_FILENO;

use io_jail::Minijail;
use kvm::IoeventAddress;
use sys_util::GuestMemory;
use sys_util;

use hw;
use kernel_cmdline;
use vm_control::VmRequest;

/// Errors for device manager.
#[derive(Debug)]
pub enum Error {
    /// Could not create the mmio device to wrap a VirtioDevice.
    CreateMmioDevice(sys_util::Error),
    /// Failed to clone a queue's ioeventfd.
    CloneIoeventFd(sys_util::Error),
    /// Failed to clone the mmio irqfd.
    CloneIrqFd(sys_util::Error),
    /// Failed to initialize proxy device for jailed device.
    ProxyDeviceCreation(io::Error),
    /// Appending to kernel command line failed.
    Cmdline(kernel_cmdline::Error),
    /// No more IRQs are available.
    IrqsExhausted,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::CreateMmioDevice(ref e) => write!(f, "failed to create mmio device: {:?}", e),
            &Error::CloneIoeventFd(ref e) => write!(f, "failed to clone ioeventfd: {:?}", e),
            &Error::CloneIrqFd(ref e) => write!(f, "failed to clone irqfd: {:?}", e),
            &Error::ProxyDeviceCreation(ref e) => write!(f, "failed to create proxy device: {}", e),
            &Error::Cmdline(ref e) => {
                write!(f, "unable to add device to kernel command line: {}", e)
            }
            &Error::IrqsExhausted => {
                write!(f, "no more IRQs are available")
            }
        }
    }
}

type Result<T> = ::std::result::Result<T, Error>;

const MAX_IRQ: u32 = 15;

/// Manages the complexities of adding a device.
pub struct DeviceManager {
    pub bus: hw::Bus,
    pub vm_requests: Vec<VmRequest>,
    guest_mem: GuestMemory,
    mmio_len: u64,
    mmio_base: u64,
    irq: u32,
}

impl DeviceManager {
    /// Create a new DeviceManager.
    pub fn new(guest_mem: GuestMemory, mmio_len: u64, mmio_base: u64, irq_base: u32) -> DeviceManager {
        DeviceManager {
            guest_mem: guest_mem,
            vm_requests: Vec::new(),
            mmio_len: mmio_len,
            mmio_base: mmio_base,
            irq: irq_base,
            bus: hw::Bus::new(),
        }
    }

    /// Register a device to be used via MMIO transport.
    pub fn register_mmio(&mut self,
                         device: Box<hw::virtio::VirtioDevice>,
                         jail: Option<Minijail>,
                         cmdline: &mut kernel_cmdline::Cmdline)
                         -> Result<()> {
        if self.irq > MAX_IRQ {
            return Err(Error::IrqsExhausted);
        }

        // List of FDs to keep open in the child after it forks.
        let mut keep_fds: Vec<RawFd> = device.keep_fds();
        keep_fds.push(STDERR_FILENO);

        let mmio_device = hw::virtio::MmioDevice::new(self.guest_mem.clone(), device)
            .map_err(Error::CreateMmioDevice)?;
        for (i, queue_evt) in mmio_device.queue_evts().iter().enumerate() {
            let io_addr = IoeventAddress::Mmio(self.mmio_base +
                                               hw::virtio::NOITFY_REG_OFFSET as u64);
            self.vm_requests.push(VmRequest::RegisterIoevent(queue_evt
                                                                 .try_clone()
                                                                 .map_err(Error::CloneIoeventFd)?,
                                                             io_addr,
                                                             i as u32));
            keep_fds.push(queue_evt.as_raw_fd());
        }

        if let Some(interrupt_evt) = mmio_device.interrupt_evt() {
            self.vm_requests.push(VmRequest::RegisterIrqfd(interrupt_evt
                                                               .try_clone()
                                                               .map_err(Error::CloneIrqFd)?,
                                                           self.irq));
            keep_fds.push(interrupt_evt.as_raw_fd());
        }

        if let Some(jail) = jail {
            let proxy_dev = hw::ProxyDevice::new(mmio_device, move |keep_pipe| {
                keep_fds.push(keep_pipe.as_raw_fd());
                // Need to panic here as there isn't a way to recover from a
                // partly-jailed process.
                unsafe {
                    // This is OK as we have whitelisted all the FDs we need open.
                    jail.enter(Some(&keep_fds)).unwrap();
                }
            }).map_err(|e| Error::ProxyDeviceCreation(e))?;
            self.bus
                .insert(Arc::new(Mutex::new(proxy_dev)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        } else {
            self.bus
                .insert(Arc::new(Mutex::new(mmio_device)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        }

        cmdline
            .insert("virtio_mmio.device",
                    &format!("4K@0x{:08x}:{}", self.mmio_base, self.irq))
            .map_err(Error::Cmdline)?;
        self.mmio_base += self.mmio_len;
        self.irq += 1;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;
    use std::os::unix::io::RawFd;
    use sys_util::{EventFd, GuestAddress, GuestMemory};
    use device_manager;
    use kernel_cmdline;
    use hw;

    const QUEUE_SIZES: &'static [u16] = &[64];

    #[allow(dead_code)]
    struct DummyDevice {
        dummy: u32,
    }

    impl hw::virtio::VirtioDevice for DummyDevice {
        fn keep_fds(&self) -> Vec<RawFd> {
            Vec::new()
        }

        fn device_type(&self) -> u32 {
            0
        }

        fn queue_max_sizes(&self) -> &[u16] {
            QUEUE_SIZES
        }

        #[allow(unused_variables)]
        #[allow(unused_mut)]
        fn activate(&mut self,
                    mem: GuestMemory,
                    interrupt_evt: EventFd,
                    status: Arc<AtomicUsize>,
                    queues: Vec<hw::virtio::Queue>,
                    mut queue_evts: Vec<EventFd>) {
        }
    }

    #[test]
    fn register_device() {
        let start_addr1 = GuestAddress(0x0);
        let start_addr2 = GuestAddress(0x1000);
        let guest_mem = GuestMemory::new(&vec![(start_addr1, 0x1000), (start_addr2, 0x1000)])
            .unwrap();
        let mut device_manager =
            device_manager::DeviceManager::new(guest_mem, 0x1000, 0xd0000000, 5);

        let mut cmdline = kernel_cmdline::Cmdline::new(4096);
        let dummy_box = Box::new(DummyDevice { dummy: 0 });
        device_manager
            .register_mmio(dummy_box, None, &mut cmdline)
            .unwrap();
    }
}
<|endoftext|>


<commit_before>use hash_map::HashMap;
use core::hash;
use core::cmp;

pub type DispatchFn<T, U> = fn(T, &mut U) -> bool;

pub struct DispatchTable<T: cmp::Eq + hash::Hash, U> {
    table: HashMap<T, DispatchFn<T, U>>,
}

impl<T: cmp::Eq + hash::Hash, U> DispatchTable<T, U> {
    pub fn new(size: usize) -> Self {
        DispatchTable { table: HashMap::new(size) }
    }

    pub fn dispatch(&self, event: T, arguments: &mut U) -> bool {
        match self.table.get(&event) {
            Some(func) => func(event, arguments),
            None => false,
        }
    }

    pub fn register(&mut self, event: T, func: DispatchFn<T, U>) {
        self.table.insert(event, func);
    }

    pub fn unregister(&mut self, event: T) {
        self.table.remove(event);
    }
}

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

    #[derive(Hash, Eq, PartialEq, Debug)]
    enum Event {
        VMExit,
        PageFault,
    }

    fn on_vmexit(_: &Event, received: bool) -> bool {
        assert!(received);
        true
    }

    #[test]
    fn test_dispatch_table_register() {
        let mut dt = DispatchTable::<Event, bool>::new(16);
        dt.register(Event::VMExit, on_vmexit);
        assert!(dt.dispatch(&Event::VMExit, true));
        assert!(!dt.dispatch(&Event::PageFault, false));
    }

    #[test]
    fn test_dispatch_table_unregister() {
        let mut dt = DispatchTable::<Event, bool>::new(16);
        dt.register(Event::VMExit, on_vmexit);
        assert!(dt.dispatch(&Event::VMExit, true));
        assert!(!dt.dispatch(&Event::PageFault, false));
        dt.unregister(Event::VMExit);
        assert!(!dt.dispatch(&Event::VMExit, false));
    }

}
<commit_msg>Fix dispatch table tests<commit_after>use hash_map::HashMap;
use core::hash;
use core::cmp;

pub type DispatchFn<T, U> = fn(T, &mut U) -> bool;

pub struct DispatchTable<T: cmp::Eq + hash::Hash, U> {
    table: HashMap<T, DispatchFn<T, U>>,
}

impl<T: cmp::Eq + hash::Hash, U> DispatchTable<T, U> {
    pub fn new(size: usize) -> Self {
        DispatchTable { table: HashMap::new(size) }
    }

    pub fn dispatch(&self, event: T, arguments: &mut U) -> bool {
        match self.table.get(&event) {
            Some(func) => func(event, arguments),
            None => false,
        }
    }

    pub fn register(&mut self, event: T, func: DispatchFn<T, U>) {
        self.table.insert(event, func);
    }

    pub fn unregister(&mut self, event: T) {
        self.table.remove(event);
    }
}

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

    #[derive(Hash, Eq, PartialEq, Debug)]
    enum Event {
        VMExit,
        PageFault,
    }

    fn on_vmexit(_: Event, received: &mut bool) -> bool {
        *received = !(*received);
        true
    }

    #[test]
    fn test_dispatch_table_register() {
        let mut dt = DispatchTable::<Event, bool>::new(16);
        let mut arg = true;
        dt.register(Event::VMExit, on_vmexit);
        assert!(dt.dispatch(Event::VMExit, &mut arg));
        assert_eq!(arg, false);
        assert!(!dt.dispatch(Event::PageFault, &mut arg));
        assert_eq!(arg, false);
    }

    #[test]
    fn test_dispatch_table_unregister() {
        let mut dt = DispatchTable::<Event, bool>::new(16);
        let mut arg = true;
        dt.register(Event::VMExit, on_vmexit);
        assert!(dt.dispatch(Event::VMExit, &mut arg));
        assert_eq!(arg, false);
        assert!(!dt.dispatch(Event::PageFault, &mut arg));
        dt.unregister(Event::VMExit);
        assert_eq!(arg, false);
        assert!(!dt.dispatch(Event::VMExit, &mut arg));
    }
}
<|endoftext|>


<commit_before>use std::path::{PathBuf, Path};
use std::marker::PhantomData;

use perlin_core::index::posting::{DocId};

pub type Pipeline<Out, T> = Box<Fn() -> Box<Fn(DocId, &mut T, &str) -> PhantomData<Out>>>;

/// `DocumentIndex` takes some of the basic building blocks in `perlin_core`
/// and provides an abstraction that can be used to index and query documents
/// using fields, metadata, taxonomies etc
pub struct DocumentIndex {
    // We need to overwrite perlin_core's default DocIds as some Documents might contain
    // other fields than others. This counter acts as DocumentIndex global document counter.
    doc_id_counter: DocId,
    // The base path of this index.
    base_path: PathBuf
}

impl DocumentIndex{
    /// Create a new index.
    pub fn new(path: &Path) -> Self {
        DocumentIndex {
            doc_id_counter: DocId::none(),
            base_path: path.to_path_buf(),
        }
    }

    pub fn get_next_doc_id(&mut self) -> DocId {
        self.doc_id_counter.inc();
        self.doc_id_counter
    }

    /// Commits this index
    pub fn commit(&mut self) {
    }

}

#[cfg(test)]
mod tests {
    use document::PerlinDocument;
    use document_index::Pipeline;
    use perlin_core::index::posting::DocId;
    use field::Field;

    use rust_stemmers::Algorithm;

    use test_utils::create_test_dir;

    #[derive(PerlinDocument)]
    struct Test {
        text: Field<String, Test>,
        number: Field<u64, Test>,
        emails: Field<usize, Test>,
    }

    use language::{Stemmer, LowercaseFilter, IndexerFunnel, WhitespaceTokenizer};
    use language::integers::NumberFilter;

    #[test]
    fn test() {
        use perlin_core::index::posting::Posting;
        let mut t = Test::create(
            &create_test_dir("doc_index/test"),
            Some(pipeline!(Test: text
                           WhitespaceTokenizer
                           > NumberFilter
                           | [number]
                           > LowercaseFilter
                           > Stemmer(Algorithm::English))),
            None, None);

        t.index_field(DocId(0), "text", "10 birds flew over MT EVEREST");
        t.index_field(DocId(1), "text", "125 birds flew accross THE ocean");
        t.index_field(DocId(2), "text", "2514 unicorns drew a CAR on mars");
        t.commit();
        assert_eq!(t.text.index.query_atom(&"bird".to_string()), vec![Posting(DocId(0)), Posting(DocId(1))]);
        assert_eq!(t.text.index.query_atom(&"unicorn".to_string()), vec![Posting(DocId(2))]);
        assert_eq!(t.number.index.query_atom(&2514), vec![Posting(DocId(2))]);
    }
}
<commit_msg>add field for sanity check<commit_after>use std::path::{PathBuf, Path};
use std::marker::PhantomData;

use perlin_core::index::posting::{DocId};

pub type Pipeline<Out, T> = Box<Fn() -> Box<Fn(DocId, &mut T, &str) -> PhantomData<Out>>>;

/// `DocumentIndex` takes some of the basic building blocks in `perlin_core`
/// and provides an abstraction that can be used to index and query documents
/// using fields, metadata, taxonomies etc
pub struct DocumentIndex {
    // We need to overwrite perlin_core's default DocIds as some Documents might contain
    // other fields than others. This counter acts as DocumentIndex global document counter.
    doc_id_counter: DocId,
    // The base path of this index.
    base_path: PathBuf
}

impl DocumentIndex{
    /// Create a new index.
    pub fn new(path: &Path) -> Self {
        DocumentIndex {
            doc_id_counter: DocId::none(),
            base_path: path.to_path_buf(),
        }
    }

    pub fn get_next_doc_id(&mut self) -> DocId {
        self.doc_id_counter.inc();
        self.doc_id_counter
    }

    /// Commits this index
    pub fn commit(&mut self) {
    }

}

#[cfg(test)]
mod tests {
    use document::PerlinDocument;
    use document_index::Pipeline;
    use perlin_core::index::posting::DocId;
    use field::Field;

    use rust_stemmers::Algorithm;

    use test_utils::create_test_dir;

    #[derive(PerlinDocument)]
    struct Test {
        text: Field<String, Test>,
        title: Field<String, Test>,
        number: Field<u64, Test>,
        emails: Field<usize, Test>,
    }

    use language::{Stemmer, LowercaseFilter, IndexerFunnel, WhitespaceTokenizer};
    use language::integers::NumberFilter;

    #[test]
    fn test() {
        use perlin_core::index::posting::Posting;
        let mut t = Test::create(
            &create_test_dir("doc_index/test"),
            Some(pipeline!(Test: text
                           WhitespaceTokenizer
                           > NumberFilter
                           | [number]
                           > LowercaseFilter
                           > Stemmer(Algorithm::English))),
            Some(pipeline!(Test: title
                           WhitespaceTokenizer
                           > LowercaseFilter
                           > Stemmer(Algorithm::English))),
            None, None);

        t.index_field(DocId(0), "text", "10 birds flew over MT EVEREST");
        t.index_field(DocId(1), "text", "125 birds flew accross THE ocean");
        t.index_field(DocId(2), "text", "2514 unicorns drew a CAR on mars");
        t.commit();
        assert_eq!(t.text.index.query_atom(&"bird".to_string()), vec![Posting(DocId(0)), Posting(DocId(1))]);
        assert_eq!(t.text.index.query_atom(&"unicorn".to_string()), vec![Posting(DocId(2))]);
        assert_eq!(t.number.index.query_atom(&2514), vec![Posting(DocId(2))]);
    }
}
<|endoftext|>