From e9ba9a45f4d114fdde7072a910893fca2ff14735 Mon Sep 17 00:00:00 2001
From: vrenaville <vincent.renaville@camptocamp.com>
Date: Fri, 2 Aug 2024 16:07:09 +0200
Subject: [PATCH] feat: add worker mem metrics

---
 .../controllers/prometheus_metrics.py         |  3 ++
 monitoring_prometheus/models/__init__.py      |  1 +
 .../models/psutils_helpers.py                 | 39 +++++++++++++++++++
 3 files changed, 43 insertions(+)
 create mode 100644 monitoring_prometheus/models/psutils_helpers.py

diff --git a/monitoring_prometheus/controllers/prometheus_metrics.py b/monitoring_prometheus/controllers/prometheus_metrics.py
index 4aa49960..e3e4cdb4 100644
--- a/monitoring_prometheus/controllers/prometheus_metrics.py
+++ b/monitoring_prometheus/controllers/prometheus_metrics.py
@@ -5,8 +5,11 @@
 
 from odoo.http import Controller, route
 
+from ..models.psutils_helpers import get_process_info
+
 
 class PrometheusController(Controller):
     @route("/metrics", auth="public")
     def metrics(self):
+        get_process_info()
         return generate_latest()
diff --git a/monitoring_prometheus/models/__init__.py b/monitoring_prometheus/models/__init__.py
index 9a5eb718..15f6add4 100644
--- a/monitoring_prometheus/models/__init__.py
+++ b/monitoring_prometheus/models/__init__.py
@@ -1 +1,2 @@
 from . import ir_http
+from . import psutils_helpers
diff --git a/monitoring_prometheus/models/psutils_helpers.py b/monitoring_prometheus/models/psutils_helpers.py
new file mode 100644
index 00000000..5b14e337
--- /dev/null
+++ b/monitoring_prometheus/models/psutils_helpers.py
@@ -0,0 +1,39 @@
+import psutil
+from prometheus_client import Gauge
+
+MEMORY_USAGE_VMS = Gauge(
+    "odoo_worker_memory_user_vms_mb", "Memory usage in MB", ["process", "pid"]
+)
+
+MEMORY_USAGE_RSS = Gauge(
+    "odoo_worker_memory_user_rss_mb", "Memory usage in MB", ["process", "pid"]
+)
+
+
+def get_process_info():
+    for process in psutil.process_iter(
+        ["pid", "name", "memory_full_info", "cmdline", "nice"]
+    ):
+        try:
+            if process.info["memory_full_info"]:
+                if process.info["nice"] == 10:
+                    ProcessLabel = "workercron"
+                elif process.info["pid"] == 1:
+                    ProcessLabel = "dispatcher"
+                elif any("gevent" in x for x in process.info["cmdline"]):
+                    ProcessLabel = "gevent"
+                elif any("odoo" in x for x in process.info["cmdline"]):
+                    ProcessLabel = "workerhttp"
+                elif any("shell" in x for x in process.cmdline()):
+                    ProcessLabel = "OdooShell"
+                else:
+                    ProcessLabel = "other"
+                MEMORY_USAGE_VMS.labels(ProcessLabel, process.info["pid"]).set(
+                    process.info["memory_full_info"].rss // 1000000
+                )
+                MEMORY_USAGE_RSS.labels(ProcessLabel, process.info["pid"]).set(
+                    process.info["memory_full_info"].vms // 1000000
+                )
+
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+            continue