<commit_before>//! K-means Classification
//!
//! Provides implementation of K-Means classification.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::learning::k_means::KMeansClassifier;
//! use rusty_machine::learning::UnSupModel;
//!
//! let train_data = Matrix::new(3, 2, vec![1.0, 2.0, 1.0, 3.0, 1.0, 4.0]);
//! let pred_data = Matrix::new(1,2, vec![0.0, 0.0]);
//!
//! // Create model with k(=2) classes.
//! let mut model = KMeansClassifier::new(2);
//!
//! // Where train_data is a Matrix with features in columns.
//! model.train(train_data);
//!
//! // Where pred_data is a Matrix with features in columns.
//! let a = model.predict(pred_data);
//! ```
//!
//! Additionally you can control the initialization
//! algorithm and max number of iterations.
//!
//! # Initializations
//!
//! Three initialization algorithms are supported.
//!
//! ## Forgy initialization
//!
//! Choose initial centroids randomly from the data.
//!
//! ## Random Partition initialization
//!
//! Randomly assign each data point to one of k clusters.
//! The initial centroids are the mean of the data in their class.
//!
//! ## K-means++ initialization
//!
//! The [k-means++](https://en.wikipedia.org/wiki/K-means%2B%2B) scheme.

use linalg::matrix::Matrix;
use linalg::vector::Vector;
use learning::UnSupModel;
use rand::{Rng, thread_rng};

/// Initialization Algorithm enum.
pub enum InitAlgorithm {
    Forgy,
    RandomPartition,
    KPlusPlus,
}

/// K-Means Classification model.
///
/// Contains option for centroids.
/// Specifies iterations and number of classes.
pub struct KMeansClassifier {
    pub iters: usize,
    pub k: usize,
    pub centroids: Option<Matrix<f64>>,
    pub init_algorithm: InitAlgorithm,
}

impl UnSupModel<Matrix<f64>, Vector<usize>> for KMeansClassifier {
    /// Predict classes from data.
    ///
    /// Model must be trained.
    fn predict(&self, data: Matrix<f64>) -> Vector<usize> {
        if let Some(ref centroids) = self.centroids {
            return KMeansClassifier::find_closest_centroids(centroids, &data).0;
        }
        else {
            panic!("Model has not been trained.");
        }
    }

    /// Train the classifier using input data.
    fn train(&mut self, data: Matrix<f64>) {
        self.init_centroids(&data);

        for _i in 0..self.iters {
                let (idx,_) = self.get_closest_centroids(&data);
                self.update_centroids(&data, idx);
        }
    }
}

impl KMeansClassifier {
    /// Constructs untrained k-means classifier model.
    ///
    /// Requires number of classes to be specified.
    /// Defaults to 100 iterations.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::k_means::KMeansClassifier;
    ///
    /// let model = KMeansClassifier::new(5);
    /// ```
    pub fn new(k: usize) -> KMeansClassifier {
        KMeansClassifier {
            iters: 100,
            k: k,
            centroids: None,
            init_algorithm: InitAlgorithm::Forgy,
        }
    }

    /// Initialize the centroids.
    ///
    /// Used internally within model.
    fn init_centroids(&mut self, data: &Matrix<f64>) {
        match self.init_algorithm {
            InitAlgorithm::Forgy => {
                self.centroids = Some(KMeansClassifier::forgy_init(self.k, data))
            }
            InitAlgorithm::RandomPartition => {
                self.centroids = Some(KMeansClassifier::ran_partition_init(self.k, data))
            }
            InitAlgorithm::KPlusPlus => {
                self.centroids = Some(KMeansClassifier::plusplus_init(self.k, data))
            }
        }
    }

    /// Updated the centroids by computing means of assigned classes.
    ///
    /// Used internally within model.
    fn update_centroids(&mut self, data: &Matrix<f64>, classes: Vector<usize>) {
        let mut new_centroids = Vec::with_capacity(self.k * data.cols());
        for i in 0..self.k {
            let mut vec_i = Vec::new();

            for j in classes.data.iter() {
                if *j == i {
                    vec_i.push(*j);
                }
            }

            let mat_i = data.select_rows(&vec_i);
            new_centroids.extend(mat_i.mean(0).data);
        }

        self.centroids = Some(Matrix::new(self.k, data.cols(), new_centroids));
    }

    fn get_closest_centroids(&self, data: &Matrix<f64>) -> (Vector<usize>, Vector<f64>) {
        if let Some(ref c) = self.centroids {
            return KMeansClassifier::find_closest_centroids(&c, data);
        }
        else {
            panic!("Centroids not correctly initialized.");
        }
    }

    /// Find the centroid closest to each data point.
    ///
    /// Used internally within model.
    /// Returns the index of the closest centroid and the distance to it.
    fn find_closest_centroids(centroids: &Matrix<f64>, data: &Matrix<f64>) -> (Vector<usize>, Vector<f64>) {
        let mut idx = Vector::zeros(data.rows());
        let mut distances = Vector::zeros(data.rows());

        for i in 0..data.rows() {
            // This works like repmat pulling out row i repeatedly.
            let centroid_diff = centroids - data.select_rows(&vec![i; centroids.rows()]);
            let dist = &centroid_diff.elemul(&centroid_diff).sum_cols();

            // Now take argmin and this is the centroid.
            let (min_idx, min_dist) = dist.argmin();
            idx.data[i]= min_idx;
            distances.data[i] = min_dist;

        }

        (idx, distances)
    }

    /// Compute initial centroids using Forgy scheme.
    ///
    /// Selects k random points in data for centroids.
    fn forgy_init(k: usize, data: &Matrix<f64>) -> Matrix<f64> {
        assert!(k <= data.rows());

        let mut random_choices = Vec::with_capacity(k);
        let mut rng = thread_rng();
        while random_choices.len() < k {
            let r = rng.gen_range(0, data.rows());

            if !random_choices.contains(&r) {
                random_choices.push(r);
            }
        }

        data.select_rows(&random_choices)
    }

    /// Compute initial centroids using random partition.
    ///
    /// Selects centroids by assigning each point randomly to a class
    /// and computing the mean of each class.
    fn ran_partition_init(k: usize, data: &Matrix<f64>) -> Matrix<f64> {
        assert!(k <= data.rows());

        let mut random_assignments = Vec::with_capacity(data.rows());

        // Populate so we have something in each class.
        for i in 0..k {
            random_assignments.push(i);
        }

        let mut rng = thread_rng();
        for _i in k..data.rows() {
            random_assignments.push(rng.gen_range(0, k));
        }

        let mut init_centroids = Vec::with_capacity(k * data.cols());
        for i in 0..k {
            let mut vec_i = Vec::new();

            for j in random_assignments.iter() {
                if *j == i {
                    vec_i.push(*j);
                }
            }

            let mat_i = data.select_rows(&vec_i);
            init_centroids.extend(mat_i.mean(0).data);
        }

        Matrix::new(k, data.cols(), init_centroids)
    }

    /// Compute initial centroids using k-means++.
    ///
    /// Selects centroids using weighted probability from
    /// distances.
    fn plusplus_init(k: usize, data: &Matrix<f64>) -> Matrix<f64> {
        assert!(k <= data.rows());

        let mut rng = thread_rng();

        let mut init_centroids = Vec::with_capacity(k * data.cols());
        let first_cen = rng.gen_range(0usize, data.rows());

        init_centroids.append(&mut data.select_rows(&vec![first_cen]).data);

        for i in 1..k {
            let temp_centroids = Matrix::new(i, data.cols(), init_centroids.clone());
            let (_, dist) = KMeansClassifier::find_closest_centroids(&temp_centroids, &data);
            let next_cen = sample_discretely(dist);
            init_centroids.append(&mut data.select_rows(&vec![next_cen]).data)
        }

        Matrix::new(k, data.cols(), init_centroids)
    }
}

/// Sample from an unnormalized distribution.
///
///
fn sample_discretely(unnorm_dist: Vector<f64>) -> usize {
    assert!(unnorm_dist.size() > 0);

    let sum = unnorm_dist.sum();

    let rand = thread_rng().gen_range(0.0f64, sum);

    let mut tempsum = 0.0;
    for (i,p) in unnorm_dist.data.iter().enumerate() {
        tempsum += *p;

        if rand < tempsum {
            return i;
        }
    }

    panic!("No random value was sampled! There may be more clusters than unique data points.");
}<commit_msg>Cost checking to cut from iterations if converged.<commit_after>//! K-means Classification
//!
//! Provides implementation of K-Means classification.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::learning::k_means::KMeansClassifier;
//! use rusty_machine::learning::UnSupModel;
//!
//! let train_data = Matrix::new(3, 2, vec![1.0, 2.0, 1.0, 3.0, 1.0, 4.0]);
//! let pred_data = Matrix::new(1,2, vec![0.0, 0.0]);
//!
//! // Create model with k(=2) classes.
//! let mut model = KMeansClassifier::new(2);
//!
//! // Where train_data is a Matrix with features in columns.
//! model.train(train_data);
//!
//! // Where pred_data is a Matrix with features in columns.
//! let a = model.predict(pred_data);
//! ```
//!
//! Additionally you can control the initialization
//! algorithm and max number of iterations.
//!
//! # Initializations
//!
//! Three initialization algorithms are supported.
//!
//! ## Forgy initialization
//!
//! Choose initial centroids randomly from the data.
//!
//! ## Random Partition initialization
//!
//! Randomly assign each data point to one of k clusters.
//! The initial centroids are the mean of the data in their class.
//!
//! ## K-means++ initialization
//!
//! The [k-means++](https://en.wikipedia.org/wiki/K-means%2B%2B) scheme.

