<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use backup_service::start_backup_service;
use consensus::{consensus_provider::start_consensus, gen_consensus_reconfig_subscription};
use debug_interface::node_debug_service::NodeDebugService;
use executor::{db_bootstrapper::maybe_bootstrap, Executor};
use executor_types::ChunkExecutor;
use futures::{channel::mpsc::channel, executor::block_on};
use libra_config::{
    config::{NetworkConfig, NodeConfig, RoleType},
    network_id::NodeNetworkId,
    utils::get_genesis_txn,
};
use libra_json_rpc::bootstrap_from_config as bootstrap_rpc;
use libra_logger::{prelude::*, Logger};
use libra_mempool::gen_mempool_reconfig_subscription;
use libra_metrics::metric_server;
use libra_types::{
    account_config::libra_root_address, account_state::AccountState, chain_id::ChainId,
    move_resource::MoveStorage,
};
use libra_vm::LibraVM;
use libradb::LibraDB;
use network_builder::builder::NetworkBuilder;
use state_synchronizer::StateSynchronizer;
use std::{boxed::Box, convert::TryFrom, net::ToSocketAddrs, sync::Arc, thread, time::Instant};
use storage_interface::DbReaderWriter;
use storage_service::start_storage_service_with_db;
use tokio::runtime::{Builder, Runtime};

const AC_SMP_CHANNEL_BUFFER_SIZE: usize = 1_024;
const INTRA_NODE_CHANNEL_BUFFER_SIZE: usize = 1;

pub struct LibraHandle {
    _rpc: Runtime,
    _mempool: Runtime,
    _state_synchronizer: StateSynchronizer,
    network_runtimes: Vec<Runtime>,
    _consensus_runtime: Option<Runtime>,
    _debug: NodeDebugService,
    _backup: Runtime,
}

impl LibraHandle {
    pub fn shutdown(&mut self) {
        // Shutdown network runtimes to avoid panic error log after LibraHandle is dropped:
        // thread ‘network-’ panicked at ‘SelectNextSome polled after terminated’,...
        // stack backtrace:
        //    ......
        //    8: network_simple_onchain_discovery::ConfigurationChangeListener::start::{{closure}}
        //      at network/simple-onchain-discovery/src/lib.rs:175
        //    ......
        // Other runtimes don't have same problem.
        while self.network_runtimes.len() > 0 {
            self.network_runtimes.remove(0).shutdown_background();
        }
    }
}

// Fetch chain ID from on-chain resource
fn fetch_chain_id(db: &DbReaderWriter) -> ChainId {
    let blob = db
        .reader
        .get_account_state_with_proof_by_version(
            libra_root_address(),
            (&*db.reader)
                .fetch_synced_version()
                .expect("[libra-node] failed fetching synced version."),
        )
        .expect("[libra-node] failed to get Libra root address account state")
        .0
        .expect("[libra-node] missing Libra root address account state");
    AccountState::try_from(&blob)
        .expect("[libra-node] failed to convert blob to account state")
        .get_chain_id_resource()
        .expect("[libra-node] failed to get chain ID resource")
        .expect("[libra-node] missing chain ID resource")
        .chain_id()
}

fn setup_chunk_executor(db: DbReaderWriter) -> Box<dyn ChunkExecutor> {
    Box::new(Executor::<LibraVM>::new(db))
}

fn setup_debug_interface(config: &NodeConfig, logger: Option<Arc<Logger>>) -> NodeDebugService {
    let addr = format!(
        "{}:{}",
        config.debug_interface.address, config.debug_interface.admission_control_node_debug_port,
    )
    .to_socket_addrs()
    .unwrap()
    .next()
    .unwrap();

    libra_trace::set_libra_trace(&config.debug_interface.libra_trace.sampling)
        .expect("Failed to set libra trace sampling rate.");

    NodeDebugService::new(addr, logger)
}

pub fn setup_environment(node_config: &NodeConfig, logger: Option<Arc<Logger>>) -> LibraHandle {
    let debug_if = setup_debug_interface(&node_config, logger);

    let metrics_port = node_config.debug_interface.metrics_server_port;
    let metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || metric_server::start_server(metric_host, metrics_port, false));
    let public_metrics_port = node_config.debug_interface.public_metrics_server_port;
    let public_metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || {
        metric_server::start_server(public_metric_host, public_metrics_port, true)
    });

    let mut instant = Instant::now();
    let (libra_db, db_rw) = DbReaderWriter::wrap(
        LibraDB::open(
            &node_config.storage.dir(),
            false, /* readonly */
            node_config.storage.prune_window,
        )
        .expect("DB should open."),
    );
    let _simple_storage_service =
        start_storage_service_with_db(&node_config, Arc::clone(&libra_db));
    let backup_service = start_backup_service(
        node_config.storage.backup_service_address,
        Arc::clone(&libra_db),
    );

    let waypoint = node_config.base.waypoint.waypoint();
    println!("WWWW1111: {}", waypoint);
    // if there's genesis txn and waypoint, commit it if the result matches.
    if let Some(genesis) = get_genesis_txn(&node_config) {
        maybe_bootstrap::<LibraVM>(&db_rw, genesis, waypoint)
            .expect("Db-bootstrapper should not fail.");
    }

    debug!(
        "Storage service started in {} ms",
        instant.elapsed().as_millis()
    );

    instant = Instant::now();
    let chunk_executor = setup_chunk_executor(db_rw.clone());
    debug!(
        "ChunkExecutor setup in {} ms",
        instant.elapsed().as_millis()
    );
    let chain_id = fetch_chain_id(&db_rw);
    let mut network_runtimes = vec![];
    let mut state_sync_network_handles = vec![];
    let mut mempool_network_handles = vec![];
    let mut consensus_network_handles = None;
    let mut reconfig_subscriptions = vec![];

    let (mempool_reconfig_subscription, mempool_reconfig_events) =
        gen_mempool_reconfig_subscription();
    reconfig_subscriptions.push(mempool_reconfig_subscription);
    // consensus has to subscribe to ALL on-chain configs
    let (consensus_reconfig_subscription, consensus_reconfig_events) =
        gen_consensus_reconfig_subscription();
    if node_config.base.role.is_validator() {
        reconfig_subscriptions.push(consensus_reconfig_subscription);
    }

    // Gather all network configs into a single vector.
    // TODO:  consider explicitly encoding the role in the NetworkConfig
    let mut network_configs: Vec<(RoleType, &NetworkConfig)> = node_config
        .full_node_networks
        .iter()
        .map(|network_config| (RoleType::FullNode, network_config))
        .collect();
    if let Some(network_config) = node_config.validator_network.as_ref() {
        network_configs.push((RoleType::Validator, network_config));
    }

    let mut network_builders = Vec::new();

    // Instantiate every network and collect the requisite endpoints for state_sync, mempool, and consensus.
    for (idx, (role, network_config)) in network_configs.into_iter().enumerate() {
        // Perform common instantiation steps
        let mut network_builder = NetworkBuilder::create(chain_id, role, network_config);
        let network_id = network_config.network_id.clone();

        // Create the endpoints to connect the Network to StateSynchronizer.
        let (state_sync_sender, state_sync_events) = network_builder
            .add_protocol_handler(state_synchronizer::network::network_endpoint_config());
        state_sync_network_handles.push((
            NodeNetworkId::new(network_id.clone(), idx),
            state_sync_sender,
            state_sync_events,
        ));

        // Create the endpoints to connect the Network to mempool.
        let (mempool_sender, mempool_events) =
            network_builder.add_protocol_handler(libra_mempool::network::network_endpoint_config(
                // TODO:  Make this configuration option more clear.
                node_config.mempool.max_broadcasts_per_peer,
            ));
        mempool_network_handles.push((
            NodeNetworkId::new(network_id, idx),
            mempool_sender,
            mempool_events,
        ));

        match role {
            // Perform steps relevant specifically to Validator networks.
            RoleType::Validator => {
                // A valid config is allowed to have at most one ValidatorNetwork
                // TODO:  `expect_none` would be perfect here, once it is stable.
                if consensus_network_handles.is_some() {
                    panic!("There can be at most one validator network!");
                }

                consensus_network_handles =
                    Some(network_builder.add_protocol_handler(
                        consensus::network_interface::network_endpoint_config(),
                    ));
            }
            // Currently no FullNode network specific steps.
            RoleType::FullNode => (),
        }

        reconfig_subscriptions.append(network_builder.reconfig_subscriptions());

        network_builders.push(network_builder);
    }

    // Build the configured networks.
    for network_builder in &mut network_builders {
        debug!("Creating runtime for {}", network_builder.network_context());
        let runtime = Builder::new()
            .thread_name("network-")
            .threaded_scheduler()
            .enable_all()
            .build()
            .expect("Failed to start runtime. Won't be able to start networking.");
        network_builder.build(runtime.handle().clone());
        network_runtimes.push(runtime);
        debug!(
            "Network built for network context: {}",
            network_builder.network_context()
        );
    }

    // Start the configured networks.
    // TODO:  Collect all component starts at the end of this function
    for network_builder in &mut network_builders {
        network_builder.start();
    }

    // TODO set up on-chain discovery network based on UpstreamConfig.fallback_network
    // and pass network handles to mempool/state sync

    // for state sync to send requests to mempool
    let (state_sync_to_mempool_sender, state_sync_requests) =
        channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);
    let state_synchronizer = StateSynchronizer::bootstrap(
        state_sync_network_handles,
        state_sync_to_mempool_sender,
        Arc::clone(&db_rw.reader),
        chunk_executor,
        &node_config,
        waypoint,
        reconfig_subscriptions,
    );
    let (mp_client_sender, mp_client_events) = channel(AC_SMP_CHANNEL_BUFFER_SIZE);

    let rpc_runtime = bootstrap_rpc(&node_config, chain_id, libra_db.clone(), mp_client_sender);

    let mut consensus_runtime = None;
    let (consensus_to_mempool_sender, consensus_requests) = channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);

    instant = Instant::now();
    let mempool = libra_mempool::bootstrap(
        node_config,
        Arc::clone(&db_rw.reader),
        mempool_network_handles,
        mp_client_events,
        consensus_requests,
        state_sync_requests,
        mempool_reconfig_events,
    );
    debug!("Mempool started in {} ms", instant.elapsed().as_millis());

    // StateSync should be instantiated and started before Consensus to avoid a cyclic dependency:
    // network provider -> consensus -> state synchronizer -> network provider.  This has resulted
    // in a deadlock as observed in GitHub issue #749.
    if let Some((consensus_network_sender, consensus_network_events)) = consensus_network_handles {
        // Make sure that state synchronizer is caught up at least to its waypoint
        // (in case it's present). There is no sense to start consensus prior to that.
        // TODO: Note that we need the networking layer to be able to discover & connect to the
        // peers with potentially outdated network identity public keys.
        debug!("Wait until state synchronizer is initialized");
        block_on(state_synchronizer.wait_until_initialized())
            .expect("State synchronizer initialization failure");
        debug!("State synchronizer initialization complete.");

        // Initialize and start consensus.
        instant = Instant::now();
        consensus_runtime = Some(start_consensus(
            node_config,
            consensus_network_sender,
            consensus_network_events,
            state_synchronizer.create_client(),
            consensus_to_mempool_sender,
            libra_db,
            consensus_reconfig_events,
        ));
        debug!("Consensus started in {} ms", instant.elapsed().as_millis());
    }

    LibraHandle {
        network_runtimes,
        _rpc: rpc_runtime,
        _mempool: mempool,
        _state_synchronizer: state_synchronizer,
        _consensus_runtime: consensus_runtime,
        _debug: debug_if,
        _backup: backup_service,
    }
}
<commit_msg>fix lint<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use backup_service::start_backup_service;
use consensus::{consensus_provider::start_consensus, gen_consensus_reconfig_subscription};
use debug_interface::node_debug_service::NodeDebugService;
use executor::{db_bootstrapper::maybe_bootstrap, Executor};
use executor_types::ChunkExecutor;
use futures::{channel::mpsc::channel, executor::block_on};
use libra_config::{
    config::{NetworkConfig, NodeConfig, RoleType},
    network_id::NodeNetworkId,
    utils::get_genesis_txn,
};
use libra_json_rpc::bootstrap_from_config as bootstrap_rpc;
use libra_logger::{prelude::*, Logger};
use libra_mempool::gen_mempool_reconfig_subscription;
use libra_metrics::metric_server;
use libra_types::{
    account_config::libra_root_address, account_state::AccountState, chain_id::ChainId,
    move_resource::MoveStorage,
};
use libra_vm::LibraVM;
use libradb::LibraDB;
use network_builder::builder::NetworkBuilder;
use state_synchronizer::StateSynchronizer;
use std::{boxed::Box, convert::TryFrom, net::ToSocketAddrs, sync::Arc, thread, time::Instant};
use storage_interface::DbReaderWriter;
use storage_service::start_storage_service_with_db;
use tokio::runtime::{Builder, Runtime};

