<commit_before>extern crate num_cpus;
extern crate rand;
extern crate crossbeam;

use self::rand::{thread_rng, Rng};
use self::crossbeam::{scope, ScopedJoinHandle};

use std::ops::Range;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::{Mutex, RwLock, MutexGuard};
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::spawn;
use std::collections::BTreeSet;

use task::{TaskGenerator, Task};
use candidate::{WorkingCandidate, Candidate};
use context::Context;
use scaling::{ScalingFunction, proportionate};
use result::{Result as AbcResult, Error as AbcError};

/// Manages the parameters of the ABC algorithm.
pub struct HiveBuilder<Ctx: Context> {
    workers: usize,
    observers: usize,
    retries: usize,
    context: Ctx,
    threads: usize,
    scale: Box<ScalingFunction>,
}

impl<Ctx: Context> HiveBuilder<Ctx> {
    /// Creates a new hive.
    ///
    /// * `context` - Factory-like state that can be used while generating solutions.
    /// * `workers` - Number of working solution candidates to maintain at a time.
    pub fn new(context: Ctx, workers: usize) -> HiveBuilder<Ctx> {
        if workers == 0 {
            panic!("HiveBuilder must have at least one worker.");
        }

        HiveBuilder {
            workers: workers,
            observers: workers,
            retries: workers,

            context: context,
            threads: num_cpus::get(),
            scale: proportionate(),
        }
    }

    /// Sets the number of "bees" that will pick a candidate to work on at random.
    ///
    /// This defaults to the number of workers.
    pub fn set_observers(mut self, observers: usize) -> HiveBuilder<Ctx> {
        self.observers = observers;
        self
    }

    /// Sets the number of times a candidate can go unimproved before being reinitialized.
    ///
    /// This defaults to the number of workers.
    pub fn set_retries(mut self, retries: usize) -> HiveBuilder<Ctx> {
        self.retries = retries;
        self
    }

    /// Sets the number of worker threads to use while running.
    pub fn set_threads(mut self, threads: usize) -> HiveBuilder<Ctx> {
        self.threads = threads;
        self
    }

    /// Sets the scaling function for observers to use.
    pub fn set_scaling(mut self, scale: Box<ScalingFunction>) -> HiveBuilder<Ctx> {
        self.scale = scale;
        self
    }

    /// Activates the `HiveBuilder` to create a runnable object.
    pub fn build(self) -> AbcResult<Hive<Ctx>> {
        Hive::new(self)
    }

    fn new_candidate(&self) -> Candidate<Ctx::Solution> {
        let solution = self.context.make();
        let fitness = self.context.evaluate_fitness(&solution);
        Candidate::new(solution, fitness)
    }
}

/// Runs the ABC algorithm, maintaining any necessary state.
pub struct Hive<Ctx: Context> {
    hive: HiveBuilder<Ctx>,

    working: Vec<RwLock<WorkingCandidate<Ctx::Solution>>>,
    best: Mutex<Candidate<Ctx::Solution>>,
    scouting: RwLock<BTreeSet<usize>>,

    tasks: Mutex<Option<TaskGenerator>>,
    sender: Option<Mutex<Sender<Candidate<Ctx::Solution>>>>,
}

impl<Ctx: Context> Hive<Ctx> {
    fn new(hive: HiveBuilder<Ctx>) -> AbcResult<Hive<Ctx>> {
        // Start by populating the field with an initial set of solution candidates.

        // Feed the worker threads a total of N items, each signifying that
        // we need another candidate.
        let tokens: Mutex<Range<usize>> = Mutex::new(0..hive.workers);

        let candidates = Mutex::new(Vec::with_capacity(hive.workers));
        let mut handles = Vec::<ScopedJoinHandle<AbcResult<()>>>::with_capacity(hive.threads);

        try!(crossbeam::scope(|scope| {
            for _ in 0..hive.threads {
                handles.push(scope.spawn(|| {
                    while let Some(_) = {
                        let mut guard = tokens.lock().unwrap();
                        guard.next()
                    } {
                        let candidate = hive.new_candidate();
                        try!(candidates.lock()).push(candidate);
                    }
                    Ok(())
                }));
            }

            // Gather and return `Ok` iff all of the workers finished
            // successfully, otherwise abort the construction.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
        }));

        // We don't need the mutex anymore, since we're no longer populating
        // the candidate set from multiple threads.
        let mut candidates = try!(candidates.into_inner());

        // Find the current best candidate, since we want to cache the best
        // at any given moment.
        let best = {
            let (first, rest) = candidates.split_first().unwrap();
            let best_candidate = rest.iter()
                                     .fold(first, |best, next| {
                                         if next.fitness > best.fitness {
                                             next
                                         } else {
                                             best
                                         }
                                     });
            Mutex::new(best_candidate.clone())
        };

        // Wrap the candidates in a structure that will let the eventual
        // thread swarm work on them.
        let working = candidates.drain(..)
                                .map(|c| RwLock::new(WorkingCandidate::new(c, hive.retries)))
                                .collect::<Vec<RwLock<WorkingCandidate<Ctx::Solution>>>>();

        Ok(Hive {
            hive: hive,
            working: working,
            best: best,
            scouting: RwLock::new(BTreeSet::new()),
            tasks: Mutex::new(None),
            sender: None,
        })
    }

    /// Clone a snapshot of the current set of working candidates.
    ///
    /// The goal of this function is to hold a guard for each solution for as
    /// little time as possible, so we can get out of the way of the other
    /// threads. To this end, we clone the solutions, so that the thread can do
    /// its work on a snapshot.
    fn current_working(&self) -> AbcResult<Vec<Candidate<Ctx::Solution>>> {
        let mut current_working = Vec::with_capacity(self.working.len());
        for candidate_mutex in &self.working {
            let read_guard = try!(candidate_mutex.read());
            current_working.push(read_guard.candidate.clone())
        }
        Ok(current_working)
    }

    /// Returns a guard for the current best solution found by the hive.
    ///
    /// If the hive is running, you should drop the guard returned by this
    /// function as soon as convenient, since the logic of the hive can block
    /// on the availability of the associated mutex. If you plan on performing
    /// expensive computations, you should `drop` the guard as soon as
    /// possible, or acquire and clone it within a small block.
    pub fn get(&self) -> AbcResult<MutexGuard<Candidate<Ctx::Solution>>> {
        self.best.lock().map_err(AbcError::from)
    }

    /// Perform greedy selection between a new candidate and the current best.
    fn consider_improvement(&self, candidate: &Candidate<Ctx::Solution>) -> AbcResult<()> {
        let mut best_guard = try!(self.best.lock());
        if candidate.fitness > best_guard.fitness {
            *best_guard = candidate.clone();
            if let Some(mutex) = self.sender.as_ref() {
                // We're streaming, so we need to post the improved candidate.
                let sender_guard = try!(mutex.lock());
                // If this errors, the receiver was dropped, so we're done.
                if let Err(_) = sender_guard.send(candidate.clone()) {
                    try!(self.stop());
                }
            }
        }
        Ok(())
    }

    fn work_on(&self, current_working: &[Candidate<Ctx::Solution>], n: usize) -> AbcResult<()> {
        let variant_solution = self.hive.context.explore(current_working, n);
        let variant_fitness = self.hive.context.evaluate_fitness(&variant_solution);
        let variant = Candidate::new(variant_solution, variant_fitness);
        let mut write_guard = try!(self.working[n].write());
        if variant.fitness > write_guard.candidate.fitness {
            *write_guard = WorkingCandidate::new(variant, self.hive.retries);
            try!(self.consider_improvement(&write_guard.candidate));
        } else {
            write_guard.deplete();
            // Scouting has been folded into the working process
            if write_guard.expired() {
                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.insert(n);
                drop(scouting_guard);
                drop(write_guard);

                let candidate = self.hive.new_candidate();
                let mut write_guard = try!(self.working[n].write());
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
                drop(write_guard);

                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.remove(&n);
            }
        }
        Ok(())
    }

