import inspect import logging import time from uuid import uuid4 from dispatcherd.publish import submit_task from dispatcherd.processors.blocker import Blocker from dispatcherd.utils import resolve_callable from django_guid import get_guid logger = logging.getLogger('awx.main.dispatch') def serialize_task(f): return '.'.join([f.__module__, f.__name__]) class task: """ Used to decorate a function or class so that it can be run asynchronously via the task dispatcher. Tasks can be simple functions: @task() def add(a, b): return a + b ...or classes that define a `run` method: @task() class Adder: def run(self, a, b): return a + b # Tasks can be run synchronously... assert add(1, 1) == 2 assert Adder().run(1, 1) == 2 # ...or published to a queue: add.apply_async([1, 1]) Adder.apply_async([1, 1]) # Tasks can also define a specific target queue or use the special fan-out queue tower_broadcast: @task(queue='slow-tasks') def snooze(): time.sleep(10) @task(queue='tower_broadcast') def announce(): print("Run this everywhere!") # The special parameter bind_kwargs tells the main dispatcher process to add certain kwargs @task(bind_kwargs=['dispatch_time']) def print_time(dispatch_time=None): print(f"Time I was dispatched: {dispatch_time}") """ def __init__(self, queue=None, bind_kwargs=None, timeout=None, on_duplicate=None): self.queue = queue self.bind_kwargs = bind_kwargs self.timeout = timeout self.on_duplicate = on_duplicate def __call__(self, fn=None): queue = self.queue bind_kwargs = self.bind_kwargs timeout = self.timeout on_duplicate = self.on_duplicate class PublisherMixin(object): queue = None @classmethod def delay(cls, *args, **kwargs): return cls.apply_async(args, kwargs) @classmethod def get_async_body(cls, args=None, kwargs=None, uuid=None, **kw): """ Get the python dict to become JSON data in the pg_notify message This same message gets passed over the dispatcher IPC queue to workers If a task is submitted to a multiprocessing pool, skipping pg_notify, this might be used directly """ task_id = uuid or str(uuid4()) args = args or [] kwargs = kwargs or {} obj = {'uuid': task_id, 'args': args, 'kwargs': kwargs, 'task': cls.name, 'time_pub': time.time()} guid = get_guid() if guid: obj['guid'] = guid if bind_kwargs: obj['bind_kwargs'] = bind_kwargs obj.update(**kw) return obj @classmethod def apply_async(cls, args=None, kwargs=None, queue=None, uuid=None, **kw): # At this point we have the import string, and submit_task wants the method, so back to that actual_task = resolve_callable(cls.name) processor_options = () if on_duplicate is not None: processor_options = (Blocker.Params(on_duplicate=on_duplicate),) return submit_task( actual_task, args=args, kwargs=kwargs, queue=queue, uuid=uuid, timeout=timeout, processor_options=processor_options, **kw, ) # If the object we're wrapping *is* a class (e.g., RunJob), return # a *new* class that inherits from the wrapped class *and* BaseTask # In this way, the new class returned by our decorator is the class # being decorated *plus* PublisherMixin so cls.apply_async() and # cls.delay() work bases = [] ns = {'name': serialize_task(fn), 'queue': queue} if inspect.isclass(fn): bases = list(fn.__bases__) ns.update(fn.__dict__) cls = type(fn.__name__, tuple(bases + [PublisherMixin]), ns) if inspect.isclass(fn): return cls # if the object being decorated is *not* a class (it's a Python # function), make fn.apply_async and fn.delay proxy through to the # PublisherMixin we dynamically created above setattr(fn, 'name', cls.name) setattr(fn, 'apply_async', cls.apply_async) setattr(fn, 'delay', cls.delay) setattr(fn, 'get_async_body', cls.get_async_body) return fn