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

use crate::{
    chained_bft::{
        block_storage::{BlockReader, BlockStore, InsertError},
        common::{Author, Payload},
        consensus_types::{block::Block, quorum_cert::QuorumCert},
        network::ConsensusNetworkImpl,
        persistent_storage::PersistentStorage,
    },
    counters,
    state_replication::StateComputer,
    state_synchronizer::SyncStatus,
    util::mutex_map::MutexMap,
};
use crypto::HashValue;
use failure::{Fail, Result};
use logger::prelude::*;
use network::proto::BlockRetrievalStatus;
use rand::{prelude::*, Rng};
use std::{
    clone::Clone,
    sync::Arc,
    time::{Duration, Instant},
};
use termion::color::*;
use types::{account_address::AccountAddress, transaction::TransactionListWithProof};

/// SyncManager is responsible for fetching dependencies and 'catching up' for given qc/ledger info
pub struct SyncManager<T> {
    block_store: Arc<BlockStore<T>>,
    storage: Arc<dyn PersistentStorage<T>>,
    network: ConsensusNetworkImpl,
    state_computer: Arc<dyn StateComputer<Payload = T>>,
    block_mutex_map: MutexMap<HashValue>,
}

/// This struct describes where do we sync to
pub struct SyncInfo {
    /// Highest ledger info to invoke state sync for
    /// This is optional for now, because vote does not have it
    pub highest_ledger_info: QuorumCert,
    /// Quorum certificate to be inserted into block tree
    pub highest_quorum_cert: QuorumCert,
    /// Author of messages that triggered this sync.
    /// For now we sync from this peer. In future we going to use peers from quorum certs,
    /// and this field going to be mostly informational
    pub peer: Author,
}

impl<T> SyncManager<T>
where
    T: Payload,
{
    pub fn new(
        block_store: Arc<BlockStore<T>>,
        storage: Arc<dyn PersistentStorage<T>>,
        network: ConsensusNetworkImpl,
        state_computer: Arc<dyn StateComputer<Payload = T>>,
    ) -> SyncManager<T> {
        // Our counters are initialized via lazy_static, so they're not going to appear in
        // Prometheus if some conditions never happen.  Invoking get() function enforces creation.
        counters::BLOCK_RETRIEVAL_COUNT.get();
        counters::STATE_SYNC_COUNT.get();
        counters::STATE_SYNC_TXN_REPLAYED.get();
        let block_mutex_map = MutexMap::new();
        SyncManager {
            block_store,
            storage,
            network,
            state_computer,
            block_mutex_map,
        }
    }

    /// Fetches dependencies for given sync_info.quorum_cert
    /// If gap is large, performs state sync using process_highest_ledger_info
    /// Inserts sync_info.quorum_cert into block store as the last step
    pub async fn sync_to(&mut self, deadline: Instant, sync_info: SyncInfo) -> Result<()> {
        let highest_ledger_info = sync_info.highest_ledger_info.clone();

        self.process_highest_ledger_info(highest_ledger_info, sync_info.peer, deadline)
            .await?;

        self.fetch_quorum_cert(
            sync_info.highest_quorum_cert.clone(),
            sync_info.peer,
            deadline,
        )
        .await?;
        Ok(())
    }

    /// Get a chunk of transactions as a batch
    pub async fn get_chunk(
        &self,
        start_version: u64,
        target_version: u64,
        batch_size: u64,
    ) -> Result<TransactionListWithProof> {
        self.state_computer
            .get_chunk(start_version, target_version, batch_size)
            .await
    }

    pub async fn execute_and_insert_block(
        &self,
        block: Block<T>,
    ) -> std::result::Result<Arc<Block<T>>, InsertError> {
        let _guard = self.block_mutex_map.lock(block.id());
        // execute_and_insert_block has shortcut to return block if it exists
        self.block_store.execute_and_insert_block(block).await
    }

    /// Insert the quorum certificate separately from the block, used to split the processing of
    /// updating the consensus state(with qc) and deciding whether to vote(with block)
    /// The missing ancestors are going to be retrieved from the given peer. If a given peer
    /// fails to provide the missing ancestors, the qc is not going to be added.
    pub async fn fetch_quorum_cert(
        &self,
        qc: QuorumCert,
        preferred_peer: Author,
        deadline: Instant,
    ) -> std::result::Result<(), InsertError> {
        let mut lock_set = self.block_mutex_map.new_lock_set();
        let mut pending = vec![];
        let network = self.network.clone();
        let mut retriever = BlockRetriever {
            network,
            deadline,
            preferred_peer,
        };
        let mut retrieve_qc = qc.clone();
        loop {
            if lock_set
                .lock(retrieve_qc.certified_block_id())
                .await
                .is_err()
            {
                // This should not be possible because that would mean we have circular
                // dependency between signed blocks
                panic!(
                    "Can not re-acquire lock for block {} during fetch_quorum_cert",
                    retrieve_qc.certified_block_id()
                );
            }
            if self
                .block_store
                .block_exists(retrieve_qc.certified_block_id())
            {
                break;
            }
            let mut blocks = retriever.retrieve_block_for_qc(&retrieve_qc, 1).await?;
            // retrieve_block_for_qc guarantees that blocks has exactly 1 element
            let block = blocks.remove(0);
            retrieve_qc = block.quorum_cert().clone();
            pending.push(block);
        }
        // insert the qc <- block pair
        while let Some(block) = pending.pop() {
            let block_qc = block.quorum_cert().clone();
            self.block_store.insert_single_quorum_cert(block_qc).await?;
            self.block_store.execute_and_insert_block(block).await?;
        }
        self.block_store.insert_single_quorum_cert(qc).await
    }

    /// Check the highest ledger info sent by peer to see if we're behind and start a fast
    /// forward sync if the committed block doesn't exist in our tree.
    /// It works as follows:
    /// 1. request the committed 3-chain from the peer, if C2 is the highest_ledger_info
    /// we request for B0 <- C0 <- B1 <- C1 <- B2 (<- C2)
    /// 2. We persist the 3-chain to storage before start sync to ensure we could restart if we
    /// crash in the middle of the sync.
    /// 3. We prune the old tree and replace with a new tree built with the 3-chain.
    async fn process_highest_ledger_info(
        &self,
        highest_ledger_info: QuorumCert,
        peer: Author,
        deadline: Instant,
    ) -> Result<()> {
        let committed_block_id = highest_ledger_info
            .committed_block_id()
            .ok_or_else(|| format_err!("highest ledger info has no committed block"))?;
        if !self
            .block_store
            .need_sync_for_quorum_cert(committed_block_id, &highest_ledger_info)
        {
            return Ok(());
        }
        debug!(
            "Start state sync with peer: {}, to block: {}, round: {} from {}",
            peer.short_str(),
            committed_block_id,
            highest_ledger_info.certified_block_round() - 2,
            self.block_store.root()
        );
        let network = self.network.clone();
        let mut retriever = BlockRetriever {
            network,
            deadline,
            preferred_peer: peer,
        };
        let mut blocks = retriever
            .retrieve_block_for_qc(&highest_ledger_info, 3)
            .await?;
        assert_eq!(
            blocks.last().expect("should have 3-chain").id(),
            committed_block_id
        );
        let mut quorum_certs = vec![];
        quorum_certs.push(highest_ledger_info.clone());
        quorum_certs.push(blocks[0].quorum_cert().clone());
        quorum_certs.push(blocks[1].quorum_cert().clone());
        // If a node restarts in the middle of state synchronization, it is going to try to catch up
        // to the stored quorum certs as the new root.
        self.storage
            .save_tree(blocks.clone(), quorum_certs.clone())?;
        let pre_sync_instance = Instant::now();
        match self
            .state_computer
            .sync_to(highest_ledger_info.clone())
            .await
        {
            Ok(SyncStatus::Finished) => (),
            Ok(e) => panic!(
                "state synchronizer failure: {:?}, this validator will be killed as it can not \
                 recover from this error.  After the validator is restarted, synchronization will \
                 be retried.",
                e
            ),
            Err(e) => panic!(
                "state synchronizer failure: {:?}, this validator will be killed as it can not \
                 recover from this error.  After the validator is restarted, synchronization will \
                 be retried.",
                e
            ),
        };
        counters::STATE_SYNC_DURATION_MS.observe(pre_sync_instance.elapsed().as_millis() as f64);
        let root = (
            blocks.pop().expect("should have 3-chain"),
            quorum_certs.last().expect("should have 3-chain").clone(),
            highest_ledger_info.clone(),
        );
        debug!("{}Sync to{} {}", Fg(Blue), Fg(Reset), root.0);
        // ensure it's [b1, b2]
        blocks.reverse();
        self.block_store.rebuild(root, blocks, quorum_certs).await;
        Ok(())
    }
}

