Add a worker for reporting global stats to Prometheus

Fixes #1789

data/model/organization.py

@@ -137,6 +137,10 @@ def get_organizations():
   return User.select().where(User.organization == True, User.robot == False)
 
 
+def get_active_org_count():
+  return get_organizations().count()
+
+
 def add_user_as_admin(user_obj, org_obj):
   try:
     admin_role = TeamRole.get(name='admin')

data/model/repository.py

@@ -16,6 +16,9 @@ from data.database import (Repository, Namespace, RepositoryTag, Star, Image, Us
 
 logger = logging.getLogger(__name__)
 
+def get_repository_count():
+  return Repository.select().count()
+
 
 def get_public_repo_visibility():
   return _basequery.get_public_repo_visibility()

data/model/user.py

@@ -643,6 +643,10 @@ def get_active_user_count():
   return get_active_users().count()
 
 
+def get_robot_count():
+  return User.select().where(User.robot == True).count()
+
+
 def detach_external_login(user, service_name):
   try:
     service = LoginService.get(name=service_name)

util/metrics/metricqueue.py

@@ -57,6 +57,11 @@ class MetricQueue(object):
                                                           labelnames=['namespace', 'repo_name',
                                                                       'status'])
 
+    self.repository_count = prom.create_gauge('repository_count', 'Number of repositories')
+    self.user_count = prom.create_gauge('user_count', 'Number of users')
+    self.org_count = prom.create_gauge('org_count', 'Number of Organizations')
+    self.robot_count = prom.create_gauge('robot_count', 'Number of robot accounts')
+
     # Deprecated: Define an in-memory queue for reporting metrics to CloudWatch or another
     # provider.
     self._queue = None

workers/globalpromstats.py

@@ -0,0 +1,56 @@
+import logging
+import time
+
+from app import app, metric_queue
+from data.database import UseThenDisconnect
+from data import model
+from util.locking import GlobalLock, LockNotAcquiredException
+from workers.worker import Worker
+
+logger = logging.getLogger(__name__)
+
+WORKER_FREQUENCY = app.config.get('GLOBAL_PROMETHEUS_STATS_FREQUENCY', 60 * 60)
+
+
+class GlobalPrometheusStatsWorker(Worker):
+  """ Worker which reports global stats (# of users, orgs, repos, etc) to Prometheus periodically.
+  """
+  def __init__(self):
+    super(GlobalPrometheusStatsWorker, self).__init__()
+    self.add_operation(self._try_report_stats, WORKER_FREQUENCY)
+
+  def _try_report_stats(self):
+    logger.debug('Attempting to report stats')
+
+    try:
+      with GlobalLock('GLOBAL_PROM_STATS'):
+        self._report_stats()
+    except LockNotAcquiredException:
+      return
+
+  def _report_stats(self):
+    logger.debug('Reporting global stats')
+    with UseThenDisconnect(app.config):
+      # Repository count.
+      metric_queue.repository_count.set(model.repository.get_repository_count())
+
+      # User counts.
+      metric_queue.user_count.set(model.user.get_active_user_count())
+      metric_queue.org_count.set(model.organization.get_active_org_count())
+      metric_queue.robot_count.set(model.user.get_robot_count())
+
+
+def main():
+  logging.config.fileConfig('conf/logging_debug.conf', disable_existing_loggers=False)
+
+  if not app.config.get('PROMETHEUS_AGGREGATOR_URL'):
+    logger.debug('Prometheus not enabled; skipping global stats reporting')
+    while True:
+      time.sleep(100000)
+
+  worker = GlobalPrometheusStatsWorker()
+  worker.start()
+
+
+if __name__ == "__main__":
+  main()