237 lines
7.6 KiB
Python
237 lines
7.6 KiB
Python
from django.utils import timezone
|
|
from django.db import models
|
|
from django.db.models.constraints import UniqueConstraint, Q
|
|
from django.tasks.base import TaskResult, TaskResultStatus
|
|
from django.tasks import task_backends
|
|
|
|
from .base.models import BaseModel
|
|
from .utils import TIMEDELTA_PATTERN
|
|
|
|
import functools, traceback, io
|
|
import random, uuid
|
|
import asyncio
|
|
import os, threading
|
|
|
|
|
|
class Worker( BaseModel ):
|
|
id = models.UUIDField( primary_key = True, default = uuid.uuid4, editable = False )
|
|
|
|
process_id = models.IntegerField( default = os.getpid )
|
|
thread_id = models.PositiveBigIntegerField( default = threading.get_ident )
|
|
|
|
creation_datetime = models.DateTimeField( auto_now_add = True )
|
|
|
|
is_standalone = models.BooleanField( default = False )
|
|
is_coordinator = models.BooleanField( default = False )
|
|
#in_grace = models.BooleanField( default = False ) #If the worker sees this as True, it should kill itself!
|
|
|
|
last_activity = models.DateTimeField( null = True, blank = True )
|
|
|
|
def __str__( self ): return f"P{self.process_id}W{self.thread_id}"
|
|
|
|
class Meta:
|
|
constraints = [
|
|
UniqueConstraint( fields = ['is_coordinator'], condition = Q( is_coordinator = True ), name = 'unique_coordinator' ),
|
|
]
|
|
|
|
|
|
class TaskSchedule( BaseModel ):
|
|
name = models.CharField( default = "default", max_length = 200 )
|
|
|
|
is_enabled = models.BooleanField( default = True )
|
|
|
|
task_path = models.TextField() #Path to the task function
|
|
@property
|
|
def task( self ): return task_backends['default'].TASKS[self.task_path]
|
|
|
|
args = models.JSONField( default = list, blank = True )
|
|
kwargs = models.JSONField( default = dict, blank = True )
|
|
|
|
#Distinguishes Periodic and Service like Tasks
|
|
interval = models.DurationField( null = True, blank = True )
|
|
@property
|
|
def type( self ):
|
|
if self.interval is None: return "S" #Service
|
|
return "P" #Periodic
|
|
|
|
#For Periodic tasks, it's the whole execution,
|
|
#For Service like tasks, it's the gracetime after the first exit signal.
|
|
timeout = models.DurationField( default = timezone.timedelta( minutes = 5 ) )
|
|
|
|
#delay before execution for Both task types
|
|
jitter_length = models.DurationField( default = timezone.timedelta( seconds = 0 ), blank = True )
|
|
jitter_pivot = models.CharField( default = "M", max_length = 1, choices = {
|
|
"S":"Start", "M":"Middle", "E":"End",
|
|
})
|
|
|
|
def get_jitter( self ):
|
|
jitter = self.jitter_length * random.random()
|
|
match self.jitter_pivot:
|
|
case "M": jitter -= self.jitter_length / 2
|
|
case "E": jitter *= -1
|
|
return jitter
|
|
|
|
def set_jitter( self, jitter ):
|
|
if not jitter: #Set the default values
|
|
self.jitter_length = timezone.timedelta( seconds = 0 )
|
|
self.jitter_pivot = "M"
|
|
return
|
|
|
|
match = TIMEDELTA_PATTERN.match( jitter )
|
|
assert match, "Provided jitter value does not match the correct timedelta pattern! (ex: 1w2d5h30m10s500ms1000us)"
|
|
|
|
self.jitter_length = timezone.timedelta( **{
|
|
k: int(v)
|
|
for k, v in match.groupdict().items()
|
|
if v is not None
|
|
} )
|
|
self.jitter_pivot = {
|
|
"-": "S",
|
|
None: "M",
|
|
"+": "E",
|
|
}[match.group(1)]
|
|
|
|
prune_success_over = models.IntegerField( default = 10 )
|
|
prune_failed_over = models.IntegerField( default = 1000 )
|
|
|
|
class Meta:
|
|
unique_together = [
|
|
('name', 'task_path'),
|
|
]
|
|
|
|
def __str__( self ):
|
|
type_display = ( "Service Task" if self.type == "S" else "Periodic Task" )
|
|
short = self.task_path.rsplit('.')[-1]
|
|
return f"{self.name} {type_display} {short}"
|
|
|
|
def as_trace( self ):
|
|
return Trace(
|
|
schedule = self,
|
|
task_path = self.task_path,
|
|
|
|
args = self.args,
|
|
kwargs = self.kwargs
|
|
)
|
|
|
|
async def set_traces_scheduled_datetime( self, trace ):
|
|
assert trace.schedule_id == self.id, "This trace does not belong to this schedule!"
|
|
|
|
now = timezone.now()
|
|
jitter_delta = self.get_jitter()
|
|
last_trace = await self.trace_set.exclude(
|
|
status = TaskResultStatus.READY
|
|
).order_by( '-started_datetime' ).afirst()
|
|
|
|
if last_trace: #Execute now + jitter
|
|
trace.scheduled_datetime = last_trace.started_datetime + jitter_delta
|
|
if self.interval is not None: #Add the "Period" in periodic tasks
|
|
trace.scheduled_datetime += self.interval
|
|
|
|
else:
|
|
trace.scheduled_datetime = now + jitter_delta
|
|
|
|
if trace.scheduled_datetime < now: #So, in case jitter is negative
|
|
trace.scheduled_datetime = now
|
|
|
|
|
|
|
|
class Trace( BaseModel ):
|
|
id = models.UUIDField( primary_key = True, default = uuid.uuid4, editable = False )
|
|
worker = models.ForeignKey( Worker, null = True, on_delete = models.SET_NULL )
|
|
|
|
task_path = models.TextField() #Path to the task function
|
|
schedule = models.ForeignKey( TaskSchedule, null = True, on_delete = models.SET_NULL )
|
|
|
|
status = models.CharField(
|
|
default = TaskResultStatus.READY,
|
|
choices = TaskResultStatus.choices,
|
|
max_length = max( len(v) for v in TaskResultStatus.values ),
|
|
)
|
|
|
|
|
|
args = models.JSONField( default = list, blank = True )
|
|
kwargs = models.JSONField( default = dict, blank = True )
|
|
|
|
enqueued_datetime = models.DateTimeField( auto_now_add = True )
|
|
scheduled_datetime = models.DateTimeField()
|
|
started_datetime = models.DateTimeField( null = True, blank = True )
|
|
finished_datetime = models.DateTimeField( null = True, blank = True )
|
|
|
|
lifetime = models.JSONField( default = dict, blank = True )
|
|
stdout = models.TextField( default = "", blank = True )
|
|
exception_class_path = models.TextField( null = True )
|
|
traceback = models.TextField( null = True )
|
|
|
|
def set_status( self, status, desc = "" ):
|
|
|
|
if isinstance( status, BaseException ):
|
|
exception_type = type(status)
|
|
self.status = TaskResultStatus.FAILED
|
|
self.lifetime[self.status] = desc
|
|
self.traceback = traceback.format_exception( status )
|
|
self.exception_class_path = f"{exception_type.__module__}.{exception_type.__qualname__}"
|
|
return
|
|
|
|
assert status in TaskResultStatus, f"Unkown status '{status}' provided for trace!"
|
|
self.status = status
|
|
self.lifetime[self.status] = desc
|
|
|
|
return_value = models.JSONField( null = True, blank = True )
|
|
prune_protected = models.BooleanField( default = False ) #Do not delete traces with this flag.
|
|
|
|
def __str__( self ): return f"{self.get_status_display()} Trace of {self.task_path}"
|
|
|
|
|
|
@property
|
|
def task( self ): return task_backends['default'].TASKS[self.task_path]
|
|
|
|
@property
|
|
def task_error( self ):
|
|
if not self.exception_class_path: return
|
|
return TaskError(
|
|
exception_class_path = self.exception_class_path,
|
|
traceback = self.traceback,
|
|
)
|
|
|
|
@property
|
|
def task_result( self ):
|
|
error = self.task_error
|
|
return TaskResult(
|
|
task = self.task,
|
|
id = self.id,
|
|
status = self.status,
|
|
enqueued_at = self.enqueued_datetime,
|
|
started_at = self.started_datetime,
|
|
last_attempted_at = self.started_datetime,
|
|
finished_at = self.finished_datetime,
|
|
args = self.args,
|
|
kwargs = self.kwargs,
|
|
backend = self.task.backend,
|
|
errors = [error] if error else [],
|
|
worker_ids = [self.worker_id] if self.worker_id else [],
|
|
)
|
|
|
|
class Meta:
|
|
constraints = [
|
|
UniqueConstraint(
|
|
fields = ['schedule_id'],
|
|
condition = models.Q(status = TaskResultStatus.READY),
|
|
name = "unique_ready_for_each_schedule",
|
|
),
|
|
]
|
|
|
|
|
|
### Runtime methods for context aware tasks
|
|
async def commit_on_new_print( self ):
|
|
while True: #new_print event needs to be created in the worker.
|
|
await self.new_print.wait()
|
|
await self.asave( update_fields = ['stdout'] )
|
|
self.new_print.clear()
|
|
|
|
def print( self, *args, sep = " ", end = "\n", file = None, flush = True ): #We get 'file' here to fool tasks that aren't self_aware
|
|
assert hasattr(self, 'commit_on_new_print_task'), "trace.print needs to be called while a trace.commit_on_new_print task is active!"
|
|
string = sep.join( str(i) for i in args ) + end
|
|
self.stdout += string
|
|
if flush: self.new_print.set()
|
|
if getattr(self, "show_prints", False): print( string )
|