feat(task-processor): Add priority support

api/audit/tasks.py

@@ -9,18 +9,19 @@
 )
 from audit.models import AuditLog, RelatedObjectType
 from task_processor.decorators import register_task_handler
+from task_processor.models import TaskPriority
 
 logger = logging.getLogger(__name__)
 
 
-@register_task_handler()
+@register_task_handler(priority=TaskPriority.HIGHEST)
 def create_feature_state_went_live_audit_log(feature_state_id: int):
     _create_feature_state_audit_log_for_change_request(
         feature_state_id, FEATURE_STATE_WENT_LIVE_MESSAGE
     )
 
 
-@register_task_handler()
+@register_task_handler(priority=TaskPriority.HIGHEST)
 def create_feature_state_updated_by_change_request_audit_log(feature_state_id: int):
     _create_feature_state_audit_log_for_change_request(
         feature_state_id, FEATURE_STATE_UPDATED_BY_CHANGE_REQUEST_MESSAGE
@@ -57,7 +58,7 @@ def _create_feature_state_audit_log_for_change_request(
     )
 
 
-@register_task_handler()
+@register_task_handler(priority=TaskPriority.HIGHEST)
 def create_audit_log_from_historical_record(
     history_instance_id: int,
     history_user_id: typing.Optional[int],

api/edge_api/identities/edge_request_forwarder.py

@@ -7,14 +7,15 @@
 
 from environments.dynamodb.migrator import IdentityMigrator
 from task_processor.decorators import register_task_handler
+from task_processor.models import TaskPriority
 
 
 def _should_forward(project_id: int) -> bool:
     migrator = IdentityMigrator(project_id)
     return bool(migrator.is_migration_done)
 
 
-@register_task_handler(queue_size=2000)
+@register_task_handler(queue_size=2000, priority=TaskPriority.LOW)
 def forward_identity_request(
     request_method: str,
     headers: dict,
@@ -35,7 +36,7 @@ def forward_identity_request(
     requests.get(url, params=query_params, headers=headers, timeout=5)
 
 
-@register_task_handler(queue_size=2000)
+@register_task_handler(queue_size=2000, priority=TaskPriority.LOW)
 def forward_trait_request(
     request_method: str,
     headers: dict,
@@ -61,7 +62,7 @@ def forward_trait_request_sync(
     )
 
 
-@register_task_handler(queue_size=1000)
+@register_task_handler(queue_size=1000, priority=TaskPriority.LOW)
 def forward_trait_requests(
     request_method: str,
     headers: str,

api/edge_api/identities/tasks.py

@@ -6,6 +6,7 @@
 from environments.models import Environment, Webhook
 from features.models import Feature, FeatureState
 from task_processor.decorators import register_task_handler
+from task_processor.models import TaskPriority
 from users.models import FFAdminUser
 from webhooks.webhooks import WebhookEventType, call_environment_webhooks
 
@@ -71,7 +72,7 @@ def call_environment_webhook_for_feature_state_change(
     call_environment_webhooks(environment, data, event_type=event_type)
 
 
-@register_task_handler()
+@register_task_handler(priority=TaskPriority.HIGHEST)
 def sync_identity_document_features(identity_uuid: str):
     from .models import EdgeIdentity
 

api/environments/tasks.py

@@ -6,17 +6,18 @@
     send_environment_update_message_for_project,
 )
 from task_processor.decorators import register_task_handler
+from task_processor.models import TaskPriority
 
 
-@register_task_handler()
+@register_task_handler(priority=TaskPriority.HIGH)
 def rebuild_environment_document(environment_id: int):
     wrapper = DynamoEnvironmentWrapper()
     if wrapper.is_enabled:
         environment = Environment.objects.get(id=environment_id)
         wrapper.write_environment(environment)
 
 
-@register_task_handler()
+@register_task_handler(priority=TaskPriority.HIGHEST)
 def process_environment_update(audit_log_id: int):
     audit_log = AuditLog.objects.get(id=audit_log_id)
 

api/task_processor/decorators.py

@@ -9,14 +9,18 @@
 from django.utils import timezone
 
 from task_processor.exceptions import InvalidArgumentsError, TaskQueueFullError
-from task_processor.models import RecurringTask, Task
+from task_processor.models import RecurringTask, Task, TaskPriority
 from task_processor.task_registry import register_task
 from task_processor.task_run_method import TaskRunMethod
 
 logger = logging.getLogger(__name__)
 
 
-def register_task_handler(task_name: str = None, queue_size: int = None):
+def register_task_handler(
+    task_name: str = None,
+    queue_size: int = None,
+    priority: TaskPriority = TaskPriority.NORMAL,
+):
     def decorator(f: typing.Callable):
         nonlocal task_name
 
@@ -61,7 +65,7 @@ def delay(
                     logger.warning(e)
                     return
 
-                task.save()
+                task.save(priority=priority)
                 return task
 
         def run_in_thread(*, args: typing.Tuple = (), kwargs: typing.Dict = None):

api/task_processor/migrations/0008_add_get_task_to_process_function.py

@@ -16,15 +16,15 @@ class Migration(migrations.Migration):
             os.path.join(
                 os.path.dirname(__file__),
                 "sql",
-                "get_tasks_to_process.sql",
+                "0008_get_tasks_to_process.sql",
             ),
             reverse_sql="DROP FUNCTION IF EXISTS get_tasks_to_process",
         ),
         PostgresOnlyRunSQL.from_sql_file(
             os.path.join(
                 os.path.dirname(__file__),
                 "sql",
-                "get_recurring_tasks_to_process.sql",
+                "0008_get_recurring_tasks_to_process.sql",
             ),
             reverse_sql="DROP FUNCTION IF EXISTS get_recurringtasks_to_process",
         ),

api/task_processor/migrations/0010_task_priority.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.2.20 on 2023-10-13 06:04
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('task_processor', '0009_add_recurring_task_run_first_run_at'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='task',
+            name='priority',
+            field=models.PositiveSmallIntegerField(choices=[(100, 'Lower'), (75, 'Low'), (50, 'Normal'), (25, 'High'), (0, 'Highest')], default=None, null=True),
+        ),
+    ]

api/task_processor/migrations/0011_add_priority_to_get_tasks_to_process.py

@@ -0,0 +1,23 @@
+# Generated by Django 3.2.20 on 2023-10-13 04:44
+
+from django.db import migrations
+
+from core.migration_helpers import PostgresOnlyRunSQL
+import os
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("task_processor", "0010_task_priority"),
+    ]
+
+    operations = [
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0011_get_tasks_to_process.sql",
+            ),
+            reverse_sql="DROP FUNCTION IF EXISTS get_tasks_to_process",
+        ),
+    ]

api/task_processor/migrations/sql/0011_get_tasks_to_process.sql

@@ -0,0 +1,30 @@
+CREATE OR REPLACE FUNCTION get_tasks_to_process(num_tasks integer)
+RETURNS SETOF task_processor_task AS $$
+DECLARE
+    row_to_return task_processor_task;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_task
+        WHERE num_failures < 3 AND scheduled_for < NOW() AND completed = FALSE AND is_locked = FALSE
+        ORDER BY priority ASC, scheduled_for ASC, created_at ASC
+        LIMIT num_tasks
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_task
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the `is_locked` column here, the client will receive the row that is actually locked but has the `is_locked` value set to `False`.
+        row_to_return.is_locked := TRUE;
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+

api/task_processor/models.py

@@ -11,6 +11,14 @@
 from task_processor.task_registry import registered_tasks
 
 
+class TaskPriority(models.IntegerChoices):
+    LOWER = 100
+    LOW = 75
+    NORMAL = 50
+    HIGH = 25
+    HIGHEST = 0
+
+
 class AbstractBaseTask(models.Model):
     uuid = models.UUIDField(unique=True, default=uuid.uuid4)
     created_at = models.DateTimeField(auto_now_add=True)
@@ -74,6 +82,9 @@ class Task(AbstractBaseTask):
     num_failures = models.IntegerField(default=0)
     completed = models.BooleanField(default=False)
     objects = TaskManager()
+    priority = models.PositiveSmallIntegerField(
+        default=None, null=True, choices=TaskPriority.choices
+    )
 
     class Meta:
         # We have customised the migration in 0004 to only apply this change to postgres databases
@@ -90,12 +101,14 @@ class Meta:
     def create(
         cls,
         task_identifier: str,
+        priority: TaskPriority = TaskPriority.NORMAL,
         *,
         args: typing.Tuple[typing.Any] = None,
         kwargs: typing.Dict[str, typing.Any] = None,
     ) -> "Task":
         return Task(
             task_identifier=task_identifier,
+            priority=priority,
             serialized_args=cls.serialize_data(args or tuple()),
             serialized_kwargs=cls.serialize_data(kwargs or dict()),
         )

api/tests/unit/task_processor/test_unit_task_processor_processor.py

@@ -15,6 +15,7 @@
     RecurringTask,
     RecurringTaskRun,
     Task,
+    TaskPriority,
     TaskResult,
     TaskRun,
 )
@@ -248,26 +249,39 @@ def test_run_task_does_nothing_if_no_tasks(db):
 
 
 @pytest.mark.django_db(transaction=True)
-def test_run_task_runs_tasks_in_correct_order():
+def test_run_task_runs_tasks_in_correct_priority():
     # Given
     # 2 tasks
     task_1 = Task.create(
-        _create_organisation.task_identifier, args=("task 1 organisation",)
+        _create_organisation.task_identifier,
+        args=("task 1 organisation",),
+        priority=TaskPriority.HIGH,
     )
     task_1.save()
 
     task_2 = Task.create(
-        _create_organisation.task_identifier, args=("task 2 organisation",)
+        _create_organisation.task_identifier,
+        args=("task 2 organisation",),
+        priority=TaskPriority.HIGH,
     )
     task_2.save()
 
+    task_3 = Task.create(
+        _create_organisation.task_identifier,
+        args=("task 3 organisation",),
+        priority=TaskPriority.HIGHEST,
+    )
+    task_3.save()
+
     # When
     task_runs_1 = run_tasks()
     task_runs_2 = run_tasks()
+    task_runs_3 = run_tasks()
 
     # Then
-    assert task_runs_1[0].task == task_1
-    assert task_runs_2[0].task == task_2
+    assert task_runs_1[0].task == task_3
+    assert task_runs_2[0].task == task_1
+    assert task_runs_3[0].task == task_2
 
 
 @pytest.mark.django_db(transaction=True)