Merge pull request #342 from DSD-DBS/process-metrics

feat: Add process metrics using psutil

remote/Dockerfile

@@ -84,7 +84,7 @@ COPY startup.sh .startup.sh
 RUN chmod 755 .startup.sh /home/techuser/.config/openbox/autostart
 
 # Prepare idletime metric endpoint
-RUN pip install --no-cache-dir prometheus-client==0.17.1
+RUN pip install --no-cache-dir prometheus-client==0.21.0 psutil==6.1.0
 
 COPY metrics.py .metrics.py
 RUN chown techuser /home/techuser/.metrics.py

remote/metrics.py

@@ -7,9 +7,11 @@
 import logging
 import os
 import subprocess
+import typing as t
 from wsgiref import simple_server
 
 import prometheus_client
+import psutil
 
 METRICS_PORT = int(os.getenv("METRICS_PORT", "9118"))
 LOGGER = logging.getLogger(__file__)
@@ -56,8 +58,92 @@ def get_idletime(self) -> float:
         return round(current_idle_time, 2)
 
 
+class ProcessCollector(prometheus_client.registry.Collector):
+    """Track metrics of all system processes."""
+
+    def collect(self) -> t.Iterable[prometheus_client.Metric]:
+        """Collect metrics from all system processes."""
+        processes = psutil.process_iter()
+
+        process_cpu_percent_metric = (
+            prometheus_client.metrics_core.GaugeMetricFamily(
+                "process_cpu_percent",
+                "CPU percent of the process",
+                labels=["process_name"],
+            )
+        )
+
+        process_memory_usage_metric = (
+            prometheus_client.metrics_core.GaugeMetricFamily(
+                "process_memory_usage_bytes",
+                "Memory usage of the process in bytes",
+                labels=["process_name"],
+            )
+        )
+
+        process_io_counters_metric = (
+            prometheus_client.metrics_core.GaugeMetricFamily(
+                "process_io_counters_bytes",
+                "Read and write bytes by the process",
+                labels=["process_name", "io_type"],
+            )
+        )
+
+        process_num_threads_metric = (
+            prometheus_client.metrics_core.GaugeMetricFamily(
+                "process_num_threads",
+                "Number of threads of the process",
+                labels=["process_name"],
+            )
+        )
+
+        process_open_fds_metric = (
+            prometheus_client.metrics_core.GaugeMetricFamily(
+                "process_open_file_descriptors",
+                "Number of open file descriptors by the process",
+                labels=["process_name"],
+            )
+        )
+
+        for process in processes:
+            with process.oneshot():
+                process_cpu_percent_metric.add_metric(
+                    [process.name()], process.cpu_percent()
+                )
+
+                process_memory_usage_metric.add_metric(
+                    [process.name()], process.memory_info().rss
+                )
+
+                io_counters = process.io_counters()
+                if io_counters:
+                    process_io_counters_metric.add_metric(
+                        [process.name(), "read"], io_counters.read_bytes
+                    )
+                    process_io_counters_metric.add_metric(
+                        [process.name(), "write"], io_counters.write_bytes
+                    )
+
+                process_num_threads_metric.add_metric(
+                    [process.name()], process.num_threads()
+                )
+
+                if hasattr(process, "num_fds"):
+                    process_open_fds_metric.add_metric(
+                        [process.name()], process.num_fds()
+                    )
+
+        yield process_cpu_percent_metric
+        yield process_memory_usage_metric
+        yield process_io_counters_metric
+        yield process_num_threads_metric
+        yield process_open_fds_metric
+
+
 IDLETIME.set_function(IdleTimer().get_idletime)
 
+prometheus_client.REGISTRY.register(ProcessCollector())
+
 
 def start_server(
     addr: str,