/// BlockRetriever is used internally to retrieve blocks
struct BlockRetriever {
    network: ConsensusNetworkImpl,
    deadline: Instant,
    preferred_peer: Author,
}

#[derive(Debug, Fail)]
enum BlockRetrieverError {
    #[fail(display = "All peers failed")]
    AllPeersFailed,
    #[fail(display = "Round deadline reached")]
    RoundDeadlineReached,
}

impl From<BlockRetrieverError> for InsertError {
    fn from(_error: BlockRetrieverError) -> Self {
        InsertError::AncestorRetrievalError
    }
}

impl BlockRetriever {
    /// Retrieve chain of n blocks for given QC
    ///
    /// Returns Result with Vec that has a guaranteed size of num_blocks
    /// This guarantee is based on BlockRetrievalResponse::verify that ensures that number of
    /// blocks in response is equal to number of blocks requested.  This method will
    /// continue until either the round deadline is reached or the quorum certificate members all
    /// fail to return the missing chain.
    ///
    /// The first attempt of block retrieval will always be sent to preferred_peer to allow the
    /// leader to drive quorum certificate creation The other peers from the quorum certificate
    /// will be randomly tried next.  If all members of the quorum certificate are exhausted, an
    /// error is returned
    pub async fn retrieve_block_for_qc<'a, T>(
        &'a mut self,
        qc: &'a QuorumCert,
        num_blocks: u64,
    ) -> std::result::Result<Vec<Block<T>>, BlockRetrieverError>
    where
        T: Payload,
    {
        let block_id = qc.certified_block_id();
        let mut peers: Vec<&AccountAddress> = qc.ledger_info().signatures().keys().collect();
        let mut attempt = 0_u32;
        loop {
            if peers.is_empty() {
                warn!(
                    "Failed to fetch block {} in {} attempts: no more peers available",
                    block_id, attempt
                );
                return Err(BlockRetrieverError::AllPeersFailed);
            }
            let peer = self.pick_peer(attempt, &mut peers);
            attempt += 1;

            let timeout = retrieval_timeout(&self.deadline, attempt);
            let timeout = if let Some(timeout) = timeout {
                timeout
            } else {
                warn!("Failed to fetch block {} from {}, attempt {}: round deadline was reached, won't make more attempts", block_id, peer, attempt);
                return Err(BlockRetrieverError::RoundDeadlineReached);
            };
            debug!(
                "Fetching {} from {}, attempt {}",
                block_id,
                peer.short_str(),
                attempt
            );
            let response = self
                .network
                .request_block(block_id, num_blocks, peer, timeout)
                .await;
            let response = match response {
                Err(e) => {
                    warn!(
                        "Failed to fetch block {} from {}: {:?}, trying another peer",
                        block_id,
                        peer.short_str(),
                        e
                    );
                    continue;
                }
                Ok(response) => response,
            };
            if response.status != BlockRetrievalStatus::SUCCEEDED {
                warn!(
                    "Failed to fetch block {} from {}: {:?}, trying another peer",
                    block_id,
                    peer.short_str(),
                    response.status
                );
                continue;
            }
            return Ok(response.blocks);
        }
    }

    fn pick_peer(&self, attempt: u32, peers: &mut Vec<&AccountAddress>) -> AccountAddress {
        assert!(!peers.is_empty(), "pick_peer on empty peer list");

        if attempt == 0 {
            // remove preferred_peer if its in list of peers
            // (strictly speaking it is not required to be there)
            for i in 0..peers.len() {
                if *peers[i] == self.preferred_peer {
                    peers.remove(i);
                    break;
                }
            }
            return self.preferred_peer;
        }
        let peer_idx = thread_rng().gen_range(0, peers.len());
        *peers.remove(peer_idx)
    }
}

// Max timeout is 16s=RETRIEVAL_INITIAL_TIMEOUT*(2^RETRIEVAL_MAX_EXP)
const RETRIEVAL_INITIAL_TIMEOUT: Duration = Duration::from_secs(1);
const RETRIEVAL_MAX_EXP: u32 = 4;

/// Returns exponentially increasing timeout with
/// limit of RETRIEVAL_INITIAL_TIMEOUT*(2^RETRIEVAL_MAX_EXP)
fn retrieval_timeout(deadline: &Instant, attempt: u32) -> Option<Duration> {
    assert!(attempt > 0, "retrieval_timeout attempt can't be 0");
    let exp = RETRIEVAL_MAX_EXP.min(attempt - 1); // [0..RETRIEVAL_MAX_EXP]
    let request_timeout = RETRIEVAL_INITIAL_TIMEOUT * 2_u32.pow(exp);
    let deadline_timeout = deadline.checked_duration_since(Instant::now());
    if let Some(deadline_timeout) = deadline_timeout {
        Some(request_timeout.min(deadline_timeout))
    } else {
        None
    }
}
<commit_msg>[Consensus] SyncManager - import std::Result instead of failure::Result<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    chained_bft::{
        block_storage::{BlockReader, BlockStore, InsertError},
        common::{Author, Payload},
        consensus_types::{block::Block, quorum_cert::QuorumCert},
        network::ConsensusNetworkImpl,
        persistent_storage::PersistentStorage,
    },
    counters,
    state_replication::StateComputer,
    state_synchronizer::SyncStatus,
    util::mutex_map::MutexMap,
};
use crypto::HashValue;
use failure::{self, Fail};
use logger::prelude::*;
use network::proto::BlockRetrievalStatus;
use rand::{prelude::*, Rng};
use std::{
    clone::Clone,
    result::Result,
    sync::Arc,
    time::{Duration, Instant},
};
use termion::color::*;
use types::{account_address::AccountAddress, transaction::TransactionListWithProof};

