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

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

use crate::{
    code_action::on_code_action,
    completion::{on_completion, on_resolve_completion_item},
    explore_schema_for_type::{on_explore_schema_for_type, ExploreSchemaForType},
    goto_definition::{
        on_get_source_location_of_type_definition, on_goto_definition,
        GetSourceLocationOfTypeDefinition,
    },
    graphql_tools::on_graphql_execute_query,
    graphql_tools::GraphQLExecuteQuery,
    hover::on_hover,
    js_language_server::JSLanguageServer,
    lsp_process_error::LSPProcessResult,
    lsp_runtime_error::LSPRuntimeError,
    references::on_references,
    resolved_types_at_location::{on_get_resolved_types_at_location, ResolvedTypesAtLocation},
    search_schema_items::{on_search_schema_items, SearchSchemaItems},
    server::lsp_state_resources::LSPStateResources,
    shutdown::{on_exit, on_shutdown},
    status_reporter::LSPStatusReporter,
    status_updater::set_initializing_status,
    text_documents::{
        on_cancel, on_did_change_text_document, on_did_close_text_document,
        on_did_open_text_document, on_did_save_text_document,
    },
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::channel::SendError;
use log::debug;
pub use lsp_notification_dispatch::LSPNotificationDispatch;
pub use lsp_request_dispatch::LSPRequestDispatch;
use lsp_server::{
    Connection, ErrorCode, Message, Notification, Response as ServerResponse, ResponseError,
};
use lsp_types::{
    notification::{
        Cancel, DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument,
        DidSaveTextDocument, Exit,
    },
    request::{
        CodeActionRequest, Completion, GotoDefinition, HoverRequest, References,
        ResolveCompletionItem, Shutdown,
    },
    CodeActionProviderCapability, CompletionOptions, InitializeParams, ServerCapabilities,
    TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions,
};
use relay_compiler::{config::Config, NoopArtifactWriter};
use schema_documentation::{SchemaDocumentation, SchemaDocumentationLoader};
use std::{sync::Arc, thread};
use tokio::task;

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

use heartbeat::{on_heartbeat, HeartbeatRequest};

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

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

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

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

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

    config.artifact_writer = Box::new(NoopArtifactWriter);
    config.status_reporter = Box::new(LSPStatusReporter::new(
        config.root_dir.clone(),
        connection.sender.clone(),
    ));

    let lsp_state = Arc::new(LSPState::create_state(
        Arc::new(config),
        Arc::clone(&perf_logger),
        extra_data_provider,
        schema_documentation_loader,
        js_resource,
        connection.sender.clone(),
    ));
    // Watchman Subscription to handle Schema/Programs/Validation
    let lsp_state_clone = Arc::clone(&lsp_state);
    task::spawn(async move {
        LSPStateResources::new(lsp_state_clone)
            .watch()
            .await
            .unwrap();
    });

    loop {
        debug!("waiting for incoming messages...");
        match connection.receiver.recv() {
            Ok(Message::Request(request)) => {
                debug!("creating a task to handle request...");
                let lsp_state = Arc::clone(&lsp_state);
                thread::spawn(move || {
                    handle_request(lsp_state, request).unwrap();
                });
            }
            Ok(Message::Notification(notification)) => {
                debug!("creating a task to handle notification...");
                let lsp_state = Arc::clone(&lsp_state);
                thread::spawn(move || {
                    handle_notification(lsp_state, notification);
                });
            }
            Ok(_) => {
                // Ignore responses for now
            }
            Err(error) => {
                panic!("Relay Language Server receiver error {:?}", error);
            }
        }
    }
}

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

    lsp_state.send_message(Message::Response(get_response(request)))
}

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

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

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

        let response = get_response(request);

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

        lsp_request_event.stop(lsp_request_processing_time);
        lsp_request_event.complete();

        response
    }
}

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

    let notification_result = dispatch_notification(notification, &lsp_state);

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

    lsp_notification_event.stop(lsp_notification_processing_time);
    lsp_notification_event.complete();
}

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

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

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

use crate::{
    code_action::on_code_action,
    completion::{on_completion, on_resolve_completion_item},
    explore_schema_for_type::{on_explore_schema_for_type, ExploreSchemaForType},
    goto_definition::{
        on_get_source_location_of_type_definition, on_goto_definition,
        GetSourceLocationOfTypeDefinition,
    },
    graphql_tools::on_graphql_execute_query,
    graphql_tools::GraphQLExecuteQuery,
    hover::on_hover,
    js_language_server::JSLanguageServer,
    lsp_process_error::LSPProcessResult,
    lsp_runtime_error::LSPRuntimeError,
    references::on_references,
    resolved_types_at_location::{on_get_resolved_types_at_location, ResolvedTypesAtLocation},
    search_schema_items::{on_search_schema_items, SearchSchemaItems},
    server::lsp_state_resources::LSPStateResources,
    shutdown::{on_exit, on_shutdown},
    status_reporter::LSPStatusReporter,
    status_updater::set_initializing_status,
    text_documents::{
        on_cancel, on_did_change_text_document, on_did_close_text_document,
        on_did_open_text_document, on_did_save_text_document,
    },
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::channel::SendError;
use log::debug;
pub use lsp_notification_dispatch::LSPNotificationDispatch;
pub use lsp_request_dispatch::LSPRequestDispatch;
use lsp_server::{
    Connection, ErrorCode, Message, Notification, Response as ServerResponse, ResponseError,
};
use lsp_types::{
    notification::{
        Cancel, DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument,
        DidSaveTextDocument, Exit,
    },
    request::{
        CodeActionRequest, Completion, GotoDefinition, HoverRequest, References,
        ResolveCompletionItem, Shutdown,
    },
    CodeActionProviderCapability, CompletionOptions, InitializeParams, ServerCapabilities,
    TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions,
};
use relay_compiler::{config::Config, NoopArtifactWriter};
use schema_documentation::{SchemaDocumentation, SchemaDocumentationLoader};
use std::{sync::Arc, thread};
use tokio::task;

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

use heartbeat::{on_heartbeat, HeartbeatRequest};

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

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

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

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

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

    config.artifact_writer = Box::new(NoopArtifactWriter);
    config.status_reporter = Box::new(LSPStatusReporter::new(
        config.root_dir.clone(),
        connection.sender.clone(),
    ));

    let lsp_state = Arc::new(LSPState::create_state(
        Arc::new(config),
        Arc::clone(&perf_logger),
        extra_data_provider,
        schema_documentation_loader,
        js_resource,
        connection.sender.clone(),
    ));
    // Watchman Subscription to handle Schema/Programs/Validation
    let lsp_state_clone = Arc::clone(&lsp_state);
    task::spawn(async move {
        LSPStateResources::new(lsp_state_clone)
            .watch()
            .await
            .unwrap();
    });

    loop {
        debug!("waiting for incoming messages...");
        match connection.receiver.recv() {
            Ok(Message::Request(request)) => {
                debug!("creating a task to handle request...");
                let lsp_state = Arc::clone(&lsp_state);
                thread::spawn(move || {
                    handle_request(lsp_state, request).unwrap();
                });
            }
            Ok(Message::Notification(notification)) => {
                debug!("creating a task to handle notification...");
                let lsp_state = Arc::clone(&lsp_state);
                thread::spawn(move || {
                    handle_notification(lsp_state, notification);
                });
            }
            Ok(_) => {
                // Ignore responses for now
            }
            Err(error) => {
                panic!("Relay Language Server receiver error {:?}", error);
            }
        }
    }
}

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

    lsp_state.send_message(Message::Response(get_response(request)))
}

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

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

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

        let response = get_response(request);

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

        lsp_request_event.stop(lsp_request_processing_time);
        lsp_request_event.complete();

        response
    }
}

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

    let notification_result = dispatch_notification(notification, lsp_state.as_ref());

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

    lsp_notification_event.stop(lsp_notification_processing_time);
    lsp_notification_event.complete();
}

