# Copyright 2013 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

"""Classes and functions related to threading."""

import functools
import inspect
import logging
import os
import Queue
import sys
import threading
import time
import traceback


# Priorities for tasks in AutoRetryThreadPool, particular values are important.
PRIORITY_HIGH = 1 << 8
PRIORITY_MED = 2 << 8
PRIORITY_LOW = 3 << 8


class LockWithAssert(object):
  """Wrapper around (non recursive) Lock that tracks its owner."""

  def __init__(self):
    self._lock = threading.Lock()
    self._owner = None

  def __enter__(self):
    self._lock.acquire()
    assert self._owner is None
    self._owner = threading.current_thread()

  def __exit__(self, _exc_type, _exec_value, _traceback):
    self.assert_locked('Releasing unowned lock')
    self._owner = None
    self._lock.release()
    return False

  def assert_locked(self, msg=None):
    """Asserts the lock is owned by running thread."""
    assert self._owner == threading.current_thread(), msg


class ThreadPoolError(Exception):
  """Base class for exceptions raised by ThreadPool."""
  pass


class ThreadPoolEmpty(ThreadPoolError):
  """Trying to get task result from a thread pool with no pending tasks."""
  pass


class ThreadPoolClosed(ThreadPoolError):
  """Trying to do something with a closed thread pool."""
  pass


