From 4ef5ceb63095f635d64f5f0a6a7372b439061db0 Mon Sep 17 00:00:00 2001
From: schmid-jn <jonathan.schmid@4am-robotics.com>
Date: Mon, 12 Aug 2024 09:07:01 +0200
Subject: [PATCH] add prio_stop feature

---
 README.md                    |  5 +++--
 robmuxinator/robmuxinator.py | 27 +++++++++++++++++----------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 3d42f8c..b5b6068 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,7 @@ It is also helpful for local development to start multiple commands easily with
 
 - [robmuxinator](#robmuxinator)
   - [Installation](#installation)
+    - [Nix](#nix)
   - [Getting Started](#getting-started)
   - [Usage](#usage)
     - [Command Choices](#command-choices)
@@ -136,7 +137,8 @@ sessions:
 - `host: string` (optional, default: hostname of localhost): Target host of the `tmux` session.
 - `user: string` (optional, default: robot): Target user of the `tmux` session.
 - `wait_for_core: bool` (optional, default: true): Starts the session only after `roscore` is available.
-- `prio: int` (optional, default: 10): Priority of the session. Sessions with the same priority start concurrently. Smaller numbers have higher priority.
+- `prio: int` (optional, default: 10): Priority of starting the session. Sessions with the same priority start concurrently. Smaller numbers have higher priority.
+- `prio_stop: int` (optional, default: `prio`): Priority of stopping the session. Sessions with the same priority stop concurrently. Higher numbers have higher priority.
 - `locked: bool` (optional, default: false): Locked sessions won't be closed on `stop` or `restart` (only if forced).
 - `pre_condition: string` (optional): Bash command used as a condition that must be fulfilled before the session can start.
 
@@ -175,4 +177,3 @@ This is s summary of useful commands for working with `tmux`.
 
 ## License
 Apache License Version 2.0, January 2004
-
diff --git a/robmuxinator/robmuxinator.py b/robmuxinator/robmuxinator.py
index d09690d..af0f14e 100755
--- a/robmuxinator/robmuxinator.py
+++ b/robmuxinator/robmuxinator.py
@@ -494,6 +494,11 @@ def __init__(self, ssh_client, session_name, yaml_session, envs=None) -> None:
         else:
             self.prio = 10
 
+        if "prio_stop" in yaml_session:
+            self.prio_stop = int(yaml_session["prio_stop"])
+        else:
+            self.prio_stop = self.prio
+
         if "locked" in yaml_session:
             self._locked = yaml_session["locked"]
         else:
@@ -568,6 +573,7 @@ def dump(self):
         print("\tcommand: {}".format(self._command))
         print("\tpre_condition: {}".format(self._pre_condition))
         print("\tprio: {}".format(self.prio))
+        print("\tprio_stop: {}".format(self.prio_stop))
         print("\tlocked: {}".format(self._locked))
 
 
@@ -645,9 +651,9 @@ def order_sessions_by_key(sessions, key):
     ordered_array = sorted(sessions, key=operator.attrgetter(key))
     ordered_sessions = dict()
     for s in ordered_array:
-        if not s.prio in ordered_sessions:
-            ordered_sessions[s.prio] = []
-        ordered_sessions[s.prio].append(s)
+        if not getattr(s, key) in ordered_sessions:
+            ordered_sessions[getattr(s, key)] = []
+        ordered_sessions[getattr(s, key)].append(s)
     # returns a dict of arrays with prio as key and array of session with the same key(prio) as value
     return ordered_sessions
 
@@ -919,33 +925,34 @@ def main():
         logger.error(e)
         sys.exit()
 
-    ordered_sessions = order_sessions_by_key(sessions, "prio")
+    ordered_sessions_start = order_sessions_by_key(sessions, "prio")
+    ordered_sessions_stop = order_sessions_by_key(sessions, "prio_stop")
 
     if command == "start":
         # wait for other hosts
         if len(hosts) > 1 and not wait_for_hosts(hosts, timeout):
             sys.exit()
-        start_sessions(ordered_sessions)
+        start_sessions(ordered_sessions_start)
         logger.info(
             "starting took {} secs".format((datetime.now() - start).total_seconds())
         )
     elif command == "stop":
-        stop_sessions(ordered_sessions, args.force)
+        stop_sessions(ordered_sessions_stop, args.force)
         logger.info(
             "stopping took {} secs".format((datetime.now() - start).total_seconds())
         )
     elif command == "restart":
-        stop_sessions(ordered_sessions, args.force)
-        start_sessions(ordered_sessions)
+        stop_sessions(ordered_sessions_stop, args.force)
+        start_sessions(ordered_sessions_start)
         logger.info(
             "restart took {} secs".format((datetime.now() - start).total_seconds())
         )
     elif command == "shutdown":
         pre_shutdown_commands(yaml_content)
-        stop_sessions(ordered_sessions, True)
+        stop_sessions(ordered_sessions_stop, True)
         shutdown_system(hosts, timeout)
     elif command == "reboot":
         pre_reboot_commands(yaml_content)
         pre_shutdown_commands(yaml_content)
-        stop_sessions(ordered_sessions, True)
+        stop_sessions(ordered_sessions_stop, True)
         shutdown_system(hosts, timeout)