<commit_before>// Copyright 2017 Dmitry Tantsur <divius.inside@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Generic API bits for implementing new services.

use std::marker::PhantomData;

use hyper::Url;
use hyper::client::{Body, RequestBuilder as HyperRequestBuilder, Response};
use hyper::header::{Header, HeaderFormat, Headers};
pub use hyper::method::Method;
use serde::{Deserialize, Serialize};
use serde_json;

use super::{ApiError, ApiResult, ApiVersion, ApiVersionRequest, Session};
use super::auth::Method as AuthMethod;
use super::identity::protocol::AuthTokenHeader;
use super::utils;

/// Request builder with authentication, error checking and JSON support.
///
/// Partly copies the interface of hyper::client::RequestBuilder.
#[allow(missing_debug_implementations)]
pub struct RequestBuilder<'a, A: AuthMethod + 'a> {
    parent: &'a Session<A>,
    inner: HyperRequestBuilder<'a>
}

/// Type of query parameters.
#[derive(Clone, Debug)]
pub struct Query(pub Vec<(String, String)>);

/// Information about API endpoint.
#[derive(Clone, Debug)]
pub struct ServiceInfo {
    /// Root endpoint.
    pub root_url: Url,
    /// Current API version (if supported).
    pub current_version: Option<ApiVersion>,
    /// Minimum API version (if supported).
    pub minimum_version: Option<ApiVersion>
}

/// Trait representing a service type.
pub trait ServiceType {
    /// Service type to pass to the catalog.
    fn catalog_type() -> &'static str;

    /// Get basic service information.
    fn service_info<Auth: AuthMethod>(endpoint: Url, session: &Session<Auth>)
        -> ApiResult<ServiceInfo>;
}

/// Trait representing a service with API version support.
pub trait ApiVersioning {
    /// Return headers to set for this API version.
    fn api_version_headers(version: ApiVersion) -> ApiResult<Headers>;
}

/// A service-specific wrapper around Session.
#[derive(Debug)]
pub struct ServiceWrapper<'session, Auth: AuthMethod + 'session,
                          Srv: ServiceType> {
    session: &'session Session<Auth>,
    service_type: PhantomData<Srv>,
    endpoint_interface: Option<String>
}

impl Query {
    /// Empty query.
    pub fn new() -> Query {
        Query(Vec::new())
    }

    /// Add an item to the query.
    pub fn push<K, V>(&mut self, param: K, value: V)
            where K: Into<String>, V: ToString {
        self.0.push((param.into(), value.to_string()))
    }

    /// Add a strng item to the query.
    pub fn push_str<K, V>(&mut self, param: K, value: V)
            where K: Into<String>, V: Into<String> {
        self.0.push((param.into(), value.into()))
    }
}

impl<'a, Auth: AuthMethod> RequestBuilder<'a, Auth> {
    /// Wrap a request builder.
    pub fn new(inner: HyperRequestBuilder<'a>, session: &'a Session<Auth>)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            parent: session,
            inner: inner
        }
    }

    /// Send this request.
    pub fn send(self) -> ApiResult<Response> {
        let resp = try!(self.send_unchecked());
        if resp.status.is_success() {
            Ok(resp)
        } else {
            Err(ApiError::HttpError(resp.status, resp))
        }
    }

    /// Send this request without checking on status code.
    pub fn send_unchecked(self) -> ApiResult<Response> {
        let token = try!(self.parent.auth_token());
        let hdr = AuthTokenHeader(token.into());
        self.inner.header(hdr).send().map_err(From::from)
    }

    /// Send this request and parse JSON response on success.
    pub fn fetch_json<T: Deserialize>(self) -> ApiResult<T> {
        let resp = try!(self.send());
        serde_json::from_reader(resp).map_err(From::from)
    }

    /// Add body to the request.
    pub fn body<B: Into<Body<'a>>>(self, body: B)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            inner: self.inner.body(body),
            .. self
        }
    }

    /// Add additional headers to the request.
    pub fn headers(self, headers: Headers)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            inner: self.inner.headers(headers),
            .. self
        }
    }

    /// Add an individual header to the request.
    ///
    /// Note that X-Auth-Token is always overwritten with a token in use.
    pub fn header<H: Header + HeaderFormat>(self, header: H)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            inner: self.inner.header(header),
            .. self
        }
    }
}

impl<'session, Auth: AuthMethod + 'session, Srv: ServiceType>
        ServiceWrapper<'session, Auth, Srv> {
    /// Create a new wrapper for the specific service.
    pub fn new(session: &'session Session<Auth>)
            -> ServiceWrapper<'session, Auth, Srv> {
        ServiceWrapper {
            session: session,
            service_type: PhantomData,
            endpoint_interface: None
        }
    }

    /// Change the endpoint interface used for this wrapper.
    pub fn with_endpoint_interface(self, endpoint_interface: String)
            -> ServiceWrapper<'session, Auth, Srv> {
        ServiceWrapper {
            endpoint_interface: Some(endpoint_interface),
            .. self
        }
    }

    /// Construct and endpoint for the given service from the path.
    pub fn get_endpoint<P>(&self, path: P, query: Query) -> ApiResult<Url>
            where P: IntoIterator, P::Item: AsRef<str> {
        let ep = self.endpoint_interface.clone();
        let info = try!(self.session.get_service_info::<Srv>(ep));
        let mut url = utils::url::extend(info.root_url, path);
        let _ = url.query_pairs_mut().extend_pairs(query.0);
        Ok(url)
    }

    /// Make an HTTP request to the given service.
    pub fn request<P>(&'session self, method: Method, path: P, query: Query)
            -> ApiResult<RequestBuilder<'session, Auth>>
            where P: IntoIterator, P::Item: AsRef<str> {
        let url = try!(self.get_endpoint(path, query));
        let headers = self.session.service_headers::<Srv>();
        trace!("Sending HTTP {} request to {} with {:?}",
               method, url, headers);
        Ok(self.session.raw_request(method, url).headers(headers))
    }

    /// Make an HTTP request with JSON body and JSON response.
    pub fn json<P, Req, Res>(&self, method: Method, path: P, query: Query,
                             body: &Req) -> ApiResult<Res>
            where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        let str_body = try!(serde_json::to_string(body));
        let request = try!(self.request(method, path, query));
        request.body(&str_body).fetch_json()
    }

    /// Make a GET request returning a JSON.
    pub fn get_json<P, Res>(&self, path: P, query: Query) -> ApiResult<Res>
            where Res: Deserialize, P: IntoIterator, P::Item: AsRef<str> {
        try!(self.request(Method::Get, path, query)).fetch_json()
    }

    /// Make a POST request sending and returning a JSON.
    pub fn post_json<P, Req, Res>(&self, path: P, query: Query, body: &Req)
            -> ApiResult<Res> where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        self.json(Method::Post, path, query, body)
    }

    /// Make a POST request sending and returning a JSON.
    pub fn put_json<P, Req, Res>(&self, path: P, query: Query, body: &Req)
            -> ApiResult<Res> where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        self.json(Method::Put, path, query, body)
    }

    /// Make a PATCH request sending and returning a JSON.
    pub fn patch_json<P, Req, Res>(&self, path: P, query: Query, body: &Req)
            -> ApiResult<Res> where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        self.json(Method::Patch, path, query, body)
    }

    /// Make a DELETE request.
    pub fn delete<P>(&self, path: P, query: Query) -> ApiResult<Response>
            where P: IntoIterator, P::Item: AsRef<str> {
        try!(self.request(Method::Delete, path, query)).send()
    }
}

impl<'session, Auth: AuthMethod + 'session, Srv: ServiceType>
        Clone for ServiceWrapper<'session, Auth, Srv> {
    fn clone(&self) -> ServiceWrapper<'session, Auth, Srv> {
        ServiceWrapper {
            session: self.session,
            service_type: PhantomData,
            endpoint_interface: self.endpoint_interface.clone()
        }
    }
}