fn dispatch_notification(
    notification: lsp_server::Notification,
    lsp_state: &impl GlobalState,
) -> Result<(), Option<LSPRuntimeError>> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Cancel>(on_cancel)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

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


<commit_before>#[link(name="cat", vers="1.0.0", author="Seldaek")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Jordi Boggiano <j.boggiano@seld.be>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

/* last synced with: cat (GNU coreutils) 8.13 */

extern mod extra;

use std::os;
use std::io::{stdin, stderr, stdout, Writer, Reader};
use extra::getopts::groups;

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        groups::optflag("A", "show-all", "equivalent to -vET"),
        groups::optflag("b", "number-nonblank", "number nonempty output lines, overrides -n"),
        groups::optflag("e", "", "equivalent to -vE"),
        groups::optflag("E", "show-ends", "display $ at end of each line"),
        groups::optflag("n", "number", "number all output lines"),
        groups::optflag("s", "squeeze-blank", "suppress repeated empty output lines"),
        groups::optflag("t", "", "equivalent to -vT"),
        groups::optflag("T", "show-tabs", "display TAB characters as ^I"),
        groups::optflag("v", "show-nonprinting", "use ^ and M- notation, except for LF (\\n) and TAB (\\t)"),
        groups::optflag("h", "help", "display this help and exit"),
        groups::optflag("V", "version", "output version information and exit"),
    ];
    let matches = match groups::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            stderr().write_line("Invalid options");
            stderr().write_line(f.to_err_msg());
            os::set_exit_status(1);
            return
        }
    };
    if matches.opts_present([~"h", ~"help"]) {
        println("cat 1.0.0");
        println("");
        println("Usage:");
        println(fmt!("  %s [OPTION]... [FILE]...", program));
        println("");
        print(groups::usage("Concatenate FILE(s), or standard input, to standard output.", opts));
        println("");
        println("With no FILE, or when FILE is -, read standard input.");
        return;
    }
    if matches.opts_present([~"V", ~"version"]) {
        println("cat 1.0.0");
        return;
    }

    let mut number_mode = NumberNone;
    if matches.opts_present([~"n", ~"number"]) {
        number_mode = NumberAll;
    }
    if matches.opts_present([~"b", ~"number-nonblank"]) {
        number_mode = NumberNonEmpty;
    }
    let show_nonprint = matches.opts_present([~"v", ~"show-nonprinting", ~"A", ~"show-all", ~"t", ~"e"]);
    let show_ends = matches.opts_present([~"E", ~"show-ends", ~"A", ~"show-all", ~"e"]);
    let show_tabs = matches.opts_present([~"T", ~"show-tabs",  ~"A", ~"show-all", ~"t"]);
    let squeeze_blank = matches.opts_present([~"s", ~"squeeze-blank"]);
    let mut files = matches.free;
    if files.is_empty() {
        files = ~[~"-"];
    }

    exec(files, number_mode, show_nonprint, show_ends, show_tabs, squeeze_blank);
}

#[deriving(Eq)]
pub enum NumberingMode {
    NumberNone,
    NumberNonEmpty,
    NumberAll,
}

static TAB: u8 = '\t' as u8;
static CR: u8 = '\r' as u8;
static LF: u8 = '\n' as u8;

#[cfg(windows)]
fn is_newline_char(byte: u8) -> bool {
    byte == LF || byte == CR
}

#[cfg(unix)]
fn is_newline_char(byte: u8) -> bool {
    byte == LF
}

pub fn exec(files: ~[~str], number: NumberingMode, show_nonprint: bool, show_ends: bool, show_tabs: bool, squeeze_blank: bool) {
    let writer = stdout();

    if NumberNone != number || show_nonprint || show_ends || show_tabs || squeeze_blank {
        for path in files.iter() {
            let reader = match open(path.to_owned()) {
                Some(f) => f,
                None => { continue }
            };

            let mut counter: uint = 1;
            let mut at_line_start = true;
            let is_numbering = number == NumberAll || number == NumberNonEmpty;

            loop {
                let buf = reader.read_bytes(2);
                for byte in buf.iter() {
                    if at_line_start && (number == NumberAll || (number == NumberNonEmpty && !is_newline_char(*byte))) {
                        writer.write_str(fmt!("%6u  ", counter));
                        counter += 1;
                        at_line_start = false;
                    }
                    if is_numbering && *byte == LF {
                        at_line_start = true;
                    }
                    if show_tabs && *byte == TAB {
                        writer.write(bytes!("^I"));
                    } else if show_ends && *byte == LF {
                        writer.write(bytes!("$\n"));
                    } else if show_nonprint && (*byte < 32 || *byte >= 127) && !is_newline_char(*byte) {
                        let mut byte = *byte;
                        if byte >= 128 {
                            writer.write(bytes!("M-"));
                            byte = byte - 128;
                        }
                        if byte < 32 {
                            writer.write(['^' as u8, byte + 64]);
                        } else if byte == 127 {
                            writer.write(['^' as u8, byte - 64]);
                        } else {
                            writer.write([byte]);
                        }
                    } else {
                        writer.write([*byte]);
                    }
                }

                if reader.eof() {
                    break;
                }
            }
        }

        return;
    }

    // passthru mode
    for path in files.iter() {
        let reader = match open(path.to_owned()) {
            Some(f) => f,
            None => { continue }
        };

        loop {
            writer.write(reader.read_bytes(100000));
            if reader.eof() {
                break;
            }
        }
    }
}

fn open(path: ~str) -> Option<@Reader> {
    if "-" == path {
        return Some(stdin());
    }

    match std::io::file_reader(&std::path::Path(path)) {
        Ok(fd) => return Some(fd),
        Err(e) => {
            stderr().write_line(fmt!("cat: %s: %s", path, e));
            os::set_exit_status(1);
        }
    }

    None
}
<commit_msg>Use new Path "constructor"<commit_after>#[link(name="cat", vers="1.0.0", author="Seldaek")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Jordi Boggiano <j.boggiano@seld.be>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

/* last synced with: cat (GNU coreutils) 8.13 */

extern mod extra;

use std::os;
use std::io::{stdin, stderr, stdout, Writer, Reader};
use extra::getopts::groups;

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        groups::optflag("A", "show-all", "equivalent to -vET"),
        groups::optflag("b", "number-nonblank", "number nonempty output lines, overrides -n"),
        groups::optflag("e", "", "equivalent to -vE"),
        groups::optflag("E", "show-ends", "display $ at end of each line"),
        groups::optflag("n", "number", "number all output lines"),
        groups::optflag("s", "squeeze-blank", "suppress repeated empty output lines"),
        groups::optflag("t", "", "equivalent to -vT"),
        groups::optflag("T", "show-tabs", "display TAB characters as ^I"),
        groups::optflag("v", "show-nonprinting", "use ^ and M- notation, except for LF (\\n) and TAB (\\t)"),
        groups::optflag("h", "help", "display this help and exit"),
        groups::optflag("V", "version", "output version information and exit"),
    ];
    let matches = match groups::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            stderr().write_line("Invalid options");
            stderr().write_line(f.to_err_msg());
            os::set_exit_status(1);
            return
        }
    };
    if matches.opts_present([~"h", ~"help"]) {
        println("cat 1.0.0");
        println("");
        println("Usage:");
        println(fmt!("  %s [OPTION]... [FILE]...", program));
        println("");
        print(groups::usage("Concatenate FILE(s), or standard input, to standard output.", opts));
        println("");
        println("With no FILE, or when FILE is -, read standard input.");
        return;
    }
    if matches.opts_present([~"V", ~"version"]) {
        println("cat 1.0.0");
        return;
    }

    let mut number_mode = NumberNone;
    if matches.opts_present([~"n", ~"number"]) {
        number_mode = NumberAll;
    }
    if matches.opts_present([~"b", ~"number-nonblank"]) {
        number_mode = NumberNonEmpty;
    }
    let show_nonprint = matches.opts_present([~"v", ~"show-nonprinting", ~"A", ~"show-all", ~"t", ~"e"]);
    let show_ends = matches.opts_present([~"E", ~"show-ends", ~"A", ~"show-all", ~"e"]);
    let show_tabs = matches.opts_present([~"T", ~"show-tabs",  ~"A", ~"show-all", ~"t"]);
    let squeeze_blank = matches.opts_present([~"s", ~"squeeze-blank"]);
    let mut files = matches.free;
    if files.is_empty() {
        files = ~[~"-"];
    }

    exec(files, number_mode, show_nonprint, show_ends, show_tabs, squeeze_blank);
}

