Ticket #576: ThreadPool.d

/++
    author: Jeff Davey <jeffd@gwava.com>
    changed: Darryl Bleau <darrylb@gwava.com> (altered to a variadic template class)
    license: BSD style

    A basic thread queueing module.

    This specific module pre-detaches all threads up to the size of the thread pool. The idea is that state
    can then be saved on each thread's TLS for caching and other purposes.

    All threads are set to "detach" mode, so that you can close the running program and let the OS clean up
    any thread resources.

    There is also a close function, which specifically waits for all detached threads to move to a closed state
    before returning.

    Thread pools are templated variadically to the type and number of arguments that each threaded function will take.

    Usage is fairly straightforward, and both delegates and functions are supported:

    Example:
    ---
    uint a = 0;
    Mutex m = new Mutex();

    void func(void *arg)
    {            
        char[] name = Thread.getThis().name; // get the name of this thread
        int *data = cast(int *)arg; // do some fun with casting
        m.lock(); // lock something
        a++; // increment something
        m.unlock(); // unlock something
        Thread.sleep(1); // zzz...
    }

    ThreadPool!(void*) pool = new ThreadPool!(void*)(100, &func);
    if (pool)
    {
        for (uint i = 0; i < 1000; i++)
        {
            int *data = new int();
            *data = i;
            pool.start(cast(void*)(data));
        }
    }
    pool.close();
    assert(a == 1000);
    ---

    ThreadPool functions don't have to take a void*, you can declare the ThreadPool with any type of arguments that you like:
    ---
    ThreadPool!(char[],float)
    ThreadPool!(Object,uint,char[])
    ---
++/

module tetra.util.ThreadPool;

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

private class PooledThread(T...)
{
    static enum State
    {
        IDLE,
        RUNNING,
        SHUTDOWN,
        CLOSED
    }

    this(void delegate(T) dg, size_t stackSize = 0)
    {
        this._dg = dg;
        this(stackSize);
    }

    this(void (*func)(T), size_t stackSize = 0)
    {
        this._func = func;
        this(stackSize);
    }

    void join()
    {
        _thread.join();
    }

    void lock()
    {
        _lock.lock();
    }

    bool tryLock()
    {
        return _lock.tryLock();
    }

    void unlock()
    {
        _lock.unlock();
    }

    void notify()
    {
        _signal.notify();
    }

    private:
    Mutex _lock;                         // for condition/mutex pair
    Condition _signal;                     // used to signal work available
    void delegate(T) _dg = null;        // delegate to call
    void (*_func)(T) = null;            // or function to call
    Thread _thread;                        // Underlying Thread class
    T arg;                                // Arg that has been passed
    State state = State.CLOSED;         // Our current state.
    
    void run()
    {
        _lock.lock();
        scope (exit)
        {    
            this.state = State.CLOSED;    
            _lock.unlock();
        }

        for(;;)
        {
            this.state = State.IDLE;
            _signal.wait();    // wait for signal from pool
            if (this.state == State.SHUTDOWN)    // shutting down?
                break;

            if (this.state == State.RUNNING)    // let's run the work.
            {            
                try
                {    
                    if (_dg)
                        _dg(arg);
                    else if (_func)
                        _func(arg);
                }
                catch (Exception ex) {} // if the underlying function call blows up, we still want the pool to be available.
                                        // In practice, the underlying delegate/function passed, should make sure to capture
                                        // it's exceptions.
            }
        }
    }

    this(size_t stackSize)
    {
        this._lock = new Mutex();
        this._signal = new Condition(this._lock);

        _thread = new Thread(&run, stackSize);
        _thread.isDaemon = true;
        _thread.start();
    }
}

/++
    ThreadPool is the thread queueing class.
++/
class ThreadPool(T...)
{
    /++
        Construct a ThreadPool of type T

        Params:
        poolSize             =     The size of the thread pool. It will 
                                actually create this many threads on startup.

        (function/delegate) =    void func(T argument). argument is the passed
                                argument from the start function.

        stackSize            =    Size of the stack per thread. As tango
                                doesn't quite support this functionality,
                                it's not working quite yet.
    ++/    
    this(uint poolSize, void (*func)(T), size_t stackSize = 0)
    in
    {
        assert(func);
    }
    body
    {
        this._func = func;
        this(poolSize, stackSize);
    }

    this(uint poolSize, void delegate(T) dg, size_t stackSize = 0)
    in
    {
        assert(dg);
    }
    body
    {
        this._dg = dg;
        this(poolSize, stackSize);        
    }

    /++
        Start a unit of work.

        Params:
        arg                    =    A templated (T) arg that will be passed to the
                                underlying worker thread.

