ThreadWorkers.h

// WARNING: Changes to this file must be contributed back to Sawyer or else they will
//          be clobbered by the next update from Sawyer.  The Sawyer repository is at
//          https://github.com/matzke1/sawyer.




#include <Sawyer/Exception.h>
#include <Sawyer/Graph.h>
#include <Sawyer/Map.h>
#include <Sawyer/Sawyer.h>
#include <Sawyer/Stack.h>

#include <boost/foreach.hpp>
#include <boost/thread/condition_variable.hpp>
#include <boost/thread/locks.hpp>
#include <boost/thread/mutex.hpp>
#include <boost/thread/thread.hpp>
#include <boost/version.hpp>
#include <set>

namespace Sawyer {

template<class DependencyGraph, class Functor>
class ThreadWorkers {
    boost::mutex mutex_;                                // protects the following members after the constructor
    DependencyGraph dependencies_;                      // outstanding dependencies
    bool hasStarted_;                                   // set when work has started
    bool hasWaited_;                                    // set when wait() is called
    Container::Stack<size_t> workQueue_;                // outstanding work identified by vertex ID of dependency graph
    boost::condition_variable workInserted_;            // signaled when work is added to the queue or all work is consumed
    size_t nWorkers_;                                   // number of worker threads allocated
    boost::thread *workers_;                            // worker threads
    size_t nItemsStarted_;                              // number of work items started
    size_t nItemsFinished_;                             // number of work items that have been completed already
    size_t nWorkersRunning_;                            // number of workers that are currently busy doing something
    size_t nWorkersFinished_;                           // number of worker threads that have returned
    std::set<size_t> runningTasks_;                     // tasks (vertex IDs) that are running

public:
    ThreadWorkers()
        : hasStarted_(false), hasWaited_(false), nWorkers_(0), workers_(NULL), nItemsStarted_(0), nItemsFinished_(0),
          nWorkersRunning_(0), nWorkersFinished_(0) {}

    ThreadWorkers(const DependencyGraph &dependencies, size_t nWorkers, Functor functor)
        : hasStarted_(false), hasWaited_(false), nWorkers_(0), workers_(NULL), nItemsStarted_(0), nItemsFinished_(0),
          nWorkersRunning_(0), nWorkersFinished_(0) {
        try {
            run(dependencies, nWorkers, functor);
        } catch (const Exception::ContainsCycle&) {
            delete[] workers_;
            throw;                                      // destructor won't be called
        }
    }

    ~ThreadWorkers() {
        wait();
        delete[] workers_;
    }

    void start(const DependencyGraph &dependencies, size_t nWorkers, Functor functor) {
        boost::lock_guard<boost::mutex> lock(mutex_);
        if (hasStarted_)
            throw std::runtime_error("work can start only once per object");
        hasStarted_ = true;
        dependencies_ = dependencies;
        if (0 == nWorkers)
            nWorkers = boost::thread::hardware_concurrency();
        nWorkers_ = std::max((size_t)1, std::min(nWorkers, dependencies.nVertices()));
        nItemsStarted_ = nWorkersFinished_ = 0;
        runningTasks_.clear();
        fillWorkQueueNS();
        startWorkersNS(functor);
    }

    void wait() {
        boost::unique_lock<boost::mutex> lock(mutex_);
        if (!hasStarted_ || hasWaited_)
            return;
        hasWaited_ = true;
        lock.unlock();

        for (size_t i=0; i<nWorkers_; ++i)
            workers_[i].join();

        lock.lock();
        if (dependencies_.nEdges() != 0)
            throw Exception::ContainsCycle("task dependency graph contains cycle(s)");
        dependencies_.clear();
    }

    template<class Rep, class Period>
    bool tryWaitFor(const boost::chrono::duration<Rep, Period> &relTime) {
        const boost::chrono::steady_clock::time_point endAt = boost::chrono::steady_clock::now() + relTime;
        boost::unique_lock<boost::mutex> lock(mutex_);
        if (!hasStarted_ || hasWaited_)
            return true;
        lock.unlock();

        for (size_t i = 0; i < nWorkers_; ++i) {
            if (!workers_[i].try_join_until(endAt))
                return false;
        }

        lock.lock();
        hasWaited_ = true;
        if (dependencies_.nEdges() != 0)
            throw Exception::ContainsCycle("task dependency graph contains cycle(s)");
        dependencies_.clear();
        return true;
    }

    void run(const DependencyGraph &dependencies, size_t nWorkers, Functor functor) {
        start(dependencies, nWorkers, functor);
        wait();
    }