impl ServiceInfo {
    /// Pick an API version.
    pub fn pick_api_version(&self, request: ApiVersionRequest)
            -> Option<ApiVersion> {
        match request {
            ApiVersionRequest::Minimum =>
                self.minimum_version,
            ApiVersionRequest::Latest =>
                self.current_version,
            ApiVersionRequest::Exact(req) => {
                self.current_version.and_then(|max| {
                    match self.minimum_version {
                        Some(min) if req >= min && req <= max => Some(req),
                        None if req == max => Some(req),
                        _ => None
                    }
                })
            },
            ApiVersionRequest::Choice(vec) => {
                if vec.is_empty() {
                    return None;
                }

                self.current_version.and_then(|max| {
                    match self.minimum_version {
                        Some(min) => vec.into_iter().filter(|x| {
                            *x >= min && *x <= max
                        }).max(),
                        None =>vec.into_iter().find(|x| *x == max)
                    }
                })
            }
        }
    }
}

#[cfg(test)]
pub mod test {
    use hyper::Url;

    use super::super::{ApiVersion, ApiVersionRequest};
    use super::ServiceInfo;

    fn service_info(min: Option<u16>, max: Option<u16>) -> ServiceInfo {
        ServiceInfo {
            root_url: Url::parse("http://127.0.0.1").unwrap(),
            minimum_version: min.map(|x| ApiVersion(2, x)),
            current_version: max.map(|x| ApiVersion(2, x)),
        }
    }

    #[test]
    fn test_pick_version_exact() {
        let info = service_info(Some(1), Some(24));
        let version = ApiVersion(2, 22);
        let result = info.pick_api_version(ApiVersionRequest::Exact(version))
            .unwrap();
        assert_eq!(result, version);
    }

    #[test]
    fn test_pick_version_exact_mismatch() {
        let info = service_info(Some(1), Some(24));
        let version = ApiVersion(2, 25);
        let res1 = info.pick_api_version(ApiVersionRequest::Exact(version));
        assert!(res1.is_none());
        let version2 = ApiVersion(1, 11);
        let res2 = info.pick_api_version(ApiVersionRequest::Exact(version2));
        assert!(res2.is_none());
    }

    #[test]
    fn test_pick_version_exact_current_only() {
        let info = service_info(None, Some(24));
        let version = ApiVersion(2, 24);
        let result = info.pick_api_version(ApiVersionRequest::Exact(version))
            .unwrap();
        assert_eq!(result, version);
    }

    #[test]
    fn test_pick_version_exact_current_only_mismatch() {
        let info = service_info(None, Some(24));
        let version = ApiVersion(2, 22);
        let result = info.pick_api_version(ApiVersionRequest::Exact(version));
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_minimum() {
        let info = service_info(Some(1), Some(24));
        let result = info.pick_api_version(ApiVersionRequest::Minimum)
            .unwrap();
        assert_eq!(result, ApiVersion(2, 1));
    }

    #[test]
    fn test_pick_version_minimum_unknown() {
        let info = service_info(None, Some(24));
        let result = info.pick_api_version(ApiVersionRequest::Minimum);
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_latest() {
        let info = service_info(Some(1), Some(24));
        let result = info.pick_api_version(ApiVersionRequest::Latest)
            .unwrap();
        assert_eq!(result, ApiVersion(2, 24));
    }

    #[test]
    fn test_pick_version_latest_unknown() {
        let info = service_info(Some(1), None);
        let result = info.pick_api_version(ApiVersionRequest::Latest);
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_choice() {
        let info = service_info(Some(1), Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 2),
                          ApiVersion(2, 22), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice))
            .unwrap();
        assert_eq!(result, ApiVersion(2, 22));
    }

    #[test]
    fn test_pick_version_choice_mismatch() {
        let info = service_info(Some(1), Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice));
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_choice_current_only() {
        let info = service_info(None, Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 2),
                          ApiVersion(2, 24), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice))
            .unwrap();
        assert_eq!(result, ApiVersion(2, 24));
    }

    #[test]
    fn test_pick_version_choice_current_only_mismatch() {
        let info = service_info(None, Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 2),
                          ApiVersion(2, 22), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice));
        assert!(result.is_none());
    }
}
<commit_msg>feat(service): log every response as TRACE<commit_after>// Copyright 2017 Dmitry Tantsur <divius.inside@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Generic API bits for implementing new services.

use std::marker::PhantomData;

use hyper::Url;
use hyper::client::{Body, RequestBuilder as HyperRequestBuilder, Response};
use hyper::header::{Header, HeaderFormat, Headers};
pub use hyper::method::Method;
use serde::{Deserialize, Serialize};
use serde_json;

use super::{ApiError, ApiResult, ApiVersion, ApiVersionRequest, Session};
use super::auth::Method as AuthMethod;
use super::identity::protocol::AuthTokenHeader;
use super::utils;

/// Request builder with authentication, error checking and JSON support.
///
/// Partly copies the interface of hyper::client::RequestBuilder.
#[allow(missing_debug_implementations)]
pub struct RequestBuilder<'a, A: AuthMethod + 'a> {
    parent: &'a Session<A>,
    inner: HyperRequestBuilder<'a>
}

/// Type of query parameters.
#[derive(Clone, Debug)]
pub struct Query(pub Vec<(String, String)>);

/// Information about API endpoint.
#[derive(Clone, Debug)]
pub struct ServiceInfo {
    /// Root endpoint.
    pub root_url: Url,
    /// Current API version (if supported).
    pub current_version: Option<ApiVersion>,
    /// Minimum API version (if supported).
    pub minimum_version: Option<ApiVersion>
}

/// Trait representing a service type.
pub trait ServiceType {
    /// Service type to pass to the catalog.
    fn catalog_type() -> &'static str;

    /// Get basic service information.
    fn service_info<Auth: AuthMethod>(endpoint: Url, session: &Session<Auth>)
        -> ApiResult<ServiceInfo>;
}

/// Trait representing a service with API version support.
pub trait ApiVersioning {
    /// Return headers to set for this API version.
    fn api_version_headers(version: ApiVersion) -> ApiResult<Headers>;
}

/// A service-specific wrapper around Session.
#[derive(Debug)]
pub struct ServiceWrapper<'session, Auth: AuthMethod + 'session,
                          Srv: ServiceType> {
    session: &'session Session<Auth>,
    service_type: PhantomData<Srv>,
    endpoint_interface: Option<String>
}

impl Query {
    /// Empty query.
    pub fn new() -> Query {
        Query(Vec::new())
    }

    /// Add an item to the query.
    pub fn push<K, V>(&mut self, param: K, value: V)
            where K: Into<String>, V: ToString {
        self.0.push((param.into(), value.to_string()))
    }

    /// Add a strng item to the query.
    pub fn push_str<K, V>(&mut self, param: K, value: V)
            where K: Into<String>, V: Into<String> {
        self.0.push((param.into(), value.into()))
    }
}

impl<'a, Auth: AuthMethod> RequestBuilder<'a, Auth> {
    /// Wrap a request builder.
    pub fn new(inner: HyperRequestBuilder<'a>, session: &'a Session<Auth>)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            parent: session,
            inner: inner
        }
    }

    /// Send this request.
    pub fn send(self) -> ApiResult<Response> {
        let resp = try!(self.send_unchecked());
        if resp.status.is_success() {
            Ok(resp)
        } else {
            Err(ApiError::HttpError(resp.status, resp))
        }
    }

    /// Send this request without checking on status code.
    pub fn send_unchecked(self) -> ApiResult<Response> {
        let token = try!(self.parent.auth_token());
        let hdr = AuthTokenHeader(token.into());
        self.inner.header(hdr).send().map_err(From::from).map(|resp| {
            trace!("Got {} from {} with {:?}",
                   resp.status, resp.url, resp.headers);
            resp
        })
    }

    /// Send this request and parse JSON response on success.
    pub fn fetch_json<T: Deserialize>(self) -> ApiResult<T> {
        let resp = try!(self.send());
        serde_json::from_reader(resp).map_err(From::from)
    }

    /// Add body to the request.
    pub fn body<B: Into<Body<'a>>>(self, body: B)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            inner: self.inner.body(body),
            .. self
        }
    }

    /// Add additional headers to the request.
    pub fn headers(self, headers: Headers)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            inner: self.inner.headers(headers),
            .. self
        }
    }

    /// Add an individual header to the request.
    ///
    /// Note that X-Auth-Token is always overwritten with a token in use.
    pub fn header<H: Header + HeaderFormat>(self, header: H)
            -> RequestBuilder<'a, Auth> {
        RequestBuilder {
            inner: self.inner.header(header),
            .. self
        }
    }
}

