<commit_before>use de::read::BincodeRead;
use error::Result;
use serde;
use std::io::{Read, Write};
use std::marker::PhantomData;

pub(crate) use self::endian::BincodeByteOrder;
pub(crate) use self::int::IntEncoding;
pub(crate) use self::internal::*;
pub(crate) use self::limit::SizeLimit;
pub(crate) use self::trailing::TrailingBytes;

pub use self::endian::{BigEndian, LittleEndian, NativeEndian};
pub use self::int::{FixintEncoding, VarintEncoding};
pub use self::legacy::*;
pub use self::limit::{Bounded, Infinite};
pub use self::trailing::{AllowTrailing, RejectTrailing};

mod endian;
mod int;
mod legacy;
mod limit;
mod trailing;

/// The default options for bincode serialization/deserialization.
///
/// ### Defaults
/// By default bincode will use little-endian encoding for multi-byte integers, and will not
/// limit the number of serialized/deserialized bytes.
#[derive(Copy, Clone)]
pub struct DefaultOptions(Infinite);

impl DefaultOptions {
    /// Get a default configuration object.
    ///
    /// ### Default Configuration:
    ///
    /// | Byte limit | Endianness | Int Encoding | Trailing Behavior |
    /// |------------|------------|--------------|-------------------|
    /// | Unlimited  | Little     | Varint       | Reject            |
    pub fn new() -> DefaultOptions {
        DefaultOptions(Infinite)
    }
}

impl Default for DefaultOptions {
    fn default() -> Self {
        Self::new()
    }
}

impl InternalOptions for DefaultOptions {
    type Limit = Infinite;
    type Endian = LittleEndian;
    type IntEncoding = VarintEncoding;
    type Trailing = RejectTrailing;

    #[inline(always)]
    fn limit(&mut self) -> &mut Infinite {
        &mut self.0
    }
}

/// A configuration builder trait whose options Bincode will use
/// while serializing and deserializing.
///
/// ### Options
/// Endianness: The endianness with which multi-byte integers will be read/written.  *default: little endian*
///
/// Limit: The maximum number of bytes that will be read/written in a bincode serialize/deserialize. *default: unlimited*
///
/// Int Encoding: The encoding used for numbers, enum discriminants, and lengths. *default: varint*
///
/// Trailing Behavior: The behavior when there are trailing bytes left over in a slice after deserialization. *default: reject*
///
/// ### Byte Limit Details
/// The purpose of byte-limiting is to prevent Denial-Of-Service attacks whereby malicious attackers get bincode
/// deserialization to crash your process by allocating too much memory or keeping a connection open for too long.
///
/// When a byte limit is set, bincode will return `Err` on any deserialization that goes over the limit, or any
/// serialization that goes over the limit.
/// Sets the byte limit to be unlimited.
/// This is the default.
pub trait Options: InternalOptions + Sized {
    /// Sets the byte limit to be unlimited.
    /// This is the default.
    fn with_no_limit(self) -> WithOtherLimit<Self, Infinite> {
        WithOtherLimit::new(self, Infinite)
    }

    /// Sets the byte limit to `limit`.
    fn with_limit(self, limit: u64) -> WithOtherLimit<Self, Bounded> {
        WithOtherLimit::new(self, Bounded(limit))
    }

    /// Sets the endianness to little-endian
    /// This is the default.
    fn with_little_endian(self) -> WithOtherEndian<Self, LittleEndian> {
        WithOtherEndian::new(self)
    }

    /// Sets the endianness to big-endian
    fn with_big_endian(self) -> WithOtherEndian<Self, BigEndian> {
        WithOtherEndian::new(self)
    }

    /// Sets the endianness to the the machine-native endianness
    fn with_native_endian(self) -> WithOtherEndian<Self, NativeEndian> {
        WithOtherEndian::new(self)
    }

    /// Sets the length encoding to varint
    fn with_varint_encoding(self) -> WithOtherIntEncoding<Self, VarintEncoding> {
        WithOtherIntEncoding::new(self)
    }

    /// Sets the length encoding to be fixed
    fn with_fixint_encoding(self) -> WithOtherIntEncoding<Self, FixintEncoding> {
        WithOtherIntEncoding::new(self)
    }

    /// Sets the deserializer to reject trailing bytes
    fn reject_trailing_bytes(self) -> WithOtherTrailing<Self, RejectTrailing> {
        WithOtherTrailing::new(self)
    }

    /// Sets the deserializer to allow trailing bytes
    fn allow_trailing_bytes(self) -> WithOtherTrailing<Self, AllowTrailing> {
        WithOtherTrailing::new(self)
    }

    /// Serializes a serializable object into a `Vec` of bytes using this configuration
    #[inline(always)]
    fn serialize<S: ?Sized + serde::Serialize>(self, t: &S) -> Result<Vec<u8>> {
        ::internal::serialize(t, self)
    }

    /// Returns the size that an object would be if serialized using Bincode with this configuration
    #[inline(always)]
    fn serialized_size<T: ?Sized + serde::Serialize>(self, t: &T) -> Result<u64> {
        ::internal::serialized_size(t, self)
    }

    /// Serializes an object directly into a `Writer` using this configuration
    ///
    /// If the serialization would take more bytes than allowed by the size limit, an error
    /// is returned and *no bytes* will be written into the `Writer`
    #[inline(always)]
    fn serialize_into<W: Write, T: ?Sized + serde::Serialize>(self, w: W, t: &T) -> Result<()> {
        ::internal::serialize_into(w, t, self)
    }

    /// Deserializes a slice of bytes into an instance of `T` using this configuration
    #[inline(always)]
    fn deserialize<'a, T: serde::Deserialize<'a>>(self, bytes: &'a [u8]) -> Result<T> {
        ::internal::deserialize(bytes, self)
    }

    /// TODO: document
    #[doc(hidden)]
    #[inline(always)]
    fn deserialize_in_place<'a, R, T>(self, reader: R, place: &mut T) -> Result<()>
    where
        R: BincodeRead<'a>,
        T: serde::de::Deserialize<'a>,
    {
        ::internal::deserialize_in_place(reader, self, place)
    }

    /// Deserializes a slice of bytes with state `seed` using this configuration.
    #[inline(always)]
    fn deserialize_seed<'a, T: serde::de::DeserializeSeed<'a>>(
        self,
        seed: T,
        bytes: &'a [u8],
    ) -> Result<T::Value> {
        ::internal::deserialize_seed(seed, bytes, self)
    }

    /// Deserializes an object directly from a `Read`er using this configuration
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from<R: Read, T: serde::de::DeserializeOwned>(self, reader: R) -> Result<T> {
        ::internal::deserialize_from(reader, self)
    }

    /// Deserializes an object directly from a `Read`er with state `seed` using this configuration
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from_seed<'a, R: Read, T: serde::de::DeserializeSeed<'a>>(
        self,
        seed: T,
        reader: R,
    ) -> Result<T::Value> {
        ::internal::deserialize_from_seed(seed, reader, self)
    }

    /// Deserializes an object from a custom `BincodeRead`er using the default configuration.
    /// It is highly recommended to use `deserialize_from` unless you need to implement
    /// `BincodeRead` for performance reasons.
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from_custom<'a, R: BincodeRead<'a>, T: serde::de::DeserializeOwned>(
        self,
        reader: R,
    ) -> Result<T> {
        ::internal::deserialize_from_custom(reader, self)
    }

    /// Deserializes an object from a custom `BincodeRead`er with state `seed` using the default
    /// configuration. It is highly recommended to use `deserialize_from` unless you need to
    /// implement `BincodeRead` for performance reasons.
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from_custom_seed<'a, R: BincodeRead<'a>, T: serde::de::DeserializeSeed<'a>>(
        self,
        seed: T,
        reader: R,
    ) -> Result<T::Value> {
        ::internal::deserialize_from_custom_seed(seed, reader, self)
    }
}

impl<T: InternalOptions> Options for T {}

/// A configuration struct with a user-specified byte limit
#[derive(Clone, Copy)]
pub struct WithOtherLimit<O: Options, L: SizeLimit> {
    _options: O,
    pub(crate) new_limit: L,
}

/// A configuration struct with a user-specified endian order
#[derive(Clone, Copy)]
pub struct WithOtherEndian<O: Options, E: BincodeByteOrder> {
    options: O,
    _endian: PhantomData<E>,
}

/// A configuration struct with a user-specified length encoding
pub struct WithOtherIntEncoding<O: Options, I: IntEncoding> {
    options: O,
    _length: PhantomData<I>,
}

/// A configuration struct with a user-specified trailing bytes behavior.
pub struct WithOtherTrailing<O: Options, T: TrailingBytes> {
    options: O,
    _trailing: PhantomData<T>,
}

impl<O: Options, L: SizeLimit> WithOtherLimit<O, L> {
    #[inline(always)]
    pub(crate) fn new(options: O, limit: L) -> WithOtherLimit<O, L> {
        WithOtherLimit {
            _options: options,
            new_limit: limit,
        }
    }
}

impl<O: Options, E: BincodeByteOrder> WithOtherEndian<O, E> {
    #[inline(always)]
    pub(crate) fn new(options: O) -> WithOtherEndian<O, E> {
        WithOtherEndian {
            options,
            _endian: PhantomData,
        }
    }
}

impl<O: Options, I: IntEncoding> WithOtherIntEncoding<O, I> {
    #[inline(always)]
    pub(crate) fn new(options: O) -> WithOtherIntEncoding<O, I> {
        WithOtherIntEncoding {
            options,
            _length: PhantomData,
        }
    }
}

