<commit_before>use std::sync::mpsc::Sender;
use std::time::Duration;
use std::process::{Command, Stdio};
use std::thread;
use std::sync::{Arc, RwLock};
use std::io::{BufReader, BufRead};
use super::base::Widget;
use format::data::Format;
use nom::IResult;

#[derive(Debug, Clone)]
pub struct BspwmDesktop {
    pub name: String,
    pub occupied: bool,
    pub focused: bool,
    pub urgent: bool,
}

named!(bspstr<&[u8], Vec<BspwmDesktop> >,
   chain!(
       tag!("WM") ~
       take_until_and_consume!(":") ~
       d: many0!(
           chain!(
               mode: take!(1) ~
               name: take_until_and_consume!(":") ,
               || BspwmDesktop {
                   name: String::from_utf8_lossy(name).into_owned(),
                   occupied: mode == b"o" || mode == b"O",
                   focused: mode == b"F" || mode == b"O" || mode == b"U",
                   urgent: mode == b"u" || mode == b"U",
               }
           )
       ) ~
       tag!("LT") ,
       || d
   )
);

pub struct Bspwm<F: Fn(Vec<BspwmDesktop>) -> Format> {
    updater: Arc<Box<F>>,
    last_value: Arc<RwLock<Format>>,
}

impl<F> Widget for Bspwm<F> where F: Fn(Vec<BspwmDesktop>) -> Format + Sync + Send + 'static  {
    fn current_value(&self) -> Format {
        (*self.last_value).read().unwrap().clone()
    }

    fn spawn_notifier(&mut self, tx: Sender<()>) {
        let updater = self.updater.clone();
        let last_value = self.last_value.clone();
        thread::spawn(move || {
            loop {
                // Should be possible to use the socket directly...
                let bspc = Command::new("bspc").arg("control").arg("--subscribe")
                    .stdout(Stdio::piped()).spawn().expect("Couldn't run bspc");
                for line in BufReader::new(bspc.stdout.unwrap()).lines() {
                    let mut writer = last_value.write().unwrap();
                    let line = line.unwrap_or("".to_owned());
                    if let IResult::Done(_, result) = bspstr(&line.into_bytes()) {
                        *writer = (*updater)(result);
                        let _ = tx.send(());
                    }
                }
                thread::sleep(Duration::from_millis(500));
            }
        });
    }
}

impl<F> Bspwm<F> where F: Fn(Vec<BspwmDesktop>) -> Format {
    pub fn new(updater: F) -> Box<Bspwm<F>> {
        Box::new(Bspwm {
            updater: Arc::new(Box::new(updater)),
            last_value: Arc::new(RwLock::new(Format::Str("".to_owned()))),
        })
    }
}
<commit_msg>fix bspwm command<commit_after>use std::sync::mpsc::Sender;
use std::time::Duration;
use std::process::{Command, Stdio};
use std::thread;
use std::sync::{Arc, RwLock};
use std::io::{BufReader, BufRead};
use super::base::Widget;
use format::data::Format;
use nom::IResult;

#[derive(Debug, Clone)]
pub struct BspwmDesktop {
    pub name: String,
    pub occupied: bool,
    pub focused: bool,
    pub urgent: bool,
}

named!(bspstr<&[u8], Vec<BspwmDesktop> >,
   chain!(
       tag!("WM") ~
       take_until_and_consume!(":") ~
       d: many0!(
           chain!(
               mode: take!(1) ~
               name: take_until_and_consume!(":") ,
               || BspwmDesktop {
                   name: String::from_utf8_lossy(name).into_owned(),
                   occupied: mode == b"o" || mode == b"O",
                   focused: mode == b"F" || mode == b"O" || mode == b"U",
                   urgent: mode == b"u" || mode == b"U",
               }
           )
       ) ~
       tag!("LT") ,
       || d
   )
);

pub struct Bspwm<F: Fn(Vec<BspwmDesktop>) -> Format> {
    updater: Arc<Box<F>>,
    last_value: Arc<RwLock<Format>>,
}

impl<F> Widget for Bspwm<F> where F: Fn(Vec<BspwmDesktop>) -> Format + Sync + Send + 'static  {
    fn current_value(&self) -> Format {
        (*self.last_value).read().unwrap().clone()
    }

    fn spawn_notifier(&mut self, tx: Sender<()>) {
        let updater = self.updater.clone();
        let last_value = self.last_value.clone();
        thread::spawn(move || {
            loop {
                // Should be possible to use the socket directly...
                let bspc = Command::new("bspc").arg("subscribe")
                    .stdout(Stdio::piped()).spawn().expect("Couldn't run bspc");
                for line in BufReader::new(bspc.stdout.unwrap()).lines() {
                    let mut writer = last_value.write().unwrap();
                    let line = line.unwrap_or("".to_owned());
                    if let IResult::Done(_, result) = bspstr(&line.into_bytes()) {
                        *writer = (*updater)(result);
                        let _ = tx.send(());
                    }
                }
                thread::sleep(Duration::from_millis(500));
            }
        });
    }
}