use linalg::matrix::Matrix;
use linalg::vector::Vector;
use learning::UnSupModel;
use rand::{Rng, thread_rng};

/// Initialization Algorithm enum.
pub enum InitAlgorithm {
    Forgy,
    RandomPartition,
    KPlusPlus,
}

/// K-Means Classification model.
///
/// Contains option for centroids.
/// Specifies iterations and number of classes.
pub struct KMeansClassifier {
    pub iters: usize,
    pub k: usize,
    pub centroids: Option<Matrix<f64>>,
    pub init_algorithm: InitAlgorithm,
}

impl UnSupModel<Matrix<f64>, Vector<usize>> for KMeansClassifier {
    /// Predict classes from data.
    ///
    /// Model must be trained.
    fn predict(&self, data: Matrix<f64>) -> Vector<usize> {
        if let Some(ref centroids) = self.centroids {
            return KMeansClassifier::find_closest_centroids(centroids, &data).0;
        }
        else {
            panic!("Model has not been trained.");
        }
    }

    /// Train the classifier using input data.
    fn train(&mut self, data: Matrix<f64>) {
        self.init_centroids(&data);
        let mut cost = 0.0;

        for _i in 0..self.iters {
                let (idx, distances) = self.get_closest_centroids(&data);
                self.update_centroids(&data, idx);

                let cost_i = distances.sum();
                if cost == cost_i {
                    break;
                }

                cost = cost_i;
        }
    }
}

impl KMeansClassifier {
    /// Constructs untrained k-means classifier model.
    ///
    /// Requires number of classes to be specified.
    /// Defaults to 100 iterations and kmeans++ initialization.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::k_means::KMeansClassifier;
    ///
    /// let model = KMeansClassifier::new(5);
    /// ```
    pub fn new(k: usize) -> KMeansClassifier {
        KMeansClassifier {
            iters: 100,
            k: k,
            centroids: None,
            init_algorithm: InitAlgorithm::KPlusPlus,
        }
    }

    /// Initialize the centroids.
    ///
    /// Used internally within model.
    fn init_centroids(&mut self, data: &Matrix<f64>) {
        match self.init_algorithm {
            InitAlgorithm::Forgy => {
                self.centroids = Some(KMeansClassifier::forgy_init(self.k, data))
            }
            InitAlgorithm::RandomPartition => {
                self.centroids = Some(KMeansClassifier::ran_partition_init(self.k, data))
            }
            InitAlgorithm::KPlusPlus => {
                self.centroids = Some(KMeansClassifier::plusplus_init(self.k, data))
            }
        }
    }

    /// Updated the centroids by computing means of assigned classes.
    ///
    /// Used internally within model.
    fn update_centroids(&mut self, data: &Matrix<f64>, classes: Vector<usize>) {
        let mut new_centroids = Vec::with_capacity(self.k * data.cols());
        for i in 0..self.k {
            let mut vec_i = Vec::new();

            for j in classes.data.iter() {
                if *j == i {
                    vec_i.push(*j);
                }
            }

            let mat_i = data.select_rows(&vec_i);
            new_centroids.extend(mat_i.mean(0).data);
        }

        self.centroids = Some(Matrix::new(self.k, data.cols(), new_centroids));
    }

    fn get_closest_centroids(&self, data: &Matrix<f64>) -> (Vector<usize>, Vector<f64>) {
        if let Some(ref c) = self.centroids {
            return KMeansClassifier::find_closest_centroids(&c, data);
        }
        else {
            panic!("Centroids not correctly initialized.");
        }
    }

    /// Find the centroid closest to each data point.
    ///
    /// Used internally within model.
    /// Returns the index of the closest centroid and the distance to it.
    fn find_closest_centroids(centroids: &Matrix<f64>, data: &Matrix<f64>) -> (Vector<usize>, Vector<f64>) {
        let mut idx = Vector::zeros(data.rows());
        let mut distances = Vector::zeros(data.rows());

        for i in 0..data.rows() {
            // This works like repmat pulling out row i repeatedly.
            let centroid_diff = centroids - data.select_rows(&vec![i; centroids.rows()]);
            let dist = &centroid_diff.elemul(&centroid_diff).sum_cols();

            // Now take argmin and this is the centroid.
            let (min_idx, min_dist) = dist.argmin();
            idx.data[i]= min_idx;
            distances.data[i] = min_dist;

        }

        (idx, distances)
    }

    /// Compute initial centroids using Forgy scheme.
    ///
    /// Selects k random points in data for centroids.
    fn forgy_init(k: usize, data: &Matrix<f64>) -> Matrix<f64> {
        assert!(k <= data.rows());

        let mut random_choices = Vec::with_capacity(k);
        let mut rng = thread_rng();
        while random_choices.len() < k {
            let r = rng.gen_range(0, data.rows());

            if !random_choices.contains(&r) {
                random_choices.push(r);
            }
        }

        data.select_rows(&random_choices)
    }

    /// Compute initial centroids using random partition.
    ///
    /// Selects centroids by assigning each point randomly to a class
    /// and computing the mean of each class.
    fn ran_partition_init(k: usize, data: &Matrix<f64>) -> Matrix<f64> {
        assert!(k <= data.rows());

        let mut random_assignments = Vec::with_capacity(data.rows());

        // Populate so we have something in each class.
        for i in 0..k {
            random_assignments.push(i);
        }

        let mut rng = thread_rng();
        for _i in k..data.rows() {
            random_assignments.push(rng.gen_range(0, k));
        }

        let mut init_centroids = Vec::with_capacity(k * data.cols());
        for i in 0..k {
            let mut vec_i = Vec::new();

            for j in random_assignments.iter() {
                if *j == i {
                    vec_i.push(*j);
                }
            }

            let mat_i = data.select_rows(&vec_i);
            init_centroids.extend(mat_i.mean(0).data);
        }

        Matrix::new(k, data.cols(), init_centroids)
    }

    /// Compute initial centroids using k-means++.
    ///
    /// Selects centroids using weighted probability from
    /// distances.
    fn plusplus_init(k: usize, data: &Matrix<f64>) -> Matrix<f64> {
        assert!(k <= data.rows());

        let mut rng = thread_rng();

        let mut init_centroids = Vec::with_capacity(k * data.cols());
        let first_cen = rng.gen_range(0usize, data.rows());

        init_centroids.append(&mut data.select_rows(&vec![first_cen]).data);

        for i in 1..k {
            let temp_centroids = Matrix::new(i, data.cols(), init_centroids.clone());
            let (_, dist) = KMeansClassifier::find_closest_centroids(&temp_centroids, &data);
            let next_cen = sample_discretely(dist);
            init_centroids.append(&mut data.select_rows(&vec![next_cen]).data)
        }

        Matrix::new(k, data.cols(), init_centroids)
    }
}

/// Sample from an unnormalized distribution.
///
///
fn sample_discretely(unnorm_dist: Vector<f64>) -> usize {
    assert!(unnorm_dist.size() > 0);

    let sum = unnorm_dist.sum();

    let rand = thread_rng().gen_range(0.0f64, sum);

    let mut tempsum = 0.0;
    for (i,p) in unnorm_dist.data.iter().enumerate() {
        tempsum += *p;

        if rand < tempsum {
            return i;
        }
    }

    panic!("No random value was sampled! There may be more clusters than unique data points.");
}<|endoftext|>


<commit_before>use super::pool::Pool;
use iron::headers::ContentType;
use iron::prelude::*;
use iron::status::Status;
use prometheus::{
    opts, register_counter, register_int_counter, register_int_gauge, Encoder, IntCounter,
    IntGauge, TextEncoder, __register_gauge, register_int_counter_vec, IntCounterVec,
    __register_counter_vec, histogram_opts, register_histogram_vec, HistogramVec,
};
use std::time::{Duration, Instant};

lazy_static::lazy_static! {
    static ref QUEUED_CRATES_COUNT: IntGauge = register_int_gauge!(
        "docsrs_queued_crates_count",
        "Number of crates in the build queue"
    )
    .unwrap();

    static ref FAILED_CRATES_COUNT: IntGauge = register_int_gauge!(
        "docsrs_failed_crates_count",
        "Number of crates that failed to build"
    )
    .unwrap();

    pub static ref TOTAL_BUILDS: IntCounter = register_int_counter!(
        "docsrs_total_builds",
        "Number of crates built"
    )
    .unwrap();

    pub static ref SUCCESSFUL_BUILDS: IntCounter = register_int_counter!(
        "docsrs_successful_builds",
        "Number of builds that successfully generated docs"
    )
    .unwrap();

    pub static ref FAILED_BUILDS: IntCounter = register_int_counter!(
        "docsrs_failed_builds",
        "Number of builds that generated a compile error"
    )
    .unwrap();

    pub static ref NON_LIBRARY_BUILDS: IntCounter = register_int_counter!(
        "docsrs_non_library_builds",
        "Number of builds that did not complete due to not being a library"
    )
    .unwrap();

    pub static ref UPLOADED_FILES_TOTAL: IntCounter = register_int_counter!(
        "docsrs_uploaded_files_total",
        "Number of files uploaded to S3 or stored in the database"
    )
    .unwrap();

    pub static ref ROUTES_VISITED: IntCounterVec = register_int_counter_vec!(
        "docsrs_routes_visited",
        "The traffic of various docs.rs routes",
        &["route"]
    )
    .unwrap();

    pub static ref RESPONSE_TIMES: HistogramVec = register_histogram_vec!(
        "docsrs_response_time",
        "The response times of various docs.rs routes",
        &["route"]
    )
    .unwrap();
}