const AC_SMP_CHANNEL_BUFFER_SIZE: usize = 1_024;
const INTRA_NODE_CHANNEL_BUFFER_SIZE: usize = 1;

pub struct LibraHandle {
    _rpc: Runtime,
    _mempool: Runtime,
    _state_synchronizer: StateSynchronizer,
    network_runtimes: Vec<Runtime>,
    _consensus_runtime: Option<Runtime>,
    _debug: NodeDebugService,
    _backup: Runtime,
}

impl LibraHandle {
    pub fn shutdown(&mut self) {
        // Shutdown network runtimes to avoid panic error log after LibraHandle is dropped:
        // thread ‘network-’ panicked at ‘SelectNextSome polled after terminated’,...
        // stack backtrace:
        //    ......
        //    8: network_simple_onchain_discovery::ConfigurationChangeListener::start::{{closure}}
        //      at network/simple-onchain-discovery/src/lib.rs:175
        //    ......
        // Other runtimes don't have same problem.
        while !self.network_runtimes.is_empty() {
            self.network_runtimes.remove(0).shutdown_background();
        }
    }
}

// Fetch chain ID from on-chain resource
fn fetch_chain_id(db: &DbReaderWriter) -> ChainId {
    let blob = db
        .reader
        .get_account_state_with_proof_by_version(
            libra_root_address(),
            (&*db.reader)
                .fetch_synced_version()
                .expect("[libra-node] failed fetching synced version."),
        )
        .expect("[libra-node] failed to get Libra root address account state")
        .0
        .expect("[libra-node] missing Libra root address account state");
    AccountState::try_from(&blob)
        .expect("[libra-node] failed to convert blob to account state")
        .get_chain_id_resource()
        .expect("[libra-node] failed to get chain ID resource")
        .expect("[libra-node] missing chain ID resource")
        .chain_id()
}

fn setup_chunk_executor(db: DbReaderWriter) -> Box<dyn ChunkExecutor> {
    Box::new(Executor::<LibraVM>::new(db))
}

fn setup_debug_interface(config: &NodeConfig, logger: Option<Arc<Logger>>) -> NodeDebugService {
    let addr = format!(
        "{}:{}",
        config.debug_interface.address, config.debug_interface.admission_control_node_debug_port,
    )
    .to_socket_addrs()
    .unwrap()
    .next()
    .unwrap();

    libra_trace::set_libra_trace(&config.debug_interface.libra_trace.sampling)
        .expect("Failed to set libra trace sampling rate.");

    NodeDebugService::new(addr, logger)
}

pub fn setup_environment(node_config: &NodeConfig, logger: Option<Arc<Logger>>) -> LibraHandle {
    let debug_if = setup_debug_interface(&node_config, logger);

    let metrics_port = node_config.debug_interface.metrics_server_port;
    let metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || metric_server::start_server(metric_host, metrics_port, false));
    let public_metrics_port = node_config.debug_interface.public_metrics_server_port;
    let public_metric_host = node_config.debug_interface.address.clone();
    thread::spawn(move || {
        metric_server::start_server(public_metric_host, public_metrics_port, true)
    });

    let mut instant = Instant::now();
    let (libra_db, db_rw) = DbReaderWriter::wrap(
        LibraDB::open(
            &node_config.storage.dir(),
            false, /* readonly */
            node_config.storage.prune_window,
        )
        .expect("DB should open."),
    );
    let _simple_storage_service =
        start_storage_service_with_db(&node_config, Arc::clone(&libra_db));
    let backup_service = start_backup_service(
        node_config.storage.backup_service_address,
        Arc::clone(&libra_db),
    );

    let waypoint = node_config.base.waypoint.waypoint();
    println!("WWWW1111: {}", waypoint);
    // if there's genesis txn and waypoint, commit it if the result matches.
    if let Some(genesis) = get_genesis_txn(&node_config) {
        maybe_bootstrap::<LibraVM>(&db_rw, genesis, waypoint)
            .expect("Db-bootstrapper should not fail.");
    }

    debug!(
        "Storage service started in {} ms",
        instant.elapsed().as_millis()
    );

    instant = Instant::now();
    let chunk_executor = setup_chunk_executor(db_rw.clone());
    debug!(
        "ChunkExecutor setup in {} ms",
        instant.elapsed().as_millis()
    );
    let chain_id = fetch_chain_id(&db_rw);
    let mut network_runtimes = vec![];
    let mut state_sync_network_handles = vec![];
    let mut mempool_network_handles = vec![];
    let mut consensus_network_handles = None;
    let mut reconfig_subscriptions = vec![];

    let (mempool_reconfig_subscription, mempool_reconfig_events) =
        gen_mempool_reconfig_subscription();
    reconfig_subscriptions.push(mempool_reconfig_subscription);
    // consensus has to subscribe to ALL on-chain configs
    let (consensus_reconfig_subscription, consensus_reconfig_events) =
        gen_consensus_reconfig_subscription();
    if node_config.base.role.is_validator() {
        reconfig_subscriptions.push(consensus_reconfig_subscription);
    }

    // Gather all network configs into a single vector.
    // TODO:  consider explicitly encoding the role in the NetworkConfig
    let mut network_configs: Vec<(RoleType, &NetworkConfig)> = node_config
        .full_node_networks
        .iter()
        .map(|network_config| (RoleType::FullNode, network_config))
        .collect();
    if let Some(network_config) = node_config.validator_network.as_ref() {
        network_configs.push((RoleType::Validator, network_config));
    }

    let mut network_builders = Vec::new();

    // Instantiate every network and collect the requisite endpoints for state_sync, mempool, and consensus.
    for (idx, (role, network_config)) in network_configs.into_iter().enumerate() {
        // Perform common instantiation steps
        let mut network_builder = NetworkBuilder::create(chain_id, role, network_config);
        let network_id = network_config.network_id.clone();

        // Create the endpoints to connect the Network to StateSynchronizer.
        let (state_sync_sender, state_sync_events) = network_builder
            .add_protocol_handler(state_synchronizer::network::network_endpoint_config());
        state_sync_network_handles.push((
            NodeNetworkId::new(network_id.clone(), idx),
            state_sync_sender,
            state_sync_events,
        ));

        // Create the endpoints to connect the Network to mempool.
        let (mempool_sender, mempool_events) =
            network_builder.add_protocol_handler(libra_mempool::network::network_endpoint_config(
                // TODO:  Make this configuration option more clear.
                node_config.mempool.max_broadcasts_per_peer,
            ));
        mempool_network_handles.push((
            NodeNetworkId::new(network_id, idx),
            mempool_sender,
            mempool_events,
        ));

        match role {
            // Perform steps relevant specifically to Validator networks.
            RoleType::Validator => {
                // A valid config is allowed to have at most one ValidatorNetwork
                // TODO:  `expect_none` would be perfect here, once it is stable.
                if consensus_network_handles.is_some() {
                    panic!("There can be at most one validator network!");
                }

                consensus_network_handles =
                    Some(network_builder.add_protocol_handler(
                        consensus::network_interface::network_endpoint_config(),
                    ));
            }
            // Currently no FullNode network specific steps.
            RoleType::FullNode => (),
        }

        reconfig_subscriptions.append(network_builder.reconfig_subscriptions());

        network_builders.push(network_builder);
    }

    // Build the configured networks.
    for network_builder in &mut network_builders {
        debug!("Creating runtime for {}", network_builder.network_context());
        let runtime = Builder::new()
            .thread_name("network-")
            .threaded_scheduler()
            .enable_all()
            .build()
            .expect("Failed to start runtime. Won't be able to start networking.");
        network_builder.build(runtime.handle().clone());
        network_runtimes.push(runtime);
        debug!(
            "Network built for network context: {}",
            network_builder.network_context()
        );
    }

    // Start the configured networks.
    // TODO:  Collect all component starts at the end of this function
    for network_builder in &mut network_builders {
        network_builder.start();
    }

    // TODO set up on-chain discovery network based on UpstreamConfig.fallback_network
    // and pass network handles to mempool/state sync

    // for state sync to send requests to mempool
    let (state_sync_to_mempool_sender, state_sync_requests) =
        channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);
    let state_synchronizer = StateSynchronizer::bootstrap(
        state_sync_network_handles,
        state_sync_to_mempool_sender,
        Arc::clone(&db_rw.reader),
        chunk_executor,
        &node_config,
        waypoint,
        reconfig_subscriptions,
    );
    let (mp_client_sender, mp_client_events) = channel(AC_SMP_CHANNEL_BUFFER_SIZE);

    let rpc_runtime = bootstrap_rpc(&node_config, chain_id, libra_db.clone(), mp_client_sender);

    let mut consensus_runtime = None;
    let (consensus_to_mempool_sender, consensus_requests) = channel(INTRA_NODE_CHANNEL_BUFFER_SIZE);

    instant = Instant::now();
    let mempool = libra_mempool::bootstrap(
        node_config,
        Arc::clone(&db_rw.reader),
        mempool_network_handles,
        mp_client_events,
        consensus_requests,
        state_sync_requests,
        mempool_reconfig_events,
    );
    debug!("Mempool started in {} ms", instant.elapsed().as_millis());

    // StateSync should be instantiated and started before Consensus to avoid a cyclic dependency:
    // network provider -> consensus -> state synchronizer -> network provider.  This has resulted
    // in a deadlock as observed in GitHub issue #749.
    if let Some((consensus_network_sender, consensus_network_events)) = consensus_network_handles {
        // Make sure that state synchronizer is caught up at least to its waypoint
        // (in case it's present). There is no sense to start consensus prior to that.
        // TODO: Note that we need the networking layer to be able to discover & connect to the
        // peers with potentially outdated network identity public keys.
        debug!("Wait until state synchronizer is initialized");
        block_on(state_synchronizer.wait_until_initialized())
            .expect("State synchronizer initialization failure");
        debug!("State synchronizer initialization complete.");

        // Initialize and start consensus.
        instant = Instant::now();
        consensus_runtime = Some(start_consensus(
            node_config,
            consensus_network_sender,
            consensus_network_events,
            state_synchronizer.create_client(),
            consensus_to_mempool_sender,
            libra_db,
            consensus_reconfig_events,
        ));
        debug!("Consensus started in {} ms", instant.elapsed().as_millis());
    }

    LibraHandle {
        network_runtimes,
        _rpc: rpc_runtime,
        _mempool: mempool,
        _state_synchronizer: state_synchronizer,
        _consensus_runtime: consensus_runtime,
        _debug: debug_if,
        _backup: backup_service,
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(ascii_ctype)]
#![feature(box_syntax)]
#![feature(conservative_impl_trait)]
#![feature(const_fn)]
#![feature(const_ptr_null)]
#![feature(const_ptr_null_mut)]
#![feature(core_intrinsics)]
#![feature(mpsc_select)]
#![feature(nonzero)]
#![feature(on_unimplemented)]
#![feature(plugin)]
#![feature(proc_macro)]
#![feature(stmt_expr_attributes)]
#![feature(try_from)]
#![feature(unboxed_closures)]
#![feature(untagged_unions)]

