<commit_before>//! Module contianing the core functionality for OAuth2 Authentication.
use crate::authenticator_delegate::{DeviceFlowDelegate, InstalledFlowDelegate};
use crate::device::DeviceFlow;
use crate::error::Error;
use crate::installed::{InstalledFlow, InstalledFlowReturnMethod};
use crate::refresh::RefreshFlow;
use crate::service_account::{ServiceAccountFlow, ServiceAccountFlowOpts, ServiceAccountKey};
use crate::storage::{self, Storage};
use crate::types::{ApplicationSecret, Token};
use private::AuthFlow;

use std::borrow::Cow;
use std::io;
use std::path::PathBuf;
use std::sync::Mutex;

/// Authenticator is responsible for fetching tokens, handling refreshing tokens,
/// and optionally persisting tokens to disk.
pub struct Authenticator<C> {
    hyper_client: hyper::Client<C>,
    storage: Storage,
    auth_flow: AuthFlow,
}

impl<C> Authenticator<C>
where
    C: hyper::client::connect::Connect + 'static,
{
    /// Return the current token for the provided scopes.
    pub async fn token<'a, T>(&'a self, scopes: &'a [T]) -> Result<Token, Error>
    where
        T: AsRef<str>,
    {
        let hashed_scopes = storage::ScopeSet::from(scopes);
        match (self.storage.get(hashed_scopes), self.auth_flow.app_secret()) {
            (Some(t), _) if !t.expired() => {
                // unexpired token found
                Ok(t)
            }
            (
                Some(Token {
                    refresh_token: Some(refresh_token),
                    ..
                }),
                Some(app_secret),
            ) => {
                // token is expired but has a refresh token.
                let token =
                    RefreshFlow::refresh_token(&self.hyper_client, app_secret, &refresh_token)
                        .await?;
                self.storage.set(hashed_scopes, token.clone()).await;
                Ok(token)
            }
            _ => {
                // no token in the cache or the token returned can't be refreshed.
                let t = self.auth_flow.token(&self.hyper_client, scopes).await?;
                self.storage.set(hashed_scopes, t.clone()).await;
                Ok(t)
            }
        }
    }
}

/// Configure an Authenticator using the builder pattern.
pub struct AuthenticatorBuilder<C, F> {
    hyper_client_builder: C,
    storage_type: StorageType,
    auth_flow: F,
}

/// Create an authenticator that uses the installed flow.
/// ```
/// # async fn foo() {
/// # use yup_oauth2::InstalledFlowReturnMethod;
/// # let custom_flow_delegate = yup_oauth2::authenticator_delegate::DefaultInstalledFlowDelegate;
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::InstalledFlowAuthenticator::builder(
///         app_secret,
///         InstalledFlowReturnMethod::HTTPRedirect,
///     )
///     .build()
///     .await
///     .expect("failed to create authenticator");
/// # }
/// ```
pub struct InstalledFlowAuthenticator;
impl InstalledFlowAuthenticator {
    /// Use the builder pattern to create an Authenticator that uses the installed flow.
    pub fn builder(
        app_secret: ApplicationSecret,
        method: InstalledFlowReturnMethod,
    ) -> AuthenticatorBuilder<DefaultHyperClient, InstalledFlow> {
        AuthenticatorBuilder::<DefaultHyperClient, _>::with_auth_flow(InstalledFlow::new(
            app_secret, method,
        ))
    }
}

/// Create an authenticator that uses the device flow.
/// ```
/// # async fn foo() {
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::DeviceFlowAuthenticator::builder(app_secret)
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
pub struct DeviceFlowAuthenticator;
impl DeviceFlowAuthenticator {
    /// Use the builder pattern to create an Authenticator that uses the device flow.
    pub fn builder(
        app_secret: ApplicationSecret,
    ) -> AuthenticatorBuilder<DefaultHyperClient, DeviceFlow> {
        AuthenticatorBuilder::<DefaultHyperClient, _>::with_auth_flow(DeviceFlow::new(app_secret))
    }
}

/// Create an authenticator that uses a service account.
/// ```
/// # async fn foo() {
/// # let service_account_key = yup_oauth2::read_service_account_key("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::ServiceAccountAuthenticator::builder(service_account_key)
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
pub struct ServiceAccountAuthenticator;
impl ServiceAccountAuthenticator {
    /// Use the builder pattern to create an Authenticator that uses a service account.
    pub fn builder(
        service_account_key: ServiceAccountKey,
    ) -> AuthenticatorBuilder<DefaultHyperClient, ServiceAccountFlowOpts> {
        AuthenticatorBuilder::<DefaultHyperClient, _>::with_auth_flow(ServiceAccountFlowOpts {
            key: service_account_key,
            subject: None,
        })
    }
}

/// ## Methods available when building any Authenticator.
/// ```
/// # async fn foo() {
/// # let custom_hyper_client = hyper::Client::new();
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::DeviceFlowAuthenticator::builder(app_secret)
///         .hyper_client(custom_hyper_client)
///         .persist_tokens_to_disk("/tmp/tokenfile.json")
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
impl<C, F> AuthenticatorBuilder<C, F> {
    async fn common_build(
        hyper_client_builder: C,
        storage_type: StorageType,
        auth_flow: AuthFlow,
    ) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        let hyper_client = hyper_client_builder.build_hyper_client();
        let storage = match storage_type {
            StorageType::Memory => Storage::Memory {
                tokens: Mutex::new(storage::JSONTokens::new()),
            },
            StorageType::Disk(path) => Storage::Disk(storage::DiskStorage::new(path).await?),
        };

        Ok(Authenticator {
            hyper_client,
            storage,
            auth_flow,
        })
    }

    fn with_auth_flow(auth_flow: F) -> AuthenticatorBuilder<DefaultHyperClient, F> {
        AuthenticatorBuilder {
            hyper_client_builder: DefaultHyperClient,
            storage_type: StorageType::Memory,
            auth_flow,
        }
    }

    /// Use the provided hyper client.
    pub fn hyper_client<NewC>(
        self,
        hyper_client: hyper::Client<NewC>,
    ) -> AuthenticatorBuilder<hyper::Client<NewC>, F> {
        AuthenticatorBuilder {
            hyper_client_builder: hyper_client,
            storage_type: self.storage_type,
            auth_flow: self.auth_flow,
        }
    }

    /// Persist tokens to disk in the provided filename.
    pub fn persist_tokens_to_disk<P: Into<PathBuf>>(self, path: P) -> AuthenticatorBuilder<C, F> {
        AuthenticatorBuilder {
            storage_type: StorageType::Disk(path.into()),
            ..self
        }
    }
}

