mklauncher: implement importances

src/machinetalk/mklauncher/mklauncher.py

@@ -37,10 +37,55 @@ def printError(msg):
     sys.stderr.write('ERROR: ' + msg + '\n')
 
 
+class LauncherImportance(object):
+    DEFAULT_IMPORTANCE = 0
+
+    def __init__(self, config_file):
+        self._config_file = config_file
+        self._importances = {}
+
+    def __setitem__(self, launcher_id, importance):
+        self._importances[launcher_id] = importance
+
+    def __getitem__(self, launcher_id):
+        if launcher_id in self._importances:
+            return self._importances[launcher_id]
+        else:
+            return LauncherImportance.DEFAULT_IMPORTANCE
+
+    def save(self):
+        config_dir = os.path.dirname(self._config_file)
+        if not os.path.exists(config_dir):
+            os.makedirs(config_dir)
+
+        cfg = configparser.ConfigParser()
+        for key in self._importances:
+            section, name = key.split(':')
+            if not cfg.has_section(section):
+                cfg.add_section(section)
+            cfg.set(section, name, self._importances[key])
+
+        with open(self._config_file, 'w') as config_file:
+            cfg.write(config_file)
+
+    def load(self):
+        self._importances = {}
+        cfg = configparser.ConfigParser()
+        cfg.read(self._config_file)
+
+        for section in cfg.sections():
+            for (key, value) in cfg.items(section):
+                self._importances['%s:%s' % (section, key)] = int(value)
+
+    def __str__(self):
+        return str(self._importances)
+
+
 class Mklauncher(object):
     def __init__(self, context, launcherDirs=None, host='',
                  svcUuid='', debug=False, name=None, hostInName=True,
-                 pollInterval=0.5, pingInterval=2.0, loopback=False):
+                 pollInterval=0.5, pingInterval=2.0, loopback=False,
+                 config_dir='~/.config/machinekit/mklauncher'):
         if launcherDirs is None:
             launcherDirs = []
 
@@ -66,14 +111,19 @@ def __init__(self, context, launcherDirs=None, host='',
         self.processes = {}  # for processes mapped to launcher
         self.terminating = set()  # set of terminating processes
 
-        launchers = self._search_launchers(self.launcherDirs)
+        launchers, ids = self._search_launchers(self.launcherDirs)
+        self._launcher_ids = {}
         for index, launcher in enumerate(launchers):
+            self._launcher_ids[index] = ids[launcher.index]
             launcher.index = index
-            self.container.launcher.add().CopyFrom(launcher)
-            self.txContainer.launcher.add().MergeFrom(launcher)
-
+            self.container.launcher.extend([launcher])
+            self.txContainer.launcher.add().CopyFrom(launcher)
         logger.debug('parsed launchers:\n%s' % str(self.container))
 
+        config_file = os.path.expanduser(os.path.join(config_dir, 'importances.ini'))
+        self._importances = LauncherImportance(config_file)
+        self._importances.load()
+
         # prepare pings
         if self.pingInterval > 0:
             self.pingRatio = math.floor(self.pingInterval / self.pollInterval)
@@ -109,6 +159,8 @@ def _search_launchers(self, directories):
         }
 
         launchers = []
+        ids = {}
+        index = 0
         for rootDir in directories:
             for root, _, files in os.walk(rootDir):
                 if INI_NAME not in files:
@@ -128,6 +180,7 @@ def _search_launchers(self, directories):
                     info.model = cfg.get(section, 'model')
                     info.variant = cfg.get(section, 'variant')
                     launcher.priority = cfg.getint(section, 'priority')
+                    launcher.importance = 0
                     launcher.info.MergeFrom(info)
                     # command data
                     launcher.command = cfg.get(section, 'command')
@@ -150,10 +203,14 @@ def _search_launchers(self, directories):
                         image.encoding = CLEARTEXT
                         image.blob = fileBuffer
                         launcher.image.MergeFrom(image)
+
+                    launcher.index = index
+                    index += 1
                     launchers.append(launcher)
+                    ids[launcher.index] = '%s:%s' % (root, section)
 
         # sort using the priority attribute before distribution
-        return sorted(launchers, key=attrgetter('priority'), reverse=True)
+        return sorted(launchers, key=attrgetter('priority'), reverse=True), ids
 
     def _start_threads(self):
         threading.Thread(target=self._process_sockets).start()
@@ -234,18 +291,22 @@ def _add_pparams_to_message(self):
         self.txContainer.pparams.MergeFrom(parameters)
 
     def _update_launcher_status(self):
-        modified = False
+        txLauncher = Launcher()  # new pb message for tx
         for launcher in self.container.launcher:
+            modified = False
             index = launcher.index
 
+            importance = self._importances[self._launcher_ids[launcher.index]]
+            if importance is not launcher.importance:
+                txLauncher.importance = importance
+                modified = True
+
             terminating = False
             if index in self.terminating:
                 terminating = True
                 self.terminating.remove(index)
 
             if index in self.processes:
