<commit_before>//! File holding the FormHashMap type and associated tests
//!
//! Author: [Boris](mailto:boris@humanenginuity.com)
//! Version: 1.1

// =======================================================================
// ATTRIBUTES
// =======================================================================
#![allow(dead_code)]

// =======================================================================
// LIBRARY IMPORTS
// =======================================================================
use std::collections::HashMap;
use std::error::Error;
use std::fmt::Debug;
use std::io::Read;
use std::marker::PhantomData;

use rocket;
use rocket::{ Request, Data };
use rocket::data::FromData;
use rocket::http::Status;
use rocket::outcome::IntoOutcome;
use rocket::request::{ FromForm, FromFormValue, FormItems };

use types::OneOrMany;

// =======================================================================
// FORMVALUES INNER STRUCT DEFINITION
// =======================================================================
/// Private structure to hold a random number of values
/// providing utility methods to `push()`or `get()` the underlying data
struct FormValues<'v, T: 'v> {
    value: Option<OneOrMany<T>>,
    _phantom: PhantomData<&'v T>,
}

/// Generic structure to hold the value in a key/value pair
/// since we can have an arbitrary number of values associated to a key
impl<'v, T: 'v> FormValues<'v, T> {
    fn new() -> FormValues<'v, T> {
        FormValues {
            value: None,
            _phantom: PhantomData
        }
    }

    /// Set/add a new value
    fn push(&mut self, new_value: T) {
        let old_value = ::std::mem::replace(&mut self.value, None);
        self._phantom = PhantomData;
        self.value = match old_value {
            None => {
                Some(OneOrMany::One(new_value))
            }
            Some(OneOrMany::One(existing_value)) => {
                Some(OneOrMany::Many(vec![existing_value, new_value]))
            }
            Some(OneOrMany::Many(mut vec)) => {
                vec.push(new_value);
                Some(OneOrMany::Many(vec))
            }
        }
    }

    /// Get a reference to this object's value
    fn get(&self) -> Option<&OneOrMany<T>> {
        self.value.as_ref()
    }

    /// Consume this object (since it takes `self`) and move out its content
    fn into_inner(self) -> Option<OneOrMany<T>> {
        self.value
    }
}

impl<'v, T: Debug + 'v> Debug for FormValues<'v, T> {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "FormValues {{ values: {:?} }}", self.value)
    }
}

// =======================================================================
// FORMHASHMAP (implement FormData and FromForm)
// =======================================================================

/// A `FromData` type that creates a map of the key/value pairs from a
/// `x-www-form-urlencoded` form `string`.
pub struct FormHashMap<'f> {
    form_string: String,
    form_hash: HashMap<&'f str, FormValues<'f, String>>,
    _phantom: PhantomData<&'f String>,
}

impl<'f> FormHashMap<'f> {
    /// Get a reference for the value (or values) associated with `key`.
    pub fn get(&'f self, key: &str) -> Option<&OneOrMany<String>> {
        // Note: The map method takes the self argument by value, consuming the original
        // `HashMap::get()` returns a reference to the value corresponding to the key
        // so we then consume *that* reference with `map`
        self.form_hash
            .get(key) // Option<&FormValues<String>> corresponding to the key
            .map(|form_value| form_value.get())
            .unwrap() // converted to an Option<&OneOrMany<String>> to 'hide' FormValues
    }

    /// Returns the raw form string that was used to parse the encapsulated
    /// object.
    pub fn raw_form_string(&self) -> &str {
        &self.form_string
    }

    // We'd like to have form objects have pointers directly to the form string.
    // This means that the form string has to live at least as long as the form object. So,
    // to enforce this, we store the form_string along with the form object.
    //
    // So far so good. Now, this means that the form_string can never be
    // deallocated while the object is alive. That implies that the
    // `form_string` value should never be moved away. We can enforce that
    // easily by 1) not making `form_string` public, and 2) not exposing any
    // `&mut self` methods that could modify `form_string`.
    fn new(form_string: String) -> Result<Self, String> {
        let long_lived_string: &'f str = unsafe {
            ::std::mem::transmute(form_string.as_str())
        };

        let mut items = FormItems::from(long_lived_string);

        // Handle parsing or decode errors
        let parsing_errors: Vec<_> = items.by_ref()
            .map(|(key, value)| (key, String::from_form_value(value)))
            .filter(|&(_, ref decoded_value)| decoded_value.is_err())
            .collect();

        if !parsing_errors.is_empty() {
            // Parsing errors => fail with invalid result
            return Err(format!("::AMIWO::CONTRIB::ROCKET::FORM_HASMAP::NEW::WARNING Unable to parse form string {}", form_string));
        }
        if !items.completed() {
            warn!("::AMIWO::CONTRIB::ROCKET::FORM_HASMAP::NEW::WARNING Form string {} couldn't be completely parsed", form_string);
        }

        Ok(FormHashMap {
            form_string: form_string,
            form_hash: FormItems::from(long_lived_string).by_ref()
                            .map(|(key, value)| (key, String::from_form_value(value)))
                            .filter(|&(_, ref decoded_value)| decoded_value.is_ok())
                            .fold(
                                HashMap::<&'f str, FormValues<String>>::new(),
                                |mut map, (key, decoded_value)| {
                                    map.entry(key).or_insert(FormValues::new()).push(decoded_value.unwrap());
                                    map
                                }
                            ),
            _phantom: PhantomData,
        })
    }
}

/// Parses a `FormHashMap` from incoming POST/... form data.
///
/// - If the content type of the request data is not
/// `application/x-www-form-urlencoded`, `Forward`s the request.
/// - If the form string is malformed, a `Failure` with status code
/// `BadRequest` is returned.
/// - Finally, if reading the incoming stream fails, returns a `Failure` with status code
/// `InternalServerError`.
/// In all failure cases, the raw form string is returned if it was able to be retrieved from the incoming stream.
///
/// All relevant warnings and errors are written to the console
impl<'f> FromData for FormHashMap<'f> {
    type Error = String;

    fn from_data(request: &Request, data: Data) -> rocket::data::Outcome<Self, Self::Error> {
        // TODO add support for application/json
        if !request.content_type().map_or(false, |ct| ct.is_form()) {
            error!("::AMIWO::CONTRIB::ROCKET::FORM_HASMAP::FROM_DATA::WARNING Form data does not have form content type.");
            return rocket::Outcome::Forward(data);
        }

        let size_limit = rocket::config::active()
            .and_then(|c| c.extras.get("limits.application")) // TODO: remove placeholder when upgrading to rocket version > 0.2.6
            // .and_then(|c| c.limits.get("application") // In next version
            .and_then(|limit| limit.as_integer())
            .unwrap_or(32768) as u64;

        let mut buffer = String::new();
        data.open()
            .take(size_limit)
            .read_to_string(&mut buffer)
            .or_else(|err| Err(format!("::AMIWO::CONTRIB::ROCKET::FORM_HASMAP::FROM_DATA::ERROR IO Error: {}", err.description())) )
            .and_then(|_| FormHashMap::new(buffer)) // Note: if ok, read_to_string() returns how many bytes where read
            .or_else(|error_message| {
                error!("{}", error_message);
                Err(error_message)
            }).into_outcome() // Note: trait implemented by Rocket FromData for Result<S,E> `fn into_outcome(self, status: Status) -> Outcome<S, E>`
    }
}

/// Parses a `FormHashMap` from incoming GET query strings.
///
/// - If the form string is malformed, a `Failure` with status code
/// `BadRequest` is returned.
/// - Finally, if reading the incoming stream fails, returns a `Failure` with status code
/// `InternalServerError`.
/// In all failure cases, the raw form string is returned if it was able to be retrieved from the incoming stream.
///
/// All relevant warnings and errors are written to the console
impl<'f> FromForm<'f> for FormHashMap<'f> {
    /// The raw form string, if it was able to be retrieved from the request.
    type Error = (Status, Option<String>);

    fn from_form_items(items: &mut FormItems<'f>) -> Result<Self, Self::Error> {
        FormHashMap::new(items.inner_str().to_string())
            .map(|map| {
                info!("::AMIWO::CONTRIB::ROCKET::FORM_HASMAP::FROM_FORM_ITEMS::INFO Successfully parsed input data => {:?}", map);
                map
            }).map_err(|invalid_string| {
                error!("::AMIWO::CONTRIB::ROCKET::FORM_HASMAP::FROM_FORM_ITEMS::ERROR The request's form string '{}' was malformed.", invalid_string);
                (Status::BadRequest, Some(invalid_string))
            })
    }
}

/// Implement Debug displaying '<internal data holding structure>' from string: <parsed string>'
impl<'f> Debug for FormHashMap<'f> {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "{:?} from string: {:?}", self.form_hash, self.form_string)
    }
}

// =======================================================================
// TESTS
// =======================================================================
#[cfg(test)]
mod tests {
    #![allow(unmounted_route)]
    #![allow(non_snake_case)]

    use super::FormHashMap;
    use types::OneOrMany;

    use rocket;
    use rocket::testing::MockRequest;
    use rocket::http::{ ContentType, Method, Status };

    #[test]
    fn FormHashMap_test_new() {
        let form_string = "a=b1&a=b2&b=c";

        match FormHashMap::new(form_string.to_string()) {
            Ok(map) => {
                assert_eq!(map.get("a"), Some(&OneOrMany::Many(vec!["b1".to_string(), "b2".to_string()])));
                assert_eq!(map.get("b"), Some(&OneOrMany::One("c".to_string())));
            },
            Err(invalid_string) => {
                panic!("Unable to parse {}", invalid_string);
            }
        }
    }

    #[test]
    fn FormHashMap_test_post_route() {
        #[post("/test", data= "<params>")]
        fn test_route(params: FormHashMap) -> &'static str {
            assert_eq!(params.get("a"), Some(&OneOrMany::Many(vec!["b1".to_string(), "b2".to_string()])));
            assert_eq!(params.get("b"), Some(&OneOrMany::One("c".to_string())));
            "It's working !"
        }

        let rocket = rocket::ignite()
            .mount("/post", routes![test_route]);

        let mut req = MockRequest::new(Method::Post, "/post/test")
            .header(ContentType::Form)
            .body("a=b1&a=b2&b=c");

        let mut response = req.dispatch_with(&rocket);
        let body_str = response.body().and_then(|b| b.into_string());

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(body_str, Some("It's working !".to_string()));
    }

    #[test]
    fn FormHashMap_test_get_route() {
        #[get("/test?<params>")]
        fn test_route(params: FormHashMap) -> &'static str {
            assert_eq!(params.get("a"), Some(&OneOrMany::Many(vec!["b1".to_string(), "b2".to_string()])));
            assert_eq!(params.get("b"), Some(&OneOrMany::One("c".to_string())));
            "It's working !"
        }

        let rocket = rocket::ignite()
            .mount("/get", routes![test_route]);

        let mut req = MockRequest::new(Method::Get, "/get/test?a=b1&a=b2&b=c");

        let mut response = req.dispatch_with(&rocket);
        let body_str = response.body().and_then(|b| b.into_string());

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(body_str, Some("It's working !".to_string()));
    }

    #[test]
    fn FormHashMap_test_get_qs_with_dot() {
        #[get("/test?<params>")]
        fn test_route(params: FormHashMap) -> &'static str {
            assert_eq!(params.get("v"), Some(&OneOrMany::One("4.7.0".to_string())));
            "It's working !"
        }

        let rocket = rocket::ignite()
            .mount("/get", routes![test_route]);

        let mut req = MockRequest::new(Method::Get, "/get/test?v=4.7.0");

        let mut response = req.dispatch_with(&rocket);
        let body_str = response.body().and_then(|b| b.into_string());

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(body_str, Some("It's working !".to_string()));
    }

    // TODO: add test lifetime
}<commit_msg>FormHashMap > refactored using 'serde_json::Value' instead of HashMap => now returns 'Value' instead of 'OneOrMany' types FormHashMap > can now parse 'application/json' data (tests ok)<commit_after>//! File holding the FormHashMap type accepting variable parameters from a Rocket request
//! exposing a simplified Map type interface to access them
//!
//! Author: [Boris](mailto:boris@humanenginuity.com)
//! Version: 2.0
//!
//! ## Release notes
//! - v1.1 : refactored using serde_json Map & Value
//! - v1.1 : implemented Index trait, renamed old `new()` method into `from_application_data`, added method `from_json_data`
//! - v1.0 : creation