/// SyncManager is responsible for fetching dependencies and 'catching up' for given qc/ledger info
pub struct SyncManager<T> {
    block_store: Arc<BlockStore<T>>,
    storage: Arc<dyn PersistentStorage<T>>,
    network: ConsensusNetworkImpl,
    state_computer: Arc<dyn StateComputer<Payload = T>>,
    block_mutex_map: MutexMap<HashValue>,
}

/// This struct describes where do we sync to
pub struct SyncInfo {
    /// Highest ledger info to invoke state sync for
    /// This is optional for now, because vote does not have it
    pub highest_ledger_info: QuorumCert,
    /// Quorum certificate to be inserted into block tree
    pub highest_quorum_cert: QuorumCert,
    /// Author of messages that triggered this sync.
    /// For now we sync from this peer. In future we going to use peers from quorum certs,
    /// and this field going to be mostly informational
    pub peer: Author,
}

impl<T> SyncManager<T>
where
    T: Payload,
{
    pub fn new(
        block_store: Arc<BlockStore<T>>,
        storage: Arc<dyn PersistentStorage<T>>,
        network: ConsensusNetworkImpl,
        state_computer: Arc<dyn StateComputer<Payload = T>>,
    ) -> SyncManager<T> {
        // Our counters are initialized via lazy_static, so they're not going to appear in
        // Prometheus if some conditions never happen.  Invoking get() function enforces creation.
        counters::BLOCK_RETRIEVAL_COUNT.get();
        counters::STATE_SYNC_COUNT.get();
        counters::STATE_SYNC_TXN_REPLAYED.get();
        let block_mutex_map = MutexMap::new();
        SyncManager {
            block_store,
            storage,
            network,
            state_computer,
            block_mutex_map,
        }
    }

    /// Fetches dependencies for given sync_info.quorum_cert
    /// If gap is large, performs state sync using process_highest_ledger_info
    /// Inserts sync_info.quorum_cert into block store as the last step
    pub async fn sync_to(&mut self, deadline: Instant, sync_info: SyncInfo) -> failure::Result<()> {
        let highest_ledger_info = sync_info.highest_ledger_info.clone();

        self.process_highest_ledger_info(highest_ledger_info, sync_info.peer, deadline)
            .await?;

        self.fetch_quorum_cert(
            sync_info.highest_quorum_cert.clone(),
            sync_info.peer,
            deadline,
        )
        .await?;
        Ok(())
    }

    /// Get a chunk of transactions as a batch
    pub async fn get_chunk(
        &self,
        start_version: u64,
        target_version: u64,
        batch_size: u64,
    ) -> failure::Result<TransactionListWithProof> {
        self.state_computer
            .get_chunk(start_version, target_version, batch_size)
            .await
    }

    pub async fn execute_and_insert_block(
        &self,
        block: Block<T>,
    ) -> Result<Arc<Block<T>>, InsertError> {
        let _guard = self.block_mutex_map.lock(block.id());
        // execute_and_insert_block has shortcut to return block if it exists
        self.block_store.execute_and_insert_block(block).await
    }

    /// Insert the quorum certificate separately from the block, used to split the processing of
    /// updating the consensus state(with qc) and deciding whether to vote(with block)
    /// The missing ancestors are going to be retrieved from the given peer. If a given peer
    /// fails to provide the missing ancestors, the qc is not going to be added.
    pub async fn fetch_quorum_cert(
        &self,
        qc: QuorumCert,
        preferred_peer: Author,
        deadline: Instant,
    ) -> Result<(), InsertError> {
        let mut lock_set = self.block_mutex_map.new_lock_set();
        let mut pending = vec![];
        let network = self.network.clone();
        let mut retriever = BlockRetriever {
            network,
            deadline,
            preferred_peer,
        };
        let mut retrieve_qc = qc.clone();
        loop {
            if lock_set
                .lock(retrieve_qc.certified_block_id())
                .await
                .is_err()
            {
                // This should not be possible because that would mean we have circular
                // dependency between signed blocks
                panic!(
                    "Can not re-acquire lock for block {} during fetch_quorum_cert",
                    retrieve_qc.certified_block_id()
                );
            }
            if self
                .block_store
                .block_exists(retrieve_qc.certified_block_id())
            {
                break;
            }
            let mut blocks = retriever.retrieve_block_for_qc(&retrieve_qc, 1).await?;
            // retrieve_block_for_qc guarantees that blocks has exactly 1 element
            let block = blocks.remove(0);
            retrieve_qc = block.quorum_cert().clone();
            pending.push(block);
        }
        // insert the qc <- block pair
        while let Some(block) = pending.pop() {
            let block_qc = block.quorum_cert().clone();
            self.block_store.insert_single_quorum_cert(block_qc).await?;
            self.block_store.execute_and_insert_block(block).await?;
        }
        self.block_store.insert_single_quorum_cert(qc).await
    }

    /// Check the highest ledger info sent by peer to see if we're behind and start a fast
    /// forward sync if the committed block doesn't exist in our tree.
    /// It works as follows:
    /// 1. request the committed 3-chain from the peer, if C2 is the highest_ledger_info
    /// we request for B0 <- C0 <- B1 <- C1 <- B2 (<- C2)
    /// 2. We persist the 3-chain to storage before start sync to ensure we could restart if we
    /// crash in the middle of the sync.
    /// 3. We prune the old tree and replace with a new tree built with the 3-chain.
    async fn process_highest_ledger_info(
        &self,
        highest_ledger_info: QuorumCert,
        peer: Author,
        deadline: Instant,
    ) -> failure::Result<()> {
        let committed_block_id = highest_ledger_info
            .committed_block_id()
            .ok_or_else(|| format_err!("highest ledger info has no committed block"))?;
        if !self
            .block_store
            .need_sync_for_quorum_cert(committed_block_id, &highest_ledger_info)
        {
            return Ok(());
        }
        debug!(
            "Start state sync with peer: {}, to block: {}, round: {} from {}",
            peer.short_str(),
            committed_block_id,
            highest_ledger_info.certified_block_round() - 2,
            self.block_store.root()
        );
        let network = self.network.clone();
        let mut retriever = BlockRetriever {
            network,
            deadline,
            preferred_peer: peer,
        };
        let mut blocks = retriever
            .retrieve_block_for_qc(&highest_ledger_info, 3)
            .await?;
        assert_eq!(
            blocks.last().expect("should have 3-chain").id(),
            committed_block_id
        );
        let mut quorum_certs = vec![];
        quorum_certs.push(highest_ledger_info.clone());
        quorum_certs.push(blocks[0].quorum_cert().clone());
        quorum_certs.push(blocks[1].quorum_cert().clone());
        // If a node restarts in the middle of state synchronization, it is going to try to catch up
        // to the stored quorum certs as the new root.
        self.storage
            .save_tree(blocks.clone(), quorum_certs.clone())?;
        let pre_sync_instance = Instant::now();
        match self
            .state_computer
            .sync_to(highest_ledger_info.clone())
            .await
        {
            Ok(SyncStatus::Finished) => (),
            Ok(e) => panic!(
                "state synchronizer failure: {:?}, this validator will be killed as it can not \
                 recover from this error.  After the validator is restarted, synchronization will \
                 be retried.",
                e
            ),
            Err(e) => panic!(
                "state synchronizer failure: {:?}, this validator will be killed as it can not \
                 recover from this error.  After the validator is restarted, synchronization will \
                 be retried.",
                e
            ),
        };
        counters::STATE_SYNC_DURATION_MS.observe(pre_sync_instance.elapsed().as_millis() as f64);
        let root = (
            blocks.pop().expect("should have 3-chain"),
            quorum_certs.last().expect("should have 3-chain").clone(),
            highest_ledger_info.clone(),
        );
        debug!("{}Sync to{} {}", Fg(Blue), Fg(Reset), root.0);
        // ensure it's [b1, b2]
        blocks.reverse();
        self.block_store.rebuild(root, blocks, quorum_certs).await;
        Ok(())
    }
}