    fn choose(&self, current_working: &[Candidate<Ctx::Solution>]) -> AbcResult<usize> {
        let fitnesses = (self.hive.scale)(current_working.iter()
                                                         .map(|candidate| candidate.fitness)
                                                         .collect::<Vec<f64>>());

        // Avoid observing candidates that are being scouted.
        let scouting_guard = try!(self.scouting.read());
        let running_totals = fitnesses.iter()
                                      .enumerate()
                                      .filter(|&(ref i, _)| !scouting_guard.contains(i))
                                      .scan(0f64, |total, (i, fitness)| {
                                          *total += *fitness;
                                          Some((i, *total))
                                      })
                                      .collect::<Vec<(usize, f64)>>();
        drop(scouting_guard);

        // Multiplying the choice point is equivalent to, and more efficient than, normalizing
        // all of the scaled fitnesses and having a choice point in [0,1)
        match running_totals.last() {
            Some(&(_, total_fitness)) => {
                let choice_point = thread_rng().next_f64() * total_fitness;
                for &(i, total) in &running_totals {
                    if total > choice_point {
                        return Ok(i);
                    }
                }
                unreachable!();
            }

            // If we are currently scouting all of the solutions, pick one at random.
            None => Ok(thread_rng().gen_range::<usize>(0, fitnesses.len())),
        }
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => {
                // If the worker's candidate is in the middle of being replaced, just skip it.
                let scouting_guard = try!(self.scouting.read());
                if scouting_guard.contains(&n) {
                    return Ok(());
                }
                n
            }
            Task::Observer(_) => try!(self.choose(&current_working)),
        };
        self.work_on(&current_working, index)
    }

    fn run(&self, tasks: TaskGenerator) -> AbcResult<()> {
        let mut guard = try!(self.tasks.lock());
        *guard = Some(tasks);
        drop(guard);

        let mut handles: Vec<ScopedJoinHandle<AbcResult<()>>> = Vec::new();

        scope(|scope| {
            for _ in 0..self.hive.threads {
                handles.push(scope.spawn(|| {
                    loop {
                        let mut guard = try!(self.tasks.lock());
                        let task = guard.as_mut().and_then(|gen| gen.next());
                        drop(guard);

                        match task {
                            Some(t) => try!(self.execute(&t)),
                            None => return Ok(()),
                        };
                    }
                }));
            }

            // Returns `Ok(())` only if all threads join cleanly, and the task
            // cycle is successfully cleared away.
            //
            // We avoid `try!` because we want all of the following logic to
            // execute unconditionally.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
                   .and(self.tasks
                            .lock()
                            .map(|mut tasks_guard| *tasks_guard = None)
                            .map_err(AbcError::from))
        })
    }

    /// Runs for a fixed number of rounds, then return the best solution found.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok` with a `Candidate`.
    pub fn run_for_rounds(&self, rounds: usize) -> AbcResult<Candidate<Ctx::Solution>> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers).max_rounds(rounds);
        try!(self.run(tasks));
        self.get().map(|guard| guard.clone())
    }

    /// Run indefinitely.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok(())`.
    pub fn run_forever(&self) -> AbcResult<()> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
        self.run(tasks)
    }

    /// Stops a running hive.
    ///
    /// If a worker thread has panicked, this returns `Err(abc::Error)`.
    pub fn stop(&self) -> AbcResult<()> {
        let mut tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_mut().map_or((), |t| t.stop()))
    }

    /// Each new best candidate will be sent to `sender`.
    ///
    /// This is kept in a separate function so that the hive can be borrowed
    /// while running.
    pub fn set_sender(&mut self, sender: Sender<Candidate<Ctx::Solution>>) {
        if let Ok(best_guard) = self.best.lock() {
            sender.send(best_guard.clone()).unwrap_or(());
        }
        self.sender = Some(Mutex::new(sender));
    }

    /// Returns the current round of a running hive.
    ///
    /// If a worker thread has panicked and poisoned the task generator lock,
    /// `get_round` will return `Err(abc::Error)`. If the hive has not been
    /// run, `get_round` will return `Ok(None)`.
    ///
    /// If the hive is running, this will return `Ok(Some(n))`. `n` will start
    /// at 0, and increment each time every task in the round has been claimed
    /// (though not necessarily completed) by a worker thread.
    pub fn get_round(&self) -> AbcResult<Option<usize>> {
        let tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_ref().map(|tasks| tasks.round))
    }

    /// Get a reference to the hive's context.
    pub fn context(&self) -> &Ctx {
        &self.hive.context
    }
}

impl<Ctx: Context + 'static> Hive<Ctx> {
    /// Runs indefinitely in the background, providing a stream of results.
    ///
    /// This method consumes the hive, which will run until the `HiveBuilder`
    /// object is dropped. It returns an `mpsc::Receiver`, which receives a
    /// `Candidate` each time the hive improves on its best solution.
    pub fn stream(mut self) -> Receiver<Candidate<Ctx::Solution>> {
        let (sender, receiver) = channel();
        spawn(move || {
            self.set_sender(sender);
            let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
            self.run(tasks)
        });
        receiver
    }
}

impl<Ctx: Context> Debug for Hive<Ctx>
    where Ctx::Solution: Debug
{
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        for mutex in (&self.working).iter() {
            let working = mutex.read().unwrap();
            try!(write!(f, "..{:?}..\n", working.candidate));
        }
        let best_candidate = self.get().unwrap();
        write!(f, ">>{:?}<<", *best_candidate)
    }
}

impl<Ctx: Context> Drop for Hive<Ctx> {
    fn drop(&mut self) {
        self.stop().unwrap_or(())
    }
}
<commit_msg>Drop row lock while checking for best in scouting<commit_after>extern crate num_cpus;
extern crate rand;
extern crate crossbeam;

use self::rand::{thread_rng, Rng};
use self::crossbeam::{scope, ScopedJoinHandle};

use std::ops::Range;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::{Mutex, RwLock, MutexGuard};
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::spawn;
use std::collections::BTreeSet;

use task::{TaskGenerator, Task};
use candidate::{WorkingCandidate, Candidate};
use context::Context;
use scaling::{ScalingFunction, proportionate};
use result::{Result as AbcResult, Error as AbcError};

/// Manages the parameters of the ABC algorithm.
pub struct HiveBuilder<Ctx: Context> {
    workers: usize,
    observers: usize,
    retries: usize,
    context: Ctx,
    threads: usize,
    scale: Box<ScalingFunction>,
}

impl<Ctx: Context> HiveBuilder<Ctx> {
    /// Creates a new hive.
    ///
    /// * `context` - Factory-like state that can be used while generating solutions.
    /// * `workers` - Number of working solution candidates to maintain at a time.
    pub fn new(context: Ctx, workers: usize) -> HiveBuilder<Ctx> {
        if workers == 0 {
            panic!("HiveBuilder must have at least one worker.");
        }

        HiveBuilder {
            workers: workers,
            observers: workers,
            retries: workers,

            context: context,
            threads: num_cpus::get(),
            scale: proportionate(),
        }
    }