// =======================================================================
// ATTRIBUTES
// =======================================================================
#![allow(dead_code)]

// =======================================================================
// LIBRARY IMPORTS
// =======================================================================
use std::convert::AsRef;
use std::error::Error;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::io::Read;
use std::ops::Index;

use rocket;
use rocket::{ Request, Data };
use rocket::data::FromData;
use rocket::http::Status;
use rocket::outcome::IntoOutcome;
use rocket::request::{ FromForm, FromFormValue, FormItems };

use serde_json;
use serde_json::Value;
use serde_json::map::Map;

use error::GenericError;
use traits::Pushable;

// =======================================================================
// STRUCT & TRAIT DEFINITION
// =======================================================================
/// A `FromData` type that creates a map of the key/value pairs from a
/// `x-www-form-urlencoded` or `json` form string
pub struct FormHashMap<'s> {
    form_string: String,
    map: Map<String, Value>,
    _phantom: PhantomData<&'s str>,
}

// =======================================================================
// IMPLEMENTATION
// =======================================================================

impl<'s> FormHashMap<'s> {
    /// Get a reference for the value (or values) associated with `key`.
    pub fn get<T: AsRef<str>>(&self, key: T) -> Option<&Value> {
        self.map.get(key.as_ref())
    }

    /// Returns the raw form string that was used to parse the encapsulated
    /// object.
    pub fn raw_form_string(&self) -> &str {
        &self.form_string
    }

    /// Build a FormHashMap from application data (i.e. content type application/x-www-form-urlencoded)
    /// Uses Rocket's `FormItems::from<'f>(&'f str)` to parse the form's String
    fn from_application_data(form_string: String) -> Result<Self, GenericError> {
        let long_lived_string: &'s str = unsafe {
            ::std::mem::transmute(form_string.as_str())
        };
        let items = FormItems::from(long_lived_string);

        // Handle parsing or decode errors
        let parsing_errors: Vec<_> = FormItems::from(long_lived_string)
            .map(|(key, value)| (key, String::from_form_value(value)))
            .filter(|&(_, ref decoded_value)| decoded_value.is_err())
            .collect();

        if !parsing_errors.is_empty() {
            return amiwo_error!( format!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_APPLICATION_DATA::WARNING Unable to parse form string {} [parsing errors = {:?}]", form_string, parsing_errors) );
        }
        if !items.completed() {
            warn!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_APPLICATION_DATA::WARNING Form string {} couldn't be completely parsed", form_string);
        }

        Ok(FormHashMap {
            form_string: form_string,
            map: FormItems::from(long_lived_string)
                .map(|(key, value)| (key, String::from_form_value(value)))
                .filter(|&(_, ref decoded_value)| decoded_value.is_ok())
                .fold(
                    Map::new(),
                    |mut map, (key, decoded_value)| {
                        map.entry(key).or_insert(Value::Null).push(Value::String(decoded_value.unwrap()));
                        map
                    }
                ),
            _phantom: PhantomData,
        })
    }

    /// Build a FormHashMap from JSON data (i.e. content type application/json)
    /// Uses serde_json's `serde_json::from_str<'a, T>(&'a str)` to parse the form's String
    fn from_json_data(form_string: String) -> Result<Self, GenericError> {
        let long_lived_string: &'s str = unsafe {
            ::std::mem::transmute(form_string.as_str())
        };
        serde_json::from_str(long_lived_string)
            .or_else(|err| amiwo_error!(
                format!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_JSON_DATA::ERROR Error parsing string {} > {}", form_string, &err.description()),
                GenericError::Serde(err)
            )).and_then(|value : Value| {
                if value.is_object() {
                    Ok(FormHashMap {
                        form_string: form_string,
                        map: value.as_object().unwrap().clone(),
                        _phantom: PhantomData,
                    })
                } else {
                    amiwo_error!(format!(":AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_JSON_DATA::ERROR Invalid JSON data {}", form_string))
                }
            })
    }

    // We'd like to have form objects have pointers directly to the form string.
    // This means that the form string has to live at least as long as the form object. So,
    // to enforce this, we store the form_string along with the form object.
    //
    // So far so good. Now, this means that the form_string can never be
    // deallocated while the object is alive. That implies that the
    // `form_string` value should never be moved away. We can enforce that
    // easily by 1) not making `form_string` public, and 2) not exposing any
    // `&mut self` methods that could modify `form_string`.
    fn new(content_type: &str, form_string: String) -> Result<Self, GenericError> {
        match content_type {
            "application" => FormHashMap::from_application_data(form_string),
            "json" => FormHashMap::from_json_data(form_string),
            _ => amiwo_error!(format!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::NEW::ERROR Unsupported content type {}", content_type)),
        }
    }
}

// =======================================================================
// EXTERNAL TRAITS IMPLEMENTATION
// =======================================================================
/// Parses a `FormHashMap` from incoming POST/... form data.
///
/// - If the content type of the request data is not
/// `application/x-www-form-urlencoded` or `application/json`, `Forward`s the request.
/// - If the form string is malformed, a `Failure` with status code
/// `BadRequest` is returned.
/// - Finally, if reading the incoming stream fails, returns a `Failure` with status code
/// `InternalServerError`.
/// In all failure cases, the raw form string is returned if it was able to be retrieved from the incoming stream.
///
/// All relevant warnings and errors are written to the console
impl<'f> FromData for FormHashMap<'f> {
    type Error = GenericError;

    fn from_data(request: &Request, data: Data) -> rocket::data::Outcome<Self, Self::Error> {
        if !request.content_type().map_or(false, |ct| ct.is_form() || ct.is_json()) {
            error!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_DATA::WARNING Form data does not have application/x-www-form-urlencoded or application/json content type.");
            return rocket::Outcome::Forward(data);
        }

        let content_type = request.content_type().map_or("unsupported content type", |ct| if ct.is_form() { "application" } else { "json" });

        let size_limit = rocket::config::active()
            .and_then(|c| c.extras.get(&("limits.".to_string() + content_type))) // TODO: remove placeholder when upgrading to rocket version > 0.2.6
            // .and_then(|c| c.limits.get("application") // In next version
            .and_then(|limit| limit.as_integer())
            .unwrap_or_else(|| if content_type == "json" { 1<<20 } else { 32768 }) as u64;

        let mut buffer = String::new();
        data.open()
            .take(size_limit)
            .read_to_string(&mut buffer)
            .or_else(|err| amiwo_error!(format!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_DATA::ERROR IO Error: {}", err.description())) )
            .and_then(|_| FormHashMap::new(content_type, buffer))
            .or_else(|error_message| {
                error!("{}", error_message);
                Err(error_message)
            }).into_outcome() // Note: trait implemented by Rocket FromData for Result<S,E> `fn into_outcome(self, status: Status) -> Outcome<S, E>`
    }
}

/// Parses a `FormHashMap` from incoming GET query strings.
///
/// - If the form string is malformed, a `Failure` with status code
/// `BadRequest` is returned.
/// - Finally, if reading the incoming stream fails, returns a `Failure` with status code
/// `InternalServerError`.
/// In all failure cases, the raw form string is returned if it was able to be retrieved from the incoming stream.
///
/// All relevant warnings and errors are written to the console
impl<'f> FromForm<'f> for FormHashMap<'f> {
    /// The raw form string, if it was able to be retrieved from the request.
    type Error = (Status, Option<GenericError>);

    fn from_form_items(items: &mut FormItems<'f>) -> Result<Self, Self::Error> {
        FormHashMap::from_application_data(items.inner_str().to_string())
            .map(|map| {
                info!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_FORM_ITEMS::INFO Successfully parsed input data => {:?}", map);
                map
            }).map_err(|invalid_string| {
                error!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_FORM_ITEMS::ERROR The request's form string '{}' was malformed.", invalid_string);
                ( Status::BadRequest, Some(GenericError::Basic(format!("::AMIWO::CONTRIB::ROCKET::FORM_HASHMAP::FROM_FORM_ITEMS::ERROR The request's form string '{}' was malformed.", invalid_string))) )
            })
    }
}

/// Implement Debug displaying '<internal data holding structure>' from string: <parsed string>'
impl<'f> Debug for FormHashMap<'f> {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "{:?} from string: {:?}", self.map, self.form_string)
    }
}

/// Access an element of this type. Panics if the key is not defined
impl<'a, I: AsRef<str>> Index<I> for FormHashMap<'a> {
    type Output = Value;

    fn index(&self, index: I) -> &Self::Output {
        &self.map[index.as_ref()]
    }
}

// =======================================================================
// UNIT TESTS
// =======================================================================
#[cfg(test)]
mod tests {
    #![allow(unmounted_route)]
    #![allow(non_snake_case)]

    use super::FormHashMap;

    use rocket;
    use rocket::testing::MockRequest;
    use rocket::http::{ ContentType, Method, Status };

    #[test]
    fn FormHashMap_test_new() {
        let form_string = "a=b1&a=b2&b=c";

        match FormHashMap::from_application_data(form_string.to_string()) {
            Ok(map) => {
                assert_eq!(map.get("a"), Some(&json!(["b1", "b2"])));
                assert_eq!(map.get("b"), Some(&json!("c")));
                assert_eq!(map["b"], json!("c"));
                assert_eq!(map["b".to_string()], json!("c"));
            },
            Err(invalid_string) => {
                panic!("Unable to parse {}", invalid_string);
            }
        }
    }

    #[test]
    fn FormHashMap_test_post_route() {
        #[post("/test", data= "<params>")]
        fn test_route(params: FormHashMap) -> &'static str {
            assert_eq!(params.get("a"), Some(&json!(["b1", "b2"])));
            assert_eq!(params.get("b"), Some(&json!("c")));
            "It's working !"
        }

        let rocket = rocket::ignite()
            .mount("/post", routes![test_route]);

        let mut req = MockRequest::new(Method::Post, "/post/test")
            .header(ContentType::Form)
            .body("a=b1&a=b2&b=c");

        let mut response = req.dispatch_with(&rocket);
        let body_str = response.body().and_then(|b| b.into_string());

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(body_str, Some("It's working !".to_string()));
    }

    #[test]
    fn FormHashMap_test_get_route() {
        #[get("/test?<params>")]
        fn test_route(params: FormHashMap) -> &'static str {
            assert_eq!(params.get("a"), Some(&json!(["b1", "b2"])));
            assert_eq!(params.get("b"), Some(&json!("c")));
            "It's working !"
        }

        let rocket = rocket::ignite()
            .mount("/get", routes![test_route]);

        let mut req = MockRequest::new(Method::Get, "/get/test?a=b1&a=b2&b=c");

        let mut response = req.dispatch_with(&rocket);
        let body_str = response.body().and_then(|b| b.into_string());

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(body_str, Some("It's working !".to_string()));
    }

    #[test]
    fn FormHashMap_test_get_qs_with_dot() {
        #[get("/test?<params>")]
        fn test_route(params: FormHashMap) -> &'static str {
            assert_eq!(params.get("v"), Some(&json!("4.7.0")));
            "It's working !"
        }

        let rocket = rocket::ignite()
            .mount("/get", routes![test_route]);

        let mut req = MockRequest::new(Method::Get, "/get/test?v=4.7.0");

        let mut response = req.dispatch_with(&rocket);
        let body_str = response.body().and_then(|b| b.into_string());

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(body_str, Some("It's working !".to_string()));
    }

    // TODO: add test lifetime
}<|endoftext|>