impl<O: Options, T: TrailingBytes> WithOtherTrailing<O, T> {
    #[inline(always)]
    pub(crate) fn new(options: O) -> WithOtherTrailing<O, T> {
        WithOtherTrailing {
            options,
            _trailing: PhantomData,
        }
    }
}

impl<O: Options, E: BincodeByteOrder + 'static> InternalOptions for WithOtherEndian<O, E> {
    type Limit = O::Limit;
    type Endian = E;
    type IntEncoding = O::IntEncoding;
    type Trailing = O::Trailing;
    #[inline(always)]
    fn limit(&mut self) -> &mut O::Limit {
        self.options.limit()
    }
}

impl<O: Options, L: SizeLimit + 'static> InternalOptions for WithOtherLimit<O, L> {
    type Limit = L;
    type Endian = O::Endian;
    type IntEncoding = O::IntEncoding;
    type Trailing = O::Trailing;
    fn limit(&mut self) -> &mut L {
        &mut self.new_limit
    }
}

impl<O: Options, I: IntEncoding + 'static> InternalOptions for WithOtherIntEncoding<O, I> {
    type Limit = O::Limit;
    type Endian = O::Endian;
    type IntEncoding = I;
    type Trailing = O::Trailing;

    fn limit(&mut self) -> &mut O::Limit {
        self.options.limit()
    }
}

impl<O: Options, T: TrailingBytes + 'static> InternalOptions for WithOtherTrailing<O, T> {
    type Limit = O::Limit;
    type Endian = O::Endian;
    type IntEncoding = O::IntEncoding;
    type Trailing = T;

    fn limit(&mut self) -> &mut O::Limit {
        self.options.limit()
    }
}

mod internal {
    use super::*;

    pub trait InternalOptions {
        type Limit: SizeLimit + 'static;
        type Endian: BincodeByteOrder + 'static;
        type IntEncoding: IntEncoding + 'static;
        type Trailing: TrailingBytes + 'static;

        fn limit(&mut self) -> &mut Self::Limit;
    }

    impl<'a, O: InternalOptions> InternalOptions for &'a mut O {
        type Limit = O::Limit;
        type Endian = O::Endian;
        type IntEncoding = O::IntEncoding;
        type Trailing = O::Trailing;

        #[inline(always)]
        fn limit(&mut self) -> &mut Self::Limit {
            (*self).limit()
        }
    }
}
<commit_msg>Fix `WithOtherTrailing` and `WithOtherIntEncoding` (#342)<commit_after>use de::read::BincodeRead;
use error::Result;
use serde;
use std::io::{Read, Write};
use std::marker::PhantomData;

pub(crate) use self::endian::BincodeByteOrder;
pub(crate) use self::int::IntEncoding;
pub(crate) use self::internal::*;
pub(crate) use self::limit::SizeLimit;
pub(crate) use self::trailing::TrailingBytes;

pub use self::endian::{BigEndian, LittleEndian, NativeEndian};
pub use self::int::{FixintEncoding, VarintEncoding};
pub use self::legacy::*;
pub use self::limit::{Bounded, Infinite};
pub use self::trailing::{AllowTrailing, RejectTrailing};

mod endian;
mod int;
mod legacy;
mod limit;
mod trailing;

/// The default options for bincode serialization/deserialization.
///
/// ### Defaults
/// By default bincode will use little-endian encoding for multi-byte integers, and will not
/// limit the number of serialized/deserialized bytes.
#[derive(Copy, Clone)]
pub struct DefaultOptions(Infinite);

impl DefaultOptions {
    /// Get a default configuration object.
    ///
    /// ### Default Configuration:
    ///
    /// | Byte limit | Endianness | Int Encoding | Trailing Behavior |
    /// |------------|------------|--------------|-------------------|
    /// | Unlimited  | Little     | Varint       | Reject            |
    pub fn new() -> DefaultOptions {
        DefaultOptions(Infinite)
    }
}

impl Default for DefaultOptions {
    fn default() -> Self {
        Self::new()
    }
}

impl InternalOptions for DefaultOptions {
    type Limit = Infinite;
    type Endian = LittleEndian;
    type IntEncoding = VarintEncoding;
    type Trailing = RejectTrailing;

    #[inline(always)]
    fn limit(&mut self) -> &mut Infinite {
        &mut self.0
    }
}

/// A configuration builder trait whose options Bincode will use
/// while serializing and deserializing.
///
/// ### Options
/// Endianness: The endianness with which multi-byte integers will be read/written.  *default: little endian*
///
/// Limit: The maximum number of bytes that will be read/written in a bincode serialize/deserialize. *default: unlimited*
///
/// Int Encoding: The encoding used for numbers, enum discriminants, and lengths. *default: varint*
///
/// Trailing Behavior: The behavior when there are trailing bytes left over in a slice after deserialization. *default: reject*
///
/// ### Byte Limit Details
/// The purpose of byte-limiting is to prevent Denial-Of-Service attacks whereby malicious attackers get bincode
/// deserialization to crash your process by allocating too much memory or keeping a connection open for too long.
///
/// When a byte limit is set, bincode will return `Err` on any deserialization that goes over the limit, or any
/// serialization that goes over the limit.
/// Sets the byte limit to be unlimited.
/// This is the default.
pub trait Options: InternalOptions + Sized {
    /// Sets the byte limit to be unlimited.
    /// This is the default.
    fn with_no_limit(self) -> WithOtherLimit<Self, Infinite> {
        WithOtherLimit::new(self, Infinite)
    }

    /// Sets the byte limit to `limit`.
    fn with_limit(self, limit: u64) -> WithOtherLimit<Self, Bounded> {
        WithOtherLimit::new(self, Bounded(limit))
    }

    /// Sets the endianness to little-endian
    /// This is the default.
    fn with_little_endian(self) -> WithOtherEndian<Self, LittleEndian> {
        WithOtherEndian::new(self)
    }

    /// Sets the endianness to big-endian
    fn with_big_endian(self) -> WithOtherEndian<Self, BigEndian> {
        WithOtherEndian::new(self)
    }

    /// Sets the endianness to the the machine-native endianness
    fn with_native_endian(self) -> WithOtherEndian<Self, NativeEndian> {
        WithOtherEndian::new(self)
    }

    /// Sets the length encoding to varint
    fn with_varint_encoding(self) -> WithOtherIntEncoding<Self, VarintEncoding> {
        WithOtherIntEncoding::new(self)
    }

    /// Sets the length encoding to be fixed
    fn with_fixint_encoding(self) -> WithOtherIntEncoding<Self, FixintEncoding> {
        WithOtherIntEncoding::new(self)
    }

    /// Sets the deserializer to reject trailing bytes
    fn reject_trailing_bytes(self) -> WithOtherTrailing<Self, RejectTrailing> {
        WithOtherTrailing::new(self)
    }

    /// Sets the deserializer to allow trailing bytes
    fn allow_trailing_bytes(self) -> WithOtherTrailing<Self, AllowTrailing> {
        WithOtherTrailing::new(self)
    }

    /// Serializes a serializable object into a `Vec` of bytes using this configuration
    #[inline(always)]
    fn serialize<S: ?Sized + serde::Serialize>(self, t: &S) -> Result<Vec<u8>> {
        ::internal::serialize(t, self)
    }

    /// Returns the size that an object would be if serialized using Bincode with this configuration
    #[inline(always)]
    fn serialized_size<T: ?Sized + serde::Serialize>(self, t: &T) -> Result<u64> {
        ::internal::serialized_size(t, self)
    }

    /// Serializes an object directly into a `Writer` using this configuration
    ///
    /// If the serialization would take more bytes than allowed by the size limit, an error
    /// is returned and *no bytes* will be written into the `Writer`
    #[inline(always)]
    fn serialize_into<W: Write, T: ?Sized + serde::Serialize>(self, w: W, t: &T) -> Result<()> {
        ::internal::serialize_into(w, t, self)
    }

    /// Deserializes a slice of bytes into an instance of `T` using this configuration
    #[inline(always)]
    fn deserialize<'a, T: serde::Deserialize<'a>>(self, bytes: &'a [u8]) -> Result<T> {
        ::internal::deserialize(bytes, self)
    }

    /// TODO: document
    #[doc(hidden)]
    #[inline(always)]
    fn deserialize_in_place<'a, R, T>(self, reader: R, place: &mut T) -> Result<()>
    where
        R: BincodeRead<'a>,
        T: serde::de::Deserialize<'a>,
    {
        ::internal::deserialize_in_place(reader, self, place)
    }

    /// Deserializes a slice of bytes with state `seed` using this configuration.
    #[inline(always)]
    fn deserialize_seed<'a, T: serde::de::DeserializeSeed<'a>>(
        self,
        seed: T,
        bytes: &'a [u8],
    ) -> Result<T::Value> {
        ::internal::deserialize_seed(seed, bytes, self)
    }

    /// Deserializes an object directly from a `Read`er using this configuration
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from<R: Read, T: serde::de::DeserializeOwned>(self, reader: R) -> Result<T> {
        ::internal::deserialize_from(reader, self)
    }

    /// Deserializes an object directly from a `Read`er with state `seed` using this configuration
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from_seed<'a, R: Read, T: serde::de::DeserializeSeed<'a>>(
        self,
        seed: T,
        reader: R,
    ) -> Result<T::Value> {
        ::internal::deserialize_from_seed(seed, reader, self)
    }

    /// Deserializes an object from a custom `BincodeRead`er using the default configuration.
    /// It is highly recommended to use `deserialize_from` unless you need to implement
    /// `BincodeRead` for performance reasons.
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from_custom<'a, R: BincodeRead<'a>, T: serde::de::DeserializeOwned>(
        self,
        reader: R,
    ) -> Result<T> {
        ::internal::deserialize_from_custom(reader, self)
    }

    /// Deserializes an object from a custom `BincodeRead`er with state `seed` using the default
    /// configuration. It is highly recommended to use `deserialize_from` unless you need to
    /// implement `BincodeRead` for performance reasons.
    ///
    /// If this returns an `Error`, `reader` may be in an invalid state.
    #[inline(always)]
    fn deserialize_from_custom_seed<'a, R: BincodeRead<'a>, T: serde::de::DeserializeSeed<'a>>(
        self,
        seed: T,
        reader: R,
    ) -> Result<T::Value> {
        ::internal::deserialize_from_custom_seed(seed, reader, self)
    }
}