-                txLauncher = Launcher()  # new pb message for tx
-                txLauncher.index = index
                 process = self.processes[index]
                 process.poll()
                 returncode = process.returncode
@@ -279,9 +340,11 @@ def _update_launcher_status(self):
                     txLauncher.terminating = False
                     modified = True
                     self.processes.pop(index, None)  # remove from watchlist
-                if modified:
-                    launcher.MergeFrom(txLauncher)
-                    self.txContainer.launcher.add().MergeFrom(txLauncher)
+            if modified:
+                launcher.MergeFrom(txLauncher)
+                txLauncher.index = index
+                self.txContainer.launcher.add().MergeFrom(txLauncher)
+                txLauncher.Clear()
 
         if self.launcherFullUpdate:
             self._add_pparams_to_message()
@@ -388,6 +451,11 @@ def _shutdown_system(self):
         except:
             return False
 
+    def _update_importance(self, launcher):
+        launcher_id = self._launcher_ids[launcher.index]
+        self._importances[launcher_id] = launcher.importance
+        self._importances.save()
+
     def _send_command_wrong_params(self, identity, note='wrong parameters'):
         self.tx.note.append(note)
         self._send_command_message(identity, pb.MT_ERROR)
@@ -464,6 +532,19 @@ def _process_command_socket(self, s):
                 self.tx.note.append("cannot shutdown system: DBus error")
                 self._send_command_message(identity, pb.MT_ERROR)
 
+        elif self.rx.type == pb.MT_LAUNCHER_SET:
+            for launcher in self.rx.launcher:
+                if not launcher.HasField('index') \
+                   or not launcher.HasField('importance'):
+                    self._send_command_wrong_params()
+                    continue
+
+                index = launcher.index
+                if index >= len(self.container.launcher):
+                    self._send_command_wrong_index(identity)
+                else:
+                    self._update_importance(launcher)
+
         else:
             self.tx.note.append("unknown command")
             self._send_command_message(identity, pb.MT_ERROR)

src/machinetalk/mklauncher/test_mklauncher.py

@@ -5,7 +5,7 @@
 import_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../../lib/python/')
 sys.path.insert(0, os.path.abspath(import_path))
 
-from mklauncher import Mklauncher
+from mklauncher import Mklauncher, LauncherImportance
 
 
 @pytest.fixture
@@ -30,6 +30,11 @@ def single_launcher_file(tmpdir):
     return [str(tmpdir)]
 
 
+@pytest.fixture
+def config_dir(tmpdir):
+    return str(tmpdir.join('config'))
+
+
 def test_reading_single_launcher_file_works(context, single_launcher_file):
     launcher = Mklauncher(context, launcherDirs=single_launcher_file)
 
@@ -39,3 +44,66 @@ def test_reading_single_launcher_file_works(context, single_launcher_file):
     assert launchers[0].description == 'My super demo'
     assert launchers[0].command == 'python run.py'
     assert launchers[0].info.variant == 'default'
+
+
+@pytest.fixture
+def valid_importance_file(tmpdir):
+    data = '''
+[/foo/bar/baz]
+myconfig = 10
+anotherconfig = 2
+'''
+    ini = tmpdir.join('importances.ini')
+    ini.write(data)
+    return str(ini)
+
+
+def test_reading_launcher_importances_works(valid_importance_file):
+    importances = LauncherImportance(valid_importance_file)
+
+    importances.load()
+
+    print(importances._importances)
+    assert importances['/foo/bar/baz:myconfig'] == 10
+    assert importances['/foo/bar/baz:anotherconfig'] == 2
+
+
+def test_writing_launcher_importances_works(tmpdir):
+    save_file = tmpdir.join('test/output.ini')
+    importances = LauncherImportance(str(save_file))
+
+    importances['/my/config/path/:config1'] = 10
+    importances['/my/config/path/:config2'] = 2
+    importances['/home/alexander/:another_config'] = 0
+    importances.save()
+
+    assert os.path.exists(str(save_file))
+    data = save_file.read()
+    assert '[/my/config/path/]' in data
+    assert '[/home/alexander/]' in data
+    assert 'config1 = 10' in data
+    assert 'config2 = 2' in data
+    assert 'another_config = 0' in data
+
+
+def test_reading_launcher_importances_with_non_existing_file_does_not_throw_error(tmpdir):
+    save_file = tmpdir.join('save.ini')
+    importances = LauncherImportance(str(save_file))
+
+    importances.load()
+
+
+def test_updating_launcher_importance_works(context, single_launcher_file, config_dir):
+    launcher = Mklauncher(context, launcherDirs=single_launcher_file, config_dir=config_dir)
+
+    from machinetalk.protobuf.config_pb2 import Launcher
+    msg = Launcher()
+    msg.index = 0
+    msg.importance = 5
+
+    launcher._update_importance(msg)
+    launcher._update_launcher_status()
+
+    launchers = launcher.container.launcher
+    assert len(launchers) == 1
+    assert launchers[0].importance == 5