<commit_before>//! Handles to an operator's input and output streams.

use std::rc::Rc;
use std::cell::RefCell;
use progress::Timestamp;
use progress::count_map::CountMap;
use progress::frontier::MutableAntichain;
use dataflow::channels::pullers::Counter as PullCounter;
use dataflow::channels::pushers::Counter as PushCounter;
use dataflow::channels::pushers::buffer::{Buffer, Session};
use dataflow::channels::Content;
use timely_communication::Push;

use dataflow::operators::Capability;
use dataflow::operators::capability::mint as mint_capability;

/// Handle to an operator's input stream.
pub struct InputHandle<'a, T: Timestamp, D: 'a> {
    pull_counter: &'a mut PullCounter<T, D>,
    internal: Rc<RefCell<CountMap<T>>>,
}

/// Handle to an operator's input stream and frontier.
pub struct FrontieredInputHandle<'a, T: Timestamp, D: 'a> {
    handle: InputHandle<'a, T, D>,
    frontier: &'a MutableAntichain<T>,
}

impl<'a, T: Timestamp, D> InputHandle<'a, T, D> {
    /// Reads the next input buffer (at some timestamp `t`) and a corresponding capability for `t`.
    /// The timestamp `t` of the input buffer can be retrieved by invoking `.time()` on the capability.
    /// Returns `None` when there's no more data available.
    #[inline]
    pub fn next(&mut self) -> Option<(Capability<T>, &mut Content<D>)> {
        let internal = &mut self.internal;
        self.pull_counter.next().map(|(&time, content)| {
            (mint_capability(time, internal.clone()), content)
        })
    }

    /// Repeatedly calls `logic` till exhaustion of the available input data.
    /// `logic` receives a capability and an input buffer.
    ///
    /// #Examples
    /// ```
    /// use timely::dataflow::operators::{ToStream, Unary};
    /// use timely::dataflow::channels::pact::Pipeline;
    ///
    /// timely::example(|scope| {
    ///     (0..10).to_stream(scope)
    ///            .unary_stream(Pipeline, "example", |input, output| {
    ///                input.for_each(|cap, data| {
    ///                    output.session(&cap).give_content(data);
    ///                });
    ///            });
    /// });
    /// ```
    #[inline]
    pub fn for_each<F: FnMut(Capability<T>, &mut Content<D>)>(&mut self, mut logic: F) {
        while let Some((cap, data)) = self.next() {
            ::logging::log(&::logging::GUARDED_MESSAGE, true);
            logic(cap, data);
            ::logging::log(&::logging::GUARDED_MESSAGE, false);
        }
    }

}

impl<'a, T: Timestamp, D> FrontieredInputHandle<'a, T, D> {
    /// Reads the next input buffer (at some timestamp `t`) and a corresponding capability for `t`.
    /// The timestamp `t` of the input buffer can be retrieved by invoking `.time()` on the capability.
    /// Returns `None` when there's no more data available.
    #[inline]
    pub fn next(&mut self) -> Option<(Capability<T>, &mut Content<D>)> {
        self.handle.next()
    }

    /// Repeatedly calls `logic` till exhaustion of the available input data.
    /// `logic` receives a capability and an input buffer.
    ///
    /// #Examples
    /// ```
    /// use timely::dataflow::operators::{ToStream, Unary};
    /// use timely::dataflow::channels::pact::Pipeline;
    ///
    /// timely::example(|scope| {
    ///     (0..10).to_stream(scope)
    ///            .unary_stream(Pipeline, "example", |input, output| {
    ///                input.for_each(|cap, data| {
    ///                    output.session(&cap).give_content(data);
    ///                });
    ///            });
    /// });
    /// ```
    #[inline]
    pub fn for_each<F: FnMut(Capability<T>, &mut Content<D>)>(&mut self, mut logic: F) {
        self.handle.for_each(logic)
    }

    /// Inspect the frontier associated with this input.
    #[inline]
    pub fn frontier(&self) -> &'a MutableAntichain<T> {
        self.frontier
    }
}

/// Constructs an input handle.
/// Declared separately so that it can be kept private when `InputHandle` is re-exported.
pub fn new_input_handle<'a, T: Timestamp, D: 'a>(pull_counter: &'a mut PullCounter<T, D>, internal: Rc<RefCell<CountMap<T>>>) -> InputHandle<'a, T, D> {
    InputHandle {
        pull_counter: pull_counter,
        internal: internal,
    }
}

pub fn new_frontier_input_handle<'a, T: Timestamp, D: 'a>(pull_counter: &'a mut PullCounter<T, D>, internal: Rc<RefCell<CountMap<T>>>, frontier: &'a MutableAntichain<T>) -> FrontieredInputHandle<'a, T, D> {
    FrontieredInputHandle {
        handle: new_input_handle(pull_counter, internal),
        frontier: frontier,
    }
}

/// Handle to an operator's output stream.
pub struct OutputHandle<'a, T: Timestamp, D: 'a, P: Push<(T, Content<D>)>+'a> {
    push_buffer: &'a mut Buffer<T, D, PushCounter<T, D, P>>,
}

impl<'a, T: Timestamp, D, P: Push<(T, Content<D>)>> OutputHandle<'a, T, D, P> {
    /// Obtains a session that can send data at the timestamp associated with capability `cap`.
    ///
    /// In order to send data at a future timestamp, obtain a capability for the new timestamp
    /// first, as show in the example.
    ///
    /// #Examples
    /// ```
    /// use timely::dataflow::operators::{ToStream, Unary};
    /// use timely::dataflow::channels::pact::Pipeline;
    ///
    /// timely::example(|scope| {
    ///     (0..10).to_stream(scope)
    ///            .unary_stream(Pipeline, "example", |input, output| {
    ///                input.for_each(|cap, data| {
    ///                    let mut time = cap.time();
    ///                    time.inner += 1;
    ///                    output.session(&cap.delayed(&time)).give_content(data);
    ///                });
    ///            });
    /// });
    /// ```
    pub fn session<'b>(&'b mut self, cap: &'b Capability<T>) -> Session<'b, T, D, PushCounter<T, D, P>> where 'a: 'b {
        self.push_buffer.session(cap)
    }
}

/// Constructs an output handle.
/// Declared separately so that it can be kept private when `OutputHandle` is re-exported.
pub fn new_output_handle<'a, T: Timestamp, D, P: Push<(T, Content<D>)>>(push_buffer: &'a mut Buffer<T, D, PushCounter<T, D, P>>) -> OutputHandle<'a, T, D, P> {
    OutputHandle {
        push_buffer: push_buffer,
    }
}

<commit_msg>changed mutability to remove warning<commit_after>//! Handles to an operator's input and output streams.

use std::rc::Rc;
use std::cell::RefCell;
use progress::Timestamp;
use progress::count_map::CountMap;
use progress::frontier::MutableAntichain;
use dataflow::channels::pullers::Counter as PullCounter;
use dataflow::channels::pushers::Counter as PushCounter;
use dataflow::channels::pushers::buffer::{Buffer, Session};
use dataflow::channels::Content;
use timely_communication::Push;

use dataflow::operators::Capability;
use dataflow::operators::capability::mint as mint_capability;

/// Handle to an operator's input stream.
pub struct InputHandle<'a, T: Timestamp, D: 'a> {
    pull_counter: &'a mut PullCounter<T, D>,
    internal: Rc<RefCell<CountMap<T>>>,
}

/// Handle to an operator's input stream and frontier.
pub struct FrontieredInputHandle<'a, T: Timestamp, D: 'a> {
    handle: InputHandle<'a, T, D>,
    frontier: &'a MutableAntichain<T>,
}

impl<'a, T: Timestamp, D> InputHandle<'a, T, D> {
    /// Reads the next input buffer (at some timestamp `t`) and a corresponding capability for `t`.
    /// The timestamp `t` of the input buffer can be retrieved by invoking `.time()` on the capability.
    /// Returns `None` when there's no more data available.
    #[inline]
    pub fn next(&mut self) -> Option<(Capability<T>, &mut Content<D>)> {
        let internal = &mut self.internal;
        self.pull_counter.next().map(|(&time, content)| {
            (mint_capability(time, internal.clone()), content)
        })
    }

    /// Repeatedly calls `logic` till exhaustion of the available input data.
    /// `logic` receives a capability and an input buffer.
    ///
    /// #Examples
    /// ```
    /// use timely::dataflow::operators::{ToStream, Unary};
    /// use timely::dataflow::channels::pact::Pipeline;
    ///
    /// timely::example(|scope| {
    ///     (0..10).to_stream(scope)
    ///            .unary_stream(Pipeline, "example", |input, output| {
    ///                input.for_each(|cap, data| {
    ///                    output.session(&cap).give_content(data);
    ///                });
    ///            });
    /// });
    /// ```
    #[inline]
    pub fn for_each<F: FnMut(Capability<T>, &mut Content<D>)>(&mut self, mut logic: F) {
        while let Some((cap, data)) = self.next() {
            ::logging::log(&::logging::GUARDED_MESSAGE, true);
            logic(cap, data);
            ::logging::log(&::logging::GUARDED_MESSAGE, false);
        }
    }

}

impl<'a, T: Timestamp, D> FrontieredInputHandle<'a, T, D> {
    /// Reads the next input buffer (at some timestamp `t`) and a corresponding capability for `t`.
    /// The timestamp `t` of the input buffer can be retrieved by invoking `.time()` on the capability.
    /// Returns `None` when there's no more data available.
    #[inline]
    pub fn next(&mut self) -> Option<(Capability<T>, &mut Content<D>)> {
        self.handle.next()
    }

    /// Repeatedly calls `logic` till exhaustion of the available input data.
    /// `logic` receives a capability and an input buffer.
    ///
    /// #Examples
    /// ```
    /// use timely::dataflow::operators::{ToStream, Unary};
    /// use timely::dataflow::channels::pact::Pipeline;
    ///
    /// timely::example(|scope| {
    ///     (0..10).to_stream(scope)
    ///            .unary_stream(Pipeline, "example", |input, output| {
    ///                input.for_each(|cap, data| {
    ///                    output.session(&cap).give_content(data);
    ///                });
    ///            });
    /// });
    /// ```
    #[inline]
    pub fn for_each<F: FnMut(Capability<T>, &mut Content<D>)>(&mut self, logic: F) {
        self.handle.for_each(logic)
    }

    /// Inspect the frontier associated with this input.
    #[inline]
    pub fn frontier(&self) -> &'a MutableAntichain<T> {
        self.frontier
    }
}

/// Constructs an input handle.
/// Declared separately so that it can be kept private when `InputHandle` is re-exported.
pub fn new_input_handle<'a, T: Timestamp, D: 'a>(pull_counter: &'a mut PullCounter<T, D>, internal: Rc<RefCell<CountMap<T>>>) -> InputHandle<'a, T, D> {
    InputHandle {
        pull_counter: pull_counter,
        internal: internal,
    }
}

pub fn new_frontier_input_handle<'a, T: Timestamp, D: 'a>(pull_counter: &'a mut PullCounter<T, D>, internal: Rc<RefCell<CountMap<T>>>, frontier: &'a MutableAntichain<T>) -> FrontieredInputHandle<'a, T, D> {
    FrontieredInputHandle {
        handle: new_input_handle(pull_counter, internal),
        frontier: frontier,
    }
}