#[deriving(Eq)]
pub enum NumberingMode {
    NumberNone,
    NumberNonEmpty,
    NumberAll,
}

static TAB: u8 = '\t' as u8;
static CR: u8 = '\r' as u8;
static LF: u8 = '\n' as u8;

#[cfg(windows)]
fn is_newline_char(byte: u8) -> bool {
    byte == LF || byte == CR
}

#[cfg(unix)]
fn is_newline_char(byte: u8) -> bool {
    byte == LF
}

pub fn exec(files: ~[~str], number: NumberingMode, show_nonprint: bool, show_ends: bool, show_tabs: bool, squeeze_blank: bool) {
    let writer = stdout();

    if NumberNone != number || show_nonprint || show_ends || show_tabs || squeeze_blank {
        for path in files.iter() {
            let reader = match open(path.to_owned()) {
                Some(f) => f,
                None => { continue }
            };

            let mut counter: uint = 1;
            let mut at_line_start = true;
            let is_numbering = number == NumberAll || number == NumberNonEmpty;

            loop {
                let buf = reader.read_bytes(2);
                for byte in buf.iter() {
                    if at_line_start && (number == NumberAll || (number == NumberNonEmpty && !is_newline_char(*byte))) {
                        writer.write_str(fmt!("%6u  ", counter));
                        counter += 1;
                        at_line_start = false;
                    }
                    if is_numbering && *byte == LF {
                        at_line_start = true;
                    }
                    if show_tabs && *byte == TAB {
                        writer.write(bytes!("^I"));
                    } else if show_ends && *byte == LF {
                        writer.write(bytes!("$\n"));
                    } else if show_nonprint && (*byte < 32 || *byte >= 127) && !is_newline_char(*byte) {
                        let mut byte = *byte;
                        if byte >= 128 {
                            writer.write(bytes!("M-"));
                            byte = byte - 128;
                        }
                        if byte < 32 {
                            writer.write(['^' as u8, byte + 64]);
                        } else if byte == 127 {
                            writer.write(['^' as u8, byte - 64]);
                        } else {
                            writer.write([byte]);
                        }
                    } else {
                        writer.write([*byte]);
                    }
                }

                if reader.eof() {
                    break;
                }
            }
        }

        return;
    }

    // passthru mode
    for path in files.iter() {
        let reader = match open(path.to_owned()) {
            Some(f) => f,
            None => { continue }
        };

        loop {
            writer.write(reader.read_bytes(100000));
            if reader.eof() {
                break;
            }
        }
    }
}

fn open(path: ~str) -> Option<@Reader> {
    if "-" == path {
        return Some(stdin());
    }

    match std::io::file_reader(&std::path::Path::new(path.as_slice())) {
        Ok(fd) => return Some(fd),
        Err(e) => {
            stderr().write_line(fmt!("cat: %s: %s", path, e));
            os::set_exit_status(1);
        }
    }

    None
}
<|endoftext|>


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

pub mod logger;
pub mod workspace;
pub mod postprocessor;

pub use protocol::jobsrv::JobState;

use std::ffi::OsString;
use std::fs;
use std::ops::{Deref, DerefMut};
use std::path::Path;
use std::process::{Command, Stdio};
use std::str::FromStr;
use std::sync::{mpsc, Arc, RwLock};
use std::thread::{self, JoinHandle};

use hab_core::package::archive::PackageArchive;
use hab_core::package::install::PackageInstall;
use hab_core::package::PackageIdent;
use hab_net::server::ZMQ_CONTEXT;
use protobuf::{parse_from_bytes, Message};
use protocol::jobsrv as proto;
use zmq;

use self::logger::Logger;
use self::postprocessor::PostProcessor;
use self::workspace::Workspace;
use config::Config;
use error::{Error, Result};
use vcs;

/// In-memory zmq address of Job RunnerMgr
const INPROC_ADDR: &'static str = "inproc://runner";
/// Protocol message to indicate the Job Runner has received a work request
const WORK_ACK: &'static str = "A";
/// Protocol message to indicate the Job Runner has completed a work request
const WORK_COMPLETE: &'static str = "C";

lazy_static! {
    // JW TODO: expose public API functions in the core crate to check if the Rust proccess which
    // is currently executing is, itself, packaged by Habitat. If so, then we should expose another
    // public API function to load the PackageInstall of a dep for the given `origin`/`name`
    // combination. If we can't, then we should fall back to the latest of core/hab-studio because
    // that means we're just in a dev shell.
    static ref STUDIO_PKG: PackageIdent = PackageIdent::from_str("core/hab-studio").unwrap();
}

#[derive(Debug)]
pub struct Job(proto::Job);

impl Job {
    pub fn new(job: proto::Job) -> Self {
        Job(job)
    }

    pub fn vcs(&self) -> &vcs::RemoteSource {
        if self.0.get_project().has_git() {
            return self.0.get_project().get_git();
        }
        unreachable!("unknown vcs associated with job's project");
    }
}