class ThreadPool(object):
  """Multithreaded worker pool with priority support.

  When the priority of tasks match, it works in strict FIFO mode.
  """
  QUEUE_CLASS = Queue.PriorityQueue

  def __init__(self, initial_threads, max_threads, queue_size, prefix=None):
    """Immediately starts |initial_threads| threads.

    Arguments:
      initial_threads: Number of threads to start immediately. Can be 0 if it is
                       uncertain that threads will be needed.
      max_threads: Maximum number of threads that will be started when all the
                   threads are busy working. Often the number of CPU cores.
      queue_size: Maximum number of tasks to buffer in the queue. 0 for
                  unlimited queue. A non-zero value may make add_task()
                  blocking.
      prefix: Prefix to use for thread names. Pool's threads will be
              named '<prefix>-<thread index>'.
    """
    prefix = prefix or 'tp-0x%0x' % id(self)
    logging.debug(
        'New ThreadPool(%d, %d, %d): %s', initial_threads, max_threads,
        queue_size, prefix)
    assert initial_threads <= max_threads
    assert max_threads <= 1024

    self.tasks = self.QUEUE_CLASS(queue_size)
    self._max_threads = max_threads
    self._prefix = prefix

    # Used to assign indexes to tasks.
    self._num_of_added_tasks_lock = threading.Lock()
    self._num_of_added_tasks = 0

    # Lock that protected everything below (including conditional variable).
    self._lock = threading.Lock()

    # Condition 'bool(_outputs) or bool(_exceptions) or _pending_count == 0'.
    self._outputs_exceptions_cond = threading.Condition(self._lock)
    self._outputs = []
    self._exceptions = []

    # Number of pending tasks (queued or being processed now).
    self._pending_count = 0

    # List of threads.
    self._workers = []
    # Number of threads that are waiting for new tasks.
    self._ready = 0
    # Number of threads already added to _workers, but not yet running the loop.
    self._starting = 0
    # True if close was called. Forbids adding new tasks.
    self._is_closed = False

    for _ in range(initial_threads):
      self._add_worker()

  def _add_worker(self):
    """Adds one worker thread if there isn't too many. Thread-safe."""
    with self._lock:
      if len(self._workers) >= self._max_threads or self._is_closed:
        return False
      worker = threading.Thread(
        name='%s-%d' % (self._prefix, len(self._workers)), target=self._run)
      self._workers.append(worker)
      self._starting += 1
    logging.debug('Starting worker thread %s', worker.name)
    worker.daemon = True
    worker.start()
    return True

  def add_task(self, priority, func, *args, **kwargs):
    """Adds a task, a function to be executed by a worker.

    Arguments:
    - priority: priority of the task versus others. Lower priority takes
                precedence.
    - func: function to run. Can either return a return value to be added to the
            output list or be a generator which can emit multiple values.
    - args and kwargs: arguments to |func|. Note that if func mutates |args| or
                       |kwargs| and that the task is retried, see
                       AutoRetryThreadPool, the retry will use the mutated
                       values.

    Returns:
      Index of the item added, e.g. the total number of enqueued items up to
      now.
    """
    assert isinstance(priority, int)
    assert callable(func)
    with self._lock:
      if self._is_closed:
        raise ThreadPoolClosed('Can not add a task to a closed ThreadPool')
      start_new_worker = (
        # Pending task count plus new task > number of available workers.
        self.tasks.qsize() + 1 > self._ready + self._starting and
        # Enough slots.
        len(self._workers) < self._max_threads
      )
      self._pending_count += 1
    with self._num_of_added_tasks_lock:
      self._num_of_added_tasks += 1
      index = self._num_of_added_tasks
    self.tasks.put((priority, index, func, args, kwargs))
    if start_new_worker:
      self._add_worker()
    return index

  def _run(self):
    """Worker thread loop. Runs until a None task is queued."""
    # Thread has started, adjust counters.
    with self._lock:
      self._starting -= 1
      self._ready += 1
    while True:
      try:
        task = self.tasks.get()
      finally:
        with self._lock:
          self._ready -= 1
      try:
        if task is None:
          # We're done.
          return
        _priority, _index, func, args, kwargs = task
        if inspect.isgeneratorfunction(func):
          for out in func(*args, **kwargs):
            self._output_append(out)
        else:
          out = func(*args, **kwargs)
          self._output_append(out)
      except Exception as e:
        logging.warning('Caught exception: %s', e)
        exc_info = sys.exc_info()
        logging.info(''.join(traceback.format_tb(exc_info[2])))
        with self._outputs_exceptions_cond:
          self._exceptions.append(exc_info)
          self._outputs_exceptions_cond.notifyAll()
      finally:
        try:
          # Mark thread as ready again, mark task as processed. Do it before
          # waking up threads waiting on self.tasks.join(). Otherwise they might
          # find ThreadPool still 'busy' and perform unnecessary wait on CV.
          with self._outputs_exceptions_cond:
            self._ready += 1
            self._pending_count -= 1
            if self._pending_count == 0:
              self._outputs_exceptions_cond.notifyAll()
          self.tasks.task_done()
        except Exception as e:
          # We need to catch and log this error here because this is the root
          # function for the thread, nothing higher will catch the error.
          logging.exception('Caught exception while marking task as done: %s',
                            e)

  def _output_append(self, out):
    if out is not None:
      with self._outputs_exceptions_cond:
        self._outputs.append(out)
        self._outputs_exceptions_cond.notifyAll()

  def join(self):
    """Extracts all the results from each threads unordered.

    Call repeatedly to extract all the exceptions if desired.

    Note: will wait for all work items to be done before returning an exception.
    To get an exception early, use get_one_result().
    """
    # TODO(maruel): Stop waiting as soon as an exception is caught.
    self.tasks.join()
    with self._outputs_exceptions_cond:
      if self._exceptions:
        e = self._exceptions.pop(0)
        raise e[0], e[1], e[2]
      out = self._outputs
      self._outputs = []
    return out

  def get_one_result(self):
    """Returns the next item that was generated or raises an exception if one
    occurred.

    Raises:
      ThreadPoolEmpty - no results available.
    """
    # Get first available result.
    for result in self.iter_results():
      return result
    # No results -> tasks queue is empty.
    raise ThreadPoolEmpty('Task queue is empty')

  def iter_results(self):
    """Yields results as they appear until all tasks are processed."""
    while True:
      # Check for pending results.
      result = None
      self._on_iter_results_step()
      with self._outputs_exceptions_cond:
        if self._exceptions:
          e = self._exceptions.pop(0)
          raise e[0], e[1], e[2]
        if self._outputs:
          # Remember the result to yield it outside of the lock.
          result = self._outputs.pop(0)
        else:
          # No pending tasks -> all tasks are done.
          if not self._pending_count:
            return
          # Some task is queued, wait for its result to appear.
          # Use non-None timeout so that process reacts to Ctrl+C and other
          # signals, see http://bugs.python.org/issue8844.
          self._outputs_exceptions_cond.wait(timeout=0.1)
          continue
      yield result

  def close(self):
    """Closes all the threads."""
    # Ensure no new threads can be started, self._workers is effectively
    # a constant after that and can be accessed outside the lock.
    with self._lock:
      if self._is_closed:
        raise ThreadPoolClosed('Can not close already closed ThreadPool')
      self._is_closed = True
    for _ in range(len(self._workers)):
      # Enqueueing None causes the worker to stop.
      self.tasks.put(None)
    for t in self._workers:
      # 'join' without timeout blocks signal handlers, spin with timeout.
      while t.is_alive():
        t.join(30)
    logging.debug(
      'Thread pool \'%s\' closed: spawned %d threads total',
      self._prefix, len(self._workers))

  def abort(self):
    """Empties the queue.

    To be used when the pool should stop early, like when Ctrl-C was detected.

    Returns:
      Number of tasks cancelled.
    """
    index = 0
    while True:
      try:
        self.tasks.get_nowait()
        self.tasks.task_done()
        index += 1
      except Queue.Empty:
        return index

  def _on_iter_results_step(self):
    pass

  def __enter__(self):
    """Enables 'with' statement."""
    return self

  def __exit__(self, _exc_type, _exc_value, _traceback):
    """Enables 'with' statement."""
    self.close()


