Fix #1497: Add rate limit to batch request result generation

interface/batch/util.py

@@ -908,6 +908,20 @@ def patch_request(request, batch_request):
         return general_server_error_response("Problem cancelling the batch request.")
 
 
+def request_already_generating(request_id):
+    """Check the cache and rabbitmq to see if there is a request for generating batch request results."""
+
+    try:
+        lock_id = redis_id.batch_results_request_lock.id.format(request_id)
+        if cache.get(lock_id):
+            log.debug("There is already a request for generating this batch request results.")
+            return True
+    except BaseException:
+        log.exception("Failed to check batch request results generating status.")
+
+    return False
+
+
 def get_request(request, batch_request, user):
     provide_progress = request.GET.get("progress")
     provide_progress = provide_progress and provide_progress.lower() == "true"
@@ -919,6 +933,15 @@ def get_request(request, batch_request, user):
         ).count()
         res["request"]["progress"] = f"{finished_domains}/{total_domains}"
         res["request"]["num_domains"] = total_domains
-    if batch_request.status == BatchRequestStatus.done and not batch_request.has_report_file():
+
+    if (
+        batch_request.status == BatchRequestStatus.done
+        and not batch_request.has_report_file()
+        and not request_already_generating(batch_request.request_id)
+    ):
         batch_async_generate_results.delay(user=user, batch_request=batch_request, site_url=get_site_url(request))
+
+        lock_id = redis_id.batch_results_request_lock.id.format(batch_request.request_id)
+        cache.add(lock_id, True)
+
     return api_response(res)

interface/redis_id.py

@@ -85,6 +85,9 @@
 # Lock for generating batch results
 batch_results_lock = REDIS_RECORD("batch:results:gen:{}:{}", None)
 
+# Lock for requesting generate batch results, rate limit 1/10min
+batch_results_request_lock = REDIS_RECORD("batch:results:req:{}", 60 * 10)
+
 # Lock for batch scheduler
 batch_scheduler_lock = REDIS_RECORD("batch:scheduler:lock", 60 * 5)
 

interface/test/batch/test_batch.py

@@ -1,9 +1,12 @@
 import json
 import uuid
-
-
-from checks.models import BatchRequestType, BatchUser
-from interface.batch.util import register_request
+from django.core.cache import cache
+from checks.models import BatchRequest, BatchRequestType, BatchUser, BatchRequestStatus
+from interface.batch.util import get_request, register_request
+from django.test.client import RequestFactory
+import pytest
+from interface.batch.util import batch_async_generate_results
+from interface import redis_id
 
 
 def test_convert_batch_request_type():
@@ -40,3 +43,40 @@ def test_register_batch_request(db):
 
 #  todo: add testcase where data is available and an API response is created.
 #  this requires a scan to be ran etc.
+
+
+@pytest.mark.withoutresponses
+def test_batch_request_result_generation(db, client, mocker):
+    """There can only be one result generate task in the queue per batch request id."""
+
+    request = RequestFactory().get("/")
+
+    # mock putting task on the queue as queue/worker dynamics makes testing unreliable
+    mocker.patch("interface.batch.util.batch_async_generate_results.delay")
+    batch_async_generate_results.delay.return_value = "123"
+
+    # create dummy batch request for testing
+    test_user = BatchUser.objects.create(username="test_user")
+    batch_request = BatchRequest.objects.create(user=test_user, name="test_batch_request", type=BatchRequestType.web)
+
+    get_request(request, batch_request, test_user)
+
+    #  batch_async_generate_results task should not be queued if the batch request is not done yet
+    assert batch_async_generate_results.delay.call_count == 0
+
+    batch_request.status = BatchRequestStatus.done
+    batch_request.save()
+
+    # if batch request is done, a batch_async_generate_results task should be put on the queue to generate the results
+    get_request(request, batch_request, test_user)
+    assert batch_async_generate_results.delay.call_count == 1
+
+    # there should not be an additional task put on the queue when one is already present
+    get_request(request, batch_request, test_user)
+    assert batch_async_generate_results.delay.call_count == 1
+
+    # if the cache expires a new batch_async_generate_results task can be added
+    lock_id = redis_id.batch_results_request_lock.id.format(batch_request.request_id)
+    cache.delete(lock_id)
+    get_request(request, batch_request, test_user)
+    assert batch_async_generate_results.delay.call_count == 2