Add a queue capacity reporter plugin to the queue. Move the queue definitions to app. Add a cloudwatch reporter to the dockerfile build queue.

data/queue.py

@@ -11,18 +11,53 @@ MINIMUM_EXTENSION = timedelta(seconds=20)
 
 
 class WorkQueue(object):
-  def __init__(self, queue_name, canonical_name_match_list=None):
-    self.queue_name = queue_name
+  def __init__(self, queue_name, canonical_name_match_list=None, reporter=None):
+    self._queue_name = queue_name
+    self._reporter = reporter
 
     if canonical_name_match_list is None:
-      self.canonical_name_match_list = []
+      self._canonical_name_match_list = []
     else:
-      self.canonical_name_match_list = canonical_name_match_list
+      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 _name_match_query(self):
+    return '%s%%' % self._canonical_name([self._queue_name] + self._canonical_name_match_list)
+
+  def _report_queue_metrics(self):
+    if self._reporter is None:
+      return
+
+    now = datetime.now()
+    name_match_query = self._name_match_query()
+
+    total_jobs = (QueueItem
+      .select(QueueItem.queue_name)
+      .where(QueueItem.queue_name ** name_match_query,
+             QueueItem.available_after <= now,
+             ((QueueItem.available == True) | (QueueItem.processing_expires > now) |
+              (QueueItem.retries_remaining > 0)))
+      .distinct()
+      .count())
+
+    running = self._running_jobs(now, name_match_query).distinct().count()
+
+    self._reporter(running, total_jobs)
+
+  def update_metrics(self):
+    with transaction_factory(db):
+      self._report_queue_metrics()
+
   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,
@@ -30,7 +65,7 @@ class WorkQueue(object):
     """
 
     params = {
-      'queue_name': self._canonical_name([self.queue_name] + canonical_name_list),
+      'queue_name': self._canonical_name([self._queue_name] + canonical_name_list),
       'body': message,
       'retries_remaining': retries_remaining,
     }
@@ -39,7 +74,9 @@ class WorkQueue(object):
       available_date = datetime.now() + timedelta(seconds=available_after)
       params['available_after'] = available_date
 
-    QueueItem.create(**params)
+    with transaction_factory(db):
+      QueueItem.create(**params)
+      self._report_queue_metrics()
 
   def get(self, processing_time=300):
     """
@@ -48,15 +85,10 @@ class WorkQueue(object):
     """
     now = datetime.now()
 
-    name_match_query = '%s%%' % self._canonical_name([self.queue_name] +
-                                                     self.canonical_name_match_list)
+    name_match_query = self._name_match_query()
 
     with transaction_factory(db):
-      running = (QueueItem
-        .select(QueueItem.queue_name)
-        .where(QueueItem.available == False,
-               QueueItem.processing_expires > now,
-               QueueItem.queue_name ** name_match_query))
+      running = self._running_jobs(now, name_match_query)
 
       avail = QueueItem.select().where(QueueItem.queue_name ** name_match_query,
                                        QueueItem.available_after <= now,
@@ -67,6 +99,8 @@ class WorkQueue(object):
 
       found = list(avail.limit(1).order_by(QueueItem.id))
 
+      item = None
+
       if found:
         item = found[0]
         item.available = False
@@ -74,24 +108,26 @@ class WorkQueue(object):
         item.retries_remaining -= 1
         item.save()
 
-        return item
+      self._report_queue_metrics()
 
-      return None
+      return item
 
-  @staticmethod
-  def complete(completed_item):
-    completed_item.delete_instance()
+  def complete(self, completed_item):
+    with transaction_factory(db):
+      completed_item.delete_instance()
+      self._report_queue_metrics()
 
-  @staticmethod
-  def incomplete(incomplete_item, retry_after=300, restore_retry=False):
-    retry_date = datetime.now() + timedelta(seconds=retry_after)
-    incomplete_item.available_after = retry_date
-    incomplete_item.available = True
+  def incomplete(self, incomplete_item, retry_after=300, restore_retry=False):
+    with transaction_factory(db):
+      retry_date = datetime.now() + timedelta(seconds=retry_after)
+      incomplete_item.available_after = retry_date
+      incomplete_item.available = True
 
-    if restore_retry:
-      incomplete_item.retries_remaining += 1
+      if restore_retry:
+        incomplete_item.retries_remaining += 1
 
-    incomplete_item.save()
+      incomplete_item.save()
+      self._report_queue_metrics()
 
   @staticmethod
   def extend_processing(queue_item, seconds_from_now):
@@ -101,8 +137,3 @@ class WorkQueue(object):
     if new_expiration - queue_item.processing_expires > MINIMUM_EXTENSION:
       queue_item.processing_expires = new_expiration
       queue_item.save()
-
-
-image_diff_queue = WorkQueue(app.config['DIFFS_QUEUE_NAME'])
-dockerfile_build_queue = WorkQueue(app.config['DOCKERFILE_BUILD_QUEUE_NAME'])
-webhook_queue = WorkQueue(app.config['WEBHOOK_QUEUE_NAME'])