class AutoRetryThreadPool(ThreadPool):
  """Automatically retries enqueued operations on exception."""
  # See also PRIORITY_* module-level constants.
  INTERNAL_PRIORITY_BITS = (1<<8) - 1

  def __init__(self, exceptions, retries, *args, **kwargs):
    """
    Arguments:
      exceptions: list of exception classes that can be retried on.
      retries: maximum number of retries to do.
    """
    assert exceptions and all(issubclass(e, Exception) for e in exceptions), (
        exceptions)
    assert 1 <= retries <= self.INTERNAL_PRIORITY_BITS
    super(AutoRetryThreadPool, self).__init__(*args, **kwargs)
    self._swallowed_exceptions = tuple(exceptions)
    self._retries = retries

  def add_task(self, priority, func, *args, **kwargs):
    """Tasks added must not use the lower priority bits since they are reserved
    for retries.
    """
    assert (priority & self.INTERNAL_PRIORITY_BITS) == 0
    return super(AutoRetryThreadPool, self).add_task(
        priority,
        self._task_executer,
        priority,
        None,
        func,
        *args,
        **kwargs)

  def add_task_with_channel(self, channel, priority, func, *args, **kwargs):
    """Tasks added must not use the lower priority bits since they are reserved
    for retries.
    """
    assert (priority & self.INTERNAL_PRIORITY_BITS) == 0
    return super(AutoRetryThreadPool, self).add_task(
        priority,
        self._task_executer,
        priority,
        channel,
        func,
        *args,
        **kwargs)

  def _task_executer(self, priority, channel, func, *args, **kwargs):
    """Wraps the function and automatically retry on exceptions."""
    try:
      result = func(*args, **kwargs)
      if channel is None:
        return result
      channel.send_result(result)
    # pylint: disable=catching-non-exception
    except self._swallowed_exceptions as e:
      # Retry a few times, lowering the priority.
      actual_retries = priority & self.INTERNAL_PRIORITY_BITS
      if actual_retries < self._retries:
        priority += 1
        logging.debug(
            'Swallowed exception \'%s\'. Retrying at lower priority %X',
            e, priority)
        super(AutoRetryThreadPool, self).add_task(
            priority,
            self._task_executer,
            priority,
            channel,
            func,
            *args,
            **kwargs)
        return
      if channel is None:
        raise
      channel.send_exception()
    except Exception:
      if channel is None:
        raise
      channel.send_exception()


class IOAutoRetryThreadPool(AutoRetryThreadPool):
  """Thread pool that automatically retries on IOError.

  Supposed to be used for IO bound tasks, and thus default maximum number of
  worker threads is independent of number of CPU cores.
  """
  # Initial and maximum number of worker threads.
  INITIAL_WORKERS = 2
  MAX_WORKERS = 16 if sys.maxsize > 2L**32 else 8
  RETRIES = 5

  def __init__(self):
    super(IOAutoRetryThreadPool, self).__init__(
        [IOError],
        self.RETRIES,
        self.INITIAL_WORKERS,
        self.MAX_WORKERS,
        0,
        'io')