impl<F> Bspwm<F> where F: Fn(Vec<BspwmDesktop>) -> Format {
    pub fn new(updater: F) -> Box<Bspwm<F>> {
        Box::new(Bspwm {
            updater: Arc::new(Box::new(updater)),
            last_value: Arc::new(RwLock::new(Format::Str("".to_owned()))),
        })
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#[doc(hidden)]
#[macro_export]
macro_rules! as_item { ($i:item) => {$i} }

// Inserts a placeholder doc comment for the module if it's missing
#[doc(hidden)]
#[macro_export]
macro_rules! add_mod_doc {
    // If nothing left, return
    (
        @rec
        { $(#[$done:meta])* }
        { }
        $i:item
    ) => {
        $(#[$done])*
        #[doc="A module containing an rpc service and client stub."]
        $i
    };

    // If we find a doc attribute, return
    (
        @rec
        { $(#[$done:meta])* }
        {
            #[doc=$doc:expr]
            $(#[$rest:meta])*
        }
        $i:item
    ) => {
        $(#[$done])*
        #[doc=$doc]
        $(#[$rest])*
        $i
    };

    // If we don't find a doc attribute, keep going
    (
        @rec
        { $(#[$($done:tt)*])* }
        {
            #[$($attr:tt)*]
            $($rest:tt)*
        }
        $i:item
    ) => {
        add_mod_doc! {
            @rec
            { $(#[$($done)*])* #[$($attr)*] }
            { $($rest)* }
            $i
        }
    };

    // Entry
    (
        { $(#[$($attr:tt)*])* }
        $i:item
    ) => {
        add_mod_doc! {
            @rec
            {}
            { $(#[$($attr)*])* }
            $i
        }
    };
}

// Required because if-let can't be used with irrefutable patterns, so it needs
// to be special cased.
#[doc(hidden)]
#[macro_export]
macro_rules! client_methods {
    (
        { $(#[$attr:meta])* }
        $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    ) => (
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> $crate::Result<$out> {
            let reply = try!((self.0).rpc(&request_variant!($fn_name $($arg),*)));
            let __Reply::$fn_name(reply) = reply;
            Ok(reply)
        }
    );
    ($(
            { $(#[$attr:meta])* }
            $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    )*) => ( $(
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> $crate::Result<$out> {
            let reply = try!((self.0).rpc(&request_variant!($fn_name $($arg),*)));
            if let __Reply::$fn_name(reply) = reply {
                Ok(reply)
            } else {
                panic!("Incorrect reply variant returned from protocol::Clientrpc; expected `{}`, but got {:?}", stringify!($fn_name), reply);
            }
        }
    )*);
}

// Required because if-let can't be used with irrefutable patterns, so it needs
// to be special cased.
#[doc(hidden)]
#[macro_export]
macro_rules! async_client_methods {
    (
        { $(#[$attr:meta])* }
        $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    ) => (
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> Future<$out> {
            fn mapper(reply: __Reply) -> $out {
                let __Reply::$fn_name(reply) = reply;
                reply
            }
            let reply = (self.0).rpc_async(&request_variant!($fn_name $($arg),*));
            Future {
                future: reply,
                mapper: mapper,
            }
        }
    );
    ($(
            { $(#[$attr:meta])* }
            $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    )*) => ( $(
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> Future<$out> {
            fn mapper(reply: __Reply) -> $out {
                if let __Reply::$fn_name(reply) = reply {
                    reply
                } else {
                    panic!("Incorrect reply variant returned from protocol::Clientrpc; expected `{}`, but got {:?}", stringify!($fn_name), reply);
                }
            }
            let reply = (self.0).rpc_async(&request_variant!($fn_name $($arg),*));
            Future {
                future: reply,
                mapper: mapper,
            }
        }
    )*);
}

// Required because enum variants with no fields can't be suffixed by parens
#[doc(hidden)]
#[macro_export]
macro_rules! define_request {
    ($(@($($finished:tt)*))* --) => (as_item!(
            #[allow(non_camel_case_types)]
            #[derive(Debug, Serialize, Deserialize)]
            enum __Request { $($($finished)*),* }
    ););
    ($(@$finished:tt)* -- $name:ident() $($req:tt)*) =>
        (define_request!($(@$finished)* @($name) -- $($req)*););
    ($(@$finished:tt)* -- $name:ident $args: tt $($req:tt)*) =>
        (define_request!($(@$finished)* @($name $args) -- $($req)*););
    ($($started:tt)*) => (define_request!(-- $($started)*););
}

// Required because enum variants with no fields can't be suffixed by parens
#[doc(hidden)]
#[macro_export]
macro_rules! request_variant {
    ($x:ident) => (__Request::$x);
    ($x:ident $($y:ident),+) => (__Request::$x($($y),+));
}

// The main macro that creates RPC services.
#[macro_export]
macro_rules! rpc {
    (
        $(#[$($service_attr:tt)*])*
        mod $server:ident {

            service {
                $(
                    $(#[$attr:meta])*
                    rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
                )*
            }
        }
    ) => {
        rpc! {
            $(#[$($service_attr)*])*
            mod $server {

                items { }

                service {
                    $(
                        $(#[$attr])*
                        rpc $fn_name($($arg: $in_),*) -> $out;
                    )*
                }
            }
        }
    };

    (
        // Names the service
        $(#[$($service_attr:tt)*])*
        mod $server:ident {

            // Include any desired or required items. Conflicts can arise with the following names:
            // 1. Service
            // 2. Client
            // 3. serve
            // 4. __Reply
            // 5. __Request
            items { $($i:item)* }

            // List any rpc methods: rpc foo(arg1: Arg1, ..., argN: ArgN) -> Out
            service {
                $(
                    $(#[$attr:meta])*
                    rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
                )*
            }
        }
    ) => {

        add_mod_doc! {
            { $(#[$($service_attr)*])* }
            pub mod $server {

                $($i)*

                #[doc="The provided RPC service."]
                pub trait Service: Send + Sync {
                    $(
                        $(#[$attr])*
                        fn $fn_name(&self, $($arg:$in_),*) -> $out;
                    )*
                }

                impl<P, S> Service for P
                    where P: Send + Sync + ::std::ops::Deref<Target=S>,
                          S: Service
                {
                    $(
                        $(#[$attr])*
                        fn $fn_name(&self, $($arg:$in_),*) -> $out {
                            Service::$fn_name(&**self, $($arg),*)
                        }
                    )*
                }

                define_request!($($fn_name($($in_),*))*);

                #[allow(non_camel_case_types)]
                #[derive(Debug, Serialize, Deserialize)]
                enum __Reply {
                    $(
                        $fn_name($out),
                    )*
                }

                /// An asynchronous RPC call
                pub struct Future<T> {
                    future: $crate::protocol::Future<__Reply>,
                    mapper: fn(__Reply) -> T,
                }

                impl<T> Future<T> {
                    /// Block until the result of the RPC call is available
                    pub fn get(self) -> $crate::Result<T> {
                        self.future.get().map(self.mapper)
                    }
                }

                #[doc="The client stub that makes RPC calls to the server."]
                pub struct Client($crate::protocol::Client<__Request, __Reply>);

                impl Client {
                    #[doc="Create a new client that connects to the given address."]
                    pub fn new<A>(addr: A, timeout: ::std::option::Option<::std::time::Duration>)
                        -> $crate::Result<Self>
                        where A: ::std::net::ToSocketAddrs,
                    {
                        let inner = try!($crate::protocol::Client::new(addr, timeout));
                        Ok(Client(inner))
                    }

                    client_methods!(
                        $(
                            { $(#[$attr])* }
                            $fn_name($($arg: $in_),*) -> $out
                        )*
                    );
                }

                #[doc="The client stub that makes asynchronous RPC calls to the server."]
                pub struct AsyncClient($crate::protocol::Client<__Request, __Reply>);

                impl AsyncClient {
                    #[doc="Create a new asynchronous client that connects to the given address."]
                    pub fn new<A>(addr: A, timeout: ::std::option::Option<::std::time::Duration>)
                        -> $crate::Result<Self>
                        where A: ::std::net::ToSocketAddrs,
                    {
                        let inner = try!($crate::protocol::Client::new(addr, timeout));
                        Ok(AsyncClient(inner))
                    }

                    async_client_methods!(
                        $(
                            { $(#[$attr])* }
                            $fn_name($($arg: $in_),*) -> $out
                        )*
                    );
                }

                struct __Server<S: 'static + Service>(S);

                impl<S> $crate::protocol::Serve for __Server<S>
                    where S: 'static + Service
                {
                    type Request = __Request;
                    type Reply = __Reply;
                    fn serve(&self, request: __Request) -> __Reply {
                        match request {
                            $(
                                request_variant!($fn_name $($arg),*) =>
                                    __Reply::$fn_name((self.0).$fn_name($($arg),*)),
                             )*
                        }
                    }
                }

                #[doc="Start a running service."]
                pub fn serve<A, S>(addr: A,
                                   service: S,
                                   read_timeout: ::std::option::Option<::std::time::Duration>)
                    -> $crate::Result<$crate::protocol::ServeHandle>
                    where A: ::std::net::ToSocketAddrs,
                          S: 'static + Service
                {
                    let server = ::std::sync::Arc::new(__Server(service));
                    Ok(try!($crate::protocol::serve_async(addr, server, read_timeout)))
                }
            }
        }
    }
}

#[cfg(test)]
#[allow(dead_code)]
mod test {
    extern crate env_logger;
    use std::time::Duration;
    use test::Bencher;

    fn test_timeout() -> Option<Duration> {
        Some(Duration::from_secs(5))
    }

    rpc! {
        #[deny(missing_docs)]
        #[doc="Hello"]
        mod my_server {
            items {
                #[derive(PartialEq, Debug, Serialize, Deserialize)]
                pub struct Foo {
                    pub message: String
                }
            }

            service {
                rpc hello(foo: Foo) -> Foo;
                rpc add(x: i32, y: i32) -> i32;
            }
        }
    }

    use self::my_server::*;

    struct Server;
    impl Service for Server {
        fn hello(&self, s: Foo) -> Foo {
            Foo { message: format!("Hello, {}", &s.message) }
        }

        fn add(&self, x: i32, y: i32) -> i32 {
            x + y
        }
    }

    #[test]
    fn serve_arc_server() {
        serve("localhost:0", ::std::sync::Arc::new(Server), None)
            .unwrap()
            .shutdown();
    }

    #[test]
    fn simple_test() {
        println!("Starting");
        let addr = "127.0.0.1:9000";
        let shutdown = my_server::serve(addr, Server, test_timeout()).unwrap();
        let client = Client::new(addr, None).unwrap();
        assert_eq!(3, client.add(1, 2).unwrap());
        let foo = Foo { message: "Adam".into() };
        let want = Foo { message: format!("Hello, {}", &foo.message) };
        assert_eq!(want, client.hello(Foo { message: "Adam".into() }).unwrap());
        drop(client);
        shutdown.shutdown();
    }

    // Tests a service definition with a fn that takes no args
    rpc! {
        mod qux {
            service {
                rpc hello() -> String;
            }
        }
    }

    // Tests a service definition with an import
    rpc! {
        mod foo {
            items {
                use std::collections::HashMap;
            }

            service {
                #[doc="Hello bob"]
                #[inline(always)]
                rpc baz(s: String) -> HashMap<String, String>;
            }
        }
    }

    // Tests a service definition with an attribute but no doc comment
    rpc! {
        #[deny(missing_docs)]
        mod bar {
            items {
                use std::collections::HashMap;
            }

            service {
                #[inline(always)]
                rpc baz(s: String) -> HashMap<String, String>;
            }
        }
    }

    // Tests a service definition with an attribute and a doc comment
    rpc! {
        #[deny(missing_docs)]
        #[doc="Hello bob"]
        #[allow(unused)]
        mod baz {
            items {
                use std::collections::HashMap;

                #[derive(Debug)]
                pub struct Debuggable;
            }

            service {
                #[doc="Hello bob"]
                #[inline(always)]
                rpc baz(s: String) -> HashMap<String, String>;
            }
        }
    }

    #[test]
    fn debug() {
        println!("{:?}", baz::Debuggable);
    }

    rpc! {
        mod hello {
            service {
                rpc hello(s: String) -> String;
            }
        }
    }

    struct HelloServer;
    impl hello::Service for HelloServer {
        fn hello(&self, s: String) -> String {
            format!("Hello, {}!", s)
        }
    }

    #[bench]
    fn hello(bencher: &mut Bencher) {
        let _ = env_logger::init();
        let handle = hello::serve("localhost:0", HelloServer, None).unwrap();
        let client = hello::AsyncClient::new(handle.local_addr(), None).unwrap();
        let concurrency = 100;
        let mut rpcs = Vec::with_capacity(concurrency);
        bencher.iter(|| {
            for _ in 0..concurrency {
                rpcs.push(client.hello("Bob".into()));
            }
            for _ in 0..concurrency {
                rpcs.pop().unwrap().get().unwrap();
            }
        });
        drop(client);
        handle.shutdown();
    }
}
<commit_msg>Add async test<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#[doc(hidden)]
#[macro_export]
macro_rules! as_item { ($i:item) => {$i} }

// Inserts a placeholder doc comment for the module if it's missing
#[doc(hidden)]
#[macro_export]
macro_rules! add_mod_doc {
    // If nothing left, return
    (
        @rec
        { $(#[$done:meta])* }
        { }
        $i:item
    ) => {
        $(#[$done])*
        #[doc="A module containing an rpc service and client stub."]
        $i
    };

    // If we find a doc attribute, return
    (
        @rec
        { $(#[$done:meta])* }
        {
            #[doc=$doc:expr]
            $(#[$rest:meta])*
        }
        $i:item
    ) => {
        $(#[$done])*
        #[doc=$doc]
        $(#[$rest])*
        $i
    };

    // If we don't find a doc attribute, keep going
    (
        @rec
        { $(#[$($done:tt)*])* }
        {
            #[$($attr:tt)*]
            $($rest:tt)*
        }
        $i:item
    ) => {
        add_mod_doc! {
            @rec
            { $(#[$($done)*])* #[$($attr)*] }
            { $($rest)* }
            $i
        }
    };

    // Entry
    (
        { $(#[$($attr:tt)*])* }
        $i:item
    ) => {
        add_mod_doc! {
            @rec
            {}
            { $(#[$($attr)*])* }
            $i
        }
    };
}

// Required because if-let can't be used with irrefutable patterns, so it needs
// to be special cased.
#[doc(hidden)]
#[macro_export]
macro_rules! client_methods {
    (
        { $(#[$attr:meta])* }
        $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    ) => (
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> $crate::Result<$out> {
            let reply = try!((self.0).rpc(&request_variant!($fn_name $($arg),*)));
            let __Reply::$fn_name(reply) = reply;
            Ok(reply)
        }
    );
    ($(
            { $(#[$attr:meta])* }
            $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    )*) => ( $(
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> $crate::Result<$out> {
            let reply = try!((self.0).rpc(&request_variant!($fn_name $($arg),*)));
            if let __Reply::$fn_name(reply) = reply {
                Ok(reply)
            } else {
                panic!("Incorrect reply variant returned from protocol::Clientrpc; expected `{}`, but got {:?}", stringify!($fn_name), reply);
            }
        }
    )*);
}

// Required because if-let can't be used with irrefutable patterns, so it needs
// to be special cased.
#[doc(hidden)]
#[macro_export]
macro_rules! async_client_methods {
    (
        { $(#[$attr:meta])* }
        $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    ) => (
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> Future<$out> {
            fn mapper(reply: __Reply) -> $out {
                let __Reply::$fn_name(reply) = reply;
                reply
            }
            let reply = (self.0).rpc_async(&request_variant!($fn_name $($arg),*));
            Future {
                future: reply,
                mapper: mapper,
            }
        }
    );
    ($(
            { $(#[$attr:meta])* }
            $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty
    )*) => ( $(
        $(#[$attr])*
        pub fn $fn_name(&self, $($arg: $in_),*) -> Future<$out> {
            fn mapper(reply: __Reply) -> $out {
                if let __Reply::$fn_name(reply) = reply {
                    reply
                } else {
                    panic!("Incorrect reply variant returned from protocol::Clientrpc; expected `{}`, but got {:?}", stringify!($fn_name), reply);
                }
            }
            let reply = (self.0).rpc_async(&request_variant!($fn_name $($arg),*));
            Future {
                future: reply,
                mapper: mapper,
            }
        }
    )*);
}

// Required because enum variants with no fields can't be suffixed by parens
#[doc(hidden)]
#[macro_export]
macro_rules! define_request {
    ($(@($($finished:tt)*))* --) => (as_item!(
            #[allow(non_camel_case_types)]
            #[derive(Debug, Serialize, Deserialize)]
            enum __Request { $($($finished)*),* }
    ););
    ($(@$finished:tt)* -- $name:ident() $($req:tt)*) =>
        (define_request!($(@$finished)* @($name) -- $($req)*););
    ($(@$finished:tt)* -- $name:ident $args: tt $($req:tt)*) =>
        (define_request!($(@$finished)* @($name $args) -- $($req)*););
    ($($started:tt)*) => (define_request!(-- $($started)*););
}

// Required because enum variants with no fields can't be suffixed by parens
#[doc(hidden)]
#[macro_export]
macro_rules! request_variant {
    ($x:ident) => (__Request::$x);
    ($x:ident $($y:ident),+) => (__Request::$x($($y),+));
}

// The main macro that creates RPC services.
#[macro_export]
macro_rules! rpc {
    (
        $(#[$($service_attr:tt)*])*
        mod $server:ident {

            service {
                $(
                    $(#[$attr:meta])*
                    rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
                )*
            }
        }
    ) => {
        rpc! {
            $(#[$($service_attr)*])*
            mod $server {

                items { }

                service {
                    $(
                        $(#[$attr])*
                        rpc $fn_name($($arg: $in_),*) -> $out;
                    )*
                }
            }
        }
    };

    (
        // Names the service
        $(#[$($service_attr:tt)*])*
        mod $server:ident {

            // Include any desired or required items. Conflicts can arise with the following names:
            // 1. Service
            // 2. Client
            // 3. serve
            // 4. __Reply
            // 5. __Request
            items { $($i:item)* }

            // List any rpc methods: rpc foo(arg1: Arg1, ..., argN: ArgN) -> Out
            service {
                $(
                    $(#[$attr:meta])*
                    rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
                )*
            }
        }
    ) => {

        add_mod_doc! {
            { $(#[$($service_attr)*])* }
            pub mod $server {

                $($i)*

                #[doc="The provided RPC service."]
                pub trait Service: Send + Sync {
                    $(
                        $(#[$attr])*
                        fn $fn_name(&self, $($arg:$in_),*) -> $out;
                    )*
                }

                impl<P, S> Service for P
                    where P: Send + Sync + ::std::ops::Deref<Target=S>,
                          S: Service
                {
                    $(
                        $(#[$attr])*
                        fn $fn_name(&self, $($arg:$in_),*) -> $out {
                            Service::$fn_name(&**self, $($arg),*)
                        }
                    )*
                }

                define_request!($($fn_name($($in_),*))*);

                #[allow(non_camel_case_types)]
                #[derive(Debug, Serialize, Deserialize)]
                enum __Reply {
                    $(
                        $fn_name($out),
                    )*
                }

                /// An asynchronous RPC call
                pub struct Future<T> {
                    future: $crate::protocol::Future<__Reply>,
                    mapper: fn(__Reply) -> T,
                }

                impl<T> Future<T> {
                    /// Block until the result of the RPC call is available
                    pub fn get(self) -> $crate::Result<T> {
                        self.future.get().map(self.mapper)
                    }
                }

                #[doc="The client stub that makes RPC calls to the server."]
                pub struct Client($crate::protocol::Client<__Request, __Reply>);

                impl Client {
                    #[doc="Create a new client that connects to the given address."]
                    pub fn new<A>(addr: A, timeout: ::std::option::Option<::std::time::Duration>)
                        -> $crate::Result<Self>
                        where A: ::std::net::ToSocketAddrs,
                    {
                        let inner = try!($crate::protocol::Client::new(addr, timeout));
                        Ok(Client(inner))
                    }

                    client_methods!(
                        $(
                            { $(#[$attr])* }
                            $fn_name($($arg: $in_),*) -> $out
                        )*
                    );
                }

                #[doc="The client stub that makes asynchronous RPC calls to the server."]
                pub struct AsyncClient($crate::protocol::Client<__Request, __Reply>);

                impl AsyncClient {
                    #[doc="Create a new asynchronous client that connects to the given address."]
                    pub fn new<A>(addr: A, timeout: ::std::option::Option<::std::time::Duration>)
                        -> $crate::Result<Self>
                        where A: ::std::net::ToSocketAddrs,
                    {
                        let inner = try!($crate::protocol::Client::new(addr, timeout));
                        Ok(AsyncClient(inner))
                    }

                    async_client_methods!(
                        $(
                            { $(#[$attr])* }
                            $fn_name($($arg: $in_),*) -> $out
                        )*
                    );
                }

                struct __Server<S: 'static + Service>(S);

                impl<S> $crate::protocol::Serve for __Server<S>
                    where S: 'static + Service
                {
                    type Request = __Request;
                    type Reply = __Reply;
                    fn serve(&self, request: __Request) -> __Reply {
                        match request {
                            $(
                                request_variant!($fn_name $($arg),*) =>
                                    __Reply::$fn_name((self.0).$fn_name($($arg),*)),
                             )*
                        }
                    }
                }

                #[doc="Start a running service."]
                pub fn serve<A, S>(addr: A,
                                   service: S,
                                   read_timeout: ::std::option::Option<::std::time::Duration>)
                    -> $crate::Result<$crate::protocol::ServeHandle>
                    where A: ::std::net::ToSocketAddrs,
                          S: 'static + Service
                {
                    let server = ::std::sync::Arc::new(__Server(service));
                    Ok(try!($crate::protocol::serve_async(addr, server, read_timeout)))
                }
            }
        }
    }
}

#[cfg(test)]
#[allow(dead_code)]
mod test {
    extern crate env_logger;
    use std::time::Duration;
    use test::Bencher;

    fn test_timeout() -> Option<Duration> {
        Some(Duration::from_secs(5))
    }

    rpc! {
        #[deny(missing_docs)]
        #[doc="Hello"]
        mod my_server {
            items {
                #[derive(PartialEq, Debug, Serialize, Deserialize)]
                pub struct Foo {
                    pub message: String
                }
            }

            service {
                rpc hello(foo: Foo) -> Foo;
                rpc add(x: i32, y: i32) -> i32;
            }
        }
    }

    use self::my_server::*;

    struct Server;
    impl Service for Server {
        fn hello(&self, s: Foo) -> Foo {
            Foo { message: format!("Hello, {}", &s.message) }
        }

        fn add(&self, x: i32, y: i32) -> i32 {
            x + y
        }
    }

    #[test]
    fn serve_arc_server() {
        serve("localhost:0", ::std::sync::Arc::new(Server), None)
            .unwrap()
            .shutdown();
    }

    #[test]
    fn simple() {
        let handle = my_server::serve( "localhost:0", Server, test_timeout()).unwrap();
        let client = Client::new(handle.local_addr(), None).unwrap();
        assert_eq!(3, client.add(1, 2).unwrap());
        let foo = Foo { message: "Adam".into() };
        let want = Foo { message: format!("Hello, {}", &foo.message) };
        assert_eq!(want, client.hello(Foo { message: "Adam".into() }).unwrap());
        drop(client);
        handle.shutdown();
    }

    #[test]
    fn simple_async() {
        let handle = my_server::serve("localhost:0", Server, test_timeout()).unwrap();
        let client = AsyncClient::new(handle.local_addr(), None).unwrap();
        assert_eq!(3, client.add(1, 2).get().unwrap());
        let foo = Foo { message: "Adam".into() };
        let want = Foo { message: format!("Hello, {}", &foo.message) };
        assert_eq!(want, client.hello(Foo { message: "Adam".into() }).get().unwrap());
        drop(client);
        handle.shutdown();
    }

    // Tests a service definition with a fn that takes no args
    rpc! {
        mod qux {
            service {
                rpc hello() -> String;
            }
        }
    }

    // Tests a service definition with an import
    rpc! {
        mod foo {
            items {
                use std::collections::HashMap;
            }

            service {
                #[doc="Hello bob"]
                #[inline(always)]
                rpc baz(s: String) -> HashMap<String, String>;
            }
        }
    }

    // Tests a service definition with an attribute but no doc comment
    rpc! {
        #[deny(missing_docs)]
        mod bar {
            items {
                use std::collections::HashMap;
            }

            service {
                #[inline(always)]
                rpc baz(s: String) -> HashMap<String, String>;
            }
        }
    }

    // Tests a service definition with an attribute and a doc comment
    rpc! {
        #[deny(missing_docs)]
        #[doc="Hello bob"]
        #[allow(unused)]
        mod baz {
            items {
                use std::collections::HashMap;

                #[derive(Debug)]
                pub struct Debuggable;
            }

            service {
                #[doc="Hello bob"]
                #[inline(always)]
                rpc baz(s: String) -> HashMap<String, String>;
            }
        }
    }

    #[test]
    fn debug() {
        println!("{:?}", baz::Debuggable);
    }

    rpc! {
        mod hello {
            service {
                rpc hello(s: String) -> String;
            }
        }
    }

    struct HelloServer;
    impl hello::Service for HelloServer {
        fn hello(&self, s: String) -> String {
            format!("Hello, {}!", s)
        }
    }

    #[bench]
    fn hello(bencher: &mut Bencher) {
        let _ = env_logger::init();
        let handle = hello::serve("localhost:0", HelloServer, None).unwrap();
        let client = hello::AsyncClient::new(handle.local_addr(), None).unwrap();
        let concurrency = 100;
        let mut rpcs = Vec::with_capacity(concurrency);
        bencher.iter(|| {
            for _ in 0..concurrency {
                rpcs.push(client.hello("Bob".into()));
            }
            for _ in 0..concurrency {
                rpcs.pop().unwrap().get().unwrap();
            }
        });
        drop(client);
        handle.shutdown();
    }
}
<|endoftext|>


<commit_before>use std::net::{UdpSocket,SocketAddr};
use std::io::{self,Write,Error,ErrorKind};
use std::sync::Mutex;
use std::thread;
use std::time::Duration;
use std::str;

use network::buffer::Buffer;

lazy_static! {
  pub static ref METRICS: Mutex<ProxyMetrics> = Mutex::new(ProxyMetrics::new(String::from("sozu")));
}

pub struct ProxyMetrics {
  pub buffer: Buffer,
  pub prefix: String,
  remote: Option<(SocketAddr, UdpSocket)>,
}

impl ProxyMetrics {
  pub fn new(prefix: String) -> Self {
    ProxyMetrics {
      buffer: Buffer::with_capacity(2048),
      prefix: prefix,
      remote:   None,
    }
  }

  pub fn run() -> thread::JoinHandle<()> {
    thread::spawn(move || {
      loop {
        thread::sleep(Duration::from_millis(500));
        METRICS.lock().unwrap().send();
      }
    })
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.remote = Some((addr, socket));
  }

  pub fn send(&mut self) -> io::Result<usize> {
    if self.buffer.available_data() >= 512 {
      if let Some((ref addr, ref socket)) = self.remote {
        match socket.send_to(self.buffer.data(), addr) {
          Ok(sz) => {
            self.buffer.consume(sz);
            Ok(sz)
          },
          Err(e) => {
            Err(e)
          }
        }
      } else {
        Err(Error::new(ErrorKind::NotConnected, "metrics socket not set up"))
      }
    } else {
      Err(Error::new(ErrorKind::Other, "no data to send"))
    }
  }

  fn emit(&mut self, metric: &str) -> io::Result<usize> {
    self.buffer.write(metric.as_bytes())
  }

  pub fn count(&mut self, key: &str, count: i64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|c\n", &self.prefix, key, count))
  }

  pub fn incr(&mut self, key: &str) -> io::Result<()> {
    self.count(key, 1)
  }

  pub fn decr(&mut self, key: &str) -> io::Result<()> {
    self.count(key, -1)
  }

  pub fn time(&mut self, key: &str, time: u64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|ms\n", self.prefix, key, time))
  }

  pub fn gauge(&mut self, key: &str, value: u64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|g\n", self.prefix, key, value))
  }

  pub fn meter(&mut self, key: &str, value: u64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|m\n", self.prefix, key, value))
  }
}

#[macro_export]
macro_rules! count (
  ($key:expr, $value: expr) => {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.buffer.write_fmt(format_args!("{}.{}:{}|c\n", *$crate::logging::TAG, $key, $value));
    metrics.send();
  }
);

#[macro_export]
macro_rules! incr (
  ($key:expr) => (count!($key, 1);)
);

#[macro_export]
macro_rules! decr (
  ($key:expr) => (count!($key, -1);)
);

#[macro_export]
macro_rules! time (
  ($key:expr, $value: expr) => {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.buffer.write_fmt(format_args!("{}.{}:{}|ms\n", *$crate::logging::TAG, $key, $value));
    metrics.send();
  }
);

#[macro_export]
macro_rules! gauge (
  ($key:expr, $value: expr) => {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.buffer.write_fmt(format_args!("{}.{}:{}|g\n", *$crate::logging::TAG, $key, $value));
    metrics.send();
  }
);

#[macro_export]
macro_rules! meter (
  ($key:expr, $value: expr) =>  {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.buffer.write_fmt(format_args!("{}.{}:{}|m\n", *$crate::logging::TAG, $key, $value));
    metrics.send();
  }
);
<commit_msg>reduce the code quantity in macros<commit_after>use std::str;
use std::thread;
use std::sync::Mutex;
use std::time::Duration;
use std::fmt::Arguments;
use std::net::{UdpSocket,SocketAddr};
use std::io::{self,Write,Error,ErrorKind};

use network::buffer::Buffer;

lazy_static! {
  pub static ref METRICS: Mutex<ProxyMetrics> = Mutex::new(ProxyMetrics::new(String::from("sozu")));
}

pub struct ProxyMetrics {
  pub buffer: Buffer,
  pub prefix: String,
  remote: Option<(SocketAddr, UdpSocket)>,
}

impl ProxyMetrics {
  pub fn new(prefix: String) -> Self {
    ProxyMetrics {
      buffer: Buffer::with_capacity(2048),
      prefix: prefix,
      remote:   None,
    }
  }

  pub fn run() -> thread::JoinHandle<()> {
    thread::spawn(move || {
      loop {
        thread::sleep(Duration::from_millis(500));
        METRICS.lock().unwrap().send();
      }
    })
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.remote = Some((addr, socket));
  }

  pub fn write(&mut self, args: Arguments) {
    //FIXME: error handling
    self.buffer.write_fmt(args);
    self.send();
  }

  pub fn send(&mut self) -> io::Result<usize> {
    if self.buffer.available_data() >= 512 {
      if let Some((ref addr, ref socket)) = self.remote {
        match socket.send_to(self.buffer.data(), addr) {
          Ok(sz) => {
            self.buffer.consume(sz);
            Ok(sz)
          },
          Err(e) => {
            Err(e)
          }
        }
      } else {
        Err(Error::new(ErrorKind::NotConnected, "metrics socket not set up"))
      }
    } else {
      Err(Error::new(ErrorKind::Other, "no data to send"))
    }
  }

  fn emit(&mut self, metric: &str) -> io::Result<usize> {
    self.buffer.write(metric.as_bytes())
  }

  pub fn count(&mut self, key: &str, count: i64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|c\n", &self.prefix, key, count))
  }

  pub fn incr(&mut self, key: &str) -> io::Result<()> {
    self.count(key, 1)
  }

  pub fn decr(&mut self, key: &str) -> io::Result<()> {
    self.count(key, -1)
  }

  pub fn time(&mut self, key: &str, time: u64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|ms\n", self.prefix, key, time))
  }

  pub fn gauge(&mut self, key: &str, value: u64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|g\n", self.prefix, key, value))
  }

  pub fn meter(&mut self, key: &str, value: u64) -> io::Result<()> {
    self.buffer.write_fmt(format_args!("{}.{}:{}|m\n", self.prefix, key, value))
  }
}

#[macro_export]
macro_rules! count (
  ($key:expr, $value: expr) => {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.write(format_args!("{}.{}:{}|c\n", *$crate::logging::TAG, $key, $value));
  }
);

#[macro_export]
macro_rules! incr (
  ($key:expr) => (count!($key, 1);)
);

#[macro_export]
macro_rules! decr (
  ($key:expr) => (count!($key, -1);)
);

#[macro_export]
macro_rules! time (
  ($key:expr, $value: expr) => {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.write(format_args!("{}.{}:{}|ms\n", *$crate::logging::TAG, $key, $value));
  }
);

#[macro_export]
macro_rules! gauge (
  ($key:expr, $value: expr) => {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.write(format_args!("{}.{}:{}|g\n", *$crate::logging::TAG, $key, $value));
  }
);

#[macro_export]
macro_rules! meter (
  ($key:expr, $value: expr) =>  {
    let mut metrics = ::network::metrics::METRICS.lock().unwrap();
    metrics.write(format_args!("{}.{}:{}|m\n", *$crate::logging::TAG, $key, $value));
  }
);
<|endoftext|>


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

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

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

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

/// Configuration for the dummy datalink backend. Contains `std::sync::mpsc`
/// channels that are used to communicate with the fake network.
#[derive(Debug)]
pub struct Config {
    /// The fake network will pull packets (or errors) form this `Receiver`
    pub in_packets_rx: Receiver<io::Result<Box<[u8]>>>,

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

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

    out_packets_rx: Option<Receiver<Box<[u8]>>>,
}

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

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

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

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

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

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

struct MockEthernetDataLinkSender {
    out_packets: Sender<Box<[u8]>>,
}

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

    fn send_to(&mut self,
               packet: &EthernetPacket,
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        let buffer = packet.packet().to_vec();
        self.out_packets.send(buffer.into_boxed_slice()).unwrap_or(());
        Some(Ok(()))
    }
}

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

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

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

impl<'a> EthernetDataLinkChannelIterator<'a> for MockEthernetDataLinkChannelIterator {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        match self.in_packets.recv() {
            Ok(result) => {
                // A network event happened. Might be a packet or a simulated error
                match result {
                    Ok(buffer) => {
                        self.used_packets.push(buffer);
                        let buffer_ref = &*self.used_packets[self.used_packets.len() - 1];
                        let packet = EthernetPacket::new(buffer_ref).unwrap();
                        Ok(packet)
                    }
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // The channel supplying fake packets is broken. The user lost/destroyed their
                // inject_handle. This means there will never be any more packets sent to this
                // dummy network. To simulate an idle network we block and sleep forever here.
                loop {
                    thread::sleep(time::Duration::new(10, 0));
                }
            }
        }
    }
}

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

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

#[cfg(test)]
mod tests {
    use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
    use std::io;
    use std::thread::{sleep, spawn};
    use std::time::Duration;

    use packet::{MutablePacket, Packet};
    use packet::ethernet::EthernetPacket;
    use datalink::{EthernetDataLinkReceiver, EthernetDataLinkSender};
    use datalink::Channel::Ethernet;

    #[test]
    fn send_too_small_packet_size() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that it fails to send with too small packet sizes
        assert!(tx.build_and_send(1, 0, &mut |_| {}).is_none());
    }

    #[test]
    fn send_nothing() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending zero packets yields zero packets
        tx.build_and_send(0,
                            20,
                            &mut |_| {
                                panic!("Should not be called");
                            })
            .unwrap()
            .unwrap();
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_one_packet() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending one packet yields one packet
        tx.build_and_send(1,
                            20,
                            &mut |mut pkg| {
                                assert_eq!(pkg.packet().len(), 20);
                                pkg.packet_mut()[0] = 9;
                                pkg.packet_mut()[19] = 201;
                            })
            .unwrap()
            .unwrap();
        let pkg = read_handle.try_recv().expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[0], 9);
        assert_eq!(pkg[19], 201);
    }

    #[test]
    fn send_multiple_packets() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending multiple packets does the correct thing
        let mut closure_counter = 0;
        tx.build_and_send(3,
                            20,
                            &mut |mut pkg| {
                                pkg.packet_mut()[0] = closure_counter;
                                closure_counter += 1;
                            })
            .unwrap()
            .unwrap();
        for i in 0..3 {
            let pkg = read_handle.try_recv().expect("Expected a packet");
            assert_eq!(pkg[0], i);
        }
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_to() {
        let (_, read_handle, mut tx, _) = create_net();
        let mut buffer = vec![0; 20];
        buffer[1] = 34;
        buffer[18] = 76;
        let pkg = EthernetPacket::new(&buffer[..]).unwrap();

        tx.send_to(&pkg, None).unwrap().unwrap();
        let pkg = read_handle.try_recv().expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[1], 34);
        assert_eq!(pkg[18], 76);
    }

    #[test]
    fn read_nothing() {
        let (inject_handle, _, _, mut rx) = create_net();
        let (control_tx, control_rx) = mpsc::channel();
        spawn(move || {
            let mut rx_iter = rx.iter();
            rx_iter.next().expect("Should not happen 1");
            control_tx.send(()).expect("Should not happen 2");
        });
        sleep(Duration::new(0, 1_000_000));
        match control_rx.try_recv() {
            Ok(_) => panic!("Nothing should have arrived"),
            Err(TryRecvError::Disconnected) => panic!("Thread should not have quit"),
            Err(TryRecvError::Empty) => (),
        }
    }

    #[test]
    fn read_one_pkg() {
        let (inject_handle, _, _, mut rx) = create_net();

        let buffer = vec![0; 20];
        inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();

        let mut rx_iter = rx.iter();
        let pkg = rx_iter.next().expect("Expected a packet");
        assert_eq!(pkg.packet().len(), 20);
    }

    #[test]
    fn read_multiple_pkgs() {
        let (inject_handle, _, _, mut rx) = create_net();

        for i in 0..3 {
            let buffer = vec![i; 20];
            inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();
        }

        let mut rx_iter = rx.iter();
        {
            let pkg1 = rx_iter.next().expect("Expected a packet");
            assert_eq!(pkg1.packet()[0], 0);
        }
        {
            let pkg2 = rx_iter.next().expect("Expected a packet");
            assert_eq!(pkg2.packet()[0], 1);
        }
        {
            let pkg3 = rx_iter.next().expect("Expected a packet");
            assert_eq!(pkg3.packet()[0], 2);
        }
    }

    fn create_net()
        -> (Sender<io::Result<Box<[u8]>>>,
            Receiver<Box<[u8]>>,
            Box<EthernetDataLinkSender>,
            Box<EthernetDataLinkReceiver>) {
        let interface = super::dummy_interface(56);
        let mut config = super::Config::default();
        let inject_handle = config.inject_handle().unwrap();
        let read_handle = config.read_handle().unwrap();

        let channel = super::channel(&interface, config);
        let (tx, rx) = match channel {
            Ok(Ethernet(tx, rx)) => (tx, rx),
            _ => panic!("Not a valid channel returned"),
        };
        (inject_handle, read_handle, tx, rx)
    }
}
<commit_msg>Add documentation to dummy backend<commit_after>// Copyright (c) 2016 Linus Färnstrand <faern@faern.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

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

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

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

/// Configuration for the dummy datalink backend. Contains `std::sync::mpsc`
/// channels that are used to communicate with the fake network.
#[derive(Debug)]
pub struct Config {
    in_packets_rx: Receiver<io::Result<Box<[u8]>>>,
    in_packets_tx: Option<Sender<io::Result<Box<[u8]>>>>,

    out_packets_tx: Sender<Box<[u8]>>,
    out_packets_rx: Option<Receiver<Box<[u8]>>>,
}

impl Config {
    /// Creates a new `Config` with the given channels as the backing network.
    /// When using this constructor `inject_handle` and `read_handle` will return `None`.
    /// Those handles must be kept track of elsewhere.
    ///
    /// The `EthernetDataLinkChannelIterator` created by the dummy backend will read packets from
    /// `in_packets`. Both network errors and data can be sent on this channel.
    /// When the `in_packets` channel is closed (`Sender` is dropped)
    /// `EthernetDataLinkChannelIterator::next()` will sleep forever, simlating an idle network.
    ///
    /// The `EthernetDataLinkSender` created by the dummy backend will send all packets sent
    /// through `build_and_send()` and `send_to()` to the `out_packets` channel.
    pub fn new(in_packets: Receiver<io::Result<Box<[u8]>>>,
               out_packets: Sender<Box<[u8]>>)
        -> Config {
        Config {
            in_packets_rx: in_packets,
            in_packets_tx: None,
            out_packets_tx: out_packets,
            out_packets_rx: None,
        }
    }

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

    /// Get the `Receiver` handle where packets sent to the fake network can be read.
    /// Only usable with `Config`s generated from `default()`
    pub fn read_handle(&mut self) -> Option<Receiver<Box<[u8]>>> {
        self.out_packets_rx.take()
    }
}

impl<'a> From<&'a datalink::Config> for Config {
    /// Will not use the `datalink::Config`. This will simply call `dummy::Config::default()`.
    fn from(_config: &datalink::Config) -> Config {
        Config::default()
    }
}

impl Default for Config {
    /// Creates a default config with one input and one output channel. The handles used to inject
    /// to and read form the network can be fetched with `inject_handle()` and `read_handle()`
    fn default() -> Config {
        let (in_tx, in_rx) = mpsc::channel();
        let (out_tx, out_rx) = mpsc::channel();
        Config {
            in_packets_rx: in_rx,
            in_packets_tx: Some(in_tx),
            out_packets_tx: out_tx,
            out_packets_rx: Some(out_rx),
        }
    }
}

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

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

struct MockEthernetDataLinkSender {
    out_packets: Sender<Box<[u8]>>,
}

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

    fn send_to(&mut self,
               packet: &EthernetPacket,
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        let buffer = packet.packet().to_vec();
        self.out_packets.send(buffer.into_boxed_slice()).unwrap_or(());
        Some(Ok(()))
    }
}

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

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

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

impl<'a> EthernetDataLinkChannelIterator<'a> for MockEthernetDataLinkChannelIterator {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        match self.in_packets.recv() {
            Ok(result) => {
                // A network event happened. Might be a packet or a simulated error
                match result {
                    Ok(buffer) => {
                        self.used_packets.push(buffer);
                        let buffer_ref = &*self.used_packets[self.used_packets.len() - 1];
                        let packet = EthernetPacket::new(buffer_ref).unwrap();
                        Ok(packet)
                    }
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // The channel supplying fake packets is broken. The user lost/destroyed their
                // inject_handle. This means there will never be any more packets sent to this
                // dummy network. To simulate an idle network we block and sleep forever here.
                loop {
                    thread::sleep(time::Duration::new(10, 0));
                }
            }
        }
    }
}

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

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

#[cfg(test)]
mod tests {
    use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
    use std::io;
    use std::thread::{sleep, spawn};
    use std::time::Duration;

    use packet::{MutablePacket, Packet};
    use packet::ethernet::EthernetPacket;
    use datalink::{EthernetDataLinkReceiver, EthernetDataLinkSender};
    use datalink::Channel::Ethernet;

    #[test]
    fn send_too_small_packet_size() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that it fails to send with too small packet sizes
        assert!(tx.build_and_send(1, 0, &mut |_| {}).is_none());
    }

    #[test]
    fn send_nothing() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending zero packets yields zero packets
        tx.build_and_send(0,
                            20,
                            &mut |_| {
                                panic!("Should not be called");
                            })
            .unwrap()
            .unwrap();
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_one_packet() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending one packet yields one packet
        tx.build_and_send(1,
                            20,
                            &mut |mut pkg| {
                                assert_eq!(pkg.packet().len(), 20);
                                pkg.packet_mut()[0] = 9;
                                pkg.packet_mut()[19] = 201;
                            })
            .unwrap()
            .unwrap();
        let pkg = read_handle.try_recv().expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[0], 9);
        assert_eq!(pkg[19], 201);
    }

    #[test]
    fn send_multiple_packets() {
        let (_, read_handle, mut tx, _) = create_net();
        // Check that sending multiple packets does the correct thing
        let mut closure_counter = 0;
        tx.build_and_send(3,
                            20,
                            &mut |mut pkg| {
                                pkg.packet_mut()[0] = closure_counter;
                                closure_counter += 1;
                            })
            .unwrap()
            .unwrap();
        for i in 0..3 {
            let pkg = read_handle.try_recv().expect("Expected a packet");
            assert_eq!(pkg[0], i);
        }
        assert!(read_handle.try_recv().is_err());
    }

    #[test]
    fn send_to() {
        let (_, read_handle, mut tx, _) = create_net();
        let mut buffer = vec![0; 20];
        buffer[1] = 34;
        buffer[18] = 76;
        let pkg = EthernetPacket::new(&buffer[..]).unwrap();

        tx.send_to(&pkg, None).unwrap().unwrap();
        let pkg = read_handle.try_recv().expect("Expected one packet to be sent");
        assert!(read_handle.try_recv().is_err());
        assert_eq!(pkg.len(), 20);
        assert_eq!(pkg[1], 34);
        assert_eq!(pkg[18], 76);
    }

    #[test]
    fn read_nothing() {
        let (inject_handle, _, _, mut rx) = create_net();
        let (control_tx, control_rx) = mpsc::channel();
        spawn(move || {
            let mut rx_iter = rx.iter();
            rx_iter.next().expect("Should not happen 1");
            control_tx.send(()).expect("Should not happen 2");
        });
        sleep(Duration::new(0, 1_000_000));
        match control_rx.try_recv() {
            Ok(_) => panic!("Nothing should have arrived"),
            Err(TryRecvError::Disconnected) => panic!("Thread should not have quit"),
            Err(TryRecvError::Empty) => (),
        }
    }

    #[test]
    fn read_one_pkg() {
        let (inject_handle, _, _, mut rx) = create_net();

        let buffer = vec![0; 20];
        inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();

        let mut rx_iter = rx.iter();
        let pkg = rx_iter.next().expect("Expected a packet");
        assert_eq!(pkg.packet().len(), 20);
    }

    #[test]
    fn read_multiple_pkgs() {
        let (inject_handle, _, _, mut rx) = create_net();

        for i in 0..3 {
            let buffer = vec![i; 20];
            inject_handle.send(Ok(buffer.into_boxed_slice())).unwrap();
        }

        let mut rx_iter = rx.iter();
        {
            let pkg1 = rx_iter.next().expect("Expected a packet");
            assert_eq!(pkg1.packet()[0], 0);
        }
        {
            let pkg2 = rx_iter.next().expect("Expected a packet");
            assert_eq!(pkg2.packet()[0], 1);
        }
        {
            let pkg3 = rx_iter.next().expect("Expected a packet");
            assert_eq!(pkg3.packet()[0], 2);
        }
    }

    fn create_net()
        -> (Sender<io::Result<Box<[u8]>>>,
            Receiver<Box<[u8]>>,
            Box<EthernetDataLinkSender>,
            Box<EthernetDataLinkReceiver>) {
        let interface = super::dummy_interface(56);
        let mut config = super::Config::default();
        let inject_handle = config.inject_handle().unwrap();
        let read_handle = config.read_handle().unwrap();

        let channel = super::channel(&interface, config);
        let (tx, rx) = match channel {
            Ok(Ethernet(tx, rx)) => (tx, rx),
            _ => panic!("Not a valid channel returned"),
        };
        (inject_handle, read_handle, tx, rx)
    }
}
<|endoftext|>


<commit_before>use core::sync::atomic::{self, AtomicUsize, Ordering};

use crate::Backoff;

/// A simple stamped lock.
///
/// The state is represented as two `AtomicUsize`: `state_hi` for high bits and `state_lo` for low
/// bits.
pub struct SeqLock {
    /// The high bits of the current state of the lock.
    state_hi: AtomicUsize,

    /// The low bits of the current state of the lock.
    ///
    /// All bits except the least significant one hold the current stamp. When locked, the state_lo
    /// equals 1 and doesn't contain a valid stamp.
    state_lo: AtomicUsize,
}

impl SeqLock {
    pub const fn new() -> Self {
        Self {
            state_hi: AtomicUsize::new(0),
            state_lo: AtomicUsize::new(0),
        }
    }

    /// If not locked, returns the current stamp.
    ///
    /// This method should be called before optimistic reads.
    #[inline]
    pub fn optimistic_read(&self) -> Option<(usize, usize)> {
        // The acquire loads from `state_hi` and `state_lo` synchronize with the release stores in
        // `SeqLockWriteGuard::drop`.
        //
        // As a consequence, we can make sure that (1) all writes within the era of `state_hi - 1`
        // happens before now; and therefore, (2) if `state_lo` is even, all writes within the
        // critical section of (`state_hi`, `state_lo`) happens before now.
        let state_hi = self.state_hi.load(Ordering::Acquire);
        let state_lo = self.state_lo.load(Ordering::Acquire);
        if state_lo == 1 {
            None
        } else {
            Some((state_hi, state_lo))
        }
    }

    /// Returns `true` if the current stamp is equal to `stamp`.
    ///
    /// This method should be called after optimistic reads to check whether they are valid. The
    /// argument `stamp` should correspond to the one returned by method `optimistic_read`.
    #[inline]
    pub fn validate_read(&self, stamp: (usize, usize)) -> bool {
        // Thanks to the fence, if we're noticing any modification to the data at the critical
        // section of `(a, b)`, then the critical section's write of 1 to state_lo should be
        // visible.
        atomic::fence(Ordering::Acquire);

        // So if `state_lo` coincides with `stamp.1`, then either (1) we're noticing no modification
        // to the data after the critical section of `(stamp.0, stamp.1)`, or (2) `state_lo` wrapped
        // around.
        //
        // If (2) is the case, the acquire ordering ensures we see the new value of `state_hi`.
        let state_lo = self.state_lo.load(Ordering::Acquire);

        // If (2) is the case and `state_hi` coincides with `stamp.0`, then `state_hi` also wrapped
        // around, which we give up to correctly validate the read.
        let state_hi = self.state_hi.load(Ordering::Relaxed);

        // Except for the case that both `state_hi` and `state_lo` wrapped around, the following
        // condition implies that we're noticing no modification to the data after the critical
        // section of `(stamp.0, stamp.1)`.
        (state_hi, state_lo) == stamp
    }

    /// Grabs the lock for writing.
    #[inline]
    pub fn write(&'static self) -> SeqLockWriteGuard {
        let backoff = Backoff::new();
        loop {
            let previous = self.state_lo.swap(1, Ordering::Acquire);

            if previous != 1 {
                // To synchronize with the acquire fence in `validate_read` via any modification to
                // the data at the critical section of `(state_hi, previous)`.
                atomic::fence(Ordering::Release);

                return SeqLockWriteGuard {
                    lock: self,
                    state_lo: previous,
                };
            }

            backoff.snooze();
        }
    }
}

/// An RAII guard that releases the lock and increments the stamp when dropped.
pub struct SeqLockWriteGuard {
    /// The parent lock.
    lock: &'static SeqLock,

    /// The stamp before locking.
    state_lo: usize,
}

impl SeqLockWriteGuard {
    /// Releases the lock without incrementing the stamp.
    #[inline]
    pub fn abort(self) {
        self.lock.state_lo.store(self.state_lo, Ordering::Release);
    }
}

impl Drop for SeqLockWriteGuard {
    #[inline]
    fn drop(&mut self) {
        let state_lo = self.state_lo.wrapping_add(2);

        // Increase the high bits if the low bits wrap around.
        //
        // Release ordering for synchronizing with `optimistic_read`.
        if state_lo == 0 {
            let state_hi = self.lock.state_hi.load(Ordering::Relaxed);
            self.lock
                .state_hi
                .store(state_hi.wrapping_add(1), Ordering::Release);
        }

        // Release the lock and increment the stamp.
        //
        // Release ordering for synchronizing with `optimistic_read`.
        self.lock.state_lo.store(state_lo, Ordering::Release);
    }
}

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

    #[test]
    fn test_abort() {
        static LK: SeqLock = SeqLock::new();
        let before = LK.optimistic_read().unwrap();
        {
            let guard = LK.write();
            guard.abort();
        }
        let after = LK.optimistic_read().unwrap();
        assert_eq!(before, after, "aborted write does not update the stamp");
    }
}
<commit_msg>Fix a bug in seq_lock_wide (closes #595)<commit_after>use core::mem;
use core::sync::atomic::{self, AtomicUsize, Ordering};

use crate::Backoff;

/// A simple stamped lock.
///
/// The state is represented as two `AtomicUsize`: `state_hi` for high bits and `state_lo` for low
/// bits.
pub struct SeqLock {
    /// The high bits of the current state of the lock.
    state_hi: AtomicUsize,

    /// The low bits of the current state of the lock.
    ///
    /// All bits except the least significant one hold the current stamp. When locked, the state_lo
    /// equals 1 and doesn't contain a valid stamp.
    state_lo: AtomicUsize,
}

impl SeqLock {
    pub const fn new() -> Self {
        Self {
            state_hi: AtomicUsize::new(0),
            state_lo: AtomicUsize::new(0),
        }
    }

    /// If not locked, returns the current stamp.
    ///
    /// This method should be called before optimistic reads.
    #[inline]
    pub fn optimistic_read(&self) -> Option<(usize, usize)> {
        // The acquire loads from `state_hi` and `state_lo` synchronize with the release stores in
        // `SeqLockWriteGuard::drop`.
        //
        // As a consequence, we can make sure that (1) all writes within the era of `state_hi - 1`
        // happens before now; and therefore, (2) if `state_lo` is even, all writes within the
        // critical section of (`state_hi`, `state_lo`) happens before now.
        let state_hi = self.state_hi.load(Ordering::Acquire);
        let state_lo = self.state_lo.load(Ordering::Acquire);
        if state_lo == 1 {
            None
        } else {
            Some((state_hi, state_lo))
        }
    }

    /// Returns `true` if the current stamp is equal to `stamp`.
    ///
    /// This method should be called after optimistic reads to check whether they are valid. The
    /// argument `stamp` should correspond to the one returned by method `optimistic_read`.
    #[inline]
    pub fn validate_read(&self, stamp: (usize, usize)) -> bool {
        // Thanks to the fence, if we're noticing any modification to the data at the critical
        // section of `(a, b)`, then the critical section's write of 1 to state_lo should be
        // visible.
        atomic::fence(Ordering::Acquire);

        // So if `state_lo` coincides with `stamp.1`, then either (1) we're noticing no modification
        // to the data after the critical section of `(stamp.0, stamp.1)`, or (2) `state_lo` wrapped
        // around.
        //
        // If (2) is the case, the acquire ordering ensures we see the new value of `state_hi`.
        let state_lo = self.state_lo.load(Ordering::Acquire);

        // If (2) is the case and `state_hi` coincides with `stamp.0`, then `state_hi` also wrapped
        // around, which we give up to correctly validate the read.
        let state_hi = self.state_hi.load(Ordering::Relaxed);

        // Except for the case that both `state_hi` and `state_lo` wrapped around, the following
        // condition implies that we're noticing no modification to the data after the critical
        // section of `(stamp.0, stamp.1)`.
        (state_hi, state_lo) == stamp
    }

    /// Grabs the lock for writing.
    #[inline]
    pub fn write(&'static self) -> SeqLockWriteGuard {
        let backoff = Backoff::new();
        loop {
            let previous = self.state_lo.swap(1, Ordering::Acquire);

            if previous != 1 {
                // To synchronize with the acquire fence in `validate_read` via any modification to
                // the data at the critical section of `(state_hi, previous)`.
                atomic::fence(Ordering::Release);

                return SeqLockWriteGuard {
                    lock: self,
                    state_lo: previous,
                };
            }

            backoff.snooze();
        }
    }
}

/// An RAII guard that releases the lock and increments the stamp when dropped.
pub struct SeqLockWriteGuard {
    /// The parent lock.
    lock: &'static SeqLock,

    /// The stamp before locking.
    state_lo: usize,
}

impl SeqLockWriteGuard {
    /// Releases the lock without incrementing the stamp.
    #[inline]
    pub fn abort(self) {
        self.lock.state_lo.store(self.state_lo, Ordering::Release);
        mem::forget(self);
    }
}

impl Drop for SeqLockWriteGuard {
    #[inline]
    fn drop(&mut self) {
        let state_lo = self.state_lo.wrapping_add(2);

        // Increase the high bits if the low bits wrap around.
        //
        // Release ordering for synchronizing with `optimistic_read`.
        if state_lo == 0 {
            let state_hi = self.lock.state_hi.load(Ordering::Relaxed);
            self.lock
                .state_hi
                .store(state_hi.wrapping_add(1), Ordering::Release);
        }

        // Release the lock and increment the stamp.
        //
        // Release ordering for synchronizing with `optimistic_read`.
        self.lock.state_lo.store(state_lo, Ordering::Release);
    }
}

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

    #[test]
    fn test_abort() {
        static LK: SeqLock = SeqLock::new();
        let before = LK.optimistic_read().unwrap();
        {
            let guard = LK.write();
            guard.abort();
        }
        let after = LK.optimistic_read().unwrap();
        assert_eq!(before, after, "aborted write does not update the stamp");
    }
}
<|endoftext|>