django-nativemojo 0.1.10__py3-none-any.whl → 0.1.15__py3-none-any.whl

mojo/apps/tasks/manager.py

@@ -1,3 +1,4 @@
+from posix import lockf
 from venv import create
 from mojo.helpers import redis
 from objict import objict
@@ -6,7 +7,7 @@ import time
 
 
 class TaskManager:
-    def __init__(self, channels, prefix="taskit"):
+    def __init__(self, channels, prefix="mojo:tasks"):
         """
         Initialize the TaskManager with Redis connection, channels, and a key prefix.
 
@@ -17,11 +18,13 @@ class TaskManager:
         self.channels = channels
         self.prefix = prefix
 
-    def take_out_the_dead(self):
-        for channel in self.channels:
+    def take_out_the_dead(self, local=False):
+        # Get channels from Redis instead of using self.channels
+        channels = self.channels if local else self.get_all_channels()
+        for channel in channels:
             # this will remove all expired tasks by default
-            self.get_pending()
-            self.get_running()
+            self.get_pending(channel)
+            self.get_running(channel)
 
     def get_completed_key(self, channel):
         """
@@ -41,6 +44,23 @@ class TaskManager:
         """
         return f"{self.prefix}:p:{channel}"
 
+    def get_global_pending_key(self):
+        """
+        Get the Redis key for pending tasks in a channel.
+
+        :param channel: Channel name.
+        :return: Redis key for pending tasks.
+        """
+        return f"{self.prefix}:pending"
+
+    def get_channels_key(self):
+        """
+        Get the Redis key for tracking all channels.
+
+        :return: Redis key for channels set.
+        """
+        return f"{self.prefix}:channels"
+
     def get_error_key(self, channel):
         """
         Get the Redis key for tasks with errors in a channel.
@@ -77,6 +97,40 @@ class TaskManager:
         """
         return f"{self.prefix}:c:{channel}"
 
+    def get_runners_key(self):
+        return f"{self.prefix}:runners"
+
+    def add_channel(self, channel):
+        """
+        Add a channel to the global channels set.
+
+        :param channel: Channel name.
+        """
+        self.redis.sadd(self.get_channels_key(), channel)
+
+    def remove_channel(self, channel):
+        """
+        Remove a channel from the global channels set.
+
+        :param channel: Channel name.
+        """
+        self.redis.srem(self.get_channels_key(), channel)
+
+    def remove_all_channels(self):
+        """
+        Remove all channels from the global channels set.
+        """
+        self.redis.delete(self.get_channels_key())
+
+    def get_all_channels(self):
+        """
+        Get all channels from the global channels set.
+
+        :return: List of channel names.
+        """
+        channels = self.redis.smembers(self.get_channels_key())
+        return [channel.decode('utf-8') for channel in channels]
+
     def get_task(self, task_id):
         """
         Retrieve a task from Redis using its task ID.
@@ -85,6 +139,8 @@ class TaskManager:
         :return: Task data as an objict, or None if not found.
         """
         task_data_raw = self.redis.get(self.get_task_key(task_id))
+        if not task_data_raw:
+            return None
         return objict.from_json(task_data_raw, ignore_errors=True)
 
     def save_task(self, task_data, expires=1800):
@@ -114,6 +170,8 @@ class TaskManager:
         :param channel: Channel name. Default is "default".
         :return: True if operation is successful.
         """
+        # Add channel to global channels set
+        self.add_channel(channel)
         self.redis.sadd(self.get_pending_key(channel), task_id)
         return True
 