/// Handle to an operator's output stream.
pub struct OutputHandle<'a, T: Timestamp, D: 'a, P: Push<(T, Content<D>)>+'a> {
    push_buffer: &'a mut Buffer<T, D, PushCounter<T, D, P>>,
}

impl<'a, T: Timestamp, D, P: Push<(T, Content<D>)>> OutputHandle<'a, T, D, P> {
    /// Obtains a session that can send data at the timestamp associated with capability `cap`.
    ///
    /// In order to send data at a future timestamp, obtain a capability for the new timestamp
    /// first, as show in the example.
    ///
    /// #Examples
    /// ```
    /// use timely::dataflow::operators::{ToStream, Unary};
    /// use timely::dataflow::channels::pact::Pipeline;
    ///
    /// timely::example(|scope| {
    ///     (0..10).to_stream(scope)
    ///            .unary_stream(Pipeline, "example", |input, output| {
    ///                input.for_each(|cap, data| {
    ///                    let mut time = cap.time();
    ///                    time.inner += 1;
    ///                    output.session(&cap.delayed(&time)).give_content(data);
    ///                });
    ///            });
    /// });
    /// ```
    pub fn session<'b>(&'b mut self, cap: &'b Capability<T>) -> Session<'b, T, D, PushCounter<T, D, P>> where 'a: 'b {
        self.push_buffer.session(cap)
    }
}

/// Constructs an output handle.
/// Declared separately so that it can be kept private when `OutputHandle` is re-exported.
pub fn new_output_handle<'a, T: Timestamp, D, P: Push<(T, Content<D>)>>(push_buffer: &'a mut Buffer<T, D, PushCounter<T, D, P>>) -> OutputHandle<'a, T, D, P> {
    OutputHandle {
        push_buffer: push_buffer,
    }
}

<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod lsp_notification_dispatch;
mod lsp_request_dispatch;
mod lsp_state;
mod lsp_state_resources;

use crate::{
    code_action::on_code_action,
    completion::on_completion,
    goto_definition::{
        on_get_source_location_of_type_definition, on_goto_definition,
        GetSourceLocationOfTypeDefinition,
    },
    graphql_tools::on_graphql_execute_query,
    graphql_tools::GraphQLExecuteQuery,
    hover::on_hover,
    js_language_server::JSLanguageServer,
    lsp::{
        set_initializing_status, CompletionOptions, Connection, GotoDefinition, HoverRequest,
        InitializeParams, Message, ServerCapabilities, ServerResponse, TextDocumentSyncCapability,
        TextDocumentSyncKind, WorkDoneProgressOptions,
    },
    lsp_process_error::{LSPProcessError, LSPProcessResult},
    lsp_runtime_error::LSPRuntimeError,
    references::on_references,
    resolved_types_at_location::{on_get_resolved_types_at_location, ResolvedTypesAtLocation},
    search_schema_items::{on_search_schema_items, SearchSchemaItems},
    shutdown::{on_exit, on_shutdown},
    status_reporting::{LSPStatusReporter, StatusReportingArtifactWriter},
    text_documents::{
        on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
        on_did_save_text_document,
    },
    ExtensionConfig,
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::{SendError, Sender};
use log::debug;
use lsp_notification_dispatch::LSPNotificationDispatch;
use lsp_request_dispatch::LSPRequestDispatch;
use lsp_server::{ErrorCode, Notification, ResponseError};
use lsp_types::{
    notification::{
        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Exit,
    },
    request::{CodeActionRequest, Completion, References, Shutdown},
    CodeActionProviderCapability,
};
use relay_compiler::{config::Config, NoopArtifactWriter};
use std::sync::Arc;

pub use crate::LSPExtraDataProvider;
pub use lsp_state::LSPState;
pub use lsp_state::{Schemas, SourcePrograms};

/// Initializes an LSP connection, handling the `initialize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> LSPProcessResult<InitializeParams> {
    let server_capabilities = ServerCapabilities {
        // Enable text document syncing so we can know when files are opened/changed/saved/closed
        text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full)),

        completion_provider: Some(CompletionOptions {
            resolve_provider: Some(true),
            trigger_characters: Some(vec!["(".into(), "\n".into(), ",".into(), "@".into()]),
            work_done_progress_options: WorkDoneProgressOptions {
                work_done_progress: None,
            },
        }),

        hover_provider: Some(true),
        definition_provider: Some(true),
        references_provider: Some(true),
        code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
        ..Default::default()
    };

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub async fn run<TPerfLogger: PerfLogger + 'static>(
    connection: Connection,
    mut config: Config,
    extension_config: ExtensionConfig,
    _params: InitializeParams,
    perf_logger: Arc<TPerfLogger>,
    extra_data_provider: Box<dyn LSPExtraDataProvider + Send + Sync>,
    js_resource: Box<dyn JSLanguageServer<TPerfLogger>>,
) -> LSPProcessResult<()>
where
    TPerfLogger: PerfLogger + 'static,
{
    debug!(
        "Running language server with config root {:?}",
        config.root_dir
    );
    set_initializing_status(&connection.sender);

    config.artifact_writer = if extension_config.no_artifacts {
        Box::new(NoopArtifactWriter)
    } else {
        Box::new(StatusReportingArtifactWriter::new(
            connection.sender.clone(),
            config.artifact_writer,
        ))
    };
    config.status_reporter = Box::new(LSPStatusReporter::new(
        config.root_dir.clone(),
        connection.sender.clone(),
    ));

    let mut lsp_state = LSPState::create_state(
        Arc::new(config),
        Arc::clone(&perf_logger),
        extra_data_provider,
        js_resource,
        &extension_config,
        connection.sender.clone(),
    );

    for msg in connection.receiver {
        debug!("LSP message received {:?}", msg);
        match msg {
            Message::Request(req) => {
                handle_request(&mut lsp_state, req, &connection.sender, &perf_logger)
                    .map_err(LSPProcessError::from)?;
            }
            Message::Notification(notification) => {
                handle_notification(&mut lsp_state, notification, &perf_logger);
            }
            _ => {
                // Ignore responses for now
            }
        }
    }

    Ok(())
}

fn handle_request<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    request: lsp_server::Request,
    sender: &Sender<Message>,
    perf_logger: &Arc<TPerfLogger>,
) -> Result<(), SendError<Message>> {
    let get_server_response_bound = |req| dispatch_request(req, lsp_state);
    let get_response = with_request_logging(perf_logger, get_server_response_bound);

    sender.send(Message::Response(get_response(request)))
}

fn dispatch_request<TPerfLogger: PerfLogger + 'static>(
    request: lsp_server::Request,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> ServerResponse {
    let get_response = || -> Result<_, ServerResponse> {
        let request = LSPRequestDispatch::new(request, lsp_state)
            .on_request_sync::<ResolvedTypesAtLocation>(on_get_resolved_types_at_location)?
            .on_request_sync::<SearchSchemaItems>(on_search_schema_items)?
            .on_request_sync::<GetSourceLocationOfTypeDefinition>(
                on_get_source_location_of_type_definition,
            )?
            .on_request_sync::<HoverRequest>(on_hover)?
            .on_request_sync::<GotoDefinition>(on_goto_definition)?
            .on_request_sync::<References>(on_references)?
            .on_request_sync::<Completion>(on_completion)?
            .on_request_sync::<CodeActionRequest>(on_code_action)?
            .on_request_sync::<Shutdown>(on_shutdown)?
            .on_request_sync::<GraphQLExecuteQuery>(on_graphql_execute_query)?
            .request();

        // If we have gotten here, we have not handled the request
        Ok(ServerResponse {
            id: request.id,
            result: None,
            error: Some(ResponseError {
                code: ErrorCode::MethodNotFound as i32,
                data: None,
                message: format!("No handler registered for method '{}'", request.method),
            }),
        })
    };
    get_response().unwrap_or_else(|response| response)
}

fn with_request_logging<'a, TPerfLogger: PerfLogger + 'static>(
    perf_logger: &'a Arc<TPerfLogger>,
    get_response: impl FnOnce(lsp_server::Request) -> ServerResponse + 'a,
) -> impl FnOnce(lsp_server::Request) -> ServerResponse + 'a {
    move |request| {
        let lsp_request_event = perf_logger.create_event("lsp_message");
        lsp_request_event.string("lsp_method", request.method.clone());
        lsp_request_event.string("lsp_type", "request".to_string());
        let lsp_request_processing_time = lsp_request_event.start("lsp_message_processing_time");

        let response = get_response(request);

        if response.result.is_some() {
            lsp_request_event.string("lsp_outcome", "success".to_string());
        } else if let Some(error) = &response.error {
            lsp_request_event.string("lsp_outcome", "error".to_string());
            if error.code != ErrorCode::RequestCanceled as i32 {
                lsp_request_event.string("lsp_error_message", error.message.to_string());
            }
            if let Some(data) = &error.data {
                lsp_request_event.string("lsp_error_data", data.to_string());
            }
        }
        // N.B. we don't handle the case where the ServerResponse has neither a result nor
        // an error, which is an invalid state.

        lsp_request_event.stop(lsp_request_processing_time);
        perf_logger.complete_event(lsp_request_event);

        response
    }
}

fn handle_notification<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    notification: Notification,
    perf_logger: &Arc<TPerfLogger>,
) {
    let lsp_notification_event = perf_logger.create_event("lsp_message");
    lsp_notification_event.string("lsp_method", notification.method.clone());
    lsp_notification_event.string("lsp_type", "notification".to_string());
    let lsp_notification_processing_time =
        lsp_notification_event.start("lsp_message_processing_time");

    let notification_result = dispatch_notification(notification, lsp_state);

    match notification_result {
        Ok(()) => {
            // The notification is not handled
            lsp_notification_event.string("lsp_outcome", "error".to_string());
        }
        Err(err) => {
            if let Some(err) = err {
                lsp_notification_event.string("lsp_outcome", "error".to_string());
                if let LSPRuntimeError::UnexpectedError(message) = err {
                    lsp_notification_event.string("lsp_error_message", message);
                }
            } else {
                lsp_notification_event.string("lsp_outcome", "success".to_string());
            }
        }
    }

    lsp_notification_event.stop(lsp_notification_processing_time);
    perf_logger.complete_event(lsp_notification_event);
}

fn dispatch_notification<TPerfLogger: PerfLogger + 'static>(
    notification: lsp_server::Notification,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> Result<(), Option<LSPRuntimeError>> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

    // If we have gotten here, we have not handled the notification
    debug!(
        "Error: no handler registered for notification '{}'",
        notification.method
    );
    Ok(())
}
<commit_msg>Panic on receiver error<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod lsp_notification_dispatch;
mod lsp_request_dispatch;
mod lsp_state;
mod lsp_state_resources;

