from __future__ import print_function from six.moves import queue as Queue from six.moves import range from functools import partial import threading import time from tqdm import tqdm DEFAULT_THREADS = 20 class ThreadedQueue(object): """Grant threaded task processing to any derived class.""" def __init__(self, n_threads, queue_size=0, progress=None): self._n_threads = n_threads self._queue = Queue.Queue(maxsize=queue_size) # 0 = infinite size self._error_queue = Queue.Queue(maxsize=queue_size) self._threads = () self._terminate = threading.Event() self._processed_lock = threading.Lock() self.processed = 0 self._inserted = 0 self.with_progress = progress self.start_threads(n_threads) @property def pending(self): return self._queue.qsize() def put(self, fn): """ Enqueue a task function for processing. Requires: fn: a function object that takes one argument that is the interface associated with each thread. e.g. def download(api): results.append(api.download()) self.put(download) Returns: self """ self._inserted += 1 self._queue.put(fn, block=True) return self def start_threads(self, n_threads): """ Terminate existing threads and create a new set if the thread number doesn't match the desired number. Required: n_threads: (int) number of threads to spawn Returns: self """ if n_threads == len(self._threads): return self # Terminate all previous tasks with the existing # event object, then create a new one for the next # generation of threads. The old object will hang # around in memory until the threads actually terminate # after another iteration. self._terminate.set() self._terminate = threading.Event() threads = [] for _ in range(n_threads): worker = threading.Thread( target=self._consume_queue, args=(self._terminate,) ) worker.daemon = True worker.start() threads.append(worker) self._threads = tuple(threads) return self def are_threads_alive(self): """Returns: boolean indicating if any threads are alive""" return any(map(lambda t: t.is_alive(), self._threads)) def kill_threads(self): """Kill all threads.""" self._terminate.set() while self.are_threads_alive(): time.sleep(0.001) self._threads = () return self def _initialize_interface(self): """ This is used to initialize the interfaces used in each thread. You should reimplement it in subclasses. For example, return an API object, file handle, or network connection. The functions you pass into the self._queue will get it as the first parameter. e.g. an implementation in a subclass. def _initialize_interface(self): return HTTPConnection() def other_function(self): def threaded_file_read(connection): # do stuff self._queue.put(threaded_file_handle) Returns: Interface object used in threads """ return None def _close_interface(self, interface): """Allows derived classes to clean up after a thread finishes.""" pass def _consume_queue(self, terminate_evt): """ This is the main thread function that consumes functions that are inside the _queue object. To use, execute self._queue(fn), where fn is a function that performs some kind of network IO or otherwise benefits from threading and is independent. terminate_evt is automatically passed in on thread creation and is a common event for this generation of threads. The threads will terminate when the event is set and the queue burns down. Returns: void """ interface = self._initialize_interface() while not terminate_evt.is_set(): try: fn = self._queue.get(block=True, timeout=0.01) except Queue.Empty: continue # periodically check if the thread is supposed to die fn = partial(fn, interface) try: self._consume_queue_execution(fn) except Exception as err: self._error_queue.put(err) self._close_interface(interface) def _consume_queue_execution(self, fn): """ The actual task execution in each thread. This is broken out so that exceptions can be caught in derived classes and allow them to manipulate the errant task, e.g. putting it back in the queue for a retry. Every task processed will automatically be marked complete. Required: [0] fn: A curried function that includes the interface as its first argument. Returns: void """ # `finally` fires after all success or exceptions # exceptions are handled in derived classes # and uncaught ones are caught as a last resort # in _consume_queue to be raised on the main thread. try: fn() finally: with self._processed_lock: self.processed += 1 self._queue.task_done() def _check_errors(self): try: err = self._error_queue.get(block=False) self._error_queue.task_done() self.kill_threads() raise err except Queue.Empty: pass def wait(self, progress=None): """ Allow background threads to process until the task queue is empty. If there are no threads, in theory the queue should always be empty as processing happens immediately on the main thread. Optional: progress: (bool or str) show a tqdm progress bar optionally with a description if a string is provided Returns: self (for chaining) Raises: The first exception recieved from threads """ if not len(self._threads): return self desc = None if type(progress) is str: desc = progress last = self._inserted with tqdm(total=self._inserted, disable=(not progress), desc=desc) as pbar: # Allow queue to consume, but check up on # progress and errors every tenth of a second while not self._queue.empty(): size = self._queue.qsize() delta = last - size if delta != 0: # We should crash on negative numbers pbar.update(delta) last = size self._check_errors() time.sleep(0.1) # Wait until all tasks in the queue are # fully processed. queue.task_done must be # called for each task. self._queue.join() self._check_errors() final = self._inserted - last if final: pbar.update(final) if self._queue.empty(): self._inserted = 0 return self def __del__(self): self.wait() # if no threads were set the queue is always empty self.kill_threads() def __enter__(self): if self.__class__ is ThreadedQueue and self._n_threads == 0: raise ValueError("Using 0 threads in base class ThreadedQueue with statement will never exit.") self.start_threads(self._n_threads) return self def __exit__(self, exception_type, exception_value, traceback): self.wait(progress=self.with_progress) self.kill_threads()