impl<T: InternalOptions> Options for T {}

/// A configuration struct with a user-specified byte limit
#[derive(Clone, Copy)]
pub struct WithOtherLimit<O: Options, L: SizeLimit> {
    _options: O,
    pub(crate) new_limit: L,
}

/// A configuration struct with a user-specified endian order
#[derive(Clone, Copy)]
pub struct WithOtherEndian<O: Options, E: BincodeByteOrder> {
    options: O,
    _endian: PhantomData<E>,
}

/// A configuration struct with a user-specified length encoding
#[derive(Clone, Copy)]
pub struct WithOtherIntEncoding<O: Options, I: IntEncoding> {
    options: O,
    _length: PhantomData<I>,
}

/// A configuration struct with a user-specified trailing bytes behavior.
#[derive(Clone, Copy)]
pub struct WithOtherTrailing<O: Options, T: TrailingBytes> {
    options: O,
    _trailing: PhantomData<T>,
}

impl<O: Options, L: SizeLimit> WithOtherLimit<O, L> {
    #[inline(always)]
    pub(crate) fn new(options: O, limit: L) -> WithOtherLimit<O, L> {
        WithOtherLimit {
            _options: options,
            new_limit: limit,
        }
    }
}

impl<O: Options, E: BincodeByteOrder> WithOtherEndian<O, E> {
    #[inline(always)]
    pub(crate) fn new(options: O) -> WithOtherEndian<O, E> {
        WithOtherEndian {
            options,
            _endian: PhantomData,
        }
    }
}

impl<O: Options, I: IntEncoding> WithOtherIntEncoding<O, I> {
    #[inline(always)]
    pub(crate) fn new(options: O) -> WithOtherIntEncoding<O, I> {
        WithOtherIntEncoding {
            options,
            _length: PhantomData,
        }
    }
}

impl<O: Options, T: TrailingBytes> WithOtherTrailing<O, T> {
    #[inline(always)]
    pub(crate) fn new(options: O) -> WithOtherTrailing<O, T> {
        WithOtherTrailing {
            options,
            _trailing: PhantomData,
        }
    }
}

impl<O: Options, E: BincodeByteOrder + 'static> InternalOptions for WithOtherEndian<O, E> {
    type Limit = O::Limit;
    type Endian = E;
    type IntEncoding = O::IntEncoding;
    type Trailing = O::Trailing;
    #[inline(always)]
    fn limit(&mut self) -> &mut O::Limit {
        self.options.limit()
    }
}

impl<O: Options, L: SizeLimit + 'static> InternalOptions for WithOtherLimit<O, L> {
    type Limit = L;
    type Endian = O::Endian;
    type IntEncoding = O::IntEncoding;
    type Trailing = O::Trailing;
    fn limit(&mut self) -> &mut L {
        &mut self.new_limit
    }
}

impl<O: Options, I: IntEncoding + 'static> InternalOptions for WithOtherIntEncoding<O, I> {
    type Limit = O::Limit;
    type Endian = O::Endian;
    type IntEncoding = I;
    type Trailing = O::Trailing;

    fn limit(&mut self) -> &mut O::Limit {
        self.options.limit()
    }
}

impl<O: Options, T: TrailingBytes + 'static> InternalOptions for WithOtherTrailing<O, T> {
    type Limit = O::Limit;
    type Endian = O::Endian;
    type IntEncoding = O::IntEncoding;
    type Trailing = T;

    fn limit(&mut self) -> &mut O::Limit {
        self.options.limit()
    }
}

mod internal {
    use super::*;

    pub trait InternalOptions {
        type Limit: SizeLimit + 'static;
        type Endian: BincodeByteOrder + 'static;
        type IntEncoding: IntEncoding + 'static;
        type Trailing: TrailingBytes + 'static;

        fn limit(&mut self) -> &mut Self::Limit;
    }

    impl<'a, O: InternalOptions> InternalOptions for &'a mut O {
        type Limit = O::Limit;
        type Endian = O::Endian;
        type IntEncoding = O::IntEncoding;
        type Trailing = O::Trailing;

        #[inline(always)]
        fn limit(&mut self) -> &mut Self::Limit {
            (*self).limit()
        }
    }
}
<|endoftext|>


<commit_before>use crate::task::{waker_ref, ArcWake};
use futures_core::future::{FusedFuture, Future};
use futures_core::task::{Context, Poll, Waker};
use slab::Slab;
use std::cell::UnsafeCell;
use std::fmt;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{Acquire, SeqCst};
use std::sync::{Arc, Mutex, Weak};

/// Future for the [`shared`](super::FutureExt::shared) method.
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Shared<Fut: Future> {
    inner: Option<Arc<Inner<Fut>>>,
    waker_key: usize,
}

struct Inner<Fut: Future> {
    future_or_output: UnsafeCell<FutureOrOutput<Fut>>,
    notifier: Arc<Notifier>,
}

struct Notifier {
    state: AtomicUsize,
    wakers: Mutex<Option<Slab<Option<Waker>>>>,
}

/// A weak reference to a [`Shared`] that can be upgraded much like an `Arc`.
pub struct WeakShared<Fut: Future>(Weak<Inner<Fut>>);

// The future itself is polled behind the `Arc`, so it won't be moved
// when `Shared` is moved.
impl<Fut: Future> Unpin for Shared<Fut> {}

impl<Fut: Future> fmt::Debug for Shared<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Shared")
            .field("inner", &self.inner)
            .field("waker_key", &self.waker_key)
            .finish()
    }
}

impl<Fut: Future> fmt::Debug for Inner<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Inner").finish()
    }
}

impl<Fut: Future> fmt::Debug for WeakShared<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("WeakShared").finish()
    }
}

enum FutureOrOutput<Fut: Future> {
    Future(Fut),
    Output(Fut::Output),
}

unsafe impl<Fut> Send for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{
}

unsafe impl<Fut> Sync for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{
}

const IDLE: usize = 0;
const POLLING: usize = 1;
const COMPLETE: usize = 2;
const POISONED: usize = 3;

const NULL_WAKER_KEY: usize = usize::max_value();

impl<Fut: Future> Shared<Fut> {
    pub(super) fn new(future: Fut) -> Self {
        let inner = Inner {
            future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)),
            notifier: Arc::new(Notifier {
                state: AtomicUsize::new(IDLE),
                wakers: Mutex::new(Some(Slab::new())),
            }),
        };

        Self {
            inner: Some(Arc::new(inner)),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    /// Returns [`Some`] containing a reference to this [`Shared`]'s output if
    /// it has already been computed by a clone or [`None`] if it hasn't been
    /// computed yet or this [`Shared`] already returned its output from
    /// [`poll`](Future::poll).
    pub fn peek(&self) -> Option<&Fut::Output> {
        if let Some(inner) = self.inner.as_ref() {
            match inner.notifier.state.load(SeqCst) {
                COMPLETE => unsafe { return Some(inner.output()) },
                POISONED => panic!("inner future panicked during poll"),
                _ => {}
            }
        }
        None
    }

    /// Creates a new [`WeakShared`] for this [`Shared`].
    ///
    /// Returns [`None`] if it has already been polled to completion.
    pub fn downgrade(&self) -> Option<WeakShared<Fut>> {
        if let Some(inner) = self.inner.as_ref() {
            return Some(WeakShared(Arc::downgrade(inner)));
        }
        None
    }
}

impl<Fut> Inner<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn output(&self) -> &Fut::Output {
        match &*self.future_or_output.get() {
            FutureOrOutput::Output(ref item) => item,
            FutureOrOutput::Future(_) => unreachable!(),
        }
    }
    /// Registers the current task to receive a wakeup when we are awoken.
    fn record_waker(&self, waker_key: &mut usize, cx: &mut Context<'_>) {
        let mut wakers_guard = self.notifier.wakers.lock().unwrap();

        let wakers = match wakers_guard.as_mut() {
            Some(wakers) => wakers,
            None => return,
        };

        let new_waker = cx.waker();

        if *waker_key == NULL_WAKER_KEY {
            *waker_key = wakers.insert(Some(new_waker.clone()));
        } else {
            match wakers[*waker_key] {
                Some(ref old_waker) if new_waker.will_wake(old_waker) => {}
                // Could use clone_from here, but Waker doesn't specialize it.
                ref mut slot => *slot = Some(new_waker.clone()),
            }
        }
        debug_assert!(*waker_key != NULL_WAKER_KEY);
    }

    /// Safety: callers must first ensure that `inner.state`
    /// is `COMPLETE`
    unsafe fn take_or_clone_output(self: Arc<Self>) -> Fut::Output {
        match Arc::try_unwrap(self) {
            Ok(inner) => match inner.future_or_output.into_inner() {
                FutureOrOutput::Output(item) => item,
                FutureOrOutput::Future(_) => unreachable!(),
            },
            Err(inner) => inner.output().clone(),
        }
    }
}