/// BlockRetriever is used internally to retrieve blocks
struct BlockRetriever {
    network: ConsensusNetworkImpl,
    deadline: Instant,
    preferred_peer: Author,
}

#[derive(Debug, Fail)]
enum BlockRetrieverError {
    #[fail(display = "All peers failed")]
    AllPeersFailed,
    #[fail(display = "Round deadline reached")]
    RoundDeadlineReached,
}

impl From<BlockRetrieverError> for InsertError {
    fn from(_error: BlockRetrieverError) -> Self {
        InsertError::AncestorRetrievalError
    }
}

impl BlockRetriever {
    /// Retrieve chain of n blocks for given QC
    ///
    /// Returns Result with Vec that has a guaranteed size of num_blocks
    /// This guarantee is based on BlockRetrievalResponse::verify that ensures that number of
    /// blocks in response is equal to number of blocks requested.  This method will
    /// continue until either the round deadline is reached or the quorum certificate members all
    /// fail to return the missing chain.
    ///
    /// The first attempt of block retrieval will always be sent to preferred_peer to allow the
    /// leader to drive quorum certificate creation The other peers from the quorum certificate
    /// will be randomly tried next.  If all members of the quorum certificate are exhausted, an
    /// error is returned
    pub async fn retrieve_block_for_qc<'a, T>(
        &'a mut self,
        qc: &'a QuorumCert,
        num_blocks: u64,
    ) -> Result<Vec<Block<T>>, BlockRetrieverError>
    where
        T: Payload,
    {
        let block_id = qc.certified_block_id();
        let mut peers: Vec<&AccountAddress> = qc.ledger_info().signatures().keys().collect();
        let mut attempt = 0_u32;
        loop {
            if peers.is_empty() {
                warn!(
                    "Failed to fetch block {} in {} attempts: no more peers available",
                    block_id, attempt
                );
                return Err(BlockRetrieverError::AllPeersFailed);
            }
            let peer = self.pick_peer(attempt, &mut peers);
            attempt += 1;

            let timeout = retrieval_timeout(&self.deadline, attempt);
            let timeout = if let Some(timeout) = timeout {
                timeout
            } else {
                warn!("Failed to fetch block {} from {}, attempt {}: round deadline was reached, won't make more attempts", block_id, peer, attempt);
                return Err(BlockRetrieverError::RoundDeadlineReached);
            };
            debug!(
                "Fetching {} from {}, attempt {}",
                block_id,
                peer.short_str(),
                attempt
            );
            let response = self
                .network
                .request_block(block_id, num_blocks, peer, timeout)
                .await;
            let response = match response {
                Err(e) => {
                    warn!(
                        "Failed to fetch block {} from {}: {:?}, trying another peer",
                        block_id,
                        peer.short_str(),
                        e
                    );
                    continue;
                }
                Ok(response) => response,
            };
            if response.status != BlockRetrievalStatus::SUCCEEDED {
                warn!(
                    "Failed to fetch block {} from {}: {:?}, trying another peer",
                    block_id,
                    peer.short_str(),
                    response.status
                );
                continue;
            }
            return Ok(response.blocks);
        }
    }

    fn pick_peer(&self, attempt: u32, peers: &mut Vec<&AccountAddress>) -> AccountAddress {
        assert!(!peers.is_empty(), "pick_peer on empty peer list");

        if attempt == 0 {
            // remove preferred_peer if its in list of peers
            // (strictly speaking it is not required to be there)
            for i in 0..peers.len() {
                if *peers[i] == self.preferred_peer {
                    peers.remove(i);
                    break;
                }
            }
            return self.preferred_peer;
        }
        let peer_idx = thread_rng().gen_range(0, peers.len());
        *peers.remove(peer_idx)
    }
}

// Max timeout is 16s=RETRIEVAL_INITIAL_TIMEOUT*(2^RETRIEVAL_MAX_EXP)
const RETRIEVAL_INITIAL_TIMEOUT: Duration = Duration::from_secs(1);
const RETRIEVAL_MAX_EXP: u32 = 4;

/// Returns exponentially increasing timeout with
/// limit of RETRIEVAL_INITIAL_TIMEOUT*(2^RETRIEVAL_MAX_EXP)
fn retrieval_timeout(deadline: &Instant, attempt: u32) -> Option<Duration> {
    assert!(attempt > 0, "retrieval_timeout attempt can't be 0");
    let exp = RETRIEVAL_MAX_EXP.min(attempt - 1); // [0..RETRIEVAL_MAX_EXP]
    let request_timeout = RETRIEVAL_INITIAL_TIMEOUT * 2_u32.pow(exp);
    let deadline_timeout = deadline.checked_duration_since(Instant::now());
    if let Some(deadline_timeout) = deadline_timeout {
        Some(request_timeout.min(deadline_timeout))
    } else {
        None
    }
}
<|endoftext|>


<commit_before>use libc::*;

use std::io::process::{Command, ProcessOutput, ProcessExit, Process};
use std::io::process::StdioContainer::*;
use std::io::{IoError, IoErrorKind, IoResult};
use std::io::IoErrorKind::*;
use std::collections::VecMap;

use controls::*;
use constants::*;
use termios::*;
use signal::*;

// start off as null pointer
static mut uglobal_term:*mut TermState = 0 as *mut TermState;

unsafe extern fn term_sigint(_:c_int, _:*const SigInfo,
                             _:*const c_void) {
    let term:&mut TermState = match uglobal_term.as_mut() {
        Some(v) => v,
        None => {
            // this handler shouldn't be called when Term isn't active
            panic!("Term signal interrupt called when Term not active");
        }
    };
    term.controls.outs("\nInterrupt\n");
    // pass on to foreground job, if there is one
    match term.fg_job {
        None => {/* nothing */},
        Some(id) => {
            match term.interrupt_job(&id) {
                Err(e) => {
                    term.controls.errf(format_args!("Could not interrupt job: {}", e));
                },
                Ok(_) => {/* nothing */}
            }
        }
    }
}