    /// Sets the number of "bees" that will pick a candidate to work on at random.
    ///
    /// This defaults to the number of workers.
    pub fn set_observers(mut self, observers: usize) -> HiveBuilder<Ctx> {
        self.observers = observers;
        self
    }

    /// Sets the number of times a candidate can go unimproved before being reinitialized.
    ///
    /// This defaults to the number of workers.
    pub fn set_retries(mut self, retries: usize) -> HiveBuilder<Ctx> {
        self.retries = retries;
        self
    }

    /// Sets the number of worker threads to use while running.
    pub fn set_threads(mut self, threads: usize) -> HiveBuilder<Ctx> {
        self.threads = threads;
        self
    }

    /// Sets the scaling function for observers to use.
    pub fn set_scaling(mut self, scale: Box<ScalingFunction>) -> HiveBuilder<Ctx> {
        self.scale = scale;
        self
    }

    /// Activates the `HiveBuilder` to create a runnable object.
    pub fn build(self) -> AbcResult<Hive<Ctx>> {
        Hive::new(self)
    }

    fn new_candidate(&self) -> Candidate<Ctx::Solution> {
        let solution = self.context.make();
        let fitness = self.context.evaluate_fitness(&solution);
        Candidate::new(solution, fitness)
    }
}

/// Runs the ABC algorithm, maintaining any necessary state.
pub struct Hive<Ctx: Context> {
    hive: HiveBuilder<Ctx>,

    working: Vec<RwLock<WorkingCandidate<Ctx::Solution>>>,
    best: Mutex<Candidate<Ctx::Solution>>,
    scouting: RwLock<BTreeSet<usize>>,

    tasks: Mutex<Option<TaskGenerator>>,
    sender: Option<Mutex<Sender<Candidate<Ctx::Solution>>>>,
}

impl<Ctx: Context> Hive<Ctx> {
    fn new(hive: HiveBuilder<Ctx>) -> AbcResult<Hive<Ctx>> {
        // Start by populating the field with an initial set of solution candidates.

        // Feed the worker threads a total of N items, each signifying that
        // we need another candidate.
        let tokens: Mutex<Range<usize>> = Mutex::new(0..hive.workers);

        let candidates = Mutex::new(Vec::with_capacity(hive.workers));
        let mut handles = Vec::<ScopedJoinHandle<AbcResult<()>>>::with_capacity(hive.threads);

        try!(crossbeam::scope(|scope| {
            for _ in 0..hive.threads {
                handles.push(scope.spawn(|| {
                    while let Some(_) = {
                        let mut guard = tokens.lock().unwrap();
                        guard.next()
                    } {
                        let candidate = hive.new_candidate();
                        try!(candidates.lock()).push(candidate);
                    }
                    Ok(())
                }));
            }

            // Gather and return `Ok` iff all of the workers finished
            // successfully, otherwise abort the construction.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
        }));

        // We don't need the mutex anymore, since we're no longer populating
        // the candidate set from multiple threads.
        let mut candidates = try!(candidates.into_inner());

        // Find the current best candidate, since we want to cache the best
        // at any given moment.
        let best = {
            let (first, rest) = candidates.split_first().unwrap();
            let best_candidate = rest.iter()
                                     .fold(first, |best, next| {
                                         if next.fitness > best.fitness {
                                             next
                                         } else {
                                             best
                                         }
                                     });
            Mutex::new(best_candidate.clone())
        };

        // Wrap the candidates in a structure that will let the eventual
        // thread swarm work on them.
        let working = candidates.drain(..)
                                .map(|c| RwLock::new(WorkingCandidate::new(c, hive.retries)))
                                .collect::<Vec<RwLock<WorkingCandidate<Ctx::Solution>>>>();

        Ok(Hive {
            hive: hive,
            working: working,
            best: best,
            scouting: RwLock::new(BTreeSet::new()),
            tasks: Mutex::new(None),
            sender: None,
        })
    }

    /// Clone a snapshot of the current set of working candidates.
    ///
    /// The goal of this function is to hold a guard for each solution for as
    /// little time as possible, so we can get out of the way of the other
    /// threads. To this end, we clone the solutions, so that the thread can do
    /// its work on a snapshot.
    fn current_working(&self) -> AbcResult<Vec<Candidate<Ctx::Solution>>> {
        let mut current_working = Vec::with_capacity(self.working.len());
        for candidate_mutex in &self.working {
            let read_guard = try!(candidate_mutex.read());
            current_working.push(read_guard.candidate.clone())
        }
        Ok(current_working)
    }

    /// Returns a guard for the current best solution found by the hive.
    ///
    /// If the hive is running, you should drop the guard returned by this
    /// function as soon as convenient, since the logic of the hive can block
    /// on the availability of the associated mutex. If you plan on performing
    /// expensive computations, you should `drop` the guard as soon as
    /// possible, or acquire and clone it within a small block.
    pub fn get(&self) -> AbcResult<MutexGuard<Candidate<Ctx::Solution>>> {
        self.best.lock().map_err(AbcError::from)
    }

    /// Perform greedy selection between a new candidate and the current best.
    fn consider_improvement(&self, candidate: &Candidate<Ctx::Solution>) -> AbcResult<()> {
        let mut best_guard = try!(self.best.lock());
        if candidate.fitness > best_guard.fitness {
            *best_guard = candidate.clone();
            if let Some(mutex) = self.sender.as_ref() {
                // We're streaming, so we need to post the improved candidate.
                let sender_guard = try!(mutex.lock());
                // If this errors, the receiver was dropped, so we're done.
                if let Err(_) = sender_guard.send(candidate.clone()) {
                    try!(self.stop());
                }
            }
        }
        Ok(())
    }

    fn work_on(&self, current_working: &[Candidate<Ctx::Solution>], n: usize) -> AbcResult<()> {
        let variant_solution = self.hive.context.explore(current_working, n);
        let variant_fitness = self.hive.context.evaluate_fitness(&variant_solution);
        let variant = Candidate::new(variant_solution, variant_fitness);
        let mut write_guard = try!(self.working[n].write());
        if variant.fitness > write_guard.candidate.fitness {
            *write_guard = WorkingCandidate::new(variant, self.hive.retries);
            try!(self.consider_improvement(&write_guard.candidate));
        } else {
            write_guard.deplete();
            // Scouting has been folded into the working process
            if write_guard.expired() {
                {
                    let mut scouting_guard = try!(self.scouting.write());
                    scouting_guard.insert(n);
                }
                drop(write_guard);

                let candidate = self.hive.new_candidate();
                try!(self.consider_improvement(&candidate));
                {
                    let mut write_guard = try!(self.working[n].write());
                    *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                }

                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.remove(&n);
            }
        }
        Ok(())
    }