#![deny(unsafe_code)]
#![allow(non_snake_case)]

#![doc = "The script crate contains all matters DOM."]

#![plugin(script_plugins)]

extern crate angle;
extern crate app_units;
extern crate audio_video_metadata;
extern crate base64;
#[macro_use]
extern crate bitflags;
extern crate bluetooth_traits;
extern crate byteorder;
extern crate canvas_traits;
extern crate caseless;
extern crate cookie as cookie_rs;
extern crate core;
#[macro_use] extern crate cssparser;
#[macro_use] extern crate deny_public_fields;
extern crate devtools_traits;
extern crate dom_struct;
#[macro_use]
extern crate domobject_derive;
extern crate encoding;
extern crate euclid;
extern crate fnv;
extern crate gleam;
extern crate half;
#[macro_use] extern crate heapsize;
#[macro_use] extern crate heapsize_derive;
#[macro_use] extern crate html5ever;
#[macro_use]
extern crate hyper;
extern crate hyper_serde;
extern crate image;
extern crate ipc_channel;
#[macro_use]
extern crate js;
#[macro_use]
extern crate jstraceable_derive;
#[macro_use]
extern crate lazy_static;
extern crate libc;
#[macro_use]
extern crate log;
extern crate metrics;
#[macro_use]
extern crate mime;
extern crate mime_guess;
extern crate msg;
extern crate net_traits;
extern crate num_traits;
extern crate offscreen_gl_context;
extern crate open;
extern crate parking_lot;
extern crate phf;
#[macro_use]
extern crate profile_traits;
extern crate ref_filter_map;
extern crate ref_slice;
extern crate regex;
extern crate script_layout_interface;
extern crate script_traits;
extern crate selectors;
extern crate serde;
extern crate servo_arc;
#[macro_use] extern crate servo_atoms;
extern crate servo_config;
extern crate servo_geometry;
extern crate servo_rand;
extern crate servo_url;
extern crate smallvec;
#[macro_use]
extern crate style;
extern crate style_traits;
extern crate swapper;
extern crate time;
#[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))]
extern crate tinyfiledialogs;
extern crate unicode_segmentation;
extern crate url;
extern crate utf8;
extern crate uuid;
extern crate webrender_api;
extern crate webvr_traits;
extern crate xml5ever;

mod body;
pub mod clipboard_provider;
mod devtools;
pub mod document_loader;
#[macro_use]
mod dom;
pub mod fetch;
mod layout_image;
mod mem;
mod microtask;
mod network_listener;
pub mod script_runtime;
#[allow(unsafe_code)]
pub mod script_thread;
mod serviceworker_manager;
mod serviceworkerjob;
mod stylesheet_loader;
mod task_source;
pub mod test;
pub mod textinput;
mod timers;
mod unpremultiplytable;
mod webdriver_handlers;

/// A module with everything layout can use from script.
///
/// Try to keep this small!
///
/// TODO(emilio): A few of the FooHelpers can go away, presumably...
pub mod layout_exports {
    pub use dom::bindings::inheritance::{CharacterDataTypeId, ElementTypeId};
    pub use dom::bindings::inheritance::{HTMLElementTypeId, NodeTypeId};
    pub use dom::bindings::js::LayoutJS;
    pub use dom::characterdata::LayoutCharacterDataHelpers;
    pub use dom::document::{Document, LayoutDocumentHelpers, PendingRestyle};
    pub use dom::element::{Element, LayoutElementHelpers, RawLayoutElementHelpers};
    pub use dom::node::{CAN_BE_FRAGMENTED, HAS_DIRTY_DESCENDANTS, IS_IN_DOC};
    pub use dom::node::{HANDLED_SNAPSHOT, HAS_SNAPSHOT};
    pub use dom::node::{LayoutNodeHelpers, Node};
    pub use dom::text::Text;
}

use dom::bindings::codegen::RegisterBindings;
use dom::bindings::proxyhandler;
use script_traits::SWManagerSenders;
use serviceworker_manager::ServiceWorkerManager;

#[cfg(target_os = "linux")]
#[allow(unsafe_code)]
fn perform_platform_specific_initialization() {
    use std::mem;
    // 4096 is default max on many linux systems
    const MAX_FILE_LIMIT: libc::rlim_t = 4096;

    // Bump up our number of file descriptors to save us from impending doom caused by an onslaught
    // of iframes.
    unsafe {
        let mut rlim: libc::rlimit = mem::uninitialized();
        match libc::getrlimit(libc::RLIMIT_NOFILE, &mut rlim) {
            0 => {
                if rlim.rlim_cur >= MAX_FILE_LIMIT {
                    // we have more than enough
                    return;
                }

                rlim.rlim_cur = match rlim.rlim_max {
                    libc::RLIM_INFINITY => MAX_FILE_LIMIT,
                    _ => {
                        if rlim.rlim_max < MAX_FILE_LIMIT {
                            rlim.rlim_max
                        } else {
                            MAX_FILE_LIMIT
                        }
                    }
                };
                match libc::setrlimit(libc::RLIMIT_NOFILE, &rlim) {
                    0 => (),
                    _ => warn!("Failed to set file count limit"),
                };
            },
            _ => warn!("Failed to get file count limit"),
        };
    }
}

#[cfg(not(target_os = "linux"))]
fn perform_platform_specific_initialization() {}

pub fn init_service_workers(sw_senders: SWManagerSenders) {
    // Spawn the service worker manager passing the constellation sender
    ServiceWorkerManager::spawn_manager(sw_senders);
}

#[allow(unsafe_code)]
pub fn init() {
    unsafe {
        proxyhandler::init();

        // Create the global vtables used by the (generated) DOM
        // bindings to implement JS proxies.
        RegisterBindings::RegisterProxyHandlers();
    }

    perform_platform_specific_initialization();
}
<commit_msg>Remove unnecessary 'stmt_expr_attributes' feature flag<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(ascii_ctype)]
#![feature(box_syntax)]
#![feature(conservative_impl_trait)]
#![feature(const_fn)]
#![feature(const_ptr_null)]
#![feature(const_ptr_null_mut)]
#![feature(core_intrinsics)]
#![feature(mpsc_select)]
#![feature(nonzero)]
#![feature(on_unimplemented)]
#![feature(plugin)]
#![feature(proc_macro)]
#![feature(try_from)]
#![feature(unboxed_closures)]
#![feature(untagged_unions)]

#![deny(unsafe_code)]
#![allow(non_snake_case)]

#![doc = "The script crate contains all matters DOM."]

#![plugin(script_plugins)]

