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

mojo/apps/tasks/runner.py

@@ -1,12 +1,16 @@
 from importlib import import_module
 from concurrent.futures import ThreadPoolExecutor
 from .manager import TaskManager
-from mojo.tasks import manager
+from mojo.apps.tasks import manager
 import os
 from mojo.helpers import logit
 from mojo.helpers import daemon
 from mojo.helpers import paths
+from mojo.apps import metrics
 import time
+import socket
+import threading
+import json
 
 
 class TaskEngine(daemon.Daemon):
@@ -23,13 +27,174 @@ class TaskEngine(daemon.Daemon):
             max_workers (int, optional): The maximum number of threads available for task execution. Defaults to 5.
         """
         super().__init__("taskit", os.path.join(paths.VAR_ROOT, "taskit"))
+        self.hostname = socket.gethostname()
         self.manager = manager.TaskManager(channels)
         self.channels = channels
         if "broadcast" not in self.channels:
             self.channels.append("broadcast")
+
+        # Add hostname-specific channel for this runner
+        self.runner_channel = f"runner_{self.hostname}"
+        if self.runner_channel not in self.channels:
+            self.channels.append(self.runner_channel)
+
         self.max_workers = max_workers
         self.executor = None
-        self.logger = logit.get_logger("taskit", "taskit.log")
+        self.logger = logit.get_logger("tasks", "tasks.log")
+        self.ping_thread = None
+        self.ping_interval = 30  # seconds
+        self.started_at = time.time()
+
+    def register_runner(self):
+        """
+        Register this runner as active in the system.
+        """
+        runner_data = {
+            'hostname': self.hostname,
+            'started_at': self.started_at,
+            'max_workers': self.max_workers,
+            'channels': self.channels,
+            'last_ping': time.time(),
+            'status': 'active'
+        }
+        self.manager.redis.hset(
+            self.manager.get_runners_key(),
+            self.hostname,
+            json.dumps(runner_data)
+        )
+        self.logger.info(f"Registered runner {self.hostname}")
+
+    def unregister_runner(self):
+        """
+        Unregister this runner from the active runners list.
+        """
+        self.manager.redis.hdel(self.manager.get_runners_key(), self.hostname)
+        self.logger.info(f"Unregistered runner {self.hostname}")
+
+    def update_runner_status(self, status_data=None):
+        """
+        Update the status of this runner.
+        """
+        if status_data is None:
+            status_data = {}
+
+        runner_data = {
+            'hostname': self.hostname,
+            'last_ping': time.time(),
+            'status': 'active',
+            'started_at': self.started_at,
+            'max_workers': self.max_workers,
+            'channels': self.channels,
+            **status_data
+        }
+        self.manager.redis.hset(
+            self.manager.get_runners_key(),
+            self.hostname,
+            json.dumps(runner_data)
+        )
+
+
+    def ping_runners(self):
+        """
+        Send ping messages to all active runners to check their status.
+        """
+        active_runners = self.manager.get_active_runners()
+        for hostname in active_runners.keys():
+            if hostname != self.hostname:  # Don't ping ourselves
+                ping_message = {
+                    'type': 'ping',
+                    'from': self.hostname,
+                    'timestamp': time.time()
+                }
+                runner_channel = f"runner_{hostname}"
+                self.manager.redis.publish(
+                    self.manager.get_channel_key(runner_channel),
+                    json.dumps(ping_message)
+                )
+
+    def handle_ping_request(self, message_data):
+        """
+        Handle incoming ping requests and send response.
+        """
+        ping_data = json.loads(message_data)
+        response = {
+            'type': 'ping_response',
+            'from': self.hostname,
+            'to': ping_data['from'],
+            'timestamp': time.time(),
+            'status': self.get_runner_status()
+        }
+
+        # Send response to the requesting runner's channel
+        requester_channel = f"runner_{ping_data['from']}"
+        self.manager.redis.publish(
+            self.manager.get_channel_key(requester_channel),
+            json.dumps(response)
+        )
+
+    def handle_ping_response(self, message_data):
+        """
+        Handle ping responses from other runners.
+        """
+        response_data = json.loads(message_data)
+        self.logger.info(f"Received ping response from {response_data['from']}")
+        # Update the runner's status in our active runners list
+        self.manager.redis.hset(
+            self.manager.get_runners_key(),
+            response_data['from'],
+            json.dumps(response_data['status'])
+        )
+
+    def get_runner_status(self):
+        """
+        Get the current status of this runner.
+
+        Returns:
+            dict: Status information for this runner.
+        """
+        active_threads = 0
+        if self.executor and hasattr(self.executor, '_threads'):
+            active_threads = len([t for t in self.executor._threads if t.is_alive()])
+
+        return {
+            'hostname': self.hostname,
+            'status': 'active',
+            'max_workers': self.max_workers,
+            'active_threads': active_threads,
+            'channels': self.channels,
+            'last_ping': time.time(),
+            'uptime': time.time() - getattr(self, 'start_time', time.time())
+        }
+
+    def start_ping_thread(self):
+        """
+        Start the background thread that periodically pings other runners.
+        """
+        def ping_loop():
+            while self.running:
+                try:
+                    self.ping_runners()
+                    self.update_runner_status()
+                    time.sleep(self.ping_interval)
+                except Exception as e:
+                    self.logger.error(f"Error in ping loop: {e}")
+                    time.sleep(5)
+
+        self.ping_thread = threading.Thread(target=ping_loop, daemon=True)
+        self.ping_thread.start()
+
+    def cleanup_stale_runners(self):
+        """
+        Remove runners that haven't been seen for a while.
+        """
+        cutoff_time = time.time() - (self.ping_interval * 3)  # 3 missed pings
+        active_runners = self.manager.get_active_runners()
+
+        for hostname, runner_data in active_runners.items():
+            last_ping = runner_data.get('last_ping', 0)
+            if last_ping < cutoff_time:
+                self.logger.info(f"Removing stale runner: {hostname}")
+                self.manager.redis.hdel(self.manager.get_runners_key(), hostname)
 
     def reset_running_tasks(self):
         """