pub struct TermState {
    pub controls: Controls,
    tios: Termios,
    old_tios: Termios,
    jobs: VecMap<(String, Process)>,
    fg_job: Option<usize>
}

impl TermState {
    pub fn new() -> TermState {
        let mut controls = Controls::new();
        let mut tios = match Termios::get() {
            Some(t) => t,
            None => {
                controls.err("Warning: Could not get terminal mode\n");
                Termios::new()
            }
        };
        let old_tios = tios.clone();
        tios.fdisable(0, 0, ICANON|ECHO, 0);

        return TermState {
            controls: controls,
            tios: tios,
            old_tios: old_tios,
            jobs: VecMap::new(),
            fg_job: None
        };
    }

    pub fn update_terminal(&mut self) {
        if !Termios::set(&self.tios) {
            self.controls.err("Warning: Could not set terminal mode\n");
        }
    }

    pub fn restore_terminal(&mut self) {
        if !Termios::set(&self.old_tios) {
            self.controls.err("Warning: Could not set terminal mode\n");
        }
    }

    fn set_pointer(&mut self) {
        unsafe {
            if !uglobal_term.is_null() {
                panic!("Tried to set Term location twice");
            }
            uglobal_term = self as *mut TermState;
        }
    }

    fn unset_pointer(&mut self) {
        unsafe {
            if uglobal_term.is_null() {
                panic!("Tried to unset Term location twice");
            }
            uglobal_term = 0 as *mut TermState;
        }
    }

    fn handle_sigint(&mut self) {
        self.set_pointer();
        let mut sa = SigAction {
            handler: term_sigint,
            mask: [0; SIGSET_NWORDS],
            flags: SA_RESTART | SA_SIGINFO,
            restorer: 0 // null pointer
        };
        let mask = full_sigset().expect("Could not get a full sigset");
        sa.mask = mask;
        if !signal_handle(SIGINT, &sa) {
            self.controls.err("Warning: could not set handler for SIGINT\n");
        }
    }

    fn unhandle_sigint(&mut self) {
        self.unset_pointer();
        if !signal_ignore(SIGINT) {
            self.controls.err("Warning: could not unset handler for SIGINT\n");
        }
    }

    pub fn interrupt_job(&mut self, id:&usize) -> IoResult<()> {
        match self.jobs.get_mut(id) {
            None => return Err(IoError {
                kind: OtherIoError,
                desc: "Job not found",
                detail: None
            }),
            Some(&mut (_, ref mut job)) => {
                return job.signal(SIGINT as isize);
            }
        }
    }

    fn find_jobs_hole(&self) -> usize {
        // find a hole in the job map
        if self.jobs.len() < 20 {
            // with less than twenty jobs just be lazy
            return self.jobs.len() + 1;
        }
        let mut last = 0;
        for key in self.jobs.keys() {
            if key - last != 1 {
                // we've found a hole
                return key - 1;
            } else {
                last = key;
            }
        }
        // job list is full
        return last + 1;
    }

    pub fn run_job(&mut self, name:&String, args:&Vec<String>) -> Result<(usize, String), String> {
        let mut process = Command::new(name);
        process.args(args.as_slice());
        // running as job means no input/output handles
        process.stdin(Ignored);
        process.stdout(Ignored);
        process.stderr(Ignored);
        // all others are redirected
        let mut child = match process.spawn() {
            Err(e) => {
                return Err(format!("Couldn't spawn {}: {}", name, e));
            },
            Ok(v) => v
        };
        // set wait timeout to zero so we can check for process exit quickly
        child.set_timeout(Some(0));
        let id = self.find_jobs_hole();
        // panic if we've overwritten a job
        match self.jobs.insert(id.clone(), (name.clone(), child)) {
            Some(_) => panic!("Overwrote job"),
            _ => {/* nothing */}
        }
        return Ok((id, name.clone()));
    }

    pub fn get_jobs(&mut self) -> Vec<(usize, String, &Process)> {
        let mut out = vec![];
        for (id, &(ref name, ref child)) in self.jobs.iter() {
            out.push((id, name.clone(), child));
        }
        return out;
    }

    pub fn clean_jobs(&mut self) -> Vec<(usize, String, IoResult<ProcessExit>)> {
        let mut out = vec![];
        let mut remove = vec![];
        for (id, &mut (ref mut name, ref mut child)) in self.jobs.iter_mut() {
             match child.wait() {
                 Err(IoError{kind: IoErrorKind::TimedOut, desc: _, detail: _}) => {
                     // this is expected, do nothing
                 },
                 v => {
                     // all other outputs mean the job is done
                     // if it isn't it'll be cleaned up in drop
                     child.set_timeout(None);
                     remove.push(id);
                     out.push((id, name.clone(), v));
                 }
             }
        }
        for id in remove.iter() {
            self.jobs.remove(id);
        }
        return out;
    }

    pub fn run_command(&mut self, name:&String, args:&Vec<String>) -> Result<ProcessExit, String> {
        let mut process = Command::new(name);
        process.args(args.as_slice());
        process.stdout(InheritFd(STDOUT));
        process.stdin(InheritFd(STDIN));
        process.stderr(InheritFd(STDERR));
        // set terminal settings for process
        self.restore_terminal();
        // push job into jobs
        let id = self.find_jobs_hole();
        // handle interrupts
        self.handle_sigint();
        let val = (name.clone(), match process.spawn() {
            Err(e) => {
                self.update_terminal();
                return Err(format!("Couldn't spawn {}: {}", name, e));
            },
            Ok(v) => v
        });
        // insert job
        self.jobs.insert(id, val);
        // set forground job
        self.fg_job = Some(id);
        let out = match self.jobs.get_mut(&id).unwrap().1.wait() {
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                self.update_terminal();
                return Err(format!("Couldn't wait for child to exit: {}", e));
            },
            Ok(v) => v
        };
        // unset forground job
        self.fg_job = None;
        // remove job
        self.jobs.remove(&id);
        // unhandle sigint
        self.unhandle_sigint();
        // restore settings for Wash
        self.update_terminal();
        return Ok(out);
    }

    pub fn run_command_directed(&mut self, name:&String,
                                args:&Vec<String>) -> Result<ProcessOutput, String> {
        let mut process = Command::new(name);
        process.args(args.as_slice());
        let id = self.find_jobs_hole();
        // handle sigint
        self.handle_sigint();
        let val = (name.clone(), match process.spawn() {
            Err(e) => {
                self.update_terminal();
                return Err(format!("Couldn't spawn {}: {}", name, e));
            },
            Ok(v) => v
        });
        // put job in jobs list
        self.jobs.insert(id, val);
        // set foreground job
        self.fg_job = Some(id);
        let out = match self.jobs.get_mut(&id).unwrap().1.stdout.as_mut().unwrap().read_to_end() {
            Ok(v) => v,
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                return Err(format!("Couldn't get stdout: {}", e));
            }
        };
        let err = match self.jobs.get_mut(&id).unwrap().1.stderr.as_mut().unwrap().read_to_end() {
            Ok(v) => v,
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                return Err(format!("Couldn't get stderr: {}", e));
            }
        };
        let output = match self.jobs.get_mut(&id).unwrap().1.wait() {
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                self.update_terminal();
                return Err(format!("Couldn't wait for child to exit: {}", e));
            },
            Ok(v) => ProcessOutput {
                status: v,
                output: out,
                error: err
            }
        };
        // unset foreground job
        self.fg_job = None;
        // remove job
        self.jobs.remove(&id);
        // unhandle sigint
        self.unhandle_sigint();
        return Ok(output);
    }
}

