From df7f65003ffd78ca5df2f87512787fd0ac0a9a21 Mon Sep 17 00:00:00 2001
From: Joel <oldsneerjaw@gmail.com>
Date: Fri, 14 Aug 2020 08:03:36 -0700
Subject: [PATCH] Output request queue time for Django (#55)

HireFire includes a strategy for scaling web dynos based on Heroku request queue
times. This change adds Django support for queue time scaling according to the
official docs: https://help.hirefire.io/article/49-logplex-queue-time.
---
 README.rst                               | 17 ++++++++++
 hirefire/contrib/django/middleware.py    | 43 ++++++++++++++++++++++++
 tests/contrib/django/test_middleware.py  |  6 ++++
 tests/contrib/django/testapp/settings.py |  1 +
 4 files changed, 67 insertions(+)

diff --git a/README.rst b/README.rst
index b38773a8..fc1d4cc3 100644
--- a/README.rst
+++ b/README.rst
@@ -147,6 +147,23 @@ Setting up HireFire support for Django is easy:
 
    .. _`application page`: https://manager.hirefire.io/applications
 
+#. Add ``'hirefire.contrib.django.middleware.QueueTimeMiddleware'`` to your
+   ``MIDDLEWARE`` setting to enable HireFire's `support`_ for scaling
+   according to Heroku request queue times (optional).
+
+   .. code-block:: python
+
+     # Use ``MIDDLEWARE_CLASSES`` prior to Django 1.10
+     MIDDLEWARE = [
+         'hirefire.contrib.django.middleware.HireFireMiddleware',
+         # ...
+     ]
+
+   Make sure to place it before any other item in the list/tuple so that
+   request queue time is calculated as accurately as possible.
+
+   .. _`support`: https://help.hirefire.io/article/49-logplex-queue-time
+
 #. Check that the middleware has been correctly setup by opening the
    following URL in a browser::
 
diff --git a/hirefire/contrib/django/middleware.py b/hirefire/contrib/django/middleware.py
index 750c2f3b..401482b9 100644
--- a/hirefire/contrib/django/middleware.py
+++ b/hirefire/contrib/django/middleware.py
@@ -2,10 +2,12 @@
 
 import os
 import re
+from logging import getLogger
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.http import HttpResponse, JsonResponse
+from django.utils import timezone
 
 try:
     # Django >= 1.10
@@ -24,6 +26,8 @@ def setting(name, default=None):
     return os.environ.get(name, getattr(settings, name, None)) or default
 
 
+logger = getLogger('hirefire')
+
 TOKEN = setting('HIREFIRE_TOKEN', 'development')
 PROCS = setting('HIREFIRE_PROCS', [])
 USE_CONCURRENCY = setting('HIREFIRE_USE_CONCURRENCY', False)
@@ -86,3 +90,42 @@ def process_request(self, request):
 
         elif self.info_path.match(path):
             return self.info(request)
+
+
+class QueueTimeMiddleware(MiddlewareMixin):
+    """
+    The Django middleware that outputs Heroku request queue times to stdout
+
+    The behavior and format for request queue time scaling are described in
+    HireFire's docs: https://help.hirefire.io/article/49-logplex-queue-time
+    """
+    def process_request(self, request):
+        """
+        Outputs the Heroku request queue time to stdout
+        """
+        request_start_header_value = request.META.get("HTTP_X_REQUEST_START")
+        if request_start_header_value is not None:
+            try:
+                request_start_timestamp_ms = int(request_start_header_value)
+            except ValueError:
+                logger.warning(
+                    'Received an invalid "X-Request-Start" header value from '
+                    'Heroku: "%s"',
+                    request_start_header_value,
+                )
+            else:
+                # The timestamp generated by timestamp() is the number of
+                # seconds since the Unix epoch, so it needs to be multiplied
+                # and truncated as an integer to convert it to milliseconds
+                now_timestamp_ms = int(timezone.now().timestamp() * 1000)
+
+                # There may be some clock drift between the Heroku router and
+                # the dyno that is running this middleware; if the calculated
+                # queue time is negative, treat it as zero instead
+                queue_time_ms = (
+                    now_timestamp_ms - request_start_timestamp_ms
+                    if now_timestamp_ms >= request_start_timestamp_ms
+                    else 0
+                )
+
+                print("[hirefire:router] queue={}ms".format(queue_time_ms))
diff --git a/tests/contrib/django/test_middleware.py b/tests/contrib/django/test_middleware.py
index 9164a04f..8e662325 100644
--- a/tests/contrib/django/test_middleware.py
+++ b/tests/contrib/django/test_middleware.py
@@ -9,3 +9,9 @@ def test_token(self, client, settings):
 
         response = client.get('/hirefire/not-the-token-%s/info' % settings.HIREFIRE_TOKEN)
         assert response.status_code == 404
+
+
+class TestQueueTimeMiddleware:
+    def test_queue_time(self, client):
+        response = client.get('/hirefire/test', HTTP_X_REQUEST_START='946733845303')
+        assert response.status_code == 200
diff --git a/tests/contrib/django/testapp/settings.py b/tests/contrib/django/testapp/settings.py
index c1eee272..d6fc130f 100644
--- a/tests/contrib/django/testapp/settings.py
+++ b/tests/contrib/django/testapp/settings.py
@@ -40,6 +40,7 @@
 ]
 
 MIDDLEWARE = [
+    'hirefire.contrib.django.middleware.QueueTimeMiddleware',
     'django.middleware.security.SecurityMiddleware',
     'django.contrib.sessions.middleware.SessionMiddleware',
     'django.middleware.common.CommonMiddleware',