    fn choose(&self, current_working: &[Candidate<Ctx::Solution>]) -> AbcResult<usize> {
        let fitnesses = (self.hive.scale)(current_working.iter()
                                                         .map(|candidate| candidate.fitness)
                                                         .collect::<Vec<f64>>());

        // Avoid observing candidates that are being scouted.
        let running_totals = {
            let scouting_guard = try!(self.scouting.read());
            fitnesses.iter()
                     .enumerate()
                     .filter(|&(ref i, _)| !scouting_guard.contains(i))
                     .scan(0f64, |total, (i, fitness)| {
                         *total += *fitness;
                         Some((i, *total))
                     })
                     .collect::<Vec<(usize, f64)>>()
        };

        // Multiplying the choice point is equivalent to, and more efficient than, normalizing
        // all of the scaled fitnesses and having a choice point in [0,1)
        match running_totals.last() {
            Some(&(_, total_fitness)) => {
                let choice_point = thread_rng().next_f64() * total_fitness;
                for &(i, total) in &running_totals {
                    if total > choice_point {
                        return Ok(i);
                    }
                }
                unreachable!();
            }

            // If we are currently scouting all of the solutions, pick one at random.
            None => Ok(thread_rng().gen_range::<usize>(0, fitnesses.len())),
        }
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => {
                // If the worker's candidate is in the middle of being replaced, just skip it.
                let scouting_guard = try!(self.scouting.read());
                if scouting_guard.contains(&n) {
                    return Ok(());
                }
                n
            }
            Task::Observer(_) => try!(self.choose(&current_working)),
        };
        self.work_on(&current_working, index)
    }

    fn run(&self, tasks: TaskGenerator) -> AbcResult<()> {
        {
            let mut guard = try!(self.tasks.lock());
            *guard = Some(tasks);
        }

        let mut handles: Vec<ScopedJoinHandle<AbcResult<()>>> = Vec::new();

        scope(|scope| {
            for _ in 0..self.hive.threads {
                handles.push(scope.spawn(|| {
                    loop {
                        let task = {
                            let mut guard = try!(self.tasks.lock());
                            guard.as_mut().and_then(|gen| gen.next())
                        };

                        match task {
                            Some(t) => try!(self.execute(&t)),
                            None => return Ok(()),
                        };
                    }
                }));
            }

            // Returns `Ok(())` only if all threads join cleanly, and the task
            // cycle is successfully cleared away.
            //
            // We avoid `try!` because we want all of the following logic to
            // execute unconditionally.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
                   .and(self.tasks
                            .lock()
                            .map(|mut tasks_guard| *tasks_guard = None)
                            .map_err(AbcError::from))
        })
    }

    /// Runs for a fixed number of rounds, then return the best solution found.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok` with a `Candidate`.
    pub fn run_for_rounds(&self, rounds: usize) -> AbcResult<Candidate<Ctx::Solution>> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers).max_rounds(rounds);
        try!(self.run(tasks));
        self.get().map(|guard| guard.clone())
    }

    /// Run indefinitely.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok(())`.
    pub fn run_forever(&self) -> AbcResult<()> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
        self.run(tasks)
    }

    /// Stops a running hive.
    ///
    /// If a worker thread has panicked, this returns `Err(abc::Error)`.
    pub fn stop(&self) -> AbcResult<()> {
        let mut tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_mut().map_or((), |t| t.stop()))
    }

    /// Each new best candidate will be sent to `sender`.
    ///
    /// This is kept in a separate function so that the hive can be borrowed
    /// while running.
    pub fn set_sender(&mut self, sender: Sender<Candidate<Ctx::Solution>>) {
        if let Ok(best_guard) = self.best.lock() {
            sender.send(best_guard.clone()).unwrap_or(());
        }
        self.sender = Some(Mutex::new(sender));
    }

    /// Returns the current round of a running hive.
    ///
    /// If a worker thread has panicked and poisoned the task generator lock,
    /// `get_round` will return `Err(abc::Error)`. If the hive has not been
    /// run, `get_round` will return `Ok(None)`.
    ///
    /// If the hive is running, this will return `Ok(Some(n))`. `n` will start
    /// at 0, and increment each time every task in the round has been claimed
    /// (though not necessarily completed) by a worker thread.
    pub fn get_round(&self) -> AbcResult<Option<usize>> {
        let tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_ref().map(|tasks| tasks.round))
    }

    /// Get a reference to the hive's context.
    pub fn context(&self) -> &Ctx {
        &self.hive.context
    }
}

impl<Ctx: Context + 'static> Hive<Ctx> {
    /// Runs indefinitely in the background, providing a stream of results.
    ///
    /// This method consumes the hive, which will run until the `HiveBuilder`
    /// object is dropped. It returns an `mpsc::Receiver`, which receives a
    /// `Candidate` each time the hive improves on its best solution.
    pub fn stream(mut self) -> Receiver<Candidate<Ctx::Solution>> {
        let (sender, receiver) = channel();
        spawn(move || {
            self.set_sender(sender);
            let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
            self.run(tasks)
        });
        receiver
    }
}

impl<Ctx: Context> Debug for Hive<Ctx>
    where Ctx::Solution: Debug
{
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        for mutex in (&self.working).iter() {
            let working = mutex.read().unwrap();
            try!(write!(f, "..{:?}..\n", working.candidate));
        }
        let best_candidate = self.get().unwrap();
        write!(f, ">>{:?}<<", *best_candidate)
    }
}

impl<Ctx: Context> Drop for Hive<Ctx> {
    fn drop(&mut self) {
        self.stop().unwrap_or(())
    }
}
<|endoftext|>


<commit_before>use crate::rparser::*;
use crate::{gen_get_variants, Variant};
use httparse::{Request, Response, EMPTY_HEADER};
use std::cmp::Ordering;

pub struct HTTPBuilder {}
impl RBuilder for HTTPBuilder {
    fn build(&self) -> Box<dyn RParser> {
        Box::new(HTTPParser::new(b"HTTP"))
    }
    fn get_l4_probe(&self) -> Option<ProbeL4> {
        Some(http_probe)
    }
}

#[derive(Default)]
pub struct HTTPParser<'a> {
    _name: Option<&'a [u8]>,

    // request
    pub version: Option<String>,
    pub method: Option<String>,
    pub uri: Option<String>,
    pub user_agent: Option<String>,
    pub cookie: Option<String>,
    // response
    pub code: Option<u16>,
    pub content_length: Option<usize>,
    pub body: Vec<u8>,
}

impl<'a> HTTPParser<'a> {
    pub fn new(name: &'a [u8]) -> Self {
        HTTPParser {
            _name: Some(name),
            ..HTTPParser::default()
        }
    }
}

impl<'a> RParser for HTTPParser<'a> {
    fn parse(&mut self, i: &[u8], direction: u8) -> u32 {
        // apache usually sets a limit of 100 headers max
        const NUM_OF_HEADERS: usize = 20;
        let mut headers = [EMPTY_HEADER; NUM_OF_HEADERS];
        if direction == STREAM_TOSERVER {
            let mut req = Request::new(&mut headers[..]);
            let status = req.parse(i);
            trace!("status {:?}", status);
            trace!("request: {:?}", req);
            if let Some(version) = req.version {
                self.version = Some(format!("HTTP/1.{}", version));
            }
            if let Some(method) = req.method {
                self.method = Some(method.to_owned());
            }
            if let Some(uri) = req.path {
                self.uri = Some(uri.to_owned());
            }
            for hdr in &headers {
                let name = hdr.name.to_lowercase();
                match name.as_ref() {
                    "user-agent" => {
                        let s = String::from_utf8_lossy(hdr.value).into_owned();
                        self.user_agent = Some(s);
                    }
                    "cookie" => {
                        let s = String::from_utf8_lossy(hdr.value).into_owned();
                        self.cookie = Some(s);
                    },
                    "host" => {
                        let s = String::from_utf8_lossy(hdr.value);
                        debug!("host: {}", s);
                    },
                    _ => (),
                }
            }
        } else {
            // check if continuation of a response
            if let Some(sz) = self.content_length {
                trace!("queueing data for response body");
                self.body.extend_from_slice(i);
                // check if finished
                match self.body.len().cmp(&sz) {
                    Ordering::Less =>
                        /* not yet complete */
                        {}
                    Ordering::Equal => {
                        debug!("finished");
                    }
                    Ordering::Greater => {
                        warn!(
                            "Response body {} larger than content-length {}",
                            self.body.len(),
                            sz
                        );
                    }
                }
                return R_STATUS_OK;
            }
            let mut resp = Response::new(&mut headers[..]);
            let status = resp.parse(i);
            trace!("status: {:?}", status);
            trace!("response headers: {:?}", resp);
            if let Some(code) = resp.code {
                self.code = Some(code);
            }
            if let Ok(httparse::Status::Complete(sz)) = status {
                self.body = i[sz..].to_vec();
            }
            // check for chunked encoding (Transfer-Encoding: Chunked)
            //     if yes, see https://en.wikipedia.org/wiki/Chunked_transfer_encoding
            // check if we have a Content-Length
            for hdr in &headers {
                let name = hdr.name.to_lowercase();
                match name.as_ref() {
                    "content-length" => {
                        if let Ok(s) = std::str::from_utf8(hdr.value) {
                            if let Ok(length) = str::parse::<usize>(s) {
                                debug!("Content-Length: {}", length);
                                self.content_length = Some(length);
                                if length > self.body.len() {
                                    debug!("expecting more bytes");
                                }
                                continue;
                            }
                        }
                        warn!("Invalid encoding of Content-Length header");
                    }
                    _ => (),
                }
            }
        }
        R_STATUS_OK
    }

