quay/data/queue.py

from datetime import datetime, timedelta

from data.database import QueueItem, db, db_for_update
from util.morecollections import AttrDict


MINIMUM_EXTENSION = timedelta(seconds=20)

class NoopWith:
  def __enter__(self):
    pass

  def __exit__(self, type, value, traceback):
    pass

class MetricQueueReporter(object):
  def __init__(self, metric_queue):
    self._metric_queue = metric_queue

  def __call__(self, currently_processing, running_count, total_count):
    need_capacity_count = total_count - running_count
    self._metric_queue.put('BuildCapacityShortage', need_capacity_count, unit='Count')

    building_percent = 100 if currently_processing else 0
    self._metric_queue.put('PercentBuilding', building_percent, unit='Percent')

class WorkQueue(object):
  def __init__(self, queue_name, transaction_factory,
               canonical_name_match_list=None, reporter=None, metric_queue=None):
    self._queue_name = queue_name
    self._reporter = reporter
    self._metric_queue = metric_queue
    self._transaction_factory = transaction_factory
    self._currently_processing = False

    if canonical_name_match_list is None:
      self._canonical_name_match_list = []
    else:
      self._canonical_name_match_list = canonical_name_match_list

  @staticmethod
  def _canonical_name(name_list):
    return '/'.join(name_list) + '/'

  def _running_jobs(self, now, name_match_query):
    return (QueueItem
            .select(QueueItem.queue_name)
            .where(QueueItem.available == False,
                   QueueItem.processing_expires > now,
                   QueueItem.queue_name ** name_match_query))

  def _available_jobs(self, now, name_match_query):
    return self._available_jobs_where(QueueItem.select(), now, name_match_query)

  def _available_jobs_where(self, query, now, name_match_query):
    return query.where(QueueItem.queue_name ** name_match_query, QueueItem.available_after <= now,
                   ((QueueItem.available == True) | (QueueItem.processing_expires <= now)),
                   QueueItem.retries_remaining > 0)

  def _available_jobs_not_running(self, now, name_match_query, running_query):
    return (self
            ._available_jobs(now, name_match_query)
            .where(~(QueueItem.queue_name << running_query)))

  def _name_match_query(self):
    return '%s%%' % self._canonical_name([self._queue_name] + self._canonical_name_match_list)

  def _item_by_id_for_update(self, queue_id):
    return db_for_update(QueueItem.select().where(QueueItem.id == queue_id)).get()

  def get_metrics(self):
    now = datetime.utcnow()
    name_match_query = self._name_match_query()

    running_query = self._running_jobs(now, name_match_query)
    running_count = running_query.distinct().count()

    available_query = self._available_jobs(now, name_match_query)
    available_count = available_query.select(QueueItem.queue_name).distinct().count()

    available_not_running_query = self._available_jobs_not_running(now, name_match_query,
                                                                   running_query)
    available_not_running_count = (available_not_running_query.select(QueueItem.queue_name)
                                                              .distinct().count())

    return (running_count, available_not_running_count, available_count)

  def update_metrics(self):
    if self._reporter is None and self._metric_queue is None:
      return

    (running_count, available_not_running_count, available_count) = self.get_metrics()

    if self._metric_queue:
      dim = {'queue': self._queue_name}
      self._metric_queue.put('Running', running_count, dimensions=dim)
      self._metric_queue.put('AvailableNotRunning', available_not_running_count, dimensions=dim)
      self._metric_queue.put('Available', available_count, dimensions=dim)

    if self._reporter:
      self._reporter(self._currently_processing, running_count,
                     running_count + available_not_running_count)

  def has_retries_remaining(self, item_id):
    """ Returns whether the queue item with the given id has any retries remaining. If the
        queue item does not exist, returns False. """
    with self._transaction_factory(db):
      try:
        return QueueItem.get(id=item_id).retries_remaining > 0
      except QueueItem.DoesNotExist:
        return False

  def put(self, canonical_name_list, message, available_after=0, retries_remaining=5):
    """
    Put an item, if it shouldn't be processed for some number of seconds,
    specify that amount as available_after. Returns the ID of the queue item added.
    """

    params = {
      'queue_name': self._canonical_name([self._queue_name] + canonical_name_list),
      'body': message,
      'retries_remaining': retries_remaining,
    }

    available_date = datetime.utcnow() + timedelta(seconds=available_after or 0)
    params['available_after'] = available_date

    with self._transaction_factory(db):
      r = str(QueueItem.create(**params).id)
    if self._metric_queue:
      self._metric_queue.put('Added', 1, dimensions={'queue': self._queue_name})
    return r

  def get(self, processing_time=300):
    """
    Get an available item and mark it as unavailable for the default of five
    minutes. The result of this method must always be composed of simple
    python objects which are JSON serializable for network portability reasons.
    """
    now = datetime.utcnow()

    name_match_query = self._name_match_query()

    running = self._running_jobs(now, name_match_query)
    avail = self._available_jobs_not_running(now, name_match_query, running)

    item = None
    try:
      # The previous solution to this used a select for update in a
      # transaction to prevent multiple instances from processing the
      # same queue item. This suffered performance problems. This solution
      # instead has instances attempt to update the potential queue item to be
      # unavailable. However, since their update clause is restricted to items
      # that are available=False, only one instance's update will succeed, and
      # it will have a changed row count of 1. Instances that have 0 changed
      # rows know that another instance is already handling that item.
      db_item = avail.order_by(QueueItem.id).get()
      changed_query = (QueueItem.update(
          available=False,
          processing_expires=now + timedelta(seconds=processing_time),
          retries_remaining=QueueItem.retries_remaining-1,
        )
        .where(QueueItem.id == db_item.id))

      changed_query = self._available_jobs_where(changed_query, now, name_match_query)
      changed = changed_query.execute()
      if changed == 1:
        item = AttrDict({
          'id': db_item.id,
          'body': db_item.body,
          'retries_remaining': db_item.retries_remaining - 1,
        })
        self._currently_processing = True
    except QueueItem.DoesNotExist:
      self._currently_processing = False

    # Return a view of the queue item rather than an active db object
    return item

  def cancel(self, item_id):
    """ Attempts to cancel the queue item with the given ID from the queue. Returns true on success
        and false if the queue item could not be canceled. A queue item can only be canceled if
        if is available and has retries remaining.
    """

    with self._transaction_factory(db):
      # Load the build queue item for update.
      try:
        queue_item = db_for_update(QueueItem.select()
                                            .where(QueueItem.id == item_id)).get()
      except QueueItem.DoesNotExist:
        return False

      # Check the queue item.
      if not queue_item.available or queue_item.retries_remaining == 0:
        return False

      # Delete the queue item.
      queue_item.delete_instance(recursive=True)
      return True

  def complete(self, completed_item):
    with self._transaction_factory(db):
      try:
        completed_item_obj = self._item_by_id_for_update(completed_item.id)
      except QueueItem.DoesNotExist:
        self._currently_processing = False
        return

      completed_item_obj.delete_instance(recursive=True)
      self._currently_processing = False

  def incomplete(self, incomplete_item, retry_after=300, restore_retry=False):
    with self._transaction_factory(db):
      retry_date = datetime.utcnow() + timedelta(seconds=retry_after)

      try:
        incomplete_item_obj = self._item_by_id_for_update(incomplete_item.id)
        incomplete_item_obj.available_after = retry_date
        incomplete_item_obj.available = True

        if restore_retry:
          incomplete_item_obj.retries_remaining += 1

        incomplete_item_obj.save()
        self._currently_processing = False
        return incomplete_item_obj.retries_remaining > 0
      except QueueItem.DoesNotExist:
        return False

  def extend_processing(self, item, seconds_from_now, minimum_extension=MINIMUM_EXTENSION,
                        updated_data=None):
    with self._transaction_factory(db):
      try:
        queue_item = self._item_by_id_for_update(item.id)
        new_expiration = datetime.utcnow() + timedelta(seconds=seconds_from_now)
        has_change = False

        # Only actually write the new expiration to the db if it moves the expiration some minimum
        if new_expiration - queue_item.processing_expires > minimum_extension:
          queue_item.processing_expires = new_expiration
          has_change = True

        if updated_data is not None:
          queue_item.body = updated_data
          has_change = True

        if has_change:
          queue_item.save()

      except QueueItem.DoesNotExist:
        return