Ticket #576: pool.d

/***************************************
* Authors: Jeff Davey, jeffd@gwava.com
* Thread pooling library
*/

module pool;

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

private import tango.io.Stdout;


private class PooledThread
{
    enum State
    {
        IDLE,
        RUNNING,
        SHUTDOWN,
        CLOSED
    };

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

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

    State state()
    {
        return this._state;
    }

    State state(State newState)
    {
        return this._state = newState;
    }

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

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

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

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

    void* arg()
    {
        return this._arg;
    }

    void* arg(void* newArg)
    {
        return this._arg = newArg;
    }
    
    void notify()
    {
        _signal.notify();
    }

    private:
    Mutex _lock;
    Condition _signal;
    void delegate(void *) _dg = null;
    void (*_func)(void *) = null;
    void* _arg = null;
    Thread _thread;
    State _state = State.CLOSED;

    
    void run()
    {
        _lock.lock();
        scope (exit) _lock.unlock();

        for(;;)
        {
            this._state = State.IDLE;
            _signal.wait();
            if (this._state == State.SHUTDOWN)
            {
                this._state = State.CLOSED;
                break;
            }

            if (this._state == State.RUNNING)
            {
                if (_dg)
                    _dg(_arg);
                else if (_func)
                    _func(_arg);
            }
        }

        return 0;
    }

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

//      _thread = new Thread(&run, stackSize); // TODO: WE NEED STACKSIZE!
        _thread = new Thread(&run);
        _thread.start();
    }

}

class ThreadPool
{
    this(uint poolSize, void (*func)(void *), size_t stackSize = 0)
    in
    {
        assert(func);
    }
    body
    {
        this._func = func;
        this(poolSize, stackSize);
    }

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

    bool start(void *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.State.IDLE)
                    {
                        int *data = cast(int *)arg;
                        _pool[i].state = PooledThread.State.RUNNING;
                        _pool[i].arg = arg;
                        _pool[i].notify();
                        rtn = true;
                        break;
                    }
                }
            }

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

        return rtn;
    }


    void close()
    {
        if (open)
        {
            for (uint i = 0; i < _pool.length; i++)         
            {
                _pool[i].lock();
                if (_pool[i].state != PooledThread.State.RUNNING)
                {
                    _pool[i].state = PooledThread.State.SHUTDOWN;       
                    _pool[i].notify();
                }
                _pool[i].unlock();
            }
            for (uint i = 0; i < _pool.length; i++)
                _pool[i].join();
            open = false;
        }
    }


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

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

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

    unittest
    {   
        uint a = 0;
        Mutex m = new Mutex();

        void func(void *arg)
        {           
            char[] name = Thread.getThis().name;
            int *data = cast(int *)arg;
            m.lock();
            Stdout.format("{}", *data).newline;
            a++;
            m.unlock();
            Thread.sleep(1);
            return 0;
        }


        ThreadPool pool = new ThreadPool(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);
    }
}