impl<Fut> FusedFuture for Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    fn is_terminated(&self) -> bool {
        self.inner.is_none()
    }
}

impl<Fut> Future for Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    type Output = Fut::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = &mut *self;

        let inner = this
            .inner
            .take()
            .expect("Shared future polled again after completion");

        // Fast path for when the wrapped future has already completed
        if inner.notifier.state.load(Acquire) == COMPLETE {
            // Safety: We're in the COMPLETE state
            return unsafe { Poll::Ready(inner.take_or_clone_output()) };
        }

        inner.record_waker(&mut this.waker_key, cx);

        match inner
            .notifier
            .state
            .compare_exchange(IDLE, POLLING, SeqCst, SeqCst)
            .unwrap_or_else(|x| x)
        {
            IDLE => {
                // Lock acquired, fall through
            }
            POLLING => {
                // Another task is currently polling, at this point we just want
                // to ensure that the waker for this task is registered
                this.inner = Some(inner);
                return Poll::Pending;
            }
            COMPLETE => {
                // Safety: We're in the COMPLETE state
                return unsafe { Poll::Ready(inner.take_or_clone_output()) };
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => unreachable!(),
        }

        let waker = waker_ref(&inner.notifier);
        let mut cx = Context::from_waker(&waker);

        struct Reset<'a>(&'a AtomicUsize);

        impl Drop for Reset<'_> {
            fn drop(&mut self) {
                use std::thread;

                if thread::panicking() {
                    self.0.store(POISONED, SeqCst);
                }
            }
        }

        let _reset = Reset(&inner.notifier.state);

        let output = {
            let future = unsafe {
                match &mut *inner.future_or_output.get() {
                    FutureOrOutput::Future(fut) => Pin::new_unchecked(fut),
                    _ => unreachable!(),
                }
            };

            match future.poll(&mut cx) {
                Poll::Pending => {
                    if inner
                        .notifier
                        .state
                        .compare_exchange(POLLING, IDLE, SeqCst, SeqCst)
                        .is_ok()
                    {
                        // Success
                        drop(_reset);
                        this.inner = Some(inner);
                        return Poll::Pending;
                    } else {
                        unreachable!()
                    }
                }
                Poll::Ready(output) => output,
            }
        };

        unsafe {
            *inner.future_or_output.get() = FutureOrOutput::Output(output);
        }

        inner.notifier.state.store(COMPLETE, SeqCst);

        // Wake all tasks and drop the slab
        let mut wakers_guard = inner.notifier.wakers.lock().unwrap();
        let mut wakers = wakers_guard.take().unwrap();
        for opt_waker in wakers.drain() {
            if let Some(waker) = opt_waker {
                waker.wake();
            }
        }

        drop(_reset); // Make borrow checker happy
        drop(wakers_guard);

        // Safety: We're in the COMPLETE state
        unsafe { Poll::Ready(inner.take_or_clone_output()) }
    }
}

impl<Fut> Clone for Shared<Fut>
where
    Fut: Future,
{
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Drop for Shared<Fut>
where
    Fut: Future,
{
    fn drop(&mut self) {
        if self.waker_key != NULL_WAKER_KEY {
            if let Some(ref inner) = self.inner {
                if let Ok(mut wakers) = inner.notifier.wakers.lock() {
                    if let Some(wakers) = wakers.as_mut() {
                        wakers.remove(self.waker_key);
                    }
                }
            }
        }
    }
}

impl ArcWake for Notifier {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        let wakers = &mut *arc_self.wakers.lock().unwrap();
        if let Some(wakers) = wakers.as_mut() {
            for (_key, opt_waker) in wakers {
                if let Some(waker) = opt_waker.take() {
                    waker.wake();
                }
            }
        }
    }
}

impl<Fut: Future> WeakShared<Fut>
{
    /// Attempts to upgrade this [`WeakShared`] into a [`Shared`].
    ///
    /// Returns [`None`] if all clones of the [`Shared`] have been dropped or polled
    /// to completion.
    pub fn upgrade(&self) -> Option<Shared<Fut>> {
        Some(Shared {
            inner: Some(self.0.upgrade()?),
            waker_key: NULL_WAKER_KEY,
        })
    }
}
<commit_msg>Add strong/weak pointer counts to Shared (#2346)<commit_after>use crate::task::{waker_ref, ArcWake};
use futures_core::future::{FusedFuture, Future};
use futures_core::task::{Context, Poll, Waker};
use slab::Slab;
use std::cell::UnsafeCell;
use std::fmt;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{Acquire, SeqCst};
use std::sync::{Arc, Mutex, Weak};

/// Future for the [`shared`](super::FutureExt::shared) method.
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Shared<Fut: Future> {
    inner: Option<Arc<Inner<Fut>>>,
    waker_key: usize,
}

struct Inner<Fut: Future> {
    future_or_output: UnsafeCell<FutureOrOutput<Fut>>,
    notifier: Arc<Notifier>,
}

struct Notifier {
    state: AtomicUsize,
    wakers: Mutex<Option<Slab<Option<Waker>>>>,
}

/// A weak reference to a [`Shared`] that can be upgraded much like an `Arc`.
pub struct WeakShared<Fut: Future>(Weak<Inner<Fut>>);

// The future itself is polled behind the `Arc`, so it won't be moved
// when `Shared` is moved.
impl<Fut: Future> Unpin for Shared<Fut> {}

impl<Fut: Future> fmt::Debug for Shared<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Shared")
            .field("inner", &self.inner)
            .field("waker_key", &self.waker_key)
            .finish()
    }
}

impl<Fut: Future> fmt::Debug for Inner<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Inner").finish()
    }
}

impl<Fut: Future> fmt::Debug for WeakShared<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("WeakShared").finish()
    }
}

enum FutureOrOutput<Fut: Future> {
    Future(Fut),
    Output(Fut::Output),
}

unsafe impl<Fut> Send for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{
}

unsafe impl<Fut> Sync for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{
}

const IDLE: usize = 0;
const POLLING: usize = 1;
const COMPLETE: usize = 2;
const POISONED: usize = 3;

const NULL_WAKER_KEY: usize = usize::max_value();

impl<Fut: Future> Shared<Fut> {
    pub(super) fn new(future: Fut) -> Self {
        let inner = Inner {
            future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)),
            notifier: Arc::new(Notifier {
                state: AtomicUsize::new(IDLE),
                wakers: Mutex::new(Some(Slab::new())),
            }),
        };

        Self {
            inner: Some(Arc::new(inner)),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    /// Returns [`Some`] containing a reference to this [`Shared`]'s output if
    /// it has already been computed by a clone or [`None`] if it hasn't been
    /// computed yet or this [`Shared`] already returned its output from
    /// [`poll`](Future::poll).
    pub fn peek(&self) -> Option<&Fut::Output> {
        if let Some(inner) = self.inner.as_ref() {
            match inner.notifier.state.load(SeqCst) {
                COMPLETE => unsafe { return Some(inner.output()) },
                POISONED => panic!("inner future panicked during poll"),
                _ => {}
            }
        }
        None
    }

    /// Creates a new [`WeakShared`] for this [`Shared`].
    ///
    /// Returns [`None`] if it has already been polled to completion.
    pub fn downgrade(&self) -> Option<WeakShared<Fut>> {
        if let Some(inner) = self.inner.as_ref() {
            return Some(WeakShared(Arc::downgrade(inner)));
        }
        None
    }

    /// Gets the number of strong pointers to this allocation.
    ///
    /// Returns [`None`] if it has already been polled to completion.
    ///
    /// # Safety
    ///
    /// This method by itself is safe, but using it correctly requires extra care. Another thread
    /// can change the strong count at any time, including potentially between calling this method
    /// and acting on the result.
    pub fn strong_count(&self) -> Option<usize> {
        self.inner.as_ref().map(|arc| Arc::strong_count(arc))
    }

    /// Gets the number of weak pointers to this allocation.
    ///
    /// Returns [`None`] if it has already been polled to completion.
    ///
    /// # Safety
    ///
    /// This method by itself is safe, but using it correctly requires extra care. Another thread
    /// can change the weak count at any time, including potentially between calling this method
    /// and acting on the result.
    pub fn weak_count(&self) -> Option<usize> {
        self.inner.as_ref().map(|arc| Arc::weak_count(arc))
    }
}

impl<Fut> Inner<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn output(&self) -> &Fut::Output {
        match &*self.future_or_output.get() {
            FutureOrOutput::Output(ref item) => item,
            FutureOrOutput::Future(_) => unreachable!(),
        }
    }
    /// Registers the current task to receive a wakeup when we are awoken.
    fn record_waker(&self, waker_key: &mut usize, cx: &mut Context<'_>) {
        let mut wakers_guard = self.notifier.wakers.lock().unwrap();

        let wakers = match wakers_guard.as_mut() {
            Some(wakers) => wakers,
            None => return,
        };

        let new_waker = cx.waker();

        if *waker_key == NULL_WAKER_KEY {
            *waker_key = wakers.insert(Some(new_waker.clone()));
        } else {
            match wakers[*waker_key] {
                Some(ref old_waker) if new_waker.will_wake(old_waker) => {}
                // Could use clone_from here, but Waker doesn't specialize it.
                ref mut slot => *slot = Some(new_waker.clone()),
            }
        }
        debug_assert!(*waker_key != NULL_WAKER_KEY);
    }

    /// Safety: callers must first ensure that `inner.state`
    /// is `COMPLETE`
    unsafe fn take_or_clone_output(self: Arc<Self>) -> Fut::Output {
        match Arc::try_unwrap(self) {
            Ok(inner) => match inner.future_or_output.into_inner() {
                FutureOrOutput::Output(item) => item,
                FutureOrOutput::Future(_) => unreachable!(),
            },
            Err(inner) => inner.output().clone(),
        }
    }
}