use crate::{
    code_action::on_code_action,
    completion::on_completion,
    goto_definition::{
        on_get_source_location_of_type_definition, on_goto_definition,
        GetSourceLocationOfTypeDefinition,
    },
    graphql_tools::on_graphql_execute_query,
    graphql_tools::GraphQLExecuteQuery,
    hover::on_hover,
    js_language_server::JSLanguageServer,
    lsp::{
        set_initializing_status, CompletionOptions, Connection, GotoDefinition, HoverRequest,
        InitializeParams, Message, ServerCapabilities, ServerResponse, TextDocumentSyncCapability,
        TextDocumentSyncKind, WorkDoneProgressOptions,
    },
    lsp_process_error::{LSPProcessError, LSPProcessResult},
    lsp_runtime_error::LSPRuntimeError,
    references::on_references,
    resolved_types_at_location::{on_get_resolved_types_at_location, ResolvedTypesAtLocation},
    search_schema_items::{on_search_schema_items, SearchSchemaItems},
    shutdown::{on_exit, on_shutdown},
    status_reporting::{LSPStatusReporter, StatusReportingArtifactWriter},
    text_documents::{
        on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
        on_did_save_text_document,
    },
    ExtensionConfig,
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::{SendError, Sender};
use log::debug;
use lsp_notification_dispatch::LSPNotificationDispatch;
use lsp_request_dispatch::LSPRequestDispatch;
use lsp_server::{ErrorCode, Notification, ResponseError};
use lsp_types::{
    notification::{
        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Exit,
    },
    request::{CodeActionRequest, Completion, References, Shutdown},
    CodeActionProviderCapability,
};
use relay_compiler::{config::Config, NoopArtifactWriter};
use std::sync::Arc;

pub use crate::LSPExtraDataProvider;
pub use lsp_state::LSPState;
pub use lsp_state::{Schemas, SourcePrograms};

/// Initializes an LSP connection, handling the `initialize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> LSPProcessResult<InitializeParams> {
    let server_capabilities = ServerCapabilities {
        // Enable text document syncing so we can know when files are opened/changed/saved/closed
        text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full)),

        completion_provider: Some(CompletionOptions {
            resolve_provider: Some(true),
            trigger_characters: Some(vec!["(".into(), "\n".into(), ",".into(), "@".into()]),
            work_done_progress_options: WorkDoneProgressOptions {
                work_done_progress: None,
            },
        }),

        hover_provider: Some(true),
        definition_provider: Some(true),
        references_provider: Some(true),
        code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
        ..Default::default()
    };

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub async fn run<TPerfLogger: PerfLogger + 'static>(
    connection: Connection,
    mut config: Config,
    extension_config: ExtensionConfig,
    _params: InitializeParams,
    perf_logger: Arc<TPerfLogger>,
    extra_data_provider: Box<dyn LSPExtraDataProvider + Send + Sync>,
    js_resource: Box<dyn JSLanguageServer<TPerfLogger>>,
) -> LSPProcessResult<()>
where
    TPerfLogger: PerfLogger + 'static,
{
    debug!(
        "Running language server with config root {:?}",
        config.root_dir
    );
    set_initializing_status(&connection.sender);

    config.artifact_writer = if extension_config.no_artifacts {
        Box::new(NoopArtifactWriter)
    } else {
        Box::new(StatusReportingArtifactWriter::new(
            connection.sender.clone(),
            config.artifact_writer,
        ))
    };
    config.status_reporter = Box::new(LSPStatusReporter::new(
        config.root_dir.clone(),
        connection.sender.clone(),
    ));

    let mut lsp_state = LSPState::create_state(
        Arc::new(config),
        Arc::clone(&perf_logger),
        extra_data_provider,
        js_resource,
        &extension_config,
        connection.sender.clone(),
    );

    loop {
        debug!("waiting for incoming messages...");
        match connection.receiver.recv() {
            Ok(Message::Request(req)) => {
                handle_request(&mut lsp_state, req, &connection.sender, &perf_logger)
                    .map_err(LSPProcessError::from)?;
            }
            Ok(Message::Notification(notification)) => {
                handle_notification(&mut lsp_state, notification, &perf_logger);
            }
            Ok(_) => {
                // Ignore responses for now
            }
            Err(error) => {
                panic!("Relay Language Server receiver error {:?}", error);
            }
        }
    }
}

fn handle_request<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    request: lsp_server::Request,
    sender: &Sender<Message>,
    perf_logger: &Arc<TPerfLogger>,
) -> Result<(), SendError<Message>> {
    debug!("request received {:?}", request);
    let get_server_response_bound = |req| dispatch_request(req, lsp_state);
    let get_response = with_request_logging(perf_logger, get_server_response_bound);

    sender.send(Message::Response(get_response(request)))
}

fn dispatch_request<TPerfLogger: PerfLogger + 'static>(
    request: lsp_server::Request,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> ServerResponse {
    let get_response = || -> Result<_, ServerResponse> {
        let request = LSPRequestDispatch::new(request, lsp_state)
            .on_request_sync::<ResolvedTypesAtLocation>(on_get_resolved_types_at_location)?
            .on_request_sync::<SearchSchemaItems>(on_search_schema_items)?
            .on_request_sync::<GetSourceLocationOfTypeDefinition>(
                on_get_source_location_of_type_definition,
            )?
            .on_request_sync::<HoverRequest>(on_hover)?
            .on_request_sync::<GotoDefinition>(on_goto_definition)?
            .on_request_sync::<References>(on_references)?
            .on_request_sync::<Completion>(on_completion)?
            .on_request_sync::<CodeActionRequest>(on_code_action)?
            .on_request_sync::<Shutdown>(on_shutdown)?
            .on_request_sync::<GraphQLExecuteQuery>(on_graphql_execute_query)?
            .request();

        // If we have gotten here, we have not handled the request
        Ok(ServerResponse {
            id: request.id,
            result: None,
            error: Some(ResponseError {
                code: ErrorCode::MethodNotFound as i32,
                data: None,
                message: format!("No handler registered for method '{}'", request.method),
            }),
        })
    };
    get_response().unwrap_or_else(|response| response)
}

fn with_request_logging<'a, TPerfLogger: PerfLogger + 'static>(
    perf_logger: &'a Arc<TPerfLogger>,
    get_response: impl FnOnce(lsp_server::Request) -> ServerResponse + 'a,
) -> impl FnOnce(lsp_server::Request) -> ServerResponse + 'a {
    move |request| {
        let lsp_request_event = perf_logger.create_event("lsp_message");
        lsp_request_event.string("lsp_method", request.method.clone());
        lsp_request_event.string("lsp_type", "request".to_string());
        let lsp_request_processing_time = lsp_request_event.start("lsp_message_processing_time");

        let response = get_response(request);

        if response.result.is_some() {
            lsp_request_event.string("lsp_outcome", "success".to_string());
        } else if let Some(error) = &response.error {
            lsp_request_event.string("lsp_outcome", "error".to_string());
            if error.code != ErrorCode::RequestCanceled as i32 {
                lsp_request_event.string("lsp_error_message", error.message.to_string());
            }
            if let Some(data) = &error.data {
                lsp_request_event.string("lsp_error_data", data.to_string());
            }
        }
        // N.B. we don't handle the case where the ServerResponse has neither a result nor
        // an error, which is an invalid state.

        lsp_request_event.stop(lsp_request_processing_time);
        perf_logger.complete_event(lsp_request_event);

        response
    }
}

fn handle_notification<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    notification: Notification,
    perf_logger: &Arc<TPerfLogger>,
) {
    debug!("notification received {:?}", notification);
    let lsp_notification_event = perf_logger.create_event("lsp_message");
    lsp_notification_event.string("lsp_method", notification.method.clone());
    lsp_notification_event.string("lsp_type", "notification".to_string());
    let lsp_notification_processing_time =
        lsp_notification_event.start("lsp_message_processing_time");

    let notification_result = dispatch_notification(notification, lsp_state);

    match notification_result {
        Ok(()) => {
            // The notification is not handled
            lsp_notification_event.string("lsp_outcome", "error".to_string());
        }
        Err(err) => {
            if let Some(err) = err {
                lsp_notification_event.string("lsp_outcome", "error".to_string());
                if let LSPRuntimeError::UnexpectedError(message) = err {
                    lsp_notification_event.string("lsp_error_message", message);
                }
            } else {
                lsp_notification_event.string("lsp_outcome", "success".to_string());
            }
        }
    }

    lsp_notification_event.stop(lsp_notification_processing_time);
    perf_logger.complete_event(lsp_notification_event);
}

fn dispatch_notification<TPerfLogger: PerfLogger + 'static>(
    notification: lsp_server::Notification,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> Result<(), Option<LSPRuntimeError>> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

    // If we have gotten here, we have not handled the notification
    debug!(
        "Error: no handler registered for notification '{}'",
        notification.method
    );
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// 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.

use num_traits::{Float, FloatConst, Zero, One};
use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, Sub, SubAssign};
use core::cmp::{Eq, PartialEq};
use core::hash::{Hash};
use crate::trig::Trig;
use crate::approxeq::ApproxEq;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

/// An angle in radians
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Angle<T> {
    pub radians: T,
}

impl<T> Angle<T> {
    #[inline]
    pub fn radians(radians: T) -> Self {
        Angle { radians }
    }

    #[inline]
    pub fn get(self) -> T {
        self.radians
    }
}

impl<T> Angle<T>
where
    T: Trig,
{
    #[inline]
    pub fn degrees(deg: T) -> Self {
        Angle {
            radians: T::degrees_to_radians(deg),
        }
    }

    #[inline]
    pub fn to_degrees(self) -> T {
        T::radians_to_degrees(self.radians)
    }
}

impl<T> Angle<T>
where
    T: Rem<Output = T> + Sub<Output = T> + Add<Output = T> + Zero + FloatConst + PartialOrd + Copy,
{
    /// Returns this angle in the [0..2*PI[ range.
    pub fn positive(&self) -> Self {
        let two_pi = T::PI() + T::PI();
        let mut a = self.radians % two_pi;
        if a < T::zero() {
            a = a + two_pi;
        }
        Angle::radians(a)
    }

    /// Returns this angle in the ]-PI..PI] range.
    pub fn signed(&self) -> Self {
        Angle::pi() - (Angle::pi() - *self).positive()
    }

}

impl<T> Angle<T>
where
    T: Rem<Output = T> + Mul<Output = T> + Sub<Output = T> + Add<Output = T> + One + FloatConst + Copy,
{
    /// Returns the shortest signed angle between two angles.
    ///
    /// Takes wrapping and signs into account.
    pub fn angle_to(&self, to: Self) -> Self {
        let two = T::one() + T::one();
        let max = T::PI() * two;
        let d = (to.radians - self.radians) % max;

        Angle::radians(two * d % max - d)
    }

    /// Linear interpolation between two angles, using the shortest path.
    pub fn lerp(&self, other: Self, t: T) -> Self {
        *self + self.angle_to(other) * t
    }
}

impl<T> Angle<T>
where
    T: Float,
{
    /// Returns (sin(self), cos(self)).
    pub fn sin_cos(self) -> (T, T) {
        self.radians.sin_cos()
    }
}

impl<T> Angle<T>
where
    T: Zero,
{
    pub fn zero() -> Self {
        Angle::radians(T::zero())
    }
}

impl<T> Angle<T>
where
    T: FloatConst + Add<Output = T>,
{
    pub fn pi() -> Self {
        Angle::radians(T::PI())
    }

    pub fn two_pi() -> Self {
        Angle::radians(T::PI() + T::PI())
    }

    pub fn frac_pi_2() -> Self {
        Angle::radians(T::FRAC_PI_2())
    }

    pub fn frac_pi_3() -> Self {
        Angle::radians(T::FRAC_PI_3())
    }

    pub fn frac_pi_4() -> Self {
        Angle::radians(T::FRAC_PI_4())
    }
}

impl<T: Add<T, Output = T>> Add for Angle<T> {
    type Output = Angle<T>;
    fn add(self, other: Angle<T>) -> Angle<T> {
        Angle::radians(self.radians + other.radians)
    }
}

impl<T: AddAssign<T>> AddAssign for Angle<T> {
    fn add_assign(&mut self, other: Angle<T>) {
        self.radians += other.radians;
    }
}

impl<T: Sub<T, Output = T>> Sub<Angle<T>> for Angle<T> {
    type Output = Angle<T>;
    fn sub(self, other: Angle<T>) -> <Self as Sub>::Output {
        Angle::radians(self.radians - other.radians)
    }
}

