353 lines
10 KiB
Python
353 lines
10 KiB
Python
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published
|
|
# by the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from typing import Any, Callable, Dict, List
|
|
|
|
import schedule
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# TODO - move this File where all jobs are persisted
|
|
JOBS_FILE = f"{os.getenv('WORKING_DIR')}\\scheduling\\jobs.json"
|
|
|
|
# Ensure directory exists
|
|
os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True)
|
|
|
|
# Registry of functions that can be scheduled
|
|
FUNCTION_MAP: Dict[str, Callable] = {}
|
|
|
|
# -------------------------------
|
|
# Function Registration
|
|
# -------------------------------
|
|
|
|
|
|
def register_function(name: str, func: Callable):
|
|
"""
|
|
Register a function so it can be called by name later.
|
|
Example:
|
|
register_function("say_hello", say_hello)
|
|
"""
|
|
FUNCTION_MAP[name] = func
|
|
|
|
|
|
# -------------------------------
|
|
# Persistence Helpers
|
|
# -------------------------------
|
|
|
|
|
|
def load_jobs() -> List[Dict[str, Any]]:
|
|
"""Load jobs from the JSON file, or return [] if none exist."""
|
|
if not os.path.exists(JOBS_FILE):
|
|
return []
|
|
with open(JOBS_FILE, "r") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def _atomic_save(path: str, data: Any):
|
|
"""Write JSON atomically to avoid partial writes."""
|
|
tmp = f"{path}.tmp"
|
|
with open(tmp, "w") as f:
|
|
json.dump(data, f, indent=4)
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def save_jobs(jobs: List[Dict[str, Any]]):
|
|
"""Save jobs to the JSON file (overwrite)."""
|
|
_atomic_save(JOBS_FILE, jobs)
|
|
|
|
|
|
# -------------------------------
|
|
# Uniqueness Helpers
|
|
# -------------------------------
|
|
|
|
|
|
def job_in_store(job_id: str) -> bool:
|
|
"""Check if a job id exists in the persisted JSON file."""
|
|
return any(j.get("id") == job_id for j in load_jobs())
|
|
|
|
|
|
def job_in_scheduler(job_id: str) -> bool:
|
|
"""
|
|
Check if a job with this tag exists in the in-memory scheduler.
|
|
Uses schedule.get_jobs(tag=...) if available, otherwise scans tags.
|
|
"""
|
|
try:
|
|
jobs = schedule.get_jobs(tag=job_id) # schedule >= 1.2.0
|
|
return len(jobs) > 0
|
|
except TypeError:
|
|
# Fallback for older versions
|
|
return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs)
|
|
|
|
|
|
def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
|
|
"""
|
|
Ensure the job_id is unique across persistence and in-memory schedule.
|
|
on_conflict:
|
|
- "error": raise ValueError if exists.
|
|
- "skip" : print and return False.
|
|
- "replace": remove existing (in-memory + JSON), then continue.
|
|
"""
|
|
exists = job_in_store(job_id) or job_in_scheduler(job_id)
|
|
if not exists:
|
|
return True
|
|
if on_conflict == "error":
|
|
raise ValueError(f"Job id '{job_id}' already exists.")
|
|
elif on_conflict == "skip":
|
|
logger.info(f"Job '{job_id}' already exists. Skipping creation.")
|
|
return False
|
|
elif on_conflict == "replace":
|
|
# Clear from scheduler
|
|
schedule.clear(job_id)
|
|
# Remove from persistence
|
|
jobs = [j for j in load_jobs() if j.get("id") != job_id]
|
|
save_jobs(jobs)
|
|
return True
|
|
else:
|
|
raise ValueError(f"Unsupported on_conflict policy: {on_conflict}")
|
|
|
|
|
|
# -------------------------------
|
|
# Internal scheduling (no persistence)
|
|
# -------------------------------
|
|
|
|
|
|
def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None):
|
|
args = args or []
|
|
kwargs = kwargs or {}
|
|
|
|
def job_wrapper():
|
|
"""Executes the job once, then removes it."""
|
|
if func_name not in FUNCTION_MAP:
|
|
logger.error(f"Function '{func_name}' is not registered.")
|
|
return
|
|
FUNCTION_MAP[func_name](*args, **kwargs)
|
|
# Remove from persistence
|
|
jobs = load_jobs()
|
|
jobs = [j for j in jobs if j["id"] != job_id]
|
|
save_jobs(jobs)
|
|
# Clear from in-memory schedule
|
|
schedule.clear(job_id)
|
|
|
|
delay_seconds = run_at_timestamp - time.time()
|
|
if delay_seconds <= 0:
|
|
logger.info(f"Job {job_id} scheduled in the past. Skipping.")
|
|
return
|
|
# Schedule via schedule library
|
|
schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id)
|
|
|
|
|
|
def _schedule_recurring(
|
|
job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None
|
|
):
|
|
args = args or []
|
|
kwargs = kwargs or {}
|
|
|
|
def job_wrapper():
|
|
if func_name not in FUNCTION_MAP:
|
|
logger.error(f"Function '{func_name}' is not registered.")
|
|
return
|
|
FUNCTION_MAP[func_name](*args, **kwargs)
|
|
|
|
if unit == "seconds":
|
|
schedule.every(interval).seconds.do(job_wrapper).tag(job_id)
|
|
elif unit == "minutes":
|
|
schedule.every(interval).minutes.do(job_wrapper).tag(job_id)
|
|
elif unit == "hours":
|
|
schedule.every(interval).hours.do(job_wrapper).tag(job_id)
|
|
elif unit == "days":
|
|
schedule.every(interval).days.do(job_wrapper).tag(job_id)
|
|
else:
|
|
raise ValueError(f"Unsupported unit: {unit}")
|
|
|
|
|
|
# -------------------------------
|
|
# Public APIs (with uniqueness + persistence)
|
|
# -------------------------------
|
|
|
|
|
|
def run_once_job(
|
|
job_id: str,
|
|
func_name: str,
|
|
run_at_timestamp: float,
|
|
args=None,
|
|
kwargs=None,
|
|
*,
|
|
replace: bool = False,
|
|
persist: bool = True,
|
|
):
|
|
"""
|
|
Schedule a job to run once at a specific timestamp.
|
|
replace: if True, replace existing job with same id; otherwise print and skip.
|
|
persist: if False, do not write to JSON (used by reload_jobs()).
|
|
"""
|
|
if persist:
|
|
policy = "replace" if replace else "skip"
|
|
if not ensure_unique(job_id, on_conflict=policy):
|
|
return
|
|
elif job_in_scheduler(job_id):
|
|
schedule.clear(job_id)
|
|
_schedule_once(job_id, func_name, run_at_timestamp, args, kwargs)
|
|
if persist:
|
|
jobs = [j for j in load_jobs() if j["id"] != job_id]
|
|
jobs.append(
|
|
{
|
|
"id": job_id,
|
|
"type": "once",
|
|
"run_at": run_at_timestamp,
|
|
"function": func_name,
|
|
"args": args or [],
|
|
"kwargs": kwargs or {},
|
|
}
|
|
)
|
|
save_jobs(jobs)
|
|
|
|
|
|
def recurring_job(
|
|
job_id: str,
|
|
func_name: str,
|
|
interval: int,
|
|
unit: str,
|
|
args=None,
|
|
kwargs=None,
|
|
*,
|
|
replace: bool = False,
|
|
persist: bool = True,
|
|
):
|
|
"""
|
|
Schedule a recurring job.
|
|
replace: if True, replace existing job with same id; otherwise print and skip.
|
|
persist: if False, do not write to JSON (used by reload_jobs()).
|
|
"""
|
|
if persist:
|
|
policy = "replace" if replace else "skip"
|
|
if not ensure_unique(job_id, on_conflict=policy):
|
|
return
|
|
elif job_in_scheduler(job_id):
|
|
schedule.clear(job_id)
|
|
_schedule_recurring(job_id, func_name, interval, unit, args, kwargs)
|
|
if persist:
|
|
jobs = [j for j in load_jobs() if j["id"] != job_id]
|
|
jobs.append(
|
|
{
|
|
"id": job_id,
|
|
"type": "recurring",
|
|
"interval": interval,
|
|
"unit": unit,
|
|
"function": func_name,
|
|
"args": args or [],
|
|
"kwargs": kwargs or {},
|
|
}
|
|
)
|
|
save_jobs(jobs)
|
|
|
|
|
|
def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float = 1.0):
|
|
"""
|
|
Find all jobs whose ID contains the given PID substring and reschedule them to run sooner.
|
|
"""
|
|
jobs = load_jobs()
|
|
matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")]
|
|
|
|
if not matched_jobs:
|
|
logger.info(f"No jobs found containing PID substring '{pid_substring}'.")
|
|
return
|
|
|
|
logger.info(f"Found {len(matched_jobs)} job(s) containing '{pid_substring}':")
|
|
|
|
for job in matched_jobs:
|
|
job_id = job["id"]
|
|
logger.debug(f" - Prioritizing job: {job_id}")
|
|
|
|
# Clear existing job from scheduler
|
|
schedule.clear(job_id)
|
|
|
|
# Reschedule based on job type
|
|
if job["type"] == "once":
|
|
run_once_job(
|
|
job_id,
|
|
job["function"],
|
|
time.time() + new_delay_seconds,
|
|
job.get("args"),
|
|
job.get("kwargs"),
|
|
replace=True,
|
|
persist=True,
|
|
)
|
|
elif job["type"] == "recurring":
|
|
recurring_job(
|
|
job_id,
|
|
job["function"],
|
|
job["interval"],
|
|
job["unit"],
|
|
job.get("args"),
|
|
job.get("kwargs"),
|
|
replace=True,
|
|
persist=True,
|
|
)
|
|
else:
|
|
logger.warning(f"Unknown job type for job '{job_id}'")
|
|
|
|
|
|
# -------------------------------
|
|
# Reload Saved Jobs
|
|
# -------------------------------
|
|
|
|
|
|
def reload_jobs():
|
|
"""Reload jobs from JSON and reschedule them (no re-persist)."""
|
|
jobs = load_jobs()
|
|
for job in jobs:
|
|
if job["type"] == "once":
|
|
if job["run_at"] > time.time():
|
|
run_once_job(
|
|
job["id"],
|
|
job["function"],
|
|
job["run_at"],
|
|
job.get("args"),
|
|
job.get("kwargs"),
|
|
persist=False,
|
|
)
|
|
elif job["type"] == "recurring":
|
|
recurring_job(
|
|
job["id"],
|
|
job["function"],
|
|
job["interval"],
|
|
job["unit"],
|
|
job.get("args"),
|
|
job.get("kwargs"),
|
|
persist=False,
|
|
)
|
|
|
|
|
|
# -------------------------------
|
|
# Scheduler Loop
|
|
# -------------------------------
|
|
|
|
|
|
def start_scheduler():
|
|
"""
|
|
Start the scheduler loop (blocking).
|
|
Call this once in main to begin.
|
|
"""
|
|
try:
|
|
while True:
|
|
schedule.run_pending()
|
|
time.sleep(0.5)
|
|
except KeyboardInterrupt:
|
|
logger.critical("Scheduler stopped.")
|