commit 06ab3029794648152a9831aa1dc107405dce4aa8
Author: Daniel Tsvetkov <daniel@pi2.dev>
Date:   Sat Dec 21 15:03:39 2024 +0200

    Initial commit

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..ed8ebf5
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__
\ No newline at end of file
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..51c4feb
--- /dev/null
+++ b/README.md
@@ -0,0 +1,33 @@
+# Monitor Memory and Disk usage
+
+The scripts in this repo monitors for high memory usage and low disk space available.
+
+It sends a notification using `notify-send` to the user.
+
+For memory, when the condition is met, a signal is sent (e.g. `SIGSTOP`, `SIGKILL`).
+
+## Configure
+
+In `config.py` you could configure how often to monitor for memory usage and disk usage and what are the limits for each.
+
+You could also configure whether to send the `STOP` signal or `KILL` signal (or something else). If you send the `STOP` signal, the process can be continued using `SIGCONT` signal.
+
+You can also configure the icon used by `notify-send`.
+
+## Install
+
+Change the `WorkingDirectory` param in the two service files.
+
+Copy the `mem_monitor.service` and `disk_monitor.service` files into `/etc/systemd/system`.
+
+Then run:
+
+```bash
+systemctl daemon-reload
+systemctl enable mem_monitor.service
+systemctl enable disk_monitor.service
+systemctl start mem_monitor.service
+systemctl start disk_monitor.service
+```
+
+Check if there might be any errors using `journalctl -fu mem_monitor.service` or `journalctl -fu disk_monitor.service`.
diff --git a/config.py b/config.py
new file mode 100644
index 0000000..ddd3934
--- /dev/null
+++ b/config.py
@@ -0,0 +1,12 @@
+# Memory monitor configs
+MEM_MIN_AVAILABLE_G = 1.0
+TIME_SLEEP_MEM_S = 1
+SIGNAL = 'STOP' # 'STOP' or 'KILL'
+
+STR_CONT_OPT = "\n\n Run 'kill -SIGCONT <PID>' to resume" if SIGNAL == 'STOP' else ""
+
+# Disk space monitor
+MIN_FREE_DISK_SPACE_G = 2.0
+TIME_SLEEP_DISK_S = 60
+
+ICON_NOTIFY = '/home/pi2/.local/share/icons/Flat-Remix-Blue-Dark/apps/scalable/dialog-warning.svg'
diff --git a/disk_monitor.py b/disk_monitor.py
new file mode 100644
index 0000000..332f1ea
--- /dev/null
+++ b/disk_monitor.py
@@ -0,0 +1,61 @@
+"""
+Monitors disk usage and notifies on low
+
+Need to pass the path argument as first, e.g.:
+
+python disk_monitor.py /
+"""
+import argparse
+import logging
+import shutil
+import time
+
+
+import config
+import lib
+
+logging.basicConfig()
+logger = logging.getLogger()
+
+
+def setup_logging_level(debug=False):
+    log_level = logging.DEBUG if debug else logging.ERROR
+    logger.setLevel(log_level)
+    logger.debug("Debugging enabled")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('path', help="Path to monitor")
+    parser.add_argument('--debug', dest='debug', action='store_true')
+    return parser.parse_args()
+
+
+def get_free_disk_space_g(path):
+    free_disk_space = shutil.disk_usage(path).free
+    free_disk_space_g = round(free_disk_space / (1024 ** 3), 1)
+    return free_disk_space_g
+
+
+def loop(args):
+    is_low_disk = False
+    while True:
+        free_disk_space_g = get_free_disk_space_g(args.path)
+        logging.debug(f"Free disk space: {free_disk_space_g} G")
+        if free_disk_space_g < config.MIN_FREE_DISK_SPACE_G:
+            if not is_low_disk:
+                message = f"Less than {config.MIN_FREE_DISK_SPACE_G}G of free disk space available on {args.path}."
+                logging.warning(message)
+                lib.send_notification("Disk Space Alert", message)
+            is_low_disk = True
+        time.sleep(config.TIME_SLEEP_DISK_S)
+
+
+def main():
+    args = parse_args()
+    setup_logging_level(args.debug)
+    loop(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/disk_monitor.service b/disk_monitor.service
new file mode 100644
index 0000000..f501f9b
--- /dev/null
+++ b/disk_monitor.service
@@ -0,0 +1,13 @@
+[Unit]
+Description=Disk Monitor
+After=network.target
+
+[Service]
+User=root
+WorkingDirectory=/home/pi2/workspace/memmon
+ExecStart=/usr/bin/python3 disk_monitor.py /
+Restart=always
+Environment="PYTHONUNBUFFERED=1" "DISPLAY=:0" "XAUTHORITY=/home/pi2/.Xauthority"
+
+[Install]
+WantedBy=multi-user.target
\ No newline at end of file
diff --git a/lib.py b/lib.py
new file mode 100644
index 0000000..03c311f
--- /dev/null
+++ b/lib.py
@@ -0,0 +1,13 @@
+import logging
+import subprocess
+
+import config
+
+def run_os_command(cmd):
+    logging.debug(f'COMMAND: {cmd}')
+    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+    return result.stdout
+
+
+def send_notification(line1, line2):
+    run_os_command(f'notify-send -w -u critical -i "{config.ICON_NOTIFY}" "{line1}" "{line2}"')
diff --git a/mem_monitor.py b/mem_monitor.py
new file mode 100644
index 0000000..4634981
--- /dev/null
+++ b/mem_monitor.py
@@ -0,0 +1,72 @@
+"""
+Monitors memory usage, notifies and stops/kills the top memory usage process.
+
+Test with:
+
+stress --vm-bytes 7G --vm-hang 0 --vm 1
+"""
+import argparse
+import logging
+import time
+
+import config
+import lib
+
+logging.basicConfig()
+logger = logging.getLogger()
+
+
+def setup_logging_level(debug=False):
+    log_level = logging.DEBUG if debug else logging.ERROR
+    logger.setLevel(log_level)
+    logger.debug("Debugging enabled")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('path', help="Path to monitor")
+    parser.add_argument('--debug', dest='debug', action='store_true')
+    return parser.parse_args()
+
+
+def get_available_mem_g():
+    mem_available = lib.run_os_command(
+        "grep MemAvailable /proc/meminfo | awk '{print $2}'")
+    mem_available = int(mem_available)
+    mem_available_g = round(mem_available / (1024 ** 2), 1)
+    return mem_available_g
+
+
+def get_top_pid_comm():
+    top_pid_comm = lib.run_os_command(
+        "ps -eo pid,pmem,comm --sort=-pmem | head -n 2 | awk '{print $1 \" \" $3}' | tail -1")
+    return top_pid_comm.split()
+
+
+def stop_process(pid):
+    lib.run_os_command(f"kill -{config.SIGNAL} {pid}")
+
+
+def main():
+    is_low_mem = False
+    while True:
+        mem_available_g = get_available_mem_g()
+        logging.debug(f"Memory Available: {mem_available_g} G")
+        top_pid, top_comm = get_top_pid_comm()
+        if mem_available_g < config.MEM_MIN_AVAILABLE_G:
+            if not is_low_mem:
+                message = f"Less than {config.MEM_MIN_AVAILABLE_G}G of memory available. \n\n Sent signal {config.SIGNAL} {top_comm} ({top_pid}) {config.STR_CONT_OPT} \n\n Run 'kill -SIGKILL <PID>' to kill immediately."
+                logging.warning(message)
+                stop_process(int(top_pid))
+                lib.send_notification("Memory Alert", message)
+            is_low_mem = True
+        else:
+            is_low_mem = False
+
+        logging.debug(
+            f"Top process by memory usage: PID {top_pid}, Command: {top_comm}")
+        time.sleep(config.TIME_SLEEP_MEM_S)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/mem_monitor.service b/mem_monitor.service
new file mode 100644
index 0000000..86b1104
--- /dev/null
+++ b/mem_monitor.service
@@ -0,0 +1,13 @@
+[Unit]
+Description=Memory Monitor
+After=network.target
+
+[Service]
+User=root
+WorkingDirectory=/home/pi2/workspace/memmon
+ExecStart=/usr/bin/python3 mem_monitor.py
+Restart=always
+Environment="PYTHONUNBUFFERED=1" "DISPLAY=:0" "XAUTHORITY=/home/pi2/.Xauthority"
+
+[Install]
+WantedBy=multi-user.target
\ No newline at end of file