Support for bind=True kwarg (#53)

* Support for bind=True kwarg

* Add integration tests for bind=True

* Add docs

Author: AlanCoding

README.md

@@ -20,6 +20,8 @@ You will use dispatcher to trigger a background task over pg_notify.
 Both your *background dispatcher service* and your *task publisher* process must have
 python configured so that your task is importable.
 
+For more options, see `docs/usage.md`.
+
 #### Library
 
 The dispatcher `@task()` decorator is used to register tasks.

dispatcher/publish.py

@@ -10,17 +10,24 @@
 
 class DispatcherDecorator:
     def __init__(
-        self, registry: DispatcherMethodRegistry, *, queue: Optional[str] = None, on_duplicate: Optional[str] = None, timeout: Optional[float] = None
+        self,
+        registry: DispatcherMethodRegistry,
+        *,
+        bind: bool = False,
+        queue: Optional[str] = None,
+        on_duplicate: Optional[str] = None,
+        timeout: Optional[float] = None,
     ) -> None:
         self.registry = registry
+        self.bind = bind
         self.queue = queue
         self.on_duplicate = on_duplicate
         self.timeout = timeout
 
     def __call__(self, fn: DispatcherCallable, /) -> DispatcherCallable:
         "Concrete task decorator, registers method and glues on some methods from the registry"
 
-        dmethod = self.registry.register(fn, queue=self.queue, on_duplicate=self.on_duplicate, timeout=self.timeout)
+        dmethod = self.registry.register(fn, bind=self.bind, queue=self.queue, on_duplicate=self.on_duplicate, timeout=self.timeout)
 
         setattr(fn, 'apply_async', dmethod.apply_async)
         setattr(fn, 'delay', dmethod.delay)
@@ -30,6 +37,7 @@ def __call__(self, fn: DispatcherCallable, /) -> DispatcherCallable:
 
 def task(
     *,
+    bind: bool = False,
     queue: Optional[str] = None,
     on_duplicate: Optional[str] = None,
     timeout: Optional[float] = None,
@@ -72,4 +80,4 @@ def announce():
     # The on_duplicate kwarg controls behavior when multiple instances of the task running
     # options are documented in dispatcher.utils.DuplicateBehavior
     """
-    return DispatcherDecorator(registry, queue=queue, on_duplicate=on_duplicate, timeout=timeout)
+    return DispatcherDecorator(registry, bind=bind, queue=queue, on_duplicate=on_duplicate, timeout=timeout)

dispatcher/registry.py

@@ -54,7 +54,9 @@ def publication_defaults(self) -> dict:
     def delay(self, *args, **kwargs) -> Tuple[dict, str]:
         return self.apply_async(args, kwargs)
 
-    def get_async_body(self, args=None, kwargs=None, uuid=None, on_duplicate: Optional[str] = None, timeout: Optional[float] = 0.0, delay: float = 0.0) -> dict:
+    def get_async_body(
+        self, args=None, kwargs=None, uuid=None, bind: bool = False, on_duplicate: Optional[str] = None, timeout: Optional[float] = 0.0, delay: float = 0.0
+    ) -> dict:
         """
         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
@@ -66,6 +68,8 @@ def get_async_body(self, args=None, kwargs=None, uuid=None, on_duplicate: Option
 
         # TODO: callback to add other things, guid in case of AWX
 
+        if bind:
+            body['bind'] = bind
         if on_duplicate:
             body['on_duplicate'] = on_duplicate
         if delay:

dispatcher/utils.py

@@ -25,7 +25,7 @@ def resolve_callable(task: str) -> Optional[Callable]:
     That is out of scope of this method now.
     This is mainly used by the worker.
     """
-    if task.startswith('lambda:'):
+    if task.startswith('lambda'):
         return eval(task)
 
     if MODULE_METHOD_DELIMITER not in task:

dispatcher/worker/task.py

@@ -35,6 +35,18 @@ def exit_gracefully(self, *args, **kwargs):
         self.kill_now = True
 
 
+class DispatcherBoundMethods:
+    """
+    If you use the task decorator with the bind=True argument,
+    an object of this type will be passed in.
+    This contains public methods for users of the dispatcher to call.
+    """
+
+    def __init__(self, worker_id, message):
+        self.worker_id = worker_id
+        self.uuid = message.get('uuid', '<unknown>')
+
+
 class TaskWorker:
     """
     A worker implementation that deserializes task messages and runs native
@@ -50,8 +62,8 @@ class TaskWorker:
     Previously this initialized pre-fork, making init logic unusable.
     """
 
-    def __init__(self, worker_id):
-        self.worker_id = worker_id
+    def __init__(self, worker_id: int):
+        self.worker_id: int = worker_id
         self.ppid = os.getppid()
         self.pid = os.getpid()
         self.signal_handler = WorkerSignalHandler(worker_id)
@@ -68,18 +80,27 @@ def should_exit(self) -> bool:
     def get_uuid(self, message):
         return message.get('uuid', '<unknown>')
 
+    def produce_binder(self, message: dict) -> DispatcherBoundMethods:
+        """
+        Return the object with public callbacks to pass to the task
+        """
+        return DispatcherBoundMethods(self.worker_id, message)
+
     def run_callable(self, message):
         """
-        Given some AMQP message, import the correct Python code and run it.
+        Import the Python code and run it.
         """
         task = message['task']
-        args = message.get('args', [])
+        args = message.get('args', []).copy()
         kwargs = message.get('kwargs', {})
         _call = registry.get_method(task).get_callable()
 
         # don't print kwargs, they often contain launch-time secrets
         logger.debug(f'task (uuid={self.get_uuid(message)}) starting {task}(*{args}) on worker {self.worker_id}')
 
+        if message.get('bind') is True:
+            args = [self.produce_binder(message)] + args
+
         try:
             return _call(*args, **kwargs)
         except DispatcherCancel:

docs/task_options.md

@@ -0,0 +1,111 @@
+## Task Options
+
+You can specify additional details about the behavior of a task
+by applying additional options.
+
+There are normally 2 days to do this
+ - pass to the `@task` decorator
+ - pass to the `.apply_async` method for a one-off use
+
+The structure of use of those options is illustrated below,
+with `options` being keyword-based additional options.
+
+```python
+from dispatcher.publish import task
+
+@task(queue='test_channel', **options)
+def print_hello():
+    print('hello world!!')
+```
+
+For example, to set a timeout of 1 second on the task:
+
+```python
+from dispatcher.publish import task
+
+@task(queue='test_channel', **options)
+def print_hello():
+    print('hello world!!')
+```
+
+For the one-off use, `.apply_async` can take options,
+but `.delay` cannot, because of the argument structure.
+Using `.delay` inherently runs the task with the task default options.
+
+```python
+from test_methods import print_hello
+
+print_hello.apply_async(args=[], kwargs={}, **options)
+```
+
+For the timeout seconds example:
+
+```python
+from test_methods import print_hello
+
+print_hello.apply_async(args=[], kwargs={}, timeout=2)
+```
+
+The `apply_async` options will take precedence over the
+task default options (those passed into the decorator).
+
+### Task Options Manifest
+
+This section documents specific options.
+These follow a "standard" pattern, meaning that they
+can be used in both of the ways described above.
+
+#### Bind
+
+If `bind=True` is passed (default is `False`), then
+additional argument is inserted at the start of the
+argument list to the method. Like:
+
+```python
+@task(bind=True)
+def hello_world(dispatcher, *args, **kwargs):
+    print(f'I see the dispatcher object {dispatcher}')
+```
+
+The `dispatcher` object contains public methods
+which allow interaction with the parent process.
+Available methods will expand in the future,
+right now it offers:
+
+ - `uuid` - the internal id of this task call in dispatcher
+ - `worker_id` - the id of the worker running this task
+
+#### Queue
+
+The queue or channel this task is submitted to.
+For instance, the pg_notify channel.
+This can be a callable to get this dynamically.
+
+#### on_duplicate
+
+This option helps to manage capacity, controlling
+ - task "shedding"
+ - task queuing
+
+Depending on the value, a task submission will be ignored
+if certain conditions are met, "shedding", or queued if all
+workers are busy.
+
+ - parallel - multiple tasks (running the given `@task` method) are allowed at the same time. Tasks queue if no free workers are available.
+  - discard - if a task is already being ran or is queued, any new submissions of this task are ignored.
+  - serial - only 1 task (running the given method) will be ran at a single time in the local dispatcher service. Additional submissions are queued, so all submissions will be ran eventually.
+  - queue_one - for idempotent tasks, only 1 task (running the given method) will be ran at a single time, and an additional submission is queued. However, only 1 task will be held in the queue, and additional submissions are discarded. This assures _timely_ running of an idempotent task.
+
+### Unusual Options
+
+These do not follow the standard pattern for some reason.
+Usually for testing.
+
+#### registry
+
+The dispatcher uses a global task registry.
+To enable isolated testing `@task` can take a custom
+(meaning non-global) registry.
+
+There is no real multi-registry feature,
+and additional custom code hooks would be needed to make this work.

tests/integration/test_bind.py

@@ -0,0 +1,43 @@
+import asyncio
+import json
+import logging
+
+import pytest
+
+from tests.data import methods as test_methods
+
+ASSERT_UUID = 'lambda dispatcher: dispatcher.uuid'
+
+
+@pytest.mark.asyncio
+async def test_bind_uuid_matches(apg_dispatcher, pg_message, caplog):
+    assert apg_dispatcher.pool.finished_count == 0
+
+    clearing_task = asyncio.create_task(apg_dispatcher.pool.events.work_cleared.wait())
+    with caplog.at_level(logging.DEBUG):
+        await pg_message(json.dumps({
+            'task': ASSERT_UUID,
+            'uuid': 'hello-world-12345543221',
+            'bind': True
+        }))
+        await asyncio.wait_for(clearing_task, timeout=3)
+
+    assert 'result: hello-world-12345543221' in caplog.text
+
+    assert apg_dispatcher.pool.finished_count == 1
+
+@pytest.mark.asyncio
+async def test_bind_not_set(apg_dispatcher, pg_message, caplog):
+    assert apg_dispatcher.pool.finished_count == 0
+
+    clearing_task = asyncio.create_task(apg_dispatcher.pool.events.work_cleared.wait())
+    with caplog.at_level(logging.DEBUG):
+        await pg_message(json.dumps({
+            'task': ASSERT_UUID,
+            'uuid': 'hello-world-12345543221'
+        }))
+        await asyncio.wait_for(clearing_task, timeout=3)
+
+    assert 'result: hello-world-12345543221' not in caplog.text
+
+    assert apg_dispatcher.pool.finished_count == 1

tools/test_methods.py

@@ -18,6 +18,12 @@ def sleep_serial(seconds=1):
     time.sleep(seconds)
 
 
+@task(queue='test_channel', bind=True)
+def hello_world_binder(binder):
+    print(f'Values in binder {vars(binder)}')
+    print(f'Hello world, from worker {binder.worker_id} running task {binder.uuid}')
+
+
 @task(queue='test_channel')
 def print_hello():
     print('hello world!!')

tools/write_messages.py

@@ -15,7 +15,7 @@
 
 sys.path.append(tools_dir)
 
-from test_methods import sleep_function, sleep_discard, task_has_timeout
+from test_methods import sleep_function, sleep_discard, task_has_timeout, hello_world_binder
 
 # Database connection details
 CONNECTION_STRING = "dbname=dispatch_db user=dispatch password=dispatching host=localhost port=55777"
@@ -111,6 +111,9 @@ def main():
     print('demo of task_has_timeout that times out due to decorator use')
     task_has_timeout.apply_async(config={'conninfo': CONNECTION_STRING})
 
+    print('demo of using bind=True, with hello_world_binder')
+    hello_world_binder.apply_async(config={'conninfo': CONNECTION_STRING})
+
 if __name__ == "__main__":
     logging.basicConfig(level='ERROR', stream=sys.stdout)
     main()