77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
import logging
|
|
from multiprocessing.pool import ThreadPool
|
|
from threading import Lock
|
|
from typing import Union
|
|
|
|
from backend import app, LrcException
|
|
from backend.models import Recorder
|
|
from backend.tools.simple_state_checker import check_capture_agent_state, ping_capture_agent
|
|
|
|
cron_log_handler = logging.FileHandler(app.config.get('CRON_LOG_FILE'))
|
|
cron_logger = logging.getLogger("mal.cron")
|
|
cron_logger.addHandler(cron_log_handler)
|
|
logging.getLogger("apscheduler.scheduler").addHandler(cron_log_handler)
|
|
logging.getLogger("apscheduler.executors.default").addHandler(cron_log_handler)
|
|
|
|
recorder_jobs_lock = Lock()
|
|
recorder_jobs = set()
|
|
|
|
NUM_THREADS = 8
|
|
|
|
|
|
def add_recorder_to_state_check(recorder: Union[int, Recorder]):
|
|
if isinstance(recorder, int):
|
|
recorder = Recorder.get_by_identifier(recorder)
|
|
if recorder is None:
|
|
cron_logger.warning(
|
|
"Could not add recorder to state check, as specified id could not be found / recorder is None")
|
|
raise LrcException("Recorder is None / could not be found!")
|
|
recorder_jobs_lock.acquire()
|
|
recorder_jobs.add(recorder)
|
|
recorder_jobs_lock.release()
|
|
|
|
|
|
def remove_recorder_from_state_check(recorder: Union[int, Recorder]):
|
|
if isinstance(recorder, int):
|
|
recorder = Recorder.get_by_identifier(recorder)
|
|
if recorder is None:
|
|
cron_logger.warning(
|
|
"Could not remove recorder from state check, as specified id could not be found / recorder is None")
|
|
raise LrcException("Recorder is None / could not be found (and therefor not removed)!")
|
|
recorder_jobs_lock.acquire()
|
|
recorder_jobs.remove(recorder)
|
|
recorder_jobs_lock.release()
|
|
|
|
|
|
def check_recorder_state():
|
|
recorder_jobs_lock.acquire()
|
|
recorders = list(recorder_jobs)
|
|
recorder_jobs_lock.release()
|
|
|
|
recorder_states = {r['name']: {'state_ok': False, 'msg': 'unknown state!'} for r in recorders}
|
|
|
|
with ThreadPool(NUM_THREADS) as pool:
|
|
results = [pool.apply_async(check_capture_agent_state, (recorder,)) for recorder in recorders]
|
|
try:
|
|
state_results = [res.get(timeout=12) for res in results]
|
|
except TimeoutError as e:
|
|
cron_logger.error("Timeout while getting capture agent state! {}".format(e))
|
|
|
|
for r in state_results:
|
|
if r[0]: # ok :)
|
|
recorder_states[r[2]] = {'state_ok': True}
|
|
else:
|
|
recorder_states[r[2]]['msg'] = r[1]
|
|
|
|
with ThreadPool(NUM_THREADS) as pool:
|
|
results = [pool.apply_async(ping_capture_agent, (recorder,)) for recorder in recorders]
|
|
try:
|
|
ping_results = [res.get(timeout=12) for res in results]
|
|
except TimeoutError as e:
|
|
cron_logger.error("Timeout while pinging capture agent! {}".format(e))
|
|
|
|
for r in ping_results:
|
|
if not r[0]: # ok :)
|
|
recorder_states[r[2]]['msg'] = r[1]
|
|
|