@@ -208,8 +266,8 @@ class TaskManager:
         """
         task_data = self.get_task(task_id)
         if task_data:
-            self.redis.srem(self.get_running_key(task_data.channel), task_data.id)
-            self.redis.srem(self.get_pending_key(task_data.channel), task_data.id)
+            self.remove_from_running(task_data.id, task_data.channel)
+            self.remove_from_pending(task_data.id, task_data.channel)
             self.redis.delete(self.get_task_key(task_id))
             return True
         return False
@@ -223,9 +281,11 @@ class TaskManager:
         task_data_raw = self.redis.get(self.get_task_key(task_id))
         task_data = objict.from_json(task_data_raw, ignore_errors=True)
         if task_data:
-            self.redis.srem(self.get_running_key(task_data.channel), task_data.id)
-            self.redis.srem(self.get_pending_key(task_data.channel), task_data.id)
-            self.redis.delete(self.get_task_key(task_id))
+            task_data.status = "cancelled"
+            self.remove_from_pending(task_data.id, task_data.channel)
+            self.save_task(task_data)
+            return True
+        return False
 
     def get_running_ids(self, channel="default"):
         """
@@ -337,6 +397,8 @@ class TaskManager:
         :param expires: Expiration time in seconds. Default is 1800.
         :return: Task ID of the published task.
         """
+        from mojo.apps import metrics
+
         task_data = objict(channel=channel, function=function, data=objict.from_dict(data))
         task_data.id = str(uuid.uuid4()).replace('-', '')
         task_data.created = time.time()
@@ -345,62 +407,66 @@ class TaskManager:
         self.add_to_pending(task_data.id, channel)
         self.save_task(task_data, expires)
         self.redis.publish(self.get_channel_key(channel), task_data.id)
-        # metrics.record("taskit:p:global", category="taskit")
-        # metrics.record(f"taskit:p:{channel}", category="taskit_channels")
+        metrics.record("tasks_pub", category="tasks")
+        metrics.record(f"tasks_pub_{channel}", category=f"tasks_{channel}")
         return task_data.id
 
-    def get_all_pending_ids(self):
+    def get_all_pending_ids(self, local=False):
         """
         Get all pending task IDs across all channels.
 
         :return: List of all pending task IDs.
         """
         pending_ids = []
-        for channel in self.channels:
+        channels = self.channels if local else self.get_all_channels()
+        for channel in channels:
             pending_ids.extend(self.get_pending_ids(channel))
         return pending_ids
 
-    def get_all_running_ids(self):
+    def get_all_running_ids(self, local=False):
         """
         Get all running task IDs across all channels.
 
         :return: List of all running task IDs.
         """
         running_ids = []
-        for channel in self.channels:
+        channels = self.channels if local else self.get_all_channels()
+        for channel in channels:
             running_ids.extend(self.get_running_ids(channel))
         return running_ids
 
-    def get_all_completed_ids(self):
+    def get_all_completed_ids(self, local=False):
         """
         Get all completed task IDs across all channels.
 
         :return: List of all completed task IDs.
         """
         completed_ids = []
-        for channel in self.channels:
+        channels = self.channels if local else self.get_all_channels()
+        for channel in channels:
             completed_ids.extend(self.get_completed_ids(channel))
         return completed_ids
 
-    def get_all_error_ids(self):
+    def get_all_error_ids(self, local=False):
         """
         Get all error task IDs across all channels.
 
         :return: List of all error task IDs.
         """
         error_ids = []
-        for channel in self.channels:
+        channels = self.channels if local else self.get_all_channels()
+        for channel in channels:
             error_ids.extend(self.get_error_ids(channel))
         return error_ids
 
-    def get_all_pending(self, include_data=False):
+    def get_all_pending(self, include_data=False, local=False):
         """
         Get all pending tasks as objects.
 
         :return: List of pending task objects.
         """
         pending_tasks = []
-        for task_id in self.get_all_pending_ids():
+        for task_id in self.get_all_pending_ids(local=local):
             task = self.get_task(task_id)
             if task:
                 if not include_data:
@@ -410,14 +476,14 @@ class TaskManager:
                 self.remove_from_pending(task_id)
         return pending_tasks
 
-    def get_all_running(self, include_data=False):
+    def get_all_running(self, include_data=False, local=False):
         """
         Get all running tasks as objects.
 
         :return: List of running task objects.
         """
         running_tasks = []
-        for task_id in self.get_all_running_ids():
+        for task_id in self.get_all_running_ids(local=local):
             task = self.get_task(task_id)
             if task:
                 if not include_data:
@@ -427,14 +493,14 @@ class TaskManager:
                 self.remove_from_running(task_id)
         return running_tasks
 
-    def get_all_completed(self, include_data=False):
+    def get_all_completed(self, include_data=False, local=False):
         """
         Get all completed tasks as objects.
 
         :return: List of completed task objects.
         """
         completed_tasks = []
-        for task_id in self.get_all_completed_ids():
+        for task_id in self.get_all_completed_ids(local=local):
             task = self.get_task(task_id)
             if task:
                 if not include_data:
@@ -446,14 +512,14 @@ class TaskManager:
         completed_tasks.sort(key=lambda x: x.created, reverse=True)
         return completed_tasks
 
-    def get_all_errors(self):
+    def get_all_errors(self, local=False):
         """
         Get all error tasks as objects.
 
         :return: List of error task objects.
         """
         error_tasks = []
-        for task_id in self.get_all_error_ids():
+        for task_id in self.get_all_error_ids(local=local):
             task = self.get_task(task_id)
             if task:
                 error_tasks.append(task)
@@ -463,27 +529,116 @@ class TaskManager:
         error_tasks.sort(key=lambda x: x.created, reverse=True)
         return error_tasks
 
-    def get_status(self, simple=False):
+    def get_channel_status(self, channel):
+        status = objict()
+        status.pending = len(self.get_pending_ids(channel))
+        status.running = len(self.get_running_ids(channel))
+        status.completed = len(self.get_completed_ids(channel))
+        status.errors = len(self.get_error_ids(channel))
+        return status
+
+    def get_status(self, simple=False, local=False):
             """
             Get the status of tasks across all channels, including pending and running tasks.
 
             :return: Status object containing counts of pending and running tasks per channel.
             """
-            status = objict(pending=0, running=0, completed=0, errors=0, channels=objict())
-            for channel in self.channels:
-                pending = self.get_pending_ids(channel)
-                running = self.get_running_ids(channel)
-                completed = self.get_completed_ids(channel)
-                errors = self.get_error_ids(channel)
-                status.pending += len(pending)
-                status.running += len(running)
-                status.completed += len(completed)
-                status.errors += len(errors)
+            status = objict(pending=0, running=0, completed=0, errors=0)
+            if not simple:
+                status.channels = objict()
+            channels = self.channels if local else self.get_all_channels()
+            # Use channels from Redis instead of self.channels
+            for channel in channels:
+                cstatus = self.get_channel_status(channel)
+                status.pending += cstatus.pending
+                status.running += cstatus.running
+                status.completed += cstatus.completed
+                status.errors += cstatus.errors
                 if not simple:
-                    cstats = objict()
-                    cstats.pending = len(pending)
-                    cstats.running = len(running)
-                    cstats.completed = len(completed)
-                    cstats.errors = len(errors)
-                    status.channels[channel] = cstats
+                    status.channels[channel] = cstatus
+            status["runners"] = self.get_active_runners()
             return status
+
+    def get_all_runners(self):
+        """
+        Get all runners.
+
+        Returns:
+            dict: Dictionary of runners with their status.
+        """
+        runners = {}
+        raw_runners = self.redis.hgetall(self.get_runners_key())
+        for hostname, data in raw_runners.items():
+            try:
+                runner = objict.from_json(data.decode())
+                runner["ping_age"] = time.time() - runner["last_ping"]
+                if runner["ping_age"] > 30:
+                    runner["status"] = "timeout"
+                runners[hostname.decode()] = runner
+            except Exception:
+                continue
+        return runners
+
+    def get_active_runners(self):
+        """
+        Get all active runners.
+
+        Returns:
+            dict: Dictionary of active runners with their status.
+        """
+        return self.get_all_runners()
+
+    def remove_runner(self, hostname):
+        """
+        Remove a runner from the list of active runners.
+
+        Args:
+            hostname (str): The hostname of the runner to remove.
+        """
+        self.redis.hdel(self.get_runners_key(), hostname)
+
+    def clear_runners(self, ping_age=30):
+        raw_runners = self.get_all_runners()
+        for runner in raw_runners.values():
+            try:
+                runner["ping_age"] = time.time() - runner["last_ping"]
+                if runner["ping_age"] > ping_age:
+                    self.remove_runner(runner["hostname"])
+            except Exception:
+                continue
+
+    def clear_running_tasks(self):
+        """
+        Reset tasks that are stuck in a running state by moving them back to the pending state.
+        """
+        for channel in self.channels:
+            for task_id in self.get_running_ids(channel):
+                # self.logger.info(f"moving task {task_id} from running to pending")
+                self.remove_from_running(task_id, channel)
+                self.add_to_pending(task_id, channel)
+
+    def clear_pending_tasks(self):
+        """
+        Reset tasks that are stuck in a pending state by moving them back to the pending state.
+        """
+        for channel in self.channels:
+            for task_id in self.get_pending_ids(channel):
+                # self.logger.info(f"moving task {task_id} from running to pending")
+                self.remove_from_pending(task_id, channel)
+
+    def clear_channel(self, channel):
+        for task_id in self.get_running_ids(channel):
+            self.remove_from_running(task_id, channel)
+        for task_id in self.get_pending_ids(channel):
+            self.remove_from_pending(task_id, channel)
+        for task_id in self.get_completed_ids(channel):
+            self.remove_from_completed(task_id, channel)
+        for task_id in self.get_error_ids(channel):
+            self.remove_from_errors(task_id, channel)
+
+    def clear_local_queues(self):
+        """
+        Reset tasks that are stuck in a pending state by moving them back to the pending state.
+        """
+        for channel in self.channels:
+            self.clear_channel(channel)

mojo/apps/tasks/rest/tasks.py

@@ -1,19 +1,30 @@
 from mojo import decorators as md
 from mojo.helpers.response import JsonResponse
 # from django.http import JsonResponse
-import mojo.tasks
+from mojo.apps import tasks
 
-@md.URL('status')
+@md.GET('status')
 def api_status(request):
-    tman = mojo.tasks.get_manager()
-    return JsonResponse(tman.get_status())
+    tman = tasks.get_manager()
+    return JsonResponse(dict(status=True, data=tman.get_status()))
+
+
+@md.GET('runners')
+def api_task_runners(request):
+    tman = tasks.get_manager()
+    runners = [r for r in tman.get_active_runners().values()]
+    for r in runners:
+        r['id'] = r['hostname']
+    return JsonResponse(dict(status=True, data=runners, size=len(runners), count=len(runners)))
+
 
 @md.URL('pending')
 def api_pending(request):
-    tman = mojo.tasks.get_manager()
+    tman = tasks.get_manager()
     pending = tman.get_all_pending()
     size = len(pending)
     response = {
+        'status': True,
         'count': size,
         'page': 0,
         'size': size,
@@ -23,10 +34,11 @@ def api_pending(request):
 
 @md.URL('completed')
 def api_completed(request):
-    tman = mojo.tasks.get_manager()
-    completed = tman.get_all_completed()
+    tman = tasks.get_manager()
+    completed = tman.get_all_completed(include_data=True)
     size = len(completed)
     response = {
+        'status': True,
         'count': size,
         'page': 0,
         'size': size,
@@ -36,10 +48,11 @@ def api_completed(request):
 
 @md.URL('running')
 def api_running(request):
-    tman = mojo.tasks.get_manager()
-    running = tman.get_all_running()
+    tman = tasks.get_manager()
+    running = tman.get_all_running(include_data=True)
     size = len(running)
     response = {
+        'status': True,
         'count': size,
         'page': 0,
         'size': size,
@@ -50,10 +63,11 @@ def api_running(request):
 
 @md.URL('errors')
 def api_errors(request):
-    tman = mojo.tasks.get_manager()
+    tman = tasks.get_manager()
     errors = tman.get_all_errors()
     size = len(errors)
     response = {
+        'status': True,
         'count': size,
         'page': 0,
         'size': size,