class Progress(object):
  """Prints progress and accepts updates thread-safely."""
  def __init__(self, columns):
    """Creates a Progress bar that will updates asynchronously from the worker
    threads.

    Arguments:
      columns: list of tuple(name, initialvalue), defines both the number of
               columns and their initial values.
    """
    assert all(
        len(c) == 2 and isinstance(c[0], str) and isinstance(c[1], int)
        for c in columns), columns
    # Members to be used exclusively in the primary thread.
    self.use_cr_only = True
    self.unfinished_commands = set()
    self.start = time.time()
    self._last_printed_line = ''
    self._columns = [c[1] for c in columns]
    self._columns_lookup = dict((c[0], i) for i, c in enumerate(columns))
    # Setting it to True forces a print on the first print_update() call.
    self._value_changed = True

    # To be used in all threads.
    self._queued_updates = Queue.Queue()

  def update_item(self, name, raw=False, **kwargs):
    """Queue information to print out.

    Arguments:
      name: string to print out to describe something that was completed.
      raw: if True, prints the data without the header.
      raw: if True, prints the data without the header.
      <kwargs>: argument name is a name of a column. it's value is the increment
                to the column, value is usually 0 or 1.
    """
    assert isinstance(name, basestring), repr(name)
    assert isinstance(raw, bool), repr(raw)
    assert all(isinstance(v, int) for v in kwargs.itervalues()), repr(kwargs)
    args = [(self._columns_lookup[k], v) for k, v in kwargs.iteritems() if v]
    self._queued_updates.put((name, raw, args))

  def print_update(self):
    """Prints the current status."""
    # Flush all the logging output so it doesn't appear within this output.
    for handler in logging.root.handlers:
      handler.flush()

    got_one = False
    while True:
      try:
        name, raw, args = self._queued_updates.get_nowait()
      except Queue.Empty:
        break

      for k, v in args:
        self._columns[k] += v
      self._value_changed = bool(args)
      if not name:
        # Even if raw=True, there's nothing to print.
        continue

      got_one = True
      if raw:
        # Prints the data as-is.
        self._last_printed_line = ''
        sys.stdout.write('\n%s\n' % name.strip('\n'))
      else:
        line, self._last_printed_line = self._gen_line(name)
        sys.stdout.write(line)

    if not got_one and self._value_changed:
      # Make sure a line is printed in that case where statistics changes.
      line, self._last_printed_line = self._gen_line('')
      sys.stdout.write(line)
      got_one = True
    self._value_changed = False
    if got_one:
      # Ensure that all the output is flushed to prevent it from getting mixed
      # with other output streams (like the logging streams).
      sys.stdout.flush()

    if self.unfinished_commands:
      logging.debug('Waiting for the following commands to finish:\n%s',
                    '\n'.join(self.unfinished_commands))

  def _gen_line(self, name):
    """Generates the line to be printed."""
    next_line = ('[%s] %6.2fs %s') % (
        self._render_columns(), time.time() - self.start, name)
    # Fill it with whitespace only if self.use_cr_only is set.
    prefix = ''
    if self.use_cr_only and self._last_printed_line:
      prefix = '\r'
    if self.use_cr_only:
      suffix = ' ' * max(0, len(self._last_printed_line) - len(next_line))
    else:
      suffix = '\n'
    return '%s%s%s' % (prefix, next_line, suffix), next_line

  def _render_columns(self):
    """Renders the columns."""
    columns_as_str = map(str, self._columns)
    max_len = max(map(len, columns_as_str))
    return '/'.join(i.rjust(max_len) for i in columns_as_str)