<commit_msg>Delete ^C printed by kernel on SIGINT<commit_after>use libc::*;

use std::io::process::{Command, ProcessOutput, ProcessExit, Process};
use std::io::process::StdioContainer::*;
use std::io::{IoError, IoErrorKind, IoResult};
use std::io::IoErrorKind::*;
use std::collections::VecMap;

use controls::*;
use constants::*;
use termios::*;
use signal::*;

// start off as null pointer
static mut uglobal_term:*mut TermState = 0 as *mut TermState;

unsafe extern fn term_sigint(_:c_int, _:*const SigInfo,
                             _:*const c_void) {
    let term:&mut TermState = match uglobal_term.as_mut() {
        Some(v) => v,
        None => {
            // this handler shouldn't be called when Term isn't active
            panic!("Term signal interrupt called when Term not active");
        }
    };
    // delete "^C"
    term.controls.outc(BS);
    term.controls.outc(BS);
    term.controls.outc(SPC);
    term.controls.outc(SPC);
    term.controls.outs("\nInterrupt\n");
    // pass on to foreground job, if there is one
    match term.fg_job {
        None => {/* nothing */},
        Some(id) => {
            match term.interrupt_job(&id) {
                Err(e) => {
                    term.controls.errf(format_args!("Could not interrupt job: {}", e));
                },
                Ok(_) => {/* nothing */}
            }
        }
    }
}

pub struct TermState {
    pub controls: Controls,
    tios: Termios,
    old_tios: Termios,
    jobs: VecMap<(String, Process)>,
    fg_job: Option<usize>
}

impl TermState {
    pub fn new() -> TermState {
        let mut controls = Controls::new();
        let mut tios = match Termios::get() {
            Some(t) => t,
            None => {
                controls.err("Warning: Could not get terminal mode\n");
                Termios::new()
            }
        };
        let old_tios = tios.clone();
        tios.fdisable(0, 0, ICANON|ECHO, 0);

        return TermState {
            controls: controls,
            tios: tios,
            old_tios: old_tios,
            jobs: VecMap::new(),
            fg_job: None
        };
    }

    pub fn update_terminal(&mut self) {
        if !Termios::set(&self.tios) {
            self.controls.err("Warning: Could not set terminal mode\n");
        }
    }

    pub fn restore_terminal(&mut self) {
        if !Termios::set(&self.old_tios) {
            self.controls.err("Warning: Could not set terminal mode\n");
        }
    }

    fn set_pointer(&mut self) {
        unsafe {
            if !uglobal_term.is_null() {
                panic!("Tried to set Term location twice");
            }
            uglobal_term = self as *mut TermState;
        }
    }

    fn unset_pointer(&mut self) {
        unsafe {
            if uglobal_term.is_null() {
                panic!("Tried to unset Term location twice");
            }
            uglobal_term = 0 as *mut TermState;
        }
    }

    fn handle_sigint(&mut self) {
        self.set_pointer();
        let mut sa = SigAction {
            handler: term_sigint,
            mask: [0; SIGSET_NWORDS],
            flags: SA_RESTART | SA_SIGINFO,
            restorer: 0 // null pointer
        };
        let mask = full_sigset().expect("Could not get a full sigset");
        sa.mask = mask;
        if !signal_handle(SIGINT, &sa) {
            self.controls.err("Warning: could not set handler for SIGINT\n");
        }
    }

    fn unhandle_sigint(&mut self) {
        self.unset_pointer();
        if !signal_ignore(SIGINT) {
            self.controls.err("Warning: could not unset handler for SIGINT\n");
        }
    }

    pub fn interrupt_job(&mut self, id:&usize) -> IoResult<()> {
        match self.jobs.get_mut(id) {
            None => return Err(IoError {
                kind: OtherIoError,
                desc: "Job not found",
                detail: None
            }),
            Some(&mut (_, ref mut job)) => {
                return job.signal(SIGINT as isize);
            }
        }
    }

    fn find_jobs_hole(&self) -> usize {
        // find a hole in the job map
        if self.jobs.len() < 20 {
            // with less than twenty jobs just be lazy
            return self.jobs.len() + 1;
        }
        let mut last = 0;
        for key in self.jobs.keys() {
            if key - last != 1 {
                // we've found a hole
                return key - 1;
            } else {
                last = key;
            }
        }
        // job list is full
        return last + 1;
    }

    pub fn run_job(&mut self, name:&String, args:&Vec<String>) -> Result<(usize, String), String> {
        let mut process = Command::new(name);
        process.args(args.as_slice());
        // running as job means no input/output handles
        process.stdin(Ignored);
        process.stdout(Ignored);
        process.stderr(Ignored);
        // all others are redirected
        let mut child = match process.spawn() {
            Err(e) => {
                return Err(format!("Couldn't spawn {}: {}", name, e));
            },
            Ok(v) => v
        };
        // set wait timeout to zero so we can check for process exit quickly
        child.set_timeout(Some(0));
        let id = self.find_jobs_hole();
        // panic if we've overwritten a job
        match self.jobs.insert(id.clone(), (name.clone(), child)) {
            Some(_) => panic!("Overwrote job"),
            _ => {/* nothing */}
        }
        return Ok((id, name.clone()));
    }

    pub fn get_jobs(&mut self) -> Vec<(usize, String, &Process)> {
        let mut out = vec![];
        for (id, &(ref name, ref child)) in self.jobs.iter() {
            out.push((id, name.clone(), child));
        }
        return out;
    }

    pub fn clean_jobs(&mut self) -> Vec<(usize, String, IoResult<ProcessExit>)> {
        let mut out = vec![];
        let mut remove = vec![];
        for (id, &mut (ref mut name, ref mut child)) in self.jobs.iter_mut() {
             match child.wait() {
                 Err(IoError{kind: IoErrorKind::TimedOut, desc: _, detail: _}) => {
                     // this is expected, do nothing
                 },
                 v => {
                     // all other outputs mean the job is done
                     // if it isn't it'll be cleaned up in drop
                     child.set_timeout(None);
                     remove.push(id);
                     out.push((id, name.clone(), v));
                 }
             }
        }
        for id in remove.iter() {
            self.jobs.remove(id);
        }
        return out;
    }