impl Deref for Job {
    type Target = proto::Job;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Job {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

pub struct Runner {
    workspace: Workspace,
    auth_token: String,
    logger: Option<Logger>,
}

impl Runner {
    pub fn new(job: Job, config: &Config) -> Self {
        Runner {
            auth_token: config.auth_token.clone(),
            workspace: Workspace::new(config.data_path.clone(), job),
            logger: None,
        }
    }

    pub fn job(&self) -> &Job {
        &self.workspace.job
    }

    pub fn job_mut(&mut self) -> &mut Job {
        &mut self.workspace.job
    }

    pub fn logger(&mut self) -> &mut Logger {
        self.logger.as_mut().expect("logger not initialized")
    }

    pub fn run(mut self) -> Job {
        if let Some(err) = self.setup().err() {
            error!("WORKSPACE SETUP ERR={:?}", err);
            return self.fail();
        }
        // JW TODO: How are we going to get the secret keys for this thing?
        if let Some(err) = self.job().vcs().clone(&self.workspace.src()).err() {
            error!("CLONE ERROR={}", err);
            return self.fail();
        }
        let mut archive = match self.build() {
            Ok(archive) => archive,
            Err(err) => {
                error!("STUDIO ERR={}", err);
                return self.fail();
            }
        };

        let mut post_processor = PostProcessor::new(&self.workspace);

        if !post_processor.run(&mut archive, &self.auth_token) {
            // JW TODO: We should shelve the built artifacts and allow a retry on post-processing.
            // If the job is killed then we can kill the shelved artifacts.
            return self.fail();
        }

        if let Some(err) = fs::remove_dir_all(self.workspace.out()).err() {
            error!("unable to remove out directory ({}), ERR={:?}",
                   self.workspace.out().display(),
                   err)
        }
        self.complete()
    }

    fn build(&mut self) -> Result<PackageArchive> {
        let args = vec![OsString::from("-s"),
                        OsString::from(self.workspace.src()),
                        OsString::from("-r"),
                        OsString::from(self.workspace.studio()),
                        OsString::from("-k"),
                        OsString::from("core"),
                        OsString::from("build"),
                        OsString::from(Path::new(self.job().get_project().get_plan_path())
                            .parent()
                            .unwrap())];
        let command = studio_cmd();
        debug!("building, cmd={:?}, args={:?}", command, args);
        let mut child = Command::new(command)
            .args(&args)
            .env_clear()
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .expect("failed to spawn child");
        self.logger().pipe(&mut child);
        let exit_status = child.wait().expect("failed to wait on child");
        debug!("build complete, status={:?}", exit_status);
        if exit_status.success() {
            try!(fs::rename(self.workspace.src().join("results"), self.workspace.out()));
            self.workspace.last_built()
        } else {
            Err(Error::BuildFailure(exit_status.code().unwrap_or(-1)))
        }
    }

    fn complete(mut self) -> Job {
        self.teardown().err().map(|e| error!("{}", e));
        self.workspace.job.set_state(JobState::Complete);
        self.workspace.job
    }

    fn fail(mut self) -> Job {
        self.teardown().err().map(|e| error!("{}", e));
        self.workspace.job.set_state(JobState::Failed);
        self.workspace.job
    }

    fn post_process(&self, archive: &mut PackageArchive) -> bool {
        // JW TODO: In the future we'll support multiple and configurable post processors, but for
        // now let's just publish to the public depot
        //
        // Things to solve right now
        // * Where do we get the token for authentication?
        //      * Should the workers ask for a lease from the JobSrv?
        let client =
            depot_client::Client::new(&hab_core::url::default_depot_url(), PRODUCT, VERSION, None)
                .unwrap();
        if let Some(err) = client.x_put_package(archive, &self.auth_token).err() {
            error!("post processing error, ERR={:?}", err);
        }
        if let Some(err) = fs::remove_dir_all(self.workspace.out()).err() {
            error!("unable to remove out directory ({}), ERR={:?}",
                   self.workspace.out().display(),
                   err)
        }
        true
    }

    fn setup(&mut self) -> Result<()> {
        if let Some(err) = fs::create_dir_all(self.workspace.src()).err() {
            return Err(Error::WorkspaceSetup(format!("{}", self.workspace.src().display()), err));
        }
        self.logger = Some(Logger::init(&self.workspace));
        Ok(())
    }

    fn teardown(&mut self) -> Result<()> {
        let args = vec![OsString::from("-s"),
                        OsString::from(self.workspace.src()),
                        OsString::from("-r"),
                        OsString::from(self.workspace.studio()),
                        OsString::from("rm"),
                        OsString::from(Path::new(self.job().get_project().get_plan_path())
                            .parent()
                            .unwrap())];
        let command = studio_cmd();
        debug!("removing studio, cmd={:?}, args={:?}", command, args);
        let mut child = Command::new(command)
            .args(&args)
            .env_clear()
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .expect("failed to spawn child");
        self.logger().pipe(&mut child);
        let exit_status = child.wait().expect("failed to wait on child");
        debug!("studio removal complete, status={:?}", exit_status);
        if exit_status.success() {
            if let Some(err) = fs::remove_dir_all(self.workspace.src()).err() {
                return Err(Error::WorkspaceTeardown(format!("{}", self.workspace.src().display()),
                                                    err));
            }
            Ok(())
        } else {
            Err(Error::BuildFailure(exit_status.code().unwrap_or(-1)))
        }
    }
}

/// Client for sending and receiving messages to and from the Job Runner
pub struct RunnerCli {
    sock: zmq::Socket,
    msg: zmq::Message,
}

impl RunnerCli {
    /// Create a new Job Runner client
    pub fn new() -> Self {
        let sock = (**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER).unwrap();
        RunnerCli {
            sock: sock,
            msg: zmq::Message::new().unwrap(),
        }
    }

    /// Return a poll item used in `zmq::poll` for awaiting messages on multiple sockets
    pub fn as_poll_item<'a>(&'a self, events: i16) -> zmq::PollItem<'a> {
        self.sock.as_poll_item(events)
    }

    /// Connect to the Job Runner
    pub fn connect(&mut self) -> Result<()> {
        try!(self.sock.connect(INPROC_ADDR));
        Ok(())
    }

    /// Wait until client receives a work received acknowledgement by the Runner and return
    /// the assigned JobID.
    pub fn recv_ack(&mut self) -> Result<&zmq::Message> {
        try!(self.sock.recv(&mut self.msg, 0));
        if Some(WORK_ACK) != self.msg.as_str() {
            unreachable!("wk:run:1, received unexpected response from runner");
        }
        try!(self.sock.recv(&mut self.msg, 0));
        Ok(&self.msg)
    }

    /// Wait until client receives a work complete message by the Runner and return an encoded
    /// representation of the job.
    pub fn recv_complete(&mut self) -> Result<&zmq::Message> {
        try!(self.sock.recv(&mut self.msg, 0));
        if Some(WORK_COMPLETE) != self.msg.as_str() {
            unreachable!("wk:run:2, received unexpected response from runner");
        }
        try!(self.sock.recv(&mut self.msg, 0));
        Ok(&self.msg)
    }

    /// Send a message to the Job Runner
    pub fn send(&mut self, msg: &zmq::Message) -> Result<()> {
        try!(self.sock.send(&*msg, 0));
        Ok(())
    }
}

/// Receives work notifications from a `RunnerCli` and performs long-running tasks in a
/// separate thread.
pub struct RunnerMgr {
    sock: zmq::Socket,
    msg: zmq::Message,
    config: Arc<RwLock<Config>>,
}

impl RunnerMgr {
    /// Start the Job Runner
    pub fn start(config: Arc<RwLock<Config>>) -> Result<JoinHandle<()>> {
        let (tx, rx) = mpsc::sync_channel(0);
        let handle = thread::Builder::new()
            .name("runner".to_string())
            .spawn(move || {
                let mut runner = Self::new(config).unwrap();
                runner.run(tx).unwrap();
            })
            .unwrap();
        match rx.recv() {
            Ok(()) => Ok(handle),
            Err(e) => panic!("runner thread startup error, err={}", e),
        }
    }

    fn new(config: Arc<RwLock<Config>>) -> Result<Self> {
        let sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER));
        Ok(RunnerMgr {
            sock: sock,
            msg: zmq::Message::new().unwrap(),
            config: config,
        })
    }

    // Main loop for server
    fn run(&mut self, rz: mpsc::SyncSender<()>) -> Result<()> {
        try!(self.sock.bind(INPROC_ADDR));
        rz.send(()).unwrap();
        loop {
            let job = try!(self.recv_job());
            try!(self.send_ack(&job));
            try!(self.execute_job(job));
        }
    }

    fn execute_job(&mut self, job: Job) -> Result<()> {
        let runner = {
            Runner::new(job, &self.config.read().unwrap())
        };
        debug!("executing work, job={:?}", runner.job());
        let job = runner.run();
        self.send_complete(&job)
    }

    fn recv_job(&mut self) -> Result<Job> {
        try!(self.sock.recv(&mut self.msg, 0));
        let job: proto::Job = parse_from_bytes(&self.msg).unwrap();
        Ok(Job::new(job))
    }

    fn send_ack(&mut self, job: &Job) -> Result<()> {
        debug!("received work, job={:?}", job);
        try!(self.sock.send_str(WORK_ACK, zmq::SNDMORE));
        try!(self.sock.send(&*job.write_to_bytes().unwrap(), 0));
        Ok(())
    }

    fn send_complete(&mut self, job: &Job) -> Result<()> {
        debug!("work complete, job={:?}", job);
        try!(self.sock.send_str(WORK_COMPLETE, zmq::SNDMORE));
        try!(self.sock.send(&*job.write_to_bytes().unwrap(), 0));
        Ok(())
    }
}