class QueueWithProgress(Queue.PriorityQueue):
  """Implements progress support in join()."""
  def __init__(self, progress, *args, **kwargs):
    Queue.PriorityQueue.__init__(self, *args, **kwargs)
    self.progress = progress

  def task_done(self):
    """Contrary to Queue.task_done(), it wakes self.all_tasks_done at each task
    done.
    """
    with self.all_tasks_done:
      try:
        unfinished = self.unfinished_tasks - 1
        if unfinished < 0:
          raise ValueError('task_done() called too many times')
        self.unfinished_tasks = unfinished
        # This is less efficient, because we want the Progress to be updated.
        self.all_tasks_done.notify_all()
      except Exception as e:
        logging.exception('task_done threw an exception.\n%s', e)

  def wake_up(self):
    """Wakes up all_tasks_done.

    Unlike task_done(), do not substract one from self.unfinished_tasks.
    """
    # TODO(maruel): This is highly inefficient, since the listener is awaken
    # twice; once per output, once per task. There should be no relationship
    # between the number of output and the number of input task.
    with self.all_tasks_done:
      self.all_tasks_done.notify_all()

  def join(self):
    """Calls print_update() whenever possible."""
    self.progress.print_update()
    with self.all_tasks_done:
      while self.unfinished_tasks:
        self.progress.print_update()
        # Use a short wait timeout so updates are printed in a timely manner.
        # TODO(maruel): Find a way so Progress.queue and self.all_tasks_done
        # share the same underlying event so no polling is necessary.
        self.all_tasks_done.wait(0.1)
      self.progress.print_update()


class ThreadPoolWithProgress(ThreadPool):
  QUEUE_CLASS = QueueWithProgress

  def __init__(self, progress, *args, **kwargs):
    self.QUEUE_CLASS = functools.partial(self.QUEUE_CLASS, progress)
    super(ThreadPoolWithProgress, self).__init__(*args, **kwargs)

  def _output_append(self, out):
    """Also wakes up the listener on new completed test_case."""
    super(ThreadPoolWithProgress, self)._output_append(out)
    self.tasks.wake_up()

  def _on_iter_results_step(self):
    self.tasks.progress.print_update()


class DeadlockDetector(object):
  """Context manager that can detect deadlocks.

  It will dump stack frames of all running threads if its 'ping' method isn't
  called in time.

  Usage:
    with DeadlockDetector(timeout=60) as detector:
      for item in some_work():
        ...
        detector.ping()
        ...

  Arguments:
    timeout - maximum allowed time between calls to 'ping'.
  """

  def __init__(self, timeout):
    self.timeout = timeout
    self._thread = None
    # Thread stop condition. Also lock for shared variables below.
    self._stop_cv = threading.Condition()
    self._stop_flag = False
    # Time when 'ping' was called last time.
    self._last_ping = None
    # True if pings are coming on time.
    self._alive = True

  def __enter__(self):
    """Starts internal watcher thread."""
    assert self._thread is None
    self.ping()
    self._thread = threading.Thread(name='deadlock-detector', target=self._run)
    self._thread.daemon = True
    self._thread.start()
    return self

  def __exit__(self, *_args):
    """Stops internal watcher thread."""
    assert self._thread is not None
    with self._stop_cv:
      self._stop_flag = True
      self._stop_cv.notify()
    self._thread.join()
    self._thread = None
    self._stop_flag = False

  def ping(self):
    """Notify detector that main thread is still running.

    Should be called periodically to inform the detector that everything is
    running as it should.
    """
    with self._stop_cv:
      self._last_ping = time.time()
      self._alive = True

  def _run(self):
    """Loop that watches for pings and dumps threads state if ping is late."""
    with self._stop_cv:
      while not self._stop_flag:
        # Skipped deadline? Dump threads and switch to 'not alive' state.
        if self._alive and time.time() > self._last_ping + self.timeout:
          self.dump_threads(time.time() - self._last_ping, True)
          self._alive = False

        # Pings are on time?
        if self._alive:
          # Wait until the moment we need to dump stack traces.
          # Most probably some other thread will call 'ping' to move deadline
          # further in time. We don't bother to wake up after each 'ping',
          # only right before initial expected deadline.
          self._stop_cv.wait(self._last_ping + self.timeout - time.time())
        else:
          # Skipped some pings previously. Just periodically silently check
          # for new pings with some arbitrary frequency.
          self._stop_cv.wait(self.timeout * 0.1)

  @staticmethod
  def dump_threads(timeout=None, skip_current_thread=False):
    """Dumps stack frames of all running threads."""
    all_threads = threading.enumerate()
    current_thread_id = threading.current_thread().ident

    # Collect tracebacks: thread name -> traceback string.
    tracebacks = {}

    # pylint: disable=W0212
    for thread_id, frame in sys._current_frames().iteritems():
      # Don't dump deadlock detector's own thread, it's boring.
      if thread_id == current_thread_id and skip_current_thread:
        continue

      # Try to get more informative symbolic thread name.
      name = 'untitled'
      for thread in all_threads:
        if thread.ident == thread_id:
          name = thread.name
          break
      name += ' #%d' % (thread_id,)
      tracebacks[name] = ''.join(traceback.format_stack(frame))

    # Function to print a message. Makes it easier to change output destination.
    def output(msg):
      logging.warning(msg.rstrip())

    # Print tracebacks, sorting them by thread name. That way a thread pool's
    # threads will be printed as one group.
    output('=============== Potential deadlock detected ===============')
    if timeout is not None:
      output('No pings in last %d sec.' % (timeout,))
    output('Dumping stack frames for all threads:')
    for name in sorted(tracebacks):
      output('Traceback for \'%s\':\n%s' % (name, tracebacks[name]))
    output('===========================================================')


