text: #!/usr/bin/env python
  """A simple thread pool for the Google Response Rig.
  
  This file defines a simple thread pool that is used throughout this
  project for parallelizing data store accesses. This thread pool is
  rather lightweight and optimized to be used in combination with the
  GRR data_store modules. It is not meant to be general purpose - if you
  need a generalized thread pool, you should probably use a better
  suited alternative implementation.
  
  If during creation not all new worker threads can be spawned by the
  ThreadPool, a log entry will be written but execution will continue
  using a smaller pool of workers. In this case, consider reducing the
  --threadpool_size.
  
  Example usage:
  >>> def PrintMsg(value):
  >>>   print "Message: %s" % value
  >>> for _ in range(10):
  >>>   SharedPool().AddTask(PrintMsg, ("Hello World!", ))
  >>> SharedPool().Join()
  
  """
  
  import itertools
  import logging
  import queue
  import threading
  import time
  
  import psutil
  
  from grr_response_core.lib import utils
  from grr_response_core.lib.util import collection
  from grr_response_core.lib.util import random
  from grr_response_core.stats import metrics
  
  STOP_MESSAGE = "Stop message"
  
  THREADPOOL_OUTSTANDING_TASKS = metrics.Gauge(
      "threadpool_outstanding_tasks", int, fields=[("pool_name", str)])
  THREADPOOL_THREADS = metrics.Gauge(
      "threadpool_threads", int, fields=[("pool_name", str)])
  THREADPOOL_CPU_USE = metrics.Gauge(
      "threadpool_cpu_use", float, fields=[("pool_name", str)])
  THREADPOOL_TASK_EXCEPTIONS = metrics.Counter(
      "threadpool_task_exceptions", fields=[("pool_name", str)])
  THREADPOOL_WORKING_TIME = metrics.Event(
      "threadpool_working_time", fields=[("pool_name", str)])
  THREADPOOL_QUEUEING_TIME = metrics.Event(
      "threadpool_queueing_time", fields=[("pool_name", str)])
  
  
  class Error(Exception):
    pass
  
  
  class ThreadPoolNotStartedError(Error):
    """Raised when a task is added to a not-yet-started pool."""
  
  
  class DuplicateThreadpoolError(Error):
    """Raised when a thread pool with the same name already exists."""
  
  
  class Full(Error):
    """Raised when the threadpool is full."""
  
  
  class _WorkerThread(threading.Thread):
    """The workers used in the ThreadPool class."""
  
    def __init__(self, message_queue, pool, name):
      # pyformat: disable
      """Initializer.
  
      This creates a new worker object for the ThreadPool class.
  
      Args:
        message_queue: A queue.Queue object used by the ThreadPool class to
          communicate with the workers. When a new task arrives, the ThreadPool
          notifies the workers by putting a message into this queue that has the
          format (target, args, name, queueing_time).
  
          target - A callable, the function to call.
          args - A tuple of positional arguments to target. Keyword arguments
                 are not supported.
          name - A name for this task. If None, it will be unique generated by
                 the threading library.
          queueing_time - The timestamp when this task was queued as returned by
                          time.time().
  
          Or, alternatively, the message in the queue can be STOP_MESSAGE
          which indicates that the worker should terminate.
  
        pool: The thread pool this worker belongs to.
        name: A name for this worker thread.
      """
      # pyformat: enable
      super().__init__(name=name)
  
      self.pool = pool
      self._queue = message_queue
      self.daemon = True
      self.idle = True
      self.started = time.time()
  
    def ProcessTask(self, target, args, name, queueing_time):
      """Processes the tasks."""
  
      if self.pool.name:
        time_in_queue = time.time() - queueing_time
        THREADPOOL_QUEUEING_TIME.RecordEvent(
            time_in_queue, fields=[self.pool.name])
  
        start_time = time.time()
      try:
        target(*args)
      # We can't let a worker die because one of the tasks it has to process
      # throws an exception. Therefore, we catch every error that is
      # raised in the call to target().
      except Exception as e:  # pylint: disable=broad-except
        if self.pool.name:
          THREADPOOL_TASK_EXCEPTIONS.Increment(fields=[self.pool.name])
        logging.exception("Caught exception in worker thread (%s): %s", name,
                          str(e))
  
      if self.pool.name:
        total_time = time.time() - start_time
        THREADPOOL_WORKING_TIME.RecordEvent(total_time, fields=[self.pool.name])
  
    def _RemoveFromPool(self):
      """Remove ourselves from the pool.
  
      Returns:
        True if removal was possible, and False if it was not possible.
      """
      with self.pool.lock:
  
        # Pool is shutting down, we can't interfere.
        if not self.pool.started:
          return False
  
        # Keep a minimum number of threads in the pool.
        if len(self.pool) <= self.pool.min_threads:
          return False
  
        # Remove us from our pool.
        self.pool._RemoveWorker(self.name)  # pylint: disable=protected-access
  
        return True
  
    def run(self):
      """This overrides the Thread.run method.
  
      This method checks in an endless loop if new tasks are available
      in the queue and processes them.
      """
      while True:
        if self.pool.name:
          self.idle = True
  
        try:
          # Wait 60 seconds for a message, otherwise exit. This ensures that the
          # threadpool will be trimmed down when load is light.
          task = self._queue.get(timeout=60)
  
          if self.pool.name:
            self.idle = False
  
          try:
            # The pool told us to quit, likely because it is stopping.
            if task == STOP_MESSAGE:
              return
  
            self.ProcessTask(*task)
          finally:
            self._queue.task_done()
  
        except queue.Empty:
          if self._RemoveFromPool():
            return
  
        # Try to trim old threads down when they get too old. This helps the
        # thread pool size to shrink, even when it is not idle (e.g. if it is CPU
        # bound) since threads are forced to exit, but new threads will not be
        # created if the utilization is too high - resulting in a reduction of
        # threadpool size under CPU load.
        if time.time() - self.started > 600 and self._RemoveFromPool():
          return
  
  
  THREADPOOL = None
  
  
  class ThreadPool(object):
    """A thread pool implementation.
  
    The thread pool starts with the minimum number of threads. As tasks are added,
    they are added to a queue and once this is full, more threads are added until
    we reach max_threads or this process's CPU utilization approaches 100%. Since
    Python uses a global lock (GIL) it is not possible for the interpreter to use
    more than 100% of a single core. Any additional threads actually reduce
    performance due to thread switching overheads. Therefore we ensure that the
    thread pool is not too loaded at any one time.
  
    When threads are idle longer than 60 seconds they automatically exit. This
    ensures that our memory footprint is reduced when load is light.
  
    Note that this class should not be instantiated directly, but the Factory
    should be used.
    """
    # A global dictionary of pools, keyed by pool name.
    POOLS = {}
    factory_lock = threading.Lock()
  
    JOIN_TIMEOUT_DECISECONDS = 600
  
    @classmethod
    def Factory(cls, name, min_threads, max_threads=None):
      """Creates a new thread pool with the given name.
  
      If the thread pool of this name already exist, we just return the existing
      one. This allows us to have different pools with different characteristics
      used by different parts of the code, at the same time.
  
      Args:
        name: The name of the required pool.
        min_threads: The number of threads in the pool.
        max_threads: The maximum number of threads to grow the pool to. If not set
          we do not grow the pool.
  
      Returns:
        A threadpool instance.
      """
      with cls.factory_lock:
        result = cls.POOLS.get(name)
        if result is None:
          cls.POOLS[name] = result = cls(
              name, min_threads, max_threads=max_threads)
  
        return result
  
    def __init__(self, name, min_threads, max_threads=None):
      """This creates a new thread pool using min_threads workers.
  
      Args:
        name: A prefix to identify this thread pool in the exported stats.
        min_threads: The minimum number of worker threads this pool should have.
        max_threads: The maximum number of threads to grow the pool to. If not set
          we do not grow the pool.
  
      Raises:
        threading.ThreadError: If no threads can be spawned at all, ThreadError
                               will be raised.
        DuplicateThreadpoolError: This exception is raised if a thread pool with
                                  the desired name already exists.
      """
      self.min_threads = min_threads
      if max_threads is None or max_threads < min_threads:
        max_threads = min_threads
  
      self.max_threads = max_threads
      self._queue = queue.Queue(maxsize=max_threads)
      self.name = name
      self.started = False
      self.process = psutil.Process()
  
      # A reference for all our workers. Keys are thread names, and values are the
      # _WorkerThread instance.
      self._workers = {}
      # Read-only copy of self._workers that is thread-safe for reading.
      self._workers_ro_copy = {}
      self.lock = threading.RLock()
  
      if not self.name:
        raise ValueError("Unnamed thread pools not allowed.")
  
      if self.name in self.POOLS:
        raise DuplicateThreadpoolError(
            "A thread pool with the name %s already exists." % name)
  
      THREADPOOL_OUTSTANDING_TASKS.SetCallback(
          self._queue.qsize, fields=[self.name])
      THREADPOOL_THREADS.SetCallback(lambda: len(self), fields=[self.name])
      THREADPOOL_CPU_USE.SetCallback(self.CPUUsage, fields=[self.name])
  
    def __del__(self):
      if self.started:
        self.Stop()
  
    @property
    def pending_tasks(self):
      # This is thread safe as self._queue is thread safe.
      return self._queue.qsize()
  
    @property
    def busy_threads(self):
      return len([x for x in self._workers_ro_copy.values() if not x.idle])
  
    def __len__(self):
      return len(self._workers_ro_copy)
  
    @utils.Synchronized
    def Start(self):
      """This starts the worker threads."""
      if not self.started:
        self.started = True
        for _ in range(self.min_threads):
          self._AddWorker()
  
    @utils.Synchronized
    def _AddWorker(self):
      worker_name = self.name + "-%d" % random.UInt32()
      worker = _WorkerThread(self._queue, self, worker_name)
      worker.start()
  
      self._workers[worker.name] = worker
      self._workers_ro_copy = self._workers.copy()
  
    @utils.Synchronized
    def _RemoveWorker(self, key):
      self._workers.pop(key, None)
      self._workers_ro_copy = self._workers.copy()
  
    @utils.Synchronized
    def Stop(self, join_timeout=600):
      """This stops all the worker threads."""
      if not self.started:
        logging.warning("Tried to stop a thread pool that was not running.")
        return
  
      # Remove all workers from the pool.
      workers = list(self._workers.values())
      self._workers = {}
      self._workers_ro_copy = {}
  
      # Send a stop message to all the workers. We need to be careful here to not
      # send messages while we are still counting. If threads that haven't been
      # counted yet pick up a message and exit, the count will be off and the
      # shutdown process will deadlock.
      stop_messages_needed = 0
      for worker in workers:
        if worker.is_alive():
          stop_messages_needed += 1
  
      for _ in range(stop_messages_needed):
        self._queue.put(STOP_MESSAGE)
  
      self.started = False
      self.Join()
  
      # Wait for the threads to all exit now.
      for worker in workers:
        worker.join(join_timeout)
        if worker.is_alive():
          raise RuntimeError("Threadpool worker did not finish in time.")
  
    def AddTask(self,
                target,
                args=(),
                name="Unnamed task",
                blocking=True,
                inline=True):
      """Adds a task to be processed later.
  
      Args:
        target: A callable which should be processed by one of the workers.
        args: A tuple of arguments to target.
        name: The name of this task. Used to identify tasks in the log.
        blocking: If True we block until the task is finished, otherwise we raise
          queue.Full
        inline: If set, process the task inline when the queue is full. This
          implies no blocking. Specifying inline helps if the worker tasks are
          blocked because it still ensures some progress is made. However, this
          can generally block the calling thread even after the threadpool is
          available again and therefore decrease efficiency.
  
      Raises:
        ThreadPoolNotStartedError: if the pool was not started yet.
        queue.Full: if the pool is full and can not accept new jobs.
      """
      if not self.started:
        raise ThreadPoolNotStartedError(self.name)
  
      # This pool should have no worker threads - just run the task inline.
      if self.max_threads == 0:
        target(*args)
        return
  
      if inline:
        blocking = False
  
      with self.lock:
        while True:
          # This check makes sure that the threadpool will add new workers
          # even if the queue is not full. This is needed for a scenario when
          # a fresh threadpool is created (say, with min_threads=1 and
          # max_threads=10) and 2 long-running tasks are added. The code below
          # will spawn a new worker for a second long-running task.
          if len(self) < self.max_threads:
            try:
              self._AddWorker()
            except (RuntimeError, threading.ThreadError) as e:
              logging.error(
                  "Threadpool exception: "
                  "Could not spawn worker threads: %s", e)
  
          try:
            # Push the task on the queue but raise if unsuccessful.
            self._queue.put((target, args, name, time.time()), block=False)
            return
          except queue.Full:
            # We increase the number of active threads if we do not exceed the
            # maximum _and_ our process CPU utilization is not too high. This
            # ensures that if the workers are waiting on IO we add more workers,
            # but we do not waste workers when tasks are CPU bound.
            if len(self) < self.max_threads:
              try:
                self._AddWorker()
                continue
  
              # If we fail to add a worker we should keep going anyway.
              except (RuntimeError, threading.ThreadError) as e:
                logging.error(
                    "Threadpool exception: "
                    "Could not spawn worker threads: %s", e)
  
            # If we need to process the task inline just break out of the loop,
            # therefore releasing the lock and run the task inline.
            if inline:
              break
  
            # We should block and try again soon.
            elif blocking:
              try:
                self._queue.put((target, args, name, time.time()),
                                block=True,
                                timeout=1)
                return
              except queue.Full:
                continue
  
            else:
              raise Full()
  
      # We don't want to hold the lock while running the task inline
      if inline:
        target(*args)
  
    def CPUUsage(self):
      return self.process.cpu_percent(0)
  
    def Join(self):
      """Waits until all outstanding tasks are completed."""
  
      for _ in range(self.JOIN_TIMEOUT_DECISECONDS):
        if self._queue.empty() and not self.busy_threads:
          return
        time.sleep(0.1)
  
      raise ValueError("Timeout during Join() for threadpool %s." % self.name)
  
  
  class MockThreadPool(object):
    """A mock thread pool which runs all jobs serially."""
  
    def __init__(self, name, min_threads, max_threads=None, ignore_errors=True):
      _ = name
      _ = min_threads
      _ = max_threads
      self.ignore_errors = ignore_errors
  
    def AddTask(self, target, args, name="Unnamed task"):
      _ = name
      try:
        target(*args)
        # The real threadpool can not raise from a task. We emulate this here.
      except Exception as e:  # pylint: disable=broad-except
        logging.exception("MockThreadPool worker raised %s", e)
        if not self.ignore_errors:
          raise
  
    @classmethod
    def Factory(cls, name, min_threads, max_threads=None):
      return cls(name, min_threads, max_threads=max_threads)
  
    def Start(self):
      pass
  
    def Stop(self):
      pass
  
    def Join(self):
      pass
  
  
  class BatchConverter(object):
    """Generic class that does multi-threaded values conversion.
  
    BatchConverter converts a set of values to a set of different values in
    batches using a threadpool.
    """
  
    def __init__(self,
                 batch_size=1000,
                 threadpool_prefix="batch_processor",
                 threadpool_size=10):
      """BatchProcessor constructor.
  
      Args:
        batch_size: All the values will be processed in batches of this size.
        threadpool_prefix: Prefix that will be used in thread pool's threads
          names.
        threadpool_size: Size of a thread pool that will be used. If
          threadpool_size is 0, no threads will be used and all conversions will
          be done in the current thread.
      """
      super().__init__()
      self.batch_size = batch_size
      self.threadpool_prefix = threadpool_prefix
      self.threadpool_size = threadpool_size
  
    def ConvertBatch(self, batch):
      """ConvertBatch is called for every batch to do the conversion.
  
      Args:
        batch: Batch to convert.
  
      Returns:
        List with converted values.
      """
      raise NotImplementedError()
  
    def Convert(self, values, start_index=0, end_index=None):
      """Converts given collection to exported values.
  
      This method uses a threadpool to do the conversion in parallel. It
      blocks for up to one hour until everything is converted.
  
      Args:
        values: Iterable object with values to convert.
        start_index: Start from this index in the collection.
        end_index: Finish processing on the (index - 1) element of the collection.
          If None, work till the end of the collection.
  
      Returns:
        Nothing. ConvertedBatch() should handle the results.
      """
      if not values:
        return
  
      try:
        total_batch_count = len(values) // self.batch_size
      except TypeError:
        total_batch_count = -1
  
      pool = ThreadPool.Factory(self.threadpool_prefix, self.threadpool_size)
      val_iterator = itertools.islice(values, start_index, end_index)
  
      pool.Start()
      try:
        for batch_index, batch in enumerate(
            collection.Batch(val_iterator, self.batch_size)):
          logging.debug("Processing batch %d out of %d", batch_index,
                        total_batch_count)
  
          pool.AddTask(
              target=self.ConvertBatch,
              args=(batch,),
              name="batch_%d" % batch_index,
              inline=False)
  
      finally:
        pool.Stop(join_timeout=3600)
  