thread_pool.py

Go to the documentation of this file.

# -*- coding: utf-8 -*-

# python libraries
import threading
import abc
from Queue import Queue
from Queue import Empty

# local libraries
from rospeex_core import exceptions as ext
from rospeex_core import logging_util


logger = logging_util.get_logger(__name__)

__all__ = ['ThreadPool']


class IRequest(object):
    """ Request data
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def request(self, *args, **kwargs):
        pass    # pragma: no cover

    @abc.abstractmethod
    def has_error(self):
        pass    # pragma: no cover

    @abc.abstractmethod
    def response(self):
        pass    # pragma: no cover

    @abc.abstractmethod
    def result(self):
        pass    # pragma: no cover


class Worker(threading.Thread):
    """ Thrad worker class
    """
    def __init__(self, tasks, finish_tasks):
        """ init worker thread
        @param tasks:
        @param finish_tasks:
        """
        threading.Thread.__init__(self)
        self._tasks = tasks
        self._finish_tasks = finish_tasks
        self._stop_request = threading.Event()
        # self.daemon = True
        self.start()

    def run(self):
        """ run thread
        @return: None
        """
        while not self._stop_request.isSet():
            try:
                request, args, kwargs = self._tasks.get(timeout=0.01)

                logger.debug('start task')
                request.request(*args, **kwargs)
                self._finish_tasks.put_nowait(request)
                self._tasks.task_done()
                logger.debug('finish task')

            except Empty:
                pass

    def join(self, timeout=None):
        """ end worker threaed
        @return: None
        """
        self._stop_request.set()
        self._tasks.join()
        super(Worker, self).join(timeout)


class ThreadPool(object):
    """ ThradPool class """
    def __init__(self, num_thread):
        """ init class
        @param num_thread: number of worker thread
        """
        self._tasks = Queue()
        self._finish_tasks = Queue()
        self._workers = [
            Worker(self._tasks, self._finish_tasks) for i in range(num_thread)
        ]

    def add_request(self, request, *args, **kwargs):
        """ add request
        @param request:
        @param args:
        @param kwargs:
        @return: None
        """
        if isinstance(request, IRequest):
            self._tasks.put_nowait([request, args, kwargs])
        else:
            msg = '{} is not IRequest instance.'.format(type(request))
            raise ext.InvalidRequestTypeException(msg)

    def wait_completion(self):
        """ wait completion
        @return: None
        """
        self._tasks.join()

    def check_completion(self):
        """ check completion
        @return: True for complete tasks / False for NOT complete tasks
        """
        if self._tasks.qsize() == 0:
            return True
        return False

    def get_finish_task(self, wait=True):
        """ get finish task
        @return: finish task
        """
        task = None
        try:
            task = self._finish_tasks.get(wait)
            self._finish_tasks.task_done()

        except Empty:
            logger.info('finish task que is empty')
        return task

    def check_finish_task(self):
        """
        @return: None
        """
        if self._finish_tasks.qsize() == 0:
            return True
        return False

    def end(self):
        """ end woker thread
        @return: None
        """
        logger.debug('end thread pool')
        self.wait_completion()

        for worker in self._workers:
            worker.join()