impl<Fut> FusedFuture for Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    fn is_terminated(&self) -> bool {
        self.inner.is_none()
    }
}

impl<Fut> Future for Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    type Output = Fut::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = &mut *self;

        let inner = this
            .inner
            .take()
            .expect("Shared future polled again after completion");

        // Fast path for when the wrapped future has already completed
        if inner.notifier.state.load(Acquire) == COMPLETE {
            // Safety: We're in the COMPLETE state
            return unsafe { Poll::Ready(inner.take_or_clone_output()) };
        }

        inner.record_waker(&mut this.waker_key, cx);

        match inner
            .notifier
            .state
            .compare_exchange(IDLE, POLLING, SeqCst, SeqCst)
            .unwrap_or_else(|x| x)
        {
            IDLE => {
                // Lock acquired, fall through
            }
            POLLING => {
                // Another task is currently polling, at this point we just want
                // to ensure that the waker for this task is registered
                this.inner = Some(inner);
                return Poll::Pending;
            }
            COMPLETE => {
                // Safety: We're in the COMPLETE state
                return unsafe { Poll::Ready(inner.take_or_clone_output()) };
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => unreachable!(),
        }

        let waker = waker_ref(&inner.notifier);
        let mut cx = Context::from_waker(&waker);

        struct Reset<'a>(&'a AtomicUsize);

        impl Drop for Reset<'_> {
            fn drop(&mut self) {
                use std::thread;

                if thread::panicking() {
                    self.0.store(POISONED, SeqCst);
                }
            }
        }

        let _reset = Reset(&inner.notifier.state);

        let output = {
            let future = unsafe {
                match &mut *inner.future_or_output.get() {
                    FutureOrOutput::Future(fut) => Pin::new_unchecked(fut),
                    _ => unreachable!(),
                }
            };

            match future.poll(&mut cx) {
                Poll::Pending => {
                    if inner
                        .notifier
                        .state
                        .compare_exchange(POLLING, IDLE, SeqCst, SeqCst)
                        .is_ok()
                    {
                        // Success
                        drop(_reset);
                        this.inner = Some(inner);
                        return Poll::Pending;
                    } else {
                        unreachable!()
                    }
                }
                Poll::Ready(output) => output,
            }
        };

        unsafe {
            *inner.future_or_output.get() = FutureOrOutput::Output(output);
        }

        inner.notifier.state.store(COMPLETE, SeqCst);

        // Wake all tasks and drop the slab
        let mut wakers_guard = inner.notifier.wakers.lock().unwrap();
        let mut wakers = wakers_guard.take().unwrap();
        for opt_waker in wakers.drain() {
            if let Some(waker) = opt_waker {
                waker.wake();
            }
        }

        drop(_reset); // Make borrow checker happy
        drop(wakers_guard);

        // Safety: We're in the COMPLETE state
        unsafe { Poll::Ready(inner.take_or_clone_output()) }
    }
}

impl<Fut> Clone for Shared<Fut>
where
    Fut: Future,
{
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Drop for Shared<Fut>
where
    Fut: Future,
{
    fn drop(&mut self) {
        if self.waker_key != NULL_WAKER_KEY {
            if let Some(ref inner) = self.inner {
                if let Ok(mut wakers) = inner.notifier.wakers.lock() {
                    if let Some(wakers) = wakers.as_mut() {
                        wakers.remove(self.waker_key);
                    }
                }
            }
        }
    }
}

impl ArcWake for Notifier {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        let wakers = &mut *arc_self.wakers.lock().unwrap();
        if let Some(wakers) = wakers.as_mut() {
            for (_key, opt_waker) in wakers {
                if let Some(waker) = opt_waker.take() {
                    waker.wake();
                }
            }
        }
    }
}

impl<Fut: Future> WeakShared<Fut>
{
    /// Attempts to upgrade this [`WeakShared`] into a [`Shared`].
    ///
    /// Returns [`None`] if all clones of the [`Shared`] have been dropped or polled
    /// to completion.
    pub fn upgrade(&self) -> Option<Shared<Fut>> {
        Some(Shared {
            inner: Some(self.0.upgrade()?),
            waker_key: NULL_WAKER_KEY,
        })
    }
}
<|endoftext|>


<commit_before>
use rusoto_ecs;

