#!/usr/bin/env python

#

# Copyright 2007 Google Inc.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#     http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""Provides thread-pool-like functionality for workers accessing App Engine.

The pool adapts to slow or timing out requests by reducing the number of

active workers, or increasing the number when requests latency reduces.

"""

import logging

import Queue

import sys

import threading

import time

import traceback

from google.appengine.tools.requeue import ReQueue

logger = logging.getLogger('google.appengine.tools.adaptive_thread_pool')

_THREAD_SHOULD_EXIT = '_THREAD_SHOULD_EXIT'

INITIAL_BACKOFF = 1.0

BACKOFF_FACTOR = 2.0

class Error(Exception):

  """Base-class for exceptions in this module."""

class WorkItemError(Error):

  """Error while processing a WorkItem."""

class RetryException(Error):

  """A non-fatal exception that indicates that a work item should be retried."""

def InterruptibleSleep(sleep_time):

  """Puts thread to sleep, checking this threads exit_flag four times a second.

  Args:

    sleep_time: Time to sleep.

  """

  slept = 0.0

  epsilon = .0001

  thread = threading.currentThread()

  while slept < sleep_time - epsilon:

    remaining = sleep_time - slept

    this_sleep_time = min(remaining, 0.25)

    time.sleep(this_sleep_time)

    slept += this_sleep_time

    if thread.exit_flag:

      return

class WorkerThread(threading.Thread):

  """A WorkerThread to execute WorkItems.

  Attributes:

    exit_flag: A boolean indicating whether this thread should stop

      its work and exit.

  """

  def __init__(self, thread_pool, thread_gate, name=None):

    """Initialize a WorkerThread instance.

    Args:

      thread_pool: An AdaptiveThreadPool instance.

      thread_gate: A ThreadGate instance.

      name: A name for this WorkerThread.

    """

    threading.Thread.__init__(self)

    self.setDaemon(True)

    self.exit_flag = False

    self.__error = None

    self.__traceback = None

    self.__thread_pool = thread_pool

    self.__work_queue = thread_pool.requeue

    self.__thread_gate = thread_gate

    if not name:

      self.__name = 'Anonymous_' + self.__class__.__name__

    else:

      self.__name = name

  def run(self):

    """Perform the work of the thread."""

    logger.debug('[%s] %s: started', self.getName(), self.__class__.__name__)

    try:

      self.WorkOnItems()

    except:

      self.SetError()

    logger.debug('[%s] %s: exiting', self.getName(), self.__class__.__name__)

  def SetError(self):

    """Sets the error and traceback information for this thread.

    This must be called from an exception handler.

    """

    if not self.__error:

      exc_info = sys.exc_info()

      self.__error = exc_info[1]

      self.__traceback = exc_info[2]

      logger.exception('[%s] %s:', self.getName(), self.__class__.__name__)

  def WorkOnItems(self):

    """Perform the work of a WorkerThread."""

    while not self.exit_flag:

      item = None

      self.__thread_gate.StartWork()

      try:

        status, instruction = WorkItem.FAILURE, ThreadGate.DECREASE

        try:

          if self.exit_flag:

            instruction = ThreadGate.HOLD

            break

          try:

            item = self.__work_queue.get(block=True, timeout=1.0)

          except Queue.Empty:

            instruction = ThreadGate.HOLD

            continue

          if item == _THREAD_SHOULD_EXIT or self.exit_flag:

            status, instruction = WorkItem.SUCCESS, ThreadGate.HOLD

            break

          logger.debug('[%s] Got work item %s', self.getName(), item)

          status, instruction = item.PerformWork(self.__thread_pool)

        except RetryException:

          status, instruction = WorkItem.RETRY, ThreadGate.HOLD

        except:

          self.SetError()

          raise

      finally:

        try:

          if item:

            if status == WorkItem.SUCCESS:

              self.__work_queue.task_done()

            elif status == WorkItem.RETRY:

              try:

                self.__work_queue.reput(item, block=False)

              except Queue.Full:

                logger.error('[%s] Failed to reput work item.', self.getName())

                raise Error('Failed to reput work item')

            else:

              if not self.__error:

                if item.error:

                  self.__error = item.error

                  self.__traceback = item.traceback

                else:

                  self.__error = WorkItemError(

                      'Fatal error while processing %s' % item)

                raise self.__error

        finally:

          self.__thread_gate.FinishWork(instruction=instruction)

  def CheckError(self):

    """If an error is present, then log it."""

    if self.__error:

      logger.error('Error in %s: %s', self.getName(), self.__error)

      if self.__traceback:

        logger.debug('%s', ''.join(traceback.format_exception(

            self.__error.__class__,

            self.__error,

            self.__traceback)))

  def __str__(self):

    return self.__name

class AdaptiveThreadPool(object):

  """A thread pool which processes WorkItems from a queue.

  Attributes:

    requeue: The requeue instance which holds work items for this

      thread pool.

  """

  def __init__(self,

               num_threads,

               queue_size=None,

               base_thread_name=None,

               worker_thread_factory=WorkerThread,

               queue_factory=Queue.Queue):

    """Initialize an AdaptiveThreadPool.

    An adaptive thread pool executes WorkItems using a number of

    WorkerThreads.  WorkItems represent items of work that may

    succeed, soft fail, or hard fail. In addition, a completed work

    item can signal this AdaptiveThreadPool to enable more or fewer

    threads.  Initially one thread is active.  Soft failures are

    reqeueud to be retried.  Hard failures cause this

    AdaptiveThreadPool to shut down entirely.  See the WorkItem class

    for more details.

    Args:

      num_threads: The number of threads to use.

      queue_size: The size of the work item queue to use.

      base_thread_name: A string from which worker thread names are derived.

      worker_thread_factory: A factory which procudes WorkerThreads.

      queue_factory: Used for dependency injection.

    """

    if queue_size is None:

      queue_size = num_threads

    self.requeue = ReQueue(queue_size, queue_factory=queue_factory)

    self.__thread_gate = ThreadGate(num_threads)

    self.__num_threads = num_threads

    self.__threads = []

    for i in xrange(num_threads):

      thread = worker_thread_factory(self, self.__thread_gate)

      if base_thread_name:

        base = base_thread_name

      else:

        base = thread.__class__.__name__

      thread.name = '%s-%d' % (base, i)

      self.__threads.append(thread)

      thread.start()

  def num_threads(self):

    """Return the number of threads in this thread pool."""

    return self.__num_threads

  def Threads(self):

    """Yields the registered threads."""

    for thread in self.__threads:

      yield thread

  def SubmitItem(self, item, block=True, timeout=0.0):

    """Submit a WorkItem to the AdaptiveThreadPool.

    Args:

      item: A WorkItem instance.

      block: Whether to block on submitting if the submit queue is full.

      timeout: Time wait for room in the queue if block is True, 0.0 to

        block indefinitely.

    Raises:

      Queue.Full if the submit queue is full.

    """

    self.requeue.put(item, block=block, timeout=timeout)

  def QueuedItemCount(self):

    """Returns the number of items currently in the queue."""

    return self.requeue.qsize()

  def Shutdown(self):

    """Shutdown the thread pool.

    Tasks may remain unexecuted in the submit queue.

    """

    while not self.requeue.empty():

      try:

        unused_item = self.requeue.get_nowait()

        self.requeue.task_done()

      except Queue.Empty:

        pass

    for thread in self.__threads:

      thread.exit_flag = True

      self.requeue.put(_THREAD_SHOULD_EXIT)

    self.__thread_gate.EnableAllThreads()

  def Wait(self):

    """Wait until all work items have been completed."""

    self.requeue.join()

  def JoinThreads(self):

    """Wait for all threads to exit."""

    for thread in self.__threads:

      logger.debug('Waiting for %s to exit' % str(thread))

      thread.join()

  def CheckErrors(self):

    """Output logs for any errors that occurred in the worker threads."""

    for thread in self.__threads:

      thread.CheckError()

class ThreadGate(object):

  """Manage the number of active worker threads.

  The ThreadGate limits the number of threads that are simultaneously

  active in order to implement adaptive rate control.

  Initially the ThreadGate allows only one thread to be active.  For

  each successful work item, another thread is activated and for each

  failed item, the number of active threads is reduced by one.  When only

  one thread is active, failures will cause exponential backoff.

  For example, a ThreadGate instance, thread_gate can be used in a number

  of threads as so:

  # Block until this thread is enabled for work.

  thread_gate.StartWork()

  try:

    status = DoSomeWorkInvolvingLimitedSharedResources()

    suceeded = IsStatusGood(status)

    badly_failed = IsStatusVeryBad(status)

  finally:

    if suceeded:

      # Suceeded, add more simultaneously enabled threads to the task.

      thread_gate.FinishWork(instruction=ThreadGate.INCREASE)

    elif badly_failed:

      # Failed, or succeeded but with high resource load, reduce number of

      # workers.

      thread_gate.FinishWork(instruction=ThreadGate.DECREASE)

    else:

      # We succeeded, but don't want to add more workers to the task.

      thread_gate.FinishWork(instruction=ThreadGate.HOLD)

  the thread_gate will enable and disable/backoff threads in response to

  resource load conditions.

  StartWork can block indefinitely. FinishWork, while not

  lock-free, should never block absent a demonic scheduler.

  """

  INCREASE = 'increase'

  HOLD = 'hold'

  DECREASE = 'decrease'

  def __init__(self,

               num_threads,

               sleep=InterruptibleSleep):

    """Constructor for ThreadGate instances.

    Args:

      num_threads: The total number of threads using this gate.

      sleep: Used for dependency injection.

    """

    self.__enabled_count = 1

    self.__lock = threading.Lock()

    self.__thread_semaphore = threading.Semaphore(self.__enabled_count)

    self.__num_threads = num_threads

    self.__backoff_time = 0

    self.__sleep = sleep

  def num_threads(self):

    return self.__num_threads

  def EnableThread(self):

    """Enable one more worker thread."""

    self.__lock.acquire()

    try:

      self.__enabled_count += 1

    finally:

      self.__lock.release()

    self.__thread_semaphore.release()

  def EnableAllThreads(self):

    """Enable all worker threads."""

    for unused_idx in xrange(self.__num_threads - self.__enabled_count):

      self.EnableThread()

  def StartWork(self):

    """Starts a critical section in which the number of workers is limited.

    Starts a critical section which allows self.__enabled_count

    simultaneously operating threads. The critical section is ended by

    calling self.FinishWork().

    """

    self.__thread_semaphore.acquire()

    if self.__backoff_time > 0.0:

      if not threading.currentThread().exit_flag:

        logger.info('Backing off due to errors: %.1f seconds',

                    self.__backoff_time)

        self.__sleep(self.__backoff_time)

  def FinishWork(self, instruction=None):

    """Ends a critical section started with self.StartWork()."""

    if not instruction or instruction == ThreadGate.HOLD:

      self.__thread_semaphore.release()

    elif instruction == ThreadGate.INCREASE:

      if self.__backoff_time > 0.0:

        logger.info('Resetting backoff to 0.0')

        self.__backoff_time = 0.0

      do_enable = False

      self.__lock.acquire()

      try:

        if self.__num_threads > self.__enabled_count:

          do_enable = True

          self.__enabled_count += 1

      finally:

        self.__lock.release()

      if do_enable:

        logger.debug('Increasing active thread count to %d',

                     self.__enabled_count)

        self.__thread_semaphore.release()

      self.__thread_semaphore.release()

    elif instruction == ThreadGate.DECREASE:

      do_disable = False

      self.__lock.acquire()

      try:

        if self.__enabled_count > 1:

          do_disable = True

          self.__enabled_count -= 1

        else:

          if self.__backoff_time == 0.0:

            self.__backoff_time = INITIAL_BACKOFF

          else:

            self.__backoff_time *= BACKOFF_FACTOR

      finally:

        self.__lock.release()

        if do_disable:

          logger.debug('Decreasing the number of active threads to %d',

                       self.__enabled_count)

        else:

          self.__thread_semaphore.release()

class WorkItem(object):

  """Holds a unit of work."""

  SUCCESS = 'success'

  RETRY = 'retry'

  FAILURE = 'failure'

  def __init__(self, name):

    self.__name = name

  def PerformWork(self, thread_pool):

    """Perform the work of this work item and report the results.

    Args:

      thread_pool: The AdaptiveThreadPool instance associated with this

        thread.

    Returns:

      A tuple (status, instruction) of the work status and an instruction

      for the ThreadGate.

    """

    raise NotImplementedError

  def __str__(self):

    return self.__name