/// ## Methods available when building a device flow Authenticator.
/// ```
/// # async fn foo() {
/// # let custom_flow_delegate = yup_oauth2::authenticator_delegate::DefaultDeviceFlowDelegate;
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::DeviceFlowAuthenticator::builder(app_secret)
///         .device_code_url("foo")
///         .flow_delegate(Box::new(custom_flow_delegate))
///         .grant_type("foo")
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
impl<C> AuthenticatorBuilder<C, DeviceFlow> {
    /// Use the provided device code url.
    pub fn device_code_url(self, url: impl Into<Cow<'static, str>>) -> Self {
        AuthenticatorBuilder {
            auth_flow: DeviceFlow {
                device_code_url: url.into(),
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Use the provided DeviceFlowDelegate.
    pub fn flow_delegate(self, flow_delegate: Box<dyn DeviceFlowDelegate>) -> Self {
        AuthenticatorBuilder {
            auth_flow: DeviceFlow {
                flow_delegate,
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Use the provided grant type.
    pub fn grant_type(self, grant_type: impl Into<Cow<'static, str>>) -> Self {
        AuthenticatorBuilder {
            auth_flow: DeviceFlow {
                grant_type: grant_type.into(),
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Create the authenticator.
    pub async fn build(self) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        Self::common_build(
            self.hyper_client_builder,
            self.storage_type,
            AuthFlow::DeviceFlow(self.auth_flow),
        )
        .await
    }
}

/// ## Methods available when building an installed flow Authenticator.
/// ```
/// # async fn foo() {
/// # use yup_oauth2::InstalledFlowReturnMethod;
/// # let custom_flow_delegate = yup_oauth2::authenticator_delegate::DefaultInstalledFlowDelegate;
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::InstalledFlowAuthenticator::builder(
///         app_secret,
///         InstalledFlowReturnMethod::HTTPRedirect,
///     )
///     .flow_delegate(Box::new(custom_flow_delegate))
///     .build()
///     .await
///     .expect("failed to create authenticator");
/// # }
/// ```
impl<C> AuthenticatorBuilder<C, InstalledFlow> {
    /// Use the provided InstalledFlowDelegate.
    pub fn flow_delegate(self, flow_delegate: Box<dyn InstalledFlowDelegate>) -> Self {
        AuthenticatorBuilder {
            auth_flow: InstalledFlow {
                flow_delegate,
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Create the authenticator.
    pub async fn build(self) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        Self::common_build(
            self.hyper_client_builder,
            self.storage_type,
            AuthFlow::InstalledFlow(self.auth_flow),
        )
        .await
    }
}

/// ## Methods available when building a service account authenticator.
/// ```
/// # async fn foo() {
/// # let service_account_key = yup_oauth2::read_service_account_key("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::ServiceAccountAuthenticator::builder(
///         service_account_key,
///     )
///     .subject("mysubject")
///     .build()
///     .await
///     .expect("failed to create authenticator");
/// # }
/// ```
impl<C> AuthenticatorBuilder<C, ServiceAccountFlowOpts> {
    /// Use the provided subject.
    pub fn subject(self, subject: impl Into<String>) -> Self {
        AuthenticatorBuilder {
            auth_flow: ServiceAccountFlowOpts {
                subject: Some(subject.into()),
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Create the authenticator.
    pub async fn build(self) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        let service_account_auth_flow = ServiceAccountFlow::new(self.auth_flow)?;
        Self::common_build(
            self.hyper_client_builder,
            self.storage_type,
            AuthFlow::ServiceAccountFlow(service_account_auth_flow),
        )
        .await
    }
}

mod private {
    use crate::device::DeviceFlow;
    use crate::error::Error;
    use crate::installed::InstalledFlow;
    use crate::service_account::ServiceAccountFlow;
    use crate::types::{ApplicationSecret, Token};

    pub enum AuthFlow {
        DeviceFlow(DeviceFlow),
        InstalledFlow(InstalledFlow),
        ServiceAccountFlow(ServiceAccountFlow),
    }

    impl AuthFlow {
        pub(crate) fn app_secret(&self) -> Option<&ApplicationSecret> {
            match self {
                AuthFlow::DeviceFlow(device_flow) => Some(&device_flow.app_secret),
                AuthFlow::InstalledFlow(installed_flow) => Some(&installed_flow.app_secret),
                AuthFlow::ServiceAccountFlow(_) => None,
            }
        }

        pub(crate) async fn token<'a, C, T>(
            &'a self,
            hyper_client: &'a hyper::Client<C>,
            scopes: &'a [T],
        ) -> Result<Token, Error>
        where
            T: AsRef<str>,
            C: hyper::client::connect::Connect + 'static,
        {
            match self {
                AuthFlow::DeviceFlow(device_flow) => device_flow.token(hyper_client, scopes).await,
                AuthFlow::InstalledFlow(installed_flow) => {
                    installed_flow.token(hyper_client, scopes).await
                }
                AuthFlow::ServiceAccountFlow(service_account_flow) => {
                    service_account_flow.token(hyper_client, scopes).await
                }
            }
        }
    }
}

/// A trait implemented for any hyper::Client as well as the DefaultHyperClient.
pub trait HyperClientBuilder {
    /// The hyper connector that the resulting hyper client will use.
    type Connector: hyper::client::connect::Connect + 'static;

    /// Create a hyper::Client
    fn build_hyper_client(self) -> hyper::Client<Self::Connector>;
}

/// The builder value used when the default hyper client should be used.
pub struct DefaultHyperClient;
impl HyperClientBuilder for DefaultHyperClient {
    type Connector = hyper_rustls::HttpsConnector<hyper::client::connect::HttpConnector>;

    fn build_hyper_client(self) -> hyper::Client<Self::Connector> {
        hyper::Client::builder()
            .keep_alive(false)
            .build::<_, hyper::Body>(hyper_rustls::HttpsConnector::new())
    }
}

impl<C> HyperClientBuilder for hyper::Client<C>
where
    C: hyper::client::connect::Connect + 'static,
{
    type Connector = C;

    fn build_hyper_client(self) -> hyper::Client<C> {
        self
    }
}

enum StorageType {
    Memory,
    Disk(PathBuf),
}
<commit_msg>Add a test to ensure that Authenticator is Send+Sync<commit_after>//! Module contianing the core functionality for OAuth2 Authentication.
use crate::authenticator_delegate::{DeviceFlowDelegate, InstalledFlowDelegate};
use crate::device::DeviceFlow;
use crate::error::Error;
use crate::installed::{InstalledFlow, InstalledFlowReturnMethod};
use crate::refresh::RefreshFlow;
use crate::service_account::{ServiceAccountFlow, ServiceAccountFlowOpts, ServiceAccountKey};
use crate::storage::{self, Storage};
use crate::types::{ApplicationSecret, Token};
use private::AuthFlow;

use std::borrow::Cow;
use std::io;
use std::path::PathBuf;
use std::sync::Mutex;

/// Authenticator is responsible for fetching tokens, handling refreshing tokens,
/// and optionally persisting tokens to disk.
pub struct Authenticator<C> {
    hyper_client: hyper::Client<C>,
    storage: Storage,
    auth_flow: AuthFlow,
}

impl<C> Authenticator<C>
where
    C: hyper::client::connect::Connect + 'static,
{
    /// Return the current token for the provided scopes.
    pub async fn token<'a, T>(&'a self, scopes: &'a [T]) -> Result<Token, Error>
    where
        T: AsRef<str>,
    {
        let hashed_scopes = storage::ScopeSet::from(scopes);
        match (self.storage.get(hashed_scopes), self.auth_flow.app_secret()) {
            (Some(t), _) if !t.expired() => {
                // unexpired token found
                Ok(t)
            }
            (
                Some(Token {
                    refresh_token: Some(refresh_token),
                    ..
                }),
                Some(app_secret),
            ) => {
                // token is expired but has a refresh token.
                let token =
                    RefreshFlow::refresh_token(&self.hyper_client, app_secret, &refresh_token)
                        .await?;
                self.storage.set(hashed_scopes, token.clone()).await;
                Ok(token)
            }
            _ => {
                // no token in the cache or the token returned can't be refreshed.
                let t = self.auth_flow.token(&self.hyper_client, scopes).await?;
                self.storage.set(hashed_scopes, t.clone()).await;
                Ok(t)
            }
        }
    }
}

/// Configure an Authenticator using the builder pattern.
pub struct AuthenticatorBuilder<C, F> {
    hyper_client_builder: C,
    storage_type: StorageType,
    auth_flow: F,
}

/// Create an authenticator that uses the installed flow.
/// ```
/// # async fn foo() {
/// # use yup_oauth2::InstalledFlowReturnMethod;
/// # let custom_flow_delegate = yup_oauth2::authenticator_delegate::DefaultInstalledFlowDelegate;
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::InstalledFlowAuthenticator::builder(
///         app_secret,
///         InstalledFlowReturnMethod::HTTPRedirect,
///     )
///     .build()
///     .await
///     .expect("failed to create authenticator");
/// # }
/// ```
pub struct InstalledFlowAuthenticator;
impl InstalledFlowAuthenticator {
    /// Use the builder pattern to create an Authenticator that uses the installed flow.
    pub fn builder(
        app_secret: ApplicationSecret,
        method: InstalledFlowReturnMethod,
    ) -> AuthenticatorBuilder<DefaultHyperClient, InstalledFlow> {
        AuthenticatorBuilder::<DefaultHyperClient, _>::with_auth_flow(InstalledFlow::new(
            app_secret, method,
        ))
    }
}

/// Create an authenticator that uses the device flow.
/// ```
/// # async fn foo() {
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::DeviceFlowAuthenticator::builder(app_secret)
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
pub struct DeviceFlowAuthenticator;
impl DeviceFlowAuthenticator {
    /// Use the builder pattern to create an Authenticator that uses the device flow.
    pub fn builder(
        app_secret: ApplicationSecret,
    ) -> AuthenticatorBuilder<DefaultHyperClient, DeviceFlow> {
        AuthenticatorBuilder::<DefaultHyperClient, _>::with_auth_flow(DeviceFlow::new(app_secret))
    }
}

/// Create an authenticator that uses a service account.
/// ```
/// # async fn foo() {
/// # let service_account_key = yup_oauth2::read_service_account_key("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::ServiceAccountAuthenticator::builder(service_account_key)
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
pub struct ServiceAccountAuthenticator;
impl ServiceAccountAuthenticator {
    /// Use the builder pattern to create an Authenticator that uses a service account.
    pub fn builder(
        service_account_key: ServiceAccountKey,
    ) -> AuthenticatorBuilder<DefaultHyperClient, ServiceAccountFlowOpts> {
        AuthenticatorBuilder::<DefaultHyperClient, _>::with_auth_flow(ServiceAccountFlowOpts {
            key: service_account_key,
            subject: None,
        })
    }
}

/// ## Methods available when building any Authenticator.
/// ```
/// # async fn foo() {
/// # let custom_hyper_client = hyper::Client::new();
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::DeviceFlowAuthenticator::builder(app_secret)
///         .hyper_client(custom_hyper_client)
///         .persist_tokens_to_disk("/tmp/tokenfile.json")
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
impl<C, F> AuthenticatorBuilder<C, F> {
    async fn common_build(
        hyper_client_builder: C,
        storage_type: StorageType,
        auth_flow: AuthFlow,
    ) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        let hyper_client = hyper_client_builder.build_hyper_client();
        let storage = match storage_type {
            StorageType::Memory => Storage::Memory {
                tokens: Mutex::new(storage::JSONTokens::new()),
            },
            StorageType::Disk(path) => Storage::Disk(storage::DiskStorage::new(path).await?),
        };

        Ok(Authenticator {
            hyper_client,
            storage,
            auth_flow,
        })
    }

    fn with_auth_flow(auth_flow: F) -> AuthenticatorBuilder<DefaultHyperClient, F> {
        AuthenticatorBuilder {
            hyper_client_builder: DefaultHyperClient,
            storage_type: StorageType::Memory,
            auth_flow,
        }
    }

    /// Use the provided hyper client.
    pub fn hyper_client<NewC>(
        self,
        hyper_client: hyper::Client<NewC>,
    ) -> AuthenticatorBuilder<hyper::Client<NewC>, F> {
        AuthenticatorBuilder {
            hyper_client_builder: hyper_client,
            storage_type: self.storage_type,
            auth_flow: self.auth_flow,
        }
    }

    /// Persist tokens to disk in the provided filename.
    pub fn persist_tokens_to_disk<P: Into<PathBuf>>(self, path: P) -> AuthenticatorBuilder<C, F> {
        AuthenticatorBuilder {
            storage_type: StorageType::Disk(path.into()),
            ..self
        }
    }
}

/// ## Methods available when building a device flow Authenticator.
/// ```
/// # async fn foo() {
/// # let custom_flow_delegate = yup_oauth2::authenticator_delegate::DefaultDeviceFlowDelegate;
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::DeviceFlowAuthenticator::builder(app_secret)
///         .device_code_url("foo")
///         .flow_delegate(Box::new(custom_flow_delegate))
///         .grant_type("foo")
///         .build()
///         .await
///         .expect("failed to create authenticator");
/// # }
/// ```
impl<C> AuthenticatorBuilder<C, DeviceFlow> {
    /// Use the provided device code url.
    pub fn device_code_url(self, url: impl Into<Cow<'static, str>>) -> Self {
        AuthenticatorBuilder {
            auth_flow: DeviceFlow {
                device_code_url: url.into(),
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Use the provided DeviceFlowDelegate.
    pub fn flow_delegate(self, flow_delegate: Box<dyn DeviceFlowDelegate>) -> Self {
        AuthenticatorBuilder {
            auth_flow: DeviceFlow {
                flow_delegate,
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Use the provided grant type.
    pub fn grant_type(self, grant_type: impl Into<Cow<'static, str>>) -> Self {
        AuthenticatorBuilder {
            auth_flow: DeviceFlow {
                grant_type: grant_type.into(),
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Create the authenticator.
    pub async fn build(self) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        Self::common_build(
            self.hyper_client_builder,
            self.storage_type,
            AuthFlow::DeviceFlow(self.auth_flow),
        )
        .await
    }
}

/// ## Methods available when building an installed flow Authenticator.
/// ```
/// # async fn foo() {
/// # use yup_oauth2::InstalledFlowReturnMethod;
/// # let custom_flow_delegate = yup_oauth2::authenticator_delegate::DefaultInstalledFlowDelegate;
/// # let app_secret = yup_oauth2::read_application_secret("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::InstalledFlowAuthenticator::builder(
///         app_secret,
///         InstalledFlowReturnMethod::HTTPRedirect,
///     )
///     .flow_delegate(Box::new(custom_flow_delegate))
///     .build()
///     .await
///     .expect("failed to create authenticator");
/// # }
/// ```
impl<C> AuthenticatorBuilder<C, InstalledFlow> {
    /// Use the provided InstalledFlowDelegate.
    pub fn flow_delegate(self, flow_delegate: Box<dyn InstalledFlowDelegate>) -> Self {
        AuthenticatorBuilder {
            auth_flow: InstalledFlow {
                flow_delegate,
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Create the authenticator.
    pub async fn build(self) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        Self::common_build(
            self.hyper_client_builder,
            self.storage_type,
            AuthFlow::InstalledFlow(self.auth_flow),
        )
        .await
    }
}

/// ## Methods available when building a service account authenticator.
/// ```
/// # async fn foo() {
/// # let service_account_key = yup_oauth2::read_service_account_key("/tmp/foo").await.unwrap();
///     let authenticator = yup_oauth2::ServiceAccountAuthenticator::builder(
///         service_account_key,
///     )
///     .subject("mysubject")
///     .build()
///     .await
///     .expect("failed to create authenticator");
/// # }
/// ```
impl<C> AuthenticatorBuilder<C, ServiceAccountFlowOpts> {
    /// Use the provided subject.
    pub fn subject(self, subject: impl Into<String>) -> Self {
        AuthenticatorBuilder {
            auth_flow: ServiceAccountFlowOpts {
                subject: Some(subject.into()),
                ..self.auth_flow
            },
            ..self
        }
    }

    /// Create the authenticator.
    pub async fn build(self) -> io::Result<Authenticator<C::Connector>>
    where
        C: HyperClientBuilder,
    {
        let service_account_auth_flow = ServiceAccountFlow::new(self.auth_flow)?;
        Self::common_build(
            self.hyper_client_builder,
            self.storage_type,
            AuthFlow::ServiceAccountFlow(service_account_auth_flow),
        )
        .await
    }
}

mod private {
    use crate::device::DeviceFlow;
    use crate::error::Error;
    use crate::installed::InstalledFlow;
    use crate::service_account::ServiceAccountFlow;
    use crate::types::{ApplicationSecret, Token};

    pub enum AuthFlow {
        DeviceFlow(DeviceFlow),
        InstalledFlow(InstalledFlow),
        ServiceAccountFlow(ServiceAccountFlow),
    }

    impl AuthFlow {
        pub(crate) fn app_secret(&self) -> Option<&ApplicationSecret> {
            match self {
                AuthFlow::DeviceFlow(device_flow) => Some(&device_flow.app_secret),
                AuthFlow::InstalledFlow(installed_flow) => Some(&installed_flow.app_secret),
                AuthFlow::ServiceAccountFlow(_) => None,
            }
        }

        pub(crate) async fn token<'a, C, T>(
            &'a self,
            hyper_client: &'a hyper::Client<C>,
            scopes: &'a [T],
        ) -> Result<Token, Error>
        where
            T: AsRef<str>,
            C: hyper::client::connect::Connect + 'static,
        {
            match self {
                AuthFlow::DeviceFlow(device_flow) => device_flow.token(hyper_client, scopes).await,
                AuthFlow::InstalledFlow(installed_flow) => {
                    installed_flow.token(hyper_client, scopes).await
                }
                AuthFlow::ServiceAccountFlow(service_account_flow) => {
                    service_account_flow.token(hyper_client, scopes).await
                }
            }
        }
    }
}

/// A trait implemented for any hyper::Client as well as the DefaultHyperClient.
pub trait HyperClientBuilder {
    /// The hyper connector that the resulting hyper client will use.
    type Connector: hyper::client::connect::Connect + 'static;

    /// Create a hyper::Client
    fn build_hyper_client(self) -> hyper::Client<Self::Connector>;
}

/// The builder value used when the default hyper client should be used.
pub struct DefaultHyperClient;
impl HyperClientBuilder for DefaultHyperClient {
    type Connector = hyper_rustls::HttpsConnector<hyper::client::connect::HttpConnector>;

    fn build_hyper_client(self) -> hyper::Client<Self::Connector> {
        hyper::Client::builder()
            .keep_alive(false)
            .build::<_, hyper::Body>(hyper_rustls::HttpsConnector::new())
    }
}

impl<C> HyperClientBuilder for hyper::Client<C>
where
    C: hyper::client::connect::Connect + 'static,
{
    type Connector = C;

    fn build_hyper_client(self) -> hyper::Client<C> {
        self
    }
}

enum StorageType {
    Memory,
    Disk(PathBuf),
}

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

    #[test]
    fn ensure_send_sync() {
        fn is_send_sync<T: Send + Sync>() {}
        is_send_sync::<Authenticator<<DefaultHyperClient as HyperClientBuilder>::Connector>>()
    }
}
<|endoftext|>


<commit_before>//!
//! # Scheduler
//!

use std::mem;
use std::sync::*;
use std::thread::*;
use std::sync::mpsc::*;
use std::collections::vec_deque::*;

use num_cpus;

const MIN_THREADS: usize = 8;

lazy_static! {
    static ref SCHEDULER: Arc<Scheduler> = Arc::new(Scheduler::new());
}

///
/// The default maximum number of threads in a scheduler
///
fn initial_max_threads() -> usize {
    MIN_THREADS.max(num_cpus::get()*2)
}

///
/// Trait implemented by things that can be scheduled as a job
///
trait ScheduledJob : Send {
    /// Runs this particular job
    fn run(&mut self);
}

///
/// Basic job is just a FnOnce
///
struct Job<TFn>
where TFn: Send+FnOnce() -> () {
    action: Option<TFn>
}

impl<TFn> Job<TFn>
where TFn: Send+FnOnce() -> () {
    fn new(action: TFn) -> Job<TFn> {
        Job { action: Some(action) }
    }
}

impl<TFn> ScheduledJob for Job<TFn>
where TFn: Send+FnOnce() -> () {
    fn run(&mut self) {
        // Consume the action when it's run
        let mut action = None;
        mem::swap(&mut action, &mut self.action);

        if let Some(action) = action {
            action();
        } else {
            panic!("Cannot schedule an action twice");
        }
    }
}

///
/// The scheduler is used to schedule tasks onto a pool of threads
///
pub struct Scheduler {
    /// The queues that are active in the scheduler
    queues: Arc<Mutex<Vec<Weak<JobQueue>>>>,

    /// Active threads
    threads: Mutex<Vec<SchedulerThread>>,

    /// The maximum number of threads permitted in this scheduler
    max_threads: Mutex<usize>
}

///
/// A job queue provides a list of jobs to perform in order
///
pub struct JobQueue {
    /// Jobs scheduled on this queue
    queue: Mutex<VecDeque<Box<ScheduledJob>>>,

    /// True while this queue is running
    running: Mutex<bool>
}

///
/// A scheduler thread reads from the scheduler queue
///
struct SchedulerThread {
    /// The jobs that this thread should run
    jobs: Sender<Box<ScheduledJob>>,

    /// The thread itself
    thread: JoinHandle<()>,

    /// Flag that indicates that this thread is busy
    busy: Arc<Mutex<bool>>
}

impl SchedulerThread {
    ///
    /// Creates a new scheduler thread
    ///
    fn new() -> SchedulerThread {
        // All the thread does is run jobs from its channel
        let (jobs_in, jobs_out): (Sender<Box<ScheduledJob>>, Receiver<Box<ScheduledJob>>) = channel();
        let thread = spawn(move || {
            while let Ok(mut job) = jobs_out.recv() {
                job.run();
            }
        });

        SchedulerThread {
            jobs:   jobs_in,
            thread: thread,
            busy:   Arc::new(Mutex::new(false))
        }
    }

    ///
    /// Schedules a job to be run on this thread
    ///
    fn run<Job: 'static+ScheduledJob>(&self, job: Job) {
        self.jobs.send(Box::new(job)).unwrap();
    }
}

impl JobQueue {
    ///
    /// Creates a new job queue
    ///
    fn new() -> JobQueue {
        JobQueue {
            queue:      Mutex::new(VecDeque::new()),
            running:    Mutex::new(false)
        }
    }

    ///
    /// Adds a new job to this queue, returns true if the queue is dormant and needs to be started
    ///
    fn queue<TJob: 'static+ScheduledJob>(&self, job: TJob) -> bool {
        let mut queue   = self.queue.lock().unwrap();
        let running     = self.running.lock().unwrap();

        queue.push_back(Box::new(job));

        !*running
    }

    ///
    /// If there are any jobs waiting, dequeues the next one
    ///
    fn dequeue(&self) -> Option<Box<ScheduledJob>> {
        let mut queue = self.queue.lock().unwrap();

        queue.pop_front()
    }

    ///
    /// Runs jobs on this queue until there are none left, marking the job as inactive when done
    ///
    fn drain(&self) {
        let mut done = false;

        while !done {
            // Run jobs until the queue is drained
            while let Some(mut job) = self.dequeue() {
                job.run();
            }

            // Mark as not running any more
            {
                let queue       = self.queue.lock().unwrap();
                let mut running = self.running.lock().unwrap();

                // If the queue is empty at the point where we obtain the 'running' lock, we can deactivate ourselves
                if queue.len() == 0 {
                    *running    = false;
                    done        = true;
                }
            }
        }
    }
}

impl Scheduler {
    ///
    /// Creates a new scheduler
    ///
    /// (There's only actually one scheduler)
    ///
    fn new() -> Scheduler {
        let result = Scheduler {
            queues:         Arc::new(Mutex::new(vec![])),
            threads:        Mutex::new(vec![]),
            max_threads:    Mutex::new(initial_max_threads())
        };

        result
    }

    ///
    /// Finds the next queue that should be run. If this returns successfully, the queue will
    /// be marked as running.
    ///
    fn next_to_run(queues: &Arc<Mutex<Vec<Weak<JobQueue>>>>) -> Option<Arc<JobQueue>> {
        // Search the queues...
        let queues = queues.lock().unwrap();

        // Find a queue where is_running is false
        for q in queues.iter() {
            // Queue is a weak ref, check that it's still in use
            if let Some(q) = q.upgrade() {
                // If the queue is not running, then mark it as running and return it
                let mut is_running = q.running.lock().unwrap();
                if !*is_running {
                    // Clone here is necessary because the is_running update is borrowing q
                    *is_running = true;
                    return Some(q.clone());
                }
            }
        }

        None
    }

    ///
    /// Attempts to schedule a task on a dormant thread
    ///
    fn schedule_dormant<NextJob, RunJob, JobData>(&self, next_job: NextJob, job: RunJob) -> bool
    where RunJob: 'static+Send+Fn(JobData) -> (), NextJob: 'static+Send+Fn() -> Option<JobData> {
        let threads = self.threads.lock().unwrap();

        // Find the first thread that is not marked as busy and schedule this task on it
        for thread in threads.iter() {
            let mut busy = thread.busy.lock().unwrap();

            if !*busy {
                // Clone the busy mutex so we can return this thread to readiness
                let also_busy = thread.busy.clone();

                // This thread is busy
                *busy = true;
                thread.run(Job::new(move || {
                    let mut done = false;

                    while !done {
                        // Obtain the next job. The thread is not busy once there are no longer any jobs
                        // We hold the mutex while this is going on to avoid a race condition when a thread is going dormant
                        let job_data = {
                            let mut busy = also_busy.lock().unwrap();
                            let job_data = next_job();

                            // If there's no next job, then this thread is no longer busy
                            if job_data.is_none() {
                                *busy = false;
                            }

                            job_data
                        };

                        // Run the job if there is one, stop the thread if there is not
                        if let Some(job_data) = job_data {
                            job(job_data);
                        } else {
                            done = true;
                        }
                    }
                }));

                return true;
            }
        }

        // No dormant threads were found
        false
    }

    ///
    /// If we're running fewer than the maximum number of threads, try to spawn a new one
    ///
    fn spawn_thread_if_less_than_maximum(&self) -> bool {
        let max_threads = { *self.max_threads.lock().unwrap() };
        let mut threads = self.threads.lock().unwrap();

        if threads.len() < max_threads {
            // Create a new thread
            let new_thread = SchedulerThread::new();
            threads.push(new_thread);

            true
        } else {
            // Can't spawn a new thread
            false
        }
    }

    ///
    /// Wakes a thread to run a dormant queue. Returns true if a thread was woken up
    ///
    fn schedule_thread(&self) -> bool {
        // Find a dormant thread and activate it
        let queues = self.queues.clone();

        // Schedule work on this dormant thread
        if !self.schedule_dormant(move || Self::next_to_run(&queues), move |work| work.drain()) {
            // Try to create a new thread
            if self.spawn_thread_if_less_than_maximum() {
                // Try harder to schedule this task if a thread was created
                self.schedule_thread()
            } else {
                // Couldn't schedule on an existing thread or create a new one
                false
            }
        } else {
            // Successfully scheduled
            true
        }
    }

    ///
    /// Spawns a thread in this scheduler
    ///
    pub fn spawn_thread(&self) {
        let new_thread = SchedulerThread::new();
        self.threads.lock().unwrap().push(new_thread);
    }

    ///
    /// Creates a new job queue for this scheduler
    ///
    pub fn create_job_queue(&self) -> Arc<JobQueue> {
        // Create the  new queue
        let new_queue = Arc::new(JobQueue::new());

        let mut queues = self.queues.lock().unwrap();

        // Replace an existing weak queue if possible
        let weak_pos = queues.iter().position(|q| q.upgrade().is_none());

        if let Some(weak_pos) = weak_pos {
            // Replace a queue that was present but is no longer in use
            queues[weak_pos] = Arc::downgrade(&new_queue);
        } else {
            // Add to the queues managed by this object (as a weak reference)
            queues.push(Arc::downgrade(&new_queue));
        }

        new_queue
    }

    ///
    /// Schedules a job on this scheduler
    ///
    pub fn async<TFn: 'static+Send+FnOnce() -> ()>(&self, queue: &Arc<JobQueue>, job: TFn) {
        if queue.queue(Job::new(job)) {
            // A true result indicates that the job was scheduled but the queue is not running. Try to schedule a thread if this occurs.
            self.schedule_thread();
        }
    }
}

///
/// Retrieves the global scheduler
///
pub fn scheduler<'a>() -> &'a Scheduler {
    &SCHEDULER
}

///
/// Creates a scheduler queue
///
pub fn queue() -> Arc<JobQueue> {
    scheduler().create_job_queue()
}

///
/// Creates a scheduler queue
///
pub fn async<TFn: 'static+Send+FnOnce() -> ()>(queue: &Arc<JobQueue>, job: TFn) {
    scheduler().async(queue, job)
}

#[cfg(test)]
mod test {
    use super::*;
    use std::time::*;

    fn timeout<TFn: 'static+Send+FnOnce() -> ()>(action: TFn, millis: u64) {
        let (tx, rx)    = channel();
        let (tx1, tx2)  = (tx.clone(), tx.clone());

        spawn(move || {
            action();
            tx1.send(true).ok();
        });

        spawn(move || {
            sleep(Duration::from_millis(millis));
            tx2.send(false).ok();
        });

        if rx.recv().unwrap() == false {
            panic!("Timeout");
        }
    }

    #[test]
    fn can_schedule_async() {
        timeout(|| {
            let (tx, rx)    = channel();
            let queue       = queue();

            async(&queue, move || {
                tx.send(42).unwrap();
            });

            assert!(rx.recv().unwrap() == 42);
        }, 100);
    }

    #[test]
    fn can_schedule_after_queue_released() {
        timeout(|| {
            {
                let (tx, rx)    = channel();
                let queue1      = queue();

                async(&queue1, move || {
                    tx.send(42).unwrap();
                });

                assert!(rx.recv().unwrap() == 42);
            }

            {
                let (tx, rx)    = channel();
                let queue2      = queue();

                async(&queue2, move || {
                    tx.send(43).unwrap();
                });

                assert!(rx.recv().unwrap() == 43);
            }
        }, 100);
    }

    #[test]
    fn will_schedule_in_order() {
        timeout(|| {
            let (tx, rx)    = channel();
            let queue       = queue();

            let (tx1, tx2)  = (tx.clone(), tx.clone());

            async(&queue, move || {
                sleep(Duration::from_millis(100));
                tx1.send(1).unwrap();
            });
            async(&queue, move || {
                tx2.send(2).unwrap();
            });

            assert!(rx.recv().unwrap() == 1);
            assert!(rx.recv().unwrap() == 2);
        }, 500);
    }

    #[test]
    fn will_schedule_separate_queues_in_parallel() {
        timeout(|| {
            let (tx, rx)        = channel();
            let queue1          = queue();
            let queue2          = queue();
            let queue2_has_run  = Arc::new(Mutex::new(false));

            let queue1_check = queue2_has_run.clone();

            async(&queue1, move || {
                // The other task needs to start within 100ms for this to work
                sleep(Duration::from_millis(100));
                tx.send(*queue1_check.lock().unwrap()).unwrap();
            });
            async(&queue2, move || {
                *queue2_has_run.lock().unwrap() = true;
            });

            assert!(rx.recv().unwrap() == true);
        }, 500);
    }
}
<commit_msg>Add some functions to manage the maximum number of threads the scheduler is running<commit_after>//!
//! # Scheduler
//!

use std::mem;
use std::sync::*;
use std::thread::*;
use std::sync::mpsc::*;
use std::collections::vec_deque::*;

use num_cpus;

const MIN_THREADS: usize = 8;

lazy_static! {
    static ref SCHEDULER: Arc<Scheduler> = Arc::new(Scheduler::new());
}

///
/// The default maximum number of threads in a scheduler
///
fn initial_max_threads() -> usize {
    MIN_THREADS.max(num_cpus::get()*2)
}

///
/// Trait implemented by things that can be scheduled as a job
///
trait ScheduledJob : Send {
    /// Runs this particular job
    fn run(&mut self);
}

///
/// Basic job is just a FnOnce
///
struct Job<TFn>
where TFn: Send+FnOnce() -> () {
    action: Option<TFn>
}

impl<TFn> Job<TFn>
where TFn: Send+FnOnce() -> () {
    fn new(action: TFn) -> Job<TFn> {
        Job { action: Some(action) }
    }
}

impl<TFn> ScheduledJob for Job<TFn>
where TFn: Send+FnOnce() -> () {
    fn run(&mut self) {
        // Consume the action when it's run
        let mut action = None;
        mem::swap(&mut action, &mut self.action);

        if let Some(action) = action {
            action();
        } else {
            panic!("Cannot schedule an action twice");
        }
    }
}

///
/// The scheduler is used to schedule tasks onto a pool of threads
///
pub struct Scheduler {
    /// The queues that are active in the scheduler
    queues: Arc<Mutex<Vec<Weak<JobQueue>>>>,

    /// Active threads
    threads: Mutex<Vec<SchedulerThread>>,

    /// The maximum number of threads permitted in this scheduler
    max_threads: Mutex<usize>
}

///
/// A job queue provides a list of jobs to perform in order
///
pub struct JobQueue {
    /// Jobs scheduled on this queue
    queue: Mutex<VecDeque<Box<ScheduledJob>>>,

    /// True while this queue is running
    running: Mutex<bool>
}

///
/// A scheduler thread reads from the scheduler queue
///
struct SchedulerThread {
    /// The jobs that this thread should run
    jobs: Sender<Box<ScheduledJob>>,

    /// The thread itself
    thread: JoinHandle<()>,

    /// Flag that indicates that this thread is busy
    busy: Arc<Mutex<bool>>
}

impl SchedulerThread {
    ///
    /// Creates a new scheduler thread
    ///
    fn new() -> SchedulerThread {
        // All the thread does is run jobs from its channel
        let (jobs_in, jobs_out): (Sender<Box<ScheduledJob>>, Receiver<Box<ScheduledJob>>) = channel();
        let thread = spawn(move || {
            while let Ok(mut job) = jobs_out.recv() {
                job.run();
            }
        });

        SchedulerThread {
            jobs:   jobs_in,
            thread: thread,
            busy:   Arc::new(Mutex::new(false))
        }
    }

    ///
    /// Schedules a job to be run on this thread
    ///
    fn run<Job: 'static+ScheduledJob>(&self, job: Job) {
        self.jobs.send(Box::new(job)).unwrap();
    }
}

impl JobQueue {
    ///
    /// Creates a new job queue
    ///
    fn new() -> JobQueue {
        JobQueue {
            queue:      Mutex::new(VecDeque::new()),
            running:    Mutex::new(false)
        }
    }

    ///
    /// Adds a new job to this queue, returns true if the queue is dormant and needs to be started
    ///
    fn queue<TJob: 'static+ScheduledJob>(&self, job: TJob) -> bool {
        let mut queue   = self.queue.lock().unwrap();
        let running     = self.running.lock().unwrap();

        queue.push_back(Box::new(job));

        !*running
    }

    ///
    /// If there are any jobs waiting, dequeues the next one
    ///
    fn dequeue(&self) -> Option<Box<ScheduledJob>> {
        let mut queue = self.queue.lock().unwrap();

        queue.pop_front()
    }

    ///
    /// Runs jobs on this queue until there are none left, marking the job as inactive when done
    ///
    fn drain(&self) {
        let mut done = false;

        while !done {
            // Run jobs until the queue is drained
            while let Some(mut job) = self.dequeue() {
                job.run();
            }

            // Mark as not running any more
            {
                let queue       = self.queue.lock().unwrap();
                let mut running = self.running.lock().unwrap();

                // If the queue is empty at the point where we obtain the 'running' lock, we can deactivate ourselves
                if queue.len() == 0 {
                    *running    = false;
                    done        = true;
                }
            }
        }
    }
}

impl Scheduler {
    ///
    /// Creates a new scheduler
    ///
    /// (There's only actually one scheduler)
    ///
    fn new() -> Scheduler {
        let result = Scheduler {
            queues:         Arc::new(Mutex::new(vec![])),
            threads:        Mutex::new(vec![]),
            max_threads:    Mutex::new(initial_max_threads())
        };

        result
    }

    ///
    /// Changes the maximum number of threads this scheduler can spawn (existing threads
    /// are not despawned by this method)
    ///
    pub fn set_max_threads(&self, max_threads: usize) {
        // Update the maximum number of threads we can spawn
        { *self.max_threads.lock().unwrap() = max_threads };
    }

    ///
    /// Despawns threads if we're running more than the maximum number
    ///
    /// Must not be called from a scheduler thread (as it waits for the threads to despawn)
    ///
    pub fn despawn_threads_if_overloaded(&self) {
        let max_threads = { *self.max_threads.lock().unwrap() };
        let to_despawn  = {
            // Transfer the threads from the threads vector to our _to_despawn variable
            // This is then dropped outside the mutex (so we don't block if one of the threads doesn't stop)
            let mut to_despawn  = vec![];
            let mut threads     = self.threads.lock().unwrap();

            while threads.len() > max_threads {
                to_despawn.push(threads.pop().unwrap().thread);
            }

            to_despawn
        };

        // Wait for the threads to despawn
        to_despawn.into_iter().for_each(|join_handle| { join_handle.join().ok(); });
    }

    ///
    /// Finds the next queue that should be run. If this returns successfully, the queue will
    /// be marked as running.
    ///
    fn next_to_run(queues: &Arc<Mutex<Vec<Weak<JobQueue>>>>) -> Option<Arc<JobQueue>> {
        // Search the queues...
        let queues = queues.lock().unwrap();

        // Find a queue where is_running is false
        for q in queues.iter() {
            // Queue is a weak ref, check that it's still in use
            if let Some(q) = q.upgrade() {
                // If the queue is not running, then mark it as running and return it
                let mut is_running = q.running.lock().unwrap();
                if !*is_running {
                    // Clone here is necessary because the is_running update is borrowing q
                    *is_running = true;
                    return Some(q.clone());
                }
            }
        }

        None
    }

    ///
    /// Attempts to schedule a task on a dormant thread
    ///
    fn schedule_dormant<NextJob, RunJob, JobData>(&self, next_job: NextJob, job: RunJob) -> bool
    where RunJob: 'static+Send+Fn(JobData) -> (), NextJob: 'static+Send+Fn() -> Option<JobData> {
        let threads = self.threads.lock().unwrap();

        // Find the first thread that is not marked as busy and schedule this task on it
        for thread in threads.iter() {
            let mut busy = thread.busy.lock().unwrap();

            if !*busy {
                // Clone the busy mutex so we can return this thread to readiness
                let also_busy = thread.busy.clone();

                // This thread is busy
                *busy = true;
                thread.run(Job::new(move || {
                    let mut done = false;

                    while !done {
                        // Obtain the next job. The thread is not busy once there are no longer any jobs
                        // We hold the mutex while this is going on to avoid a race condition when a thread is going dormant
                        let job_data = {
                            let mut busy = also_busy.lock().unwrap();
                            let job_data = next_job();

                            // If there's no next job, then this thread is no longer busy
                            if job_data.is_none() {
                                *busy = false;
                            }

                            job_data
                        };

                        // Run the job if there is one, stop the thread if there is not
                        if let Some(job_data) = job_data {
                            job(job_data);
                        } else {
                            done = true;
                        }
                    }
                }));

                return true;
            }
        }

        // No dormant threads were found
        false
    }

    ///
    /// If we're running fewer than the maximum number of threads, try to spawn a new one
    ///
    fn spawn_thread_if_less_than_maximum(&self) -> bool {
        let max_threads = { *self.max_threads.lock().unwrap() };
        let mut threads = self.threads.lock().unwrap();

        if threads.len() < max_threads {
            // Create a new thread
            let new_thread = SchedulerThread::new();
            threads.push(new_thread);

            true
        } else {
            // Can't spawn a new thread
            false
        }
    }

    ///
    /// Wakes a thread to run a dormant queue. Returns true if a thread was woken up
    ///
    fn schedule_thread(&self) -> bool {
        // Find a dormant thread and activate it
        let queues = self.queues.clone();

        // Schedule work on this dormant thread
        if !self.schedule_dormant(move || Self::next_to_run(&queues), move |work| work.drain()) {
            // Try to create a new thread
            if self.spawn_thread_if_less_than_maximum() {
                // Try harder to schedule this task if a thread was created
                self.schedule_thread()
            } else {
                // Couldn't schedule on an existing thread or create a new one
                false
            }
        } else {
            // Successfully scheduled
            true
        }
    }

    ///
    /// Spawns a thread in this scheduler
    ///
    pub fn spawn_thread(&self) {
        let new_thread = SchedulerThread::new();
        self.threads.lock().unwrap().push(new_thread);
    }

    ///
    /// Creates a new job queue for this scheduler
    ///
    pub fn create_job_queue(&self) -> Arc<JobQueue> {
        // Create the  new queue
        let new_queue = Arc::new(JobQueue::new());

        let mut queues = self.queues.lock().unwrap();

        // Replace an existing weak queue if possible
        let weak_pos = queues.iter().position(|q| q.upgrade().is_none());

        if let Some(weak_pos) = weak_pos {
            // Replace a queue that was present but is no longer in use
            queues[weak_pos] = Arc::downgrade(&new_queue);
        } else {
            // Add to the queues managed by this object (as a weak reference)
            queues.push(Arc::downgrade(&new_queue));
        }

        new_queue
    }

    ///
    /// Schedules a job on this scheduler
    ///
    pub fn async<TFn: 'static+Send+FnOnce() -> ()>(&self, queue: &Arc<JobQueue>, job: TFn) {
        if queue.queue(Job::new(job)) {
            // A true result indicates that the job was scheduled but the queue is not running. Try to schedule a thread if this occurs.
            self.schedule_thread();
        }
    }
}

///
/// Retrieves the global scheduler
///
pub fn scheduler<'a>() -> &'a Scheduler {
    &SCHEDULER
}

///
/// Creates a scheduler queue
///
pub fn queue() -> Arc<JobQueue> {
    scheduler().create_job_queue()
}

///
/// Creates a scheduler queue
///
pub fn async<TFn: 'static+Send+FnOnce() -> ()>(queue: &Arc<JobQueue>, job: TFn) {
    scheduler().async(queue, job)
}

#[cfg(test)]
mod test {
    use super::*;
    use std::time::*;

    fn timeout<TFn: 'static+Send+FnOnce() -> ()>(action: TFn, millis: u64) {
        let (tx, rx)    = channel();
        let (tx1, tx2)  = (tx.clone(), tx.clone());

        spawn(move || {
            action();
            tx1.send(true).ok();
        });

        spawn(move || {
            sleep(Duration::from_millis(millis));
            tx2.send(false).ok();
        });

        if rx.recv().unwrap() == false {
            panic!("Timeout");
        }
    }

    #[test]
    fn can_schedule_async() {
        timeout(|| {
            let (tx, rx)    = channel();
            let queue       = queue();

            async(&queue, move || {
                tx.send(42).unwrap();
            });

            assert!(rx.recv().unwrap() == 42);
        }, 100);
    }

    #[test]
    fn can_schedule_after_queue_released() {
        timeout(|| {
            {
                let (tx, rx)    = channel();
                let queue1      = queue();

                async(&queue1, move || {
                    tx.send(42).unwrap();
                });

                assert!(rx.recv().unwrap() == 42);
            }

            {
                let (tx, rx)    = channel();
                let queue2      = queue();

                async(&queue2, move || {
                    tx.send(43).unwrap();
                });

                assert!(rx.recv().unwrap() == 43);
            }
        }, 100);
    }

    #[test]
    fn will_schedule_in_order() {
        timeout(|| {
            let (tx, rx)    = channel();
            let queue       = queue();

            let (tx1, tx2)  = (tx.clone(), tx.clone());

            async(&queue, move || {
                sleep(Duration::from_millis(100));
                tx1.send(1).unwrap();
            });
            async(&queue, move || {
                tx2.send(2).unwrap();
            });

            assert!(rx.recv().unwrap() == 1);
            assert!(rx.recv().unwrap() == 2);
        }, 500);
    }

    #[test]
    fn will_schedule_separate_queues_in_parallel() {
        timeout(|| {
            let (tx, rx)        = channel();
            let queue1          = queue();
            let queue2          = queue();
            let queue2_has_run  = Arc::new(Mutex::new(false));

            let queue1_check = queue2_has_run.clone();

            async(&queue1, move || {
                // The other task needs to start within 100ms for this to work
                sleep(Duration::from_millis(100));
                tx.send(*queue1_check.lock().unwrap()).unwrap();
            });
            async(&queue2, move || {
                *queue2_has_run.lock().unwrap() = true;
            });

            assert!(rx.recv().unwrap() == true);
        }, 500);
    }

    #[test]
    fn will_despawn_extra_threads() {
        // As we join with the threads, we'll timeout if any of the spawned threads fail to end
        timeout(|| {
            let scheduler = scheduler();

            // Maximum of 10 threads, but we'll spawn 20
            scheduler.set_max_threads(10);
            for _ in 1..20 {
                scheduler.spawn_thread();
            }

            scheduler.despawn_threads_if_overloaded();
        }, 500);
    }
}
<|endoftext|>


<commit_before>use futures::Future;
use latch::{Latch, CountLatch};
use log::Event::*;
use job::HeapJob;
use std::any::Any;
use std::marker::PhantomData;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicPtr, Ordering};
use registry::{in_worker, WorkerThread};
use unwind;

mod future;
use self::future::RayonFuture;

#[cfg(test)]
mod test;

pub struct Scope<'scope> {
    /// thread where `scope()` was executed (note that individual jobs
    /// may be executing on different worker threads, though they
    /// should always be within the same pool of threads)
    owner_thread: *const WorkerThread,

    /// if some job panicked, the error is stored here; it will be
    /// propagated to the one who created the scope
    panic: AtomicPtr<Box<Any + Send + 'static>>,

    /// latch to set when the counter drops to zero (and hence this scope is complete)
    job_completed_latch: CountLatch,

    /// you can think of a scope as containing a list of closures to
    /// execute, all of which outlive `'scope`
    marker: PhantomData<Box<FnOnce(&Scope<'scope>) + 'scope>>,
}

/// Create a "fork-join" scope `s` and invokes the closure with a
/// reference to `s`. This closure can then spawn asynchronous tasks
/// into `s`. Those tasks may run asynchronously with respect to the
/// closure; they may themselves spawn additional tasks into `s`. When
/// the closure returns, it will block until all tasks that have been
/// spawned into `s` complete.
///
/// `scope()` is a more flexible building block compared to `join()`,
/// since a loop can be used to spawn any number of tasks without
/// recursing. However, that flexibility comes at a performance price:
/// tasks spawned using `scope()` must be allocated onto the heap,
/// whereas `join()` can make exclusive use of the stack. **Prefer
/// `join()` (or, even better, parallel iterators) where possible.**
///
/// ### Example
///
/// The Rayon `join()` function launches two closures and waits for them
/// to stop. One could implement `join()` using a scope like so, although
/// it would be less efficient than the real implementation:
///
/// ```rust
/// pub fn join<A,B,RA,RB>(oper_a: A, oper_b: B) -> (RA, RB)
///     where A: FnOnce() -> RA + Send,
///           B: FnOnce() -> RB + Send,
///           RA: Send,
///           RB: Send,
/// {
///     let mut result_a: Option<RA> = None;
///     let mut result_b: Option<RB> = None;
///     rayon::scope(|s| {
///         s.spawn(|_| result_a = Some(oper_a()));
///         s.spawn(|_| result_b = Some(oper_b()));
///     });
///     (result_a.unwrap(), result_b.unwrap())
/// }
/// ```
///
/// ### A note on threading
///
/// The closure given to `scope()` executes in the Rayon thread-pool,
/// as do those given to `spawn()`. This means that you can't access
/// thread-local variables (well, you can, but they may have
/// unexpected values).
///
/// ### Task execution
///
/// Task execution potentially starts as soon as `spawn()` is called.
/// The task will end sometime before `scope()` returns. Note that the
/// *closure* given to scope may return much earlier. In general
/// the lifetime of a scope created like `scope(body) goes something like this:
///
/// - Scope begins when `scope(body)` is called
/// - Scope body `body()` is invoked
///     - Scope tasks may be spawned
/// - Scope body returns
/// - Scope tasks execute, possibly spawning more tasks
/// - Once all tasks are done, scope ends and `scope()` returns
///
/// To see how and when tasks are joined, consider this example:
///
/// ```rust
/// // point start
/// rayon::scope(|s| {
///     s.spawn(|s| { // task s.1
///         s.spawn(|s| { // task s.1.1
///             rayon::scope(|t| {
///                 t.spawn(|_| ()); // task t.1
///                 t.spawn(|_| ()); // task t.2
///             });
///         });
///     });
///     s.spawn(|s| { // task 2
///     });
///     // point mid
/// });
/// // point end
/// ```
///
/// The various tasks that are run will execute roughly like so:
///
/// ```notrust
/// | (start)
/// |
/// | (scope `s` created)
/// +--------------------+ (task s.1)
/// +-------+ (task s.2) |
/// |       |            +---+ (task s.1.1)
/// |       |            |   |
/// |       |            |   | (scope `t` created)
/// |       |            |   +----------------+ (task t.1)
/// |       |            |   +---+ (task t.2) |
/// | (mid) |            |   |   |            |
/// :       |            |   + <-+------------+ (scope `t` ends)
/// :       |            |   |
/// |<------+------------+---+ (scope `s` ends)
/// |
/// | (end)
/// ```
///
/// The point here is that everything spawned into scope `s` will
/// terminate (at latest) at the same point -- right before the
/// original call to `rayon::scope` returns. This includes new
/// subtasks created by other subtasks (e.g., task `s.1.1`). If a new
/// scope is created (such as `t`), the things spawned into that scope
/// will be joined before that scope returns, which in turn occurs
/// before the creating task (task `s.1.1` in this case) finishes.
///
/// ### Accessing stack data
///
/// In general, spawned tasks may access stack data in place that
/// outlives the scope itself. Other data must be fully owned by the
/// spawned task.
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     s.spawn(|_| {
///         // We can access `ok` because outlives the scope `s`.
///         println!("ok: {:?}", ok);
///
///         // If we just try to use `bad` here, the closure will borrow `bad`
///         // (because we are just printing it out, and that only requires a
///         // borrow), which will result in a compilation error. Read on
///         // for options.
///         // println!("bad: {:?}", bad);
///    });
/// });
/// ```
///
/// As the comments example above suggest, to reference `bad` we must
/// take ownership of it. One way to do this is to detach the closure
/// from the surrounding stack frame, using the `move` keyword. This
/// will cause it to take ownership of *all* the variables it touches,
/// in this case including both `ok` *and* `bad`:
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     s.spawn(move |_| {
///         println!("ok: {:?}", ok);
///         println!("bad: {:?}", bad);
///     });
///
///     // That closure is fine, but now we can't use `ok` anywhere else,
///     // since it is owend by the previous task:
///     // s.spawn(|_| println!("ok: {:?}", ok));
/// });
/// ```
///
/// While this works, it could be a problem if we want to use `ok` elsewhere.
/// There are two choices. We can keep the closure as a `move` closure, but
/// instead of referencing the variable `ok`, we create a shadowed variable that
/// is a borrow of `ok` and capture *that*:
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     let ok: &Vec<i32> = &ok; // shadow the original `ok`
///     s.spawn(move |_| {
///         println!("ok: {:?}", ok); // captures the shadowed version
///         println!("bad: {:?}", bad);
///     });
///
///     // Now we too can use the shadowed `ok`, since `&Vec<i32>` references
///     // can be shared freely. Note that we need a `move` closure here though,
///     // because otherwise we'd be trying to borrow the shadowed `ok`,
///     // and that doesn't outlive `scope`.
///     s.spawn(move |_| println!("ok: {:?}", ok));
/// });
/// ```
///
/// Another option is not to use the `move` keyword but instead to take ownership
/// of individual variables:
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     s.spawn(|_| {
///         // Transfer ownership of `bad` into a local variable (also named `bad`).
///         // This will force the closure to take ownership of `bad` from the environment.
///         let bad = bad;
///         println!("ok: {:?}", ok); // `ok` is only borrowed.
///         println!("bad: {:?}", bad); // refers to our local variable, above.
///     });
///
///     s.spawn(|_| println!("ok: {:?}", ok)); // we too can borrow `ok`
/// });
/// ```
///
/// ### Panics
///
/// If a panic occurs, either in the closure given to `scope()` or in
/// any of the spawned jobs, that panic will be propagated and the
/// call to `scope()` will panic. If multiple panics occurs, it is
/// non-deterministic which of their panic values will propagate.
/// Regardless, once a task is spawned using `scope.spawn()`, it will
/// execute, even if the spawning task should later panic. `scope()`
/// returns once all spawned jobs have completed, and any panics are
/// propagated at that point.
pub fn scope<'scope, OP, R>(op: OP) -> R
    where OP: for<'s> FnOnce(&'s Scope<'scope>) -> R + 'scope + Send, R: Send,
{
    in_worker(|owner_thread| {
        unsafe {
            let scope: Scope<'scope> = Scope {
                owner_thread: owner_thread as *const WorkerThread as *mut WorkerThread,
                panic: AtomicPtr::new(ptr::null_mut()),
                job_completed_latch: CountLatch::new(),
                marker: PhantomData,
            };
            let result = scope.execute_job_closure(op);
            scope.steal_till_jobs_complete();
            result.unwrap() // only None if `op` panicked, and that would have been propagated
        }
    })
}

impl<'scope> Scope<'scope> {
    /// Spawns a job into the fork-join scope `self`. This job will
    /// execute sometime before the fork-join scope completes.  The
    /// job is specified as a closure, and this closure receives its
    /// own reference to `self` as argument. This can be used to
    /// inject new jobs into `self`.
    pub fn spawn<BODY>(&self, body: BODY)
        where BODY: FnOnce(&Scope<'scope>) + 'scope
    {
        unsafe {
            self.job_completed_latch.increment();
            let job_ref = Box::new(HeapJob::new(move || self.execute_job(body)))
                .as_job_ref();
            let worker_thread = WorkerThread::current();

            // the `Scope` is not send or sync, and we only give out
            // pointers to it from within a worker thread
            debug_assert!(!WorkerThread::current().is_null());

            let worker_thread = &*worker_thread;
            worker_thread.push(job_ref);
        }
    }

    pub fn spawn_future<F>(&self, future: F) -> RayonFuture<F::Item, F::Error>
        where F: Future + Send + 'scope
    {
        self.job_completed_latch.increment();

        // we assert that we have the future `F` type will remain
        // valid until `job_completed_latch` is fully set
        let future = unsafe { future::new_rayon_future(future, &self.job_completed_latch) };

        future
    }

    /// Executes `func` as a job, either aborting or executing as
    /// appropriate.
    ///
    /// Unsafe because it must be executed on a worker thread.
    unsafe fn execute_job<FUNC>(&self, func: FUNC)
        where FUNC: FnOnce(&Scope<'scope>) + 'scope
    {
        let _: Option<()> = self.execute_job_closure(func);
    }

    /// Executes `func` as a job in scope. Adjusts the "job completed"
    /// counters and also catches any panic and stores it into
    /// `scope`.
    ///
    /// Unsafe because this must be executed on a worker thread.
    unsafe fn execute_job_closure<FUNC, R>(&self, func: FUNC) -> Option<R>
        where FUNC: FnOnce(&Scope<'scope>) -> R + 'scope
    {
        match unwind::halt_unwinding(move || func(self)) {
            Ok(r) => { self.job_completed_ok(); Some(r) }
            Err(err) => { self.job_panicked(err); None }
        }
    }

    unsafe fn job_panicked(&self, err: Box<Any + Send + 'static>) {
        // capture the first error we see, free the rest
        let nil = ptr::null_mut();
        let mut err = Box::new(err); // box up the fat ptr
        if self.panic.compare_exchange(nil, &mut *err, Ordering::Release, Ordering::Relaxed).is_ok() {
            log!(JobPanickedErrorStored { owner_thread: (*self.owner_thread).index() });
            mem::forget(err); // ownership now transferred into self.panic
        } else {
            log!(JobPanickedErrorNotStored { owner_thread: (*self.owner_thread).index() });
        }

        self.job_completed_latch.set();
    }

    unsafe fn job_completed_ok(&self) {
        log!(JobCompletedOk { owner_thread: (*self.owner_thread).index() });
        self.job_completed_latch.set();
    }

    unsafe fn steal_till_jobs_complete(&self) {
        // wait for job counter to reach 0:
        (*self.owner_thread).wait_until(&self.job_completed_latch);

        // propagate panic, if any occurred; at this point, all
        // outstanding jobs have completed, so we can use a relaxed
        // ordering:
        let panic = self.panic.swap(ptr::null_mut(), Ordering::Relaxed);
        if !panic.is_null() {
            log!(ScopeCompletePanicked { owner_thread: (*self.owner_thread).index() });
            let value: Box<Box<Any + Send + 'static>> = mem::transmute(panic);
            unwind::resume_unwinding(*value);
        } else {
            log!(ScopeCompleteNoPanic { owner_thread: (*self.owner_thread).index() });
        }
    }
}
<commit_msg>make some future-related stuff dependent on the unstable feature<commit_after>#[cfg(feature = "unstable")]
use futures::Future;
use latch::{Latch, CountLatch};
use log::Event::*;
use job::HeapJob;
use std::any::Any;
use std::marker::PhantomData;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicPtr, Ordering};
use registry::{in_worker, WorkerThread};
use unwind;

#[cfg(feature = "unstable")]
mod future;
#[cfg(feature = "unstable")]
use self::future::RayonFuture;

#[cfg(test)]
mod test;

pub struct Scope<'scope> {
    /// thread where `scope()` was executed (note that individual jobs
    /// may be executing on different worker threads, though they
    /// should always be within the same pool of threads)
    owner_thread: *const WorkerThread,

    /// if some job panicked, the error is stored here; it will be
    /// propagated to the one who created the scope
    panic: AtomicPtr<Box<Any + Send + 'static>>,

    /// latch to set when the counter drops to zero (and hence this scope is complete)
    job_completed_latch: CountLatch,

    /// you can think of a scope as containing a list of closures to
    /// execute, all of which outlive `'scope`
    marker: PhantomData<Box<FnOnce(&Scope<'scope>) + 'scope>>,
}

/// Create a "fork-join" scope `s` and invokes the closure with a
/// reference to `s`. This closure can then spawn asynchronous tasks
/// into `s`. Those tasks may run asynchronously with respect to the
/// closure; they may themselves spawn additional tasks into `s`. When
/// the closure returns, it will block until all tasks that have been
/// spawned into `s` complete.
///
/// `scope()` is a more flexible building block compared to `join()`,
/// since a loop can be used to spawn any number of tasks without
/// recursing. However, that flexibility comes at a performance price:
/// tasks spawned using `scope()` must be allocated onto the heap,
/// whereas `join()` can make exclusive use of the stack. **Prefer
/// `join()` (or, even better, parallel iterators) where possible.**
///
/// ### Example
///
/// The Rayon `join()` function launches two closures and waits for them
/// to stop. One could implement `join()` using a scope like so, although
/// it would be less efficient than the real implementation:
///
/// ```rust
/// pub fn join<A,B,RA,RB>(oper_a: A, oper_b: B) -> (RA, RB)
///     where A: FnOnce() -> RA + Send,
///           B: FnOnce() -> RB + Send,
///           RA: Send,
///           RB: Send,
/// {
///     let mut result_a: Option<RA> = None;
///     let mut result_b: Option<RB> = None;
///     rayon::scope(|s| {
///         s.spawn(|_| result_a = Some(oper_a()));
///         s.spawn(|_| result_b = Some(oper_b()));
///     });
///     (result_a.unwrap(), result_b.unwrap())
/// }
/// ```
///
/// ### A note on threading
///
/// The closure given to `scope()` executes in the Rayon thread-pool,
/// as do those given to `spawn()`. This means that you can't access
/// thread-local variables (well, you can, but they may have
/// unexpected values).
///
/// ### Task execution
///
/// Task execution potentially starts as soon as `spawn()` is called.
/// The task will end sometime before `scope()` returns. Note that the
/// *closure* given to scope may return much earlier. In general
/// the lifetime of a scope created like `scope(body) goes something like this:
///
/// - Scope begins when `scope(body)` is called
/// - Scope body `body()` is invoked
///     - Scope tasks may be spawned
/// - Scope body returns
/// - Scope tasks execute, possibly spawning more tasks
/// - Once all tasks are done, scope ends and `scope()` returns
///
/// To see how and when tasks are joined, consider this example:
///
/// ```rust
/// // point start
/// rayon::scope(|s| {
///     s.spawn(|s| { // task s.1
///         s.spawn(|s| { // task s.1.1
///             rayon::scope(|t| {
///                 t.spawn(|_| ()); // task t.1
///                 t.spawn(|_| ()); // task t.2
///             });
///         });
///     });
///     s.spawn(|s| { // task 2
///     });
///     // point mid
/// });
/// // point end
/// ```
///
/// The various tasks that are run will execute roughly like so:
///
/// ```notrust
/// | (start)
/// |
/// | (scope `s` created)
/// +--------------------+ (task s.1)
/// +-------+ (task s.2) |
/// |       |            +---+ (task s.1.1)
/// |       |            |   |
/// |       |            |   | (scope `t` created)
/// |       |            |   +----------------+ (task t.1)
/// |       |            |   +---+ (task t.2) |
/// | (mid) |            |   |   |            |
/// :       |            |   + <-+------------+ (scope `t` ends)
/// :       |            |   |
/// |<------+------------+---+ (scope `s` ends)
/// |
/// | (end)
/// ```
///
/// The point here is that everything spawned into scope `s` will
/// terminate (at latest) at the same point -- right before the
/// original call to `rayon::scope` returns. This includes new
/// subtasks created by other subtasks (e.g., task `s.1.1`). If a new
/// scope is created (such as `t`), the things spawned into that scope
/// will be joined before that scope returns, which in turn occurs
/// before the creating task (task `s.1.1` in this case) finishes.
///
/// ### Accessing stack data
///
/// In general, spawned tasks may access stack data in place that
/// outlives the scope itself. Other data must be fully owned by the
/// spawned task.
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     s.spawn(|_| {
///         // We can access `ok` because outlives the scope `s`.
///         println!("ok: {:?}", ok);
///
///         // If we just try to use `bad` here, the closure will borrow `bad`
///         // (because we are just printing it out, and that only requires a
///         // borrow), which will result in a compilation error. Read on
///         // for options.
///         // println!("bad: {:?}", bad);
///    });
/// });
/// ```
///
/// As the comments example above suggest, to reference `bad` we must
/// take ownership of it. One way to do this is to detach the closure
/// from the surrounding stack frame, using the `move` keyword. This
/// will cause it to take ownership of *all* the variables it touches,
/// in this case including both `ok` *and* `bad`:
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     s.spawn(move |_| {
///         println!("ok: {:?}", ok);
///         println!("bad: {:?}", bad);
///     });
///
///     // That closure is fine, but now we can't use `ok` anywhere else,
///     // since it is owend by the previous task:
///     // s.spawn(|_| println!("ok: {:?}", ok));
/// });
/// ```
///
/// While this works, it could be a problem if we want to use `ok` elsewhere.
/// There are two choices. We can keep the closure as a `move` closure, but
/// instead of referencing the variable `ok`, we create a shadowed variable that
/// is a borrow of `ok` and capture *that*:
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     let ok: &Vec<i32> = &ok; // shadow the original `ok`
///     s.spawn(move |_| {
///         println!("ok: {:?}", ok); // captures the shadowed version
///         println!("bad: {:?}", bad);
///     });
///
///     // Now we too can use the shadowed `ok`, since `&Vec<i32>` references
///     // can be shared freely. Note that we need a `move` closure here though,
///     // because otherwise we'd be trying to borrow the shadowed `ok`,
///     // and that doesn't outlive `scope`.
///     s.spawn(move |_| println!("ok: {:?}", ok));
/// });
/// ```
///
/// Another option is not to use the `move` keyword but instead to take ownership
/// of individual variables:
///
/// ```rust
/// let ok: Vec<i32> = vec![1, 2, 3];
/// rayon::scope(|s| {
///     let bad: Vec<i32> = vec![4, 5, 6];
///     s.spawn(|_| {
///         // Transfer ownership of `bad` into a local variable (also named `bad`).
///         // This will force the closure to take ownership of `bad` from the environment.
///         let bad = bad;
///         println!("ok: {:?}", ok); // `ok` is only borrowed.
///         println!("bad: {:?}", bad); // refers to our local variable, above.
///     });
///
///     s.spawn(|_| println!("ok: {:?}", ok)); // we too can borrow `ok`
/// });
/// ```
///
/// ### Panics
///
/// If a panic occurs, either in the closure given to `scope()` or in
/// any of the spawned jobs, that panic will be propagated and the
/// call to `scope()` will panic. If multiple panics occurs, it is
/// non-deterministic which of their panic values will propagate.
/// Regardless, once a task is spawned using `scope.spawn()`, it will
/// execute, even if the spawning task should later panic. `scope()`
/// returns once all spawned jobs have completed, and any panics are
/// propagated at that point.
pub fn scope<'scope, OP, R>(op: OP) -> R
    where OP: for<'s> FnOnce(&'s Scope<'scope>) -> R + 'scope + Send, R: Send,
{
    in_worker(|owner_thread| {
        unsafe {
            let scope: Scope<'scope> = Scope {
                owner_thread: owner_thread as *const WorkerThread as *mut WorkerThread,
                panic: AtomicPtr::new(ptr::null_mut()),
                job_completed_latch: CountLatch::new(),
                marker: PhantomData,
            };
            let result = scope.execute_job_closure(op);
            scope.steal_till_jobs_complete();
            result.unwrap() // only None if `op` panicked, and that would have been propagated
        }
    })
}

impl<'scope> Scope<'scope> {
    /// Spawns a job into the fork-join scope `self`. This job will
    /// execute sometime before the fork-join scope completes.  The
    /// job is specified as a closure, and this closure receives its
    /// own reference to `self` as argument. This can be used to
    /// inject new jobs into `self`.
    pub fn spawn<BODY>(&self, body: BODY)
        where BODY: FnOnce(&Scope<'scope>) + 'scope
    {
        unsafe {
            self.job_completed_latch.increment();
            let job_ref = Box::new(HeapJob::new(move || self.execute_job(body)))
                .as_job_ref();
            let worker_thread = WorkerThread::current();

            // the `Scope` is not send or sync, and we only give out
            // pointers to it from within a worker thread
            debug_assert!(!WorkerThread::current().is_null());

            let worker_thread = &*worker_thread;
            worker_thread.push(job_ref);
        }
    }

    #[cfg(feature = "unstable")]
    pub fn spawn_future<F>(&self, future: F) -> RayonFuture<F::Item, F::Error>
        where F: Future + Send + 'scope
    {
        self.job_completed_latch.increment();

        // we assert that we have the future `F` type will remain
        // valid until `job_completed_latch` is fully set
        let future = unsafe { future::new_rayon_future(future, &self.job_completed_latch) };

        future
    }

    /// Executes `func` as a job, either aborting or executing as
    /// appropriate.
    ///
    /// Unsafe because it must be executed on a worker thread.
    unsafe fn execute_job<FUNC>(&self, func: FUNC)
        where FUNC: FnOnce(&Scope<'scope>) + 'scope
    {
        let _: Option<()> = self.execute_job_closure(func);
    }

    /// Executes `func` as a job in scope. Adjusts the "job completed"
    /// counters and also catches any panic and stores it into
    /// `scope`.
    ///
    /// Unsafe because this must be executed on a worker thread.
    unsafe fn execute_job_closure<FUNC, R>(&self, func: FUNC) -> Option<R>
        where FUNC: FnOnce(&Scope<'scope>) -> R + 'scope
    {
        match unwind::halt_unwinding(move || func(self)) {
            Ok(r) => { self.job_completed_ok(); Some(r) }
            Err(err) => { self.job_panicked(err); None }
        }
    }

    unsafe fn job_panicked(&self, err: Box<Any + Send + 'static>) {
        // capture the first error we see, free the rest
        let nil = ptr::null_mut();
        let mut err = Box::new(err); // box up the fat ptr
        if self.panic.compare_exchange(nil, &mut *err, Ordering::Release, Ordering::Relaxed).is_ok() {
            log!(JobPanickedErrorStored { owner_thread: (*self.owner_thread).index() });
            mem::forget(err); // ownership now transferred into self.panic
        } else {
            log!(JobPanickedErrorNotStored { owner_thread: (*self.owner_thread).index() });
        }

        self.job_completed_latch.set();
    }

    unsafe fn job_completed_ok(&self) {
        log!(JobCompletedOk { owner_thread: (*self.owner_thread).index() });
        self.job_completed_latch.set();
    }

    unsafe fn steal_till_jobs_complete(&self) {
        // wait for job counter to reach 0:
        (*self.owner_thread).wait_until(&self.job_completed_latch);

        // propagate panic, if any occurred; at this point, all
        // outstanding jobs have completed, so we can use a relaxed
        // ordering:
        let panic = self.panic.swap(ptr::null_mut(), Ordering::Relaxed);
        if !panic.is_null() {
            log!(ScopeCompletePanicked { owner_thread: (*self.owner_thread).index() });
            let value: Box<Box<Any + Send + 'static>> = mem::transmute(panic);
            unwind::resume_unwinding(*value);
        } else {
            log!(ScopeCompleteNoPanic { owner_thread: (*self.owner_thread).index() });
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/)

#![cfg(not(test))]
#![cfg(feature="cargo-fmt")]

extern crate getopts;
extern crate rustc_serialize;

use std::path::PathBuf;
use std::process::Command;
use std::env;
use std::str;

use getopts::Options;
use rustc_serialize::json::Json;

fn main() {
    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("q", "quiet", "no output printed to stdout");
    opts.optflag("v", "verbose", "use verbose output");

    let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
        Ok(m) => m,
        Err(e) => {
            print_usage(&opts, &e.to_string());
            return;
        }
    };

    let (verbose, quiet) = (matches.opt_present("v"), matches.opt_present("q"));

    if verbose && quiet {
        print_usage(&opts, "quiet mode and verbose mode are not compatible");
        return;
    }

    if matches.opt_present("h") {
        print_usage(&opts, "");
    } else {
        format_crate(&opts, verbose, quiet);
    }
}

fn print_usage(opts: &Options, reason: &str) {
    let msg = format!("{}\nusage: cargo fmt [options]", reason);
    println!("{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
              Arguments after `--` are passed to rustfmt.",
             opts.usage(&msg));
}

fn format_crate(opts: &Options, verbose: bool, quiet: bool) {
    let targets = match get_targets() {
        Ok(t) => t,
        Err(e) => {
            print_usage(opts, &e.to_string());
            return;
        }
    };

    // Currently only bin and lib files get formatted
    let files: Vec<_> = targets.into_iter()
                               .filter(|t| t.kind.is_lib() | t.kind.is_bin())
                               .inspect(|t| {
                                   if verbose {
                                       println!("[{:?}] {:?}", t.kind, t.path)
                                   }
                               })
                               .map(|t| t.path)
                               .collect();

    format_files(&files, &get_fmt_args(), verbose, quiet)
        .unwrap_or_else(|e| print_usage(opts, &e.to_string()));
}

fn get_fmt_args() -> Vec<String> {
    // All arguments after -- are passed to rustfmt
    env::args().skip_while(|a| a != "--").skip(1).collect()
}

#[derive(Debug)]
enum TargetKind {
    Lib, // dylib, staticlib, lib
    Bin, // bin
    Other, // test, plugin,...
}

impl TargetKind {
    fn is_lib(&self) -> bool {
        match self {
            &TargetKind::Lib => true,
            _ => false,
        }
    }

    fn is_bin(&self) -> bool {
        match self {
            &TargetKind::Bin => true,
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct Target {
    path: PathBuf,
    kind: TargetKind,
}

// Returns a vector of all compile targets of a crate
fn get_targets() -> Result<Vec<Target>, std::io::Error> {
    let mut targets: Vec<Target> = vec![];
    let output = try!(Command::new("cargo").arg("read-manifest").output());
    if output.status.success() {
        // None of the unwraps should fail if output of `cargo read-manifest` is correct
        let data = &String::from_utf8(output.stdout).unwrap();
        let json = Json::from_str(data).unwrap();
        let jtargets = json.find("targets").unwrap().as_array().unwrap();
        for jtarget in jtargets {
            targets.push(target_from_json(jtarget));
        }

        Ok(targets)
    } else {
        // This happens when cargo-fmt is not used inside a crate
        Err(std::io::Error::new(std::io::ErrorKind::NotFound,
                                str::from_utf8(&output.stderr).unwrap()))
    }
}

fn target_from_json(jtarget: &Json) -> Target {
    let jtarget = jtarget.as_object().unwrap();
    let path = PathBuf::from(jtarget.get("src_path").unwrap().as_string().unwrap());
    let kinds = jtarget.get("kind").unwrap().as_array().unwrap();
    let kind = match kinds[0].as_string().unwrap() {
        "bin" => TargetKind::Bin,
        "lib" | "dylib" | "staticlib" => TargetKind::Lib,
        _ => TargetKind::Other,
    };

    Target {
        path: path,
        kind: kind,
    }
}

fn format_files(files: &Vec<PathBuf>,
                fmt_args: &Vec<String>,
                verbose: bool,
                quiet: bool)
                -> Result<(), std::io::Error> {
    let stdout = if quiet {
        std::process::Stdio::null()
    } else {
        std::process::Stdio::inherit()
    };
    if verbose {
        print!("rustfmt");
        for a in fmt_args.iter() {
            print!(" {}", a);
        }
        for f in files.iter() {
            print!(" {}", f.display());
        }
        println!("");
    }
    let mut command = try!(Command::new("rustfmt")
                               .stdout(stdout)
                               .args(files)
                               .args(fmt_args)
                               .spawn());
    try!(command.wait());

    Ok(())
}
<commit_msg>cargo-fmt: don't return zero on failure<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/)

#![cfg(not(test))]
#![cfg(feature="cargo-fmt")]
#![deny(warnings)]

extern crate getopts;
extern crate rustc_serialize;

use std::env;
use std::io::Write;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::str;

use getopts::Options;
use rustc_serialize::json::Json;

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

fn execute() -> i32 {
    let success = 0;
    let failure = 1;

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("q", "quiet", "no output printed to stdout");
    opts.optflag("v", "verbose", "use verbose output");

    let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
        Ok(m) => m,
        Err(e) => {
            print_usage(&opts, &e.to_string());
            return failure;
        }
    };

    let (verbose, quiet) = (matches.opt_present("v"), matches.opt_present("q"));

    if verbose && quiet {
        print_usage(&opts, "quiet mode and verbose mode are not compatible");
        return failure;
    }

    if matches.opt_present("h") {
        print_usage(&opts, "");
        return success;
    }

    match format_crate(verbose, quiet) {
        Err(e) => {
            print_usage(&opts, &e.to_string());
            failure
        }
        Ok(status) => {
            if status.success() {
                success
            } else {
                status.code().unwrap_or(failure)
            }
        }
    }
}

fn print_usage(opts: &Options, reason: &str) {
    let msg = format!("{}\nusage: cargo fmt [options]", reason);
    println!("{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
              Arguments after `--` are passed to rustfmt.",
             opts.usage(&msg));
}

fn format_crate(verbose: bool, quiet: bool) -> Result<ExitStatus, std::io::Error> {
    let targets = try!(get_targets());

    // Currently only bin and lib files get formatted
    let files: Vec<_> = targets.into_iter()
                               .filter(|t| t.kind.is_lib() | t.kind.is_bin())
                               .inspect(|t| {
                                   if verbose {
                                       println!("[{:?}] {:?}", t.kind, t.path)
                                   }
                               })
                               .map(|t| t.path)
                               .collect();

    format_files(&files, &get_fmt_args(), verbose, quiet)
}

fn get_fmt_args() -> Vec<String> {
    // All arguments after -- are passed to rustfmt
    env::args().skip_while(|a| a != "--").skip(1).collect()
}

#[derive(Debug)]
enum TargetKind {
    Lib, // dylib, staticlib, lib
    Bin, // bin
    Other, // test, plugin,...
}

impl TargetKind {
    fn is_lib(&self) -> bool {
        match self {
            &TargetKind::Lib => true,
            _ => false,
        }
    }

    fn is_bin(&self) -> bool {
        match self {
            &TargetKind::Bin => true,
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct Target {
    path: PathBuf,
    kind: TargetKind,
}

// Returns a vector of all compile targets of a crate
fn get_targets() -> Result<Vec<Target>, std::io::Error> {
    let mut targets: Vec<Target> = vec![];
    let output = try!(Command::new("cargo").arg("read-manifest").output());
    if output.status.success() {
        // None of the unwraps should fail if output of `cargo read-manifest` is correct
        let data = &String::from_utf8(output.stdout).unwrap();
        let json = Json::from_str(data).unwrap();
        let jtargets = json.find("targets").unwrap().as_array().unwrap();
        for jtarget in jtargets {
            targets.push(target_from_json(jtarget));
        }

        Ok(targets)
    } else {
        // This happens when cargo-fmt is not used inside a crate
        Err(std::io::Error::new(std::io::ErrorKind::NotFound,
                                str::from_utf8(&output.stderr).unwrap()))
    }
}

fn target_from_json(jtarget: &Json) -> Target {
    let jtarget = jtarget.as_object().unwrap();
    let path = PathBuf::from(jtarget.get("src_path").unwrap().as_string().unwrap());
    let kinds = jtarget.get("kind").unwrap().as_array().unwrap();
    let kind = match kinds[0].as_string().unwrap() {
        "bin" => TargetKind::Bin,
        "lib" | "dylib" | "staticlib" => TargetKind::Lib,
        _ => TargetKind::Other,
    };

    Target {
        path: path,
        kind: kind,
    }
}

fn format_files(files: &Vec<PathBuf>,
                fmt_args: &Vec<String>,
                verbose: bool,
                quiet: bool)
                -> Result<ExitStatus, std::io::Error> {
    let stdout = if quiet {
        std::process::Stdio::null()
    } else {
        std::process::Stdio::inherit()
    };
    if verbose {
        print!("rustfmt");
        for a in fmt_args.iter() {
            print!(" {}", a);
        }
        for f in files.iter() {
            print!(" {}", f.display());
        }
        println!("");
    }
    let mut command = try!(Command::new("rustfmt")
                               .stdout(stdout)
                               .args(files)
                               .args(fmt_args)
                               .spawn());
    command.wait()
}
<|endoftext|>