impl<'session, Auth: AuthMethod + 'session, Srv: ServiceType>
        ServiceWrapper<'session, Auth, Srv> {
    /// Create a new wrapper for the specific service.
    pub fn new(session: &'session Session<Auth>)
            -> ServiceWrapper<'session, Auth, Srv> {
        ServiceWrapper {
            session: session,
            service_type: PhantomData,
            endpoint_interface: None
        }
    }

    /// Change the endpoint interface used for this wrapper.
    pub fn with_endpoint_interface(self, endpoint_interface: String)
            -> ServiceWrapper<'session, Auth, Srv> {
        ServiceWrapper {
            endpoint_interface: Some(endpoint_interface),
            .. self
        }
    }

    /// Construct and endpoint for the given service from the path.
    pub fn get_endpoint<P>(&self, path: P, query: Query) -> ApiResult<Url>
            where P: IntoIterator, P::Item: AsRef<str> {
        let ep = self.endpoint_interface.clone();
        let info = try!(self.session.get_service_info::<Srv>(ep));
        let mut url = utils::url::extend(info.root_url, path);
        let _ = url.query_pairs_mut().extend_pairs(query.0);
        Ok(url)
    }

    /// Make an HTTP request to the given service.
    pub fn request<P>(&'session self, method: Method, path: P, query: Query)
            -> ApiResult<RequestBuilder<'session, Auth>>
            where P: IntoIterator, P::Item: AsRef<str> {
        let url = try!(self.get_endpoint(path, query));
        let headers = self.session.service_headers::<Srv>();
        trace!("Sending HTTP {} request to {} with {:?}",
               method, url, headers);
        Ok(self.session.raw_request(method, url).headers(headers))
    }

    /// Make an HTTP request with JSON body and JSON response.
    pub fn json<P, Req, Res>(&self, method: Method, path: P, query: Query,
                             body: &Req) -> ApiResult<Res>
            where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        let str_body = try!(serde_json::to_string(body));
        let request = try!(self.request(method, path, query));
        request.body(&str_body).fetch_json()
    }

    /// Make a GET request returning a JSON.
    pub fn get_json<P, Res>(&self, path: P, query: Query) -> ApiResult<Res>
            where Res: Deserialize, P: IntoIterator, P::Item: AsRef<str> {
        try!(self.request(Method::Get, path, query)).fetch_json()
    }

    /// Make a POST request sending and returning a JSON.
    pub fn post_json<P, Req, Res>(&self, path: P, query: Query, body: &Req)
            -> ApiResult<Res> where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        self.json(Method::Post, path, query, body)
    }

    /// Make a POST request sending and returning a JSON.
    pub fn put_json<P, Req, Res>(&self, path: P, query: Query, body: &Req)
            -> ApiResult<Res> where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        self.json(Method::Put, path, query, body)
    }

    /// Make a PATCH request sending and returning a JSON.
    pub fn patch_json<P, Req, Res>(&self, path: P, query: Query, body: &Req)
            -> ApiResult<Res> where Req: Serialize, Res: Deserialize,
            P: IntoIterator, P::Item: AsRef<str> {
        self.json(Method::Patch, path, query, body)
    }

    /// Make a DELETE request.
    pub fn delete<P>(&self, path: P, query: Query) -> ApiResult<Response>
            where P: IntoIterator, P::Item: AsRef<str> {
        try!(self.request(Method::Delete, path, query)).send()
    }
}

impl<'session, Auth: AuthMethod + 'session, Srv: ServiceType>
        Clone for ServiceWrapper<'session, Auth, Srv> {
    fn clone(&self) -> ServiceWrapper<'session, Auth, Srv> {
        ServiceWrapper {
            session: self.session,
            service_type: PhantomData,
            endpoint_interface: self.endpoint_interface.clone()
        }
    }
}

impl ServiceInfo {
    /// Pick an API version.
    pub fn pick_api_version(&self, request: ApiVersionRequest)
            -> Option<ApiVersion> {
        match request {
            ApiVersionRequest::Minimum =>
                self.minimum_version,
            ApiVersionRequest::Latest =>
                self.current_version,
            ApiVersionRequest::Exact(req) => {
                self.current_version.and_then(|max| {
                    match self.minimum_version {
                        Some(min) if req >= min && req <= max => Some(req),
                        None if req == max => Some(req),
                        _ => None
                    }
                })
            },
            ApiVersionRequest::Choice(vec) => {
                if vec.is_empty() {
                    return None;
                }

                self.current_version.and_then(|max| {
                    match self.minimum_version {
                        Some(min) => vec.into_iter().filter(|x| {
                            *x >= min && *x <= max
                        }).max(),
                        None =>vec.into_iter().find(|x| *x == max)
                    }
                })
            }
        }
    }
}

#[cfg(test)]
pub mod test {
    use hyper::Url;

    use super::super::{ApiVersion, ApiVersionRequest};
    use super::ServiceInfo;

    fn service_info(min: Option<u16>, max: Option<u16>) -> ServiceInfo {
        ServiceInfo {
            root_url: Url::parse("http://127.0.0.1").unwrap(),
            minimum_version: min.map(|x| ApiVersion(2, x)),
            current_version: max.map(|x| ApiVersion(2, x)),
        }
    }

    #[test]
    fn test_pick_version_exact() {
        let info = service_info(Some(1), Some(24));
        let version = ApiVersion(2, 22);
        let result = info.pick_api_version(ApiVersionRequest::Exact(version))
            .unwrap();
        assert_eq!(result, version);
    }

    #[test]
    fn test_pick_version_exact_mismatch() {
        let info = service_info(Some(1), Some(24));
        let version = ApiVersion(2, 25);
        let res1 = info.pick_api_version(ApiVersionRequest::Exact(version));
        assert!(res1.is_none());
        let version2 = ApiVersion(1, 11);
        let res2 = info.pick_api_version(ApiVersionRequest::Exact(version2));
        assert!(res2.is_none());
    }

    #[test]
    fn test_pick_version_exact_current_only() {
        let info = service_info(None, Some(24));
        let version = ApiVersion(2, 24);
        let result = info.pick_api_version(ApiVersionRequest::Exact(version))
            .unwrap();
        assert_eq!(result, version);
    }

    #[test]
    fn test_pick_version_exact_current_only_mismatch() {
        let info = service_info(None, Some(24));
        let version = ApiVersion(2, 22);
        let result = info.pick_api_version(ApiVersionRequest::Exact(version));
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_minimum() {
        let info = service_info(Some(1), Some(24));
        let result = info.pick_api_version(ApiVersionRequest::Minimum)
            .unwrap();
        assert_eq!(result, ApiVersion(2, 1));
    }

    #[test]
    fn test_pick_version_minimum_unknown() {
        let info = service_info(None, Some(24));
        let result = info.pick_api_version(ApiVersionRequest::Minimum);
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_latest() {
        let info = service_info(Some(1), Some(24));
        let result = info.pick_api_version(ApiVersionRequest::Latest)
            .unwrap();
        assert_eq!(result, ApiVersion(2, 24));
    }

    #[test]
    fn test_pick_version_latest_unknown() {
        let info = service_info(Some(1), None);
        let result = info.pick_api_version(ApiVersionRequest::Latest);
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_choice() {
        let info = service_info(Some(1), Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 2),
                          ApiVersion(2, 22), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice))
            .unwrap();
        assert_eq!(result, ApiVersion(2, 22));
    }

    #[test]
    fn test_pick_version_choice_mismatch() {
        let info = service_info(Some(1), Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice));
        assert!(result.is_none());
    }