@@ -38,8 +203,8 @@ class TaskEngine(daemon.Daemon):
         for channel in self.channels:
             for task_id in self.manager.get_running_ids(channel):
                 self.logger.info(f"moving task {task_id} from running to pending")
-                self.manager.remove_from_running(channel, task_id)
-                self.manager.add_to_pending(channel, task_id)
+                self.manager.remove_from_running(task_id, channel)
+                self.manager.add_to_pending(task_id, channel)
 
     def queue_pending_tasks(self):
         """
@@ -56,7 +221,23 @@ class TaskEngine(daemon.Daemon):
         Args:
             message (dict): A dictionary with message data containing task information.
         """
-        self.queue_task(message['data'].decode())
+        message_data = message['data'].decode()
+
+        # Check if this is a ping/status message
+        try:
+            parsed_message = json.loads(message_data)
+            if isinstance(parsed_message, dict) and 'type' in parsed_message:
+                if parsed_message['type'] == 'ping':
+                    self.handle_ping_request(message_data)
+                    return
+                elif parsed_message['type'] == 'ping_response':
+                    self.handle_ping_response(message_data)
+                    return
+        except (json.JSONDecodeError, TypeError):
+            pass
+
+        # If not a ping message, treat as a task
+        self.queue_task(message_data)
 
     def on_run_task(self, task_id):
         """
@@ -71,7 +252,10 @@ class TaskEngine(daemon.Daemon):
         if not task_data:
             # this task has expired or no longer exists
             self.logger.info(f"Task {task_id} has expired or no longer exists")
-            tman.remove_from_pending(task_id)
+            metrics.record("tasks_expired", category="tasks")
+            # try and remove any pending dead tasks
+            self.manager.channels = self.channels
+            self.manager.take_out_the_dead(local=True)
             return
         self.logger.info(f"Executing task {task_id}")
         function_path = task_data.get('function')
@@ -83,15 +267,28 @@ class TaskEngine(daemon.Daemon):
 
         try:
             task_data.started_at = time.time()
-            func(task_data)
+            task_data._thread_id = threading.current_thread().ident
+            tdata = task_data.get("data", {})
+            if tdata and "args" in tdata and "kwargs" in tdata:
+                args = tdata["args"]
+                kwargs = tdata["kwargs"]
+                # self.logger.info(f"Executing task {task_id} with args {args} and kwargs {kwargs}")
+                func(*args, **kwargs)
+            else:
+                # self.logger.info(f"Executing task {task_id} with no arguments")
+                func(task_data)
             task_data.completed_at = time.time()
             task_data.elapsed_time = task_data.completed_at - task_data.started_at
+            if "_thread_id" in task_data:
+                del task_data["_thread_id"]
             tman.save_task(task_data)
             tman.add_to_completed(task_data)
+            metrics.record("tasks_completed", category="tasks")
             self.logger.info(f"Task {task_id} completed after {task_data.elapsed_time} seconds")
         except Exception as e:
-            self.logger.error(f"Error executing task {task_id}: {str(e)}")
+            self.logger.exception(f"Error executing task {task_id}: {str(e)}")
             tman.add_to_errors(task_data, str(e))
+            metrics.record("tasks_errors", category="tasks")
         finally:
             tman.remove_from_running(task_id, task_data.channel)
 
@@ -105,31 +302,68 @@ class TaskEngine(daemon.Daemon):
         self.logger.info(f"adding task {task_id}")
         self.executor.submit(self.on_run_task, task_id)
 
-    def wait_for_all_tasks_to_complete(self, timeout=30):
+
+    def _clear_queued_tasks(self):
+        import queue
+        q = self.executor._work_queue
+        removed = 0
+        try:
+            while True:
+                q.get_nowait()
+                removed += 1
+        except queue.Empty:
+            pass
+        return removed
+
+    def _wait_for_active_tasks(self, timeout=5.0):
         """
