Update ProcessStats query by using API instead of parsing ps command. (…

…#103) (#154)

Now ProcessStats query parses ps command and populate STATE_DB, "ps" command is somewhat fragile, its format might be impacted by relevant environment change. This PR uses /proc instead which is stable and verified via systemctl procdockerstatsd.service.

Co-authored-by: Feng-msft <[email protected]>

azure-pipelines.yml

@@ -64,6 +64,7 @@ stages:
     - script: |
         set -xe
         sudo pip3 install enum34
+        sudo pip install psutil
         sudo pip3 install swsssdk-2.0.1-py3-none-any.whl
         sudo pip3 install sonic_py_common-1.0-py3-none-any.whl
         sudo pip3 install sonic_yang_mgmt-1.0-py3-none-any.whl

scripts/procdockerstatsd

@@ -5,11 +5,12 @@ Daemon which periodically gathers process and docker statistics and pushes the d
 '''
 
 import os
+import psutil
 import re
 import subprocess
 import sys
 import time
-from datetime import datetime
+from datetime import datetime, timedelta
 
 from sonic_py_common import daemon_base
 from swsscommon import swsscommon
@@ -136,18 +137,30 @@ class ProcDockerStats(daemon_base.DaemonBase):
         return True
 
     def update_processstats_command(self):
-        cmd0 = ["ps", "-eo", "uid,pid,ppid,%mem,%cpu,stime,tty,time,cmd", "--sort", "-%cpu"]
-        cmd1 = ["head", "-1024"]
-        exitcode, data = getstatusoutput_noshell_pipe(cmd0, cmd1)
-        if any(exitcode):
-            cmd = ' | '.join([' '.join(cmd0), ' '.join(cmd1)])
-            self.log_error("Error running command '{}'".format(cmd))
-            data = None
-        processdata = self.format_process_cmd_output(data)
-        value = ""
+        processdata = []
+        for process in psutil.process_iter(['pid', 'ppid', 'memory_percent', 'cpu_percent', 'create_time', 'cmdline']):
+            try:
+                uid = process.uids()[0]
+                pid = process.pid
+                ppid = process.ppid()
+                mem = process.memory_percent()
+                cpu = process.cpu_percent()
+                stime = process.create_time()
+                stime_formatted = datetime.fromtimestamp(stime).strftime("%b%d")
+                tty = process.terminal()
+                ttime = process.cpu_times()
+                time_formatted = str(timedelta(seconds=int(ttime.user + ttime.system)))
+                cmd = ' '.join(process.cmdline())
+
+                row = {'PID': pid, 'UID': uid, 'PPID': ppid, '%CPU': cpu, '%MEM': mem, 'STIME': stime_formatted,  'TT': tty, 'TIME': time_formatted, 'CMD': cmd}
+                processdata.append(row)
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                pass
+
         # wipe out all data before updating with new values
         self.state_db.delete_all_by_pattern('STATE_DB', 'PROCESS_STATS|*')
-        for row in processdata[0:]:
+
+        for row in processdata:
             cid = row.get('PID')
             if cid:
                 value = 'PROCESS_STATS|{}'.format(cid)
@@ -157,7 +170,7 @@ class ProcDockerStats(daemon_base.DaemonBase):
                 self.update_state_db(value, 'PPID', ppid)
                 cpu = row.get('%CPU')
                 self.update_state_db(value, '%CPU', str(cpu))
-                mem = row.get('%MEM')
+                mem = round(row.get('%MEM'), 1)
                 self.update_state_db(value, '%MEM', str(mem))
                 stime = row.get('STIME')
                 self.update_state_db(value, 'STIME', stime)

tests/procdockerstatsd_test.py

@@ -1,5 +1,6 @@
 import sys
 import os
+import psutil
 import pytest
 from unittest.mock import call, patch
 from swsscommon import swsscommon
@@ -18,6 +19,62 @@
 procdockerstatsd_path = os.path.join(scripts_path, 'procdockerstatsd')
 procdockerstatsd = load_module_from_source('procdockerstatsd', procdockerstatsd_path)
 
+class MockProcess:
+    def __init__(self, uids, pid, ppid, memory_percent, cpu_percent, create_time, cmdline, user_time, system_time):
+        self._uids = uids
+        self._pid = pid
+        self._ppid = ppid
+        self._memory_percent = memory_percent
+        self._cpu_percent = cpu_percent
+        self._create_time = create_time
+        self._terminal = cmdline
+        self._cmdline = cmdline
+        self._user_time = user_time
+        self._system_time = system_time
+
+    def uids(self):
+        return self._uids
+
+    def pid(self):
+        return self._pid
+
+    def ppid(self):
+        return self._ppid
+
+    def memory_percent(self):
+        return self._memory_percent
+
+    def cpu_percent(self):
+        return self._cpu_percent
+
+    def create_time(self):
+        return self._create_time
+
+    def terminal(self):
+        return self._terminal
+
+    def cmdline(self):
+        return self._cmdline
+
+    def cpu_times(self):
+        class CPUTimes:
+            def __init__(self, user_time, system_time):
+                self.user = user_time
+                self.system = system_time
+
+            def __getitem__(self, index):
+                if index == 0:
+                    return self.user
+                else:
+                    return self.system
+
+            def __iter__(self):
+                yield self.user
+                yield self.system
+
+        return CPUTimes(self._user_time, self._system_time)
+
+
 class TestProcDockerStatsDaemon(object):
     def test_convert_to_bytes(self):
         test_data = [
@@ -51,11 +108,17 @@ def test_run_command(self):
         assert output is None
 
     def test_update_processstats_command(self):
-        expected_calls = [call(["ps", "-eo", "uid,pid,ppid,%mem,%cpu,stime,tty,time,cmd", "--sort", "-%cpu"], ["head", "-1024"])]
         pdstatsd = procdockerstatsd.ProcDockerStats(procdockerstatsd.SYSLOG_IDENTIFIER)
-        with patch("procdockerstatsd.getstatusoutput_noshell_pipe", return_value=([0, 0], 'output')) as mock_cmd:
+
+        # Create a list of mocked processes
+        mocked_processes = [
+            MockProcess(uids=[1000], pid=1234, ppid=5678, memory_percent=10.5, cpu_percent=20.5, create_time=1234567890, cmdline=['python', 'script.py'], user_time=1.5, system_time=2.0),
+            MockProcess(uids=[1000], pid=5678, ppid=0, memory_percent=5.5, cpu_percent=15.5, create_time=9876543210, cmdline=['bash', 'script.sh'], user_time=3.5, system_time=4.0)
+        ]
+
+        with patch("procdockerstatsd.psutil.process_iter", return_value=mocked_processes) as mock_process_iter:
             pdstatsd.update_processstats_command()
-            mock_cmd.assert_has_calls(expected_calls)
+            mock_process_iter.assert_called_once()
 
     @patch('procdockerstatsd.getstatusoutput_noshell_pipe', return_value=([0, 0], ''))
     def test_update_fipsstats_command(self, mock_cmd):