iot-dboard/backend/scheduler.py

"""Background scheduler for periodic tasks.

All user-scoped jobs iterate over every active user (with that user's
timezone/context active) so each user's alarms, reminders, calendars,
pomodoro, countdowns and weather are serviced independently.
"""
import threading

from apscheduler.schedulers.background import BackgroundScheduler
from backend.calendar_mgr import sync_all_subscriptions, get_upcoming_reminders
from backend.weather_mgr import fetch_weather, fetch_hourly
from backend.alarm_mgr import check_alarms
from backend.countdown_mgr import check_countdowns
from backend.pomodoro_mgr import check_and_trigger
from backend.models import get_db
from backend.time_util import set_thread_user, clear_thread_user


_scheduler: BackgroundScheduler | None = None

# ─── Scheduler single-instance guard ─────────
# With Gunicorn every worker imports the app; the background scheduler must run
# in exactly ONE process. A PostgreSQL session-level advisory lock decides which
# worker owns it (the lock is held by an open connection and released when that
# process exits, at which point another worker takes over).
_SCHEDULER_LOCK_KEY = 0x494F5444424F4152  # "IOTDBOAR"
_scheduler_lock_conn = None
_watchdog: "threading.Thread | None" = None
_WATCHDOG_INTERVAL = 30  # seconds


def _lock_connect(db_url: str):
    """Open the lock connection with TCP keepalive + connect timeout.
    Without keepalives a network blip could leave PostgreSQL holding the lock
    for ~2h (default tcp keepalive), blocking every worker from taking over.
    (libpq URI query params do NOT accept tcp_keepalives_* — pass them as
    psycopg2 keyword args instead.)"""
    import psycopg2
    return psycopg2.connect(
        db_url,
        connect_timeout=5,
        keepalives_idle=30,
        keepalives_interval=10,
        keepalives_count=3,
    )


def try_acquire_scheduler_lock(db_url: str) -> bool:
    """Try to become the scheduler owner. Returns True for the winner."""
    global _scheduler_lock_conn
    conn = None
    try:
        conn = _lock_connect(db_url)
        cur = conn.cursor()
        cur.execute("SELECT pg_try_advisory_lock(%s)", (_SCHEDULER_LOCK_KEY,))
        ok = bool(cur.fetchone()[0])
        cur.close()
        if ok:
            _scheduler_lock_conn = conn
            return True
        conn.close()
        return False
    except Exception:
        try:
            if _scheduler_lock_conn:
                _scheduler_lock_conn.close()
                _scheduler_lock_conn = None
        except Exception:
            pass
        try:
            if conn is not None:
                conn.close()  # close the partially-used local connection
        except Exception:
            pass
        return False


def release_scheduler_lock():
    """Release the advisory lock (also closes the holding connection)."""
    global _scheduler_lock_conn
    if _scheduler_lock_conn:
        try:
            _scheduler_lock_conn.close()
        except Exception:
            pass
        _scheduler_lock_conn = None


def _watchdog_loop(app, db_url: str):
    """Keeps the scheduler alive across worker restarts/freezes:
    - lock owner: periodically checks the lock connection is healthy;
    - non-owner: periodically retries to take over the lock.
    """
    global _scheduler
    import time
    while True:
        time.sleep(_WATCHDOG_INTERVAL)
        try:
            if _scheduler_lock_conn is not None:
                # Owner: health check on the lock connection, and make sure the
                # scheduler thread itself is still running (start failure or a
                # dead thread must not hold the lock forever).
                cur = _scheduler_lock_conn.cursor()
                cur.execute("SELECT 1")
                cur.fetchone()
                cur.close()
                if _scheduler is None or not _scheduler.running:
                    if _scheduler:
                        try:
                            stop_scheduler()
                        except Exception:
                            pass
                    start_scheduler(app)
                    app.logger.info("Scheduler restarted (lock held but not running)")
            elif _scheduler is None:
                # Non-owner: try to take over (previous owner died/froze).
                if try_acquire_scheduler_lock(db_url):
                    start_scheduler(app)
                    app.logger.info("Scheduler started (watchdog takeover)")
        except Exception:
            # Lock connection became unhealthy → drop the lock and re-race.
            release_scheduler_lock()
            if _scheduler:
                try:
                    stop_scheduler()
                except Exception:
                    pass
            try:
                if try_acquire_scheduler_lock(db_url):
                    start_scheduler(app)
                    app.logger.info("Scheduler re-acquired lock after watchdog failure")
            except Exception:
                pass


def start_watchdog(app, db_url: str):
    """Start the scheduler watchdog thread (idempotent)."""
    global _watchdog
    if _watchdog is not None and _watchdog.is_alive():
        return
    _watchdog = threading.Thread(target=_watchdog_loop, args=(app, db_url), daemon=True)
    _watchdog.start()


def _each_user(app, fn):
    """Run fn(user_id) for every active user inside the app context."""
    with app.app_context():
        db = get_db(app.config.get("DB_PATH", "data.db"))
        rows = db.execute("SELECT id FROM users WHERE status='active' ORDER BY id").fetchall()
        db.close()
        for u in rows:
            uid = u["id"]
            try:
                set_thread_user(uid)
                fn(uid)
            except Exception as e:
                app.logger.error(f"Scheduler error for user {uid}: {e}")
            finally:
                clear_thread_user()


def start_scheduler(app):
    global _scheduler
    _scheduler = BackgroundScheduler()

    @_scheduler.scheduled_job("interval", seconds=10, id="check_alarms")
    def _check_alarms():
        _each_user(app, check_alarms)

    @_scheduler.scheduled_job("interval", minutes=1, id="check_reminders")
    def _check_reminders():
        def _run(uid):
            due = get_upcoming_reminders(uid)
            if due:
                app.event_bus.publish("reminder_due", {"events": due}, user_id=uid)
                app.logger.info(f"Reminders due for user {uid}: {len(due)} events")
        _each_user(app, _run)

    @_scheduler.scheduled_job("interval", minutes=30, id="sync_calendars")
    def _sync_calendars():
        with app.app_context():
            sync_all_subscriptions()

    @_scheduler.scheduled_job("cron", minute="0,30", id="fetch_weather")
    def _fetch_weather():
        _each_user(app, fetch_weather)

    @_scheduler.scheduled_job("interval", seconds=1, id="check_pomodoro")
    def _check_pomodoro():
        _each_user(app, check_and_trigger)

    @_scheduler.scheduled_job("cron", minute="1", id="fetch_hourly")
    def _fetch_hourly():
        _each_user(app, fetch_hourly)

    @_scheduler.scheduled_job("interval", seconds=1, id="check_countdowns")
    def _check_countdowns():
        def _run(uid):
            triggered = check_countdowns(uid)
            if triggered:
                app.logger.info(f"Countdowns triggered for user {uid}: {[c['label'] for c in triggered]}")
        _each_user(app, _run)

    _scheduler.start()


def stop_scheduler():
    global _scheduler
    if _scheduler:
        try:
            _scheduler.shutdown(wait=False)
        finally:
            _scheduler = None