source: trunk/yat/utility/Scheduler.cc @ 3407

// $Id: Scheduler.cc 3407 2015-04-07 04:04:32Z peter $

/*
  Copyright (C) 2014, 2015 Peter Johansson

  This file is part of the yat library, http://dev.thep.lu.se/yat

  The yat library is free software; you can redistribute it and/or
  modify it under the terms of the GNU General Public License as
  published by the Free Software Foundation; either version 3 of the
  License, or (at your option) any later version.

  The yat library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with yat. If not, see <http://www.gnu.org/licenses/>.
*/

#include <config.h>

#include "Scheduler.h"

#include <boost/exception.hpp>

#include <cassert>

namespace theplu {
namespace yat {
namespace utility {

  Scheduler::Scheduler(unsigned int threads)
    : running_jobs_(0), job_counter_(0)
  {
    assert(threads);
    for (size_t i=0; i<threads; ++i)
      workers_.create_thread(Worker(queue_, completed_));
  }


  void Scheduler::add_dependency(boost::shared_ptr<Job> job,
                                 boost::shared_ptr<Job> prerequisite)
  {
    assert(job->status_ == Job::pristine);
    job->prerequisite_.insert(prerequisite);
    prerequisite->observers_.push_back(job);
  }


  void Scheduler::interrupt(void)
  {
    workers_.interrupt_all();
  }


  void Scheduler::post_process(JobPtr job)
  {
    --running_jobs_;
    assert(running_jobs_>=0);
    job->status_ = Job::completed;

    // check if an error occurred
    if (job->error_) {
      interrupt();
      boost::rethrow_exception(job->error_);
    }

    // for convenience
    const std::vector<JobPtr>& vec = job->observers_;
    // notify observers, jobs that have 'job' as prerequisite
    for (size_t i=0; i<vec.size(); ++i) {
      vec[i]->prerequisite_.erase(job);
      if (vec[i]->prerequisite_.empty())
        queue(vec[i]);
    }
  }


  void Scheduler::process(JobPtr job)
  {
    if (job->status_!=Job::pristine)
      return;
    job->status_ = Job::prepared;
    job->id_ = job_counter_;
    ++job_counter_;

    // If we have prerequisite that need to be run first, process them

    typedef std::set<JobPtr>::iterator iterator;
    for (iterator j=job->prerequisite_.begin(); j!=job->prerequisite_.end();) {
      if ((*j)->status_ == Job::completed)
        job->prerequisite_.erase(j++);
      else {
        process(*j);
        ++j;
      }
    }

    // If all prerequisite are finished, send job to queue
    if (job->prerequisite_.empty()) {
      queue(job);
    }
  }


  void Scheduler::queue(JobPtr job)
  {
    job->status_ = Job::running;
    ++running_jobs_;
    assert(running_jobs_>0);
    queue_.push(job);
  }


  void Scheduler::submit(JobPtr job)
  {
    // If any job has been completed since last submission,
    // post-process the completed job
    JobPtr completed_job;
    while (completed_.try_pop(completed_job))
      post_process(completed_job);
    process(job);
  }


  void Scheduler::wait(void)
  {
    while (running_jobs_ || completed_.size()) {
      JobPtr job;
      completed_.pop(job);
      post_process(job);
    }

    // kill workers
    boost::shared_ptr<Job> end;
    queue_.push(end);

    // wait for workers to finish
    workers_.join_all();
  }


  // Scheduler::Job

  Scheduler::Job::Job(unsigned int prio)
    : status_(pristine), priority_(prio), id_(0) {}


  Scheduler::Job::~Job(void)
  {}


  unsigned int Scheduler::Job::priority(void) const
  {
    return priority_;
  }


  bool Scheduler::LowerPriority::operator()(const JobPtr& lhs,
                                            const JobPtr& rhs) const
  {
    if (rhs.get() == NULL)
      return false;
    if (lhs.get() == NULL)
      return true;

    if (lhs->priority() != rhs->priority())
      return lhs->priority() < rhs->priority();

    return lhs->id_ > rhs->id_;
  }


  // Scheduler::Worker

  Scheduler::Worker::Worker(JobQueue& q, Queue<JobPtr>& c)
    : queue_(q), completed_(c)
  {
  }


  void Scheduler::Worker::operator()(void)
  {
    while (true) {
      boost::shared_ptr<Job> job;
      // get next job
      queue_.pop(job);
      // NULL job indicates poison pill
      if (job.get()==NULL) {
        // kill other workers too
        queue_.push(job);
        break;
      }
      // action
      try {
        (*job)();
      }
      catch (...) {
        job->error_ = boost::current_exception();
        // return job to scheduler so it can act on the error
        completed_.push(job);
        // exit work and go home
        return;
      }
      // return job to scheduler
      completed_.push(job);

      // Make sure we can be interrupted
      boost::this_thread::interruption_point();
    }
  }

}}}