#[derive(Debug, Serialize, Deserialize)]
pub struct Service {
    pub name: String,
    pub desired_count: i64,
    pub deployment_configuration: Option<DeploymentConfiguration>,
    pub load_balancers: Option<LoadBalancers>,
    pub task_definition: TaskDefinition,
    pub role: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskDefinition {
    pub family: String,
    pub container_definitions: ContainerDefinitions,
    pub task_role_arn: Option<String>,
    pub network_mode: Option<NetworkMode>,
    pub volumes: Option<Vec<Volume>>,
}

pub type NetworkMode = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Volume {
    pub host: Option<HostVolumeProperties>,
    pub name: Option<String>,
}
impl Volume {
    pub fn to_rusoto(&self) -> rusoto_ecs::Volume {
        rusoto_ecs::Volume {
            host: self.host.as_ref().map(|e| e.to_rusoto()),
            name: self.name.to_owned(),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct HostVolumeProperties {
    pub source_path: Option<String>,
}

impl HostVolumeProperties {
    pub fn to_rusoto(&self) -> rusoto_ecs::HostVolumeProperties {
        rusoto_ecs::HostVolumeProperties {
            source_path: self.source_path.to_owned(),
        }
    }
}

// rusoto compatible structs

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct DeploymentConfiguration {
    pub maximum_percent: Option<i64>,
    pub minimum_healthy_percent: Option<i64>,
}
impl DeploymentConfiguration {
    pub fn to_rusoto(&self) -> rusoto_ecs::DeploymentConfiguration {
        rusoto_ecs::DeploymentConfiguration {
            maximum_percent: self.maximum_percent,
            minimum_healthy_percent: self.minimum_healthy_percent,
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct LoadBalancer {
    pub container_name: Option<String>,
    pub container_port: Option<i64>,
    pub load_balancer_name: Option<String>,
    pub target_group_arn: Option<String>,
}

impl LoadBalancer {
    pub fn to_rusoto(&self) -> rusoto_ecs::LoadBalancer {
        rusoto_ecs::LoadBalancer {
            container_name: self.container_name.to_owned(),
            container_port: self.container_port,
            load_balancer_name: self.load_balancer_name.to_owned(),
            target_group_arn: self.target_group_arn.to_owned(),
        }
    }
}

pub type LoadBalancers = Vec<LoadBalancer>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct ContainerDefinition {
    pub command: Option<StringList>,
    pub cpu: Option<Integer>,
    pub disable_networking: Option<BoxedBoolean>,
    pub dns_search_domains: Option<StringList>,
    pub dns_servers: Option<StringList>,
    pub docker_labels: Option<DockerLabelsMap>,
    pub docker_security_options: Option<StringList>,
    pub entry_point: Option<StringList>,
    pub environment: Option<EnvironmentVariables>,
    pub essential: Option<BoxedBoolean>,
    pub extra_hosts: Option<HostEntryList>,
    pub hostname: Option<String>,
    pub image: Option<String>,
    pub links: Option<StringList>,
    pub log_configuration: Option<LogConfiguration>,
    pub memory: Option<BoxedInteger>,
    pub memory_reservation: Option<BoxedInteger>,
    pub mount_points: Option<MountPointList>,
    pub name: Option<String>,
    pub port_mappings: Option<PortMappingList>,
    pub privileged: Option<BoxedBoolean>,
    pub readonly_root_filesystem: Option<BoxedBoolean>,
    pub ulimits: Option<UlimitList>,
    pub user: Option<String>,
    pub volumes_from: Option<VolumeFromList>,
    pub working_directory: Option<String>,
}
impl ContainerDefinition {
    pub fn to_rusoto(&self) -> rusoto_ecs::ContainerDefinition {
        rusoto_ecs::ContainerDefinition {
            command: self.command.to_owned(),
            cpu: self.cpu,
            disable_networking: self.disable_networking,
            dns_search_domains: self.dns_search_domains.to_owned(),
            dns_servers: self.dns_servers.to_owned(),
            docker_labels: self.docker_labels.to_owned(),
            docker_security_options: self.docker_security_options.to_owned(),
            entry_point: self.entry_point.to_owned(),
            environment: self.environment
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            essential: self.essential,
            extra_hosts: self.extra_hosts
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            hostname: self.hostname.to_owned(),
            image: self.image.to_owned(),
            links: self.links.to_owned(),
            log_configuration: self.log_configuration.as_ref().map(|e| e.to_rusoto()),
            memory: self.memory,
            memory_reservation: self.memory_reservation,
            mount_points: self.mount_points
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            name: self.name.to_owned(),
            port_mappings: self.port_mappings
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            privileged: self.privileged,
            readonly_root_filesystem: self.readonly_root_filesystem,
            ulimits: self.ulimits
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            user: self.user.to_owned(),
            volumes_from: self.volumes_from
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            working_directory: self.working_directory.to_owned(),
        }
    }
}

pub type ContainerDefinitions = Vec<ContainerDefinition>;

pub type BoxedBoolean = bool;
pub type BoxedInteger = i64;
pub type Integer = i64;
pub type StringList = Vec<String>;

pub type DockerLabelsMap = ::std::collections::HashMap<String, String>;
pub type EnvironmentVariables = Vec<KeyValuePair>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct KeyValuePair {
    pub name: Option<String>,
    pub value: Option<String>,
}
impl KeyValuePair {
    pub fn to_rusoto(&self) -> rusoto_ecs::KeyValuePair {
        rusoto_ecs::KeyValuePair {
            name: self.name.to_owned(),
            value: self.value.to_owned(),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct HostEntry {
    pub hostname: String,
    pub ip_address: String,
}
impl HostEntry {
    pub fn to_rusoto(&self) -> rusoto_ecs::HostEntry {
        rusoto_ecs::HostEntry {
            hostname: self.hostname.to_owned(),
            ip_address: self.ip_address.to_owned(),
        }
    }
}

pub type HostEntryList = Vec<HostEntry>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct MountPoint {
    pub container_path: Option<String>,
    pub read_only: Option<BoxedBoolean>,
    pub source_volume: Option<String>,
}
impl MountPoint {
    pub fn to_rusoto(&self) -> rusoto_ecs::MountPoint {
        rusoto_ecs::MountPoint {
            container_path: self.container_path.to_owned(),
            read_only: self.read_only,
            source_volume: self.source_volume.to_owned(),
        }
    }
}

pub type MountPointList = Vec<MountPoint>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct PortMapping {
    pub container_port: Option<BoxedInteger>,
    pub host_port: Option<BoxedInteger>,
    pub protocol: Option<TransportProtocol>,
}
impl PortMapping {
    pub fn to_rusoto(&self) -> rusoto_ecs::PortMapping {
        rusoto_ecs::PortMapping {
            container_port: self.container_port,
            host_port: self.host_port,
            protocol: self.protocol.to_owned(),
        }
    }
}

pub type PortMappingList = Vec<PortMapping>;
pub type TransportProtocol = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Ulimit {
    pub hard_limit: Integer,
    pub name: UlimitName,
    pub soft_limit: Integer,
}
impl Ulimit {
    pub fn to_rusoto(&self) -> rusoto_ecs::Ulimit {
        rusoto_ecs::Ulimit {
            hard_limit: self.hard_limit,
            name: self.name.to_owned(),
            soft_limit: self.soft_limit,
        }
    }
}

pub type UlimitList = Vec<Ulimit>;
pub type UlimitName = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct VolumeFrom {
    pub read_only: Option<BoxedBoolean>,
    pub source_container: Option<String>,
}
impl VolumeFrom {
    pub fn to_rusoto(&self) -> rusoto_ecs::VolumeFrom {
        rusoto_ecs::VolumeFrom {
            read_only: self.read_only,
            source_container: self.source_container.to_owned(),
        }
    }
}

pub type VolumeFromList = Vec<VolumeFrom>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct LogConfiguration {
    pub log_driver: LogDriver,
    pub options: Option<LogConfigurationOptionsMap>,
}
impl LogConfiguration {
    pub fn to_rusoto(&self) -> rusoto_ecs::LogConfiguration {
        rusoto_ecs::LogConfiguration {
            log_driver: self.log_driver.to_owned(),
            options: self.options.to_owned(),
        }
    }
}

pub type LogConfigurationOptionsMap = ::std::collections::HashMap<String, String>;
pub type LogDriver = String;
<commit_msg>Support `health_check` and `linux_parameter` configuration in container definition<commit_after>use rusoto_ecs;

#[derive(Debug, Serialize, Deserialize)]
pub struct Service {
    pub name: String,
    pub desired_count: i64,
    pub deployment_configuration: Option<DeploymentConfiguration>,
    pub load_balancers: Option<LoadBalancers>,
    pub task_definition: TaskDefinition,
    pub role: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskDefinition {
    pub family: String,
    pub container_definitions: ContainerDefinitions,
    pub task_role_arn: Option<String>,
    pub network_mode: Option<NetworkMode>,
    pub volumes: Option<Vec<Volume>>,
}

pub type NetworkMode = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Volume {
    pub host: Option<HostVolumeProperties>,
    pub name: Option<String>,
}
impl Volume {
    pub fn to_rusoto(&self) -> rusoto_ecs::Volume {
        rusoto_ecs::Volume {
            host: self.host.as_ref().map(|e| e.to_rusoto()),
            name: self.name.to_owned(),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct HostVolumeProperties {
    pub source_path: Option<String>,
}

impl HostVolumeProperties {
    pub fn to_rusoto(&self) -> rusoto_ecs::HostVolumeProperties {
        rusoto_ecs::HostVolumeProperties {
            source_path: self.source_path.to_owned(),
        }
    }
}

// rusoto compatible structs

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct DeploymentConfiguration {
    pub maximum_percent: Option<i64>,
    pub minimum_healthy_percent: Option<i64>,
}
impl DeploymentConfiguration {
    pub fn to_rusoto(&self) -> rusoto_ecs::DeploymentConfiguration {
        rusoto_ecs::DeploymentConfiguration {
            maximum_percent: self.maximum_percent,
            minimum_healthy_percent: self.minimum_healthy_percent,
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct LoadBalancer {
    pub container_name: Option<String>,
    pub container_port: Option<i64>,
    pub load_balancer_name: Option<String>,
    pub target_group_arn: Option<String>,
}

impl LoadBalancer {
    pub fn to_rusoto(&self) -> rusoto_ecs::LoadBalancer {
        rusoto_ecs::LoadBalancer {
            container_name: self.container_name.to_owned(),
            container_port: self.container_port,
            load_balancer_name: self.load_balancer_name.to_owned(),
            target_group_arn: self.target_group_arn.to_owned(),
        }
    }
}

pub type LoadBalancers = Vec<LoadBalancer>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct ContainerDefinition {
    pub command: Option<StringList>,
    pub cpu: Option<Integer>,
    pub disable_networking: Option<BoxedBoolean>,
    pub dns_search_domains: Option<StringList>,
    pub dns_servers: Option<StringList>,
    pub docker_labels: Option<DockerLabelsMap>,
    pub docker_security_options: Option<StringList>,
    pub entry_point: Option<StringList>,
    pub environment: Option<EnvironmentVariables>,
    pub essential: Option<BoxedBoolean>,
    pub extra_hosts: Option<HostEntryList>,
    pub hostname: Option<String>,
    pub image: Option<String>,
    pub links: Option<StringList>,
    pub log_configuration: Option<LogConfiguration>,
    pub memory: Option<BoxedInteger>,
    pub memory_reservation: Option<BoxedInteger>,
    pub mount_points: Option<MountPointList>,
    pub name: Option<String>,
    pub port_mappings: Option<PortMappingList>,
    pub privileged: Option<BoxedBoolean>,
    pub readonly_root_filesystem: Option<BoxedBoolean>,
    pub ulimits: Option<UlimitList>,
    pub user: Option<String>,
    pub volumes_from: Option<VolumeFromList>,
    pub working_directory: Option<String>,
    pub health_check: Option<HealthCheck>,
    pub linux_parameters: Option<LinuxParameters>,
}
impl ContainerDefinition {
    pub fn to_rusoto(&self) -> rusoto_ecs::ContainerDefinition {
        rusoto_ecs::ContainerDefinition {
            command: self.command.to_owned(),
            cpu: self.cpu,
            disable_networking: self.disable_networking,
            dns_search_domains: self.dns_search_domains.to_owned(),
            dns_servers: self.dns_servers.to_owned(),
            docker_labels: self.docker_labels.to_owned(),
            docker_security_options: self.docker_security_options.to_owned(),
            entry_point: self.entry_point.to_owned(),
            environment: self
                .environment
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            essential: self.essential,
            extra_hosts: self
                .extra_hosts
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            hostname: self.hostname.to_owned(),
            image: self.image.to_owned(),
            links: self.links.to_owned(),
            log_configuration: self.log_configuration.as_ref().map(|e| e.to_rusoto()),
            memory: self.memory,
            memory_reservation: self.memory_reservation,
            mount_points: self
                .mount_points
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            name: self.name.to_owned(),
            port_mappings: self
                .port_mappings
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            privileged: self.privileged,
            readonly_root_filesystem: self.readonly_root_filesystem,
            ulimits: self
                .ulimits
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            user: self.user.to_owned(),
            volumes_from: self
                .volumes_from
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            working_directory: self.working_directory.to_owned(),
            health_check: self.health_check.as_ref().map(|e| e.to_rusoto()),
            linux_parameters: self.linux_parameters.as_ref().map(|e| e.to_rusoto()),
        }
    }
}

pub type ContainerDefinitions = Vec<ContainerDefinition>;

pub type BoxedBoolean = bool;
pub type BoxedInteger = i64;
pub type Integer = i64;
pub type StringList = Vec<String>;

pub type DockerLabelsMap = ::std::collections::HashMap<String, String>;
pub type EnvironmentVariables = Vec<KeyValuePair>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct KeyValuePair {
    pub name: Option<String>,
    pub value: Option<String>,
}
impl KeyValuePair {
    pub fn to_rusoto(&self) -> rusoto_ecs::KeyValuePair {
        rusoto_ecs::KeyValuePair {
            name: self.name.to_owned(),
            value: self.value.to_owned(),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct HostEntry {
    pub hostname: String,
    pub ip_address: String,
}
impl HostEntry {
    pub fn to_rusoto(&self) -> rusoto_ecs::HostEntry {
        rusoto_ecs::HostEntry {
            hostname: self.hostname.to_owned(),
            ip_address: self.ip_address.to_owned(),
        }
    }
}

pub type HostEntryList = Vec<HostEntry>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct MountPoint {
    pub container_path: Option<String>,
    pub read_only: Option<BoxedBoolean>,
    pub source_volume: Option<String>,
}
impl MountPoint {
    pub fn to_rusoto(&self) -> rusoto_ecs::MountPoint {
        rusoto_ecs::MountPoint {
            container_path: self.container_path.to_owned(),
            read_only: self.read_only,
            source_volume: self.source_volume.to_owned(),
        }
    }
}

pub type MountPointList = Vec<MountPoint>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct PortMapping {
    pub container_port: Option<BoxedInteger>,
    pub host_port: Option<BoxedInteger>,
    pub protocol: Option<TransportProtocol>,
}
impl PortMapping {
    pub fn to_rusoto(&self) -> rusoto_ecs::PortMapping {
        rusoto_ecs::PortMapping {
            container_port: self.container_port,
            host_port: self.host_port,
            protocol: self.protocol.to_owned(),
        }
    }
}

pub type PortMappingList = Vec<PortMapping>;
pub type TransportProtocol = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Ulimit {
    pub hard_limit: Integer,
    pub name: UlimitName,
    pub soft_limit: Integer,
}
impl Ulimit {
    pub fn to_rusoto(&self) -> rusoto_ecs::Ulimit {
        rusoto_ecs::Ulimit {
            hard_limit: self.hard_limit,
            name: self.name.to_owned(),
            soft_limit: self.soft_limit,
        }
    }
}

pub type UlimitList = Vec<Ulimit>;
pub type UlimitName = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct VolumeFrom {
    pub read_only: Option<BoxedBoolean>,
    pub source_container: Option<String>,
}
impl VolumeFrom {
    pub fn to_rusoto(&self) -> rusoto_ecs::VolumeFrom {
        rusoto_ecs::VolumeFrom {
            read_only: self.read_only,
            source_container: self.source_container.to_owned(),
        }
    }
}

pub type VolumeFromList = Vec<VolumeFrom>;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct LogConfiguration {
    pub log_driver: LogDriver,
    pub options: Option<LogConfigurationOptionsMap>,
}
impl LogConfiguration {
    pub fn to_rusoto(&self) -> rusoto_ecs::LogConfiguration {
        rusoto_ecs::LogConfiguration {
            log_driver: self.log_driver.to_owned(),
            options: self.options.to_owned(),
        }
    }
}

pub type LogConfigurationOptionsMap = ::std::collections::HashMap<String, String>;
pub type LogDriver = String;

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct HealthCheck {
    pub command: Vec<String>,
    pub interval: Option<i64>,
    pub retries: Option<i64>,
    pub start_period: Option<i64>,
    pub timeout: Option<i64>,
}
impl HealthCheck {
    pub fn to_rusoto(&self) -> rusoto_ecs::HealthCheck {
        rusoto_ecs::HealthCheck {
            command: self.command.to_owned(),
            interval: self.interval,
            retries: self.retries,
            start_period: self.start_period,
            timeout: self.timeout,
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct LinuxParameters {
    pub capabilities: Option<KernelCapabilities>,
    pub devices: Option<Vec<Device>>,
    pub init_process_enabled: Option<bool>,
    pub shared_memory_size: Option<i64>,
    pub tmpfs: Option<Vec<Tmpfs>>,
}
impl LinuxParameters {
    pub fn to_rusoto(&self) -> rusoto_ecs::LinuxParameters {
        rusoto_ecs::LinuxParameters {
            capabilities: self.capabilities.as_ref().map(|e| e.to_rusoto()),
            devices: self
                .devices
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
            init_process_enabled: self.init_process_enabled,
            shared_memory_size: self.shared_memory_size,
            tmpfs: self
                .tmpfs
                .as_ref()
                .map(|e| e.iter().map(|e0| e0.to_rusoto()).collect()),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct KernelCapabilities {
    pub add: Option<Vec<String>>,
    pub drop: Option<Vec<String>>,
}
impl KernelCapabilities {
    pub fn to_rusoto(&self) -> rusoto_ecs::KernelCapabilities {
        rusoto_ecs::KernelCapabilities {
            add: self.add.to_owned(),
            drop: self.drop.to_owned(),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Device {
    pub container_path: Option<String>,
    pub host_path: String,
    pub permissions: Option<Vec<String>>,
}
impl Device {
    pub fn to_rusoto(&self) -> rusoto_ecs::Device {
        rusoto_ecs::Device {
            container_path: self.container_path.to_owned(),
            host_path: self.host_path.to_owned(),
            permissions: self.permissions.to_owned(),
        }
    }
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Tmpfs {
    pub container_path: String,
    pub mount_options: Option<Vec<String>>,
    pub size: i64,
}
impl Tmpfs {
    pub fn to_rusoto(&self) -> rusoto_ecs::Tmpfs {
        rusoto_ecs::Tmpfs {
            container_path: self.container_path.to_owned(),
            mount_options: self.mount_options.to_owned(),
            size: self.size,
        }
    }
}
<|endoftext|>


<commit_before>use crate::{
    channel::Channel,
    channels::Channels,
    configuration::Configuration,
    connection_closer::ConnectionCloser,
    connection_properties::ConnectionProperties,
    connection_status::{ConnectionState, ConnectionStatus, ConnectionStep},
    executor::{DefaultExecutor, Executor},
    frames::Frames,
    internal_rpc::{InternalRPC, InternalRPCHandle},
    io_loop::IoLoop,
    reactor::DefaultReactorBuilder,
    socket_state::{SocketState, SocketStateHandle},
    tcp::{AMQPUriTcpExt, HandshakeResult, TLSConfig},
    thread::ThreadHandle,
    types::ShortUInt,
    uri::AMQPUri,
    Error, Promise, PromiseChain, Result,
};
use amq_protocol::frame::{AMQPFrame, ProtocolVersion};
use log::{log_enabled, Level::Trace};
use std::{fmt, io, sync::Arc};

pub struct Connection {
    configuration: Configuration,
    status: ConnectionStatus,
    channels: Channels,
    io_loop: ThreadHandle,
    closer: Arc<ConnectionCloser>,
}

impl Connection {
    fn new(
        waker: SocketStateHandle,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
    ) -> Self {
        let configuration = Configuration::default();
        let status = ConnectionStatus::default();
        let channels = Channels::new(
            configuration.clone(),
            status.clone(),
            waker,
            internal_rpc.clone(),
            frames,
            executor,
        );
        let closer = Arc::new(ConnectionCloser::new(status.clone(), internal_rpc));
        let connection = Self {
            configuration,
            status,
            channels,
            io_loop: ThreadHandle::default(),
            closer,
        };

        connection.channels.create_zero();
        connection
    }

    /// Connect to an AMQP Server
    pub async fn connect(uri: &str, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, TLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_with_config(
        uri: &str,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri(uri: AMQPUri, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, TLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri_with_identity(
        uri: AMQPUri,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    pub async fn create_channel(&self) -> Result<Channel> {
        if !self.status.connected() {
            return Err(Error::InvalidConnectionState(self.status.state()));
        }
        let channel = self.channels.create(self.closer.clone())?;
        channel.clone().channel_open(channel).await
    }

    /// Block current thread while the connection is still active.
    /// This is useful when you only have a consumer and nothing else keeping your application
    /// "alive".
    pub fn run(self) -> Result<()> {
        let io_loop = self.io_loop.clone();
        drop(self);
        io_loop.wait("io loop")
    }

    pub fn on_error<E: FnMut(Error) + Send + 'static>(&self, handler: E) {
        self.channels.set_error_handler(handler);
    }

    pub fn configuration(&self) -> &Configuration {
        &self.configuration
    }

    pub fn status(&self) -> &ConnectionStatus {
        &self.status
    }

    pub async fn close(&self, reply_code: ShortUInt, reply_text: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0
                .connection_close(reply_code, reply_text, 0, 0)
                .await
        } else {
            Ok(())
        }
    }

    /// Block all consumers and publishers on this connection
    pub async fn block(&self, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_blocked(reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Unblock all consumers and publishers on this connection
    pub async fn unblock(&self) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_unblocked().await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Update the secret used by some authentication module such as OAuth2
    pub async fn update_secret(&self, new_secret: &str, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_update_secret(new_secret, reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    pub fn connector(
        uri: AMQPUri,
        handshake_result: HandshakeResult,
        mut options: ConnectionProperties,
    ) -> PromiseChain<Connection> {
        let executor = options
            .executor
            .take()
            .unwrap_or_else(|| Arc::new(DefaultExecutor::default()));
        let reactor_builder = options
            .reactor_builder
            .take()
            .unwrap_or_else(|| Arc::new(DefaultReactorBuilder));
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let frames = Frames::default();
        let conn = Connection::new(waker, internal_rpc.handle(), frames.clone(), executor);
        let status = conn.status.clone();
        let configuration = conn.configuration.clone();
        status.set_vhost(&uri.vhost);
        status.set_username(&uri.authority.userinfo.username);
        if let Some(frame_max) = uri.query.frame_max {
            configuration.set_frame_max(frame_max);
        }
        if let Some(channel_max) = uri.query.channel_max {
            configuration.set_channel_max(channel_max);
        }
        if let Some(heartbeat) = uri.query.heartbeat {
            configuration.set_heartbeat(heartbeat);
        }
        let (promise, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise.set_marker("ProtocolHeader".into());
        }
        let channels = conn.channels.clone();
        if let Some(channel0) = channels.get(0) {
            channel0.send_frame(
                AMQPFrame::ProtocolHeader(ProtocolVersion::amqp_0_9_1()),
                resolver,
                None,
            )
        };
        let (promise, resolver) = PromiseChain::after(promise);
        if log_enabled!(Trace) {
            promise.set_marker("ProtocolHeader.Ok".into());
        }
        let io_loop_handle = conn.io_loop.clone();
        status.set_state(ConnectionState::Connecting);
        status.set_connection_step(ConnectionStep::ProtocolHeader(
            resolver,
            conn,
            uri.authority.userinfo.into(),
            uri.query.auth_mechanism.unwrap_or_default(),
            options,
        ));
        IoLoop::new(
            status,
            configuration,
            channels,
            internal_rpc,
            frames,
            socket_state,
            io_loop_handle,
            handshake_result,
            &*reactor_builder,
        )
        .and_then(IoLoop::start)
        .map(|()| promise)
        .unwrap_or_else(|err| PromiseChain::new_with_data(Err(err)))
    }
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Connection")
            .field("configuration", &self.configuration)
            .field("status", &self.status)
            .field("channels", &self.channels)
            .finish()
    }
}

/// Trait providing a method to connect to an AMQP server
pub trait Connect {
    /// connect to an AMQP server
    fn connect(
        self,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> PromiseChain<Connection>;
}

impl Connect for AMQPUri {
    fn connect(
        self,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> PromiseChain<Connection> {
        let stream = AMQPUriTcpExt::connect_with_config(&self, config);
        Connection::connector(self, stream, options)
    }
}

impl Connect for &str {
    fn connect(
        self,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> PromiseChain<Connection> {
        match self.parse::<AMQPUri>() {
            Ok(uri) => Connect::connect(uri, options, config),
            Err(err) => {
                PromiseChain::new_with_data(Err(io::Error::new(io::ErrorKind::Other, err).into()))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use env_logger;

    use super::*;
    use crate::channel_receiver_state::ChannelReceiverState;
    use crate::channel_status::ChannelState;
    use crate::types::ShortString;
    use crate::BasicProperties;
    use amq_protocol::frame::AMQPContentHeader;
    use amq_protocol::protocol::{basic, AMQPClass};

    #[test]
    fn basic_consume_small_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let executor = Arc::new(DefaultExecutor::default());
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 2,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::ReceivingContent(
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
                2,
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let body_frame = AMQPFrame::Body(channel.id(), "{}".as_bytes().to_vec());
            conn.channels.handle_frame(body_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }

    #[test]
    fn basic_consume_empty_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let executor = Arc::new(DefaultExecutor::default());
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 0,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }
}
<commit_msg>make Connection::connector an async fn<commit_after>use crate::{
    channel::Channel,
    channels::Channels,
    configuration::Configuration,
    connection_closer::ConnectionCloser,
    connection_properties::ConnectionProperties,
    connection_status::{ConnectionState, ConnectionStatus, ConnectionStep},
    executor::{DefaultExecutor, Executor},
    frames::Frames,
    internal_rpc::{InternalRPC, InternalRPCHandle},
    io_loop::IoLoop,
    reactor::DefaultReactorBuilder,
    socket_state::{SocketState, SocketStateHandle},
    tcp::{AMQPUriTcpExt, HandshakeResult, TLSConfig},
    thread::ThreadHandle,
    types::ShortUInt,
    uri::AMQPUri,
    Error, Promise, Result,
};
use amq_protocol::frame::{AMQPFrame, ProtocolVersion};
use log::{log_enabled, Level::Trace};
use std::{fmt, future::Future, io, pin::Pin, sync::Arc};

pub struct Connection {
    configuration: Configuration,
    status: ConnectionStatus,
    channels: Channels,
    io_loop: ThreadHandle,
    closer: Arc<ConnectionCloser>,
}

impl Connection {
    fn new(
        waker: SocketStateHandle,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
    ) -> Self {
        let configuration = Configuration::default();
        let status = ConnectionStatus::default();
        let channels = Channels::new(
            configuration.clone(),
            status.clone(),
            waker,
            internal_rpc.clone(),
            frames,
            executor,
        );
        let closer = Arc::new(ConnectionCloser::new(status.clone(), internal_rpc));
        let connection = Self {
            configuration,
            status,
            channels,
            io_loop: ThreadHandle::default(),
            closer,
        };

        connection.channels.create_zero();
        connection
    }

    /// Connect to an AMQP Server
    pub async fn connect(uri: &str, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, TLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_with_config(
        uri: &str,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri(uri: AMQPUri, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, TLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri_with_identity(
        uri: AMQPUri,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    pub async fn create_channel(&self) -> Result<Channel> {
        if !self.status.connected() {
            return Err(Error::InvalidConnectionState(self.status.state()));
        }
        let channel = self.channels.create(self.closer.clone())?;
        channel.clone().channel_open(channel).await
    }

    /// Block current thread while the connection is still active.
    /// This is useful when you only have a consumer and nothing else keeping your application
    /// "alive".
    pub fn run(self) -> Result<()> {
        let io_loop = self.io_loop.clone();
        drop(self);
        io_loop.wait("io loop")
    }

    pub fn on_error<E: FnMut(Error) + Send + 'static>(&self, handler: E) {
        self.channels.set_error_handler(handler);
    }

    pub fn configuration(&self) -> &Configuration {
        &self.configuration
    }

    pub fn status(&self) -> &ConnectionStatus {
        &self.status
    }

    pub async fn close(&self, reply_code: ShortUInt, reply_text: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0
                .connection_close(reply_code, reply_text, 0, 0)
                .await
        } else {
            Ok(())
        }
    }

    /// Block all consumers and publishers on this connection
    pub async fn block(&self, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_blocked(reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Unblock all consumers and publishers on this connection
    pub async fn unblock(&self) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_unblocked().await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Update the secret used by some authentication module such as OAuth2
    pub async fn update_secret(&self, new_secret: &str, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_update_secret(new_secret, reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    pub async fn connector(
        uri: AMQPUri,
        handshake_result: HandshakeResult,
        mut options: ConnectionProperties,
    ) -> Result<Connection> {
        let executor = options
            .executor
            .take()
            .unwrap_or_else(|| Arc::new(DefaultExecutor::default()));
        let reactor_builder = options
            .reactor_builder
            .take()
            .unwrap_or_else(|| Arc::new(DefaultReactorBuilder));
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let frames = Frames::default();
        let conn = Connection::new(waker, internal_rpc.handle(), frames.clone(), executor);
        let status = conn.status.clone();
        let configuration = conn.configuration.clone();
        status.set_vhost(&uri.vhost);
        status.set_username(&uri.authority.userinfo.username);
        if let Some(frame_max) = uri.query.frame_max {
            configuration.set_frame_max(frame_max);
        }
        if let Some(channel_max) = uri.query.channel_max {
            configuration.set_channel_max(channel_max);
        }
        if let Some(heartbeat) = uri.query.heartbeat {
            configuration.set_heartbeat(heartbeat);
        }
        let (promise_out, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise_out.set_marker("ProtocolHeader".into());
        }
        let channels = conn.channels.clone();
        if let Some(channel0) = channels.get(0) {
            channel0.send_frame(
                AMQPFrame::ProtocolHeader(ProtocolVersion::amqp_0_9_1()),
                resolver,
                None,
            )
        };
        let (promise_in, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise_in.set_marker("ProtocolHeader.Ok".into());
        }
        let io_loop_handle = conn.io_loop.clone();
        status.set_state(ConnectionState::Connecting);
        status.set_connection_step(ConnectionStep::ProtocolHeader(
            resolver,
            conn,
            uri.authority.userinfo.into(),
            uri.query.auth_mechanism.unwrap_or_default(),
            options,
        ));
        IoLoop::new(
            status,
            configuration,
            channels,
            internal_rpc,
            frames,
            socket_state,
            io_loop_handle,
            handshake_result,
            &*reactor_builder,
        )
        .and_then(IoLoop::start)?;
        promise_out.await?;
        promise_in.await
    }
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Connection")
            .field("configuration", &self.configuration)
            .field("status", &self.status)
            .field("channels", &self.channels)
            .finish()
    }
}

/// Trait providing a method to connect to an AMQP server
pub trait Connect {
    /// connect to an AMQP server
    fn connect(
        self,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Pin<Box<dyn Future<Output = Result<Connection>>>>;
}

impl Connect for AMQPUri {
    fn connect(
        self,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Pin<Box<dyn Future<Output = Result<Connection>>>> {
        let stream = AMQPUriTcpExt::connect_with_config(&self, config);
        Box::pin(Connection::connector(self, stream, options))
    }
}

impl Connect for &str {
    fn connect(
        self,
        options: ConnectionProperties,
        config: TLSConfig<'_, '_, '_>,
    ) -> Pin<Box<dyn Future<Output = Result<Connection>>>> {
        match self.parse::<AMQPUri>() {
            Ok(uri) => Connect::connect(uri, options, config),
            Err(err) => Box::pin(async { Err(io::Error::new(io::ErrorKind::Other, err).into()) }),
        }
    }
}

#[cfg(test)]
mod tests {
    use env_logger;

    use super::*;
    use crate::channel_receiver_state::ChannelReceiverState;
    use crate::channel_status::ChannelState;
    use crate::types::ShortString;
    use crate::BasicProperties;
    use amq_protocol::frame::AMQPContentHeader;
    use amq_protocol::protocol::{basic, AMQPClass};

    #[test]
    fn basic_consume_small_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let executor = Arc::new(DefaultExecutor::default());
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 2,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::ReceivingContent(
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
                2,
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let body_frame = AMQPFrame::Body(channel.id(), "{}".as_bytes().to_vec());
            conn.channels.handle_frame(body_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }

    #[test]
    fn basic_consume_empty_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let executor = Arc::new(DefaultExecutor::default());
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 0,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }
}
<|endoftext|>