pub fn metrics_handler(req: &mut Request) -> IronResult<Response> {
    let conn = extension!(req, Pool).get();

    QUEUED_CRATES_COUNT.set(
        ctry!(conn.query("SELECT COUNT(*) FROM queue WHERE attempt < 5;", &[]))
            .get(0)
            .get(0),
    );
    FAILED_CRATES_COUNT.set(
        ctry!(conn.query("SELECT COUNT(*) FROM queue WHERE attempt >= 5;", &[]))
            .get(0)
            .get(0),
    );

    let mut buffer = Vec::new();
    let families = prometheus::gather();
    ctry!(TextEncoder::new().encode(&families, &mut buffer));

    let mut resp = Response::with(buffer);
    resp.status = Some(Status::Ok);
    resp.headers
        .set(ContentType("text/plain; version=0.0.4".parse().unwrap()));
    Ok(resp)
}

/// Converts a `Duration` to seconds, used by prometheus internally
#[inline]
fn duration_to_seconds(d: Duration) -> f64 {
    let nanos = f64::from(d.subsec_nanos()) / 1e9;
    d.as_secs() as f64 + nanos
}

#[derive(Debug, Clone)]
pub struct ResponseRecorder {
    start_time: Instant,
    route: Option<String>,
}

impl ResponseRecorder {
    #[inline]
    pub fn new() -> Self {
        Self {
            start_time: Instant::now(),
            route: None,
        }
    }

    #[inline]
    pub fn route(&mut self, route: impl Into<String>) {
        self.route = Some(route.into());
    }
}