    gen_get_variants! {HTTPParser, "http.",
        // version     => |s| Some(Variant::from(&s.version)),
        version        => map_as_ref,
        method         => map_as_ref,
        uri            => map_as_ref,
        user_agent     => map_as_ref,
        cookie         => map_as_ref,
        code           => map,
        content_length => map,
    }
}

pub fn http_probe(i: &[u8], _l4info: &L4Info) -> ProbeResult {
    // number of headers to parse at once
    const NUM_OF_HEADERS: usize = 4;

    if i.len() < 6 {
        return ProbeResult::Unsure;
    }
    // check if first characters match start of "request-line"
    match &i[..4] {
        b"OPTI" | b"GET " | b"HEAD" | b"POST" | b"PUT " | b"PATC" | b"COPY" | b"MOVE" | b"DELE"
        | b"LINK" | b"UNLI" | b"TRAC" | b"WRAP" => (),
        _ => return ProbeResult::NotForUs,
    }
    // try parsing request
    let mut headers = [EMPTY_HEADER; NUM_OF_HEADERS];
    let mut req = Request::new(&mut headers[..]);
    let status = req.parse(i);
    if let Err(e) = status {
        if e != httparse::Error::TooManyHeaders {
            trace!(
                "data could be HTTP, but got error {:?} while parsing",
                status
            );
            return ProbeResult::Unsure;
        }
    }
    // println!("probe {:?}", status);
    // println!("probe {:?}", req);
    ProbeResult::Certain
}

#[cfg(test)]
mod test {
    use super::http_probe;
    use crate::probe::{L4Info, ProbeResult};

    #[test]
    fn http_parse() {
        let l4_info = L4Info {
            src_port: 1234,
            dst_port: 80,
            l4_proto: 6,
        };
        println!("http request");
        let data = b"GET /foo/:foo.com HTTP/1.1\r\nHost: example.org\r\n\r\n";
        assert_eq!(http_probe(data, &l4_info), ProbeResult::Certain);
        println!("binary");
        let data = b"\x00\x30\x00\x00\x00\x00";
        assert_eq!(http_probe(data, &l4_info), ProbeResult::NotForUs);
    }
}
<commit_msg>HTTP: add content-type<commit_after>use crate::rparser::*;
use crate::{gen_get_variants, Variant};
use httparse::{Request, Response, EMPTY_HEADER};
use std::cmp::Ordering;

pub struct HTTPBuilder {}
impl RBuilder for HTTPBuilder {
    fn build(&self) -> Box<dyn RParser> {
        Box::new(HTTPParser::new(b"HTTP"))
    }
    fn get_l4_probe(&self) -> Option<ProbeL4> {
        Some(http_probe)
    }
}

#[derive(Default)]
pub struct HTTPParser<'a> {
    _name: Option<&'a [u8]>,

    // request
    pub version: Option<String>,
    pub method: Option<String>,
    pub uri: Option<String>,
    pub user_agent: Option<String>,
    pub cookie: Option<String>,
    // response
    pub code: Option<u16>,
    pub content_length: Option<usize>,
    pub content_type: Option<String>,
    pub body: Vec<u8>,
}

impl<'a> HTTPParser<'a> {
    pub fn new(name: &'a [u8]) -> Self {
        HTTPParser {
            _name: Some(name),
            ..HTTPParser::default()
        }
    }
}

impl<'a> RParser for HTTPParser<'a> {
    #[allow(clippy::cognitive_complexity)]
    fn parse(&mut self, i: &[u8], direction: u8) -> u32 {
        // apache usually sets a limit of 100 headers max
        const NUM_OF_HEADERS: usize = 20;
        let mut headers = [EMPTY_HEADER; NUM_OF_HEADERS];
        if direction == STREAM_TOSERVER {
            let mut req = Request::new(&mut headers[..]);
            let status = req.parse(i);
            trace!("status {:?}", status);
            trace!("request: {:?}", req);
            if let Some(version) = req.version {
                self.version = Some(format!("HTTP/1.{}", version));
            }
            if let Some(method) = req.method {
                self.method = Some(method.to_owned());
            }
            if let Some(uri) = req.path {
                self.uri = Some(uri.to_owned());
            }
            for hdr in &headers {
                let name = hdr.name.to_lowercase();
                match name.as_ref() {
                    "user-agent" => {
                        let s = String::from_utf8_lossy(hdr.value).into_owned();
                        self.user_agent = Some(s);
                    }
                    "cookie" => {
                        let s = String::from_utf8_lossy(hdr.value).into_owned();
                        self.cookie = Some(s);
                    },
                    "host" => {
                        let s = String::from_utf8_lossy(hdr.value);
                        debug!("host: {}", s);
                    },
                    _ => (),
                }
            }
        } else {
            // check if continuation of a response
            if let Some(sz) = self.content_length {
                trace!("queueing data for response body");
                self.body.extend_from_slice(i);
                // check if finished
                match self.body.len().cmp(&sz) {
                    Ordering::Less =>
                        /* not yet complete */
                        {}
                    Ordering::Equal => {
                        debug!("finished");
                    }
                    Ordering::Greater => {
                        warn!(
                            "Response body {} larger than content-length {}",
                            self.body.len(),
                            sz
                        );
                    }
                }
                return R_STATUS_OK;
            }
            let mut resp = Response::new(&mut headers[..]);
            let status = resp.parse(i);
            trace!("status: {:?}", status);
            trace!("response headers: {:?}", resp);
            if let Some(code) = resp.code {
                self.code = Some(code);
            }
            if let Ok(httparse::Status::Complete(sz)) = status {
                self.body = i[sz..].to_vec();
            }
            // check for chunked encoding (Transfer-Encoding: Chunked)
            //     if yes, see https://en.wikipedia.org/wiki/Chunked_transfer_encoding
            // check if we have a Content-Length
            for hdr in &headers {
                let name = hdr.name.to_lowercase();
                match name.as_ref() {
                    "content-length" => {
                        if let Ok(s) = std::str::from_utf8(hdr.value) {
                            if let Ok(length) = str::parse::<usize>(s) {
                                debug!("Content-Length: {}", length);
                                self.content_length = Some(length);
                                if length > self.body.len() {
                                    debug!("expecting more bytes");
                                }
                                continue;
                            }
                        }
                        warn!("Invalid encoding of Content-Length header");
                    }
                    "content-type" => {
                        let s = String::from_utf8_lossy(hdr.value).into_owned();
                        self.content_type = Some(s);
                    }
                    _ => (),
                }
            }
        }
        R_STATUS_OK
    }