extern crate angle;
extern crate app_units;
extern crate audio_video_metadata;
extern crate base64;
#[macro_use]
extern crate bitflags;
extern crate bluetooth_traits;
extern crate byteorder;
extern crate canvas_traits;
extern crate caseless;
extern crate cookie as cookie_rs;
extern crate core;
#[macro_use] extern crate cssparser;
#[macro_use] extern crate deny_public_fields;
extern crate devtools_traits;
extern crate dom_struct;
#[macro_use]
extern crate domobject_derive;
extern crate encoding;
extern crate euclid;
extern crate fnv;
extern crate gleam;
extern crate half;
#[macro_use] extern crate heapsize;
#[macro_use] extern crate heapsize_derive;
#[macro_use] extern crate html5ever;
#[macro_use]
extern crate hyper;
extern crate hyper_serde;
extern crate image;
extern crate ipc_channel;
#[macro_use]
extern crate js;
#[macro_use]
extern crate jstraceable_derive;
#[macro_use]
extern crate lazy_static;
extern crate libc;
#[macro_use]
extern crate log;
extern crate metrics;
#[macro_use]
extern crate mime;
extern crate mime_guess;
extern crate msg;
extern crate net_traits;
extern crate num_traits;
extern crate offscreen_gl_context;
extern crate open;
extern crate parking_lot;
extern crate phf;
#[macro_use]
extern crate profile_traits;
extern crate ref_filter_map;
extern crate ref_slice;
extern crate regex;
extern crate script_layout_interface;
extern crate script_traits;
extern crate selectors;
extern crate serde;
extern crate servo_arc;
#[macro_use] extern crate servo_atoms;
extern crate servo_config;
extern crate servo_geometry;
extern crate servo_rand;
extern crate servo_url;
extern crate smallvec;
#[macro_use]
extern crate style;
extern crate style_traits;
extern crate swapper;
extern crate time;
#[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))]
extern crate tinyfiledialogs;
extern crate unicode_segmentation;
extern crate url;
extern crate utf8;
extern crate uuid;
extern crate webrender_api;
extern crate webvr_traits;
extern crate xml5ever;

mod body;
pub mod clipboard_provider;
mod devtools;
pub mod document_loader;
#[macro_use]
mod dom;
pub mod fetch;
mod layout_image;
mod mem;
mod microtask;
mod network_listener;
pub mod script_runtime;
#[allow(unsafe_code)]
pub mod script_thread;
mod serviceworker_manager;
mod serviceworkerjob;
mod stylesheet_loader;
mod task_source;
pub mod test;
pub mod textinput;
mod timers;
mod unpremultiplytable;
mod webdriver_handlers;

/// A module with everything layout can use from script.
///
/// Try to keep this small!
///
/// TODO(emilio): A few of the FooHelpers can go away, presumably...
pub mod layout_exports {
    pub use dom::bindings::inheritance::{CharacterDataTypeId, ElementTypeId};
    pub use dom::bindings::inheritance::{HTMLElementTypeId, NodeTypeId};
    pub use dom::bindings::js::LayoutJS;
    pub use dom::characterdata::LayoutCharacterDataHelpers;
    pub use dom::document::{Document, LayoutDocumentHelpers, PendingRestyle};
    pub use dom::element::{Element, LayoutElementHelpers, RawLayoutElementHelpers};
    pub use dom::node::{CAN_BE_FRAGMENTED, HAS_DIRTY_DESCENDANTS, IS_IN_DOC};
    pub use dom::node::{HANDLED_SNAPSHOT, HAS_SNAPSHOT};
    pub use dom::node::{LayoutNodeHelpers, Node};
    pub use dom::text::Text;
}

use dom::bindings::codegen::RegisterBindings;
use dom::bindings::proxyhandler;
use script_traits::SWManagerSenders;
use serviceworker_manager::ServiceWorkerManager;

#[cfg(target_os = "linux")]
#[allow(unsafe_code)]
fn perform_platform_specific_initialization() {
    use std::mem;
    // 4096 is default max on many linux systems
    const MAX_FILE_LIMIT: libc::rlim_t = 4096;

    // Bump up our number of file descriptors to save us from impending doom caused by an onslaught
    // of iframes.
    unsafe {
        let mut rlim: libc::rlimit = mem::uninitialized();
        match libc::getrlimit(libc::RLIMIT_NOFILE, &mut rlim) {
            0 => {
                if rlim.rlim_cur >= MAX_FILE_LIMIT {
                    // we have more than enough
                    return;
                }

                rlim.rlim_cur = match rlim.rlim_max {
                    libc::RLIM_INFINITY => MAX_FILE_LIMIT,
                    _ => {
                        if rlim.rlim_max < MAX_FILE_LIMIT {
                            rlim.rlim_max
                        } else {
                            MAX_FILE_LIMIT
                        }
                    }
                };
                match libc::setrlimit(libc::RLIMIT_NOFILE, &rlim) {
                    0 => (),
                    _ => warn!("Failed to set file count limit"),
                };
            },
            _ => warn!("Failed to get file count limit"),
        };
    }
}

#[cfg(not(target_os = "linux"))]
fn perform_platform_specific_initialization() {}

pub fn init_service_workers(sw_senders: SWManagerSenders) {
    // Spawn the service worker manager passing the constellation sender
    ServiceWorkerManager::spawn_manager(sw_senders);
}

#[allow(unsafe_code)]
pub fn init() {
    unsafe {
        proxyhandler::init();

        // Create the global vtables used by the (generated) DOM
        // bindings to implement JS proxies.
        RegisterBindings::RegisterProxyHandlers();
    }

    perform_platform_specific_initialization();
}
<|endoftext|>


<commit_before>use crate::rational::Rational64;
use std::any::Any;
use std::sync::Arc;

/// Timestamp information.
#[derive(Debug, Clone, Default)]
pub struct TimeInfo {
    /// Presentation timestamp.
    pub pts: Option<i64>,
    /// Decode timestamp.
    pub dts: Option<i64>,
    /// Duration (in timebase units).
    pub duration: Option<u64>,
    /// Timebase numerator/denominator.
    pub timebase: Option<Rational64>,
    /// Timebase user private data.
    pub user_private: Option<Arc<dyn Any + Send + Sync>>,
}
<commit_msg>Clarify timeinfo documentation<commit_after>use crate::rational::Rational64;
use std::any::Any;
use std::sync::Arc;

/// Timestamp information for frames and packets.
#[derive(Debug, Clone, Default)]
pub struct TimeInfo {
    /// Presentation timestamp.
    pub pts: Option<i64>,
    /// Decode timestamp.
    pub dts: Option<i64>,
    /// Duration (in timebase units).
    pub duration: Option<u64>,
    /// Timebase numerator/denominator (i.e 1/75th of a second).
    ///
    /// Its value does not vary among frames/packets, since it is
    /// computed and defined at stream level.
    pub timebase: Option<Rational64>,
    /// Timebase user private data.
    pub user_private: Option<Arc<dyn Any + Send + Sync>>,
}
<|endoftext|>


<commit_before>use base64::write::EncoderWriter;
use indexmap::IndexMap;
use prometheus::proto::{MetricFamily, MetricType};
use prometheus::{Encoder, Error};
use serde::ser::SerializeSeq;
use serde::{Serialize, Serializer as _};
use serde_json::Serializer;
use std::cell::Cell;
use std::io::Write;
use std::rc::Rc;

const CHUNKS_MAX_SIZE_BYTES: usize = 5000;

/// The `LogEncoder` struct encodes Prometheus metrics in the format [`crates-io-heroku-metrics`]
/// expects metrics to be logged. This can be used to forward instance metrics to it, allowing them
/// to be scraped by the Rust infrastructure monitoring system.
///
/// The metrics are encoded in the format [Vector expects them][vector], and the list of metrics is
/// json-encoded and then base64-encoded. The whole thing is prefixed with a predefined string to
/// let [`crates-io-heroku-metrics`] find it easily.
///
/// This is needed mostly for crates.io hosted on Heroku. Deployments of crates.io on other
/// platforms shouldn't need this.
///
/// [`crates-io-heroku-metrics`]: https://github.com/rust-lang/crates-io-heroku-metrics
/// [vector]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/
#[derive(Debug, Clone, Copy)]
pub struct LogEncoder(());

impl LogEncoder {
    pub fn new() -> Self {
        Self(())
    }
}

impl Encoder for LogEncoder {
    fn encode<W: Write>(&self, families: &[MetricFamily], dest: &mut W) -> prometheus::Result<()> {
        let events = families_to_json_events(families);

        let chunks = serialize_and_split_list(events.iter(), CHUNKS_MAX_SIZE_BYTES)
            .map_err(|e| Error::Msg(e.to_string()))?;

        for chunk in chunks {
            dest.write_all(b"crates-io-heroku-metrics:ingest ")?;
            dest.write_all(&chunk)?;
            dest.write_all(b"\n")?;
        }

        Ok(())
    }

    fn format_type(&self) -> &str {
        "crates-io-heroku-metrics log encoding"
    }
}

fn families_to_json_events(families: &[MetricFamily]) -> Vec<VectorEvent<'_>> {
    let mut events = Vec::new();
    for family in families {
        for metric in family.get_metric() {
            let data = match family.get_field_type() {
                MetricType::COUNTER => VectorMetricData::Counter {
                    value: metric.get_counter().get_value(),
                },
                MetricType::GAUGE => VectorMetricData::Gauge {
                    value: metric.get_gauge().get_value(),
                },
                MetricType::HISTOGRAM => {
                    let histogram = metric.get_histogram();

                    // We need to convert from cumulative counts (used by the Prometheus library)
                    // to plain counts (used by Vector).
                    let mut buckets = Vec::new();
                    let mut last_cumulative_count = 0;
                    for bucket in histogram.get_bucket() {
                        buckets.push(VectorHistogramBucket {
                            upper_limit: bucket.get_upper_bound(),
                            count: bucket.get_cumulative_count() - last_cumulative_count,
                        });
                        last_cumulative_count = bucket.get_cumulative_count();
                    }

                    VectorMetricData::AggregatedHistogram {
                        count: histogram.get_sample_count(),
                        sum: histogram.get_sample_sum(),
                        buckets,
                    }
                }
                other => {
                    panic!("unsupported metric type: {:?}", other)
                }
            };
            events.push(VectorEvent {
                metric: VectorMetric {
                    data,
                    kind: "absolute",
                    name: family.get_name(),
                    tags: metric
                        .get_label()
                        .iter()
                        .map(|p| (p.get_name(), p.get_value()))
                        .collect(),
                },
            });
        }
    }
    events
}

