Test and improve concurrency (#69)

Fixes #33

Also fixes and a potential deadlocking issue.

Author: RealOrangeOne

README.md

@@ -114,9 +114,6 @@ TASKS = {
 
 Finally, you can run `manage.py db_worker` to run tasks as they're created. Check the `--help` for more options.
 
-> [!CAUTION]
-> The database backend does not work with SQLite when you are running multiple worker processes - tasks may be executed more than once. See [#33](https://github.com/RealOrangeOne/django-tasks/issues/33).
-
 ### Retrieving task result
 
 When enqueueing a task, you get a `TaskResult`, however it may be useful to retrieve said result from somewhere else (another request, another task etc). This can be done with `get_result` (or `aget_result`):

django_tasks/backends/base.py

@@ -1,6 +1,6 @@
 from abc import ABCMeta, abstractmethod
 from inspect import iscoroutinefunction
-from typing import Any, List, TypeVar
+from typing import Any, Iterable, TypeVar
 
 from asgiref.sync import sync_to_async
 from django.core.checks.messages import CheckMessage
@@ -101,7 +101,7 @@ async def aget_result(self, result_id: str) -> TaskResult:
             result_id=result_id
         )
 
-    def check(self, **kwargs: Any) -> List[CheckMessage]:
+    def check(self, **kwargs: Any) -> Iterable[CheckMessage]:
         raise NotImplementedError(
             "subclasses may provide a check() method to verify that task "
             "backend is configured correctly."

django_tasks/backends/database/backend.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, List, TypeVar
+from typing import TYPE_CHECKING, Any, Iterable, TypeVar
 
 from django.apps import apps
 from django.core.checks import ERROR, CheckMessage
 from django.core.exceptions import ValidationError
+from django.db import connections, router
 from typing_extensions import ParamSpec
 
 from django_tasks.backends.base import BaseTaskBackend
@@ -81,16 +82,26 @@ async def aget_result(self, result_id: str) -> TaskResult:
         except (DBTaskResult.DoesNotExist, ValidationError) as e:
             raise ResultDoesNotExist(result_id) from e
 
-    def check(self, **kwargs: Any) -> List[CheckMessage]:
-        if not apps.is_installed("django_tasks.backends.database"):
-            backend_name = self.__class__.__name__
+    def check(self, **kwargs: Any) -> Iterable[CheckMessage]:
+        from .models import DBTaskResult
 
-            return [
-                CheckMessage(
-                    ERROR,
-                    f"{backend_name} configured as django_tasks backend, but database app not installed",
-                    "Insert 'django_tasks.backends.database' in INSTALLED_APPS",
-                )
-            ]
+        backend_name = self.__class__.__name__
 
-        return []
+        if not apps.is_installed("django_tasks.backends.database"):
+            yield CheckMessage(
+                ERROR,
+                f"{backend_name} configured as django_tasks backend, but database app not installed",
+                "Insert 'django_tasks.backends.database' in INSTALLED_APPS",
+            )
+
+        db_connection = connections[router.db_for_read(DBTaskResult)]
+        if (
+            db_connection.vendor == "sqlite"
+            and hasattr(db_connection, "transaction_mode")
+            and db_connection.transaction_mode != "EXCLUSIVE"
+        ):
+            yield CheckMessage(
+                ERROR,
+                f"{backend_name} is using SQLite non-exclusive transactions",
+                f"Set settings.DATABASES[{db_connection.alias!r}]['OPTIONS']['transaction_mode'] to 'EXCLUSIVE'",
+            )

django_tasks/backends/database/management/commands/db_worker.py

@@ -7,11 +7,12 @@
 from typing import List, Optional
 
 from django.core.management.base import BaseCommand
-from django.db import transaction
+from django.db.utils import OperationalError
 
 from django_tasks import DEFAULT_TASK_BACKEND_ALIAS, tasks
 from django_tasks.backends.database.backend import DatabaseBackend
 from django_tasks.backends.database.models import DBTaskResult
+from django_tasks.backends.database.utils import exclusive_transaction
 from django_tasks.exceptions import InvalidTaskBackendError
 from django_tasks.task import DEFAULT_QUEUE_NAME, ResultStatus
 
@@ -68,8 +69,16 @@ def start(self) -> None:
 
                 # During this transaction, all "ready" tasks are locked. Therefore, it's important
                 # it be as efficient as possible.
-                with transaction.atomic():
-                    task_result = tasks.get_locked()
+                with exclusive_transaction(tasks.db):
+                    try:
+                        task_result = tasks.get_locked()
+                    except OperationalError as e:
+                        # Ignore locked databases and keep trying.
+                        # It should unlock eventually.
+                        if "database is locked" in e.args[0]:
+                            task_result = None
+                        else:
+                            raise
 
                     if task_result is not None:
                         # "claim" the task, so it isn't run by another worker process

django_tasks/backends/database/models.py

@@ -20,6 +20,8 @@
 )
 from django_tasks.utils import exception_to_dict, retry
 
+from .utils import normalize_uuid
+
 logger = logging.getLogger("django_tasks.backends.database")
 
 T = TypeVar("T")
@@ -62,7 +64,7 @@ def get_locked(self) -> Optional["DBTaskResult"]:
         """
         Get a job, locking the row and accounting for deadlocks.
         """
-        return self.select_for_update().first()
+        return self.select_for_update(skip_locked=True).first()
 
 
 class DBTaskResult(GenericBase[P, T], models.Model):
@@ -127,7 +129,7 @@ def task_result(self) -> "TaskResult[T]":
         result = TaskResult[T](
             db_result=self,
             task=self.task,
-            id=str(self.id),
+            id=normalize_uuid(self.id),
             status=ResultStatus[self.status],
             enqueued_at=self.enqueued_at,
             started_at=self.started_at,

django_tasks/backends/database/utils.py

@@ -0,0 +1,42 @@
+from contextlib import contextmanager
+from typing import Any, Generator, Optional, Union
+from uuid import UUID
+
+from django.db import transaction
+
+
+@contextmanager
+def exclusive_transaction(using: Optional[str] = None) -> Generator[Any, Any, Any]:
+    """
+    Wrapper around `transaction.atomic` which ensures transactions on SQLite are exclusive.
+
+    This functionality is built-in to Django 5.1+.
+    """
+    connection = transaction.get_connection(using)
+
+    if (
+        connection.vendor == "sqlite"
+        and getattr(connection, "transaction_mode", None) != "EXCLUSIVE"
+    ):
+        with connection.cursor() as c:
+            c.execute("BEGIN EXCLUSIVE")
+            try:
+                yield
+            finally:
+                c.execute("COMMIT")
+    else:
+        with transaction.atomic(using=using):
+            yield
+
+
+def normalize_uuid(val: Union[str, UUID]) -> str:
+    """
+    Normalize a UUID into its dashed representation.
+
+    This works around engines like MySQL which don't store values in a uuid field,
+    and thus drops the dashes.
+    """
+    if isinstance(val, str):
+        val = UUID(val)
+
+    return str(val)

django_tasks/checks.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Sequence
+from typing import Any, Iterable, Sequence
 
 from django.apps.config import AppConfig
 from django.core.checks.messages import CheckMessage
@@ -8,16 +8,11 @@
 
 def check_tasks(
     app_configs: Sequence[AppConfig] = None, **kwargs: Any
-) -> List[CheckMessage]:
+) -> Iterable[CheckMessage]:
     """Checks all registered task backends."""
 
-    errors = []
     for backend in tasks.all():
         try:
-            backend_errors = backend.check()
+            yield from backend.check()
         except NotImplementedError:
             pass
-        else:
-            errors.extend(backend_errors)
-
-    return errors

tests/settings.py

@@ -5,7 +5,7 @@
 
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-IN_TEST = sys.argv[0] == "test"
+IN_TEST = sys.argv[1] == "test"
 
 ALLOWED_HOSTS = ["*"]