    gen_get_variants! {HTTPParser, "http.",
        // version     => |s| Some(Variant::from(&s.version)),
        version        => map_as_ref,
        method         => map_as_ref,
        uri            => map_as_ref,
        user_agent     => map_as_ref,
        cookie         => map_as_ref,
        code           => map,
        content_length => map,
        content_type   => map_as_ref,
    }
}

pub fn http_probe(i: &[u8], _l4info: &L4Info) -> ProbeResult {
    // number of headers to parse at once
    const NUM_OF_HEADERS: usize = 4;

    if i.len() < 6 {
        return ProbeResult::Unsure;
    }
    // check if first characters match start of "request-line"
    match &i[..4] {
        b"OPTI" | b"GET " | b"HEAD" | b"POST" | b"PUT " | b"PATC" | b"COPY" | b"MOVE" | b"DELE"
        | b"LINK" | b"UNLI" | b"TRAC" | b"WRAP" => (),
        _ => return ProbeResult::NotForUs,
    }
    // try parsing request
    let mut headers = [EMPTY_HEADER; NUM_OF_HEADERS];
    let mut req = Request::new(&mut headers[..]);
    let status = req.parse(i);
    if let Err(e) = status {
        if e != httparse::Error::TooManyHeaders {
            trace!(
                "data could be HTTP, but got error {:?} while parsing",
                status
            );
            return ProbeResult::Unsure;
        }
    }
    // println!("probe {:?}", status);
    // println!("probe {:?}", req);
    ProbeResult::Certain
}

#[cfg(test)]
mod test {
    use super::http_probe;
    use crate::probe::{L4Info, ProbeResult};

    #[test]
    fn http_parse() {
        let l4_info = L4Info {
            src_port: 1234,
            dst_port: 80,
            l4_proto: 6,
        };
        println!("http request");
        let data = b"GET /foo/:foo.com HTTP/1.1\r\nHost: example.org\r\n\r\n";
        assert_eq!(http_probe(data, &l4_info), ProbeResult::Certain);
        println!("binary");
        let data = b"\x00\x30\x00\x00\x00\x00";
        assert_eq!(http_probe(data, &l4_info), ProbeResult::NotForUs);
    }
}
<|endoftext|>


<commit_before>//! Bindings to IOCP, I/O Completion Ports

use std::io;
use std::mem;
use std::os::windows::io::*;
use std::time::Duration;

use handle::Handle;
use winapi::*;
use kernel32::*;
use Overlapped;

/// A handle to an Windows I/O Completion Port.
#[derive(Debug)]
pub struct CompletionPort {
    handle: Handle,
}

/// A status message received from an I/O completion port.
///
/// These statuses can be created via the `new` or `empty` constructors and then
/// provided to a completion port, or they are read out of a completion port.
/// The fields of each status are read through its accessor methods.
#[derive(Clone, Copy, Debug)]
pub struct CompletionStatus(OVERLAPPED_ENTRY);

unsafe impl Send for CompletionStatus {}
unsafe impl Sync for CompletionStatus {}

impl CompletionPort {
    /// Creates a new I/O completion port with the specified concurrency value.
    ///
    /// The number of threads given corresponds to the level of concurrency
    /// allowed for threads associated with this port. Consult the Windows
    /// documentation for more information about this value.
    pub fn new(threads: u32) -> io::Result<CompletionPort> {
        let ret = unsafe {
            CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0 as *mut _,
                                   0, threads)
        };
        if ret.is_null() {
            Err(io::Error::last_os_error())
        } else {
            Ok(CompletionPort { handle: Handle::new(ret) })
        }
    }

    /// Associates a new `HANDLE` to this I/O completion port.
    ///
    /// This function will associate the given handle to this port with the
    /// given `token` to be returned in status messages whenever it receives a
    /// notification.
    ///
    /// Any object which is convertible to a `HANDLE` via the `AsRawHandle`
    /// trait can be provided to this function, such as `std::fs::File` and
    /// friends.
    pub fn add_handle<T: AsRawHandle>(&self, token: usize,
                                      t: &T) -> io::Result<()> {
        self._add(token, t.as_raw_handle())
    }

    /// Associates a new `SOCKET` to this I/O completion port.
    ///
    /// This function will associate the given socket to this port with the
    /// given `token` to be returned in status messages whenever it receives a
    /// notification.
    ///
    /// Any object which is convertible to a `SOCKET` via the `AsRawSocket`
    /// trait can be provided to this function, such as `std::net::TcpStream`
    /// and friends.
    pub fn add_socket<T: AsRawSocket>(&self, token: usize,
                                      t: &T) -> io::Result<()> {
        self._add(token, t.as_raw_socket() as HANDLE)
    }

    fn _add(&self, token: usize, handle: HANDLE) -> io::Result<()> {
        let ret = unsafe {
            CreateIoCompletionPort(handle, self.handle.raw(),
                                   token as ULONG_PTR, 0)
        };
        if ret.is_null() {
            Err(io::Error::last_os_error())
        } else {
            debug_assert_eq!(ret, self.handle.raw());
            Ok(())
        }
    }

    /// Dequeue a completion status from this I/O completion port.
    ///
    /// This function will associate the calling thread with this completion
    /// port and then wait for a status message to become available. The precise
    /// semantics on when this function returns depends on the concurrency value
    /// specified when the port was created.
    ///
    /// A timeout can optionally be specified to this function. If `None` is
    /// provided this function will not time out, and otherwise it will time out
    /// after the specified duration has passed.
    ///
    /// On success this will return the status message which was dequeued from
    /// this completion port.
    pub fn get(&self, timeout: Option<Duration>) -> io::Result<CompletionStatus> {
        let mut bytes = 0;
        let mut token = 0;
        let mut overlapped = 0 as *mut _;
        let timeout = timeout.map(::dur2timeout).unwrap_or(INFINITE);
        let ret = unsafe {
            GetQueuedCompletionStatus(self.handle.raw(),
                                      &mut bytes,
                                      &mut token,
                                      &mut overlapped,
                                      timeout)
        };
        ::cvt(ret).map(|_| {
            CompletionStatus(OVERLAPPED_ENTRY {
                dwNumberOfBytesTransferred: bytes,
                lpCompletionKey: token,
                lpOverlapped: overlapped,
                Internal: 0,
            })
        })
    }

    /// Dequeues a number of completion statuses from this I/O completion port.
    ///
    /// This function is the same as `get` except that it may return more than
    /// one status. A buffer of "zero" statuses is provided (the contents are
    /// not read) and then on success this function will return a sub-slice of
    /// statuses which represent those which were dequeued from this port. This
    /// function does not wait to fill up the entire list of statuses provided.
    ///
    /// Like with `get`, a timeout may be specified for this operation.
    pub fn get_many<'a>(&self,
                        list: &'a mut [CompletionStatus],
                        timeout: Option<Duration>)
                        -> io::Result<&'a mut [CompletionStatus]>
    {
        debug_assert_eq!(mem::size_of::<CompletionStatus>(),
                         mem::size_of::<OVERLAPPED_ENTRY>());
        let mut removed = 0;
        let timeout = timeout.map(::dur2timeout).unwrap_or(INFINITE);
        let ret = unsafe {
            GetQueuedCompletionStatusEx(self.handle.raw(),
                                        list.as_ptr() as *mut _,
                                        list.len() as ULONG,
                                        &mut removed,
                                        timeout,
                                        FALSE)
        };
        match ::cvt(ret) {
            Ok(_) => Ok(&mut list[..removed as usize]),
            Err(e) => Err(e),
        }
    }

    /// Posts a new completion status onto this I/O completion port.
    ///
    /// This function will post the given status, with custom parameters, to the
    /// port. Threads blocked in `get` or `get_many` will eventually receive
    /// this status.
    pub fn post(&self, status: CompletionStatus) -> io::Result<()> {
        let ret = unsafe {
            PostQueuedCompletionStatus(self.handle.raw(),
                                       status.0.dwNumberOfBytesTransferred,
                                       status.0.lpCompletionKey,
                                       status.0.lpOverlapped)
        };
        ::cvt(ret).map(|_| ())
    }
}

