forked from IEW-CCS/IOT_Gateway_Python
-
Notifications
You must be signed in to change notification settings - Fork 0
/
workers_manager.py
145 lines (115 loc) · 5.52 KB
/
workers_manager.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import importlib
import threading
from functools import partial
import logging
from apscheduler.schedulers.background import BackgroundScheduler
from interruptingcow import timeout
from pytz import utc
from workers_queue import _WORKERS_QUEUE
import logger
from pip import __version__ as pip_version
if int(pip_version.split('.')[0]) >= 10:
from pip._internal import main as pip_main
else:
from pip import main as pip_main
_LOGGER = logger.get(__name__)
class WorkersManager:
class Command:
def __init__(self, callback, timeout, args=(), options=dict()):
self._callback = callback
self._timeout = timeout
self._args = args
self._options = options
self._source = '{}.{}'.format(callback.__self__.__class__.__name__ if hasattr(callback, '__self__') else callback.__module__, callback.__name__)
def execute(self):
messages = []
with timeout(self._timeout, exception=TimeoutError('Execution of command {} timed out after {} seconds'.format(self._source, self._timeout))):
messages = self._callback(*self._args)
_LOGGER.debug('Execution result of command %s: %s', self._source, messages)
return messages
def __init__(self, config):
self._mqtt_callbacks = []
self._update_commands = []
self._scheduler = BackgroundScheduler(timezone=utc)
self._daemons = []
self._config = config
self._command_timeout = config.get('command_timeout', 35)
def register_workers(self):
for (worker_name, worker_config) in self._config['workers'].items():
#print("WorkerManager --> register_workers : start to import workers")
module_obj = importlib.import_module("workers.%s" % worker_name)
klass = getattr(module_obj, "%sWorker" % worker_name.title())
# if module_obj.REQUIREMENTS is not None:
# self._pip_install_helper(module_obj.REQUIREMENTS)
command_timeout = worker_config.get('command_timeout', self._command_timeout)
# constructor
worker_obj = klass(command_timeout, **worker_config['args'])
if not hasattr(worker_obj, 'status_update') and not hasattr(worker_obj, 'run'):
raise "%s cannot be initialized, it has to define run or status_update method" % worker_name
if hasattr(worker_obj, 'status_update'):
_LOGGER.debug("Added %s worker with %d seconds interval and a %d seconds timeout", repr(worker_obj), worker_config['update_interval'], worker_obj.command_timeout)
command = self.Command(worker_obj.status_update, worker_obj.command_timeout, [])
self._update_commands.append(command)
if 'update_interval' in worker_config:
job_id = '{}_interval_job'.format(worker_name)
interval_job = self._scheduler.add_job(
partial(self._queue_command, command), 'interval',
seconds=worker_config['update_interval'], id=job_id
)
self._mqtt_callbacks.append((
worker_obj.format_topic('update_interval'),
partial(self._update_interval_wrapper, command, job_id)
))
if hasattr(worker_obj, 'run'):
print("has run attr")
_LOGGER.debug("Registered %s as daemon", repr(worker_obj))
self._daemons.append(worker_obj)
if 'topic_subscription' in worker_config:
self._mqtt_callbacks.append((
worker_config['topic_subscription'],
partial(self._on_command_wrapper, worker_obj)
))
print("WorkerManager --> register_workers: _mqtt_callbacks = ", self._mqtt_callbacks)
if 'topic_subscription' in self._config:
for (callback_name, options) in self._config['topic_subscription'].items():
self._mqtt_callbacks.append((
options['topic'],
lambda client, _ , c: self._queue_if_matching_payload(self.Command(getattr(self, callback_name), self._command_timeout), c.payload, options['payload']))
)
return self
def start(self, mqtt):
mqtt.callbacks_subscription(self._mqtt_callbacks)
self._scheduler.start()
self.update_all()
for daemon in self._daemons:
threading.Thread(target=daemon.run, args=[mqtt], daemon=True).start()
def _queue_if_matching_payload(self, command, payload, expected_payload):
if payload.decode('utf-8') == expected_payload:
self._queue_command(command)
def update_all(self):
_LOGGER.debug("Updating all workers")
for command in self._update_commands:
self._queue_command(command)
@staticmethod
def _queue_command(command):
_WORKERS_QUEUE.put(command)
@staticmethod
def _pip_install_helper(package_names):
for package in package_names:
pip_main(['install', '-q', package])
logger.reset()
def _update_interval_wrapper(self, command, job_id, client, userdata, c):
_LOGGER.info("Recieved updated interval for %s with: %s", c.topic, c.payload)
try:
new_interval = int(c.payload)
self._scheduler.remove_job(job_id)
self._scheduler.add_job(
partial(self._queue_command, command), 'interval',
seconds=new_interval, id=job_id)
except ValueError:
logger.log_exception(_LOGGER, 'Ignoring invalid new interval: %s', c.payload)
def _on_command_wrapper(self, worker_obj, client, userdata, c):
_LOGGER.debug("Received command for %s on %s: %s", repr(worker_obj), c.topic, c.payload)
global_topic_prefix = userdata['global_topic_prefix']
topic = c.topic[len(global_topic_prefix+'/'):] if global_topic_prefix is not None else c.topic
self._queue_command(self.Command(worker_obj.on_command, worker_obj.command_timeout, [topic, c.payload]))