Ticket #539: tpool.py

import Queue
import time
from threading import Semaphore

class StopThread(Exception): pass

class ThreadPool:
    def __init__(self, max, mkthread_callable, spare=6, shrinkratio=2, shrink_every=10):
        '''Create a new thread pool.

        max should be the absolute maximum number of threads spawned
        mkthread_callable is some callable that takes a function 
            which should be called to retrieve each "job", and that 
            returns a threading.Thread-like object with a non-blocking
            start() method.  The thread should exit if get_job() returns
            StopThread.
        spare should be the number of "idle" connections to keep open
        shrinkratio should be the factor (current_idle - spare) gets 
            divided by to gradually shrink the thread pool
        shrink_every is the number of seconds between each shrink()
            run
        '''
        spare = min(spare, max)
        self.max = max
        self.mkthread_callable = mkthread_callable
        self.spare = spare
        self.shrinkratio = shrinkratio
        self.shrink_every = shrink_every

        self._queue = Queue.Queue()
        self._threadcount = 0
        self._idlecount = 0
        self._lastcheck = 0
        self._inclock = Semaphore()
    
    def _get_job(self):
        '''Returns a new job (or StopThread) from the shared Queue'''
        self._inclock.acquire()
        self._idlecount += 1
        self._inclock.release()

        r = self._queue.get()

        self._inclock.acquire()
        self._idlecount -= 1
        self._inclock.release()

        return r

    def _spawnthread(self):
        '''Creates a new thread'''
        th = self.mkthread_callable(self._get_job)
        th.start()
        self._threadcount += 1

    def _killthread(self):
        '''Tells one of the idle threads to die'''
        th = self.mkthread_callable(self._get_job)
        self._queue.put(StopThread)
        self._threadcount -= 1

    def tick(self):
        '''Determines whether or not enough time has elapsed
        to call shrink() again.
        '''
        if time.time() - self._lastcheck >= self.shrink_every:
            self.shrink()

    def shrink(self):
        '''Find out of a surge in demand has created more
        threads than we need (more idle threads than self.spare),
        and scale back the number if need be.
        '''
        self._lastcheck = time.time()
        surplus = self._idlecount - self.spare
        if surplus > 0:
            remove = max(1, surplus / self.shrinkratio)
            for x in xrange(remove):
                self._killthread()
        
    def run_job(self, job):
        '''Add some object to the job queue to be handled by 
        some thread in the pool.

        job should be a unit of work that the threads understand
        '''
        self.tick()
        if self._idlecount == 0 and self.max - self._threadcount:
            self._spawnthread()
        self._queue.put(job)

    def shutdown(self):
        '''Kill all threads'''
        def _dead_st(*args, **kw):
            pass
        self._spawnthread = _dead_st
        for x in xrange(self._threadcount):
            self._killthread()

if __name__ == '__main__':
    # Example
    from threading import Thread
    import random
    class Sleeper(Thread):
        def __init__(self, get_job):
            self.get_job = get_job
            Thread.__init__(self)

        def run(self):
            while True:
                job = self.get_job()
                if job == StopThread:
                    print 'thread done'
                    return
                time.sleep(job)
    pool = ThreadPool(40, Sleeper)

    for x in xrange(0, 80):
        pool.run_job(random.random())
    for x in xrange(21):
        time.sleep(1)
        pool.tick()
    print 'total shutdown...'
    pool.shutdown()
    time.sleep(1.5)
    print 'should not sleep'
    pool.run_job(15)
    print 'exit'