impl AsRawHandle for CompletionPort {
    fn as_raw_handle(&self) -> HANDLE {
        self.handle.raw()
    }
}

impl FromRawHandle for CompletionPort {
    unsafe fn from_raw_handle(handle: HANDLE) -> CompletionPort {
        CompletionPort { handle: Handle::new(handle) }
    }
}

#[cfg(feature = "unstable")]
impl IntoRawHandle for CompletionPort {
    fn into_raw_handle(self) -> HANDLE {
        self.handle.into_raw()
    }
}

impl CompletionStatus {
    /// Creates a new completion status with the provided parameters.
    ///
    /// This function is useful when creating a status to send to a port with
    /// the `post` method. The parameters are opaquely passed through and not
    /// interpreted by the system at all.
    pub fn new(bytes: u32, token: usize, overlapped: *mut Overlapped)
               -> CompletionStatus {
        CompletionStatus(OVERLAPPED_ENTRY {
            dwNumberOfBytesTransferred: bytes,
            lpCompletionKey: token as ULONG_PTR,
            lpOverlapped: overlapped as *mut _,
            Internal: 0,
        })
    }

    /// Creates a new "zero" completion status.
    ///
    /// This function is useful when creating a stack buffer or vector of
    /// completion statuses to be passed to the `get_many` function.
    pub fn zero() -> CompletionStatus {
        CompletionStatus::new(0, 0, 0 as *mut _)
    }

    /// Returns the number of bytes that were transferred for the I/O operation
    /// associated with this completion status.
    pub fn bytes_transferred(&self) -> u32 {
        self.0.dwNumberOfBytesTransferred
    }

    /// Returns the completion key value associated with the file handle whose
    /// I/O operation has completed.
    ///
    /// A completion key is a per-handle key that is specified when it is added
    /// to an I/O completion port via `add_handle` or `add_socket`.
    pub fn token(&self) -> usize {
        self.0.lpCompletionKey as usize
    }

    /// Returns a pointer to the `Overlapped` structure that was specified when
    /// the I/O operation was started.
    pub fn overlapped(&self) -> *mut Overlapped {
        self.0.lpOverlapped as *mut _
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use std::time::Duration;
    use winapi::*;

    use iocp::{CompletionPort, CompletionStatus};

    #[test]
    fn is_send_sync() {
        fn is_send_sync<T: Send + Sync>() {}
        is_send_sync::<CompletionPort>();
    }

    #[test]
    fn token_right_size() {
        assert_eq!(mem::size_of::<usize>(), mem::size_of::<ULONG_PTR>());
    }

    #[test]
    fn timeout() {
        let c = CompletionPort::new(1).unwrap();
        let err = c.get(Some(Duration::new(0, 1_000_000))).unwrap_err();
        assert_eq!(err.raw_os_error(), Some(WAIT_TIMEOUT as i32));
    }

    #[test]
    fn get() {
        let c = CompletionPort::new(1).unwrap();
        c.post(CompletionStatus::new(1, 2, 3 as *mut _)).unwrap();
        let s = c.get(None).unwrap();
        assert_eq!(s.bytes_transferred(), 1);
        assert_eq!(s.token(), 2);
        assert_eq!(s.overlapped(), 3 as *mut _);
    }

    #[test]
    fn get_many() {
        let c = CompletionPort::new(1).unwrap();

        c.post(CompletionStatus::new(1, 2, 3 as *mut _)).unwrap();
        c.post(CompletionStatus::new(4, 5, 6 as *mut _)).unwrap();

        let mut s = vec![CompletionStatus::zero(); 4];
        {
            let s = c.get_many(&mut s, None).unwrap();
            assert_eq!(s.len(), 2);
            assert_eq!(s[0].bytes_transferred(), 1);
            assert_eq!(s[0].token(), 2);
            assert_eq!(s[0].overlapped(), 3 as *mut _);
            assert_eq!(s[1].bytes_transferred(), 4);
            assert_eq!(s[1].token(), 5);
            assert_eq!(s[1].overlapped(), 6 as *mut _);
        }
        assert_eq!(s[2].bytes_transferred(), 0);
        assert_eq!(s[2].token(), 0);
        assert_eq!(s[2].overlapped(), 0 as *mut _);
    }
}
<commit_msg>Add ?Sized where possible<commit_after>//! Bindings to IOCP, I/O Completion Ports

use std::io;
use std::mem;
use std::os::windows::io::*;
use std::time::Duration;

use handle::Handle;
use winapi::*;
use kernel32::*;
use Overlapped;

/// A handle to an Windows I/O Completion Port.
#[derive(Debug)]
pub struct CompletionPort {
    handle: Handle,
}

/// A status message received from an I/O completion port.
///
/// These statuses can be created via the `new` or `empty` constructors and then
/// provided to a completion port, or they are read out of a completion port.
/// The fields of each status are read through its accessor methods.
#[derive(Clone, Copy, Debug)]
pub struct CompletionStatus(OVERLAPPED_ENTRY);

unsafe impl Send for CompletionStatus {}
unsafe impl Sync for CompletionStatus {}

impl CompletionPort {
    /// Creates a new I/O completion port with the specified concurrency value.
    ///
    /// The number of threads given corresponds to the level of concurrency
    /// allowed for threads associated with this port. Consult the Windows
    /// documentation for more information about this value.
    pub fn new(threads: u32) -> io::Result<CompletionPort> {
        let ret = unsafe {
            CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0 as *mut _,
                                   0, threads)
        };
        if ret.is_null() {
            Err(io::Error::last_os_error())
        } else {
            Ok(CompletionPort { handle: Handle::new(ret) })
        }
    }

    /// Associates a new `HANDLE` to this I/O completion port.
    ///
    /// This function will associate the given handle to this port with the
    /// given `token` to be returned in status messages whenever it receives a
    /// notification.
    ///
    /// Any object which is convertible to a `HANDLE` via the `AsRawHandle`
    /// trait can be provided to this function, such as `std::fs::File` and
    /// friends.
    pub fn add_handle<T: AsRawHandle + ?Sized>(&self, token: usize,
                                               t: &T) -> io::Result<()> {
        self._add(token, t.as_raw_handle())
    }

    /// Associates a new `SOCKET` to this I/O completion port.
    ///
    /// This function will associate the given socket to this port with the
    /// given `token` to be returned in status messages whenever it receives a
    /// notification.
    ///
    /// Any object which is convertible to a `SOCKET` via the `AsRawSocket`
    /// trait can be provided to this function, such as `std::net::TcpStream`
    /// and friends.
    pub fn add_socket<T: AsRawSocket + ?Sized>(&self, token: usize,
                                               t: &T) -> io::Result<()> {
        self._add(token, t.as_raw_socket() as HANDLE)
    }

