<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// 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} }

// 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.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # #![feature(custom_derive, plugin)]
/// # #![plugin(serde_macros)]
/// # #[macro_use] extern crate tarpc;
/// # extern crate serde;
/// # fn main() {}
/// # service! {
/// #[doc="Say hello"]
/// rpc hello(name: String) -> String;
/// # }
/// ```
///
/// Attributes can be attached to each rpc. These attributes
/// will then be attached to the generated `Service` trait's
/// corresponding method, as well as to the `Client` stub's rpcs methods.
///
/// The following items are expanded in the enclosing module:
///
/// * `Service` -- the trait defining the RPC service
/// * `Client` -- a client that makes synchronous requests to the RPC server
/// * `AsyncClient` -- a client that makes asynchronous requests to the RPC server
/// * `Future` -- a handle for asynchronously retrieving the result of an RPC
/// * `serve` -- the function that starts the RPC server
///
/// **Warning**: In addition to the above items, there are a few expanded items that
/// are considered implementation details. As with the above items, shadowing
/// these item names in the enclosing module is likely to break things in confusing
/// ways:
///
/// * `__Server` -- an implementation detail
/// * `__Request` -- an implementation detail
/// * `__Reply` -- an implementation detail
#[macro_export]
macro_rules! service {
    (
        $( $tokens:tt )*
    ) => {
        service_inner! {{
            $( $tokens )*
        }}
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! service_inner {
    // Pattern for when the next rpc has an implicit unit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* );

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> ();
        }
    };
    // Pattern for when the next rpc has an explicit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> $out;
        }
    };
    // Pattern when all return types have been expanded
    (
        { } // none left to expand
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
        )*
    ) => {
        #[doc="Defines the 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)] // because we're testing that the macro expansion compiles
mod syntax_test {
    // Tests a service definition with a fn that takes no args
    mod qux {
        service! {
            rpc hello() -> String;
        }
    }

    // Tests a service definition with an attribute.
    mod bar {
        service! {
            #[doc="Hello bob"]
            rpc baz(s: String) -> String;
        }
    }

    // Tests a service with implicit return types.
    mod no_return {
        service! {
            rpc ack();
            rpc apply(foo: String) -> i32;
            rpc bi_consume(bar: String, baz: u64);
        }
    }
}

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

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

    service! {
        rpc add(x: i32, y: i32) -> i32;
    }

    struct Server;

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

    #[test]
    fn simple() {
        let handle = serve( "localhost:0", Server, test_timeout()).unwrap();
        let client = Client::new(handle.local_addr(), None).unwrap();
        assert_eq!(3, client.add(1, 2).unwrap());
        drop(client);
        handle.shutdown();
    }

    #[test]
    fn simple_async() {
        let handle = 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());
        drop(client);
        handle.shutdown();
    }

    // Tests that a server can be wrapped in an Arc; no need to run, just compile
    #[allow(dead_code)]
    fn serve_arc_server() {
        let _ = serve("localhost:0", ::std::sync::Arc::new(Server), None);
    }
}

#[cfg(test)]
#[allow(dead_code)] // generated Client isn't used in this benchmark
mod benchmark {
    extern crate env_logger;
    use ServeHandle;
    use std::sync::{Arc, Mutex};
    use test::Bencher;

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

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

    // Prevents resource exhaustion when benching
    lazy_static! {
        static ref HANDLE: Arc<Mutex<ServeHandle>> = {
            let handle = serve("localhost:0", HelloServer, None).unwrap();
            Arc::new(Mutex::new(handle))
        };
        static ref CLIENT: Arc<Mutex<AsyncClient>> = {
            let addr = HANDLE.lock().unwrap().local_addr().clone();
            let client = AsyncClient::new(addr, None).unwrap();
            Arc::new(Mutex::new(client))
        };
    }

    #[bench]
    fn hello(bencher: &mut Bencher) {
        let _ = env_logger::init();
        let client = CLIENT.lock().unwrap();
        let concurrency = 100;
        let mut futures = Vec::with_capacity(concurrency);
        let mut count = 0;
        bencher.iter(|| {
            futures.push(client.hello("Bob".into()));
            count += 1;
            if count % concurrency == 0 {
                // We can't block on each rpc call, otherwise we'd be
                // benchmarking latency instead of throughput. It's also
                // not ideal to call more than one rpc per iteration, because
                // it makes the output of the bencher harder to parse (you have
                // to mentally divide the number by `concurrency` to get
                // the ns / iter for one rpc
                for f in futures.drain(..) {
                    f.get().unwrap();
                }
            }
        });
    }
}
<commit_msg>Add TryClone to Client/AsyncClient<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// 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} }