-        Wait for all tasks submitted to the executor to complete.
+        Waits up to `timeout` seconds for active executor threads to finish.
+        Returns True if all threads completed, False if timeout hit.
         """
-        self.executor.shutdown(wait=True, timeout=timeout)
-        # Check if there are still active threads
-        active_threads = [thread for thread in self.executor._threads if thread.is_alive()]
-        if active_threads:
-            self.logger.warning(f"shutdown issue, {len(active_threads)} tasks exceeded timeout")
-            self.executor.shutdown(wait=False)  # Stop accepting new tasks
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            active = self.manager.get_all_running_ids(local=True)
+            if len(active) == 0:
+                return True
+            time.sleep(0.01)
+        return False
+
+    def wait_for_all_tasks_to_complete(self, timeout=5):
+        """
+        Wait for all tasks submitted to the executor to complete with graceful degradation.
+        """
+        if not self.executor:
+            return
+
+        self.logger.info(f"Initiating graceful shutdown with {timeout}s timeout")
+        self.executor.shutdown(wait=False)
+        self._clear_queued_tasks()
+        result = self._wait_for_active_tasks(timeout)
+        if not result:
+            self.logger.warning("Timeout reached while waiting for active tasks to complete")
+        return result
 
     def start_listening(self):
         """
         Listen for messages on the subscribed channels and handle them as they arrive.
         """
         self.logger.info("starting with channels...", self.channels)
+        self.start_time = time.time()
+        self.register_runner()
+        self.manager.take_out_the_dead(local=True)
         self.reset_running_tasks()
         self.queue_pending_tasks()
+        self.start_ping_thread()
+
         pubsub = self.manager.redis.pubsub()
         channel_keys = {self.manager.get_channel_key(channel): self.handle_message for channel in self.channels}
         pubsub.subscribe(**channel_keys)
+
         for message in pubsub.listen():
             if not self.running:
                 self.logger.info("shutting down, waiting for tasks to complete")
                 self.wait_for_all_tasks_to_complete()
+                self.unregister_runner()
                 self.logger.info("shutdown complete")
                 return
             if message['type'] != 'message':
@@ -154,6 +388,7 @@ def get_args():
     parser.add_argument("--start", action="store_true", help="Start the daemon")
     parser.add_argument("--stop", action="store_true", help="Stop the daemon")
     parser.add_argument("--foreground", "-f", action="store_true", help="Run in foreground mode")
+    parser.add_argument("--status", action="store_true", help="Show status of all runners")
     parser.add_argument("-v", "--verbose", action="store_true",
                         help="Enable verbose logging")
     return parser, parser.parse_args()
@@ -162,8 +397,18 @@ def get_args():
 def main():
     from mojo.helpers.settings import settings
     parser, args = get_args()
-    daemon = TaskEngine(settings.TASKIT_CHANNELS)
-    if args.start:
+    daemon = TaskEngine(settings.TASK_CHANNELS)
+
+    if args.status:
+        runners = daemon.manager.get_active_runners()
+        if runners:
+            print("Active TaskEngine Runners:")
+            for hostname, data in runners.items():
+                print(f"  {hostname}: {data.get('status', 'unknown')} "
+                      f"(last ping: {time.time() - data.get('last_ping', 0):.1f}s ago)")
+        else:
+            print("No active runners found")
+    elif args.start:
         daemon.start()
     elif args.stop:
         daemon.stop()
@@ -172,3 +417,23 @@ def main():
         daemon.run()
     else:
         parser.print_help()
+
+
+
+def kill_thread(thread):
+    import ctypes
+    if not thread.is_alive():
+        return False
+
+    tid = thread.ident
+    if tid is None:
+        return False
+
+    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+        ctypes.c_long(tid), ctypes.py_object(SystemExit)
+    )
+    if res > 1:
+        # Undo if multiple threads were affected
+        ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), 0)
+        return False
+    return True

mojo/apps/tasks/task.py

@@ -0,0 +1,99 @@
+from objict import objict
+import time
+
+
+class Task(objict):
+    """
+    Task model for the Django Mojo task system.
+
+    This class represents a task that can be queued, executed, and tracked
+    through various states (pending, running, completed, error, cancelled).
+    """
+
+    def __init__(self, id=None, function=None, data=None, channel="default",
+                 expires=None, created=None, status="pending", error=None,
+                 completed_at=None, **kwargs):
+        """
+        Initialize a new Task instance.
+
+        Args:
+            id (str): Unique identifier for the task
+            function (str): Function name to be executed
+            data (dict): Data to be passed to the function
+            channel (str): Channel name for task routing
+            expires (float): Expiration timestamp
+            created (float): Creation timestamp
+            status (str): Current task status
+            error (str): Error message if task failed
+            completed_at (float): Completion timestamp
+            **kwargs: Additional attributes
+        """
+        super().__init__(**kwargs)
+
+        self.id = id
+        self.function = function
+        self.data = data or {}
+        self.channel = channel
+        self.expires = expires
+        self.created = created or time.time()
+        self.status = status
+        self.error = error
+        self.completed_at = completed_at
+
+    def is_expired(self):
+        """
+        Check if the task has expired.
+
+        Returns:
+            bool: True if task has expired, False otherwise
+        """
+        if self.expires is None:
+            return False
+        return time.time() > self.expires
+
+    def is_pending(self):
+        """Check if task is in pending state."""
+        return self.status == "pending"
+
+    def is_running(self):
+        """Check if task is in running state."""
+        return self.status == "running"
+
+    def is_completed(self):
+        """Check if task is in completed state."""
+        return self.status == "completed"
+
+    def is_error(self):
+        """Check if task is in error state."""
+        return self.status == "error"
+
+    def is_cancelled(self):
+        """Check if task is in cancelled state."""
+        return self.status == "cancelled"
+
+    def mark_as_running(self):
+        """Mark task as running."""
+        self.status = "running"
+
+    def mark_as_completed(self):
+        """Mark task as completed."""
+        self.status = "completed"
+        self.completed_at = time.time()
+
+    def mark_as_error(self, error_message):
+        """Mark task as error with error message."""
+        self.status = "error"
+        self.error = error_message
+
+    def mark_as_cancelled(self):
+        """Mark task as cancelled."""
+        self.status = "cancelled"
+
+    def __str__(self):
+        """String representation of the task."""
+        return f"Task({self.id}, {self.function}, {self.status})"
+
+    def __repr__(self):
+        """Detailed string representation of the task."""
+        return (f"Task(id='{self.id}', function='{self.function}', "
+                f"status='{self.status}', channel='{self.channel}')")

mojo/apps/tasks/tq_handlers.py

@@ -12,3 +12,121 @@ def run_error_task(task):
     logger.info("Running error task with data", task)
     time.sleep(2)
     raise Exception("Example error")
+
+
+def run_quick_task(task):
+    """Quick task for testing - completes immediately"""
+    logger.info("Running quick task with data", task)
+    return {"status": "completed", "data": task.data}
+
+
+def run_slow_task(task):
+    """Slow task for testing - takes 10 seconds"""
+    logger.info("Running slow task with data", task)
+    time.sleep(10)
+    return {"status": "completed", "duration": 10}
+
+
+def run_args_kwargs_task(*args, **kwargs):
+    """Task that receives args and kwargs directly"""
+    logger.info(f"Running args/kwargs task with args: {args}, kwargs: {kwargs}")
+    return {"args": args, "kwargs": kwargs}
+
+
+def run_data_processing_task(task):
+    """Task that processes data and returns results"""
+    logger.info("Running data processing task")
+    data = task.data
+    if not isinstance(data, dict):
+        raise ValueError("Data must be a dictionary")
+    
+    result = {
+        "processed": True,
+        "input_keys": list(data.keys()),
+        "total_items": len(data)
+    }
+    return result
+
+
+def run_counter_task(task):
+    """Task that increments a counter - for testing state changes"""
+    logger.info("Running counter task")
+    count = task.data.get("count", 0)
+    new_count = count + 1
+    logger.info(f"Counter incremented from {count} to {new_count}")
+    return {"count": new_count}
+
+
+def run_timeout_task(task):
+    """Task that times out - for testing timeout scenarios"""
+    duration = task.data.get("duration", 60)
+    logger.info(f"Running timeout task for {duration} seconds")
+    time.sleep(duration)
+    return {"completed": True}
+
+
+def run_memory_task(task):
+    """Task that uses memory - for testing resource usage"""
+    logger.info("Running memory task")
+    size = task.data.get("size", 1000000)  # 1MB default
+    data = bytearray(size)
+    logger.info(f"Allocated {size} bytes")
+    return {"allocated_bytes": size}
+
+
+def run_conditional_error_task(task):
+    """Task that conditionally raises an error based on input"""
+    logger.info("Running conditional error task")
+    should_error = task.data.get("should_error", False)
+    error_message = task.data.get("error_message", "Conditional error occurred")
+    
+    if should_error:
+        raise Exception(error_message)
+    
+    return {"status": "success", "should_error": should_error}
+
+
+def run_nested_data_task(task):
+    """Task that works with nested data structures"""
+    logger.info("Running nested data task")
+    data = task.data
+    
+    if "nested" not in data:
+        raise ValueError("Missing 'nested' key in data")
+    
+    nested = data["nested"]
+    result = {
+        "original": nested,
+        "keys": list(nested.keys()) if isinstance(nested, dict) else None,
+        "length": len(nested) if hasattr(nested, '__len__') else None
+    }
+    
+    return result
+
+
+# Test async task handlers
+def async_quick_task(message="Hello"):
+    """Async task handler for testing decorator"""
+    logger.info(f"Async quick task: {message}")
+    return f"Processed: {message}"
+
+
+def async_slow_task(duration=5):
+    """Async slow task handler for testing decorator"""
+    logger.info(f"Async slow task sleeping for {duration} seconds")
+    time.sleep(duration)
+    return f"Completed after {duration} seconds"
+
+
+def async_error_task(should_error=True, message="Async error"):
+    """Async error task handler for testing decorator"""
+    logger.info(f"Async error task - should_error: {should_error}")
+    if should_error:
+        raise Exception(message)
+    return "No error raised"
+
+
+def async_args_task(*args, **kwargs):
+    """Async task that tests args and kwargs handling"""
+    logger.info(f"Async args task - args: {args}, kwargs: {kwargs}")
+    return {"received_args": args, "received_kwargs": kwargs}

mojo/decorators/auth.py

@@ -1,5 +1,8 @@
 from functools import wraps
 import mojo.errors
+from mojo.helpers import logit
+
+logger = logit.get_logger("error", "error.log")
 
 def requires_perms(*required_perms):
     def decorator(func):
@@ -7,7 +10,9 @@ def requires_perms(*required_perms):
         def wrapper(request, *args, **kwargs):
             if not request.user.is_authenticated:
                 raise mojo.errors.PermissionDeniedException()
-            if not request.user.has_permission(required_perms):
+            perms = set(required_perms)
+            if not request.user.has_permission(perms):
+                logger.error(f"{request.user.username} is missing {perms}")
                 raise mojo.errors.PermissionDeniedException()
             return func(request, *args, **kwargs)
         return wrapper

mojo/decorators/http.py

@@ -32,7 +32,12 @@ def dispatcher(request, *args, **kwargs):
     rest.ACTIVE_REQUEST = request
     key = kwargs.pop('__mojo_rest_root_key__', None)
     if "group" in request.DATA:
-        request.group = modules.get_model_instance("account", "Group", int(request.DATA.group))
+        try:
+            request.group = modules.get_model_instance("account", "Group", int(request.DATA.group))
+            if request.group is not None:
+                request.group.touch()
+        except ValueError:
+            return JsonResponse({"error": "Invalid group ID", "code": 400}, status=400)
     method_key = f"{key}__{request.method}"
     if method_key not in URLPATTERN_METHODS:
         method_key = f"{key}__ALL"
@@ -66,7 +71,7 @@ def dispatch_error_handler(func):
                 metrics.record("api_errors", category="mojo_api", min_granularity=API_METRICS_GRANULARITY)
             # logger.exception(f"Unhandled REST Exception: {request.path}")
             logger.exception(f"Error: {str(err)}, Path: {request.path}, IP: {request.META.get('REMOTE_ADDR')}")
-            return JsonResponse({"error": str(err) }, status=500)
+            return JsonResponse({"error": str(err), "code": 500 }, status=500)
 
     return wrapper