    #[test]
    fn test_pick_version_choice_current_only() {
        let info = service_info(None, Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 2),
                          ApiVersion(2, 24), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice))
            .unwrap();
        assert_eq!(result, ApiVersion(2, 24));
    }

    #[test]
    fn test_pick_version_choice_current_only_mismatch() {
        let info = service_info(None, Some(24));
        let choice = vec![ApiVersion(2, 0), ApiVersion(2, 2),
                          ApiVersion(2, 22), ApiVersion(2, 25)];
        let result = info.pick_api_version(ApiVersionRequest::Choice(choice));
        assert!(result.is_none());
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io;
use std::ops::Drop;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};

use futures::{future, Future, BoxFuture, Sink};
use futures::sync::mpsc::Sender;
use tokio_core::reactor::Remote;
use tokio_service::Service;

use bytes::{BufMut, Bytes, BytesMut};

use admin;
use codec::Packet;
use job::Job;
use packet::PacketMagic;
use queues::{HandleJobStorage, JobQueuePriority, SharedJobStorage};
use worker::{SharedWorkers, Worker, Wake};
use constants::*;

fn new_res(ptype: u32, data: Bytes) -> Packet {
    Packet {
        magic: PacketMagic::RES,
        ptype: ptype,
        psize: data.len() as u32,
        data: data,
    }
}

fn new_noop() -> Packet {
    new_res(NOOP, Bytes::new())
}

type JobWaiters = Arc<Mutex<HashMap<Bytes, Vec<usize>>>>;
type SendersByConnId = Arc<Mutex<HashMap<usize, Sender<Packet>>>>;

pub struct GearmanService {
    pub conn_id: usize,
    pub queues: SharedJobStorage,
    pub workers: SharedWorkers,
    pub worker: Arc<Mutex<Worker>>,
    pub job_count: Arc<AtomicUsize>,
    senders_by_conn_id: SendersByConnId,
    job_waiters: JobWaiters,
    remote: Remote,
}

fn next_field(buf: &mut Bytes) -> Result<Bytes, io::Error> {
    match buf[..].iter().position(|b| *b == b'\0') {
        Some(null_pos) => {
            let value = buf.split_to(null_pos);
            buf.split_to(1);
            Ok(value)
        }
        None => Err(io::Error::new(io::ErrorKind::Other, "Can't find null")),
    }
}

impl Drop for GearmanService {
    fn drop(&mut self) {
        trace!("Dropping conn_id = {}", self.conn_id);
        self.workers.shutdown(self.conn_id);
        debug!("Dropped conn_id = {}", self.conn_id);
    }
}

impl GearmanService {
    /// Things that don't require a body should use this
    fn response_from_packet(&self, packet: &Packet) -> Packet {
        match packet.ptype {
            ADMIN_VERSION => {
                let resp_str = "OK some-rustygear-version\n";
                let mut resp_body = BytesMut::with_capacity(resp_str.len());
                resp_body.put(&resp_str[..]);
                Packet {
                    magic: PacketMagic::TEXT,
                    ptype: packet.ptype,
                    psize: resp_str.len() as u32,
                    data: resp_body.freeze(),
                }
            }
            ADMIN_STATUS => admin::admin_command_status(self.queues.clone(), self.workers.clone()),
            _ => {
                panic!("response_from_packet called with invalid ptype: {}",
                       packet.ptype)
            }
        }
    }

    fn send_to_conn_id(&self, conn_id: usize, packet: Packet) {
        let senders_by_conn_id = self.senders_by_conn_id.lock().unwrap();
        match senders_by_conn_id.get(&conn_id) {
            None => {
                panic!("No connection found for conn_id = {}", &conn_id); // XXX You can do better
            }
            Some(tx) => {
                let tx = tx.clone();
                self.remote.spawn(move |handle| {
                    handle.spawn(tx.send(packet).then(|res| {
                        match res {
                            Ok(_) => {}
                            Err(e) => error!("Send Error! {:?}", e),
                        }
                        Ok(())
                    }));
                    Ok(())
                });
            }
        }
    }

    pub fn new(conn_id: usize,
               queues: SharedJobStorage,
               workers: SharedWorkers,
               job_count: Arc<AtomicUsize>,
               senders_by_conn_id: SendersByConnId,
               job_waiters: JobWaiters,
               remote: Remote)
               -> GearmanService {
        GearmanService {
            conn_id: conn_id,
            queues: queues,
            worker: Arc::new(Mutex::new((Worker::new()))),
            workers: workers,
            job_count: job_count,
            senders_by_conn_id: senders_by_conn_id,
            job_waiters: job_waiters,
            remote: remote,
        }
    }

    fn no_response() -> Packet {
        Packet {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: 0,
            data: Bytes::new(),
        }
    }