/// Serialize a list of items into multiple Base64-encoded JSON chunks.
///
/// Our hosting platform (Heroku) limits the size of log lines, arbitrarily splitting them once
/// they reach a threshold. We can't let Heroku do the split as it doesn't know where to properly
/// do that, so we need to do the splitting ourselves.
///
/// This function takes an iterator of serializable items and returns the serialized version,
/// possibly split into multiple chunks. Each chunk is *at least* `max_size_hint` long, as the
/// function stops serializing new items in the same chunk only when the size limit is reached
/// after serializing an item.
///
/// Because of that `max_size_hint` should be lower than the upper bound we can't cross.
fn serialize_and_split_list<'a, S: Serialize + 'a>(
    items: impl Iterator<Item = &'a S>,
    max_size_hint: usize,
) -> Result<Vec<Vec<u8>>, serde_json::Error> {
    let mut items = items.peekable();

    let mut result = Vec::new();
    while items.peek().is_some() {
        let mut writer = TrackedWriter::new();
        let written_count = writer.written_count.clone();
        let mut serializer = Serializer::new(EncoderWriter::new(&mut writer, base64::STANDARD));

        let mut seq = serializer.serialize_seq(None)?;
        while let Some(next) = items.next() {
            seq.serialize_element(next)?;
            if written_count.get() >= max_size_hint {
                break;
            }
        }
        seq.end()?;
        drop(serializer);

        result.push(writer.buffer);
    }

    Ok(result)
}

struct TrackedWriter {
    buffer: Vec<u8>,
    written_count: Rc<Cell<usize>>,
}

impl TrackedWriter {
    fn new() -> Self {
        Self {
            buffer: Vec::new(),
            written_count: Rc::new(Cell::new(0)),
        }
    }
}

impl Write for TrackedWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let written = self.buffer.write(buf)?;
        self.written_count.set(self.written_count.get() + written);
        Ok(written)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.buffer.flush()
    }
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorEvent<'a> {
    metric: VectorMetric<'a>,
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorMetric<'a> {
    #[serde(flatten)]
    data: VectorMetricData,
    kind: &'a str,
    name: &'a str,
    tags: IndexMap<&'a str, &'a str>,
}

#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
enum VectorMetricData {
    AggregatedHistogram {
        buckets: Vec<VectorHistogramBucket>,
        count: u64,
        sum: f64,
    },
    Counter {
        value: f64,
    },
    Gauge {
        value: f64,
    },
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorHistogramBucket {
    upper_limit: f64,
    count: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Error;
    use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge, IntGaugeVec, Opts, Registry};

    #[test]
    fn test_counter_to_json() -> Result<(), Error> {
        let counter =
            IntCounter::with_opts(Opts::new("sample_counter", "sample_counter help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(counter.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 0.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        counter.inc_by(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 42.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_gauge_to_json() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 0.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        gauge.set(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 42.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_histogram_to_json() -> Result<(), Error> {
        let histogram = Histogram::with_opts(HistogramOpts::new(
            "sample_histogram",
            "sample_histogram help message",
        ))?;
        let registry = Registry::new();
        registry.register(Box::new(histogram.clone()))?;

        let mut value = 0.0;
        while value < 11.0 {
            histogram.observe(value);
            value += 0.001;
        }

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::AggregatedHistogram {
                        buckets: vec![
                            VectorHistogramBucket {
                                upper_limit: 0.005,
                                count: 6,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.01,
                                count: 4,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.025,
                                count: 15,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.05,
                                count: 25,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.1,
                                count: 50,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.25,
                                count: 150,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.5,
                                count: 250,
                            },
                            VectorHistogramBucket {
                                upper_limit: 1.0,
                                count: 500,
                            },
                            VectorHistogramBucket {
                                upper_limit: 2.5,
                                count: 1501,
                            },
                            VectorHistogramBucket {
                                upper_limit: 5.0,
                                count: 2499,
                            },
                            VectorHistogramBucket {
                                upper_limit: 10.0,
                                count: 5001,
                            },
                        ],
                        count: 11001,
                        sum: 60505.50000000138,
                    },
                    kind: "absolute",
                    name: "sample_histogram",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_metric_with_tags_to_json() -> Result<(), Error> {
        let gauge_vec = IntGaugeVec::new(
            Opts::new("sample_gauge", "sample_gauge help message"),
            &["label1", "label2"],
        )?;
        let registry = Registry::new();
        registry.register(Box::new(gauge_vec.clone()))?;

        gauge_vec.with_label_values(&["foo", "1"]).set(42);
        gauge_vec.with_label_values(&["bar", "2"]).set(98);

        assert_eq!(
            vec![
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 98.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "bar"), ("label2", "2")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 42.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "foo"), ("label2", "1")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
            ],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_encoding() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        let mut output = Vec::new();
        LogEncoder::new().encode(&registry.gather(), &mut output)?;

        assert_eq!(
            b"crates-io-heroku-metrics:ingest W3sibWV0cmljIjp7ImdhdWdlIjp7InZhbHVlIjowLjB9LCJraW5kIjoiYWJzb2x1dGUiLCJuYW1lIjoic2FtcGxlX2dhdWdlIiwidGFncyI6e319fV0=\n",
            output.as_slice(),
        );

        Ok(())
    }

    #[test]
    fn test_serialize_and_split_list_small() -> Result<(), Error> {
        let small = (0..10).collect::<Vec<_>>();
        let chunks = serialize_and_split_list(small.iter(), 256)?;

        assert_eq!(chunks.len(), 1);
        assert!(chunks[0].len() <= 256);
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[0])?)?,
            small,
        );

        Ok(())
    }

    #[test]
    fn test_serialize_and_split_list_long() -> Result<(), Error> {
        let small = (0..100).collect::<Vec<_>>();
        let chunks = serialize_and_split_list(small.iter(), 256)?;

        assert_eq!(chunks.len(), 2);
        assert!(chunks[0].len() >= 256);
        assert!(chunks[1].len() <= 256);
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[0])?)?,
            (0..=67).collect::<Vec<_>>(),
        );
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[1])?)?,
            (68..100).collect::<Vec<_>>(),
        );

        Ok(())
    }
}
<commit_msg>Auto merge of #3691 - pietroalbini:fix-histogram-repr, r=jtgeibel<commit_after>use base64::write::EncoderWriter;
use indexmap::IndexMap;
use prometheus::proto::{MetricFamily, MetricType};
use prometheus::{Encoder, Error};
use serde::ser::SerializeSeq;
use serde::{Serialize, Serializer as _};
use serde_json::Serializer;
use std::cell::Cell;
use std::io::Write;
use std::rc::Rc;

const CHUNKS_MAX_SIZE_BYTES: usize = 5000;

/// The `LogEncoder` struct encodes Prometheus metrics in the format [`crates-io-heroku-metrics`]
/// expects metrics to be logged. This can be used to forward instance metrics to it, allowing them
/// to be scraped by the Rust infrastructure monitoring system.
///
/// The metrics are encoded in the format [Vector expects them][vector], and the list of metrics is
/// json-encoded and then base64-encoded. The whole thing is prefixed with a predefined string to
/// let [`crates-io-heroku-metrics`] find it easily.
///
/// This is needed mostly for crates.io hosted on Heroku. Deployments of crates.io on other
/// platforms shouldn't need this.
///
/// [`crates-io-heroku-metrics`]: https://github.com/rust-lang/crates-io-heroku-metrics
/// [vector]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/
#[derive(Debug, Clone, Copy)]
pub struct LogEncoder(());

impl LogEncoder {
    pub fn new() -> Self {
        Self(())
    }
}

impl Encoder for LogEncoder {
    fn encode<W: Write>(&self, families: &[MetricFamily], dest: &mut W) -> prometheus::Result<()> {
        let events = families_to_json_events(families);

        let chunks = serialize_and_split_list(events.iter(), CHUNKS_MAX_SIZE_BYTES)
            .map_err(|e| Error::Msg(e.to_string()))?;

        for chunk in chunks {
            dest.write_all(b"crates-io-heroku-metrics:ingest ")?;
            dest.write_all(&chunk)?;
            dest.write_all(b"\n")?;
        }

        Ok(())
    }

    fn format_type(&self) -> &str {
        "crates-io-heroku-metrics log encoding"
    }
}

fn families_to_json_events(families: &[MetricFamily]) -> Vec<VectorEvent<'_>> {
    let mut events = Vec::new();
    for family in families {
        for metric in family.get_metric() {
            let data = match family.get_field_type() {
                MetricType::COUNTER => VectorMetricData::Counter {
                    value: metric.get_counter().get_value(),
                },
                MetricType::GAUGE => VectorMetricData::Gauge {
                    value: metric.get_gauge().get_value(),
                },
                MetricType::HISTOGRAM => {
                    let histogram = metric.get_histogram();

                    // We need to convert from cumulative counts (used by the Prometheus library)
                    // to plain counts (used by Vector).
                    let mut buckets = Vec::new();
                    let mut counts = Vec::new();
                    let mut last_cumulative_count = 0;
                    for bucket in histogram.get_bucket() {
                        buckets.push(bucket.get_upper_bound());
                        counts.push(bucket.get_cumulative_count() - last_cumulative_count);
                        last_cumulative_count = bucket.get_cumulative_count();
                    }

                    VectorMetricData::AggregatedHistogram {
                        count: histogram.get_sample_count(),
                        sum: histogram.get_sample_sum(),
                        buckets,
                        counts,
                    }
                }
                other => {
                    panic!("unsupported metric type: {:?}", other)
                }
            };
            events.push(VectorEvent {
                metric: VectorMetric {
                    data,
                    kind: "absolute",
                    name: family.get_name(),
                    tags: metric
                        .get_label()
                        .iter()
                        .map(|p| (p.get_name(), p.get_value()))
                        .collect(),
                },
            });
        }
    }
    events
}