fn studio_cmd() -> String {
    match PackageInstall::load(&STUDIO_PKG, None) {
        Ok(package) => format!("{}/hab-studio", package.paths().unwrap()[0].display()),
        Err(_) => {
            panic!("core/hab-studio not found! This should be available as it is a runtime \
                    dependency in the worker's plan.sh and also present in our dev Dockerfile")
        }
    }
}
<commit_msg>Copy key for appropriate origin in build studio<commit_after>// Copyright (c) 2016 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod logger;
pub mod workspace;
pub mod postprocessor;

pub use protocol::jobsrv::JobState;

use std::ffi::OsString;
use std::fs;
use std::ops::{Deref, DerefMut};
use std::path::Path;
use std::process::{Command, Stdio};
use std::str::FromStr;
use std::sync::{mpsc, Arc, RwLock};
use std::thread::{self, JoinHandle};

use hab_core::package::archive::PackageArchive;
use hab_core::package::install::PackageInstall;
use hab_core::package::PackageIdent;
use hab_net::server::ZMQ_CONTEXT;
use protobuf::{parse_from_bytes, Message};
use protocol::jobsrv as proto;
use zmq;

use self::logger::Logger;
use self::postprocessor::PostProcessor;
use self::workspace::Workspace;
use config::Config;
use error::{Error, Result};
use vcs;

/// In-memory zmq address of Job RunnerMgr
const INPROC_ADDR: &'static str = "inproc://runner";
/// Protocol message to indicate the Job Runner has received a work request
const WORK_ACK: &'static str = "A";
/// Protocol message to indicate the Job Runner has completed a work request
const WORK_COMPLETE: &'static str = "C";

lazy_static! {
    // JW TODO: expose public API functions in the core crate to check if the Rust proccess which
    // is currently executing is, itself, packaged by Habitat. If so, then we should expose another
    // public API function to load the PackageInstall of a dep for the given `origin`/`name`
    // combination. If we can't, then we should fall back to the latest of core/hab-studio because
    // that means we're just in a dev shell.
    static ref STUDIO_PKG: PackageIdent = PackageIdent::from_str("core/hab-studio").unwrap();
}

#[derive(Debug)]
pub struct Job(proto::Job);

impl Job {
    pub fn new(job: proto::Job) -> Self {
        Job(job)
    }

    pub fn vcs(&self) -> &vcs::RemoteSource {
        if self.0.get_project().has_git() {
            return self.0.get_project().get_git();
        }
        unreachable!("unknown vcs associated with job's project");
    }

    pub fn origin(&self) -> &str {
        let items = self.0.get_project().get_id().split("/").collect::<Vec<&str>>();
        assert!(items.len() == 2,
                format!("Invalid project identifier - {}",
                        self.0.get_project().get_id()));
        items[0]
    }
}

impl Deref for Job {
    type Target = proto::Job;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Job {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

pub struct Runner {
    workspace: Workspace,
    auth_token: String,
    logger: Option<Logger>,
}

impl Runner {
    pub fn new(job: Job, config: &Config) -> Self {
        Runner {
            auth_token: config.auth_token.clone(),
            workspace: Workspace::new(config.data_path.clone(), job),
            logger: None,
        }
    }

    pub fn job(&self) -> &Job {
        &self.workspace.job
    }

    pub fn job_mut(&mut self) -> &mut Job {
        &mut self.workspace.job
    }

    pub fn logger(&mut self) -> &mut Logger {
        self.logger.as_mut().expect("logger not initialized")
    }

    pub fn run(mut self) -> Job {
        if let Some(err) = self.setup().err() {
            error!("WORKSPACE SETUP ERR={:?}", err);
            return self.fail();
        }
        // JW TODO: How are we going to get the secret keys for this thing?
        if let Some(err) = self.job().vcs().clone(&self.workspace.src()).err() {
            error!("CLONE ERROR={}", err);
            return self.fail();
        }
        let mut archive = match self.build() {
            Ok(archive) => archive,
            Err(err) => {
                error!("STUDIO ERR={}", err);
                return self.fail();
            }
        };

        let mut post_processor = PostProcessor::new(&self.workspace);

        if !post_processor.run(&mut archive, &self.auth_token) {
            // JW TODO: We should shelve the built artifacts and allow a retry on post-processing.
            // If the job is killed then we can kill the shelved artifacts.
            return self.fail();
        }

        if let Some(err) = fs::remove_dir_all(self.workspace.out()).err() {
            error!("unable to remove out directory ({}), ERR={:?}",
                   self.workspace.out().display(),
                   err)
        }
        self.complete()
    }

    fn build(&mut self) -> Result<PackageArchive> {
        let args = vec![OsString::from("-s"),
                        OsString::from(self.workspace.src()),
                        OsString::from("-r"),
                        OsString::from(self.workspace.studio()),
                        OsString::from("-k"),
                        OsString::from(self.job().origin()),
                        OsString::from("build"),
                        OsString::from(Path::new(self.job().get_project().get_plan_path())
                            .parent()
                            .unwrap())];
        let command = studio_cmd();
        debug!("building, cmd={:?}, args={:?}", command, args);
        let mut child = Command::new(command)
            .args(&args)
            .env_clear()
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .expect("failed to spawn child");
        self.logger().pipe(&mut child);
        let exit_status = child.wait().expect("failed to wait on child");
        debug!("build complete, status={:?}", exit_status);
        if exit_status.success() {
            try!(fs::rename(self.workspace.src().join("results"), self.workspace.out()));
            self.workspace.last_built()
        } else {
            Err(Error::BuildFailure(exit_status.code().unwrap_or(-1)))
        }
    }

    fn complete(mut self) -> Job {
        self.teardown().err().map(|e| error!("{}", e));
        self.workspace.job.set_state(JobState::Complete);
        self.workspace.job
    }

    fn fail(mut self) -> Job {
        self.teardown().err().map(|e| error!("{}", e));
        self.workspace.job.set_state(JobState::Failed);
        self.workspace.job
    }

    fn post_process(&self, archive: &mut PackageArchive) -> bool {
        // JW TODO: In the future we'll support multiple and configurable post processors, but for
        // now let's just publish to the public depot
        //
        // Things to solve right now
        // * Where do we get the token for authentication?
        //      * Should the workers ask for a lease from the JobSrv?
        let client =
            depot_client::Client::new(&hab_core::url::default_depot_url(), PRODUCT, VERSION, None)
                .unwrap();
        if let Some(err) = client.x_put_package(archive, &self.auth_token).err() {
            error!("post processing error, ERR={:?}", err);
        }
        if let Some(err) = fs::remove_dir_all(self.workspace.out()).err() {
            error!("unable to remove out directory ({}), ERR={:?}",
                   self.workspace.out().display(),
                   err)
        }
        true
    }

    fn setup(&mut self) -> Result<()> {
        if let Some(err) = fs::create_dir_all(self.workspace.src()).err() {
            return Err(Error::WorkspaceSetup(format!("{}", self.workspace.src().display()), err));
        }
        self.logger = Some(Logger::init(&self.workspace));
        Ok(())
    }

    fn teardown(&mut self) -> Result<()> {
        let args = vec![OsString::from("-s"),
                        OsString::from(self.workspace.src()),
                        OsString::from("-r"),
                        OsString::from(self.workspace.studio()),
                        OsString::from("rm"),
                        OsString::from(Path::new(self.job().get_project().get_plan_path())
                            .parent()
                            .unwrap())];
        let command = studio_cmd();
        debug!("removing studio, cmd={:?}, args={:?}", command, args);
        let mut child = Command::new(command)
            .args(&args)
            .env_clear()
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .expect("failed to spawn child");
        self.logger().pipe(&mut child);
        let exit_status = child.wait().expect("failed to wait on child");
        debug!("studio removal complete, status={:?}", exit_status);
        if exit_status.success() {
            if let Some(err) = fs::remove_dir_all(self.workspace.src()).err() {
                return Err(Error::WorkspaceTeardown(format!("{}", self.workspace.src().display()),
                                                    err));
            }
            Ok(())
        } else {
            Err(Error::BuildFailure(exit_status.code().unwrap_or(-1)))
        }
    }
}

/// Client for sending and receiving messages to and from the Job Runner
pub struct RunnerCli {
    sock: zmq::Socket,
    msg: zmq::Message,
}

impl RunnerCli {
    /// Create a new Job Runner client
    pub fn new() -> Self {
        let sock = (**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER).unwrap();
        RunnerCli {
            sock: sock,
            msg: zmq::Message::new().unwrap(),
        }
    }

    /// Return a poll item used in `zmq::poll` for awaiting messages on multiple sockets
    pub fn as_poll_item<'a>(&'a self, events: i16) -> zmq::PollItem<'a> {
        self.sock.as_poll_item(events)
    }