        block                =    If true, the start function will wait
                                for an existing thread to become 
                                available, if all are being used. If
                                false, the start function will return 
                                immediately with a value of false if it
                                could not schedule work, or a vlue of true
                                if it could.
    
    ++/
    bool start(T arg, bool block = true)
    {
        bool rtn = false;
        for(;;)
        {
            uint i = 0;
            for (; i < _pool.length; i++)
            {
                if (_pool[i].tryLock())
                {
                    scope(exit) _pool[i].unlock();
                    if (_pool[i].state == PooledThread!(T).State.IDLE)
                    {
                        _pool[i].state = PooledThread!(T).State.RUNNING;
                        foreach(k,a; arg)
                            _pool[i].arg[k] = a;
                        _pool[i].notify();
                        rtn = true;
                        break;
                    }
                }
            }

            if (i == _pool.length)
            {
                if (block)
                    Thread.sleep(.1);
                else
                    break;
            }
            else
                break;
        }
        return rtn;
    }

    /++
      Close the thread pool. This will go through all running threads,
      tell them to shutdown, and wait for that to happen before returning.
    
      Currently, the poll time between thread checks is hard-coded to 100ms
    ++/
    void close()
    {
        if (open)
        {
            for(;;)
            {
                bool allClosed = true;
                for (uint i = 0; i < _pool.length; i++)
                {
                    _pool[i].lock();
                    if (_pool[i].state == PooledThread!(T).State.IDLE)
                    {
                        _pool[i].state = PooledThread!(T).State.SHUTDOWN;        
                        _pool[i].notify();
                    }    
                    if (_pool[i].state != PooledThread!(T).State.CLOSED)
                        allClosed = false;                                    
                    _pool[i].unlock();
                }
                if (allClosed)
                    break;
                else
                    Thread.sleep(.1);
            }
            open = false;
        }
    }

    private:
    void delegate(T) _dg = null;
    void (*_func)(T) = null;
    PooledThread!(T)[] _pool = null;
    bool open = false;

    this(uint poolSize, size_t stackSize)
    {
        _pool = new PooledThread!(T)[poolSize];

        for (uint i = 0; i < poolSize; i++)
        {
            if (_func)
                _pool[i] = new PooledThread!(T)(_func, stackSize);
            else
                _pool[i] = new PooledThread!(T)(_dg, stackSize);
        }
        open = true;
    }
}

debug (UnitTest)
{
    unittest
    {    
        int voidTemplate()
        {
            uint a = 0;
            Mutex m = new Mutex();
            void func(void *arg)
            {            
                char[] name = Thread.getThis().name;
                int *data = cast(int *)arg;
                m.lock();
                a++;
                m.unlock();
                Thread.sleep(.1);
                return 0;
            }

            ThreadPool!(void*) pool = new ThreadPool!(void*)(50, &func);
            if (pool)
            {
                for (uint i = 0; i < 1000; i++)
                {
                    int *data = new int();
                    *data = i;
                    pool.start(cast(void*)(data));
                }
            }
            pool.close();
            if (a == 1000)
                return 1;
            return 0;
        }

        int uintTemplate()
        {
            uint a = 0;
            Mutex m = new Mutex();
            void func(uint arg)
            {            
                m.lock();
                a += arg;
                m.unlock();
            }

            ThreadPool!(uint) pool = new ThreadPool!(uint)(50, &func);
            if (pool)
            {
                for (uint i = 0; i < 1000; i++)
                    pool.start(i);
            }
            pool.close();
            if (a == 499500)
                return 1;
            return 0;
        }

        int uintUintTemplate()
        {
            uint a = 0;
            uint b = 0;
            Mutex m = new Mutex();
            void func(uint first, uint second)
            {
                m.lock();
                a += first;
                b += second;
                m.unlock();
            }

            ThreadPool!(uint, uint) pool = new ThreadPool!(uint, uint)(50, &func);
            if (pool)
            {
                for (uint i = 0; i < 1000; i++)
                    pool.start(1, i);
            }
            pool.close();
            if ((a == 1000) && (b == 499500))
                return 1;
            return 0;
        }

        int objectTemplate()
        {
            class TestOb
            {
                int a;
            }
            Mutex m = new Mutex;
            TestOb a = new TestOb;
            
            void func(TestOb arg)
            {
                m.lock();
                arg.a += 1;
                m.unlock();
            }

            ThreadPool!(TestOb) pool = new ThreadPool!(TestOb)(50, &func);
            if (pool)
            {
                for (uint i = 0; i < 1000; i++)
                    pool.start(a);
            }
            pool.close();
            if (a.a == 1000)
                return 1;
            return 0;
        }

        assert(voidTemplate());
        assert(uintTemplate());
        assert(uintUintTemplate());
        assert(objectTemplate());
    }
}