From 9ace8542907552740bb2b9bbaabc4bcbfdcd87f8 Mon Sep 17 00:00:00 2001
From: Stan Triepels <1939656+GDay@users.noreply.github.com>
Date: Sat, 15 Oct 2022 01:56:41 +0200
Subject: [PATCH] Save limit per group/func/name (#12)

* feat: option to save tasks per group/func/name

* test: update tests for save limit

* fix: convert func when save-limit checking

when passed as function it needs to be converted to work

* Add warning if option is not valid

Co-authored-by: Noortheen Raja <jnoortheen@gmail.com>
---
 django_q/cluster.py            | 45 ++++++++++++++++++----------------
 django_q/conf.py               | 11 ++++++---
 django_q/tests/test_cluster.py | 41 +++++++++++++++++++++++++++++++
 docs/configure.rst             |  8 ++++++
 4 files changed, 80 insertions(+), 25 deletions(-)

diff --git a/django_q/cluster.py b/django_q/cluster.py
index ee084e0..2715182 100644
--- a/django_q/cluster.py
+++ b/django_q/cluster.py
@@ -457,6 +457,16 @@ def worker(
                 break
     logger.info(_(f"{proc_name} stopped doing work"))
 
+def get_func_repr(func):
+    # convert func to string
+    if inspect.isfunction(func):
+        return f"{func.__module__}.{func.__name__}"
+    elif inspect.ismethod(func):
+        return (
+            f"{func.__self__.__module__}."
+            f"{func.__self__.__name__}.{func.__name__}"
+        )
+    return func
 
 def save_task(task, broker: Broker):
     """
@@ -478,20 +488,19 @@ def save_task(task, broker: Broker):
         )
     # SAVE LIMIT > 0: Prune database, SAVE_LIMIT 0: No pruning
     close_old_django_connections()
+
     try:
-        if task["success"]:
-            # first apply per group success history limit
-            if "group" in task:
-                with db.transaction.atomic():
-                    qs = Success.objects.filter(group=task["group"])
-                    last = qs.select_for_update().last()
-                    if Conf.SAVE_LIMIT_PER_GROUP <= qs.count():
-                        last.delete()
-            # then apply global success history limit
-            with db.transaction.atomic():
-                last = Success.objects.select_for_update().last()
-                if Conf.SAVE_LIMIT <= Success.objects.count():
-                    last.delete()
+        filters = {}
+        if Conf.SAVE_LIMIT_PER and Conf.SAVE_LIMIT_PER in {"group", "name", "func"} and Conf.SAVE_LIMIT_PER in task:
+            value = task[Conf.SAVE_LIMIT_PER]
+            if Conf.SAVE_LIMIT_PER == "func":
+                value = get_func_repr(value)
+            filters[Conf.SAVE_LIMIT_PER] = value
+        with db.transaction.atomic():
+            last = Success.objects.filter(**filters).select_for_update().last()
+            if task["success"] and 0 < Conf.SAVE_LIMIT <= Success.objects.filter(**filters).count():
+                last.delete()
+
         # check if this task has previous results
         try:
             existing_task = Task.objects.get(id=task["id"], name=task["name"])
@@ -508,16 +517,10 @@ def save_task(task, broker: Broker):
                 and existing_task.attempt_count >= Conf.MAX_ATTEMPTS
             ):
                 broker.acknowledge(task["ack_id"])
+
         except Task.DoesNotExist:
-            func = task["func"]
             # convert func to string
-            if inspect.isfunction(func):
-                func = f"{func.__module__}.{func.__name__}"
-            elif inspect.ismethod(func):
-                func = (
-                    f"{func.__self__.__module__}."
-                    f"{func.__self__.__name__}.{func.__name__}"
-                )
+            func = get_func_repr(task["func"])
             Task.objects.create(
                 id=task["id"],
                 name=task["name"],
diff --git a/django_q/conf.py b/django_q/conf.py
index 6436eaf..f9ecd23 100644
--- a/django_q/conf.py
+++ b/django_q/conf.py
@@ -86,9 +86,12 @@ class Conf:
     # Failures are always saved
     SAVE_LIMIT = conf.get("save_limit", 250)
 
-    # Maximum number of successful tasks of the same group kept in the database. 0 saves everything. -1 saves none
-    # Failures are always saved
-    SAVE_LIMIT_PER_GROUP = conf.get("save_limit_per_group", 5)
+    # save-limit can be set per Task's "group" or "name" or "func"
+    SAVE_LIMIT_PER = conf.get("save_limit_per", None)
+
+    # Verify SAVE_LIMIT_PER is valid
+    if SAVE_LIMIT_PER not in ["group", "name", "func", None]:
+        warn(f"SAVE_LIMIT_PER ({SAVE_LIMIT_PER}) is not a valid option. Options are: 'group', 'name', 'func' and None. Default is None.")
 
     # Guard loop sleep in seconds. Should be between 0 and 60 seconds.
     GUARD_CYCLE = conf.get("guard_cycle", 0.5)
@@ -143,7 +146,7 @@ class Conf:
         warn(
             """Retry and timeout are misconfigured. Set retry larger than timeout,
         failure to do so will cause the tasks to be retriggered before completion.
-        See https://django-q.readthedocs.io/en/latest/configure.html#retry for details."""
+        See https://django-q2.readthedocs.io/en/master/configure.html#retry for details."""
         )
 
     # Sets the amount of tasks the cluster will try to pop off the broker.
diff --git a/django_q/tests/test_cluster.py b/django_q/tests/test_cluster.py
index 1755932..52bcc25 100644
--- a/django_q/tests/test_cluster.py
+++ b/django_q/tests/test_cluster.py
@@ -407,6 +407,47 @@ def test_recycle(broker, monkeypatch):
     assert Success.objects.count() == Conf.SAVE_LIMIT
     broker.delete_queue()
 
+@pytest.mark.django_db
+def test_save_limit_per_func(broker, monkeypatch):
+    # set up the Sentinel
+    broker.list_key = "test_recycle_test:q"
+    async_task("django_q.tests.tasks.hello", broker=broker)
+    async_task("django_q.tests.tasks.countdown", 2, broker=broker)
+    async_task("django_q.tests.tasks.multiply", 2, 2, broker=broker)
+    start_event = Event()
+    stop_event = Event()
+    cluster_id = uuidlib.uuid4()
+    task_queue = Queue()
+    result_queue = Queue()
+    # override settings
+    monkeypatch.setattr(Conf, "RECYCLE", 3)
+    monkeypatch.setattr(Conf, "WORKERS", 1)
+    # set a timer to stop the Sentinel
+    threading.Timer(3, stop_event.set).start()
+    for i in range(3):
+        pusher(task_queue, stop_event, broker=broker)
+    worker(task_queue, result_queue, Value("f", -1))
+    s = Sentinel(stop_event, start_event, cluster_id=cluster_id, broker=broker)
+    assert start_event.is_set()
+    assert s.status() == Conf.STOPPED
+    # worker should exit on recycle
+    # check if the work has been done
+    assert result_queue.qsize() == 3
+    # save_limit test
+    monkeypatch.setattr(Conf, "SAVE_LIMIT", 1)
+    monkeypatch.setattr(Conf, "SAVE_LIMIT_PER", "func")
+    result_queue.put("STOP")
+    # run monitor
+    monitor(result_queue)
+    assert Success.objects.count() == 3
+    assert set(Success.objects.filter().values_list('func', flat=True)) == {
+        'django_q.tests.tasks.countdown',
+        'django_q.tests.tasks.hello',
+        'django_q.tests.tasks.multiply',
+    }
+    broker.delete_queue()
+
+
 
 @pytest.mark.django_db
 def test_max_rss(broker, monkeypatch):
diff --git a/docs/configure.rst b/docs/configure.rst
index 69fb955..d5473bd 100644
--- a/docs/configure.rst
+++ b/docs/configure.rst
@@ -142,6 +142,14 @@ Limits the amount of successful tasks saved to Django.
  - Defaults to ``250``
  - Failures are always saved.
 
+save_limit_per
+~~~~~~~~~~~~~
+
+The above ``save_limit`` for successful tasks can be fine tuned per task type using
+ - Set to ``"group"`` to store the tasks per group
+ - Other possible values are ``"func"``, ``"name"``, ``None``
+ - Defaults to ``None``
+
 guard_cycle
 ~~~~~~~~~~~