    /// Connect to the Job Runner
    pub fn connect(&mut self) -> Result<()> {
        try!(self.sock.connect(INPROC_ADDR));
        Ok(())
    }

    /// Wait until client receives a work received acknowledgement by the Runner and return
    /// the assigned JobID.
    pub fn recv_ack(&mut self) -> Result<&zmq::Message> {
        try!(self.sock.recv(&mut self.msg, 0));
        if Some(WORK_ACK) != self.msg.as_str() {
            unreachable!("wk:run:1, received unexpected response from runner");
        }
        try!(self.sock.recv(&mut self.msg, 0));
        Ok(&self.msg)
    }

    /// Wait until client receives a work complete message by the Runner and return an encoded
    /// representation of the job.
    pub fn recv_complete(&mut self) -> Result<&zmq::Message> {
        try!(self.sock.recv(&mut self.msg, 0));
        if Some(WORK_COMPLETE) != self.msg.as_str() {
            unreachable!("wk:run:2, received unexpected response from runner");
        }
        try!(self.sock.recv(&mut self.msg, 0));
        Ok(&self.msg)
    }

    /// Send a message to the Job Runner
    pub fn send(&mut self, msg: &zmq::Message) -> Result<()> {
        try!(self.sock.send(&*msg, 0));
        Ok(())
    }
}

/// Receives work notifications from a `RunnerCli` and performs long-running tasks in a
/// separate thread.
pub struct RunnerMgr {
    sock: zmq::Socket,
    msg: zmq::Message,
    config: Arc<RwLock<Config>>,
}

impl RunnerMgr {
    /// Start the Job Runner
    pub fn start(config: Arc<RwLock<Config>>) -> Result<JoinHandle<()>> {
        let (tx, rx) = mpsc::sync_channel(0);
        let handle = thread::Builder::new()
            .name("runner".to_string())
            .spawn(move || {
                let mut runner = Self::new(config).unwrap();
                runner.run(tx).unwrap();
            })
            .unwrap();
        match rx.recv() {
            Ok(()) => Ok(handle),
            Err(e) => panic!("runner thread startup error, err={}", e),
        }
    }

    fn new(config: Arc<RwLock<Config>>) -> Result<Self> {
        let sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER));
        Ok(RunnerMgr {
            sock: sock,
            msg: zmq::Message::new().unwrap(),
            config: config,
        })
    }

    // Main loop for server
    fn run(&mut self, rz: mpsc::SyncSender<()>) -> Result<()> {
        try!(self.sock.bind(INPROC_ADDR));
        rz.send(()).unwrap();
        loop {
            let job = try!(self.recv_job());
            try!(self.send_ack(&job));
            try!(self.execute_job(job));
        }
    }

    fn execute_job(&mut self, job: Job) -> Result<()> {
        let runner = {
            Runner::new(job, &self.config.read().unwrap())
        };
        debug!("executing work, job={:?}", runner.job());
        let job = runner.run();
        self.send_complete(&job)
    }

    fn recv_job(&mut self) -> Result<Job> {
        try!(self.sock.recv(&mut self.msg, 0));
        let job: proto::Job = parse_from_bytes(&self.msg).unwrap();
        Ok(Job::new(job))
    }

    fn send_ack(&mut self, job: &Job) -> Result<()> {
        debug!("received work, job={:?}", job);
        try!(self.sock.send_str(WORK_ACK, zmq::SNDMORE));
        try!(self.sock.send(&*job.write_to_bytes().unwrap(), 0));
        Ok(())
    }

    fn send_complete(&mut self, job: &Job) -> Result<()> {
        debug!("work complete, job={:?}", job);
        try!(self.sock.send_str(WORK_COMPLETE, zmq::SNDMORE));
        try!(self.sock.send(&*job.write_to_bytes().unwrap(), 0));
        Ok(())
    }
}

