Merge branch 'FC-6202' into 'prep-release'

Fullctl Running Task Cancellation

See merge request fullctl/fullctl!317

Author: Matt Griswold

src/fullctl/django/models/concrete/tasks.py

@@ -22,6 +22,7 @@
 import fullctl.service_bridge.aaactl as aaactl
 import fullctl.service_bridge.auditctl as auditctl
 from fullctl.django.models.abstract.base import HandleRefModel
+from fullctl.django.tasks.context import task_execution_context
 from fullctl.django.tasks.qualifiers import Dynamic
 from fullctl.django.tasks.util import worker_id
 
@@ -32,6 +33,7 @@
     "TaskLimitError",
     "TaskAlreadyStarted",
     "ParentTaskNotFinished",
+    "TaskCancelledException",
     "Task",
     "TaskClaim",
     "TaskHeartbeat",
@@ -52,9 +54,13 @@ class TaskClaimed(IOError):
     def __init__(self, task):
         super().__init__(f"Task already claimed by another worker: {task}")
 
+
 class TaskScheduleClaimed(IOError):
     def __init__(self, task_schedule):
-        super().__init__(f"Task schedule already claimed by another worker: {task_schedule}")
+        super().__init__(
+            f"Task schedule already claimed by another worker: {task_schedule}"
+        )
+
 
 class WorkerUnqualified(IOError):
     def __init__(self, task, qualifier):
@@ -116,6 +122,25 @@ class ParentTaskNotFinished(IOError):
     pass
 
 
+class TaskCancelledException(IOError):
+    """
+    Raised when a task is cancelled during execution.
+
+    This allows graceful termination of task execution. The task execution
+    framework catches this exception and:
+    - Does NOT mark the task as failed
+    - Preserves the cancelled status
+    - Does NOT trigger error notifications
+
+    This exception should be raised by task.check_cancelled() when it
+    detects that the task status has been set to "cancelled".
+    """
+
+    def __init__(self, task):
+        super().__init__(f"Task was cancelled: {task.id}")
+        self.task = task
+
+
 class ErrorNotificationConfig(pydantic.BaseModel):
     subject: str
     message: str
@@ -520,6 +545,34 @@ def _cancel(self, reason):
         self.status = "cancelled"
         self.save()
 
+    def check_cancelled(self):
+        """
+        Check if the task has been cancelled and raise TaskCancelledException if so.
+
+        This method should be called periodically during long-running task execution
+        to allow for graceful cancellation. It refreshes the task status from the
+        database to check for cancellation requests.
+
+        Raises:
+            TaskCancelledException: If the task status is 'cancelled'
+
+        Usage example:
+            def run(self, *args, **kwargs):
+                for item in large_collection:
+                    self.check_cancelled()  # Check before processing each item
+                    process(item)
+
+        Note:
+            In most cases, you should use check_task_cancelled() from
+            fullctl.django.tasks.context instead of calling this directly.
+            The context function works anywhere in the call stack.
+        """
+        # Refresh status from database to get latest value
+        self.refresh_from_db(fields=["status"])
+
+        if self.status == "cancelled":
+            raise TaskCancelledException(self)
+
     def _complete(self, output):
 
         # if result is a dict we need to json encode it
@@ -605,11 +658,29 @@ def _run(self):
         try:
             param = self.param
             extensions.call(self, "before_run", *param["args"], **param["kwargs"])
-            output = self.run(*param["args"], **param["kwargs"])
+
+            # Wrap task execution in context so task can be accessed
+            # from anywhere in the call stack via check_task_cancelled()
+            with task_execution_context(self):
+                output = self.run(*param["args"], **param["kwargs"])
+
             extensions.call(self, "after_run", result=output)
             t_end = time.time()
             self.time = t_end - t_start
             self._complete(output)
+        except TaskCancelledException:
+            # Task was cancelled during execution - this is expected behavior
+            # Status is already set to "cancelled" by the cancel() call
+            # We don't call _fail() to avoid marking it as an error
+            t_end = time.time()
+            self.time = t_end - t_start
+            self.save()
+            log.info(
+                "Task cancelled during execution",
+                task_id=self.id,
+                task_op=self.op,
+                execution_time=self.time,
+            )
         except Exception as exc:
             self._fail(traceback.format_exc(), exc)
 
@@ -763,7 +834,7 @@ def spawn_tasks(self):
 
         if self.are_limited_tasks_pending():
             return []