impl Drop for ResponseRecorder {
    fn drop(&mut self) {
        if let Some(route) = &self.route {
            ROUTES_VISITED.with_label_values(&[route]).inc();

            let response_time = duration_to_seconds(self.start_time.elapsed());
            RESPONSE_TIMES
                .with_label_values(&[route])
                .observe(response_time);

            #[cfg(test)]
            tests::record_tests(route);
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::test::wrapper;
    use std::{
        collections::HashMap,
        sync::{
            atomic::{AtomicUsize, Ordering},
            Mutex, MutexGuard,
        },
    };

    static ROUTES_VISITED: AtomicUsize = AtomicUsize::new(0);
    lazy_static::lazy_static! {
        static ref RESPONSE_TIMES: Mutex<HashMap<String, usize>> = Mutex::new(HashMap::new());
    }

    pub fn record_tests(route: &str) {
        ROUTES_VISITED.fetch_add(1, Ordering::SeqCst);

        let mut times = RESPONSE_TIMES.lock().unwrap();
        if let Some(requests) = times.get_mut(route) {
            *requests += 1;
        } else {
            times.insert(route.to_owned(), 1);
        }
    }

    fn reset_records() {
        ROUTES_VISITED.store(0, Ordering::SeqCst);
        RESPONSE_TIMES.lock().unwrap().clear();
    }

    #[test]
    fn home_page() {
        wrapper(|env| {
            let frontend = env.frontend();

            reset_records();

            frontend.get("/").send()?;
            frontend.get("/").send()?;
            assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
            assert_eq!(RESPONSE_TIMES.lock().unwrap().get("home (found)"), Some(&2));

            reset_records();

            frontend.get("").send()?;
            frontend.get("").send()?;
            assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
            assert_eq!(RESPONSE_TIMES.lock().unwrap().get("home (found)"), Some(&2));

            Ok(())
        })
    }

    #[test]
    fn resources() {
        wrapper(|env| {
            let frontend = env.frontend();

            let routes = [
                "/style.css",
                "/index.js",
                "/menu.js",
                "/sitemap.xml",
                "/opensearch.xml",
            ];

            for route in routes.iter() {
                reset_records();

                frontend.get(route).send()?;
                frontend.get(route).send()?;

                assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
                assert_eq!(
                    RESPONSE_TIMES.lock().unwrap().get("resources (found)"),
                    Some(&2)
                );
            }

            reset_records();

            frontend.get("/robots.txt").send()?;
            frontend.get("/robots.txt").send()?;

            assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
            assert_eq!(RESPONSE_TIMES.lock().unwrap().get("resources"), Some(&2));

            Ok(())
        })
    }

    #[test]
    fn releases() {
        wrapper(|env| {
            env.db()
                .fake_release()
                .name("rcc")
                .version("0.0.0")
                .repo("https://github.com/jyn514/rcc")
                .create()?;
            env.db()
                .fake_release()
                .name("rcc")
                .version("1.0.0")
                .build_result_successful(false)
                .create()?;

            let frontend = env.frontend();

            let routes = [
                ("/releases", "recent releases (found)"),
                // ("/releases/recent", "recent releases (found)"),
                ("/releases/recent/1", "recent releases (found)"),
                ("/releases/feed", "release feed (found)"),
                ("/releases/queue", "release queue (found)"),
                // ("/releases/search", "search releases (found)"),
                // ("/releases/stars", "release stars (found)"),
                // ("/releases/stars/1", "release stars (found)"),
                // ("/releases/activity", "release activity (found)"),
                // ("/releases/failures", "build failures (found)"),
                // ("/releases/failures/1", "build failures (found)"),
                ("/releases/recent-failures", "recent failures (found)"),
                ("/releases/recent-failures/1", "recent failures (found)"),
            ];

            for (route, correct) in routes.iter() {
                reset_records();

                frontend.get(route).send()?;
                frontend.get(route).send()?;

                assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
                assert_eq!(RESPONSE_TIMES.lock().unwrap().get(*correct), Some(&2));
            }

            Ok(())
        })
    }

    #[test]
    fn crates() {
        wrapper(|env| {
            env.db()
                .fake_release()
                .name("rcc")
                .version("0.0.0")
                .create()?;
            env.db()
                .fake_release()
                .name("hexponent")
                .version("0.2.0")
                .create()?;

            let frontend = env.frontend();

            let routes = [
                ("/crate/rcc", "crate details (found)"),
                ("/crate/rcc/", "crate details (found)"),
                ("/crate/hexponent", "crate details (found)"),
                ("/crate/hexponent/", "crate details (found)"),
                ("/crate/rcc/0.0.0", "crate details (found)"),
                ("/crate/rcc/0.0.0/", "crate details (found)"),
                ("/crate/hexponent/0.2.0", "crate details (found)"),
                ("/crate/hexponent/0.2.0/", "crate details (found)"),
                ("/crate/i_dont_exist", "crate details (404)"),
                ("/crate/i_dont_exist/", "crate details (404)"),
                ("/crate/i_dont_exist/4.0.4", "crate details (404)"),
                ("/crate/i_dont_exist/4.0.4/", "crate details (404)"),
            ];

            for (route, correct) in routes.iter() {
                reset_records();

                frontend.get(route).send()?;
                frontend.get(route).send()?;

                assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
                assert_eq!(RESPONSE_TIMES.lock().unwrap().get(*correct), Some(&2));
            }

            Ok(())
        })
    }
}
<commit_msg>Removed unused import<commit_after>use super::pool::Pool;
use iron::headers::ContentType;
use iron::prelude::*;
use iron::status::Status;
use prometheus::{
    opts, register_counter, register_int_counter, register_int_gauge, Encoder, IntCounter,
    IntGauge, TextEncoder, __register_gauge, register_int_counter_vec, IntCounterVec,
    __register_counter_vec, histogram_opts, register_histogram_vec, HistogramVec,
};
use std::time::{Duration, Instant};

lazy_static::lazy_static! {
    static ref QUEUED_CRATES_COUNT: IntGauge = register_int_gauge!(
        "docsrs_queued_crates_count",
        "Number of crates in the build queue"
    )
    .unwrap();

    static ref FAILED_CRATES_COUNT: IntGauge = register_int_gauge!(
        "docsrs_failed_crates_count",
        "Number of crates that failed to build"
    )
    .unwrap();

    pub static ref TOTAL_BUILDS: IntCounter = register_int_counter!(
        "docsrs_total_builds",
        "Number of crates built"
    )
    .unwrap();

    pub static ref SUCCESSFUL_BUILDS: IntCounter = register_int_counter!(
        "docsrs_successful_builds",
        "Number of builds that successfully generated docs"
    )
    .unwrap();

    pub static ref FAILED_BUILDS: IntCounter = register_int_counter!(
        "docsrs_failed_builds",
        "Number of builds that generated a compile error"
    )
    .unwrap();

    pub static ref NON_LIBRARY_BUILDS: IntCounter = register_int_counter!(
        "docsrs_non_library_builds",
        "Number of builds that did not complete due to not being a library"
    )
    .unwrap();

    pub static ref UPLOADED_FILES_TOTAL: IntCounter = register_int_counter!(
        "docsrs_uploaded_files_total",
        "Number of files uploaded to S3 or stored in the database"
    )
    .unwrap();

    pub static ref ROUTES_VISITED: IntCounterVec = register_int_counter_vec!(
        "docsrs_routes_visited",
        "The traffic of various docs.rs routes",
        &["route"]
    )
    .unwrap();

    pub static ref RESPONSE_TIMES: HistogramVec = register_histogram_vec!(
        "docsrs_response_time",
        "The response times of various docs.rs routes",
        &["route"]
    )
    .unwrap();
}

pub fn metrics_handler(req: &mut Request) -> IronResult<Response> {
    let conn = extension!(req, Pool).get();

    QUEUED_CRATES_COUNT.set(
        ctry!(conn.query("SELECT COUNT(*) FROM queue WHERE attempt < 5;", &[]))
            .get(0)
            .get(0),
    );
    FAILED_CRATES_COUNT.set(
        ctry!(conn.query("SELECT COUNT(*) FROM queue WHERE attempt >= 5;", &[]))
            .get(0)
            .get(0),
    );

    let mut buffer = Vec::new();
    let families = prometheus::gather();
    ctry!(TextEncoder::new().encode(&families, &mut buffer));

    let mut resp = Response::with(buffer);
    resp.status = Some(Status::Ok);
    resp.headers
        .set(ContentType("text/plain; version=0.0.4".parse().unwrap()));
    Ok(resp)
}

/// Converts a `Duration` to seconds, used by prometheus internally
#[inline]
fn duration_to_seconds(d: Duration) -> f64 {
    let nanos = f64::from(d.subsec_nanos()) / 1e9;
    d.as_secs() as f64 + nanos
}

#[derive(Debug, Clone)]
pub struct ResponseRecorder {
    start_time: Instant,
    route: Option<String>,
}

impl ResponseRecorder {
    #[inline]
    pub fn new() -> Self {
        Self {
            start_time: Instant::now(),
            route: None,
        }
    }

    #[inline]
    pub fn route(&mut self, route: impl Into<String>) {
        self.route = Some(route.into());
    }
}

impl Drop for ResponseRecorder {
    fn drop(&mut self) {
        if let Some(route) = &self.route {
            ROUTES_VISITED.with_label_values(&[route]).inc();

            let response_time = duration_to_seconds(self.start_time.elapsed());
            RESPONSE_TIMES
                .with_label_values(&[route])
                .observe(response_time);

            #[cfg(test)]
            tests::record_tests(route);
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::test::wrapper;
    use std::{
        collections::HashMap,
        sync::{
            atomic::{AtomicUsize, Ordering},
            Mutex,
        },
    };

    static ROUTES_VISITED: AtomicUsize = AtomicUsize::new(0);
    lazy_static::lazy_static! {
        static ref RESPONSE_TIMES: Mutex<HashMap<String, usize>> = Mutex::new(HashMap::new());
    }

    pub fn record_tests(route: &str) {
        ROUTES_VISITED.fetch_add(1, Ordering::SeqCst);

        let mut times = RESPONSE_TIMES.lock().unwrap();
        if let Some(requests) = times.get_mut(route) {
            *requests += 1;
        } else {
            times.insert(route.to_owned(), 1);
        }
    }

    fn reset_records() {
        ROUTES_VISITED.store(0, Ordering::SeqCst);
        RESPONSE_TIMES.lock().unwrap().clear();
    }

    #[test]
    fn home_page() {
        wrapper(|env| {
            let frontend = env.frontend();

            reset_records();

            frontend.get("/").send()?;
            frontend.get("/").send()?;
            assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
            assert_eq!(RESPONSE_TIMES.lock().unwrap().get("home (found)"), Some(&2));

            reset_records();

            frontend.get("").send()?;
            frontend.get("").send()?;
            assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
            assert_eq!(RESPONSE_TIMES.lock().unwrap().get("home (found)"), Some(&2));

            Ok(())
        })
    }

    #[test]
    fn resources() {
        wrapper(|env| {
            let frontend = env.frontend();

            let routes = [
                "/style.css",
                "/index.js",
                "/menu.js",
                "/sitemap.xml",
                "/opensearch.xml",
            ];

            for route in routes.iter() {
                reset_records();

                frontend.get(route).send()?;
                frontend.get(route).send()?;

                assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
                assert_eq!(
                    RESPONSE_TIMES.lock().unwrap().get("resources (found)"),
                    Some(&2)
                );
            }

            reset_records();

            frontend.get("/robots.txt").send()?;
            frontend.get("/robots.txt").send()?;

            assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
            assert_eq!(RESPONSE_TIMES.lock().unwrap().get("resources"), Some(&2));

            Ok(())
        })
    }

    #[test]
    fn releases() {
        wrapper(|env| {
            env.db()
                .fake_release()
                .name("rcc")
                .version("0.0.0")
                .repo("https://github.com/jyn514/rcc")
                .create()?;
            env.db()
                .fake_release()
                .name("rcc")
                .version("1.0.0")
                .build_result_successful(false)
                .create()?;

            let frontend = env.frontend();

            let routes = [
                ("/releases", "recent releases (found)"),
                // ("/releases/recent", "recent releases (found)"),
                ("/releases/recent/1", "recent releases (found)"),
                ("/releases/feed", "release feed (found)"),
                ("/releases/queue", "release queue (found)"),
                // ("/releases/search", "search releases (found)"),
                // ("/releases/stars", "release stars (found)"),
                // ("/releases/stars/1", "release stars (found)"),
                // ("/releases/activity", "release activity (found)"),
                // ("/releases/failures", "build failures (found)"),
                // ("/releases/failures/1", "build failures (found)"),
                ("/releases/recent-failures", "recent failures (found)"),
                ("/releases/recent-failures/1", "recent failures (found)"),
            ];

            for (route, correct) in routes.iter() {
                reset_records();

                frontend.get(route).send()?;
                frontend.get(route).send()?;

                assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
                assert_eq!(RESPONSE_TIMES.lock().unwrap().get(*correct), Some(&2));
            }

            Ok(())
        })
    }

    #[test]
    fn crates() {
        wrapper(|env| {
            env.db()
                .fake_release()
                .name("rcc")
                .version("0.0.0")
                .create()?;
            env.db()
                .fake_release()
                .name("hexponent")
                .version("0.2.0")
                .create()?;

            let frontend = env.frontend();

            let routes = [
                ("/crate/rcc", "crate details (found)"),
                ("/crate/rcc/", "crate details (found)"),
                ("/crate/hexponent", "crate details (found)"),
                ("/crate/hexponent/", "crate details (found)"),
                ("/crate/rcc/0.0.0", "crate details (found)"),
                ("/crate/rcc/0.0.0/", "crate details (found)"),
                ("/crate/hexponent/0.2.0", "crate details (found)"),
                ("/crate/hexponent/0.2.0/", "crate details (found)"),
                ("/crate/i_dont_exist", "crate details (404)"),
                ("/crate/i_dont_exist/", "crate details (404)"),
                ("/crate/i_dont_exist/4.0.4", "crate details (404)"),
                ("/crate/i_dont_exist/4.0.4/", "crate details (404)"),
            ];

            for (route, correct) in routes.iter() {
                reset_records();

                frontend.get(route).send()?;
                frontend.get(route).send()?;

                assert_eq!(ROUTES_VISITED.load(Ordering::SeqCst), 2);
                assert_eq!(RESPONSE_TIMES.lock().unwrap().get(*correct), Some(&2));
            }

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


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

use std::io::{Read, Write};
use std::io;
use mio::Evented;
use openssl::ssl::SslStream;
use mio::tcp::TcpStream;

use irc::util::LineReader;
use irc::protocol::*;
use irc::security::AuthSession;

pub trait AsEvented {
    fn as_evented(&self) -> &Evented;
}

pub trait IrcStream: Read + Write + AsEvented + Send {}

impl IrcStream for SslStream<TcpStream> {}
impl IrcStream for TcpStream {}

impl AsEvented for TcpStream {
    fn as_evented(&self) -> &Evented {
        self
    }
}

impl AsEvented for SslStream<TcpStream> {
    fn as_evented(&self) -> &Evented {
        self.get_ref()
    }
}

impl AsEvented for Client {
    fn as_evented(&self) -> &Evented {
        self.stream.as_evented()
    }
}

pub struct Client {
    stream: Box<IrcStream>,
    line_reader: LineReader,
    nickname: Option<String>,
    username: Option<String>,
    pub auth: AuthSession,
}

impl Client {
    pub fn new(stream: Box<IrcStream>) -> Self {
        Client {
            stream: stream,
            line_reader: LineReader::new(),
            nickname: None,
            username: None,
            auth: AuthSession::new(),
        }
    }

    pub fn read_message(&mut self) -> Option<Message> {
        match self.line_reader.read(&mut self.stream) {
            Some(line) => {
                trace!("<< {}", line);
                Some(Message::from_str(line.trim()))
            },
            None => None
        }
    }

    pub fn set_nickname(&mut self, nickname: String) {
        self.nickname = Some(nickname);
    }

    pub fn join(&mut self, channel: &str) -> io::Result<usize> {
        let pfx = self.nickname.clone().unwrap();
        self.send(&Message {
            prefix: Some(pfx),
            command: Command::Join,
            args: vec![channel.to_string()],
            suffix: None
        })
    }

    pub fn pong(&mut self) -> io::Result<usize> {
        self.send(&Message::from(Command::Pong))
    }

    pub fn welcome(&mut self, message: &str) -> io::Result<usize> {
        self.send(&Message {
            prefix: Some("pto".to_string()),
            command: Command::Numeric(1),
            args: vec![message.to_string()],
            suffix: None
        })
    }

    pub fn send(&mut self, message: &Message) -> io::Result<usize> {
        trace!(">>> {}", message.to_string());
        self.stream.write(&message.to_string().trim().as_bytes())
            .and(self.stream.write("\r\n".as_bytes()))
    }
}

pub trait Server: AsEvented {
    fn accept(&mut self) -> Option<Client>;
}
<commit_msg>irc: streams: dont crash on empty lines<commit_after>/*
 * Copyright 2015-2016 Torrie Fischer <tdfischer@hackerbots.net>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::io::{Read, Write};
use std::io;
use mio::Evented;
use openssl::ssl::SslStream;
use mio::tcp::TcpStream;

use irc::util::LineReader;
use irc::protocol::*;
use irc::security::AuthSession;

pub trait AsEvented {
    fn as_evented(&self) -> &Evented;
}

pub trait IrcStream: Read + Write + AsEvented + Send {}

impl IrcStream for SslStream<TcpStream> {}
impl IrcStream for TcpStream {}

impl AsEvented for TcpStream {
    fn as_evented(&self) -> &Evented {
        self
    }
}

impl AsEvented for SslStream<TcpStream> {
    fn as_evented(&self) -> &Evented {
        self.get_ref()
    }
}

impl AsEvented for Client {
    fn as_evented(&self) -> &Evented {
        self.stream.as_evented()
    }
}

pub struct Client {
    stream: Box<IrcStream>,
    line_reader: LineReader,
    nickname: Option<String>,
    username: Option<String>,
    pub auth: AuthSession,
}

impl Client {
    pub fn new(stream: Box<IrcStream>) -> Self {
        Client {
            stream: stream,
            line_reader: LineReader::new(),
            nickname: None,
            username: None,
            auth: AuthSession::new(),
        }
    }

    pub fn read_message(&mut self) -> Option<Message> {
        match self.line_reader.read(&mut self.stream) {
            Some(line) => {
                trace!("<< {}", line);
                let stripped = line.trim();
                if stripped.len() == 0 {
                    None
                } else {
                    Some(Message::from_str(stripped))
                }
            },
            None => None
        }
    }

    pub fn set_nickname(&mut self, nickname: String) {
        self.nickname = Some(nickname);
    }

    pub fn join(&mut self, channel: &str) -> io::Result<usize> {
        let pfx = self.nickname.clone().unwrap();
        self.send(&Message {
            prefix: Some(pfx),
            command: Command::Join,
            args: vec![channel.to_string()],
            suffix: None
        })
    }

    pub fn pong(&mut self) -> io::Result<usize> {
        self.send(&Message::from(Command::Pong))
    }

    pub fn welcome(&mut self, message: &str) -> io::Result<usize> {
        self.send(&Message {
            prefix: Some("pto".to_string()),
            command: Command::Numeric(1),
            args: vec![message.to_string()],
            suffix: None
        })
    }

    pub fn send(&mut self, message: &Message) -> io::Result<usize> {
        trace!(">>> {}", message.to_string());
        self.stream.write(&message.to_string().trim().as_bytes())
            .and(self.stream.write("\r\n".as_bytes()))
    }
}

pub trait Server: AsEvented {
    fn accept(&mut self) -> Option<Client>;
}
<|endoftext|>


<commit_before>use mio::tcp::{TcpListener,TcpStream};
use mio::Token;
use std::net::SocketAddr;
use std::io::Write;

use irc::util::LineReader;
use irc::protocol::*;
use irc::security::AuthSession;

#[derive(Debug)]
pub struct Client {
    stream: TcpStream,
    line_reader: LineReader,
    nickname: Option<String>,
    username: Option<String>,
    pub auth: AuthSession,
}

impl Client {
    pub fn new(stream: TcpStream) -> Self {
        Client {
            stream: stream,
            line_reader: LineReader::new(),
            nickname: None,
            username: None,
            auth: AuthSession::new(),
        }
    }

    pub fn stream(&self) -> &TcpStream {
        &self.stream
    }

    pub fn read_message(&mut self) -> Option<Message> {
        match self.line_reader.read(&mut self.stream) {
            Some(line) => {
                let split: Vec<&str> = line.split(" ").collect();
                let mut args = Vec::new();
                for s in split[1..].iter() {
                    args.push(s.to_string());
                }
                let parsedCommand: Result<Command, Command> = split[0].parse();
                Some(Message{
                    prefix: None,
                    command: parsedCommand.ok().unwrap(),
                    args: args,
                    suffix: None
                })
            },
            None => None
        }
    }

    pub fn set_nickname(&mut self, nickname: String) {
        self.nickname = Some(nickname);
    }

    pub fn join(&mut self, channel: &str) {
        let pfx = self.nickname.clone().unwrap();
        self.send(&Message {
            prefix: Some(pfx),
            command: Command::Join,
            args: vec![channel.to_string()],
            suffix: None
        });
    }

    pub fn pong(&mut self) {
        self.send(&Message::from(Command::Pong));
    }

    pub fn welcome(&mut self, message: &str) {
        self.send(&Message {
            prefix: Some("pto".to_string()),
            command: Command::Numeric(1),
            args: vec![message.to_string()],
            suffix: None
        });
    }

    pub fn send(&mut self, message: &Message) {
        self.stream.write(&message.to_string().trim().as_bytes());
        self.stream.write("\r\n".as_bytes());
        println!("Wrote {:?}", message.to_string());
    }
}

pub struct Server {
    listener: TcpListener,
}

impl Server {
    pub fn new(addr: &SocketAddr) -> Self {
        Server {
            listener: TcpListener::bind(addr).unwrap(),
        }
    }

    pub fn accept(&mut self) -> Option<Client> {
         match self.listener.accept() {
             Ok(None) => None,
             Ok(Some((socket, _))) =>
                 Some(Client::new(socket)),
             Err(e) => panic!(e),
         }
    }

    pub fn listener(&self) -> &TcpListener {
        &self.listener
    }
}

<commit_msg>irc: streams: client: finish implementing proper tokenizing<commit_after>use mio::tcp::{TcpListener,TcpStream};
use mio::Token;
use std::net::SocketAddr;
use std::io::Write;

use irc::util::LineReader;
use irc::protocol::*;
use irc::security::AuthSession;

#[derive(Debug)]
pub struct Client {
    stream: TcpStream,
    line_reader: LineReader,
    nickname: Option<String>,
    username: Option<String>,
    pub auth: AuthSession,
}

impl Client {
    pub fn new(stream: TcpStream) -> Self {
        Client {
            stream: stream,
            line_reader: LineReader::new(),
            nickname: None,
            username: None,
            auth: AuthSession::new(),
        }
    }

    pub fn stream(&self) -> &TcpStream {
        &self.stream
    }

    fn split_parts(line: &str) -> (Option<String>, &str, Option<String>) {
        let mut prefix_end = 0;
        let mut args_end = 0;
        if line.chars().nth(0).unwrap() == ':' {
            for c in line[1..].chars() {
                prefix_end += 1;
                if c == ' ' {
                    break
                }
            }
        }

        let mut found_space = false;
        for c in line[prefix_end..].chars() {
            if c == ' ' {
                found_space = true;
            } else if c == ':' && found_space {
                args_end -= 1;
                break
            } else {
                found_space = false;
            }
            args_end += 1;
        }

        let prefix = match prefix_end {
            0 => None,
            _ => Some(line[1..prefix_end].to_string())
        };
        let args = &line[prefix_end..args_end].trim();
        let len = line.len();
        let suffix = if args_end == len {
            None
        } else {
            Some(line[args_end+2..].to_string())
        };

        (prefix, args, suffix)
    }

    pub fn read_message(&mut self) -> Option<Message> {
        match self.line_reader.read(&mut self.stream) {
            Some(line) => {
                let parts = Self::split_parts(line.trim());
                let split: Vec<&str> = parts.1.split(" ").collect();
                let mut args = Vec::new();
                for s in split[1..].iter() {
                    args.push(s.to_string());
                }
                let parsedCommand: Result<Command, Command> = split[0].parse();
                Some(Message{
                    prefix: parts.0,
                    command: parsedCommand.ok().unwrap(),
                    args: args,
                    suffix: parts.2
                })
            },
            None => None
        }
    }

    pub fn set_nickname(&mut self, nickname: String) {
        self.nickname = Some(nickname);
    }

    pub fn join(&mut self, channel: &str) {
        let pfx = self.nickname.clone().unwrap();
        self.send(&Message {
            prefix: Some(pfx),
            command: Command::Join,
            args: vec![channel.to_string()],
            suffix: None
        });
    }

    pub fn pong(&mut self) {
        self.send(&Message::from(Command::Pong));
    }

    pub fn welcome(&mut self, message: &str) {
        self.send(&Message {
            prefix: Some("pto".to_string()),
            command: Command::Numeric(1),
            args: vec![message.to_string()],
            suffix: None
        });
    }

    pub fn send(&mut self, message: &Message) {
        self.stream.write(&message.to_string().trim().as_bytes());
        self.stream.write("\r\n".as_bytes());
        println!("Wrote {:?}", message.to_string());
    }
}

pub struct Server {
    listener: TcpListener,
}

impl Server {
    pub fn new(addr: &SocketAddr) -> Self {
        Server {
            listener: TcpListener::bind(addr).unwrap(),
        }
    }

    pub fn accept(&mut self) -> Option<Client> {
         match self.listener.accept() {
             Ok(None) => None,
             Ok(Some((socket, _))) =>
                 Some(Client::new(socket)),
             Err(e) => panic!(e),
         }
    }

    pub fn listener(&self) -> &TcpListener {
        &self.listener
    }
}

<|endoftext|>


<commit_before>use ffi;
pub use ffi::{YamlEncoding, YamlScalarStyle, YamlSequenceStyle};
use std::libc::size_t;
use std::cast;
use std::ptr;
use std::c_str::CString;
use std::c_vec::CVec;

fn c_str_into_owned_bytes(c_str: *ffi::yaml_char_t) -> Option<~[u8]> {
    unsafe {
        if c_str == ptr::null() {
            None
        } else {
            Some(CString::new(c_str as *i8, false).as_bytes_no_nul().into_owned())
        }
    }
}

pub struct InternalEvent {
    pub event_mem: ffi::yaml_event_t
}

impl Drop for InternalEvent {
    fn drop(&mut self) {
        unsafe {
            ffi::yaml_event_delete(&mut self.event_mem as *mut ffi::yaml_event_t);
        }
    }
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlVersionDirective {
    pub major: int,
    pub minor: int,
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlTagDirective {
    pub handle: ~[u8],
    pub prefix: ~[u8],
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlSequenceParam {
    anchor: Option<~[u8]>,
    tag: Option<~[u8]>,
    implicit: bool,
    style: YamlSequenceStyle
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlScalarParam {
    anchor: Option<~[u8]>,
    tag: Option<~[u8]>,
    value: ~[u8],
    plain_implicit: bool,
    quoted_implicit: bool,
    style: YamlScalarStyle
}

#[deriving(Eq)]
#[deriving(Show)]
pub enum YamlEvent {
    YamlNoEvent,
    YamlStreamStartEvent(YamlEncoding),
    YamlStreamEndEvent,
    YamlDocumentStartEvent(Option<YamlVersionDirective>, ~[YamlTagDirective], bool),
    YamlDocumentEndEvent(bool),
    YamlAliasEvent(~[u8]),
    YamlScalarEvent(YamlScalarParam),
    YamlSequenceStartEvent(YamlSequenceParam),
    YamlSequenceEndEvent,
    YamlMappingStartEvent(YamlSequenceParam),
    YamlMappingEndEvent,
}

pub struct YamlEventStream<'r> {
    parser: ~YamlByteParser<'r>,
    end_stream: bool
}

impl<'r> Iterator<YamlEvent> for YamlEventStream<'r> {
    fn next(&mut self) -> Option<YamlEvent> {
        if self.end_stream {
            None
        } else {
            unsafe {
                let evt = self.parser.parse_event();
                match evt {
                    Some(YamlStreamEndEvent) => {
                        self.end_stream = true;
                    },
                    None => {
                        self.end_stream = true;
                    }
                    _ => ()
                }
                evt
            }
        }
    }
}

pub struct YamlByteParser<'r> {
    parser_mem: ffi::yaml_parser_t,
}

impl<'r> YamlByteParser<'r> {
    pub fn init(bytes: &'r [u8]) -> ~YamlByteParser<'r> {
        let mut parser = ~YamlByteParser {
            parser_mem: ffi::new_yaml_parser_t(),
        };

        unsafe {
            let res = ffi::yaml_parser_initialize(&mut parser.parser_mem as *mut ffi::yaml_parser_t);
            if res == 0 {
                fail!("failed to initialize yaml_parser_t");
            }
            ffi::yaml_parser_set_input_string(&mut parser.parser_mem as *mut ffi::yaml_parser_t, bytes.as_ptr() as *ffi::yaml_char_t, bytes.len() as size_t);
        }

        parser
    }

    unsafe fn parse_event(&mut self) -> Option<YamlEvent> {
        let mut event = InternalEvent {
            event_mem: ffi::yaml_event_t::new()
        };

        let res = ffi::yaml_parser_parse(&mut self.parser_mem as *mut ffi::yaml_parser_t, &mut event.event_mem as *mut ffi::yaml_event_t);
        if res == 0 {
            None
        } else {
            Some(match event.event_mem.event_type {
                ffi::YAML_NO_EVENT => YamlNoEvent,
                ffi::YAML_STREAM_START_EVENT => {
                    let evt_data: &ffi::yaml_stream_start_event_t = cast::transmute(&event.event_mem.data);
                    YamlStreamStartEvent(evt_data.encoding)
                },
                ffi::YAML_STREAM_END_EVENT => YamlStreamEndEvent,
                ffi::YAML_DOCUMENT_START_EVENT => {
                    let evt_data: &ffi::yaml_document_start_event_t = cast::transmute(&event.event_mem.data);
                    let vsn_dir = if evt_data.version_directive == ptr::null() {
                        None
                    } else {
                        let c_vsn_dir: &ffi::yaml_version_directive_t = cast::transmute(evt_data.version_directive);
                        Some(YamlVersionDirective { major: c_vsn_dir.major as int, minor: c_vsn_dir.minor as int })
                    };
                    let mut tag_dirs = ~[];
                    let mut tag_ptr = evt_data.tag_directives.start;
                    while tag_ptr != ptr::null() && tag_ptr != evt_data.tag_directives.end {
                        let tag_ref: &ffi::yaml_tag_directive_t = cast::transmute(tag_ptr);
                        let handle = CString::new(tag_ref.handle, false).as_bytes_no_nul().into_owned();
                        let prefix = CString::new(tag_ref.prefix, false).as_bytes_no_nul().into_owned();
                        tag_dirs.push(YamlTagDirective { handle: handle, prefix: prefix });
                        tag_ptr = tag_ptr.offset(1);
                    }
                    let implicit = evt_data.implicit != 0;

                    YamlDocumentStartEvent(vsn_dir, tag_dirs, implicit)
                },
                ffi::YAML_DOCUMENT_END_EVENT => {
                    let evt_data: &ffi::yaml_document_end_event_t = cast::transmute(&event.event_mem.data);
                    let implicit = evt_data.implicit != 0;

                    YamlDocumentEndEvent(implicit)
                },
                ffi::YAML_ALIAS_EVENT => {
                    let evt_data: &ffi::yaml_alias_event_t = cast::transmute(&event.event_mem.data);
                    let anchor = CString::new(evt_data.anchor as *i8, false).as_bytes_no_nul().into_owned();

                    YamlAliasEvent(anchor)
                },
                ffi::YAML_SCALAR_EVENT => {
                    let evt_data: &ffi::yaml_scalar_event_t = cast::transmute(&event.event_mem.data);
                    let value = CVec::new(evt_data.value as *mut u8, evt_data.length as uint).as_slice().into_owned();

                    YamlScalarEvent(YamlScalarParam {
                        anchor: c_str_into_owned_bytes(evt_data.anchor),
                        tag: c_str_into_owned_bytes(evt_data.tag),
                        value: value,
                        plain_implicit: evt_data.plain_implicit != 0,
                        quoted_implicit: evt_data.quoted_implicit != 0,
                        style: evt_data.style
                    })
                },
                ffi::YAML_SEQUENCE_START_EVENT => {
                    let evt_data: &ffi::yaml_sequence_start_event_t = cast::transmute(&event.event_mem.data);

                    YamlSequenceStartEvent(YamlSequenceParam {
                        anchor: c_str_into_owned_bytes(evt_data.anchor),
                        tag: c_str_into_owned_bytes(evt_data.tag),
                        implicit: evt_data.implicit != 0,
                        style: evt_data.style
                    })
                },
                ffi::YAML_SEQUENCE_END_EVENT => YamlSequenceEndEvent,
                ffi::YAML_MAPPING_START_EVENT => {
                    let evt_data: &ffi::yaml_mapping_start_event_t = cast::transmute(&event.event_mem.data);

                    YamlMappingStartEvent(YamlSequenceParam {
                        anchor: c_str_into_owned_bytes(evt_data.anchor),
                        tag: c_str_into_owned_bytes(evt_data.tag),
                        implicit: evt_data.implicit != 0,
                        style: evt_data.style
                    })
                },
                ffi::YAML_MAPPING_END_EVENT => YamlMappingEndEvent,
                _ => fail!("unknown event type")
            })
        }
    }

    pub fn parse(~self) -> YamlEventStream<'r> {
        YamlEventStream {
            parser: self,
            end_stream: false
        }
    }
}

impl<'r> Drop for YamlByteParser<'r> {
    fn drop(&mut self) {
        unsafe {
            ffi::yaml_parser_delete(&mut self.parser_mem as *mut ffi::yaml_parser_t);
        }
    }
}

#[test]
fn test_byte_parser() {
    let data = "[1, 2, 3]";
    let parser = YamlByteParser::init(data.as_bytes());
    let expected = ~[
        YamlStreamStartEvent(ffi::YamlUtf8Encoding),
        YamlDocumentStartEvent(None, ~[], true),
        YamlSequenceStartEvent(YamlSequenceParam{anchor: None, tag: None, implicit: true, style: ffi::YamlFlowSequenceStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[49u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[50u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[51u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlSequenceEndEvent,
        YamlDocumentEndEvent(true),
        YamlStreamEndEvent
    ];
    assert_eq!(expected, parser.parse().collect());
}

#[test]
fn test_byte_parser_mapping() {
    let data = "{\"a\": 1, \"b\":2}";
    let parser = YamlByteParser::init(data.as_bytes());
    let expected = ~[
        YamlStreamStartEvent(ffi::YamlUtf8Encoding),
        YamlDocumentStartEvent(None, ~[], true),
        YamlMappingStartEvent(YamlSequenceParam{anchor: None, tag: None, implicit: true, style: ffi::YamlFlowSequenceStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[97u8], plain_implicit: false, quoted_implicit: true, style: ffi::YamlDoubleQuotedScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[49u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[98u8], plain_implicit: false, quoted_implicit: true, style: ffi::YamlDoubleQuotedScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[50u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlMappingEndEvent,
        YamlDocumentEndEvent(true),
        YamlStreamEndEvent
    ];
    assert_eq!(expected, parser.parse().collect());
}
<commit_msg>Remove unnecessary pointer casts<commit_after>use ffi;
pub use ffi::{YamlEncoding, YamlScalarStyle, YamlSequenceStyle};
use std::libc::size_t;
use std::cast;
use std::ptr;
use std::c_str::CString;
use std::c_vec::CVec;

fn c_str_into_owned_bytes(c_str: *ffi::yaml_char_t) -> Option<~[u8]> {
    unsafe {
        if c_str == ptr::null() {
            None
        } else {
            Some(CString::new(c_str as *i8, false).as_bytes_no_nul().into_owned())
        }
    }
}

pub struct InternalEvent {
    pub event_mem: ffi::yaml_event_t
}

impl Drop for InternalEvent {
    fn drop(&mut self) {
        unsafe {
            ffi::yaml_event_delete(&mut self.event_mem);
        }
    }
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlVersionDirective {
    pub major: int,
    pub minor: int,
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlTagDirective {
    pub handle: ~[u8],
    pub prefix: ~[u8],
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlSequenceParam {
    anchor: Option<~[u8]>,
    tag: Option<~[u8]>,
    implicit: bool,
    style: YamlSequenceStyle
}

#[deriving(Eq)]
#[deriving(Show)]
pub struct YamlScalarParam {
    anchor: Option<~[u8]>,
    tag: Option<~[u8]>,
    value: ~[u8],
    plain_implicit: bool,
    quoted_implicit: bool,
    style: YamlScalarStyle
}

#[deriving(Eq)]
#[deriving(Show)]
pub enum YamlEvent {
    YamlNoEvent,
    YamlStreamStartEvent(YamlEncoding),
    YamlStreamEndEvent,
    YamlDocumentStartEvent(Option<YamlVersionDirective>, ~[YamlTagDirective], bool),
    YamlDocumentEndEvent(bool),
    YamlAliasEvent(~[u8]),
    YamlScalarEvent(YamlScalarParam),
    YamlSequenceStartEvent(YamlSequenceParam),
    YamlSequenceEndEvent,
    YamlMappingStartEvent(YamlSequenceParam),
    YamlMappingEndEvent,
}

pub struct YamlEventStream<'r> {
    parser: ~YamlByteParser<'r>,
    end_stream: bool
}

impl<'r> Iterator<YamlEvent> for YamlEventStream<'r> {
    fn next(&mut self) -> Option<YamlEvent> {
        if self.end_stream {
            None
        } else {
            unsafe {
                let evt = self.parser.parse_event();
                match evt {
                    Some(YamlStreamEndEvent) => {
                        self.end_stream = true;
                    },
                    None => {
                        self.end_stream = true;
                    }
                    _ => ()
                }
                evt
            }
        }
    }
}

pub struct YamlByteParser<'r> {
    parser_mem: ffi::yaml_parser_t,
}

impl<'r> YamlByteParser<'r> {
    pub fn init(bytes: &'r [u8]) -> ~YamlByteParser<'r> {
        let mut parser = ~YamlByteParser {
            parser_mem: ffi::new_yaml_parser_t(),
        };

        unsafe {
            let res = ffi::yaml_parser_initialize(&mut parser.parser_mem);
            if res == 0 {
                fail!("failed to initialize yaml_parser_t");
            }
            ffi::yaml_parser_set_input_string(&mut parser.parser_mem, bytes.as_ptr(), bytes.len() as size_t);
        }

        parser
    }

    unsafe fn parse_event(&mut self) -> Option<YamlEvent> {
        let mut event = InternalEvent {
            event_mem: ffi::yaml_event_t::new()
        };

        let res = ffi::yaml_parser_parse(&mut self.parser_mem, &mut event.event_mem);
        if res == 0 {
            None
        } else {
            Some(match event.event_mem.event_type {
                ffi::YAML_NO_EVENT => YamlNoEvent,
                ffi::YAML_STREAM_START_EVENT => {
                    let evt_data: &ffi::yaml_stream_start_event_t = cast::transmute(&event.event_mem.data);
                    YamlStreamStartEvent(evt_data.encoding)
                },
                ffi::YAML_STREAM_END_EVENT => YamlStreamEndEvent,
                ffi::YAML_DOCUMENT_START_EVENT => {
                    let evt_data: &ffi::yaml_document_start_event_t = cast::transmute(&event.event_mem.data);
                    let vsn_dir = if evt_data.version_directive == ptr::null() {
                        None
                    } else {
                        let c_vsn_dir: &ffi::yaml_version_directive_t = cast::transmute(evt_data.version_directive);
                        Some(YamlVersionDirective { major: c_vsn_dir.major as int, minor: c_vsn_dir.minor as int })
                    };
                    let mut tag_dirs = ~[];
                    let mut tag_ptr = evt_data.tag_directives.start;
                    while tag_ptr != ptr::null() && tag_ptr != evt_data.tag_directives.end {
                        let tag_ref: &ffi::yaml_tag_directive_t = cast::transmute(tag_ptr);
                        let handle = CString::new(tag_ref.handle, false).as_bytes_no_nul().into_owned();
                        let prefix = CString::new(tag_ref.prefix, false).as_bytes_no_nul().into_owned();
                        tag_dirs.push(YamlTagDirective { handle: handle, prefix: prefix });
                        tag_ptr = tag_ptr.offset(1);
                    }
                    let implicit = evt_data.implicit != 0;

                    YamlDocumentStartEvent(vsn_dir, tag_dirs, implicit)
                },
                ffi::YAML_DOCUMENT_END_EVENT => {
                    let evt_data: &ffi::yaml_document_end_event_t = cast::transmute(&event.event_mem.data);
                    let implicit = evt_data.implicit != 0;

                    YamlDocumentEndEvent(implicit)
                },
                ffi::YAML_ALIAS_EVENT => {
                    let evt_data: &ffi::yaml_alias_event_t = cast::transmute(&event.event_mem.data);
                    let anchor = CString::new(evt_data.anchor as *i8, false).as_bytes_no_nul().into_owned();

                    YamlAliasEvent(anchor)
                },
                ffi::YAML_SCALAR_EVENT => {
                    let evt_data: &ffi::yaml_scalar_event_t = cast::transmute(&event.event_mem.data);
                    let value = CVec::new(evt_data.value as *mut u8, evt_data.length as uint).as_slice().into_owned();

                    YamlScalarEvent(YamlScalarParam {
                        anchor: c_str_into_owned_bytes(evt_data.anchor),
                        tag: c_str_into_owned_bytes(evt_data.tag),
                        value: value,
                        plain_implicit: evt_data.plain_implicit != 0,
                        quoted_implicit: evt_data.quoted_implicit != 0,
                        style: evt_data.style
                    })
                },
                ffi::YAML_SEQUENCE_START_EVENT => {
                    let evt_data: &ffi::yaml_sequence_start_event_t = cast::transmute(&event.event_mem.data);

                    YamlSequenceStartEvent(YamlSequenceParam {
                        anchor: c_str_into_owned_bytes(evt_data.anchor),
                        tag: c_str_into_owned_bytes(evt_data.tag),
                        implicit: evt_data.implicit != 0,
                        style: evt_data.style
                    })
                },
                ffi::YAML_SEQUENCE_END_EVENT => YamlSequenceEndEvent,
                ffi::YAML_MAPPING_START_EVENT => {
                    let evt_data: &ffi::yaml_mapping_start_event_t = cast::transmute(&event.event_mem.data);

                    YamlMappingStartEvent(YamlSequenceParam {
                        anchor: c_str_into_owned_bytes(evt_data.anchor),
                        tag: c_str_into_owned_bytes(evt_data.tag),
                        implicit: evt_data.implicit != 0,
                        style: evt_data.style
                    })
                },
                ffi::YAML_MAPPING_END_EVENT => YamlMappingEndEvent,
                _ => fail!("unknown event type")
            })
        }
    }

    pub fn parse(~self) -> YamlEventStream<'r> {
        YamlEventStream {
            parser: self,
            end_stream: false
        }
    }
}

impl<'r> Drop for YamlByteParser<'r> {
    fn drop(&mut self) {
        unsafe {
            ffi::yaml_parser_delete(&mut self.parser_mem);
        }
    }
}

#[test]
fn test_byte_parser() {
    let data = "[1, 2, 3]";
    let parser = YamlByteParser::init(data.as_bytes());
    let expected = ~[
        YamlStreamStartEvent(ffi::YamlUtf8Encoding),
        YamlDocumentStartEvent(None, ~[], true),
        YamlSequenceStartEvent(YamlSequenceParam{anchor: None, tag: None, implicit: true, style: ffi::YamlFlowSequenceStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[49u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[50u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[51u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlSequenceEndEvent,
        YamlDocumentEndEvent(true),
        YamlStreamEndEvent
    ];
    assert_eq!(expected, parser.parse().collect());
}

#[test]
fn test_byte_parser_mapping() {
    let data = "{\"a\": 1, \"b\":2}";
    let parser = YamlByteParser::init(data.as_bytes());
    let expected = ~[
        YamlStreamStartEvent(ffi::YamlUtf8Encoding),
        YamlDocumentStartEvent(None, ~[], true),
        YamlMappingStartEvent(YamlSequenceParam{anchor: None, tag: None, implicit: true, style: ffi::YamlFlowSequenceStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[97u8], plain_implicit: false, quoted_implicit: true, style: ffi::YamlDoubleQuotedScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[49u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[98u8], plain_implicit: false, quoted_implicit: true, style: ffi::YamlDoubleQuotedScalarStyle}),
        YamlScalarEvent(YamlScalarParam{anchor: None, tag: None, value: ~[50u8], plain_implicit: true, quoted_implicit: false, style: ffi::YamlPlainScalarStyle}),
        YamlMappingEndEvent,
        YamlDocumentEndEvent(true),
        YamlStreamEndEvent
    ];
    assert_eq!(expected, parser.parse().collect());
}
<|endoftext|>


<commit_before>use std::collections::{HashSet, HashMap};
use std::time::{Instant, Duration};
use std::{self, io};
use torrent::Peer;
use util::random_sample;

pub struct Choker {
    unchoked: Vec<usize>,
    interested: HashSet<usize>,
    last_updated: Instant,
}

pub struct SwapRes {
    pub choked: (usize, io::Result<()>),
    pub unchoked: (usize, io::Result<()>),
}

impl Choker {
    pub fn new() -> Choker {
        Choker {
            unchoked: Vec::with_capacity(5),
            interested: HashSet::new(),
            last_updated: Instant::now()
        }
    }

    pub fn add_peer(&mut self, peer: &mut Peer) -> io::Result<()> {
        if self.unchoked.len() < 4 {
            self.unchoked.push(peer.id);
            peer.downloaded = 0;
            peer.uploaded = 0;
            peer.unchoke()
        } else {
            self.interested.insert(peer.id);
            Ok(())
        }
    }

    pub fn remove_peer(&mut self, peer: &mut Peer) {
        self.interested.remove(&peer.id);
    }

    fn update_timer(&mut self) -> Result<(), ()> {
        if self.last_updated.elapsed() < Duration::from_secs(10) || self.unchoked.len() < 5 {
            Err(())
        } else {
            self.last_updated = Instant::now();
            Ok(())
        }
    }

    pub fn update_upload(&mut self, peers: &mut HashMap<usize, Peer>) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }
        let (slowest, _) = self.unchoked.iter().enumerate().fold((0, std::usize::MAX), |(slowest, min), (idx, id)| {
            let ul = peers[id].uploaded;
            peers.get_mut(id).unwrap().uploaded = 0;
            if ul < min {
                (idx, ul)
            } else {
                (slowest, min)
            }
        });
        Some(self.swap_peer(slowest, peers))
    }

    pub fn update_download(&mut self, peers: &mut HashMap<usize, Peer>) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }

        let (slowest, _) = self.unchoked.iter().enumerate().fold((0, std::usize::MAX), |(slowest, min), (idx, id)| {
            let ul = peers[id].downloaded;
            peers.get_mut(id).unwrap().downloaded = 0;
            if ul < min {
                (idx, ul)
            } else {
                (slowest, min)
            }
        });
        Some(self.swap_peer(slowest, peers))
    }

    fn swap_peer(&mut self, idx: usize, peers: &mut HashMap<usize, Peer>) -> SwapRes {
        let id = self.unchoked.remove(idx);
        let cres = {
            let peer = peers.get_mut(&id).unwrap();
            self.interested.insert(id);
            peer.choke()
        };

        // Unchoke one random interested peer
        let random_id = *random_sample(self.interested.iter()).unwrap();
        let peer = peers.get_mut(&random_id).unwrap();
        self.interested.remove(&random_id);
        self.unchoked.push(random_id);
        let ures = peer.unchoke();
        return SwapRes { choked: (id, cres), unchoked: (random_id, ures) };
    }
}
<commit_msg>Add initial choker tests<commit_after>use std::collections::{HashSet, HashMap};
use std::time::{Instant, Duration};
use std::{self, io};
use torrent::Peer;
use util::random_sample;

pub struct Choker {
    unchoked: Vec<usize>,
    interested: HashSet<usize>,
    last_updated: Instant,
}

pub struct SwapRes {
    pub choked: (usize, io::Result<()>),
    pub unchoked: (usize, io::Result<()>),
}

impl Choker {
    pub fn new() -> Choker {
        Choker {
            unchoked: Vec::with_capacity(5),
            interested: HashSet::new(),
            last_updated: Instant::now()
        }
    }

    pub fn add_peer(&mut self, peer: &mut Peer) -> io::Result<()> {
        if self.unchoked.len() < 5 {
            self.unchoked.push(peer.id);
            peer.downloaded = 0;
            peer.uploaded = 0;
            peer.unchoke()
        } else {
            self.interested.insert(peer.id);
            Ok(())
        }
    }

    pub fn remove_peer(&mut self, peer: &mut Peer) {
        if self.unchoked.contains(&peer.id) {
        } else {
            self.interested.remove(&peer.id);
        }
    }

    fn update_timer(&mut self) -> Result<(), ()> {
        if self.last_updated.elapsed() < Duration::from_secs(10) || self.unchoked.len() < 5 {
            Err(())
        } else {
            self.last_updated = Instant::now();
            Ok(())
        }
    }

    pub fn update_upload(&mut self, peers: &mut HashMap<usize, Peer>) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }
        let (slowest, _) = self.unchoked.iter().enumerate().fold((0, std::usize::MAX), |(slowest, min), (idx, id)| {
            let ul = peers[id].uploaded;
            peers.get_mut(id).unwrap().uploaded = 0;
            if ul < min {
                (idx, ul)
            } else {
                (slowest, min)
            }
        });
        Some(self.swap_peer(slowest, peers))
    }

    pub fn update_download(&mut self, peers: &mut HashMap<usize, Peer>) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }

        let (slowest, _) = self.unchoked.iter().enumerate().fold((0, std::usize::MAX), |(slowest, min), (idx, id)| {
            let ul = peers[id].downloaded;
            peers.get_mut(id).unwrap().downloaded = 0;
            if ul < min {
                (idx, ul)
            } else {
                (slowest, min)
            }
        });
        Some(self.swap_peer(slowest, peers))
    }

    fn swap_peer(&mut self, idx: usize, peers: &mut HashMap<usize, Peer>) -> SwapRes {
        let id = self.unchoked.remove(idx);
        let cres = {
            let peer = peers.get_mut(&id).unwrap();
            self.interested.insert(id);
            peer.choke()
        };

        // Unchoke one random interested peer
        let random_id = *random_sample(self.interested.iter()).unwrap();
        let peer = peers.get_mut(&random_id).unwrap();
        self.interested.remove(&random_id);
        self.unchoked.push(random_id);
        let ures = peer.unchoke();
        return SwapRes { choked: (id, cres), unchoked: (random_id, ures) };
    }
}

mod tests {
    use super::Choker;
    use torrent::Peer;
    use socket::Socket;

    #[test]
    fn test_add_peers() {
        let mut c = Choker::new();
        for i in 0..6 {
            let mut p = Peer::new(Socket::empty());
            p.id = i;
            c.add_peer(&mut p).is_err();
        }
        assert_eq!(c.unchoked.len(), 5);
        assert_eq!(c.interested.len(), 1);
    }
}
<|endoftext|>