Ticket #576: ThreadPool.2.d

/**
 * -- description --
 *
 * Copyright: Copyright (C) 2007 Anders Halager. All rights reserved.
 * License:   BSD style: $(LICENSE)
 * Author:    Anders Halager
 */
module tango.util.ThreadPool;

import tango.core.Thread,
       tango.core.Atomic,
       tango.core.sync.Mutex,
       tango.core.sync.Condition;

class ThreadPool(Args...)
{
    alias void delegate(Args) JobD;

    this(size_t workers, size_t q_size = 0)
    {
        // pre-allocate memory for q_size jobs in the queue
        q.length = q_size;
        q.length = 0;

        m = new Mutex;
        poolActivity = new Condition(m);
        workerActivity = new Condition(m);

        priority_job.store(cast(Job*)null);
        active_jobs.store(0u);
        done.store(false);

        for (size_t i = 0; i < workers; i++)
        {
            auto thread = new Thread(&doJob);
            // Allow the OS to kill the threads if we exit the program without
            // handling them our selves
            thread.isDaemon = true;
            thread.start();
            pool ~= thread;
        }
    }

    /**
      Assign the given job to a thread immediately or block until one is
      available
     */
    void assign(JobD job, Args args)
    {
        m.lock();
        auto j = Job(job, args);
        priority_job.store(&j);
        m.unlock();
        poolActivity.notify();
        // Wait until someone has taken the job
        while (priority_job.load() !is null)
            workerActivity.wait();
    }

    /**
      Assign the given job to a thread immediately or return false if none is
      available. (Returns true if one was available)
     */
    bool tryAssign(JobD job, Args args)
    {
        if (active_jobs.load() >= pool.length)
            return false;
        assign(job, args);
        return true;
    }

    /**
      Put a job into the pool for eventual execution.

      Warning: Acts as a stack, not a queue as you would expect
     */
    void append(JobD job, Args args)
    {
        m.lock();
        q ~= Job(job, args);
        m.unlock();
        poolActivity.notify();
    }

    /// Get the number of jobs waiting to be executed
    int pendingJobs()
    {
        m.lock(); scope(exit) m.unlock();
        return q.length;
    }

    /// Finish currently executing jobs and shutdown.
    void shutdown()
    {
        done.store(true);
        q.length = 0;
        poolActivity.notifyAll();
        foreach (thread; pool)
            thread.join();
        m.lock();
        m.unlock();
    }

    /// Complete all jobs and then shutdown.
    void finish()
    {
        m.lock();
        while (q.length > 0 || active_jobs.load() > 0)
            workerActivity.wait();
        m.unlock();
        shutdown();
    }

private:
    // Our list of threads -- only used during startup and shutdown
    Thread[] pool;
    struct Job
    { 
        JobD dg;
        Args args;
    }
    // Used for storing queued jobs that will be executed eventually
    Job[] q;

    // This is to store a single job for immediate execution, which hopefully
    // means that any program using only assign and tryAssign wont need any
    // heap allocations after startup.
    Atomic!(Job*) priority_job;

    // This should be used when accessing the job queue
    Mutex m;

    // Notify is called on this condition whenever we have activity in the pool
    // that the workers might want to know about.
    Condition poolActivity;

    // Worker threads call notify on this when they are done with a job or are
    // completely done.
    // This allows a graceful shut down and is necessary since assign has to
    // wait for a job to become available
    Condition workerActivity;

    // Are we in the shutdown phase?
    Atomic!(bool) done;

    // Counter for the number of jobs currently being calculated
    Atomic!(uint) active_jobs;

    // Thread delegate:
    void doJob()
    {
        while (!done.load())
        {
            m.lock();
            while (q.length == 0 && priority_job.load() is null && !done.load())
                poolActivity.wait();
            if (done.load()) {
                m.unlock(); // not using scope(exit), need to manually unlock
                break;
            }
            Job job;
            Job* jobPtr = priority_job.load();
            if (jobPtr !is null)
            {
                job = *jobPtr;
                priority_job.store(cast(Job*)null);
                workerActivity.notify();
            }
            else
            {
                // A stack -- should be a queue
                job = q[length - 1];
                q.length = q.length - 1;
            }

            // Make sure we unlock before we start doing the calculations
            m.unlock();

            // Do the actual job
            active_jobs.increment();
            try {
                job.dg(job.args);
            } catch (Exception ex) { }
            active_jobs.decrement();

            // Tell the pool that we are done with something
            workerActivity.notify();
            Thread.yield();
        }
        // Tell the pool that we are now done
        workerActivity.notify();
    }
}