/// Serialize a list of items into multiple Base64-encoded JSON chunks.
///
/// Our hosting platform (Heroku) limits the size of log lines, arbitrarily splitting them once
/// they reach a threshold. We can't let Heroku do the split as it doesn't know where to properly
/// do that, so we need to do the splitting ourselves.
///
/// This function takes an iterator of serializable items and returns the serialized version,
/// possibly split into multiple chunks. Each chunk is *at least* `max_size_hint` long, as the
/// function stops serializing new items in the same chunk only when the size limit is reached
/// after serializing an item.
///
/// Because of that `max_size_hint` should be lower than the upper bound we can't cross.
fn serialize_and_split_list<'a, S: Serialize + 'a>(
    items: impl Iterator<Item = &'a S>,
    max_size_hint: usize,
) -> Result<Vec<Vec<u8>>, serde_json::Error> {
    let mut items = items.peekable();

    let mut result = Vec::new();
    while items.peek().is_some() {
        let mut writer = TrackedWriter::new();
        let written_count = writer.written_count.clone();
        let mut serializer = Serializer::new(EncoderWriter::new(&mut writer, base64::STANDARD));

        let mut seq = serializer.serialize_seq(None)?;
        while let Some(next) = items.next() {
            seq.serialize_element(next)?;
            if written_count.get() >= max_size_hint {
                break;
            }
        }
        seq.end()?;
        drop(serializer);

        result.push(writer.buffer);
    }

    Ok(result)
}

struct TrackedWriter {
    buffer: Vec<u8>,
    written_count: Rc<Cell<usize>>,
}

impl TrackedWriter {
    fn new() -> Self {
        Self {
            buffer: Vec::new(),
            written_count: Rc::new(Cell::new(0)),
        }
    }
}

impl Write for TrackedWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let written = self.buffer.write(buf)?;
        self.written_count.set(self.written_count.get() + written);
        Ok(written)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.buffer.flush()
    }
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorEvent<'a> {
    metric: VectorMetric<'a>,
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorMetric<'a> {
    #[serde(flatten)]
    data: VectorMetricData,
    kind: &'a str,
    name: &'a str,
    tags: IndexMap<&'a str, &'a str>,
}

#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
enum VectorMetricData {
    AggregatedHistogram {
        buckets: Vec<f64>,
        counts: Vec<u64>,
        count: u64,
        sum: f64,
    },
    Counter {
        value: f64,
    },
    Gauge {
        value: f64,
    },
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Error;
    use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge, IntGaugeVec, Opts, Registry};

    #[test]
    fn test_counter_to_json() -> Result<(), Error> {
        let counter =
            IntCounter::with_opts(Opts::new("sample_counter", "sample_counter help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(counter.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 0.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        counter.inc_by(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 42.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_gauge_to_json() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 0.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        gauge.set(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 42.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_histogram_to_json() -> Result<(), Error> {
        let histogram = Histogram::with_opts(HistogramOpts::new(
            "sample_histogram",
            "sample_histogram help message",
        ))?;
        let registry = Registry::new();
        registry.register(Box::new(histogram.clone()))?;

        let mut value = 0.0;
        while value < 11.0 {
            histogram.observe(value);
            value += 0.001;
        }

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::AggregatedHistogram {
                        buckets: vec![
                            0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
                        ],
                        counts: vec![6, 4, 15, 25, 50, 150, 250, 500, 1501, 2499, 5001,],
                        count: 11001,
                        sum: 60505.50000000138,
                    },
                    kind: "absolute",
                    name: "sample_histogram",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_metric_with_tags_to_json() -> Result<(), Error> {
        let gauge_vec = IntGaugeVec::new(
            Opts::new("sample_gauge", "sample_gauge help message"),
            &["label1", "label2"],
        )?;
        let registry = Registry::new();
        registry.register(Box::new(gauge_vec.clone()))?;

        gauge_vec.with_label_values(&["foo", "1"]).set(42);
        gauge_vec.with_label_values(&["bar", "2"]).set(98);

        assert_eq!(
            vec![
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 98.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "bar"), ("label2", "2")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 42.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "foo"), ("label2", "1")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
            ],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_encoding() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        let mut output = Vec::new();
        LogEncoder::new().encode(&registry.gather(), &mut output)?;

        assert_eq!(
            b"crates-io-heroku-metrics:ingest W3sibWV0cmljIjp7ImdhdWdlIjp7InZhbHVlIjowLjB9LCJraW5kIjoiYWJzb2x1dGUiLCJuYW1lIjoic2FtcGxlX2dhdWdlIiwidGFncyI6e319fV0=\n",
            output.as_slice(),
        );

        Ok(())
    }

    #[test]
    fn test_serialize_and_split_list_small() -> Result<(), Error> {
        let small = (0..10).collect::<Vec<_>>();
        let chunks = serialize_and_split_list(small.iter(), 256)?;

        assert_eq!(chunks.len(), 1);
        assert!(chunks[0].len() <= 256);
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[0])?)?,
            small,
        );

        Ok(())
    }

    #[test]
    fn test_serialize_and_split_list_long() -> Result<(), Error> {
        let small = (0..100).collect::<Vec<_>>();
        let chunks = serialize_and_split_list(small.iter(), 256)?;

        assert_eq!(chunks.len(), 2);
        assert!(chunks[0].len() >= 256);
        assert!(chunks[1].len() <= 256);
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[0])?)?,
            (0..=67).collect::<Vec<_>>(),
        );
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[1])?)?,
            (68..100).collect::<Vec<_>>(),
        );

        Ok(())
    }
}
<|endoftext|>


<commit_before>use errors::*;
use multirust_dist::{utils, dist};
use multirust_dist::prefix::InstallPrefix;
use multirust_dist::dist::ToolchainDesc;
use multirust_dist::manifestation::{Manifestation, Changes};
use multirust_dist::manifest::Component;
use config::Cfg;
use env_var;
use install::{self, InstallMethod};

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

use hyper;
use multirust_dist;

#[derive(Debug)]
pub struct Toolchain<'a> {
    cfg: &'a Cfg,
    name: String,
    path: PathBuf,
}

impl<'a> Toolchain<'a> {
    pub fn from(cfg: &'a Cfg, name: &str) -> Self {
        let path = cfg.toolchains_dir.join(name);
        Toolchain {
            cfg: cfg,
            name: name.to_owned(),
            path: path.clone(),
        }
    }
    pub fn name(&self) -> &str {
        &self.name
    }
    pub fn path(&self) -> &Path {
        &self.path
    }
    pub fn exists(&self) -> bool {
        utils::is_directory(&self.path)
    }
    pub fn verify(&self) -> Result<()> {
        Ok(try!(utils::assert_is_directory(&self.path)))
    }
    pub fn remove(&self) -> Result<()> {
        if self.exists() {
            self.cfg.notify_handler.call(Notification::UninstallingToolchain(&self.name));
        } else {
            self.cfg.notify_handler.call(Notification::ToolchainNotInstalled(&self.name));
            return Ok(());
        }
        if let Some(update_hash) = try!(self.update_hash()) {
            try!(utils::remove_file("update hash", &update_hash));
        }
        let handler = self.cfg.notify_handler.as_ref();
        let result = install::uninstall(&self.path, ntfy!(&handler));
        if !self.exists() {
            self.cfg.notify_handler.call(Notification::UninstalledToolchain(&self.name));
        }
        Ok(try!(result))
    }
    fn install(&self, install_method: InstallMethod) -> Result<()> {
        assert!(self.is_valid_install_method(install_method));
        if self.exists() {
            self.cfg.notify_handler.call(Notification::UpdatingToolchain(&self.name));
        } else {
            self.cfg.notify_handler.call(Notification::InstallingToolchain(&self.name));
        }
        self.cfg
            .notify_handler
            .call(Notification::ToolchainDirectory(&self.path, &self.name));
        let handler = self.cfg.notify_handler.as_ref();
        Ok(try!(install_method.run(&self.path, ntfy!(&handler))))
    }
    fn install_if_not_installed(&self, install_method: InstallMethod) -> Result<()> {
        assert!(self.is_valid_install_method(install_method));
        self.cfg.notify_handler.call(Notification::LookingForToolchain(&self.name));
        if !self.exists() {
            self.install(install_method)
        } else {
            self.cfg.notify_handler.call(Notification::UsingExistingToolchain(&self.name));
            Ok(())
        }
    }
    fn is_valid_install_method(&self, install_method: InstallMethod) -> bool {
        match install_method {
            InstallMethod::Copy(_) |
            InstallMethod::Link(_) |
            InstallMethod::Installer(_, _) => self.is_custom(),
            InstallMethod::Dist(_, _, _) => !self.is_custom(),
        }
    }
    fn update_hash(&self) -> Result<Option<PathBuf>> {
        if self.is_custom() {
            Ok(None)
        } else {
            Ok(Some(try!(self.cfg.get_hash_file(&self.name, true))))
        }
    }

    fn download_cfg(&self) -> dist::DownloadCfg {
        dist::DownloadCfg {
            dist_root: &self.cfg.dist_root_url,
            temp_cfg: &self.cfg.temp_cfg,
            notify_handler: ntfy!(&self.cfg.notify_handler),
        }
    }

    pub fn install_from_dist(&self) -> Result<()> {
        let update_hash = try!(self.update_hash());
        self.install(InstallMethod::Dist(&self.name,
                                         update_hash.as_ref().map(|p| &**p),
                                         self.download_cfg()))
    }
    pub fn install_from_dist_if_not_installed(&self) -> Result<()> {
        let update_hash = try!(self.update_hash());
        self.install_if_not_installed(InstallMethod::Dist(&self.name,
                                                          update_hash.as_ref().map(|p| &**p),
                                                          self.download_cfg()))
    }
    pub fn is_custom(&self) -> bool {
        ToolchainDesc::from_str(&self.name).is_err()
    }
    pub fn is_tracking(&self) -> bool {
        ToolchainDesc::from_str(&self.name).ok().map(|d| d.is_tracking()) == Some(true)
    }

    fn ensure_custom(&self) -> Result<()> {
        if !self.is_custom() {
            Err(Error::Install(multirust_dist::Error::InvalidToolchainName(self.name.to_string())))
        } else {
            Ok(())
        }
    }