impl<T: SubAssign<T>> SubAssign for Angle<T> {
    fn sub_assign(&mut self, other: Angle<T>) {
        self.radians -= other.radians;
    }
}

impl<T: Div<T, Output = T>> Div<Angle<T>> for Angle<T> {
    type Output = T;
    #[inline]
    fn div(self, other: Angle<T>) -> T {
        self.radians / other.radians
    }
}

impl<T: Div<T, Output = T>> Div<T> for Angle<T> {
    type Output = Angle<T>;
    #[inline]
    fn div(self, factor: T) -> Angle<T> {
        Angle::radians(self.radians / factor)
    }
}

impl<T: DivAssign<T>> DivAssign<T> for Angle<T> {
    fn div_assign(&mut self, factor: T) {
        self.radians /= factor;
    }
}

impl<T: Mul<T, Output = T>> Mul<T> for Angle<T> {
    type Output = Angle<T>;
    #[inline]
    fn mul(self, factor: T) -> Angle<T> {
        Angle::radians(self.radians * factor)
    }
}

impl<T: MulAssign<T>> MulAssign<T> for Angle<T> {
    fn mul_assign(&mut self, factor: T) {
        self.radians *= factor;
    }
}

impl<T: Neg<Output = T>> Neg for Angle<T> {
    type Output = Self;
    fn neg(self) -> Self {
        Angle::radians(-self.radians)
    }
}

impl<T: ApproxEq<T>> ApproxEq<T> for Angle<T> {
    #[inline]
    fn approx_epsilon() -> T {
        T::approx_epsilon()
    }

    #[inline]
    fn approx_eq_eps(&self, other: &Angle<T>, approx_epsilon: &T) -> bool {
        self.radians.approx_eq_eps(&other.radians, approx_epsilon)
    }
}