class TaskChannel(object):
  """Queue of results of async task execution."""

  class Timeout(Exception):
    """Raised by 'pull' in case of timeout."""

  _ITEM_RESULT = object()
  _ITEM_EXCEPTION = object()
  _ITEM_DONE = object()

  def __init__(self):
    self._queue = Queue.Queue()

  def send_result(self, result):
    """Enqueues a result of task execution."""
    self._queue.put((self._ITEM_RESULT, result))

  def send_done(self):
    """Stops the iteration."""
    self._queue.put((self._ITEM_DONE, None))

  def send_exception(self, exc_info=None):
    """Enqueue an exception raised by a task.

    Arguments:
      exc_info: If given, should be 3-tuple returned by sys.exc_info(),
                default is current value of sys.exc_info(). Use default in
                'except' blocks to capture currently processed exception.
    """
    exc_info = exc_info or sys.exc_info()
    assert isinstance(exc_info, tuple) and len(exc_info) == 3
    # Transparently passing Timeout will break 'pull' contract, since a caller
    # has no way to figure out that's an exception from the task and not from
    # 'pull' itself. Transform Timeout into generic RuntimeError with
    # explanation.
    if isinstance(exc_info[1], TaskChannel.Timeout):
      exc_info = (
          RuntimeError,
          RuntimeError('Task raised Timeout exception'),
          exc_info[2])
    self._queue.put((self._ITEM_EXCEPTION, exc_info))

  def __iter__(self):
    return self

  def next(self, timeout=None):
    """Dequeues available result or exception.

    Args:
      timeout: if not None will block no longer than |timeout| seconds and will
          raise TaskChannel.Timeout exception if no results are available.

    Returns:
      Whatever task pushes to the queue by calling 'send_result'.

    Raises:
      TaskChannel.Timeout: waiting longer than |timeout|.
      Whatever exception task raises.
    """
    # Do not ever use timeout == None, in that case signal handlers are not
    # being called (at least on Python 2.7, http://bugs.python.org/issue8844).
    while True:
      try:
        item_type, value = self._queue.get(
            timeout=timeout if timeout is not None else 30.0)
        break
      except Queue.Empty:
        if timeout is None:
          continue
        raise TaskChannel.Timeout()
    if item_type == self._ITEM_RESULT:
      return value
    if item_type == self._ITEM_EXCEPTION:
      # 'value' is captured sys.exc_info() 3-tuple. Use extended raise syntax
      # to preserve stack frame of original exception (that was raised in
      # another thread).
      assert isinstance(value, tuple) and len(value) == 3
      raise value[0], value[1], value[2]
    if item_type == self._ITEM_DONE:
      raise StopIteration()
    assert False, 'Impossible queue item type: %r' % item_type

  def wrap_task(self, task):
    """Decorator that makes a function push results into this channel."""
    @functools.wraps(task)
    def wrapped(*args, **kwargs):
      try:
        self.send_result(task(*args, **kwargs))
      except Exception:
        self.send_exception()
    return wrapped


def num_processors():
  """Returns the number of processors.

  Python on OSX 10.6 raises a NotImplementedError exception.
  """
  try:
    # Multiprocessing
    import multiprocessing
    return multiprocessing.cpu_count()
  except:  # pylint: disable=W0702
    try:
      # Mac OS 10.6
      return int(os.sysconf('SC_NPROCESSORS_ONLN'))  # pylint: disable=E1101
    except:
      # Some of the windows builders seem to get here.
      return 4