    pub fn install_from_installers(&self, installers: &[&OsStr]) -> Result<()> {
        try!(self.ensure_custom());

        try!(self.remove());

        // FIXME: This should do all downloads first, then do
        // installs, and do it all in a single transaction.
        for installer in installers {
            let installer_str = installer.to_str().unwrap_or("bogus");
            match installer_str.rfind(".") {
                Some(i) => {
                    let extension = &installer_str[i+1..];
                    if extension != "gz" {
                        return Err(Error::BadInstallerType(extension.to_string()));
                    }
                }
                None => return Err(Error::BadInstallerType(String::from("(none)")))
            }

            // FIXME: Pretty hacky
            let is_url = installer_str.starts_with("file://")
                || installer_str.starts_with("http://")
                || installer_str.starts_with("https://");
            let url = hyper::Url::parse(installer_str).ok();
            let url = if is_url { url } else { None };
            if let Some(url) = url {

                // Download to a local file
                let local_installer = try!(self.cfg.temp_cfg.new_file_with_ext("", ".tar.gz"));
                try!(utils::download_file(url,
                                          &local_installer,
                                          None,
                                          ntfy!(&self.cfg.notify_handler)));
                try!(self.install(InstallMethod::Installer(&local_installer, &self.cfg.temp_cfg)));
            } else {
                // If installer is a filename

                // No need to download
                let local_installer = Path::new(installer);

                // Install from file
                try!(self.install(InstallMethod::Installer(&local_installer, &self.cfg.temp_cfg)));
            }
        }

        Ok(())
    }

    pub fn install_from_dir(&self, src: &Path, link: bool) -> Result<()> {
        try!(self.ensure_custom());

        if link {
            self.install(InstallMethod::Link(&try!(utils::to_absolute(src))))
        } else {
            self.install(InstallMethod::Copy(src))
        }
    }

    pub fn create_command<T: AsRef<OsStr>>(&self, binary: T) -> Result<Command> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        let mut cmd = Command::new(binary);
        self.set_env(&mut cmd);
        Ok(cmd)
    }

    fn set_env(&self, cmd: &mut Command) {
        let ref cargo_home = self.cfg.multirust_dir.join("cargo");
        let ref bin_path = self.path.join("bin");

        self.set_ldpath(cmd);

        env_var::set_path("PATH", bin_path, cmd);
        env_var::set_default("CARGO_HOME", cargo_home.as_ref(), cmd);
        env_var::inc("RUST_RECURSION_COUNT", cmd);

        // FIXME: This should not be a path, but a toolchain name.
        // Not sure what's going on here.
        cmd.env("MULTIRUST_TOOLCHAIN", &self.path);
        cmd.env("MULTIRUST_HOME", &self.cfg.multirust_dir);
    }

    pub fn set_ldpath(&self, cmd: &mut Command) {
        let new_path = self.path.join("lib");

        env_var::set_path("LD_LIBRARY_PATH", &new_path, cmd);
        env_var::set_path("DYLD_LIBRARY_PATH", &new_path, cmd);
    }

    pub fn doc_path(&self, relative: &str) -> Result<PathBuf> {
        try!(self.verify());

        let parts = vec!["share", "doc", "rust", "html"];
        let mut doc_dir = self.path.clone();
        for part in parts {
            doc_dir.push(part);
        }
        doc_dir.push(relative);

        Ok(doc_dir)
    }
    pub fn open_docs(&self, relative: &str) -> Result<()> {
        try!(self.verify());

        Ok(try!(utils::open_browser(&try!(self.doc_path(relative)))))
    }

    pub fn make_default(&self) -> Result<()> {
        self.cfg.set_default(&self.name)
    }
    pub fn make_override(&self, path: &Path) -> Result<()> {
        Ok(try!(self.cfg.override_db.set(path,
                                         &self.name,
                                         &self.cfg.temp_cfg,
                                         self.cfg.notify_handler.as_ref())))
    }

    pub fn list_components(&self) -> Result<Vec<Component>> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        // FIXME: This toolchain handling is a mess. Just do it once
        // when the toolchain is created.
        let ref toolchain = self.name;
        let ref toolchain = try!(ToolchainDesc::from_str(toolchain));
        let trip = toolchain.target_triple();
        let prefix = InstallPrefix::from(self.path.to_owned());
        let manifestation = try!(Manifestation::open(prefix, &trip));

        if let Some(manifest) = try!(manifestation.load_manifest()) {
            // Return all optional components of the "rust" package for the
            // toolchain's target triple.
            let mut res = Vec::new();

            let rust_pkg = manifest.packages.get("rust")
                .expect("manifest should cantain a rust package");
            let targ_pkg = rust_pkg.targets.get(&trip)
                .expect("installed manifest should have a known target");

            for extension in &targ_pkg.extensions {
                res.push(extension.clone())
            }

            Ok(res)
        } else {
            Err(Error::ComponentsUnsupported(self.name.to_string()))
        }
    }

    pub fn add_component(&self, component: Component) -> Result<()> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        let ref toolchain = self.name;
        let ref toolchain = try!(ToolchainDesc::from_str(toolchain));
        let trip = toolchain.target_triple();
        let prefix = InstallPrefix::from(self.path.to_owned());
        let manifestation = try!(Manifestation::open(prefix, &trip));

        if let Some(manifest) = try!(manifestation.load_manifest()) {

            // Validate the component name
            let rust_pkg = manifest.packages.get("rust")
                .expect("manifest should cantain a rust package");
            let targ_pkg = rust_pkg.targets.get(&trip)
                .expect("installed manifest should have a known target");

            if targ_pkg.components.contains(&component) {
                return Err(Error::AddingRequiredComponent(self.name.to_string(), component));
            }

            if !targ_pkg.extensions.contains(&component) {
                return Err(Error::UnknownComponent(self.name.to_string(), component));
            }

            let changes = Changes {
                add_extensions: vec![component],
                remove_extensions: vec![]
            };

            try!(manifestation.update(&self.name,
                                      &manifest,
                                      changes,
                                      self.download_cfg().temp_cfg,
                                      self.download_cfg().notify_handler.clone()));

            Ok(())
        } else {
            Err(Error::ComponentsUnsupported(self.name.to_string()))
        }
    }

    pub fn remove_component(&self, component: Component) -> Result<()> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        let ref toolchain = self.name;
        let ref toolchain = try!(ToolchainDesc::from_str(toolchain));
        let trip = toolchain.target_triple();
        let prefix = InstallPrefix::from(self.path.to_owned());
        let manifestation = try!(Manifestation::open(prefix, &trip));

        if let Some(manifest) = try!(manifestation.load_manifest()) {

            // Validate the component name
            let rust_pkg = manifest.packages.get("rust")
                .expect("manifest should cantain a rust package");
            let targ_pkg = rust_pkg.targets.get(&trip)
                .expect("installed manifest should have a known target");

            if targ_pkg.components.contains(&component) {
                return Err(Error::RemovingRequiredComponent(self.name.to_string(), component));
            }

            let dist_config = try!(manifestation.read_config()).unwrap();
            if !dist_config.components.contains(&component) {
                return Err(Error::UnknownComponent(self.name.to_string(), component));
            }

            let changes = Changes {
                add_extensions: vec![],
                remove_extensions: vec![component]
            };

            try!(manifestation.update(&self.name,
                                      &manifest,
                                      changes,
                                      self.download_cfg().temp_cfg,
                                      self.download_cfg().notify_handler.clone()));

            Ok(())
        } else {
            Err(Error::ComponentsUnsupported(self.name.to_string()))
        }
    }

    pub fn binary_file(&self, name: &str) -> PathBuf {
        let mut path = self.path.clone();
        path.push("bin");
        path.push(name.to_owned() + env::consts::EXE_SUFFIX);
        path
    }
}
<commit_msg>Fix for weird nightly resolution bug. cc rust-lang/rust#32106<commit_after>use errors::*;
use multirust_dist::{utils, dist};
use multirust_dist::prefix::InstallPrefix;
use multirust_dist::dist::ToolchainDesc;
use multirust_dist::manifestation::{Manifestation, Changes};
use multirust_dist::manifest::Component;
use config::Cfg;
use env_var;
use install::{self, InstallMethod};

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

use hyper;

#[derive(Debug)]
pub struct Toolchain<'a> {
    cfg: &'a Cfg,
    name: String,
    path: PathBuf,
}

impl<'a> Toolchain<'a> {
    pub fn from(cfg: &'a Cfg, name: &str) -> Self {
        let path = cfg.toolchains_dir.join(name);
        Toolchain {
            cfg: cfg,
            name: name.to_owned(),
            path: path.clone(),
        }
    }
    pub fn name(&self) -> &str {
        &self.name
    }
    pub fn path(&self) -> &Path {
        &self.path
    }
    pub fn exists(&self) -> bool {
        utils::is_directory(&self.path)
    }
    pub fn verify(&self) -> Result<()> {
        Ok(try!(utils::assert_is_directory(&self.path)))
    }
    pub fn remove(&self) -> Result<()> {
        if self.exists() {
            self.cfg.notify_handler.call(Notification::UninstallingToolchain(&self.name));
        } else {
            self.cfg.notify_handler.call(Notification::ToolchainNotInstalled(&self.name));
            return Ok(());
        }
        if let Some(update_hash) = try!(self.update_hash()) {
            try!(utils::remove_file("update hash", &update_hash));
        }
        let handler = self.cfg.notify_handler.as_ref();
        let result = install::uninstall(&self.path, ntfy!(&handler));
        if !self.exists() {
            self.cfg.notify_handler.call(Notification::UninstalledToolchain(&self.name));
        }
        Ok(try!(result))
    }
    fn install(&self, install_method: InstallMethod) -> Result<()> {
        assert!(self.is_valid_install_method(install_method));
        if self.exists() {
            self.cfg.notify_handler.call(Notification::UpdatingToolchain(&self.name));
        } else {
            self.cfg.notify_handler.call(Notification::InstallingToolchain(&self.name));
        }
        self.cfg
            .notify_handler
            .call(Notification::ToolchainDirectory(&self.path, &self.name));
        let handler = self.cfg.notify_handler.as_ref();
        Ok(try!(install_method.run(&self.path, ntfy!(&handler))))
    }
    fn install_if_not_installed(&self, install_method: InstallMethod) -> Result<()> {
        assert!(self.is_valid_install_method(install_method));
        self.cfg.notify_handler.call(Notification::LookingForToolchain(&self.name));
        if !self.exists() {
            self.install(install_method)
        } else {
            self.cfg.notify_handler.call(Notification::UsingExistingToolchain(&self.name));
            Ok(())
        }
    }
    fn is_valid_install_method(&self, install_method: InstallMethod) -> bool {
        match install_method {
            InstallMethod::Copy(_) |
            InstallMethod::Link(_) |
            InstallMethod::Installer(_, _) => self.is_custom(),
            InstallMethod::Dist(_, _, _) => !self.is_custom(),
        }
    }
    fn update_hash(&self) -> Result<Option<PathBuf>> {
        if self.is_custom() {
            Ok(None)
        } else {
            Ok(Some(try!(self.cfg.get_hash_file(&self.name, true))))
        }
    }