#[test]
fn wrap_angles() {
    use core::f32::consts::{FRAC_PI_2, PI};

    assert!(Angle::radians(0.0).positive().approx_eq(&Angle::zero()));
    assert!(
        Angle::radians(FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(
        Angle::radians(-FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::radians(3.0 * FRAC_PI_2))
    );
    assert!(
        Angle::radians(3.0 * FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::radians(3.0 * FRAC_PI_2))
    );
    assert!(
        Angle::radians(5.0 * FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(Angle::radians(2.0 * PI).positive().approx_eq(&Angle::zero()));
    assert!(Angle::radians(-2.0 * PI).positive().approx_eq(&Angle::zero()));
    assert!(Angle::radians(PI).positive().approx_eq(&Angle::pi()));
    assert!(Angle::radians(-PI).positive().approx_eq(&Angle::pi()));

    assert!(
        Angle::radians(FRAC_PI_2)
            .signed()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(
        Angle::radians(3.0 * FRAC_PI_2)
            .signed()
            .approx_eq(&-Angle::frac_pi_2())
    );
    assert!(
        Angle::radians(5.0 * FRAC_PI_2)
            .signed()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(Angle::radians(2.0 * PI).signed().approx_eq(&Angle::zero()));
    assert!(Angle::radians(-2.0 * PI).signed().approx_eq(&Angle::zero()));
    assert!(Angle::radians(-PI).signed().approx_eq(&Angle::pi()));
    assert!(Angle::radians(PI).signed().approx_eq(&Angle::pi()));
}

#[test]
fn lerp() {
    type A = Angle<f32>;

    let a = A::radians(1.0);
    let b = A::radians(2.0);
    assert!(a.lerp(b, 0.25).approx_eq(&Angle::radians(1.25)));
    assert!(a.lerp(b, 0.5).approx_eq(&Angle::radians(1.5)));
    assert!(a.lerp(b, 0.75).approx_eq(&Angle::radians(1.75)));
    assert!(a.lerp(b + A::two_pi(), 0.75).approx_eq(&Angle::radians(1.75)));
    assert!(a.lerp(b - A::two_pi(), 0.75).approx_eq(&Angle::radians(1.75)));
    assert!(a.lerp(b + A::two_pi() * 5.0, 0.75).approx_eq(&Angle::radians(1.75)));
}
<commit_msg>Auto merge of #440 - hannobraun:angle-cast, r=nical<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// 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.

use num_traits::{Float, FloatConst, NumCast, Zero, One};
use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, Sub, SubAssign};
use core::cmp::{Eq, PartialEq};
use core::hash::{Hash};
use crate::trig::Trig;
use crate::approxeq::ApproxEq;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

/// An angle in radians
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Angle<T> {
    pub radians: T,
}

impl<T> Angle<T> {
    #[inline]
    pub fn radians(radians: T) -> Self {
        Angle { radians }
    }

    #[inline]
    pub fn get(self) -> T {
        self.radians
    }
}

impl<T> Angle<T>
where
    T: Trig,
{
    #[inline]
    pub fn degrees(deg: T) -> Self {
        Angle {
            radians: T::degrees_to_radians(deg),
        }
    }

    #[inline]
    pub fn to_degrees(self) -> T {
        T::radians_to_degrees(self.radians)
    }
}

impl<T> Angle<T>
where
    T: Rem<Output = T> + Sub<Output = T> + Add<Output = T> + Zero + FloatConst + PartialOrd + Copy,
{
    /// Returns this angle in the [0..2*PI[ range.
    pub fn positive(&self) -> Self {
        let two_pi = T::PI() + T::PI();
        let mut a = self.radians % two_pi;
        if a < T::zero() {
            a = a + two_pi;
        }
        Angle::radians(a)
    }

    /// Returns this angle in the ]-PI..PI] range.
    pub fn signed(&self) -> Self {
        Angle::pi() - (Angle::pi() - *self).positive()
    }

}

impl<T> Angle<T>
where
    T: Rem<Output = T> + Mul<Output = T> + Sub<Output = T> + Add<Output = T> + One + FloatConst + Copy,
{
    /// Returns the shortest signed angle between two angles.
    ///
    /// Takes wrapping and signs into account.
    pub fn angle_to(&self, to: Self) -> Self {
        let two = T::one() + T::one();
        let max = T::PI() * two;
        let d = (to.radians - self.radians) % max;

        Angle::radians(two * d % max - d)
    }

    /// Linear interpolation between two angles, using the shortest path.
    pub fn lerp(&self, other: Self, t: T) -> Self {
        *self + self.angle_to(other) * t
    }
}

impl<T> Angle<T>
where
    T: Float,
{
    /// Returns (sin(self), cos(self)).
    pub fn sin_cos(self) -> (T, T) {
        self.radians.sin_cos()
    }
}

impl<T> Angle<T>
where
    T: Zero,
{
    pub fn zero() -> Self {
        Angle::radians(T::zero())
    }
}

impl<T> Angle<T>
where
    T: FloatConst + Add<Output = T>,
{
    pub fn pi() -> Self {
        Angle::radians(T::PI())
    }

    pub fn two_pi() -> Self {
        Angle::radians(T::PI() + T::PI())
    }

    pub fn frac_pi_2() -> Self {
        Angle::radians(T::FRAC_PI_2())
    }

    pub fn frac_pi_3() -> Self {
        Angle::radians(T::FRAC_PI_3())
    }

    pub fn frac_pi_4() -> Self {
        Angle::radians(T::FRAC_PI_4())
    }
}

impl<T> Angle<T>
where
    T: NumCast + Copy,
{
    /// Cast from one numeric representation to another.
    #[inline]
    pub fn cast<NewT: NumCast>(&self) -> Angle<NewT> {
        self.try_cast().unwrap()
    }

    /// Fallible cast from one numeric representation to another.
    pub fn try_cast<NewT: NumCast>(&self) -> Option<Angle<NewT>> {
        NumCast::from(self.radians)
            .map(|radians| Angle { radians })
    }

    // Convenience functions for common casts.

    /// Cast angle to `f32`.
    #[inline]
    pub fn to_f32(&self) -> Angle<f32> {
        self.cast()
    }

    /// Cast angle `f64`.
    #[inline]
    pub fn to_f64(&self) -> Angle<f64> {
        self.cast()
    }
}

impl<T: Add<T, Output = T>> Add for Angle<T> {
    type Output = Angle<T>;
    fn add(self, other: Angle<T>) -> Angle<T> {
        Angle::radians(self.radians + other.radians)
    }
}

impl<T: AddAssign<T>> AddAssign for Angle<T> {
    fn add_assign(&mut self, other: Angle<T>) {
        self.radians += other.radians;
    }
}

impl<T: Sub<T, Output = T>> Sub<Angle<T>> for Angle<T> {
    type Output = Angle<T>;
    fn sub(self, other: Angle<T>) -> <Self as Sub>::Output {
        Angle::radians(self.radians - other.radians)
    }
}

impl<T: SubAssign<T>> SubAssign for Angle<T> {
    fn sub_assign(&mut self, other: Angle<T>) {
        self.radians -= other.radians;
    }
}

impl<T: Div<T, Output = T>> Div<Angle<T>> for Angle<T> {
    type Output = T;
    #[inline]
    fn div(self, other: Angle<T>) -> T {
        self.radians / other.radians
    }
}

impl<T: Div<T, Output = T>> Div<T> for Angle<T> {
    type Output = Angle<T>;
    #[inline]
    fn div(self, factor: T) -> Angle<T> {
        Angle::radians(self.radians / factor)
    }
}

impl<T: DivAssign<T>> DivAssign<T> for Angle<T> {
    fn div_assign(&mut self, factor: T) {
        self.radians /= factor;
    }
}

impl<T: Mul<T, Output = T>> Mul<T> for Angle<T> {
    type Output = Angle<T>;
    #[inline]
    fn mul(self, factor: T) -> Angle<T> {
        Angle::radians(self.radians * factor)
    }
}

impl<T: MulAssign<T>> MulAssign<T> for Angle<T> {
    fn mul_assign(&mut self, factor: T) {
        self.radians *= factor;
    }
}

impl<T: Neg<Output = T>> Neg for Angle<T> {
    type Output = Self;
    fn neg(self) -> Self {
        Angle::radians(-self.radians)
    }
}

impl<T: ApproxEq<T>> ApproxEq<T> for Angle<T> {
    #[inline]
    fn approx_epsilon() -> T {
        T::approx_epsilon()
    }

    #[inline]
    fn approx_eq_eps(&self, other: &Angle<T>, approx_epsilon: &T) -> bool {
        self.radians.approx_eq_eps(&other.radians, approx_epsilon)
    }
}

#[test]
fn wrap_angles() {
    use core::f32::consts::{FRAC_PI_2, PI};

    assert!(Angle::radians(0.0).positive().approx_eq(&Angle::zero()));
    assert!(
        Angle::radians(FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(
        Angle::radians(-FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::radians(3.0 * FRAC_PI_2))
    );
    assert!(
        Angle::radians(3.0 * FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::radians(3.0 * FRAC_PI_2))
    );
    assert!(
        Angle::radians(5.0 * FRAC_PI_2)
            .positive()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(Angle::radians(2.0 * PI).positive().approx_eq(&Angle::zero()));
    assert!(Angle::radians(-2.0 * PI).positive().approx_eq(&Angle::zero()));
    assert!(Angle::radians(PI).positive().approx_eq(&Angle::pi()));
    assert!(Angle::radians(-PI).positive().approx_eq(&Angle::pi()));

    assert!(
        Angle::radians(FRAC_PI_2)
            .signed()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(
        Angle::radians(3.0 * FRAC_PI_2)
            .signed()
            .approx_eq(&-Angle::frac_pi_2())
    );
    assert!(
        Angle::radians(5.0 * FRAC_PI_2)
            .signed()
            .approx_eq(&Angle::frac_pi_2())
    );
    assert!(Angle::radians(2.0 * PI).signed().approx_eq(&Angle::zero()));
    assert!(Angle::radians(-2.0 * PI).signed().approx_eq(&Angle::zero()));
    assert!(Angle::radians(-PI).signed().approx_eq(&Angle::pi()));
    assert!(Angle::radians(PI).signed().approx_eq(&Angle::pi()));
}

#[test]
fn lerp() {
    type A = Angle<f32>;

    let a = A::radians(1.0);
    let b = A::radians(2.0);
    assert!(a.lerp(b, 0.25).approx_eq(&Angle::radians(1.25)));
    assert!(a.lerp(b, 0.5).approx_eq(&Angle::radians(1.5)));
    assert!(a.lerp(b, 0.75).approx_eq(&Angle::radians(1.75)));
    assert!(a.lerp(b + A::two_pi(), 0.75).approx_eq(&Angle::radians(1.75)));
    assert!(a.lerp(b - A::two_pi(), 0.75).approx_eq(&Angle::radians(1.75)));
    assert!(a.lerp(b + A::two_pi() * 5.0, 0.75).approx_eq(&Angle::radians(1.75)));
}
<|endoftext|>


<commit_before>use nom::{le_u32};

use blocks;
use util;

#[derive(Debug)]
pub enum Block<'a> {
    SectionHeader(blocks::SectionHeader<'a>),
    EnhancedPacket(blocks::EnhancedPacket<'a>),
    InterfaceDescription(blocks::InterfaceDescription<'a>),
}

/// Public representation of a parsed block
#[derive(Debug)]
pub struct RawBlock<'a> {
    //  0                   1                   2                   3
    //  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  |                          Block Type                           |
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  |                      Block Total Length                       |
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  /                          Block Body                           /
    //  /          /* variable length, aligned to 32 bits */            /
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  |                      Block Total Length                       |
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    pub ty: u32,
    pub block_length: u32,
    pub body: &'a [u8],
    pub check_length: u32,
}

impl<'a> RawBlock<'a> {
    pub fn parse(self) -> Block<'a> {
        match self.ty {
            blocks::section_header::TY => Block::SectionHeader(blocks::section_header::parse(self)),
            blocks::enhanced_packet::TY => Block::EnhancedPacket(blocks::enhanced_packet::parse(self)),
            blocks::interface_description::TY => Block::InterfaceDescription(blocks::interface_description::parse(self)),
            _ => panic!("Unknown block type {:x}", self.ty),
        }
    }
}

named!(pub parse_block< &[u8],RawBlock >,
       chain!(
           ty: le_u32 ~
           block_length: le_u32 ~
           body: take!((block_length - 12) as usize) ~
           take!(util::pad_to_32bits((block_length - 12) as usize)) ~
           check_length: le_u32 ,

           ||{ RawBlock {
               ty: ty,
               block_length: block_length,
               body: body,
               check_length: check_length
           } }
           )
      );

named!(pub parse_blocks< &[u8],Vec<RawBlock> >,
       many1!(parse_block)
       );

#[cfg(test)]
use nom::IResult;

#[test]
fn test_parse_block() {
    let input = b"\n\r\r\n\x1c\x00\x00\x00M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x1c\x00\x00\x00";
    match parse_block(input) {
        IResult::Done(left, RawBlock { ty, block_length, body, check_length }) => {
            // Ignored because we do not currently parse the whole block
            assert_eq!(left, b"");
            assert_eq!(ty, 0x0A0D0D0A);
            assert_eq!(block_length, 28);
            assert_eq!(body, b"M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff");
            assert_eq!(check_length, 28);
        },
        _ => {
            assert_eq!(1, 2);
        },
    }
}

#[test]
fn test_parse_blocks() {
    let input = b"\n\r\r\n\x1c\x00\x00\x00M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x1c\x00\x00\x00\
    \n\r\r\n\x1c\x00\x00\x00M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x1c\x00\x00\x00";
    match parse_blocks(input) {
        IResult::Done(left, blocks) => {
            assert_eq!(blocks.len(), 2);
            for i in blocks {
                let RawBlock { ty, block_length, body, check_length } = i;
                // Ignored because we do not currently parse the whole block
                assert_eq!(left, b"");
                assert_eq!(ty, 0x0A0D0D0A);
                assert_eq!(block_length, 28);
                assert_eq!(body, b"M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff");
                assert_eq!(check_length, 28);
            }
        },
        _ => {
            assert_eq!(1, 2);
        },
    }
}

#[test]
fn test_parse_weird_length_block() {
    let input = b"\n\r\r\n\x1b\x00\x00\x00<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x00\x1b\x00\x00\x00";
    match parse_block(input) {
        IResult::Done(left, RawBlock { ty, block_length, body, check_length }) => {
            // Ignored because we do not currently parse the whole block
            assert_eq!(left, b"");
            assert_eq!(ty, 0x0A0D0D0A);
            assert_eq!(27, block_length);
            assert_eq!(body.len() + 12, block_length as usize);
            assert_eq!(body, b"<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff");
            assert_eq!(body.len() + 12, check_length as usize);
        },
        _ => {
            unreachable!("Couldn't parse the block");
        },
    }
}

#[test]
fn test_multiple_options() {
    let input = b"\x0a\x0d\x0d\x0a\x40\x00\x00\x00\x4d\x3c\x2b\x1a\x01\x00\x00\x00\
                \xff\xff\xff\xff\xff\xff\xff\xff\x03\x00\x0b\x00\x57\x69\x6e\x64\
                \x6f\x77\x73\x20\x58\x50\x00\x00\x04\x00\x0c\x00\x54\x65\x73\x74\
                \x30\x30\x34\x2e\x65\x78\x65\x00\x00\x00\x00\x00\x40\x00\x00\x00";
    match parse_block(input) {
        IResult::Done(left, block) => {
            assert_eq!(left, b"");
            if let Block::SectionHeader(blk) = block.parse() {
                if let Some(opts) = blk.options {
                    assert_eq!(opts.options.len(), 3);

                    let o = &opts.options[0];
                    assert_eq!(o.code, 0x03);
                    assert_eq!(o.length, 0x0b);
                    assert_eq!(&o.value[..], b"Windows XP\x00");

                    let o = &opts.options[1];
                    assert_eq!(o.code, 0x04);
                    assert_eq!(o.length, 0x0c);
                    assert_eq!(&o.value[..], b"Test004.exe\x00");

                    let o = &opts.options[2];
                    assert_eq!(o.code, 0x00);
                    assert_eq!(o.length, 0x00);
                    assert_eq!(&o.value[..], b"");
                } else {
                    unreachable!();
                }
            } else {
                unreachable!();
            }
        },
        _ => {
            panic!("Hit a codepath we shouldn't have");
        },
    }
}
<commit_msg>Return an Option from RawBlock.parse<commit_after>use nom::{le_u32};

use blocks;
use util;

#[derive(Debug)]
pub enum Block<'a> {
    SectionHeader(blocks::SectionHeader<'a>),
    EnhancedPacket(blocks::EnhancedPacket<'a>),
    InterfaceDescription(blocks::InterfaceDescription<'a>),
}

/// Public representation of a parsed block
#[derive(Debug)]
pub struct RawBlock<'a> {
    //  0                   1                   2                   3
    //  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  |                          Block Type                           |
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  |                      Block Total Length                       |
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  /                          Block Body                           /
    //  /          /* variable length, aligned to 32 bits */            /
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //  |                      Block Total Length                       |
    //  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    pub ty: u32,
    pub block_length: u32,
    pub body: &'a [u8],
    pub check_length: u32,
}

impl<'a> RawBlock<'a> {
    pub fn parse(self) -> Option<Block<'a>> {
        let block = match self.ty {
            blocks::section_header::TY => Block::SectionHeader(blocks::section_header::parse(self)),
            blocks::enhanced_packet::TY => Block::EnhancedPacket(blocks::enhanced_packet::parse(self)),
            blocks::interface_description::TY => Block::InterfaceDescription(blocks::interface_description::parse(self)),

            _ => return None,

        };

        return Some(block);
    }
}

named!(pub parse_block< &[u8],RawBlock >,
       chain!(
           ty: le_u32 ~
           block_length: le_u32 ~
           body: take!((block_length - 12) as usize) ~
           take!(util::pad_to_32bits((block_length - 12) as usize)) ~
           check_length: le_u32 ,

           ||{ RawBlock {
               ty: ty,
               block_length: block_length,
               body: body,
               check_length: check_length
           } }
           )
      );

named!(pub parse_blocks< &[u8],Vec<RawBlock> >,
       many1!(parse_block)
       );

#[cfg(test)]
use nom::IResult;

#[test]
fn test_parse_block() {
    let input = b"\n\r\r\n\x1c\x00\x00\x00M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x1c\x00\x00\x00";
    match parse_block(input) {
        IResult::Done(left, RawBlock { ty, block_length, body, check_length }) => {
            // Ignored because we do not currently parse the whole block
            assert_eq!(left, b"");
            assert_eq!(ty, 0x0A0D0D0A);
            assert_eq!(block_length, 28);
            assert_eq!(body, b"M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff");
            assert_eq!(check_length, 28);
        },
        _ => {
            assert_eq!(1, 2);
        },
    }
}

#[test]
fn test_parse_blocks() {
    let input = b"\n\r\r\n\x1c\x00\x00\x00M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x1c\x00\x00\x00\
    \n\r\r\n\x1c\x00\x00\x00M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x1c\x00\x00\x00";
    match parse_blocks(input) {
        IResult::Done(left, blocks) => {
            assert_eq!(blocks.len(), 2);
            for i in blocks {
                let RawBlock { ty, block_length, body, check_length } = i;
                // Ignored because we do not currently parse the whole block
                assert_eq!(left, b"");
                assert_eq!(ty, 0x0A0D0D0A);
                assert_eq!(block_length, 28);
                assert_eq!(body, b"M<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff");
                assert_eq!(check_length, 28);
            }
        },
        _ => {
            assert_eq!(1, 2);
        },
    }
}

#[test]
fn test_parse_weird_length_block() {
    let input = b"\n\r\r\n\x1b\x00\x00\x00<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\x00\x1b\x00\x00\x00";
    match parse_block(input) {
        IResult::Done(left, RawBlock { ty, block_length, body, check_length }) => {
            // Ignored because we do not currently parse the whole block
            assert_eq!(left, b"");
            assert_eq!(ty, 0x0A0D0D0A);
            assert_eq!(27, block_length);
            assert_eq!(body.len() + 12, block_length as usize);
            assert_eq!(body, b"<+\x1a\x01\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff");
            assert_eq!(body.len() + 12, check_length as usize);
        },
        _ => {
            unreachable!("Couldn't parse the block");
        },
    }
}

#[test]
fn test_multiple_options() {
    let input = b"\x0a\x0d\x0d\x0a\x40\x00\x00\x00\x4d\x3c\x2b\x1a\x01\x00\x00\x00\
                \xff\xff\xff\xff\xff\xff\xff\xff\x03\x00\x0b\x00\x57\x69\x6e\x64\
                \x6f\x77\x73\x20\x58\x50\x00\x00\x04\x00\x0c\x00\x54\x65\x73\x74\
                \x30\x30\x34\x2e\x65\x78\x65\x00\x00\x00\x00\x00\x40\x00\x00\x00";
    match parse_block(input) {
        IResult::Done(left, block) => {
            assert_eq!(left, b"");
            if let Some(Block::SectionHeader(blk)) = block.parse() {
                if let Some(opts) = blk.options {
                    assert_eq!(opts.options.len(), 3);

                    let o = &opts.options[0];
                    assert_eq!(o.code, 0x03);
                    assert_eq!(o.length, 0x0b);
                    assert_eq!(&o.value[..], b"Windows XP\x00");

                    let o = &opts.options[1];
                    assert_eq!(o.code, 0x04);
                    assert_eq!(o.length, 0x0c);
                    assert_eq!(&o.value[..], b"Test004.exe\x00");

                    let o = &opts.options[2];
                    assert_eq!(o.code, 0x00);
                    assert_eq!(o.length, 0x00);
                    assert_eq!(&o.value[..], b"");
                } else {
                    unreachable!();
                }
            } else {
                unreachable!();
            }
        },
        _ => {
            panic!("Hit a codepath we shouldn't have");
        },
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::thread;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::time::Duration;

use rocksdb::DB;

use super::cluster::{Simulator, Cluster};
use tikv::raftserver::server::*;
use tikv::raftserver::Result;
use tikv::util::codec::rpc;
use kvproto::raft_serverpb::{self, Message, MessageType};
use kvproto::raft_cmdpb::*;
use tikv::pd::PdClient;
use super::util;
use super::pd::TestPdClient;
use super::pd_ask::run_ask_loop;

pub struct ServerCluster {
    cluster_id: u64,
    senders: HashMap<u64, SendCh>,
    handles: HashMap<u64, thread::JoinHandle<()>>,
    addrs: HashMap<u64, SocketAddr>,
    conns: Mutex<HashMap<SocketAddr, TcpStream>>,

    msg_id: Mutex<u64>,
    pd_client: Arc<RwLock<TestPdClient>>,
}

impl ServerCluster {
    pub fn new(cluster_id: u64, pd_client: Arc<RwLock<TestPdClient>>) -> ServerCluster {
        ServerCluster {
            cluster_id: cluster_id,
            senders: HashMap::new(),
            handles: HashMap::new(),
            addrs: HashMap::new(),
            conns: Mutex::new(HashMap::new()),
            msg_id: Mutex::new(0),
            pd_client: pd_client,
        }
    }

    fn alloc_msg_id(&self) -> u64 {
        let mut msg_id = self.msg_id.lock().unwrap();
        *msg_id += 1;
        *msg_id
    }
}

impl Simulator for ServerCluster {
    #[allow(useless_format)]
    fn run_node(&mut self, node_id: u64, cfg: Config, engine: Arc<DB>) -> u64 {
        assert!(node_id == 0 || !self.handles.contains_key(&node_id));
        assert!(node_id == 0 || !self.senders.contains_key(&node_id));

        let mut event_loop = create_event_loop().unwrap();
        let mut server = Server::new(&mut event_loop, cfg, vec![engine], self.pd_client.clone())
                             .unwrap();
        let addr = server.listening_addr().unwrap();

        assert!(node_id == 0 || node_id == server.get_node_id());
        let node_id = server.get_node_id();

        // We must get real listening address and re-update pd again.
        self.pd_client
            .write()
            .unwrap()
            .put_node(self.cluster_id, util::new_node(node_id, addr.to_string()))
            .unwrap();

        let ch = server.get_sendch();

        let t = thread::spawn(move || {
            server.run(&mut event_loop).unwrap();
        });

        self.handles.insert(node_id, t);
        self.senders.insert(node_id, ch);
        self.addrs.insert(node_id, addr);

        node_id
    }

    fn stop_node(&mut self, node_id: u64) {
        let h = self.handles.remove(&node_id).unwrap();
        let ch = self.senders.remove(&node_id).unwrap();
        self.addrs.remove(&node_id).unwrap();

        ch.send(Msg::Quit).unwrap();
        h.join().unwrap();
    }

    fn get_node_ids(&self) -> HashSet<u64> {
        self.senders.keys().cloned().collect()
    }

    fn call_command(&self, request: RaftCmdRequest, timeout: Duration) -> Result<RaftCmdResponse> {
        let node_id = request.get_header().get_peer().get_node_id();
        let addr = self.addrs.get(&node_id).unwrap();

        let mut conns = self.conns.lock().unwrap();
        let mut conn = conns.entry(*addr).or_insert_with(|| TcpStream::connect(addr).unwrap());

        conn.set_write_timeout(Some(timeout)).unwrap();

        let mut msg = Message::new();
        msg.set_msg_type(MessageType::Cmd);
        msg.set_cmd_req(request);

        let msg_id = self.alloc_msg_id();
        try!(rpc::encode_msg(&mut conn, msg_id, &msg));

        conn.set_read_timeout(Some(timeout)).unwrap();

        let mut resp_msg = Message::new();
        let get_msg_id = try!(rpc::decode_msg(&mut conn, &mut resp_msg));

        assert_eq!(resp_msg.get_msg_type(), MessageType::CmdResp);
        assert_eq!(msg_id, get_msg_id);

        Ok(resp_msg.take_cmd_resp())
    }

    fn send_raft_msg(&self, raft_msg: raft_serverpb::RaftMessage) -> Result<()> {
        let node_id = raft_msg.get_to_peer().get_node_id();
        let addr = self.addrs.get(&node_id).unwrap();

        let mut conn = TcpStream::connect(addr).unwrap();

        conn.set_write_timeout(Some(Duration::from_secs(3))).unwrap();

        let mut msg = Message::new();
        msg.set_msg_type(MessageType::Raft);
        msg.set_raft(raft_msg);

        let msg_id = self.alloc_msg_id();
        try!(rpc::encode_msg(&mut conn, msg_id, &msg));

        Ok(())
    }
}

pub fn new_server_cluster(id: u64, count: usize) -> Cluster<ServerCluster> {
    let (tx, rx) = mpsc::channel();
    let pd_client = Arc::new(RwLock::new(TestPdClient::new(tx)));
    let sim = Arc::new(RwLock::new(ServerCluster::new(id, pd_client.clone())));
    run_ask_loop(pd_client.clone(), sim.clone(), rx);
    Cluster::new(id, count, sim, pd_client)
}
<commit_msg>tests: Add connection pool<commit_after>use std::collections::{HashMap, HashSet};
use std::thread;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::time::Duration;

use rocksdb::DB;

use super::cluster::{Simulator, Cluster};
use tikv::raftserver::server::*;
use tikv::raftserver::Result;
use tikv::util::codec::rpc;
use kvproto::raft_serverpb::{self, Message, MessageType};
use kvproto::raft_cmdpb::*;
use tikv::pd::PdClient;
use super::util;
use super::pd::TestPdClient;
use super::pd_ask::run_ask_loop;

pub struct ServerCluster {
    cluster_id: u64,
    senders: HashMap<u64, SendCh>,
    handles: HashMap<u64, thread::JoinHandle<()>>,
    addrs: HashMap<u64, SocketAddr>,
    conns: Mutex<HashMap<SocketAddr, Vec<TcpStream>>>,

    msg_id: Mutex<u64>,
    pd_client: Arc<RwLock<TestPdClient>>,
}

impl ServerCluster {
    pub fn new(cluster_id: u64, pd_client: Arc<RwLock<TestPdClient>>) -> ServerCluster {
        ServerCluster {
            cluster_id: cluster_id,
            senders: HashMap::new(),
            handles: HashMap::new(),
            addrs: HashMap::new(),
            conns: Mutex::new(HashMap::new()),
            msg_id: Mutex::new(0),
            pd_client: pd_client,
        }
    }

    fn alloc_msg_id(&self) -> u64 {
        let mut msg_id = self.msg_id.lock().unwrap();
        *msg_id += 1;
        *msg_id
    }

    fn get_conn(&self, addr: &SocketAddr) -> Result<TcpStream> {
        {
            let mut conns = self.conns
                                .lock()
                                .unwrap();
            let conn = conns.get_mut(addr);
            if let Some(mut pool) = conn {
                if !pool.is_empty() {
                    return Ok(pool.pop().unwrap());
                }
            }
        }

        let conn = TcpStream::connect(addr).unwrap();
        Ok(conn)
    }
}

impl Simulator for ServerCluster {
    #[allow(useless_format)]
    fn run_node(&mut self, node_id: u64, cfg: Config, engine: Arc<DB>) -> u64 {
        assert!(node_id == 0 || !self.handles.contains_key(&node_id));
        assert!(node_id == 0 || !self.senders.contains_key(&node_id));

        let mut event_loop = create_event_loop().unwrap();
        let mut server = Server::new(&mut event_loop, cfg, vec![engine], self.pd_client.clone())
                             .unwrap();
        let addr = server.listening_addr().unwrap();

        assert!(node_id == 0 || node_id == server.get_node_id());
        let node_id = server.get_node_id();

        // We must get real listening address and re-update pd again.
        self.pd_client
            .write()
            .unwrap()
            .put_node(self.cluster_id, util::new_node(node_id, addr.to_string()))
            .unwrap();

        let ch = server.get_sendch();

        let t = thread::spawn(move || {
            server.run(&mut event_loop).unwrap();
        });

        self.handles.insert(node_id, t);
        self.senders.insert(node_id, ch);
        self.addrs.insert(node_id, addr);

        node_id
    }

    fn stop_node(&mut self, node_id: u64) {
        let h = self.handles.remove(&node_id).unwrap();
        let ch = self.senders.remove(&node_id).unwrap();
        self.addrs.remove(&node_id).unwrap();

        ch.send(Msg::Quit).unwrap();
        h.join().unwrap();
    }

    fn get_node_ids(&self) -> HashSet<u64> {
        self.senders.keys().cloned().collect()
    }

    fn call_command(&self, request: RaftCmdRequest, timeout: Duration) -> Result<RaftCmdResponse> {
        let node_id = request.get_header().get_peer().get_node_id();
        let addr = self.addrs.get(&node_id).unwrap();
        let mut conn = self.get_conn(addr).unwrap();

        let mut msg = Message::new();
        msg.set_msg_type(MessageType::Cmd);
        msg.set_cmd_req(request);

        let msg_id = self.alloc_msg_id();
        try!(rpc::encode_msg(&mut conn, msg_id, &msg));

        let mut resp_msg = Message::new();
        let get_msg_id = try!(rpc::decode_msg(&mut conn, &mut resp_msg));

        {
            let mut conns = self.conns
                                .lock()
                                .unwrap();
            let p = conns.entry(*addr).or_insert(vec![]);
            p.push(conn);
        }

        assert_eq!(resp_msg.get_msg_type(), MessageType::CmdResp);
        assert_eq!(msg_id, get_msg_id);

        Ok(resp_msg.take_cmd_resp())
    }

    fn send_raft_msg(&self, raft_msg: raft_serverpb::RaftMessage) -> Result<()> {
        let node_id = raft_msg.get_to_peer().get_node_id();
        let addr = self.addrs.get(&node_id).unwrap();

        let mut conn = TcpStream::connect(addr).unwrap();

        conn.set_write_timeout(Some(Duration::from_secs(3))).unwrap();

        let mut msg = Message::new();
        msg.set_msg_type(MessageType::Raft);
        msg.set_raft(raft_msg);

        let msg_id = self.alloc_msg_id();
        try!(rpc::encode_msg(&mut conn, msg_id, &msg));

        Ok(())
    }
}

pub fn new_server_cluster(id: u64, count: usize) -> Cluster<ServerCluster> {
    let (tx, rx) = mpsc::channel();
    let pd_client = Arc::new(RwLock::new(TestPdClient::new(tx)));
    let sim = Arc::new(RwLock::new(ServerCluster::new(id, pd_client.clone())));
    run_ask_loop(pd_client.clone(), sim.clone(), rx);
    Cluster::new(id, count, sim, pd_client)
}
<|endoftext|>