Fix: replace stale mail-fetch overlap check with a self-expiring lock (#14189)

src/paperless_mail/tasks.py

@@ -1,9 +1,10 @@
 import logging
+from typing import Final
 
 from celery import Task
 from celery import shared_task
+from django.core.cache import cache
 
-from documents.models import PaperlessTask
 from paperless_mail.mail import MailAccountHandler
 from paperless_mail.mail import MailError
 from paperless_mail.models import MailAccount
@@ -11,43 +12,45 @@ from paperless_mail.models import MailRule
 
 logger = logging.getLogger("paperless.mail.tasks")
 
+# Cache-backed lock guarding overlapping mail-account processing runs; unlike
+# a PaperlessTask row, it self-heals if the owning worker dies mid-run.
+MAIL_FETCH_LOCK_KEY: Final = "paperless_mail_fetch_lock"
+# Ceiling on how long a run may hold the lock; renewed after each account.
+MAIL_FETCH_LOCK_TTL: Final = 30 * 60
+
 
 @shared_task(bind=True)
 def process_mail_accounts(self: Task, account_ids: list[int] | None = None) -> str:
-    # A scheduled check can still be running (or queued) when the next one
-    # ProcessedMail dedup only records a message once its
-    # handling has finished, so an overlapping run can still pick up the same
-    # not-yet-recorded message. Skip outright rather than race it.
-    other_mail_fetch_running = (
-        PaperlessTask.objects.filter(
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            status__in=[PaperlessTask.Status.PENDING, PaperlessTask.Status.STARTED],
-        )
-        .exclude(task_id=self.request.id)
-        .exists()
-    )
-    if other_mail_fetch_running:
+    if not cache.add(MAIL_FETCH_LOCK_KEY, self.request.id, timeout=MAIL_FETCH_LOCK_TTL):
         logger.info(
             "Mail account processing is already running; skipping this run.",
         )
         return "Skipped: mail account processing already in progress."
 
-    total_new_documents = 0
-    accounts = (
-        MailAccount.objects.filter(pk__in=account_ids)
-        if account_ids
-        else MailAccount.objects.all()
-    )
-    for account in accounts:
-        if not MailRule.objects.filter(account=account, enabled=True).exists():
-            logger.info(f"No rules enabled for account {account}. Skipping.")
-            continue
-        try:
-            total_new_documents += MailAccountHandler().handle_mail_account(account)
-        except MailError:
-            logger.exception(f"Error while processing mail account {account}")
+    try:
+        total_new_documents = 0
+        accounts = (
+            MailAccount.objects.filter(pk__in=account_ids)
+            if account_ids
+            else MailAccount.objects.all()
+        )
+        for account in accounts:
+            if not MailRule.objects.filter(account=account, enabled=True).exists():
+                logger.info(f"No rules enabled for account {account}. Skipping.")
+                continue
+            try:
+                total_new_documents += MailAccountHandler().handle_mail_account(
+                    account,
+                )
+            except MailError:
+                logger.exception(f"Error while processing mail account {account}")
+            # Renew the lock so a run still genuinely in progress doesn't
+            # lose it to the TTL partway through a long account list.
+            cache.touch(MAIL_FETCH_LOCK_KEY, MAIL_FETCH_LOCK_TTL)
 
-    if total_new_documents > 0:
-        return f"Added {total_new_documents} document(s)."
-    else:
-        return "No new documents were added."
+        if total_new_documents > 0:
+            return f"Added {total_new_documents} document(s)."
+        else:
+            return "No new documents were added."
+    finally:
+        cache.delete(MAIL_FETCH_LOCK_KEY)

src/paperless_mail/tests/test_process_mail_accounts_overlap.py

@@ -2,9 +2,8 @@ from typing import Final
 
 import pytest
 import pytest_mock
+from django.core.cache import cache
 
-from documents.models import PaperlessTask
-from documents.tests.factories import PaperlessTaskFactory
 from paperless_mail import tasks
 from paperless_mail.tests.factories import MailAccountFactory
 from paperless_mail.tests.factories import MailRuleFactory
@@ -13,6 +12,13 @@ NO_DOCUMENTS_ADDED: Final = "No new documents were added."
 SKIPPED: Final = "Skipped: mail account processing already in progress."
 
 
+@pytest.fixture(autouse=True)
+def _clear_mail_fetch_lock():
+    cache.delete(tasks.MAIL_FETCH_LOCK_KEY)
+    yield
+    cache.delete(tasks.MAIL_FETCH_LOCK_KEY)
+
+
 @pytest.mark.django_db
 @pytest.mark.usefixtures("account_with_rule")
 class TestProcessMailAccountsOverlap:
@@ -22,50 +28,20 @@ class TestProcessMailAccountsOverlap:
         account = MailAccountFactory.create()
         MailRuleFactory.create(account=account, enabled=True)
 
-    @pytest.mark.parametrize(
-        ("status", "expected_result", "expected_call_count"),
-        [
-            pytest.param(
-                PaperlessTask.Status.PENDING,
-                SKIPPED,
-                0,
-                id="pending-task-blocks",
-            ),
-            pytest.param(
-                PaperlessTask.Status.STARTED,
-                SKIPPED,
-                0,
-                id="started-task-blocks",
-            ),
-            pytest.param(
-                PaperlessTask.Status.SUCCESS,
-                NO_DOCUMENTS_ADDED,
-                1,
-                id="finished-task-does-not-block",
-            ),
-        ],
-    )
-    def test_skips_only_while_another_mail_fetch_task_runs(
+    def test_skips_while_lock_is_held(
         self,
         mocker: pytest_mock.MockerFixture,
-        status: PaperlessTask.Status,
-        expected_result: str,
-        expected_call_count: int,
     ) -> None:
         """
         GIVEN:
             - An enabled mail account with a rule
-            - Another mail fetch task row in the given status
+            - The mail fetch lock is already held by another run
         WHEN:
             - Mail accounts are processed
         THEN:
-            - Processing is skipped only if that other task is pending or running
+            - Processing is skipped and no account is handled
         """
-        PaperlessTaskFactory.create(
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
-            status=status,
-        )
+        cache.add(tasks.MAIL_FETCH_LOCK_KEY, "other-task-id", timeout=60)
 
         mocked_handle = mocker.patch.object(
             tasks.MailAccountHandler,
@@ -75,21 +51,21 @@ class TestProcessMailAccountsOverlap:
 
         result = tasks.process_mail_accounts()
 
-        assert mocked_handle.call_count == expected_call_count
-        assert result == expected_result
+        assert mocked_handle.call_count == 0
+        assert result == SKIPPED
 
-    def test_runs_when_no_other_mail_fetch_task_exists(
+    def test_runs_when_lock_is_free(
         self,
         mocker: pytest_mock.MockerFixture,
     ) -> None:
         """
         GIVEN:
             - An enabled mail account with a rule
-            - No other mail fetch task rows
+            - No mail fetch run currently holds the lock
         WHEN:
             - Mail accounts are processed
         THEN:
-            - The account is handled
+            - The account is handled and the lock is released afterwards
         """
         mocked_handle = mocker.patch.object(
             tasks.MailAccountHandler,
@@ -101,34 +77,56 @@ class TestProcessMailAccountsOverlap:
 
         mocked_handle.assert_called_once()
         assert result == NO_DOCUMENTS_ADDED
+        assert cache.get(tasks.MAIL_FETCH_LOCK_KEY) is None
 
-    def test_does_not_skip_due_to_its_own_task_row(
+    def test_releases_lock_even_if_handling_raises(
         self,
         mocker: pytest_mock.MockerFixture,
     ) -> None:
         """
         GIVEN:
             - An enabled mail account with a rule
-            - A running mail fetch task row belonging to this very task
+            - Handling the account raises an unexpected exception
         WHEN:
-            - Mail accounts are processed under that task id
+            - Mail accounts are processed
         THEN:
-            - The task does not skip itself and handles the account
+            - The lock is still released so the next run is not blocked forever
         """
-        PaperlessTaskFactory.create(
-            task_id="self-task-id",
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
-            status=PaperlessTask.Status.STARTED,
+        mocker.patch.object(
+            tasks.MailAccountHandler,
+            "handle_mail_account",
+            side_effect=RuntimeError("boom"),
         )
 
+        with pytest.raises(RuntimeError):
+            tasks.process_mail_accounts()
+
+        assert cache.get(tasks.MAIL_FETCH_LOCK_KEY) is None
+
+    def test_recovers_after_lock_ttl_expires(
+        self,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - An enabled mail account with a rule
+            - A lock left behind by a run that never released it (e.g. the
+              worker was killed mid-run) but whose TTL has since expired
+        WHEN:
+            - Mail accounts are processed
+        THEN:
+            - The account is handled instead of being permanently blocked
+        """
+        cache.add(tasks.MAIL_FETCH_LOCK_KEY, "dead-task-id", timeout=1)
+        cache.delete(tasks.MAIL_FETCH_LOCK_KEY)  # simulate TTL expiry
+
         mocked_handle = mocker.patch.object(
             tasks.MailAccountHandler,
             "handle_mail_account",
             return_value=0,
         )
 
-        result = tasks.process_mail_accounts.apply(task_id="self-task-id").result
+        result = tasks.process_mail_accounts()
 
         mocked_handle.assert_called_once()
         assert result == NO_DOCUMENTS_ADDED