From da6771a0fb224aed5226b2ba45fa4db162482065 Mon Sep 17 00:00:00 2001
From: marcsello <punkosdmarcell@rocketmail.com>
Date: Sun, 12 Dec 2021 18:48:06 +0100
Subject: [PATCH] done woek

---
 main.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 68 insertions(+), 7 deletions(-)

diff --git a/main.py b/main.py
index 1d057e1..19fce07 100644
--- a/main.py
+++ b/main.py
@@ -8,6 +8,8 @@ import sys
 import os
 import io
 import time
+import signal
+import random
 
 
 class DeviceState(enum.Enum):
@@ -15,6 +17,7 @@ class DeviceState(enum.Enum):
     BOOTING = 1
     RUNNING = 2
     CRASHED = 3
+    STOPPED = 4
 
 
 @dataclass
@@ -25,20 +28,23 @@ class DeviceInfo:
     log_fd: Optional[io.TextIOWrapper] = None
     proc: Optional[subprocess.Popen] = None
 
+
 class RealSlimShadyThread(threading.Thread):
 
     def __init__(self, device_list: list):
         super().__init__()
         self._device_list = device_list
         self._device_list_lock = threading.Lock()
+        self._active = True
 
     def run(self):
-        while True:
+        while self._active:
             time.sleep(0.5)
             with self._device_list_lock:
 
                 # Check for crashed devices
-                for device in filter(lambda dev: dev.state in [DeviceState.BOOTING, DeviceState.RUNNING], self._device_list):
+                for device in filter(lambda dev: dev.state in [DeviceState.BOOTING, DeviceState.RUNNING],
+                                     self._device_list):
                     if device.proc.pid:
                         ret = device.proc.poll()
 
@@ -48,7 +54,15 @@ class RealSlimShadyThread(threading.Thread):
                             device.state = DeviceState.CRASHED
 
                 # Start created, crashed processes
-                for device in filter(lambda dev: dev.state == [DeviceState.CREATED, DeviceState.CRASHED], self._device_list)[:2]:
+                created_or_crashed = list(
+                    filter(
+                        lambda dev: dev.state in [DeviceState.CREATED, DeviceState.CRASHED],
+                        self._device_list
+                    )
+                )
+                random.shuffle(created_or_crashed)
+
+                for device in created_or_crashed[:2]:  # start only two at a time to save on resources
 
                     if device.log_fd:
                         try:
@@ -58,33 +72,80 @@ class RealSlimShadyThread(threading.Thread):
 
                     device.log_fd = open(f"ned_logs/{device.id}.log", "at")
                     device.proc = subprocess.Popen(
-                        [device.executable, str(device.id)],
+                        device.executable.split(' ') + [str(device.id)],
                         stdout=device.log_fd,
                         preexec_fn=lambda: os.setpgrp()
                     )
 
                     device.state = DeviceState.BOOTING
 
+        # Stopping
+        with self._device_list_lock:
+            stop_start_time = time.time()
+            for dev in filter(lambda dev: dev.state in [DeviceState.BOOTING, DeviceState.RUNNING], self._device_list):
+                dev.proc.terminate()
+
+        still_running = True
+        while still_running:
+            time.sleep(0.2)
+            with self._device_list_lock:
+
+                for dev in self._device_list:
+                    if dev.state in [DeviceState.RUNNING, DeviceState.BOOTING]:
+                        ret = dev.proc.poll()
+                        if ret is not None:
+                            if ret == 0:
+                                dev.state = DeviceState.STOPPED
+                            else:
+                                dev.state = DeviceState.CRASHED
+
+                still_running = any(
+                    dev.state in [DeviceState.RUNNING, DeviceState.BOOTING] for dev in self._device_list
+                )
+
+                if still_running and ((time.time() - stop_start_time) > 15):
+                    for dev in filter(lambda dev: dev.state in [DeviceState.BOOTING, DeviceState.RUNNING],
+                                      self._device_list):
+                        dev.proc.kill()
+                        dev.state = DeviceState.CRASHED
+
     def get_device_list(self) -> list:
         with self._device_list_lock:
             return self._device_list.copy()  # swallow copy... eh
 
+    def stop(self):
+        self._active = False
+
+    @property
+    def stopping(self) -> bool:
+        return not self._active
+
 
 def main():
     executable = sys.argv[1]
     n_proc = int(sys.argv[2])
     os.makedirs("ned_logs", exist_ok=True)
 
-    device_list = [DeviceInfo(id=i, executable=executable) for i in range(n_proc)]
+    slim_shady = RealSlimShadyThread([DeviceInfo(id=i, executable=executable) for i in range(n_proc)])
 
-    while True:
+    def shudown_signal(signal, frame):
+        slim_shady.stop()
+
+    signal.signal(signal.SIGINT, shudown_signal)
+    signal.signal(signal.SIGTERM, shudown_signal)
+
+    slim_shady.start()
+    while slim_shady.is_alive():
         print("\033[2J\033[0;0f", end="", flush=False)
         for state in DeviceState:
             print(state.name, flush=False)
-            print(", ".join(str(dev.id) for dev in device_list if dev.state == state), flush=False)
+            print(", ".join(str(dev.id) for dev in slim_shady.get_device_list() if dev.state == state), flush=False)
 
             print(flush=False)
 
+        if slim_shady.stopping:
+            print("STOPPING...", flush=False)
+
         sys.stdout.flush()
         time.sleep(1)