mirror-checker/main.py

#!/usr/bin/env python3

"""
This mirror status checker determines whether CSC mirror is up-to-date with upstream
"""

import time
import sys
import requests
from multiprocessing import Pool, Manager
from typing import Optional
from time import sleep, localtime, strftime

from projects import *
import json

NUM_THREAD = 16
MAX_RETRY = 3
RETRY_TIMEOUT = 30 # In seconds

def safe_print(*args, **kwargs):
    # When run with 'chronic' and 'timeout', stdout gets suppressed
    # due to buffering. Make sure to always flush the output.
    print(*args, **kwargs, flush=True)

# Return None if no error occurs and a string for error message otherwise
def check_project(args) -> Optional[str]:
    current_time = int(time.time())
    project, data = args
    try:
        project_class = getattr(sys.modules[__name__], project)

        # Skip projects we no longer mirror
        if data[project].get('exclude', False):
            return None

        checker_result = project_class.check(data, project, current_time)

        if checker_result:
            data[project]["out_of_sync_since"] = None
            return None

        elif (data[project]["out_of_sync_since"] is not None
                and current_time - data[project]["out_of_sync_since"] > data[project]["out_of_sync_interval"]):
            now_str = strftime("%d %b %Y %H:%M:%S (local time)", localtime())
            duration = current_time - data[project]["out_of_sync_since"]
            return f"{project} out-of-sync at {now_str} for {duration}s"

        else:
            data[project]["out_of_sync_since"] = current_time
            return None

    except requests.exceptions.RequestException as err:
        return f"{project}\n{err}"

def check_project_with_retry(args) -> bool:
    project, _ = args
    errs = []
    for _ in range(MAX_RETRY):
        res = check_project(args)
        if res == None:
            safe_print(f"Success: {project} up-to-date")
            return True
        else:
            errs.append(res)
            # Do nothing, try again later
            sleep(RETRY_TIMEOUT)

    # Max try reached, print errors
    safe_print(f"Error: {project}")
    for reason in errs:
        safe_print(f"       {reason}")

    return False

def main():
    data_file = 'data.json'
    if len(sys.argv) > 1:
        data_file = sys.argv[1]

    manager = Manager()
    data = json.load(open(data_file))
    sync_data = manager.dict({k: manager.dict(v) for k, v in data.items()})

    with Pool(NUM_THREAD) as pool:
        all_pass = all(pool.imap(check_project_with_retry, ((k, sync_data) for k in data.keys())))

    with open(data_file, "w", encoding="utf-8") as file:
        json.dump({k: dict(v) for k, v in sync_data.items()}, file, indent='  ')

    sys.exit(0 if all_pass else 1)

if __name__ == "__main__":
    main()