    bool isFinished() {
        boost::lock_guard<boost::mutex> lock(mutex_);
        return !hasStarted_ || nWorkersFinished_ == nWorkers_;
    }

    size_t nStarted() {
        boost::lock_guard<boost::mutex> lock(mutex_);
        return nItemsStarted_;
    }

    size_t nFinished() {
        boost::lock_guard<boost::mutex> lock(mutex_);
        return nItemsFinished_;
    }

    std::set<size_t> runningTasks() {
        boost::lock_guard<boost::mutex> lock(mutex_);
        return runningTasks_;
    }
    
    std::pair<size_t, size_t> nWorkers() {
        boost::lock_guard<boost::mutex> lock(mutex_);
        return std::make_pair(nWorkers_-nWorkersFinished_, nWorkersRunning_);
    }
    
private:
    // Scan the dependency graph and fill the work queue with vertices that have no dependencies.
    void fillWorkQueueNS() {
        ASSERT_require(workQueue_.isEmpty());
        BOOST_FOREACH (const typename DependencyGraph::Vertex &vertex, dependencies_.vertices()) {
            if (vertex.nOutEdges() == 0)
                workQueue_.push(vertex.id());
        }
    }

    // Start worker threads
    void startWorkersNS(Functor functor) {
        workers_ = new boost::thread[nWorkers_];
        for (size_t i=0; i<nWorkers_; ++i)
            workers_[i] = boost::thread(startWorker, this, functor);
    }

    // Worker threads execute here
    static void startWorker(ThreadWorkers *self, Functor functor) {
        self->worker(functor);
    }

    void worker(Functor functor) {
        while (1) {
            // Get the next item of work
            boost::unique_lock<boost::mutex> lock(mutex_);
            while (nItemsFinished_ < nItemsStarted_ && workQueue_.isEmpty())
                workInserted_.wait(lock);
            if (nItemsFinished_ == nItemsStarted_ && workQueue_.isEmpty()) {
                ++nWorkersFinished_;
                return;
            }
            ASSERT_forbid(workQueue_.isEmpty());
            size_t workItemId = workQueue_.pop();
            typename DependencyGraph::ConstVertexIterator workVertex = dependencies_.findVertex(workItemId);
            ASSERT_require(workVertex->nOutEdges() == 0);
            typename DependencyGraph::VertexValue workItem = workVertex->value();
            ++nItemsStarted_;

            // Do the work
            ++nWorkersRunning_;
            runningTasks_.insert(workItemId);
            lock.unlock();
            functor(workItemId, workItem);
            lock.lock();
            ++nItemsFinished_;
            --nWorkersRunning_;
            runningTasks_.erase(workItemId);

            // Look for more work as we remove some dependency edges. Watch out for parallel edges (self edges not possible).
            Container::Map<size_t, typename DependencyGraph::ConstVertexIterator> candidateWorkItems;
            BOOST_FOREACH (const typename DependencyGraph::Edge &edge, workVertex->inEdges())
                candidateWorkItems.insert(edge.source()->id(), edge.source());
            dependencies_.clearInEdges(workVertex);
            size_t newWorkInserted = 0;
            BOOST_FOREACH (const typename DependencyGraph::ConstVertexIterator &candidate, candidateWorkItems.values()) {
                if (candidate->nOutEdges() == 0) {
                    workQueue_.push(candidate->id());
                    ++newWorkInserted;
                }
            }

            // Notify other workers
            if (0 == newWorkInserted) {
                if (workQueue_.isEmpty())
                    workInserted_.notify_all();
            } else if (1 == newWorkInserted) {
                // we'll do the new work ourself
            } else if (2 == newWorkInserted) {
                workInserted_.notify_one();
            } else {
                workInserted_.notify_all();
            }
        }
    }
};

template<class DependencyGraph, class Functor>
void
workInParallel(const DependencyGraph &dependencies, size_t nWorkers, Functor functor) {
    ThreadWorkers<DependencyGraph, Functor>(dependencies, nWorkers, functor);
}


template<class DependencyGraph, class Functor, class Monitor>
void
workInParallel(const DependencyGraph &dependencies, size_t nWorkers, Functor functor,
               Monitor monitor, boost::chrono::milliseconds period) {
    ThreadWorkers<DependencyGraph, Functor> workers;
    workers.start(dependencies, nWorkers, functor);
    while (!workers.tryWaitFor(period))
        monitor(dependencies, workers.nFinished(), workers.runningTasks());
}
} // namespace