fn studio_cmd() -> String {
    match PackageInstall::load(&STUDIO_PKG, None) {
        Ok(package) => format!("{}/hab-studio", package.paths().unwrap()[0].display()),
        Err(_) => {
            panic!("core/hab-studio not found! This should be available as it is a runtime \
                    dependency in the worker's plan.sh and also present in our dev Dockerfile")
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use protocol::{jobsrv, vault};

    #[test]
    fn extract_origin_from_job() {
        let mut inner = jobsrv::Job::new();
        let mut project = vault::Project::new();
        project.set_id("core/nginx".to_string());
        inner.set_project(project);
        let job = Job::new(inner);
        assert_eq!(job.origin(), "core");
    }
}
<|endoftext|>


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

use std::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{CACHE_ARTIFACT_PATH, CACHE_KEY_PATH, find_command};
use hcore::os::process;
use command::studio::enter::ARTIFACT_PATH_ENVVAR;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;

    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            mnt_prefix,
            CACHE_KEY_PATH
        ),
    ];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!(
            "{}:{}/{}",
            cache_artifact_path,
            mnt_prefix,
            CACHE_ARTIFACT_PATH
        ));
    }
    if !is_serving_windows_containers(&docker_cmd) &&
        (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let env_vars = vec![
        "HAB_AUTH_TOKEN",
        "HAB_BLDR_URL",
        "HAB_BLDR_CHANNEL",
        "HAB_ORIGIN",
        "HAB_ORIGIN_KEYS",
        "HAB_STUDIO_BACKLINE_PKG",
        "HAB_STUDIO_NOSTUDIORC",
        "HAB_STUDIO_SUP",
        "HAB_UPDATE_STRATEGY_FREQUENCY_MS",
        "http_proxy",
        "https_proxy",
        "RUST_LOG",
    ];
    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    check_mounts(&docker_cmd, volumes.iter())?;
    run_container(docker_cmd, args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images").arg(&image_identifier(docker_cmd)).arg(
        "-q",
    );
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier(docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
        .arg(&image)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd.spawn()
        .expect("Docker command failed to spawn")
        .wait_with_output()
        .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!(
            "Pulling Docker image '{}' failed with exit code: {:?}",
            &image,
            result.status
        );

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(
                image_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).to_string(),
            ));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }

    if !is_serving_windows_containers(docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    cmd_args.push(image_identifier(docker_cmd).into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd).args(&cmd_args).output().expect(
        "docker failed to start",
    );

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty() &&
        (stderr.as_str().contains("Mounts denied") ||
             stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(
    docker_cmd: PathBuf,
    args: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    J: IntoIterator<Item = T>,
    S: AsRef<OsStr>,
    T: AsRef<str>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    if !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            cmd_args.push("--tty".into());
            cmd_args.push("--interactive".into());
        }
    }
    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
        if !opts.is_empty() {
            debug!(
                "Adding extra Docker options from {} = {:?}",
                DOCKER_OPTS_ENVVAR,
                opts
            );
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }
    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd, cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<commit_msg>migrate DEBUG env var to docker studio<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{CACHE_ARTIFACT_PATH, CACHE_KEY_PATH, find_command};
use hcore::os::process;
use command::studio::enter::ARTIFACT_PATH_ENVVAR;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;

    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            mnt_prefix,
            CACHE_KEY_PATH
        ),
    ];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!(
            "{}:{}/{}",
            cache_artifact_path,
            mnt_prefix,
            CACHE_ARTIFACT_PATH
        ));
    }
    if !is_serving_windows_containers(&docker_cmd) &&
        (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let env_vars = vec![
        "DEBUG",
        "HAB_AUTH_TOKEN",
        "HAB_BLDR_URL",
        "HAB_BLDR_CHANNEL",
        "HAB_ORIGIN",
        "HAB_ORIGIN_KEYS",
        "HAB_STUDIO_BACKLINE_PKG",
        "HAB_STUDIO_NOSTUDIORC",
        "HAB_STUDIO_SUP",
        "HAB_UPDATE_STRATEGY_FREQUENCY_MS",
        "http_proxy",
        "https_proxy",
        "RUST_LOG",
    ];
    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    check_mounts(&docker_cmd, volumes.iter())?;
    run_container(docker_cmd, args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images").arg(&image_identifier(docker_cmd)).arg(
        "-q",
    );
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier(docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
        .arg(&image)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd.spawn()
        .expect("Docker command failed to spawn")
        .wait_with_output()
        .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!(
            "Pulling Docker image '{}' failed with exit code: {:?}",
            &image,
            result.status
        );

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(
                image_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).to_string(),
            ));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }

    if !is_serving_windows_containers(docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    cmd_args.push(image_identifier(docker_cmd).into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd).args(&cmd_args).output().expect(
        "docker failed to start",
    );

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty() &&
        (stderr.as_str().contains("Mounts denied") ||
             stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(
    docker_cmd: PathBuf,
    args: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    J: IntoIterator<Item = T>,
    S: AsRef<OsStr>,
    T: AsRef<str>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    if !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            cmd_args.push("--tty".into());
            cmd_args.push("--interactive".into());
        }
    }
    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
        if !opts.is_empty() {
            debug!(
                "Adding extra Docker options from {} = {:?}",
                DOCKER_OPTS_ENVVAR,
                opts
            );
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }
    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd, cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<|endoftext|>


<commit_before>use future::{Future, IntoFuture};
use stream::Stream;
use poll::Poll;
use Async;
use stack::{Stack, Drain};
use std::sync::Arc;
use task::{self, UnparkEvent};

use std::prelude::v1::*;

/// An adaptor for a stream of futures to execute the futures concurrently, if
/// possible, delivering results as they become available.
///
/// This adaptor will return their results in the order that they complete.
/// This is created by the `futures` method.
///
#[must_use = "streams do nothing unless polled"]
pub struct FuturesUnordered<F>
    where F: Future
{
    futures: Vec<Option<F>>,
    stack: Arc<Stack<usize>>,
    pending: Option<Drain<usize>>,
    active: usize,
}

/// Converts a `Vec` of `IntoFuture` elements into a `Stream` over the underlying
/// futures. The created stream will yield into values produced by those futures.
/// Similar to `buffer_unordered` the order of the produced items will not be equal
/// with the order of elements in the input vector.
///
pub fn futures_unordered<I>(futures: I) -> FuturesUnordered<<I::Item as IntoFuture>::Future>
    where I: IntoIterator,
          I::Item: IntoFuture
{
    let futures = futures.into_iter().map(IntoFuture::into_future).map(Some).collect::<Vec<_>>();
    let stack = Arc::new(Stack::new());
    for i in 0..futures.len() {
        stack.push(i);
    }
    FuturesUnordered {
        active: futures.len(),
        futures: futures,
        pending: None,
        stack: stack,
    }
}

impl<F> FuturesUnordered<F>
    where F: Future
{
    fn poll_pending(&mut self, mut drain: Drain<usize>) -> Option<Poll<Option<F::Item>, F::Error>> {
        while let Some(id) = drain.next() {
            let event = UnparkEvent::new(self.stack.clone(), id);
            let ret = match task::with_unpark_event(event, || {
                self.futures[id]
                    .as_mut()
                    .unwrap()
                    .poll()
            }) {
                Ok(Async::NotReady) => {
                    continue;
                }
                Ok(Async::Ready(val)) => Ok(Async::Ready(Some(val))),
                Err(e) => Err(e),
            };
            self.pending = Some(drain);
            self.active -= 1;
            self.futures[id] = None;
            return Some(ret);
        }
        None
    }
}

impl<F> Stream for FuturesUnordered<F>
    where F: Future
{
    type Item = F::Item;
    type Error = F::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.active == 0 {
            return Ok(Async::Ready(None));
        }
        let drain = self.pending
            .take()
            .unwrap_or_else(|| self.stack.drain());
        if let Some(ret) = self.poll_pending(drain){
            return ret;
        }
        let drain = self.stack.drain();
        if let Some(ret) = self.poll_pending(drain){
            return ret;
        }
        Ok(Async::NotReady)
    }
}
<commit_msg>Touch up future_unordered slightly<commit_after>use future::{Future, IntoFuture};
use stream::Stream;
use poll::Poll;
use Async;
use stack::{Stack, Drain};
use std::sync::Arc;
use task::{self, UnparkEvent};

use std::prelude::v1::*;

/// An adaptor for a stream of futures to execute the futures concurrently, if
/// possible, delivering results as they become available.
///
/// This adaptor will return their results in the order that they complete.
/// This is created by the `futures` method.
///
#[must_use = "streams do nothing unless polled"]
pub struct FuturesUnordered<F>
    where F: Future
{
    futures: Vec<Option<F>>,
    stack: Arc<Stack<usize>>,
    pending: Option<Drain<usize>>,
    active: usize,
}

/// Converts a list of futures into a `Stream` of results from the futures.
///
/// This function will take an list of futures (e.g. a vector, an iterator,
/// etc), and return a stream. The stream will yield items as they become
/// available on the futures internally, in the order that they become
/// available. This function is similar to `buffer_unordered` in that it may
/// return items in a different order than in the list specified.
pub fn futures_unordered<I>(futures: I) -> FuturesUnordered<<I::Item as IntoFuture>::Future>
    where I: IntoIterator,
          I::Item: IntoFuture
{
    let futures = futures.into_iter()
                         .map(IntoFuture::into_future)
                         .map(Some)
                         .collect::<Vec<_>>();
    let stack = Arc::new(Stack::new());
    for i in 0..futures.len() {
        stack.push(i);
    }
    FuturesUnordered {
        active: futures.len(),
        futures: futures,
        pending: None,
        stack: stack,
    }
}

impl<F> FuturesUnordered<F>
    where F: Future
{
    fn poll_pending(&mut self, mut drain: Drain<usize>)
                    -> Option<Poll<Option<F::Item>, F::Error>> {
        while let Some(id) = drain.next() {
            let event = UnparkEvent::new(self.stack.clone(), id);
            let ret = match task::with_unpark_event(event, || {
                self.futures[id]
                    .as_mut()
                    .unwrap()
                    .poll()
            }) {
                Ok(Async::NotReady) => continue,
                Ok(Async::Ready(val)) => Ok(Async::Ready(Some(val))),
                Err(e) => Err(e),
            };
            self.pending = Some(drain);
            self.active -= 1;
            self.futures[id] = None;
            return Some(ret)
        }
        None
    }
}