// 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.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # #![feature(custom_derive, plugin)]
/// # #![plugin(serde_macros)]
/// # #[macro_use] extern crate tarpc;
/// # extern crate serde;
/// # fn main() {}
/// # service! {
/// #[doc="Say hello"]
/// rpc hello(name: String) -> String;
/// # }
/// ```
///
/// Attributes can be attached to each rpc. These attributes
/// will then be attached to the generated `Service` trait's
/// corresponding method, as well as to the `Client` stub's rpcs methods.
///
/// The following items are expanded in the enclosing module:
///
/// * `Service` -- the trait defining the RPC service
/// * `Client` -- a client that makes synchronous requests to the RPC server
/// * `AsyncClient` -- a client that makes asynchronous requests to the RPC server
/// * `Future` -- a handle for asynchronously retrieving the result of an RPC
/// * `serve` -- the function that starts the RPC server
///
/// **Warning**: In addition to the above items, there are a few expanded items that
/// are considered implementation details. As with the above items, shadowing
/// these item names in the enclosing module is likely to break things in confusing
/// ways:
///
/// * `__Server` -- an implementation detail
/// * `__Request` -- an implementation detail
/// * `__Reply` -- an implementation detail
#[macro_export]
macro_rules! service {
    (
        $( $tokens:tt )*
    ) => {
        service_inner! {{
            $( $tokens )*
        }}
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! service_inner {
    // Pattern for when the next rpc has an implicit unit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* );

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> ();
        }
    };
    // Pattern for when the next rpc has an explicit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> $out;
        }
    };
    // Pattern when all return types have been expanded
    (
        { } // none left to expand
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
        )*
    ) => {
        #[doc="Defines the 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="Attempt to clone the client object. This might fail if the underlying TcpStream clone fails"]
            pub fn try_clone(&self) -> ::std::io::Result<Self> {
                Ok(Client(try!(self.0.try_clone())))
            }
        }

        #[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
                )*
            );

            #[doc="Attempt to clone the client object. This might fail if the underlying TcpStream clone fails"]
            pub fn try_clone(&self) -> ::std::io::Result<Self> {
                Ok(AsyncClient(try!(self.0.try_clone())))
            }
        }

        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)] // because we're testing that the macro expansion compiles
mod syntax_test {
    // Tests a service definition with a fn that takes no args
    mod qux {
        service! {
            rpc hello() -> String;
        }
    }

    // Tests a service definition with an attribute.
    mod bar {
        service! {
            #[doc="Hello bob"]
            rpc baz(s: String) -> String;
        }
    }

    // Tests a service with implicit return types.
    mod no_return {
        service! {
            rpc ack();
            rpc apply(foo: String) -> i32;
            rpc bi_consume(bar: String, baz: u64);
        }
    }
}

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

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

    service! {
        rpc add(x: i32, y: i32) -> i32;
    }

    struct Server;

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

    #[test]
    fn simple() {
        let handle = serve( "localhost:0", Server, test_timeout()).unwrap();
        let client = Client::new(handle.local_addr(), None).unwrap();
        assert_eq!(3, client.add(1, 2).unwrap());
        drop(client);
        handle.shutdown();
    }

    #[test]
    fn simple_async() {
        let handle = 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());
        drop(client);
        handle.shutdown();
    }

    #[test]
    fn try_clone() {
        let handle = serve( "localhost:0", Server, test_timeout()).unwrap();
        let client1 = Client::new(handle.local_addr(), None).unwrap();
        let client2 = client1.try_clone().unwrap();
        assert_eq!(3, client1.add(1, 2).unwrap());
        assert_eq!(3, client2.add(1, 2).unwrap());
    }

    #[test]
    fn async_try_clone() {
        let handle = serve("localhost:0", Server, test_timeout()).unwrap();
        let client1 = AsyncClient::new(handle.local_addr(), None).unwrap();
        let client2 = client1.try_clone().unwrap();
        assert_eq!(3, client1.add(1, 2).get().unwrap());
        assert_eq!(3, client2.add(1, 2).get().unwrap());
    }

    // Tests that a server can be wrapped in an Arc; no need to run, just compile
    #[allow(dead_code)]
    fn serve_arc_server() {
        let _ = serve("localhost:0", ::std::sync::Arc::new(Server), None);
    }
}