-        
+
         # try to create a claim for the schedule
         try:
             schedule_claim = TaskScheduleClaim.objects.create(
@@ -775,7 +846,9 @@ def spawn_tasks(self):
             raise TaskScheduleClaimed(self)
 
         # clear old claims
-        TaskScheduleClaim.objects.filter(task_schedule=self).exclude(id=schedule_claim.id).delete()
+        TaskScheduleClaim.objects.filter(task_schedule=self).exclude(
+            id=schedule_claim.id
+        ).delete()
 
         for task in self.tasks.all():
             if task.status in ["pending", "running"]:
@@ -799,10 +872,9 @@ def spawn_tasks(self):
                 self.status = "deactivated"
                 self.save()
 
-
-
         return tasks
 
+
 class TaskScheduleClaim(HandleRefModel):
     task_schedule = models.ForeignKey(TaskSchedule, on_delete=models.CASCADE)
     worker_id = models.CharField(max_length=255)
@@ -812,11 +884,12 @@ class Meta:
         db_table = "fullctl_task_schedule_claim"
         verbose_name = _("Task Schedule Claim")
         verbose_name_plural = _("Task Schedule Claims")
-        unique_together = (("task_schedule", "schedule_date"))
+        unique_together = ("task_schedule", "schedule_date")
 
     class HandleRef:
         tag = "task_schedule_claim"
 
+
 class Monitor(HandleRefModel):
     email = models.EmailField(
         null=True, blank=True, help_text=_("Primary alert notification email")

src/fullctl/django/tasks/context.py

@@ -0,0 +1,129 @@
+"""
+Task execution context management.
+
+Provides a thread-safe way to access the currently executing task
+from anywhere in the call stack without explicit parameter passing.
+"""
+
+from contextvars import ContextVar
+from typing import Optional
+
+__all__ = [
+    "task_execution_context",
+    "get_current_task",
+    "check_task_cancelled",
+]
+
+_current_task: ContextVar = ContextVar("current_task", default=None)
+
+
+class task_execution_context:
+    """
+    Context manager for tracking the currently executing task.
+
+    This sets the task in context-local storage, making it accessible
+    to any code in the call stack via get_current_task() without
+    requiring explicit parameter passing.
+
+    Usage:
+        with task_execution_context(task):
+            # task is now accessible via get_current_task()
+            do_work()
+
+    Example:
+        class MyTask(Task):
+            def run(self, *args, **kwargs):
+                # Context automatically set by Task._run()
+                process_data()  # Can call check_task_cancelled() anywhere
+
+    Note:
+        In most cases, you don't need to use this directly.
+        Task._run() automatically wraps execution in this context.
+    """
+
+    def __init__(self, task):
+        """
+        Initialize context manager.
+
+        Args:
+            task: The Task instance to set as current
+        """
+        self.task = task
+        self.token = None
+
+    def __enter__(self):
+        """
+        Enter the context and set current task.
+
+        Returns:
+            The task instance
+        """
+        # Set the task and get a token for cleanup
+        self.token = _current_task.set(self.task)
+        return self.task
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exit the context and restore previous task.
+
+        Args:
+            exc_type: Exception type if an exception occurred
+            exc_val: Exception value if an exception occurred
+            exc_tb: Exception traceback if an exception occurred
+
+        Returns:
+            False to propagate exceptions (don't suppress)
+        """
+        # Reset to previous value using the token
+        _current_task.reset(self.token)
+        return False  # Don't suppress exceptions
+
+
+def get_current_task() -> Optional["Task"]:  # noqa: F821
+    """
+    Get the currently executing task from context.
+
+    This can be called from anywhere within a task execution context
+    to access the task instance without explicit parameter passing.
+
+    Returns:
+        Task instance if within task execution context, None otherwise
+
+    Example:
+        def some_deep_function():
+            task = get_current_task()
+            if task:
+                print(f"Running in task {task.id}")
+            else:
+                print("Not in task context")
+    """
+    return _current_task.get()
+
+
+def check_task_cancelled():
+    """
+    Check if the current task has been cancelled.
+
+    This can be called from anywhere within a task execution context
+    to gracefully handle cancellation. If called outside a task context,
+    it does nothing (safe to use in code that may run outside tasks).
+
+    The function will check the database for the current task's status
+    and raise TaskCancelledException if the task has been cancelled.
+
+    Raises:
+        TaskCancelledException: If the current task is cancelled
+
+    Example:
+        def process_items(items):
+            for item in items:
+                check_task_cancelled()  # Check before each item
+                process(item)
+
+    Note:
+        This is safe to call even if not in a task context - it will
+        simply do nothing if no task is currently executing.
+    """
+    task = get_current_task()
+    if task:
+        task.check_cancelled()