    fn download_cfg(&self) -> dist::DownloadCfg {
        dist::DownloadCfg {
            dist_root: &self.cfg.dist_root_url,
            temp_cfg: &self.cfg.temp_cfg,
            notify_handler: ntfy!(&self.cfg.notify_handler),
        }
    }

    pub fn install_from_dist(&self) -> Result<()> {
        let update_hash = try!(self.update_hash());
        self.install(InstallMethod::Dist(&self.name,
                                         update_hash.as_ref().map(|p| &**p),
                                         self.download_cfg()))
    }
    pub fn install_from_dist_if_not_installed(&self) -> Result<()> {
        let update_hash = try!(self.update_hash());
        self.install_if_not_installed(InstallMethod::Dist(&self.name,
                                                          update_hash.as_ref().map(|p| &**p),
                                                          self.download_cfg()))
    }
    pub fn is_custom(&self) -> bool {
        ToolchainDesc::from_str(&self.name).is_err()
    }
    pub fn is_tracking(&self) -> bool {
        ToolchainDesc::from_str(&self.name).ok().map(|d| d.is_tracking()) == Some(true)
    }

    fn ensure_custom(&self) -> Result<()> {
        if !self.is_custom() {
            Err(Error::Install(::multirust_dist::Error::InvalidToolchainName(self.name.to_string())))
        } else {
            Ok(())
        }
    }

    pub fn install_from_installers(&self, installers: &[&OsStr]) -> Result<()> {
        try!(self.ensure_custom());

        try!(self.remove());

        // FIXME: This should do all downloads first, then do
        // installs, and do it all in a single transaction.
        for installer in installers {
            let installer_str = installer.to_str().unwrap_or("bogus");
            match installer_str.rfind(".") {
                Some(i) => {
                    let extension = &installer_str[i+1..];
                    if extension != "gz" {
                        return Err(Error::BadInstallerType(extension.to_string()));
                    }
                }
                None => return Err(Error::BadInstallerType(String::from("(none)")))
            }

            // FIXME: Pretty hacky
            let is_url = installer_str.starts_with("file://")
                || installer_str.starts_with("http://")
                || installer_str.starts_with("https://");
            let url = hyper::Url::parse(installer_str).ok();
            let url = if is_url { url } else { None };
            if let Some(url) = url {

                // Download to a local file
                let local_installer = try!(self.cfg.temp_cfg.new_file_with_ext("", ".tar.gz"));
                try!(utils::download_file(url,
                                          &local_installer,
                                          None,
                                          ntfy!(&self.cfg.notify_handler)));
                try!(self.install(InstallMethod::Installer(&local_installer, &self.cfg.temp_cfg)));
            } else {
                // If installer is a filename

                // No need to download
                let local_installer = Path::new(installer);

                // Install from file
                try!(self.install(InstallMethod::Installer(&local_installer, &self.cfg.temp_cfg)));
            }
        }

        Ok(())
    }

    pub fn install_from_dir(&self, src: &Path, link: bool) -> Result<()> {
        try!(self.ensure_custom());

        if link {
            self.install(InstallMethod::Link(&try!(utils::to_absolute(src))))
        } else {
            self.install(InstallMethod::Copy(src))
        }
    }

    pub fn create_command<T: AsRef<OsStr>>(&self, binary: T) -> Result<Command> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        let mut cmd = Command::new(binary);
        self.set_env(&mut cmd);
        Ok(cmd)
    }

    fn set_env(&self, cmd: &mut Command) {
        let ref cargo_home = self.cfg.multirust_dir.join("cargo");
        let ref bin_path = self.path.join("bin");

        self.set_ldpath(cmd);

        env_var::set_path("PATH", bin_path, cmd);
        env_var::set_default("CARGO_HOME", cargo_home.as_ref(), cmd);
        env_var::inc("RUST_RECURSION_COUNT", cmd);

        // FIXME: This should not be a path, but a toolchain name.
        // Not sure what's going on here.
        cmd.env("MULTIRUST_TOOLCHAIN", &self.path);
        cmd.env("MULTIRUST_HOME", &self.cfg.multirust_dir);
    }

    pub fn set_ldpath(&self, cmd: &mut Command) {
        let new_path = self.path.join("lib");

        env_var::set_path("LD_LIBRARY_PATH", &new_path, cmd);
        env_var::set_path("DYLD_LIBRARY_PATH", &new_path, cmd);
    }

    pub fn doc_path(&self, relative: &str) -> Result<PathBuf> {
        try!(self.verify());

        let parts = vec!["share", "doc", "rust", "html"];
        let mut doc_dir = self.path.clone();
        for part in parts {
            doc_dir.push(part);
        }
        doc_dir.push(relative);

        Ok(doc_dir)
    }
    pub fn open_docs(&self, relative: &str) -> Result<()> {
        try!(self.verify());

        Ok(try!(utils::open_browser(&try!(self.doc_path(relative)))))
    }

    pub fn make_default(&self) -> Result<()> {
        self.cfg.set_default(&self.name)
    }
    pub fn make_override(&self, path: &Path) -> Result<()> {
        Ok(try!(self.cfg.override_db.set(path,
                                         &self.name,
                                         &self.cfg.temp_cfg,
                                         self.cfg.notify_handler.as_ref())))
    }

    pub fn list_components(&self) -> Result<Vec<Component>> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        // FIXME: This toolchain handling is a mess. Just do it once
        // when the toolchain is created.
        let ref toolchain = self.name;
        let ref toolchain = try!(ToolchainDesc::from_str(toolchain));
        let trip = toolchain.target_triple();
        let prefix = InstallPrefix::from(self.path.to_owned());
        let manifestation = try!(Manifestation::open(prefix, &trip));

        if let Some(manifest) = try!(manifestation.load_manifest()) {
            // Return all optional components of the "rust" package for the
            // toolchain's target triple.
            let mut res = Vec::new();

            let rust_pkg = manifest.packages.get("rust")
                .expect("manifest should cantain a rust package");
            let targ_pkg = rust_pkg.targets.get(&trip)
                .expect("installed manifest should have a known target");

            for extension in &targ_pkg.extensions {
                res.push(extension.clone())
            }

            Ok(res)
        } else {
            Err(Error::ComponentsUnsupported(self.name.to_string()))
        }
    }

    pub fn add_component(&self, component: Component) -> Result<()> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        let ref toolchain = self.name;
        let ref toolchain = try!(ToolchainDesc::from_str(toolchain));
        let trip = toolchain.target_triple();
        let prefix = InstallPrefix::from(self.path.to_owned());
        let manifestation = try!(Manifestation::open(prefix, &trip));

        if let Some(manifest) = try!(manifestation.load_manifest()) {

            // Validate the component name
            let rust_pkg = manifest.packages.get("rust")
                .expect("manifest should cantain a rust package");
            let targ_pkg = rust_pkg.targets.get(&trip)
                .expect("installed manifest should have a known target");

            if targ_pkg.components.contains(&component) {
                return Err(Error::AddingRequiredComponent(self.name.to_string(), component));
            }

            if !targ_pkg.extensions.contains(&component) {
                return Err(Error::UnknownComponent(self.name.to_string(), component));
            }

            let changes = Changes {
                add_extensions: vec![component],
                remove_extensions: vec![]
            };

            try!(manifestation.update(&self.name,
                                      &manifest,
                                      changes,
                                      self.download_cfg().temp_cfg,
                                      self.download_cfg().notify_handler.clone()));

            Ok(())
        } else {
            Err(Error::ComponentsUnsupported(self.name.to_string()))
        }
    }

    pub fn remove_component(&self, component: Component) -> Result<()> {
        if !self.exists() {
            return Err(Error::ToolchainNotInstalled(self.name.to_owned()));
        }

        let ref toolchain = self.name;
        let ref toolchain = try!(ToolchainDesc::from_str(toolchain));
        let trip = toolchain.target_triple();
        let prefix = InstallPrefix::from(self.path.to_owned());
        let manifestation = try!(Manifestation::open(prefix, &trip));

        if let Some(manifest) = try!(manifestation.load_manifest()) {

            // Validate the component name
            let rust_pkg = manifest.packages.get("rust")
                .expect("manifest should cantain a rust package");
            let targ_pkg = rust_pkg.targets.get(&trip)
                .expect("installed manifest should have a known target");

            if targ_pkg.components.contains(&component) {
                return Err(Error::RemovingRequiredComponent(self.name.to_string(), component));
            }

            let dist_config = try!(manifestation.read_config()).unwrap();
            if !dist_config.components.contains(&component) {
                return Err(Error::UnknownComponent(self.name.to_string(), component));
            }

            let changes = Changes {
                add_extensions: vec![],
                remove_extensions: vec![component]
            };

            try!(manifestation.update(&self.name,
                                      &manifest,
                                      changes,
                                      self.download_cfg().temp_cfg,
                                      self.download_cfg().notify_handler.clone()));

            Ok(())
        } else {
            Err(Error::ComponentsUnsupported(self.name.to_string()))
        }
    }

    pub fn binary_file(&self, name: &str) -> PathBuf {
        let mut path = self.path.clone();
        path.push("bin");
        path.push(name.to_owned() + env::consts::EXE_SUFFIX);
        path
    }
}
<|endoftext|>