Merge pull request #29 from dabapps/queue-rate-limiting

Add rate limit to queues

Author: jordaneb

.python-version

@@ -0,0 +1 @@
+3.7.4

.travis.yml

@@ -1,9 +1,11 @@
 language: python
 sudo: false
+services:
+  - mysql
 python:
-- '2.7'
 - '3.4'
 - '3.6'
+- '3.7'
 env:
 - DJANGO_VERSION=1.8
 - DJANGO_VERSION=1.9

README.md

@@ -181,9 +181,10 @@ The `workspace` flag is optional. If supplied, it must be a valid JSON string.
 
 To start a worker:
 
-    manage.py worker [queue_name]
+    manage.py worker [queue_name] [--rate_limit]
 
-`queue_name` is optional, and will default to `default`
+- `queue_name` is optional, and will default to `default`
+- The `--rate_limit` flag is optional, and will default to `1`. It is the minimum number of seconds that must have elapsed before a subsequent job can be run.
 
 ## Testing
 

django_dbq/management/commands/worker.py

@@ -1,5 +1,6 @@
 from django.db import transaction
 from django.core.management.base import BaseCommand, CommandError
+from django.utils import timezone
 from django.utils.module_loading import import_string
 from django_dbq.models import Job
 from simplesignals.process import WorkerProcessBase
@@ -58,13 +59,19 @@ class Worker(WorkerProcessBase):
 
     process_title = "jobworker"
 
-    def __init__(self, name):
+    def __init__(self, name, rate_limit_in_seconds):
         self.queue_name = name
+        self.rate_limit_in_seconds = rate_limit_in_seconds
+        self.last_job_finished = None
         super(Worker, self).__init__()
 
     def do_work(self):
         sleep(1)
+        if self.last_job_finished and (timezone.now() - self.last_job_finished).total_seconds() < self.rate_limit_in_seconds:
+            return
+
         process_job(self.queue_name)
+        self.last_job_finished = timezone.now()
 
 
 class Command(BaseCommand):
@@ -73,6 +80,7 @@ class Command(BaseCommand):
 
     def add_arguments(self, parser):
         parser.add_argument('queue_name', nargs='?', default='default', type=str)
+        parser.add_argument('rate_limit', help='The rate limit in seconds. The default rate limit is 1 job per second.', nargs='?', default=1, type=int)
         parser.add_argument(
             '--dry-run',
             action='store_true',
@@ -89,10 +97,11 @@ def handle(self, *args, **options):
             raise CommandError("Please supply a single queue job name")
 
         queue_name = options['queue_name']
+        rate_limit_in_seconds = options['rate_limit']
 
-        self.stdout.write("Starting job worker for queue \"%s\"" % queue_name)
+        self.stdout.write("Starting job worker for queue \"%s\" with rate limit %s/s" % (queue_name, rate_limit_in_seconds))
 
-        worker = Worker(queue_name)
+        worker = Worker(queue_name, rate_limit_in_seconds)
 
         if options['dry_run']:
             return

django_dbq/tests.py

@@ -1,8 +1,13 @@
 from datetime import datetime, timedelta
+import mock
+
+import freezegun
 from django.core.management import call_command, CommandError
 from django.test import TestCase
 from django.test.utils import override_settings
-from django_dbq.management.commands.worker import process_job
+from django.utils import timezone
+
+from django_dbq.management.commands.worker import process_job, Worker
 from django_dbq.models import Job
 try:
     from StringIO import StringIO
@@ -77,6 +82,39 @@ def test_worker_with_queue_name(self):
         self.assertTrue('test_queue' in output)
 
 
+@freezegun.freeze_time()
+@mock.patch('django_dbq.management.commands.worker.sleep')
+@mock.patch('django_dbq.management.commands.worker.process_job')
+class WorkerProcessDoWorkTestCase(TestCase):
+
+    def setUp(self):
+        super().setUp()
+        self.MockWorker = mock.MagicMock()
+        self.MockWorker.queue_name = 'default'
+        self.MockWorker.rate_limit_in_seconds = 5
+        self.MockWorker.last_job_finished = None
+
+    def test_do_work_no_previous_job_run(self, mock_process_job, mock_sleep):
+        Worker.do_work(self.MockWorker)
+        self.assertEqual(mock_sleep.call_count, 1)
+        self.assertEqual(mock_process_job.call_count, 1)
+        self.assertEqual(self.MockWorker.last_job_finished, timezone.now())
+
+    def test_do_work_previous_job_too_soon(self, mock_process_job, mock_sleep):
+        self.MockWorker.last_job_finished = timezone.now() - timezone.timedelta(seconds=2)
+        Worker.do_work(self.MockWorker)
+        self.assertEqual(mock_sleep.call_count, 1)
+        self.assertEqual(mock_process_job.call_count, 0)
+        self.assertEqual(self.MockWorker.last_job_finished, timezone.now() - timezone.timedelta(seconds=2))
+
+    def test_do_work_previous_job_long_time_ago(self, mock_process_job, mock_sleep):
+        self.MockWorker.last_job_finished = timezone.now() - timezone.timedelta(seconds=7)
+        Worker.do_work(self.MockWorker)
+        self.assertEqual(mock_sleep.call_count, 1)
+        self.assertEqual(mock_process_job.call_count, 1)
+        self.assertEqual(self.MockWorker.last_job_finished, timezone.now())
+
+
 @override_settings(JOBS={'testjob': {'tasks': ['a']}})
 class JobTestCase(TestCase):
 

test-requirements.txt

@@ -1,2 +1,4 @@
 -r requirements.txt
 pymysql==0.7.11
+freezegun==0.3.12
+mock==3.0.5