bigbrother/core/scheduler.py

#!/usr/bin/env python3
"""Cron-like task scheduler with jitter.

Modules register periodic tasks with an interval and optional jitter range.
Tasks execute in a thread pool.  Used for PCAP rotation, health checks,
baseline captures, data pruning, resource checks.
"""

import time
import random
import logging
import threading
from dataclasses import dataclass, field
from typing import Callable, Optional
from concurrent.futures import ThreadPoolExecutor

logger = logging.getLogger("bb.scheduler")


@dataclass
class ScheduledTask:
    """A registered periodic task."""
    name: str
    callback: Callable
    interval: float          # seconds between executions
    jitter: float = 0.0      # max random jitter added to interval (seconds)
    run_immediately: bool = False
    enabled: bool = True

    # Runtime state
    last_run: float = 0.0
    next_run: float = 0.0
    run_count: int = 0
    error_count: int = 0
    _last_error: Optional[str] = None


class Scheduler:
    """Thread-safe periodic task scheduler with jitter support.

    Tasks are checked every second and dispatched to a thread pool.
    Jitter is re-randomized on each scheduling cycle to avoid patterns.
    """

    def __init__(self, max_workers: int = 4):
        self._tasks: dict[str, ScheduledTask] = {}
        self._lock = threading.Lock()
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._pool = ThreadPoolExecutor(max_workers=max_workers,
                                         thread_name_prefix="bb-sched")

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def start(self) -> None:
        """Start the scheduler loop."""
        if self._running:
            return
        self._running = True
        # Initialize next_run for tasks that want immediate execution
        now = time.time()
        with self._lock:
            for task in self._tasks.values():
                if task.run_immediately:
                    task.next_run = now
                elif task.next_run == 0.0:
                    task.next_run = now + task.interval + random.uniform(0, task.jitter)

        self._thread = threading.Thread(
            target=self._scheduler_loop, daemon=True, name="bb-scheduler"
        )
        self._thread.start()
        logger.info("Scheduler started (%d tasks)", len(self._tasks))

    def stop(self) -> None:
        """Stop the scheduler and wait for running tasks to finish."""
        self._running = False
        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=5.0)
        self._pool.shutdown(wait=False)
        logger.info("Scheduler stopped")

    # ------------------------------------------------------------------
    # Task management
    # ------------------------------------------------------------------

    def register(self, name: str, callback: Callable, interval: float,
                 jitter: float = 0.0, run_immediately: bool = False) -> None:
        """Register a periodic task.

        Args:
            name: Unique task name.
            callback: Callable (no arguments) to execute.
            interval: Base interval in seconds.
            jitter: Maximum random seconds added to interval.
            run_immediately: If True, run as soon as the scheduler starts.
        """
        now = time.time()
        task = ScheduledTask(
            name=name,
            callback=callback,
            interval=interval,
            jitter=jitter,
            run_immediately=run_immediately,
            next_run=now if run_immediately else now + interval + random.uniform(0, jitter),
        )
        with self._lock:
            self._tasks[name] = task
        logger.debug("Registered task: %s (interval=%.1fs, jitter=%.1fs)",
                      name, interval, jitter)

    def unregister(self, name: str) -> None:
        """Remove a task."""
        with self._lock:
            self._tasks.pop(name, None)

    def enable(self, name: str) -> None:
        """Enable a disabled task."""
        with self._lock:
            task = self._tasks.get(name)
            if task:
                task.enabled = True

    def disable(self, name: str) -> None:
        """Disable a task (stays registered but won't execute)."""
        with self._lock:
            task = self._tasks.get(name)
            if task:
                task.enabled = False

    def run_now(self, name: str) -> None:
        """Trigger immediate execution of a task."""
        with self._lock:
            task = self._tasks.get(name)
            if task:
                task.next_run = time.time()

    # ------------------------------------------------------------------
    # Scheduler loop
    # ------------------------------------------------------------------

    def _scheduler_loop(self) -> None:
        """Check tasks every second, dispatch due tasks to thread pool."""
        while self._running:
            now = time.time()
            with self._lock:
                tasks = list(self._tasks.values())

            for task in tasks:
                if not task.enabled:
                    continue
                if now >= task.next_run:
                    self._pool.submit(self._execute_task, task)
                    # Schedule next run with fresh jitter
                    jitter = random.uniform(0, task.jitter) if task.jitter > 0 else 0
                    task.next_run = now + task.interval + jitter

            time.sleep(1.0)

    def _execute_task(self, task: ScheduledTask) -> None:
        """Execute a task and update bookkeeping."""
        try:
            task.callback()
            task.run_count += 1
            task.last_run = time.time()
        except Exception as e:
            task.error_count += 1
            task._last_error = str(e)
            logger.exception("Scheduler task %s failed", task.name)

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    def get_tasks(self) -> dict:
        """Return info about all registered tasks."""
        with self._lock:
            return {
                name: {
                    "interval": t.interval,
                    "jitter": t.jitter,
                    "enabled": t.enabled,
                    "run_count": t.run_count,
                    "error_count": t.error_count,
                    "last_run": t.last_run,
                    "next_run": t.next_run,
                    "last_error": t._last_error,
                }
                for name, t in self._tasks.items()
            }