    pub fn run_command(&mut self, name:&String, args:&Vec<String>) -> Result<ProcessExit, String> {
        let mut process = Command::new(name);
        process.args(args.as_slice());
        process.stdout(InheritFd(STDOUT));
        process.stdin(InheritFd(STDIN));
        process.stderr(InheritFd(STDERR));
        // set terminal settings for process
        self.restore_terminal();
        // push job into jobs
        let id = self.find_jobs_hole();
        // handle interrupts
        self.handle_sigint();
        let val = (name.clone(), match process.spawn() {
            Err(e) => {
                self.update_terminal();
                return Err(format!("Couldn't spawn {}: {}", name, e));
            },
            Ok(v) => v
        });
        // insert job
        self.jobs.insert(id, val);
        // set forground job
        self.fg_job = Some(id);
        let out = match self.jobs.get_mut(&id).unwrap().1.wait() {
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                self.update_terminal();
                return Err(format!("Couldn't wait for child to exit: {}", e));
            },
            Ok(v) => v
        };
        // unset forground job
        self.fg_job = None;
        // remove job
        self.jobs.remove(&id);
        // unhandle sigint
        self.unhandle_sigint();
        // restore settings for Wash
        self.update_terminal();
        return Ok(out);
    }

    pub fn run_command_directed(&mut self, name:&String,
                                args:&Vec<String>) -> Result<ProcessOutput, String> {
        let mut process = Command::new(name);
        process.args(args.as_slice());
        let id = self.find_jobs_hole();
        // handle sigint
        self.handle_sigint();
        let val = (name.clone(), match process.spawn() {
            Err(e) => {
                self.update_terminal();
                return Err(format!("Couldn't spawn {}: {}", name, e));
            },
            Ok(v) => v
        });
        // put job in jobs list
        self.jobs.insert(id, val);
        // set foreground job
        self.fg_job = Some(id);
        let out = match self.jobs.get_mut(&id).unwrap().1.stdout.as_mut().unwrap().read_to_end() {
            Ok(v) => v,
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                return Err(format!("Couldn't get stdout: {}", e));
            }
        };
        let err = match self.jobs.get_mut(&id).unwrap().1.stderr.as_mut().unwrap().read_to_end() {
            Ok(v) => v,
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                return Err(format!("Couldn't get stderr: {}", e));
            }
        };
        let output = match self.jobs.get_mut(&id).unwrap().1.wait() {
            Err(e) => {
                // unset foreground job
                self.fg_job = None;
                // remove job
                self.jobs.remove(&id);
                self.update_terminal();
                return Err(format!("Couldn't wait for child to exit: {}", e));
            },
            Ok(v) => ProcessOutput {
                status: v,
                output: out,
                error: err
            }
        };
        // unset foreground job
        self.fg_job = None;
        // remove job
        self.jobs.remove(&id);
        // unhandle sigint
        self.unhandle_sigint();
        return Ok(output);
    }
}

<|endoftext|>


<commit_before>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    if let Some(client) = s3_client() {
        let res = client.get_object(GetObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.into(),
            ..Default::default()
        }).sync();

        let res = match res {
            Ok(r) => r,
            Err(_) => {
                return None;
            }
        };

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::new();
        b.read_to_end(&mut content).unwrap();

        let last_modified = res.last_modified.unwrap();
        let last_modified = time::strptime(&last_modified, "%a, %d %b %Y %H:%M:%S %Z")
            .unwrap_or_else(|e| panic!("failed to parse {:?} as timespec: {:?}", last_modified, e))
            .to_timespec();

        Some(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated: last_modified,
            content,
        })
    } else {
        let rows = conn.query("SELECT path, mime, date_updated, content
                            FROM files
                            WHERE path = $1", &[&path]).unwrap();

        if rows.len() == 0 {
            None
        } else {
            let row = rows.get(0);

            Some(Blob {
                path: row.get(0),
                mime: row.get(1),
                date_updated: row.get(2),
                content: row.get(3),
            })
        }
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());
    let mut client = s3_client();
    let mut file_list_with_mimes: Vec<(String, PathBuf)> = Vec::new();

    for file_path in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            let content: Option<Vec<u8>> = if let Some(client) = &mut client {
                let mut attempts = 0;
                loop {
                    let s3_res = client.put_object(PutObjectRequest {
                        bucket: "rust-docs-rs".into(),
                        key: bucket_path.clone(),
                        body: Some(content.clone().into()),
                        content_type: Some(mime.clone()),
                        ..Default::default()
                    }).sync();
                    attempts += 1;
                    match s3_res {
                        // we've successfully uploaded the content, so steal it;
                        // we don't want to put it in the DB
                        Ok(_) => {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                            break None;
                        },
                        // Since s3 was configured, we want to panic on failure to upload.
                        Err(e) => {
                            log::error!("failed to upload to {}: {:?}", bucket_path, e);
                            // Get a new client, in case the old one's connection is stale.
                            // AWS will kill our connection if it's alive for too long; this avoids
                            // that preventing us from building the crate entirely.
                            *client = s3_client().unwrap();
                            if attempts > 3 {
                                panic!("failed to upload 3 times, exiting");
                            } else {
                                continue;
                            }
                        },
                    }
                }
            } else {
                Some(content.clone().into())
            };

            file_list_with_mimes.push((mime.clone(), file_path.clone()));

            (
                bucket_path,
                content,
                mime,
            )
        };

        // If AWS credentials are configured, don't insert/update the database
        if client.is_none() {
            // check if file already exists in database
            let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

            let content = content.expect("content never None if client is None");

            if rows.get(0).get::<usize, i64>(0) == 0 {
                try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                                &[&path, &mime, &content]));
            } else {
                try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                WHERE path = $1",
                                &[&path, &mime, &content]));
            }
            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<commit_msg>perform s3 uploads in parallel<commit_after>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::cmp;
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;