impl<F> Stream for FuturesUnordered<F>
    where F: Future
{
    type Item = F::Item;
    type Error = F::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if self.active == 0 {
            return Ok(Async::Ready(None))
        }
        if let Some(drain) = self.pending.take() {
            if let Some(ret) = self.poll_pending(drain){
                return ret
            }
        }
        let drain = self.stack.drain();
        if let Some(ret) = self.poll_pending(drain) {
            return ret
        }
        assert!(self.active > 0);
        Ok(Async::NotReady)
    }
}
<|endoftext|>


<commit_before>extern mod extra;
extern mod common;

use std::{uint, vec, io, os};
use extra::{time, term};
use extra::term::Terminal;
use common::problem::Problem;
use common::extiter::Range;

mod problem;

static NSEC_PER_SEC: u64 = 1000000000;
fn nanosec_to_str(nsec: u64) -> ~str {
    return fmt!("%u.%09u",
         (nsec / NSEC_PER_SEC) as uint,
         (nsec % NSEC_PER_SEC) as uint);
}

fn bench<T>(f: &fn() -> T) -> (u64, T) {
    let start_time = time::precise_time_ns();
    let result     = f();
    let end_time   = time::precise_time_ns();
    return (end_time - start_time, result);
}

fn color_print(writer: @io::Writer, color: u8, s: &str) {
    let term = Terminal::new(writer);
    term.iter(|&t| t.fg(color));
    print(s);
    term.iter(|&t| t.reset());
}

fn print_result(correct: bool, name: &str, time: u64, comp_answer: &str) {
    print("[");
    if correct {
        color_print(io::stdout(), term::color_green, "OK");
    } else {
        color_print(io::stdout(), term::color_red, "NG");
    }
    println(fmt!("] %5s %13s %20s", name, nanosec_to_str(time), comp_answer));
}

struct ArgIterator<'self> {
    args: &'self [~str],
    idx: uint,
    cur_range: Range<uint>,
}

impl<'self> ArgIterator<'self> {
    pub fn new<'a>(args: &'a [~str]) -> ArgIterator<'a> {
        let mut it = ArgIterator { args: args, idx: 0, cur_range: Range::new(0u, 0) };
        it.update_range();
        return it;
    }

    priv fn update_range(&mut self) {
        self.cur_range = Range::new(0u, 0);
        if self.idx >= self.args.len() { return; }

        if !self.args[self.idx].contains_char('-') {
            let n = uint::from_str(self.args[self.idx]);
            for n.iter().advance |&n| {
                self.cur_range = Range::new(n, n + 1);
            }
            return;
        }

        let mut ns = ~[];
        for self.args[self.idx].split_iter('-').advance |ss| {
            match uint::from_str(ss) {
                Some(n) => { ns.push(n); }
                None    => { return; }
            }
        }
        if ns.len() > 2 { return; }
        self.cur_range = Range::new(ns[0], ns[1] + 1);
        println(fmt!("%?", ns));
    }
}

impl<'self> Iterator<uint> for ArgIterator<'self> {
    pub fn next(&mut self) -> Option<uint> {
        loop {
            if self.idx >= self.args.len() { return None; }
            match self.cur_range.next() {
                Some(n) => { return Some(n); }
                None    => {
                    self.idx += 1;
                    self.update_range();
                }
            }
        }
    }
}

fn main() {
    let args = os::args();
    let mut it = if args.tail().is_empty() {
        @Range::new(0, problem::problems.len())
            .transform(|i| problem::problems[i])
            as @Iterator<&Problem>
    } else {
        @ArgIterator::new(args.tail()).filter_map(|n| {
            vec::bsearch(problem::problems, |&p| p.id.cmp(&n))
                .map(|&i| problem::problems[i])
        }) as @Iterator<&Problem>
    };

    let mut total_time  = 0;
    let mut solve_cnt   = 0;
    let mut all_correct = true;

    loop {
        match it.next() {
            Some(&p) => {
                let (time, answer) = do bench { (p.solver)() };
                let correct = p.answer == answer;
                print_result(correct, p.id.to_str(), time, answer);

                total_time += time;
                solve_cnt  += 1;
                all_correct &= correct;
            }
            None => { break; }
        }
    }

    if solve_cnt > 1 {
        print_result(all_correct, "TOTAL", total_time, "");
    }
}
<commit_msg>Modify euler::main<commit_after>extern mod extra;
extern mod common;

use std::{uint, vec, io, os};
use extra::{time, term};
use extra::term::Terminal;
use common::problem::Problem;
use common::extiter::Range;

mod problem;

static NSEC_PER_SEC: u64 = 1000000000;
fn nanosec_to_str(nsec: u64) -> ~str {
    return fmt!("%u.%09u",
         (nsec / NSEC_PER_SEC) as uint,
         (nsec % NSEC_PER_SEC) as uint);
}

fn bench<T>(f: &fn() -> T) -> (u64, T) {
    let start_time = time::precise_time_ns();
    let result     = f();
    let end_time   = time::precise_time_ns();
    return (end_time - start_time, result);
}

fn color_print(writer: @io::Writer, color: term::color::Color, s: &str) {
    let term = Terminal::new(writer);
    term.iter(|&t| t.fg(color));
    print(s);
    term.iter(|&t| t.reset());
}

fn print_result(correct: bool, name: &str, time: u64, comp_answer: &str) {
    print("[");
    if correct {
        color_print(io::stdout(), term::color::green, "OK");
    } else {
        color_print(io::stdout(), term::color::red, "NG");
    }
    println(fmt!("] %5s %13s %20s", name, nanosec_to_str(time), comp_answer));
}

struct ArgIterator<'self> {
    args: &'self [~str],
    idx: uint,
    cur_range: Range<uint>,
}

impl<'self> ArgIterator<'self> {
    pub fn new<'a>(args: &'a [~str]) -> ArgIterator<'a> {
        let mut it = ArgIterator { args: args, idx: 0, cur_range: Range::new(0u, 0) };
        it.update_range();
        return it;
    }

    priv fn update_range(&mut self) {
        self.cur_range = Range::new(0u, 0);
        if self.idx >= self.args.len() { return; }

        if !self.args[self.idx].contains_char('-') {
            let n = uint::from_str(self.args[self.idx]);
            for n.iter().advance |&n| {
                self.cur_range = Range::new(n, n + 1);
            }
            return;
        }

        let mut ns = ~[];
        for self.args[self.idx].split_iter('-').advance |ss| {
            match uint::from_str(ss) {
                Some(n) => { ns.push(n); }
                None    => { return; }
            }
        }
        if ns.len() > 2 { return; }
        self.cur_range = Range::new(ns[0], ns[1] + 1);
        println(fmt!("%?", ns));
    }
}

impl<'self> Iterator<uint> for ArgIterator<'self> {
    pub fn next(&mut self) -> Option<uint> {
        loop {
            if self.idx >= self.args.len() { return None; }
            match self.cur_range.next() {
                Some(n) => { return Some(n); }
                None    => {
                    self.idx += 1;
                    self.update_range();
                }
            }
        }
    }
}

fn solve_all<T: Iterator<&'static Problem<'static>>>(mut it: T) {
    let mut total_time  = 0;
    let mut solve_cnt   = 0;
    let mut all_correct = true;

    for it.advance() |p| {
        let (time, answer) = do bench { (p.solver)() };
        let correct = p.answer == answer;
        print_result(correct, p.id.to_str(), time, answer);

        total_time += time;
        solve_cnt  += 1;
        all_correct &= correct;
    }

    if solve_cnt > 1 {
        print_result(all_correct, "TOTAL", total_time, "");
    }
}

fn main() {
    let args = os::args();
    let args = args.tail();

    if args.is_empty() {
        solve_all(Range::new(0, problem::problems.len())
                  .transform(|i| problem::problems[i]))
    } else {
        solve_all(ArgIterator::new(args)
                  .filter_map(|n| vec::bsearch(problem::problems, |&p| p.id.cmp(&n)))
                  .transform(|i| problem::problems[i]));
    };
}
<|endoftext|>