wip

django_q/puller.py

@@ -0,0 +1,68 @@
+from django_q.worker import Worker
+from django_q.signing import BadSignature, SignedPackage
+from time import sleep
+from django_q.brokers import get_broker
+import multiprocessing
+from django_q.queue_task import QueueTask
+from django.utils import timezone
+import enum
+import traceback
+
+from multiprocessing import Event, Process, Value, current_process
+from django_q.utils import close_old_django_connections
+from django.utils.translation import gettext_lazy as _
+from django_q.conf import Conf, logger, setproctitle, error_reporter, resource, psutil
+from django_q.exceptions import TimeoutException, TimeoutHandler
+from django_q.process_manager import ProcessManager
+
+
+class Puller(ProcessManager):
+    """The Puller is responsible for pulling the tasks from the broker, then return them to be picked up by the
+    guard"""
+
+    def get_target(self):
+        return self.run_puller
+
+    def stop_puller(self):
+        self.status.value = self.Status.DONE.value
+
+    def run_puller(self, status, pipe) -> None:
+        broker = get_broker()
+        proc_name = current_process().name
+        if setproctitle:
+            setproctitle.setproctitle(f"qcluster {proc_name} puller")
+        logger.info(
+            _("%(name)s pulling tasks from broker %(id)s")
+            % {"name": proc_name, "id": current_process().pid}
+        )
+        while True:
+            if status.value == Worker.Status.DONE.value:
+                logger.info("Stopping Puller")
+                break
+            try:
+                task_set = broker.dequeue()
+            except Exception:
+                logger.exception("Failed to pull task from broker")
+                # broker probably crashed. Let the sentinel handle it.
+                sleep(10)
+                break
+            if task_set:
+                logger.info(
+                    _("Found %(amount_tasks)s tasks") % {"amount_tasks": len(task_set)}
+                )
+                for task in task_set:
+                    ack_id = task[0]
+                    # unpack the task
+                    try:
+                        queue_task = SignedPackage.loads(task[1])
+                    except (TypeError, BadSignature):
+                        logger.exception("Failed to pull task from broker - bad task")
+                        broker.fail(ack_id)
+                        continue
+                    queue_task.ack_id = ack_id
+                    # send back to main process
+                    pipe.send(queue_task)
+                logger.debug(
+                    _("queueing from %(list_key)s") % {"list_key": broker.list_key}
+                )
+        logger.info(_("%(name)s stopped pushing tasks") % {"name": current_process().name})