#[cfg(test)]
#[allow(dead_code)] // generated Client isn't used in this benchmark
mod benchmark {
    extern crate env_logger;
    use ServeHandle;
    use std::sync::{Arc, Mutex};
    use test::Bencher;

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

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

    // Prevents resource exhaustion when benching
    lazy_static! {
        static ref HANDLE: Arc<Mutex<ServeHandle>> = {
            let handle = serve("localhost:0", HelloServer, None).unwrap();
            Arc::new(Mutex::new(handle))
        };
        static ref CLIENT: Arc<Mutex<AsyncClient>> = {
            let addr = HANDLE.lock().unwrap().local_addr().clone();
            let client = AsyncClient::new(addr, None).unwrap();
            Arc::new(Mutex::new(client))
        };
    }

    #[bench]
    fn hello(bencher: &mut Bencher) {
        let _ = env_logger::init();
        let client = CLIENT.lock().unwrap();
        let concurrency = 100;
        let mut futures = Vec::with_capacity(concurrency);
        let mut count = 0;
        bencher.iter(|| {
            futures.push(client.hello("Bob".into()));
            count += 1;
            if count % concurrency == 0 {
                // We can't block on each rpc call, otherwise we'd be
                // benchmarking latency instead of throughput. It's also
                // not ideal to call more than one rpc per iteration, because
                // it makes the output of the bencher harder to parse (you have
                // to mentally divide the number by `concurrency` to get
                // the ns / iter for one rpc
                for f in futures.drain(..) {
                    f.get().unwrap();
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// 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} }

// 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.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # #![feature(custom_derive, plugin)]
/// # #![plugin(serde_macros)]
/// # #[macro_use] extern crate tarpc;
/// # extern crate serde;
/// # fn main() {}
/// # service! {
/// #[doc="Say hello"]
/// rpc hello(name: String) -> String;
/// # }
/// ```
///
/// Attributes can be attached to each rpc. These attributes
/// will then be attached to the generated `Service` trait's
/// corresponding method, as well as to the `Client` stub's rpcs methods.
///
/// The following items are expanded in the enclosing module:
///
/// * `Service` -- the trait defining the RPC service
/// * `Client` -- a client that makes synchronous requests to the RPC server
/// * `AsyncClient` -- a client that makes asynchronous requests to the RPC server
/// * `Future` -- a handle for asynchronously retrieving the result of an RPC
/// * `serve` -- the function that starts the RPC server
///
/// **Warning**: In addition to the above items, there are a few expanded items that
/// are considered implementation details. As with the above items, shadowing
/// these item names in the enclosing module is likely to break things in confusing
/// ways:
///
/// * `__Server` -- an implementation detail
/// * `__Request` -- an implementation detail
/// * `__Reply` -- an implementation detail
#[macro_export]
macro_rules! service {
    (
        $( $tokens:tt )*
    ) => {
        service_inner! {{
            $( $tokens )*
        }}
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! service_inner {
    // Pattern for when the next rpc has an implicit unit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* );

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> ();
        }
    };
    // Pattern for when the next rpc has an explicit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> $out;
        }
    };
    // Pattern when all return types have been expanded
    (
        { } // none left to expand
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
        )*
    ) => {
        #[doc="Defines the 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)] // because we're testing that the macro expansion compiles
mod test {
    extern crate env_logger;
    use std::time::Duration;

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

    #[derive(PartialEq, Debug, Serialize, Deserialize)]
    pub struct Foo {
        pub message: String
    }
    mod my_server {
        use super::Foo;
        service! {
            rpc hello(foo: Foo) -> Foo;
            rpc add(x: i32, y: i32) -> i32;
        }
    }

    struct Server;
    impl my_server::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() {
        my_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 = my_server::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 = my_server::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
    mod qux {
        service! {
            rpc hello() -> String;
        }
    }

    /// Tests a service definition with an import
    mod foo {
        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
    mod bar {
        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
    #[deny(missing_docs)]
    #[allow(unused)]
    mod baz {
        use std::collections::HashMap;
        pub struct Debuggable;
        service! {
            #[doc="Hello bob"]
            #[inline(always)]
            rpc baz(s: String) -> HashMap<String, String>;
        }
    }

    mod no_return {
        service! {
            rpc ack();
            rpc apply(foo: String) -> i32;
            rpc bi_consume(bar: String, baz: u64);
        }
    }
}

#[cfg(test)]
#[allow(dead_code)] // because we're not using all expanded types
mod benchmarks {
    extern crate env_logger;
    use ServeHandle;
    use std::sync::{Arc, Mutex};
    use test::Bencher;

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

    struct HelloServer;

    impl hi::Service for HelloServer {
        fn hello(&self, s: String) -> String {
            format!("Hello, {}!", s)
        }
    }
    // Prevents resource exhaustion when benching
    lazy_static! {
        static ref HANDLE: Arc<Mutex<ServeHandle>> = {
            let handle = hi::serve("localhost:0", HelloServer, None).unwrap();
            Arc::new(Mutex::new(handle))
        };
        static ref CLIENT: Arc<Mutex<hi::AsyncClient>> = {
            let addr = HANDLE.lock().unwrap().local_addr().clone();
            let client = hi::AsyncClient::new(addr, None).unwrap();
            Arc::new(Mutex::new(client))
        };
    }

    #[bench]
    fn hello(bencher: &mut Bencher) {
        let _ = env_logger::init();
        let client = CLIENT.lock().unwrap();
        let concurrency = 100;
        let mut futures = Vec::with_capacity(concurrency);
        let mut count = 0;
        bencher.iter(|| {
            futures.push(client.hello("Bob".into()));
            count += 1;
            if count % concurrency == 0 {
                // We can't block on each rpc call, otherwise we'd be
                // benchmarking latency instead of throughput. It's also
                // not ideal to call more than one rpc per iteration, because
                // it makes the output of the bencher harder to parse (you have
                // to mentally divide the number by `concurrency` to get
                // the ns / iter for one rpc
                for f in futures.drain(..) {
                    f.get().unwrap();
                }
            }
        });
    }
}
<commit_msg>Organize macro module tests<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// 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} }

// 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.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # #![feature(custom_derive, plugin)]
/// # #![plugin(serde_macros)]
/// # #[macro_use] extern crate tarpc;
/// # extern crate serde;
/// # fn main() {}
/// # service! {
/// #[doc="Say hello"]
/// rpc hello(name: String) -> String;
/// # }
/// ```
///
/// Attributes can be attached to each rpc. These attributes
/// will then be attached to the generated `Service` trait's
/// corresponding method, as well as to the `Client` stub's rpcs methods.
///
/// The following items are expanded in the enclosing module:
///
/// * `Service` -- the trait defining the RPC service
/// * `Client` -- a client that makes synchronous requests to the RPC server
/// * `AsyncClient` -- a client that makes asynchronous requests to the RPC server
/// * `Future` -- a handle for asynchronously retrieving the result of an RPC
/// * `serve` -- the function that starts the RPC server
///
/// **Warning**: In addition to the above items, there are a few expanded items that
/// are considered implementation details. As with the above items, shadowing
/// these item names in the enclosing module is likely to break things in confusing
/// ways:
///
/// * `__Server` -- an implementation detail
/// * `__Request` -- an implementation detail
/// * `__Reply` -- an implementation detail
#[macro_export]
macro_rules! service {
    (
        $( $tokens:tt )*
    ) => {
        service_inner! {{
            $( $tokens )*
        }}
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! service_inner {
    // Pattern for when the next rpc has an implicit unit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* );

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> ();
        }
    };
    // Pattern for when the next rpc has an explicit return type
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        service_inner! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> $out;
        }
    };
    // Pattern when all return types have been expanded
    (
        { } // none left to expand
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
        )*
    ) => {
        #[doc="Defines the 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)] // because we're testing that the macro expansion compiles
mod syntax_tests {
    // Tests a service definition with a fn that takes no args
    mod qux {
        service! {
            rpc hello() -> String;
        }
    }

    // Tests a service definition with an attribute.
    mod bar {
        service! {
            #[doc="Hello bob"]
            rpc baz(s: String) -> String;
        }
    }

    // Tests a service with implicit return types.
    mod no_return {
        service! {
            rpc ack();
            rpc apply(foo: String) -> i32;
            rpc bi_consume(bar: String, baz: u64);
        }
    }
}

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

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

    service! {
        rpc add(x: i32, y: i32) -> i32;
    }

    struct Server;

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

    #[test]
    fn simple() {
        let handle = serve( "localhost:0", Server, test_timeout()).unwrap();
        let client = Client::new(handle.local_addr(), None).unwrap();
        assert_eq!(3, client.add(1, 2).unwrap());
        drop(client);
        handle.shutdown();
    }

    #[test]
    fn simple_async() {
        let handle = 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());
        drop(client);
        handle.shutdown();
    }

    // Tests that a server can be wrapped in an Arc; no need to run, just compile
    #[allow(dead_code)]
    fn serve_arc_server() {
        let _ = serve("localhost:0", ::std::sync::Arc::new(Server), None);
    }
}

#[cfg(test)]
#[allow(dead_code)] // generated Client isn't used in this benchmark
mod benchmark {
    extern crate env_logger;
    use ServeHandle;
    use std::sync::{Arc, Mutex};
    use test::Bencher;

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

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

    // Prevents resource exhaustion when benching
    lazy_static! {
        static ref HANDLE: Arc<Mutex<ServeHandle>> = {
            let handle = serve("localhost:0", HelloServer, None).unwrap();
            Arc::new(Mutex::new(handle))
        };
        static ref CLIENT: Arc<Mutex<AsyncClient>> = {
            let addr = HANDLE.lock().unwrap().local_addr().clone();
            let client = AsyncClient::new(addr, None).unwrap();
            Arc::new(Mutex::new(client))
        };
    }

    #[bench]
    fn hello(bencher: &mut Bencher) {
        let _ = env_logger::init();
        let client = CLIENT.lock().unwrap();
        let concurrency = 100;
        let mut futures = Vec::with_capacity(concurrency);
        let mut count = 0;
        bencher.iter(|| {
            futures.push(client.hello("Bob".into()));
            count += 1;
            if count % concurrency == 0 {
                // We can't block on each rpc call, otherwise we'd be
                // benchmarking latency instead of throughput. It's also
                // not ideal to call more than one rpc per iteration, because
                // it makes the output of the bencher harder to parse (you have
                // to mentally divide the number by `concurrency` to get
                // the ns / iter for one rpc
                for f in futures.drain(..) {
                    f.get().unwrap();
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin, WorkDoneProgressEnd,
    WorkDoneProgressReport,
};
use std::{
    collections::HashMap,
    io::{BufRead, BufReader},
    path::PathBuf,
    process::{Command, Stdio},
    sync::Arc,
    thread::JoinHandle,
    time::Instant,
};

mod conv;

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic, SuggestedFix};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckOptions {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    pub task_recv: Receiver<CheckTask>,
    pub state: Arc<CheckState>,
    cmd_send: Option<Sender<CheckCommand>>,
    handle: Option<JoinHandle<()>>,
}

impl CheckWatcher {
    pub fn new(options: &CheckOptions, workspace_root: PathBuf) -> CheckWatcher {
        let options = options.clone();
        let state = Arc::new(CheckState::new());

        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = std::thread::spawn(move || {
            let mut check = CheckWatcherThread::new(options, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send: Some(cmd_send), handle: Some(handle), state }
    }

    /// Returns a CheckWatcher that doesn't actually do anything
    pub fn dummy() -> CheckWatcher {
        let state = Arc::new(CheckState::new());
        CheckWatcher { task_recv: never(), cmd_send: None, handle: None, state }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        if let Some(cmd_send) = &self.cmd_send {
            cmd_send.send(CheckCommand::Update).unwrap();
        }
    }
}

impl std::ops::Drop for CheckWatcher {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Take the sender out of the option
            let recv = self.cmd_send.take();

            // Dropping the sender finishes the thread loop
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

#[derive(Clone, Debug)]
pub struct CheckState {
    diagnostic_collection: HashMap<Url, Vec<Diagnostic>>,
    suggested_fix_collection: HashMap<Url, Vec<SuggestedFix>>,
}

impl CheckState {
    fn new() -> CheckState {
        CheckState {
            diagnostic_collection: HashMap::new(),
            suggested_fix_collection: HashMap::new(),
        }
    }

    /// Clear the cached diagnostics, and schedule updating diagnostics by the
    /// server, to clear stale results.
    pub fn clear(&mut self) -> Vec<Url> {
        let cleared_files: Vec<Url> = self.diagnostic_collection.keys().cloned().collect();
        self.diagnostic_collection.clear();
        self.suggested_fix_collection.clear();
        cleared_files
    }

    pub fn diagnostics_for(&self, uri: &Url) -> Option<&[Diagnostic]> {
        self.diagnostic_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn fixes_for(&self, uri: &Url) -> Option<&[SuggestedFix]> {
        self.suggested_fix_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn add_diagnostic_with_fixes(&mut self, file_uri: Url, diagnostic: DiagnosticWithFixes) {
        for fix in diagnostic.suggested_fixes {
            self.add_suggested_fix_for_diagnostic(fix, &diagnostic.diagnostic);
        }
        self.add_diagnostic(file_uri, diagnostic.diagnostic);
    }

    fn add_diagnostic(&mut self, file_uri: Url, diagnostic: Diagnostic) {
        let diagnostics = self.diagnostic_collection.entry(file_uri).or_default();

        // If we're building multiple targets it's possible we've already seen this diagnostic
        let is_duplicate = diagnostics.iter().any(|d| are_diagnostics_equal(d, &diagnostic));
        if is_duplicate {
            return;
        }

        diagnostics.push(diagnostic);
    }

    fn add_suggested_fix_for_diagnostic(
        &mut self,
        mut suggested_fix: SuggestedFix,
        diagnostic: &Diagnostic,
    ) {
        let file_uri = suggested_fix.location.uri.clone();
        let file_suggestions = self.suggested_fix_collection.entry(file_uri).or_default();

        let existing_suggestion: Option<&mut SuggestedFix> =
            file_suggestions.iter_mut().find(|s| s == &&suggested_fix);
        if let Some(existing_suggestion) = existing_suggestion {
            // The existing suggestion also applies to this new diagnostic
            existing_suggestion.diagnostics.push(diagnostic.clone());
        } else {
            // We haven't seen this suggestion before
            suggested_fix.diagnostics.push(diagnostic.clone());
            file_suggestions.push(suggested_fix);
        }
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic(Url, DiagnosticWithFixes),

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckOptions,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherThread {
    fn new(options: CheckOptions, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // Replace with a dummy watcher first so we drop the original and wait for completion
                std::mem::replace(&mut self.watcher, WatchThread::dummy());

                // Then create the actual new watcher
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result =
                    match map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root) {
                        Some(map_result) => map_result,
                        None => return,
                    };

                let MappedRustDiagnostic { location, diagnostic, suggested_fixes } = map_result;

                let diagnostic = DiagnosticWithFixes { diagnostic, suggested_fixes };
                task_send.send(CheckTask::AddDiagnostic(location.uri, diagnostic)).unwrap();
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    suggested_fixes: Vec<SuggestedFix>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    handle: Option<JoinHandle<()>>,
    message_recv: Receiver<CheckEvent>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { handle: None, message_recv: never() }
    }

    fn new(options: &CheckOptions, workspace_root: &PathBuf) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.to_string_lossy()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let enabled = options.enable;
        let handle = std::thread::spawn(move || {
            if !enabled {
                return;
            }

            let mut command = Command::new("cargo")
                .args(&args)
                .stdout(Stdio::piped())
                .stderr(Stdio::null())
                .spawn()
                .expect("couldn't launch cargo");

            // If we trigger an error here, we will do so in the loop instead,
            // which will break out of the loop, and continue the shutdown
            let _ = message_send.send(CheckEvent::Begin);

            // We manually read a line at a time, instead of using serde's
            // stream deserializers, because the deserializer cannot recover
            // from an error, resulting in it getting stuck, because we try to
            // be resillient against failures.
            //
            // Because cargo only outputs one JSON object per line, we can
            // simply skip a line if it doesn't parse, which just ignores any
            // erroneus output.
            let stdout = BufReader::new(command.stdout.take().unwrap());
            for line in stdout.lines() {
                let line = match line {
                    Ok(line) => line,
                    Err(err) => {
                        log::error!("Couldn't read line from cargo: {}", err);
                        continue;
                    }
                };

                let message = serde_json::from_str::<cargo_metadata::Message>(&line);
                let message = match message {
                    Ok(message) => message,
                    Err(err) => {
                        log::error!(
                            "Invalid json from cargo check, ignoring ({}): {:?} ",
                            err,
                            line
                        );
                        continue;
                    }
                };

                // Skip certain kinds of messages to only spend time on what's useful
                match &message {
                    Message::CompilerArtifact(artifact) if artifact.fresh => continue,
                    Message::BuildScriptExecuted(_) => continue,
                    Message::Unknown => continue,
                    _ => {}
                }

                match message_send.send(CheckEvent::Msg(message)) {
                    Ok(()) => {}
                    Err(_err) => {
                        // The send channel was closed, so we want to shutdown
                        break;
                    }
                }
            }

            // We can ignore any error here, as we are already in the progress
            // of shutting down.
            let _ = message_send.send(CheckEvent::End);

            // It is okay to ignore the result, as it only errors if the process is already dead
            let _ = command.kill();

            // Again, we don't care about the exit status so just ignore the result
            let _ = command.wait();
        });
        WatchThread { handle: Some(handle), message_recv }
    }
}

impl std::ops::Drop for WatchThread {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Replace our reciever with dummy one, so we can drop and close the
            // one actually communicating with the thread
            let recv = std::mem::replace(&mut self.message_recv, never());

            // Dropping the original reciever initiates thread sub-process shutdown
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

fn are_diagnostics_equal(left: &Diagnostic, right: &Diagnostic) -> bool {
    left.source == right.source
        && left.severity == right.severity
        && left.range == right.range
        && left.message == right.message
}
<commit_msg>Prevent child cargo process from messing with our stdin<commit_after>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin, WorkDoneProgressEnd,
    WorkDoneProgressReport,
};
use std::{
    collections::HashMap,
    io::{BufRead, BufReader},
    path::PathBuf,
    process::{Command, Stdio},
    sync::Arc,
    thread::JoinHandle,
    time::Instant,
};

mod conv;

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic, SuggestedFix};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckOptions {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    pub task_recv: Receiver<CheckTask>,
    pub state: Arc<CheckState>,
    cmd_send: Option<Sender<CheckCommand>>,
    handle: Option<JoinHandle<()>>,
}

impl CheckWatcher {
    pub fn new(options: &CheckOptions, workspace_root: PathBuf) -> CheckWatcher {
        let options = options.clone();
        let state = Arc::new(CheckState::new());

        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = std::thread::spawn(move || {
            let mut check = CheckWatcherThread::new(options, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send: Some(cmd_send), handle: Some(handle), state }
    }

    /// Returns a CheckWatcher that doesn't actually do anything
    pub fn dummy() -> CheckWatcher {
        let state = Arc::new(CheckState::new());
        CheckWatcher { task_recv: never(), cmd_send: None, handle: None, state }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        if let Some(cmd_send) = &self.cmd_send {
            cmd_send.send(CheckCommand::Update).unwrap();
        }
    }
}

impl std::ops::Drop for CheckWatcher {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Take the sender out of the option
            let recv = self.cmd_send.take();

            // Dropping the sender finishes the thread loop
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

#[derive(Clone, Debug)]
pub struct CheckState {
    diagnostic_collection: HashMap<Url, Vec<Diagnostic>>,
    suggested_fix_collection: HashMap<Url, Vec<SuggestedFix>>,
}

impl CheckState {
    fn new() -> CheckState {
        CheckState {
            diagnostic_collection: HashMap::new(),
            suggested_fix_collection: HashMap::new(),
        }
    }

    /// Clear the cached diagnostics, and schedule updating diagnostics by the
    /// server, to clear stale results.
    pub fn clear(&mut self) -> Vec<Url> {
        let cleared_files: Vec<Url> = self.diagnostic_collection.keys().cloned().collect();
        self.diagnostic_collection.clear();
        self.suggested_fix_collection.clear();
        cleared_files
    }

    pub fn diagnostics_for(&self, uri: &Url) -> Option<&[Diagnostic]> {
        self.diagnostic_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn fixes_for(&self, uri: &Url) -> Option<&[SuggestedFix]> {
        self.suggested_fix_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn add_diagnostic_with_fixes(&mut self, file_uri: Url, diagnostic: DiagnosticWithFixes) {
        for fix in diagnostic.suggested_fixes {
            self.add_suggested_fix_for_diagnostic(fix, &diagnostic.diagnostic);
        }
        self.add_diagnostic(file_uri, diagnostic.diagnostic);
    }

    fn add_diagnostic(&mut self, file_uri: Url, diagnostic: Diagnostic) {
        let diagnostics = self.diagnostic_collection.entry(file_uri).or_default();

        // If we're building multiple targets it's possible we've already seen this diagnostic
        let is_duplicate = diagnostics.iter().any(|d| are_diagnostics_equal(d, &diagnostic));
        if is_duplicate {
            return;
        }

        diagnostics.push(diagnostic);
    }

    fn add_suggested_fix_for_diagnostic(
        &mut self,
        mut suggested_fix: SuggestedFix,
        diagnostic: &Diagnostic,
    ) {
        let file_uri = suggested_fix.location.uri.clone();
        let file_suggestions = self.suggested_fix_collection.entry(file_uri).or_default();

        let existing_suggestion: Option<&mut SuggestedFix> =
            file_suggestions.iter_mut().find(|s| s == &&suggested_fix);
        if let Some(existing_suggestion) = existing_suggestion {
            // The existing suggestion also applies to this new diagnostic
            existing_suggestion.diagnostics.push(diagnostic.clone());
        } else {
            // We haven't seen this suggestion before
            suggested_fix.diagnostics.push(diagnostic.clone());
            file_suggestions.push(suggested_fix);
        }
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic(Url, DiagnosticWithFixes),

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckOptions,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherThread {
    fn new(options: CheckOptions, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // Replace with a dummy watcher first so we drop the original and wait for completion
                std::mem::replace(&mut self.watcher, WatchThread::dummy());

                // Then create the actual new watcher
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result =
                    match map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root) {
                        Some(map_result) => map_result,
                        None => return,
                    };

                let MappedRustDiagnostic { location, diagnostic, suggested_fixes } = map_result;

                let diagnostic = DiagnosticWithFixes { diagnostic, suggested_fixes };
                task_send.send(CheckTask::AddDiagnostic(location.uri, diagnostic)).unwrap();
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    suggested_fixes: Vec<SuggestedFix>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    handle: Option<JoinHandle<()>>,
    message_recv: Receiver<CheckEvent>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { handle: None, message_recv: never() }
    }

    fn new(options: &CheckOptions, workspace_root: &PathBuf) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.to_string_lossy()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let enabled = options.enable;
        let handle = std::thread::spawn(move || {
            if !enabled {
                return;
            }

            let mut command = Command::new("cargo")
                .args(&args)
                .stdout(Stdio::piped())
                .stderr(Stdio::null())
                .stdin(Stdio::null())
                .spawn()
                .expect("couldn't launch cargo");

            // If we trigger an error here, we will do so in the loop instead,
            // which will break out of the loop, and continue the shutdown
            let _ = message_send.send(CheckEvent::Begin);

            // We manually read a line at a time, instead of using serde's
            // stream deserializers, because the deserializer cannot recover
            // from an error, resulting in it getting stuck, because we try to
            // be resillient against failures.
            //
            // Because cargo only outputs one JSON object per line, we can
            // simply skip a line if it doesn't parse, which just ignores any
            // erroneus output.
            let stdout = BufReader::new(command.stdout.take().unwrap());
            for line in stdout.lines() {
                let line = match line {
                    Ok(line) => line,
                    Err(err) => {
                        log::error!("Couldn't read line from cargo: {}", err);
                        continue;
                    }
                };

                let message = serde_json::from_str::<cargo_metadata::Message>(&line);
                let message = match message {
                    Ok(message) => message,
                    Err(err) => {
                        log::error!(
                            "Invalid json from cargo check, ignoring ({}): {:?} ",
                            err,
                            line
                        );
                        continue;
                    }
                };

                // Skip certain kinds of messages to only spend time on what's useful
                match &message {
                    Message::CompilerArtifact(artifact) if artifact.fresh => continue,
                    Message::BuildScriptExecuted(_) => continue,
                    Message::Unknown => continue,
                    _ => {}
                }

                match message_send.send(CheckEvent::Msg(message)) {
                    Ok(()) => {}
                    Err(_err) => {
                        // The send channel was closed, so we want to shutdown
                        break;
                    }
                }
            }

            // We can ignore any error here, as we are already in the progress
            // of shutting down.
            let _ = message_send.send(CheckEvent::End);

            // It is okay to ignore the result, as it only errors if the process is already dead
            let _ = command.kill();

            // Again, we don't care about the exit status so just ignore the result
            let _ = command.wait();
        });
        WatchThread { handle: Some(handle), message_recv }
    }
}

impl std::ops::Drop for WatchThread {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Replace our reciever with dummy one, so we can drop and close the
            // one actually communicating with the thread
            let recv = std::mem::replace(&mut self.message_recv, never());

            // Dropping the original reciever initiates thread sub-process shutdown
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

fn are_diagnostics_equal(left: &Diagnostic, right: &Diagnostic) -> bool {
    left.source == right.source
        && left.severity == right.severity
        && left.range == right.range
        && left.message == right.message
}
<|endoftext|>