    fn _add(&self, token: usize, handle: HANDLE) -> io::Result<()> {
        let ret = unsafe {
            CreateIoCompletionPort(handle, self.handle.raw(),
                                   token as ULONG_PTR, 0)
        };
        if ret.is_null() {
            Err(io::Error::last_os_error())
        } else {
            debug_assert_eq!(ret, self.handle.raw());
            Ok(())
        }
    }

    /// Dequeue a completion status from this I/O completion port.
    ///
    /// This function will associate the calling thread with this completion
    /// port and then wait for a status message to become available. The precise
    /// semantics on when this function returns depends on the concurrency value
    /// specified when the port was created.
    ///
    /// A timeout can optionally be specified to this function. If `None` is
    /// provided this function will not time out, and otherwise it will time out
    /// after the specified duration has passed.
    ///
    /// On success this will return the status message which was dequeued from
    /// this completion port.
    pub fn get(&self, timeout: Option<Duration>) -> io::Result<CompletionStatus> {
        let mut bytes = 0;
        let mut token = 0;
        let mut overlapped = 0 as *mut _;
        let timeout = timeout.map(::dur2timeout).unwrap_or(INFINITE);
        let ret = unsafe {
            GetQueuedCompletionStatus(self.handle.raw(),
                                      &mut bytes,
                                      &mut token,
                                      &mut overlapped,
                                      timeout)
        };
        ::cvt(ret).map(|_| {
            CompletionStatus(OVERLAPPED_ENTRY {
                dwNumberOfBytesTransferred: bytes,
                lpCompletionKey: token,
                lpOverlapped: overlapped,
                Internal: 0,
            })
        })
    }

    /// Dequeues a number of completion statuses from this I/O completion port.
    ///
    /// This function is the same as `get` except that it may return more than
    /// one status. A buffer of "zero" statuses is provided (the contents are
    /// not read) and then on success this function will return a sub-slice of
    /// statuses which represent those which were dequeued from this port. This
    /// function does not wait to fill up the entire list of statuses provided.
    ///
    /// Like with `get`, a timeout may be specified for this operation.
    pub fn get_many<'a>(&self,
                        list: &'a mut [CompletionStatus],
                        timeout: Option<Duration>)
                        -> io::Result<&'a mut [CompletionStatus]>
    {
        debug_assert_eq!(mem::size_of::<CompletionStatus>(),
                         mem::size_of::<OVERLAPPED_ENTRY>());
        let mut removed = 0;
        let timeout = timeout.map(::dur2timeout).unwrap_or(INFINITE);
        let ret = unsafe {
            GetQueuedCompletionStatusEx(self.handle.raw(),
                                        list.as_ptr() as *mut _,
                                        list.len() as ULONG,
                                        &mut removed,
                                        timeout,
                                        FALSE)
        };
        match ::cvt(ret) {
            Ok(_) => Ok(&mut list[..removed as usize]),
            Err(e) => Err(e),
        }
    }

    /// Posts a new completion status onto this I/O completion port.
    ///
    /// This function will post the given status, with custom parameters, to the
    /// port. Threads blocked in `get` or `get_many` will eventually receive
    /// this status.
    pub fn post(&self, status: CompletionStatus) -> io::Result<()> {
        let ret = unsafe {
            PostQueuedCompletionStatus(self.handle.raw(),
                                       status.0.dwNumberOfBytesTransferred,
                                       status.0.lpCompletionKey,
                                       status.0.lpOverlapped)
        };
        ::cvt(ret).map(|_| ())
    }
}

impl AsRawHandle for CompletionPort {
    fn as_raw_handle(&self) -> HANDLE {
        self.handle.raw()
    }
}

impl FromRawHandle for CompletionPort {
    unsafe fn from_raw_handle(handle: HANDLE) -> CompletionPort {
        CompletionPort { handle: Handle::new(handle) }
    }
}

#[cfg(feature = "unstable")]
impl IntoRawHandle for CompletionPort {
    fn into_raw_handle(self) -> HANDLE {
        self.handle.into_raw()
    }
}

impl CompletionStatus {
    /// Creates a new completion status with the provided parameters.
    ///
    /// This function is useful when creating a status to send to a port with
    /// the `post` method. The parameters are opaquely passed through and not
    /// interpreted by the system at all.
    pub fn new(bytes: u32, token: usize, overlapped: *mut Overlapped)
               -> CompletionStatus {
        CompletionStatus(OVERLAPPED_ENTRY {
            dwNumberOfBytesTransferred: bytes,
            lpCompletionKey: token as ULONG_PTR,
            lpOverlapped: overlapped as *mut _,
            Internal: 0,
        })
    }

    /// Creates a new "zero" completion status.
    ///
    /// This function is useful when creating a stack buffer or vector of
    /// completion statuses to be passed to the `get_many` function.
    pub fn zero() -> CompletionStatus {
        CompletionStatus::new(0, 0, 0 as *mut _)
    }

    /// Returns the number of bytes that were transferred for the I/O operation
    /// associated with this completion status.
    pub fn bytes_transferred(&self) -> u32 {
        self.0.dwNumberOfBytesTransferred
    }

    /// Returns the completion key value associated with the file handle whose
    /// I/O operation has completed.
    ///
    /// A completion key is a per-handle key that is specified when it is added
    /// to an I/O completion port via `add_handle` or `add_socket`.
    pub fn token(&self) -> usize {
        self.0.lpCompletionKey as usize
    }

    /// Returns a pointer to the `Overlapped` structure that was specified when
    /// the I/O operation was started.
    pub fn overlapped(&self) -> *mut Overlapped {
        self.0.lpOverlapped as *mut _
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use std::time::Duration;
    use winapi::*;

    use iocp::{CompletionPort, CompletionStatus};

    #[test]
    fn is_send_sync() {
        fn is_send_sync<T: Send + Sync>() {}
        is_send_sync::<CompletionPort>();
    }

    #[test]
    fn token_right_size() {
        assert_eq!(mem::size_of::<usize>(), mem::size_of::<ULONG_PTR>());
    }

    #[test]
    fn timeout() {
        let c = CompletionPort::new(1).unwrap();
        let err = c.get(Some(Duration::new(0, 1_000_000))).unwrap_err();
        assert_eq!(err.raw_os_error(), Some(WAIT_TIMEOUT as i32));
    }

    #[test]
    fn get() {
        let c = CompletionPort::new(1).unwrap();
        c.post(CompletionStatus::new(1, 2, 3 as *mut _)).unwrap();
        let s = c.get(None).unwrap();
        assert_eq!(s.bytes_transferred(), 1);
        assert_eq!(s.token(), 2);
        assert_eq!(s.overlapped(), 3 as *mut _);
    }

    #[test]
    fn get_many() {
        let c = CompletionPort::new(1).unwrap();

        c.post(CompletionStatus::new(1, 2, 3 as *mut _)).unwrap();
        c.post(CompletionStatus::new(4, 5, 6 as *mut _)).unwrap();

        let mut s = vec![CompletionStatus::zero(); 4];
        {
            let s = c.get_many(&mut s, None).unwrap();
            assert_eq!(s.len(), 2);
            assert_eq!(s[0].bytes_transferred(), 1);
            assert_eq!(s[0].token(), 2);
            assert_eq!(s[0].overlapped(), 3 as *mut _);
            assert_eq!(s[1].bytes_transferred(), 4);
            assert_eq!(s[1].token(), 5);
            assert_eq!(s[1].overlapped(), 6 as *mut _);
        }
        assert_eq!(s[2].bytes_transferred(), 0);
        assert_eq!(s[2].token(), 0);
        assert_eq!(s[2].overlapped(), 0 as *mut _);
    }
}
<|endoftext|>