From 13139d7dea1ed0850cea430f7d6d6199326c86c4 Mon Sep 17 00:00:00 2001 From: = <=> Date: Wed, 10 Sep 2025 13:59:31 -0400 Subject: [PATCH] Fixed adding duplicate jobs, now gracefully handles --- AirlockTools.py | 2 +- utils/otpfunctions.py | 20 ++-- utils/perstscheduler.py | 237 ++++++++++++++++++++++++---------------- 3 files changed, 157 insertions(+), 102 deletions(-) diff --git a/AirlockTools.py b/AirlockTools.py index fca4c2d..86b1842 100644 --- a/AirlockTools.py +++ b/AirlockTools.py @@ -63,7 +63,7 @@ def main(): register_function("monitorOTP", utils.otpfunctions.monitorOTP) - if not os.path.exists("scheduling\\jobs.json"): recurring_job("monitor", "monitorOTP", interval=60, unit="seconds", args=[url]) + if not os.path.exists("scheduling\\jobs.json"): recurring_job("monitor", "monitorOTP", interval=60, unit="seconds", args=[url, pups]) else: reload_jobs() start_scheduler() diff --git a/utils/otpfunctions.py b/utils/otpfunctions.py index a14d3b4..bf3006f 100644 --- a/utils/otpfunctions.py +++ b/utils/otpfunctions.py @@ -65,11 +65,17 @@ def monitorOTP(url, pups): getActiveOTP(url) - if not os.path.exists("OTP\\old_active_OTP"): - shutil.copy2("OTP\\PARQ\\newest_active_OTP.parquet", "OTP\\PARQ\\old_active_OTP.parquet") - old_active_OTP = pd.read_parquet("OTP\\PARQ\\old_active_OTP.parquet") - current_active_OTP = pd.read_parquet("OTP\\PARQ\\newest_active_OTP.parquet") + old_otp_path = "OTP\\PARQ\\old_active_OTP.parquet" + new_otp_path = "OTP\\PARQ\\newest_active_OTP.parquet" + + if os.path.exists(old_otp_path): + old_active_OTP = pd.read_parquet(old_otp_path) + else: + old_active_OTP = pd.DataFrame(columns=['otpid']) # Ensure expected column exists + + current_active_OTP = pd.read_parquet(new_otp_path) + newly_added = current_active_OTP[~current_active_OTP['otpid'].isin(old_active_OTP['otpid'])] still_in_OTP = old_active_OTP[old_active_OTP['otpid'].isin(current_active_OTP['otpid'])] @@ -79,14 +85,14 @@ def monitorOTP(url, pups): for _, row in newly_added.iterrows(): clientid = row['clientid'] - duration = (row['duration'] * 60) + duration = (int(row['duration']) * 60) hostname = row['hostname'] purpose = row ['purpose'] pid = row['otpid'] early = math.floor(duration * .9) - run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}, early run", "addhash", time.time() + early, args=[url, clientid, pid, pups]) - run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}, duration complete run", "addhash", time.time() + duration, args=[url, clientid, pid, pups]) + run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}, early run", "addhash", time.time() + early, [url, clientid, pid, pups], None) + run_once_job(f"Add activity hashes for {pid}, for {hostname} for the purpose: {purpose}, duration complete run", "addhash", time.time() + duration, [url, clientid, pid, pups], None) print(f"Processing: {pid} with other data: {row}") diff --git a/utils/perstscheduler.py b/utils/perstscheduler.py index d8bedb0..aa92cb1 100644 --- a/utils/perstscheduler.py +++ b/utils/perstscheduler.py @@ -1,3 +1,4 @@ + import schedule import time import json @@ -7,6 +8,9 @@ from typing import Callable, Any, List, Dict # File where all jobs are persisted JOBS_FILE = "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] = {} @@ -17,7 +21,6 @@ FUNCTION_MAP: Dict[str, Callable] = {} 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) """ @@ -34,19 +37,68 @@ def load_jobs() -> List[Dict[str, Any]]: 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).""" - with open(JOBS_FILE, "w") as f: - json.dump(jobs, f, indent=4) + _atomic_save(JOBS_FILE, jobs) # ------------------------------- -# Run Once Jobs +# Uniqueness Helpers # ------------------------------- -def run_once_job(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None): +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: """ - Schedule a job to run once at a specific timestamp. + 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": + print(f"[INFO] 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 {} @@ -56,12 +108,10 @@ def run_once_job(job_id: str, func_name: str, run_at_timestamp: float, args=None print(f"[ERROR] 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) @@ -69,38 +119,10 @@ def run_once_job(job_id: str, func_name: str, run_at_timestamp: float, args=None if delay_seconds <= 0: print(f"[WARN] Job {job_id} scheduled in the past. Skipping.") return - # Schedule via schedule library schedule.every(delay_seconds).seconds.do(job_wrapper).tag(job_id) - # Save to JSON - jobs = load_jobs() - jobs.append({ - "id": job_id, - "type": "once", - "run_at": run_at_timestamp, - "function": func_name, - "args": args, - "kwargs": kwargs - }) - save_jobs(jobs) - -# ------------------------------- -# Recurring Jobs -# ------------------------------- - -def recurring_job(job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None): - """ - Schedule a recurring job. - - Args: - job_id: Unique job name - func_name: Function to call (must be registered) - interval: Number of units between runs - unit: "seconds", "minutes", "hours", "days" - args: Positional arguments for the function - kwargs: Keyword arguments for the function - """ +def _schedule_recurring(job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None): args = args or [] kwargs = kwargs or {} @@ -110,7 +132,6 @@ def recurring_job(job_id: str, func_name: str, interval: int, unit: str, args=No return FUNCTION_MAP[func_name](*args, **kwargs) - # Pick correct scheduling unit if unit == "seconds": schedule.every(interval).seconds.do(job_wrapper).tag(job_id) elif unit == "minutes": @@ -122,41 +143,102 @@ def recurring_job(job_id: str, func_name: str, interval: int, unit: str, args=No else: raise ValueError(f"Unsupported unit: {unit}") - # Save to JSON - jobs = load_jobs() - # Ensure no duplicate job id in file - jobs = [j for j in jobs if j["id"] != job_id] - jobs.append({ - "id": job_id, - "type": "recurring", - "interval": interval, - "unit": unit, - "function": func_name, - "args": args, - "kwargs": kwargs - }) - save_jobs(jobs) +# ------------------------------- +# 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 + else: + if 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 + else: + if 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) # ------------------------------- # Reload Saved Jobs # ------------------------------- def reload_jobs(): - """Reload jobs from JSON and reschedule them.""" + """Reload jobs from JSON and reschedule them (no re-persist).""" jobs = load_jobs() for job in jobs: if job["type"] == "once": - # Only reschedule if still in the future if job["run_at"] > time.time(): - run_once_job(job["id"], job["function"], job["run_at"], job.get("args"), job.get("kwargs")) + 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") + job["id"], job["function"], job["interval"], job["unit"], + job.get("args"), job.get("kwargs"), + persist=False ) # ------------------------------- @@ -174,36 +256,3 @@ def start_scheduler(): time.sleep(0.5) except KeyboardInterrupt: print("[INFO] Scheduler stopped.") - - -""" -import time -from persistent_schedule import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler - -# Example functions -def greet(name, loud=False): - if loud: - print(f"HELLO, {name}!") - else: - print(f"Hello, {name}.") - -def add(a, b): - print(f"{a} + {b} = {a + b}") - -# Register functions -register_function("greet", greet) -register_function("add", add) - -# Reload saved jobs -reload_jobs() - -# Schedule a run-once job in 5 seconds -run_once_job("job1", "greet", time.time() + 5, args=["Alice"], kwargs={"loud": True}) - -# Schedule a recurring job every 10 seconds -recurring_job("job2", "add", interval=10, unit="seconds", args=[2, 3]) - -# Start scheduler loop -start_scheduler() - -""" \ No newline at end of file