    fn handle_can_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let workers = self.workers.clone();
        let conn_id = self.conn_id;
        debug!("CAN_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.can_do(packet.data.clone());
        workers.clone().wakeup(&mut worker, conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_cant_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        debug!("CANT_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.cant_do(&packet.data);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_grab_job_all(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(4 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                // reducer not implemented
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_ALL, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job_uniq(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(3 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_UNIQ, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(2 + j.handle.len() + j.fname.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_pre_sleep(&self) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let ref mut w = worker.lock().unwrap();
        self.workers.clone().sleep(w, self.conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_submit_job(&self,
                         priority: JobQueuePriority,
                         wait: bool,
                         packet: Packet)
                         -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let conn_id = match wait {
            true => Some(self.conn_id),
            false => None,
        };
        let mut workers = self.workers.clone();
        let job_count = self.job_count.clone();
        let remote = self.remote.clone();
        let senders_by_conn_id = self.senders_by_conn_id.clone();
        let mut fields = packet.data.clone();
        trace!("fields = {:?}", fields);
        let fname = next_field(&mut fields).unwrap();
        let unique = next_field(&mut fields).unwrap();
        trace!("  --> fname = {:?} unique = {:?}", fname, unique);
        let mut add = false;
        let handle = match queues.coalesce_unique(&unique, conn_id) {
            Some(handle) => handle,
            None => {
                {
                    for wake in workers.queue_wake(&fname) {
                        let senders_by_conn_id = senders_by_conn_id.lock().unwrap();
                        match senders_by_conn_id.get(&wake) {
                            None => {
                                debug!("No connection found to wake up for conn_id = {}", wake);
                            }
                            Some(tx) => {
                                let tx = tx.clone();
                                remote.spawn(move |handle| {
                                    handle.spawn(tx.send(new_noop()).then(|res| {
                                        match res {
                                            Ok(_) => {}
                                            Err(e) => error!("Send Error! {:?}", e),
                                        }
                                        Ok(())
                                    }));
                                    Ok(())
                                });
                            }
                        }
                    }
                }
                // H:091234567890
                let mut handle = BytesMut::with_capacity(12);
                let job_num = job_count.fetch_add(1, Ordering::Relaxed);
                debug!("job_num = {}", job_num);
                handle.extend(format!("H:{:010}", job_num).as_bytes());
                add = true;
                handle.freeze()
            }
        };
        if add {
            let job = Arc::new(Job::new(fname, unique, fields, handle.clone()));
            info!("Created job {:?}", job);
            queues.add_job(job.clone(), priority, conn_id);
        }
        // If we don't store any senders, the sender will be dropped and the rx
        // stream should end thus releasing the waiter immediately.
        let psize = handle.len() as u32;
        if wait {
            // Fetch our sender
            let mut job_waiters = self.job_waiters.lock().unwrap();
            let mut waiters = job_waiters.entry(handle.clone()).or_insert(Vec::new());
            waiters.push(self.conn_id);
        }
        future::finished(Packet {
                magic: PacketMagic::RES,
                ptype: JOB_CREATED,
                psize: psize,
                data: handle,
            })
            .boxed()
    }

    fn handle_work_complete(&self, packet: Packet) -> BoxFuture<Packet, io::Error> {
        // Search for handle
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let worker = self.worker.clone();
        let queues = self.queues.clone();
        info!("Job is complete {:?}", handle);
        let mut worker = worker.lock().unwrap();
        match worker.get_assigned_job(&handle) {
            Some(ref mut j) => {
                let mut queues = queues.lock().unwrap();
                queues.remove_job(&j.unique);
            }
            None => {
                error!("WORK_COMPLETE received but no active jobs");
            }
        }
        worker.unassign_job(&handle);
        let mut job_waiters = self.job_waiters.lock().unwrap();
        // If there are waiters, send the packet to them
        if let Some(waiters) = job_waiters.remove(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }
}

impl Service for GearmanService {
    type Request = Packet;
    type Response = Packet;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        debug!("Got a req {:?}", req);
        match req.ptype {
            ADMIN_VERSION | ADMIN_STATUS => future::ok(self.response_from_packet(&req)).boxed(),
            SUBMIT_JOB => self.handle_submit_job(PRIORITY_NORMAL, true, req),
            SUBMIT_JOB_HIGH => self.handle_submit_job(PRIORITY_HIGH, true, req),
            SUBMIT_JOB_LOW => self.handle_submit_job(PRIORITY_LOW, true, req),
            SUBMIT_JOB_BG => self.handle_submit_job(PRIORITY_NORMAL, false, req),
            SUBMIT_JOB_HIGH_BG => self.handle_submit_job(PRIORITY_HIGH, false, req),
            SUBMIT_JOB_LOW_BG => self.handle_submit_job(PRIORITY_LOW, false, req),
            PRE_SLEEP => self.handle_pre_sleep(),
            CAN_DO => self.handle_can_do(&req),
            CANT_DO => self.handle_cant_do(&req),
            GRAB_JOB => self.handle_grab_job(),
            GRAB_JOB_UNIQ => self.handle_grab_job_uniq(),
            GRAB_JOB_ALL => self.handle_grab_job_all(),
            WORK_COMPLETE => self.handle_work_complete(req),/*
            WORK_STATUS | WORK_DATA | WORK_WARNING => self.handle_work_update(),
                    ECHO_REQ => self.handle_echo_req(&req),*/
            _ => {
                error!("Unimplemented: {:?} processing packet", req);
                future::err(io::Error::new(io::ErrorKind::Other,
                                           format!("Invalid packet type {}", req.ptype)))
                    .boxed()
            }
        }
    }
}
<commit_msg>Handle work data status warning etc<commit_after>use std::collections::HashMap;
use std::io;
use std::ops::Drop;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};

use futures::{future, Future, BoxFuture, Sink};
use futures::sync::mpsc::Sender;
use tokio_core::reactor::Remote;
use tokio_service::Service;

use bytes::{BufMut, Bytes, BytesMut};

use admin;
use codec::Packet;
use job::Job;
use packet::PacketMagic;
use queues::{HandleJobStorage, JobQueuePriority, SharedJobStorage};
use worker::{SharedWorkers, Worker, Wake};
use constants::*;

fn new_res(ptype: u32, data: Bytes) -> Packet {
    Packet {
        magic: PacketMagic::RES,
        ptype: ptype,
        psize: data.len() as u32,
        data: data,
    }
}

fn new_noop() -> Packet {
    new_res(NOOP, Bytes::new())
}

type JobWaiters = Arc<Mutex<HashMap<Bytes, Vec<usize>>>>;
type SendersByConnId = Arc<Mutex<HashMap<usize, Sender<Packet>>>>;

pub struct GearmanService {
    pub conn_id: usize,
    pub queues: SharedJobStorage,
    pub workers: SharedWorkers,
    pub worker: Arc<Mutex<Worker>>,
    pub job_count: Arc<AtomicUsize>,
    senders_by_conn_id: SendersByConnId,
    job_waiters: JobWaiters,
    remote: Remote,
}

fn next_field(buf: &mut Bytes) -> Result<Bytes, io::Error> {
    match buf[..].iter().position(|b| *b == b'\0') {
        Some(null_pos) => {
            let value = buf.split_to(null_pos);
            buf.split_to(1);
            Ok(value)
        }
        None => Err(io::Error::new(io::ErrorKind::Other, "Can't find null")),
    }
}

impl Drop for GearmanService {
    fn drop(&mut self) {
        trace!("Dropping conn_id = {}", self.conn_id);
        self.workers.shutdown(self.conn_id);
        debug!("Dropped conn_id = {}", self.conn_id);
    }
}

impl GearmanService {
    /// Things that don't require a body should use this
    fn response_from_packet(&self, packet: &Packet) -> Packet {
        match packet.ptype {
            ADMIN_VERSION => {
                let resp_str = "OK some-rustygear-version\n";
                let mut resp_body = BytesMut::with_capacity(resp_str.len());
                resp_body.put(&resp_str[..]);
                Packet {
                    magic: PacketMagic::TEXT,
                    ptype: packet.ptype,
                    psize: resp_str.len() as u32,
                    data: resp_body.freeze(),
                }
            }
            ADMIN_STATUS => admin::admin_command_status(self.queues.clone(), self.workers.clone()),
            _ => {
                panic!("response_from_packet called with invalid ptype: {}",
                       packet.ptype)
            }
        }
    }

    fn send_to_conn_id(&self, conn_id: usize, packet: Packet) {
        let senders_by_conn_id = self.senders_by_conn_id.lock().unwrap();
        match senders_by_conn_id.get(&conn_id) {
            None => {
                panic!("No connection found for conn_id = {}", &conn_id); // XXX You can do better
            }
            Some(tx) => {
                let tx = tx.clone();
                self.remote.spawn(move |handle| {
                    handle.spawn(tx.send(packet).then(|res| {
                        match res {
                            Ok(_) => {}
                            Err(e) => error!("Send Error! {:?}", e),
                        }
                        Ok(())
                    }));
                    Ok(())
                });
            }
        }
    }

    pub fn new(conn_id: usize,
               queues: SharedJobStorage,
               workers: SharedWorkers,
               job_count: Arc<AtomicUsize>,
               senders_by_conn_id: SendersByConnId,
               job_waiters: JobWaiters,
               remote: Remote)
               -> GearmanService {
        GearmanService {
            conn_id: conn_id,
            queues: queues,
            worker: Arc::new(Mutex::new((Worker::new()))),
            workers: workers,
            job_count: job_count,
            senders_by_conn_id: senders_by_conn_id,
            job_waiters: job_waiters,
            remote: remote,
        }
    }

    fn no_response() -> Packet {
        Packet {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: 0,
            data: Bytes::new(),
        }
    }

    fn handle_can_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let workers = self.workers.clone();
        let conn_id = self.conn_id;
        debug!("CAN_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.can_do(packet.data.clone());
        workers.clone().wakeup(&mut worker, conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_cant_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        debug!("CANT_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.cant_do(&packet.data);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_grab_job_all(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(4 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                // reducer not implemented
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_ALL, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job_uniq(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(3 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_UNIQ, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(2 + j.handle.len() + j.fname.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_pre_sleep(&self) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let ref mut w = worker.lock().unwrap();
        self.workers.clone().sleep(w, self.conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_submit_job(&self,
                         priority: JobQueuePriority,
                         wait: bool,
                         packet: Packet)
                         -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let conn_id = match wait {
            true => Some(self.conn_id),
            false => None,
        };
        let mut workers = self.workers.clone();
        let job_count = self.job_count.clone();
        let remote = self.remote.clone();
        let senders_by_conn_id = self.senders_by_conn_id.clone();
        let mut fields = packet.data.clone();
        trace!("fields = {:?}", fields);
        let fname = next_field(&mut fields).unwrap();
        let unique = next_field(&mut fields).unwrap();
        trace!("  --> fname = {:?} unique = {:?}", fname, unique);
        let mut add = false;
        let handle = match queues.coalesce_unique(&unique, conn_id) {
            Some(handle) => handle,
            None => {
                {
                    for wake in workers.queue_wake(&fname) {
                        let senders_by_conn_id = senders_by_conn_id.lock().unwrap();
                        match senders_by_conn_id.get(&wake) {
                            None => {
                                debug!("No connection found to wake up for conn_id = {}", wake);
                            }
                            Some(tx) => {
                                let tx = tx.clone();
                                remote.spawn(move |handle| {
                                    handle.spawn(tx.send(new_noop()).then(|res| {
                                        match res {
                                            Ok(_) => {}
                                            Err(e) => error!("Send Error! {:?}", e),
                                        }
                                        Ok(())
                                    }));
                                    Ok(())
                                });
                            }
                        }
                    }
                }
                // H:091234567890
                let mut handle = BytesMut::with_capacity(12);
                let job_num = job_count.fetch_add(1, Ordering::Relaxed);
                debug!("job_num = {}", job_num);
                handle.extend(format!("H:{:010}", job_num).as_bytes());
                add = true;
                handle.freeze()
            }
        };
        if add {
            let job = Arc::new(Job::new(fname, unique, fields, handle.clone()));
            info!("Created job {:?}", job);
            queues.add_job(job.clone(), priority, conn_id);
        }
        // If we don't store any senders, the sender will be dropped and the rx
        // stream should end thus releasing the waiter immediately.
        let psize = handle.len() as u32;
        if wait {
            // Fetch our sender
            let mut job_waiters = self.job_waiters.lock().unwrap();
            let mut waiters = job_waiters.entry(handle.clone()).or_insert(Vec::new());
            waiters.push(self.conn_id);
        }
        future::finished(Packet {
                magic: PacketMagic::RES,
                ptype: JOB_CREATED,
                psize: psize,
                data: handle,
            })
            .boxed()
    }

    fn handle_work_complete(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        // Search for handle
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let worker = self.worker.clone();
        let queues = self.queues.clone();
        info!("Job is complete {:?}", handle);
        let mut worker = worker.lock().unwrap();
        match worker.get_assigned_job(&handle) {
            Some(ref mut j) => {
                let mut queues = queues.lock().unwrap();
                queues.remove_job(&j.unique);
            }
            None => {
                error!("WORK_COMPLETE received but no active jobs");
            }
        }
        worker.unassign_job(&handle);
        let mut job_waiters = self.job_waiters.lock().unwrap();
        // If there are waiters, send the packet to them
        if let Some(waiters) = job_waiters.remove(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }

    fn handle_work_update(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let job_waiters = self.job_waiters.lock().unwrap();
        if let Some(waiters) = job_waiters.get(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }
}

impl Service for GearmanService {
    type Request = Packet;
    type Response = Packet;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        debug!("Got a req {:?}", req);
        match req.ptype {
            ADMIN_VERSION | ADMIN_STATUS => future::ok(self.response_from_packet(&req)).boxed(),
            SUBMIT_JOB => self.handle_submit_job(PRIORITY_NORMAL, true, req),
            SUBMIT_JOB_HIGH => self.handle_submit_job(PRIORITY_HIGH, true, req),
            SUBMIT_JOB_LOW => self.handle_submit_job(PRIORITY_LOW, true, req),
            SUBMIT_JOB_BG => self.handle_submit_job(PRIORITY_NORMAL, false, req),
            SUBMIT_JOB_HIGH_BG => self.handle_submit_job(PRIORITY_HIGH, false, req),
            SUBMIT_JOB_LOW_BG => self.handle_submit_job(PRIORITY_LOW, false, req),
            PRE_SLEEP => self.handle_pre_sleep(),
            CAN_DO => self.handle_can_do(&req),
            CANT_DO => self.handle_cant_do(&req),
            GRAB_JOB => self.handle_grab_job(),
            GRAB_JOB_UNIQ => self.handle_grab_job_uniq(),
            GRAB_JOB_ALL => self.handle_grab_job_all(),
            WORK_COMPLETE => self.handle_work_complete(&req),
            WORK_STATUS | WORK_DATA | WORK_WARNING => self.handle_work_update(&req),/*
                    ECHO_REQ => self.handle_echo_req(&req),*/
            _ => {
                error!("Unimplemented: {:?} processing packet", req);
                future::err(io::Error::new(io::ErrorKind::Other,
                                           format!("Invalid packet type {}", req.ptype)))
                    .boxed()
            }
        }
    }
}
<|endoftext|>


<commit_before>#[doc = "Prunes branches of the document tree that contain no documentation"];

export mk_pass;

fn mk_pass() -> pass {
    run
}

type ctxt = {
    mutable have_docs: bool
};

fn run(
    _srv: astsrv::srv,
    doc: doc::cratedoc
) -> doc::cratedoc {
    let ctxt = {
        mutable have_docs: true
    };
    let fold = fold::fold({
        fold_mod: fold_mod,
        fold_fn: fold_fn,
        fold_modlist: fold_modlist,
        fold_fnlist: fold_fnlist
        with *fold::default_seq_fold(ctxt)
    });
    fold.fold_crate(fold, doc)
}

fn fold_mod(
    fold: fold::fold<ctxt>,
    doc: doc::moddoc
) -> doc::moddoc {
    let doc = fold::default_seq_fold_mod(fold, doc);
    fold.ctxt.have_docs =
        doc.brief != none
        || doc.desc != none
        || vec::is_not_empty(*doc.mods)
        || vec::is_not_empty(*doc.fns);
    ret doc;
}

fn fold_fn(
    fold: fold::fold<ctxt>,
    doc: doc::fndoc
) -> doc::fndoc {
    fold.ctxt.have_docs =
        doc.brief != none
        || doc.desc != none
        || doc.return.desc != none;
    ret doc;
}

fn fold_modlist(
    fold: fold::fold<ctxt>,
    list: doc::modlist
) -> doc::modlist {
    doc::modlist(vec::filter_map(*list) {|doc|
        let doc = fold_mod(fold, doc);
        if fold.ctxt.have_docs {
            some(doc)
        } else {
            none
        }
    })
}

fn fold_fnlist(
    fold: fold::fold<ctxt>,
    list: doc::fnlist
) -> doc::fnlist {
    doc::fnlist(vec::filter_map(*list) {|doc|
        let doc = fold_fn(fold, doc);
        if fold.ctxt.have_docs {
            some(doc)
        } else {
            none
        }
    })
}

#[cfg(test)]
mod tests {
    #[test]
    fn should_elide_undocumented_mods() {
        let source = "mod a { }";
        let srv = astsrv::mk_srv_from_str(source);
        let doc = extract::from_srv(srv, "");
        let doc = run(srv, doc);
        assert vec::is_empty(*doc.topmod.mods);
    }

    #[test]
    fn should_not_elide_undocument_mods_with_documented_mods() {
        let source = "mod a { #[doc = \"b\"] mod b { } }";
        let srv = astsrv::mk_srv_from_str(source);
        let doc = extract::from_srv(srv, "");
        let doc = attr_pass::mk_pass()(srv, doc);
        let doc = run(srv, doc);
        assert vec::is_not_empty(*doc.topmod.mods);
    }

    #[test]
    fn should_not_elide_undocument_mods_with_documented_fns() {
        let source = "mod a { #[doc = \"b\"] fn b() { } }";
        let srv = astsrv::mk_srv_from_str(source);
        let doc = extract::from_srv(srv, "");
        let doc = attr_pass::mk_pass()(srv, doc);
        let doc = run(srv, doc);
        assert vec::is_not_empty(*doc.topmod.mods);
    }

    #[test]
    fn should_elide_undocumented_fns() {
        let source = "fn a() { }";
        let srv = astsrv::mk_srv_from_str(source);
        let doc = extract::from_srv(srv, "");
        let doc = run(srv, doc);
        assert vec::is_empty(*doc.topmod.fns);
    }
}<commit_msg>rustdoc: Prune undocumented fn args<commit_after>#[doc = "Prunes branches of the document tree that contain no documentation"];

export mk_pass;

fn mk_pass() -> pass {
    run
}

type ctxt = {
    mutable have_docs: bool
};

fn run(
    _srv: astsrv::srv,
    doc: doc::cratedoc
) -> doc::cratedoc {
    let ctxt = {
        mutable have_docs: true
    };
    let fold = fold::fold({
        fold_mod: fold_mod,
        fold_fn: fold_fn,
        fold_modlist: fold_modlist,
        fold_fnlist: fold_fnlist
        with *fold::default_seq_fold(ctxt)
    });
    fold.fold_crate(fold, doc)
}

fn fold_mod(
    fold: fold::fold<ctxt>,
    doc: doc::moddoc
) -> doc::moddoc {
    let doc = fold::default_seq_fold_mod(fold, doc);
    fold.ctxt.have_docs =
        doc.brief != none
        || doc.desc != none
        || vec::is_not_empty(*doc.mods)
        || vec::is_not_empty(*doc.fns);
    ret doc;
}

fn fold_fn(
    fold: fold::fold<ctxt>,
    doc: doc::fndoc
) -> doc::fndoc {
    let doc = ~{
        args: vec::filter_map(doc.args) {|doc|
            if option::is_some(doc.desc) {
                some(doc)
            } else {
                none
            }
        }
        with *doc
    };

    fold.ctxt.have_docs =
        doc.brief != none
        || doc.desc != none
        || doc.return.desc != none;
    ret doc;
}

#[test]
fn should_elide_undocumented_arguments() {
    let source = "#[doc = \"hey\"] fn a(b: int) { }";
    let srv = astsrv::mk_srv_from_str(source);
    let doc = extract::from_srv(srv, "");
    let doc = attr_pass::mk_pass()(srv, doc);
    let doc = run(srv, doc);
    assert vec::is_empty(doc.topmod.fns[0].args);
}

fn fold_modlist(
    fold: fold::fold<ctxt>,
    list: doc::modlist
) -> doc::modlist {
    doc::modlist(vec::filter_map(*list) {|doc|
        let doc = fold_mod(fold, doc);
        if fold.ctxt.have_docs {
            some(doc)
        } else {
            none
        }
    })
}

#[test]
fn should_elide_undocumented_mods() {
    let source = "mod a { }";
    let srv = astsrv::mk_srv_from_str(source);
    let doc = extract::from_srv(srv, "");
    let doc = run(srv, doc);
    assert vec::is_empty(*doc.topmod.mods);
}

#[test]
fn should_not_elide_undocument_mods_with_documented_mods() {
    let source = "mod a { #[doc = \"b\"] mod b { } }";
    let srv = astsrv::mk_srv_from_str(source);
    let doc = extract::from_srv(srv, "");
    let doc = attr_pass::mk_pass()(srv, doc);
    let doc = run(srv, doc);
    assert vec::is_not_empty(*doc.topmod.mods);
}

#[test]
fn should_not_elide_undocument_mods_with_documented_fns() {
    let source = "mod a { #[doc = \"b\"] fn b() { } }";
    let srv = astsrv::mk_srv_from_str(source);
    let doc = extract::from_srv(srv, "");
    let doc = attr_pass::mk_pass()(srv, doc);
    let doc = run(srv, doc);
    assert vec::is_not_empty(*doc.topmod.mods);
}

fn fold_fnlist(
    fold: fold::fold<ctxt>,
    list: doc::fnlist
) -> doc::fnlist {
    doc::fnlist(vec::filter_map(*list) {|doc|
        let doc = fold_fn(fold, doc);
        if fold.ctxt.have_docs {
            some(doc)
        } else {
            none
        }
    })
}

#[test]
fn should_elide_undocumented_fns() {
    let source = "fn a() { }";
    let srv = astsrv::mk_srv_from_str(source);
    let doc = extract::from_srv(srv, "");
    let doc = run(srv, doc);
    assert vec::is_empty(*doc.topmod.fns);
}
<|endoftext|>


<commit_before>///
/// Possible error conditions from a `try_sync()` call on the scheduler
///
#[derive( Clone, PartialEq, Debug)]
pub enum TrySyncError {
    /// The queue is busy, so the function has not been executed
    Busy
}
<commit_msg>Add an Into to convert the try_sync error into an io::ErrorKind error<commit_after>use std::io::{ErrorKind};

///
/// Possible error conditions from a `try_sync()` call on the scheduler
///
#[derive(Clone, PartialEq, Debug)]
pub enum TrySyncError {
    /// The queue is busy, so the function has not been executed
    Busy
}

impl Into<ErrorKind> for TrySyncError {
    fn into(self) -> ErrorKind {
        match self {
            TrySyncError::Busy => ErrorKind::WouldBlock
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{Reader, Writer, IoResult, IoError, ConnectionFailed};
use std::io::net::tcp::TcpStream;
use std::cell::RefCell;
use std::collections::HashSet;

use cmd::cmd;
use types::{RedisResult, Okay, Error, Value, InternalIoError,
            ToRedisArgs, FromRedisValue};
use parser::Parser;

struct ActualConnection {
    sock: TcpStream,
}

/// Represents a stateful redis TCP connection.
pub struct Connection {
    con: RefCell<ActualConnection>,
    db: i64,
}

/// Represents a pubsub connection.
pub struct PubSub {
    con: Connection,
    channels: HashSet<Vec<u8>>,
}

impl ActualConnection {

    pub fn new(host: &str, port: u16) -> IoResult<ActualConnection> {
        let sock = try!(TcpStream::connect(host, port));
        Ok(ActualConnection { sock: sock })
    }

    pub fn send_bytes(&mut self, bytes: &[u8]) -> RedisResult<Value> {
        let w = &mut self.sock as &mut Writer;
        match w.write(bytes) {
            Err(err) => {
                Err(Error::simple(
                    InternalIoError(err),
                    "Could not send command because of an IO error"))
            },
            Ok(_) => Ok(Okay)
        }
    }

    pub fn read_response(&mut self) -> RedisResult<Value> {
        let mut parser = Parser::new(&mut self.sock as &mut Reader);
        parser.parse_value()
    }
}

pub fn connect(host: &str, port: u16, db: i64) -> IoResult<Connection> {
    let con = try!(ActualConnection::new(host, port));
    let rv = Connection { con: RefCell::new(con), db: db };
    if db != 0 {
        match cmd("SELECT").arg(db).query::<Value>(&rv) {
            Ok(Okay) => {}
            _ => { return Err(IoError {
                kind: ConnectionFailed,
                desc: "Redis server refused to switch database",
                detail: None,
            }); }
        }
    }
    Ok(rv)
}

pub fn connect_pubsub(host: &str, port: u16) -> IoResult<PubSub> {
    Ok(PubSub {
        con: try!(connect(host, port, 0)),
        channels: HashSet::new(),
    })
}

/// A connection is an object that represents a single redis connection.  It
/// provides basic support for sending encoded commands into a redis connection
/// and to read a response from it.  It's bound to a single database and can
/// only be created from the client.
///
/// You generally do not much with this object other than passing it to
/// `Cmd` objects.
impl Connection {

    /// Sends an already encoded (packed) command into the TCP socket and
    /// reads the single response from it.
    pub fn req_packed_command(&self, cmd: &[u8]) -> RedisResult<Value> {
        let mut con = self.con.borrow_mut();
        try!(con.send_bytes(cmd));
        con.read_response()
    }

    /// Sends an already encoded (packed) command into the TCP socket and
    /// does not read a response.  This is useful for commands like
    /// `MONITOR` which yield multiple items.  This needs to be used with
    /// care because it changes the state of the connection.
    pub unsafe fn send_packed_command(&self, cmd: &[u8]) -> RedisResult<()> {
        try!(self.con.borrow_mut().send_bytes(cmd));
        Ok(())
    }

    /// Fetches a single response from the connection.  This is useful
    /// if used in combination with `send_packed_command`.
    pub unsafe fn recv_response(&self) -> RedisResult<Value> {
        self.con.borrow_mut().read_response()
    }

    /// Sends multiple already encoded (packed) command into the TCP socket
    /// and reads `count` responses from it.  This is used to implement
    /// pipelining.
    pub fn req_packed_commands(&self, cmd: &[u8],
            offset: uint, count: uint) -> RedisResult<Vec<Value>> {
        let mut con = self.con.borrow_mut();
        try!(con.send_bytes(cmd));
        let mut rv = vec![];
        for idx in range(0, offset + count) {
            let item = try!(con.read_response());
            if idx >= offset {
                rv.push(item);
            }
        }
        Ok(rv)
    }

    /// Returns the database this connection is bound to.
    pub fn get_db(&self) -> i64 {
        self.db
    }
}

/// The pubsub object provides convenient access to the redis pubsub
/// system.  Once created you can subscribe and unsubscribe from channels
/// and listen in on messages.
///
/// Example:
///
/// ```rust,no_run
/// let client = redis::Client::open("redis://127.0.0.1/").unwrap();
/// let mut pubsub = client.pubsub();
/// pubsub.subscribe("channel_1").unwrap();
/// pubsub.subscribe("channel_2").unwrap();
///
/// loop {
///     let (channel, message) : (String, String) = pubsub.get_message();
///     println!("channel '{}': {}", channel, message);
/// }
/// ```
impl PubSub {

    fn get_channel<T: ToRedisArgs>(&mut self, channel: &T) -> Vec<u8> {
        let mut chan = vec![];
        for item in channel.to_redis_args().iter() {
            chan.push_all(item[]);
        }
        chan
    }

    /// Subscribes to a new channel.
    pub fn subscribe<T: ToRedisArgs>(&mut self, channel: T) -> RedisResult<()> {
        let chan = self.get_channel(&channel);
        let _ : () = try!(cmd("SUBSCRIBE").arg(chan[]).query(&self.con));
        self.channels.insert(chan);
        Ok(())
    }

    /// Unsubscribes from a channel.
    pub fn unsubscribe<T: ToRedisArgs>(&mut self, channel: T) -> RedisResult<()> {
        let chan = self.get_channel(&channel);
        let _ : () = try!(cmd("UNSUBSCRIBE").arg(chan[]).query(&self.con));
        self.channels.remove(&chan);
        Ok(())
    }

    /// Fetches the next message from the pubsub connection.  Blocks until
    /// a message becomes available.  This currently does not provide a
    /// wait not to block :(
    ///
    /// The return value is in the format ``(channel, message)``.
    pub fn get_message<C: FromRedisValue, M: FromRedisValue>(&self) -> RedisResult<(C, M)> {
        loop {
            unsafe {
                let (msg_type, channel, msg) : (String, C, M) =
                    try!(FromRedisValue::from_redis_value(&try!(self.con.recv_response())));
                if msg_type.as_slice() == "message" {
                    return Ok((channel, msg))
                }
            }
        }
    }
}
<commit_msg>Fixed a broken docs example<commit_after>use std::io::{Reader, Writer, IoResult, IoError, ConnectionFailed};
use std::io::net::tcp::TcpStream;
use std::cell::RefCell;
use std::collections::HashSet;

use cmd::cmd;
use types::{RedisResult, Okay, Error, Value, InternalIoError,
            ToRedisArgs, FromRedisValue};
use parser::Parser;

struct ActualConnection {
    sock: TcpStream,
}

/// Represents a stateful redis TCP connection.
pub struct Connection {
    con: RefCell<ActualConnection>,
    db: i64,
}

/// Represents a pubsub connection.
pub struct PubSub {
    con: Connection,
    channels: HashSet<Vec<u8>>,
}

impl ActualConnection {

    pub fn new(host: &str, port: u16) -> IoResult<ActualConnection> {
        let sock = try!(TcpStream::connect(host, port));
        Ok(ActualConnection { sock: sock })
    }

    pub fn send_bytes(&mut self, bytes: &[u8]) -> RedisResult<Value> {
        let w = &mut self.sock as &mut Writer;
        match w.write(bytes) {
            Err(err) => {
                Err(Error::simple(
                    InternalIoError(err),
                    "Could not send command because of an IO error"))
            },
            Ok(_) => Ok(Okay)
        }
    }

    pub fn read_response(&mut self) -> RedisResult<Value> {
        let mut parser = Parser::new(&mut self.sock as &mut Reader);
        parser.parse_value()
    }
}

pub fn connect(host: &str, port: u16, db: i64) -> IoResult<Connection> {
    let con = try!(ActualConnection::new(host, port));
    let rv = Connection { con: RefCell::new(con), db: db };
    if db != 0 {
        match cmd("SELECT").arg(db).query::<Value>(&rv) {
            Ok(Okay) => {}
            _ => { return Err(IoError {
                kind: ConnectionFailed,
                desc: "Redis server refused to switch database",
                detail: None,
            }); }
        }
    }
    Ok(rv)
}

pub fn connect_pubsub(host: &str, port: u16) -> IoResult<PubSub> {
    Ok(PubSub {
        con: try!(connect(host, port, 0)),
        channels: HashSet::new(),
    })
}

/// A connection is an object that represents a single redis connection.  It
/// provides basic support for sending encoded commands into a redis connection
/// and to read a response from it.  It's bound to a single database and can
/// only be created from the client.
///
/// You generally do not much with this object other than passing it to
/// `Cmd` objects.
impl Connection {

    /// Sends an already encoded (packed) command into the TCP socket and
    /// reads the single response from it.
    pub fn req_packed_command(&self, cmd: &[u8]) -> RedisResult<Value> {
        let mut con = self.con.borrow_mut();
        try!(con.send_bytes(cmd));
        con.read_response()
    }

    /// Sends an already encoded (packed) command into the TCP socket and
    /// does not read a response.  This is useful for commands like
    /// `MONITOR` which yield multiple items.  This needs to be used with
    /// care because it changes the state of the connection.
    pub unsafe fn send_packed_command(&self, cmd: &[u8]) -> RedisResult<()> {
        try!(self.con.borrow_mut().send_bytes(cmd));
        Ok(())
    }

    /// Fetches a single response from the connection.  This is useful
    /// if used in combination with `send_packed_command`.
    pub unsafe fn recv_response(&self) -> RedisResult<Value> {
        self.con.borrow_mut().read_response()
    }

    /// Sends multiple already encoded (packed) command into the TCP socket
    /// and reads `count` responses from it.  This is used to implement
    /// pipelining.
    pub fn req_packed_commands(&self, cmd: &[u8],
            offset: uint, count: uint) -> RedisResult<Vec<Value>> {
        let mut con = self.con.borrow_mut();
        try!(con.send_bytes(cmd));
        let mut rv = vec![];
        for idx in range(0, offset + count) {
            let item = try!(con.read_response());
            if idx >= offset {
                rv.push(item);
            }
        }
        Ok(rv)
    }

    /// Returns the database this connection is bound to.
    pub fn get_db(&self) -> i64 {
        self.db
    }
}

/// The pubsub object provides convenient access to the redis pubsub
/// system.  Once created you can subscribe and unsubscribe from channels
/// and listen in on messages.
///
/// Example:
///
/// ```rust,no_run
/// let client = redis::Client::open("redis://127.0.0.1/").unwrap();
/// let mut pubsub = client.get_pubsub().unwrap();
/// pubsub.subscribe("channel_1").unwrap();
/// pubsub.subscribe("channel_2").unwrap();
///
/// loop {
///     let (channel, message) : (String, String) = pubsub.get_message().unwrap();
///     println!("channel '{}': {}", channel, message);
/// }
/// ```
impl PubSub {

    fn get_channel<T: ToRedisArgs>(&mut self, channel: &T) -> Vec<u8> {
        let mut chan = vec![];
        for item in channel.to_redis_args().iter() {
            chan.push_all(item[]);
        }
        chan
    }

    /// Subscribes to a new channel.
    pub fn subscribe<T: ToRedisArgs>(&mut self, channel: T) -> RedisResult<()> {
        let chan = self.get_channel(&channel);
        let _ : () = try!(cmd("SUBSCRIBE").arg(chan[]).query(&self.con));
        self.channels.insert(chan);
        Ok(())
    }

    /// Unsubscribes from a channel.
    pub fn unsubscribe<T: ToRedisArgs>(&mut self, channel: T) -> RedisResult<()> {
        let chan = self.get_channel(&channel);
        let _ : () = try!(cmd("UNSUBSCRIBE").arg(chan[]).query(&self.con));
        self.channels.remove(&chan);
        Ok(())
    }

    /// Fetches the next message from the pubsub connection.  Blocks until
    /// a message becomes available.  This currently does not provide a
    /// wait not to block :(
    ///
    /// The return value is in the format ``(channel, message)``.
    pub fn get_message<C: FromRedisValue, M: FromRedisValue>(&self) -> RedisResult<(C, M)> {
        loop {
            unsafe {
                let (msg_type, channel, msg) : (String, C, M) =
                    try!(FromRedisValue::from_redis_value(&try!(self.con.recv_response())));
                if msg_type.as_slice() == "message" {
                    return Ok((channel, msg))
                }
            }
        }
    }
}
<|endoftext|>