const MAX_CONCURRENT_UPLOADS: usize = 1000;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    if let Some(client) = s3_client() {
        let res = client.get_object(GetObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.into(),
            ..Default::default()
        }).sync();

        let res = match res {
            Ok(r) => r,
            Err(_) => {
                return None;
            }
        };

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::new();
        b.read_to_end(&mut content).unwrap();

        let last_modified = res.last_modified.unwrap();
        let last_modified = time::strptime(&last_modified, "%a, %d %b %Y %H:%M:%S %Z")
            .unwrap_or_else(|e| panic!("failed to parse {:?} as timespec: {:?}", last_modified, e))
            .to_timespec();

        Some(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated: last_modified,
            content,
        })
    } else {
        let rows = conn.query("SELECT path, mime, date_updated, content
                            FROM files
                            WHERE path = $1", &[&path]).unwrap();

        if rows.len() == 0 {
            None
        } else {
            let row = rows.get(0);

            Some(Blob {
                path: row.get(0),
                mime: row.get(1),
                date_updated: row.get(2),
                content: row.get(3),
            })
        }
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());

    use std::collections::HashMap;
    let mut file_paths_and_mimes: HashMap<PathBuf, String> = HashMap::new();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();

    let mut to_upload = try!(get_file_list(&path));
    let mut batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
    let mut currently_uploading: Vec<_> = to_upload.drain(..batch_size).collect();
    let mut attempts = 0;

    while !to_upload.is_empty() || !currently_uploading.is_empty() {

        let mut futures = Vec::new();
        let client = s3_client();

        for file_path in &currently_uploading {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            if let Some(client) = &client {
                futures.push(client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }));
            } else {
                // If AWS credentials are configured, don't insert/update the database
                // check if file already exists in database
                let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&bucket_path]));

                if rows.get(0).get::<usize, i64>(0) == 0 {
                    try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                                    &[&bucket_path, &mime, &content]));
                } else {
                    try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                    WHERE path = $1",
                                    &[&bucket_path, &mime, &content]));
                }
            }

            file_paths_and_mimes.insert(file_path.clone(), mime.clone());
        }

        if !futures.is_empty() {
            attempts += 1;

            if rt.block_on(::futures::future::join_all(futures)).is_ok() {
                // this batch was successful, start another batch if there are still more files
                crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(batch_size as i64);
                batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
                currently_uploading = to_upload.drain(..batch_size).collect();
                attempts = 0;
            } else {
                // if any futures error, leave `currently_uploading` in tact so that we can retry the batch
                if attempts > 2 {
                    panic!("failed to upload 3 times, exiting");
                }
            }
        } else {
            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(batch_size as i64);
            batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
            currently_uploading = to_upload.drain(..batch_size).collect();
        }
    }

    try!(trans.commit());

    let file_list_with_mimes: Vec<(String, PathBuf)> = file_paths_and_mimes
        .into_iter()
        .map(|(file_path, mime)| (mime, file_path))
        .collect();
    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<|endoftext|>


<commit_before>use crate::{BuildMetadata, Comparator, Op, Prerelease, Version, VersionReq};
use core::fmt::{self, Debug, Display};

impl Display for Version {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt(&self.major, formatter)?;
        formatter.write_str(".")?;
        Display::fmt(&self.minor, formatter)?;
        formatter.write_str(".")?;
        Display::fmt(&self.patch, formatter)?;
        if !self.pre.is_empty() {
            formatter.write_str("-")?;
            Display::fmt(&self.pre, formatter)?;
        }
        if !self.build.is_empty() {
            formatter.write_str("+")?;
            Display::fmt(&self.build, formatter)?;
        }
        Ok(())
    }
}

impl Display for VersionReq {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        if self.comparators.is_empty() {
            return formatter.write_str("*");
        }
        for (i, comparator) in self.comparators.iter().enumerate() {
            if i > 0 {
                formatter.write_str(", ")?;
            }
            Display::fmt(comparator, formatter)?;
        }
        Ok(())
    }
}

impl Display for Comparator {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        let op = match self.op {
            Op::Exact => "=",
            Op::Greater => ">",
            Op::GreaterEq => ">=",
            Op::Less => "<",
            Op::LessEq => "<=",
            Op::Tilde => "~",
            Op::Caret => "^",
            Op::Wildcard => "",
            #[cfg(no_non_exhaustive)]
            Op::__NonExhaustive => unreachable!(),
        };
        formatter.write_str(op)?;
        Display::fmt(&self.major, formatter)?;
        if let Some(minor) = &self.minor {
            formatter.write_str(".")?;
            Display::fmt(minor, formatter)?;
            if let Some(patch) = &self.patch {
                formatter.write_str(".")?;
                Display::fmt(patch, formatter)?;
                if !self.pre.is_empty() {
                    formatter.write_str("-")?;
                    Display::fmt(&self.pre, formatter)?;
                }
            } else if self.op == Op::Wildcard {
                formatter.write_str(".*")?;
            }
        } else if self.op == Op::Wildcard {
            formatter.write_str(".*")?;
        }
        Ok(())
    }
}

impl Display for Prerelease {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(self.as_str())
    }
}

impl Display for BuildMetadata {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(self.as_str())
    }
}

impl Debug for Version {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        let mut debug = formatter.debug_struct("Version");
        debug
            .field("major", &self.major)
            .field("minor", &self.minor)
            .field("patch", &self.patch);
        if !self.pre.is_empty() {
            debug.field("pre", &self.pre);
        }
        if !self.build.is_empty() {
            debug.field("build", &self.build);
        }
        debug.finish()
    }
}

impl Debug for Prerelease {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("Prerelease(\"")?;
        Display::fmt(self, formatter)?;
        formatter.write_str("\")")?;
        Ok(())
    }
}

impl Debug for BuildMetadata {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("BuildMetadata(\"")?;
        Display::fmt(self, formatter)?;
        formatter.write_str("\")")?;
        Ok(())
    }
}
<commit_msg>Avoid inheriting formatter flags for displaying Version<commit_after>use crate::{BuildMetadata, Comparator, Op, Prerelease, Version, VersionReq};
use core::fmt::{self, Debug, Display};

impl Display for Version {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "{}.{}.{}", self.major, self.minor, self.patch)?;
        if !self.pre.is_empty() {
            write!(formatter, "-{}", self.pre)?;
        }
        if !self.build.is_empty() {
            write!(formatter, "+{}", self.build)?;
        }
        Ok(())
    }
}

impl Display for VersionReq {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        if self.comparators.is_empty() {
            return formatter.write_str("*");
        }
        for (i, comparator) in self.comparators.iter().enumerate() {
            if i > 0 {
                formatter.write_str(", ")?;
            }
            write!(formatter, "{}", comparator)?;
        }
        Ok(())
    }
}

impl Display for Comparator {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        let op = match self.op {
            Op::Exact => "=",
            Op::Greater => ">",
            Op::GreaterEq => ">=",
            Op::Less => "<",
            Op::LessEq => "<=",
            Op::Tilde => "~",
            Op::Caret => "^",
            Op::Wildcard => "",
            #[cfg(no_non_exhaustive)]
            Op::__NonExhaustive => unreachable!(),
        };
        formatter.write_str(op)?;
        write!(formatter, "{}", self.major)?;
        if let Some(minor) = &self.minor {
            write!(formatter, ".{}", minor)?;
            if let Some(patch) = &self.patch {
                write!(formatter, ".{}", patch)?;
                if !self.pre.is_empty() {
                    write!(formatter, "-{}", self.pre)?;
                }
            } else if self.op == Op::Wildcard {
                formatter.write_str(".*")?;
            }
        } else if self.op == Op::Wildcard {
            formatter.write_str(".*")?;
        }
        Ok(())
    }
}

impl Display for Prerelease {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(self.as_str())
    }
}

impl Display for BuildMetadata {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(self.as_str())
    }
}

impl Debug for Version {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        let mut debug = formatter.debug_struct("Version");
        debug
            .field("major", &self.major)
            .field("minor", &self.minor)
            .field("patch", &self.patch);
        if !self.pre.is_empty() {
            debug.field("pre", &self.pre);
        }
        if !self.build.is_empty() {
            debug.field("build", &self.build);
        }
        debug.finish()
    }
}

impl Debug for Prerelease {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "Prerelease(\"{}\")", self)
    }
}

impl Debug for BuildMetadata {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "BuildMetadata(\"{}\")", self)
    }
}
<|endoftext|>