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

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,6 +5,7 @@ Daemon which periodically gathers process and docker statistics and pushes the d
 '''
 
 import os
+import psutil
 import re
 import subprocess
 import sys
@@ -136,18 +137,28 @@ 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().real
+                pid = process.pid
+                ppid = process.ppid()
+                mem = process.memory_percent()
+                cpu = process.cpu_percent()
+                stime = process.create_time()
+                tty = process.terminal()
+                time = process.cpu_times().user + process.cpu_times().system
+                cmd = ' '.join(process.cmdline())
+
+                row = {'PID': pid, 'UID': uid, 'PPID': ppid, '%CPU': cpu, '%MEM': mem, 'STIME': stime,  'TT': tty, 'TIME': time, '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)

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,37 @@
 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):
+        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._cmdline = cmdline
+
+    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 cmdline(self):
+        return self._cmdline
+
 class TestProcDockerStatsDaemon(object):
     def test_convert_to_bytes(self):
         test_data = [
@@ -53,10 +85,57 @@ def test_run_command(self):
     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:
-            pdstatsd.update_processstats_command()
+
+        # Patch the 'psutil.process_iter' function to return a list of mock processes
+        with patch("psutil.process_iter") as mock_process_iter:
+            mock_processes = [
+                psutil.Process(pid=123, name='test_process1', status='running', cmdline=['command']),
+                psutil.Process(pid=456, name='test_process2', status='running', cmdline=['command2']),
+            ]
+            mock_process_iter.return_value = mock_processes
+
+            # Patch the 'procdockerstatsd.getstatusoutput_noshell_pipe' function
+            with patch("procdockerstatsd.getstatusoutput_noshell_pipe", return_value=([0, 0], 'output')) as mock_cmd:
+                pdstatsd.update_processstats_command()
+
+            # Perform assertions
             mock_cmd.assert_has_calls(expected_calls)
 
+    def test_update_processstats_command(self):
+        pdstatsd = procdockerstatsd.ProcDockerStats(procdockerstatsd.SYSLOG_IDENTIFIER)
+
+        # 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']),
+            MockProcess(uids=[1000], pid=5678, ppid=0, memory_percent=5.5, cpu_percent=15.5, create_time=9876543210, cmdline=['bash', 'script.sh'])
+        ]
+
+        with patch("procdockerstatsd.psutil.process_iter", return_value=mocked_processes) as mock_process_iter, \
+            patch("procdockerstatsd.getstatusoutput_noshell_pipe", return_value=([0, 0], 'output')) as mock_cmd:
+            pdstatsd.update_processstats_command()
+            mock_cmd.assert_called_once()
+            mock_process_iter.assert_called_once()
+
+            # Assert the expected calls to self.update_state_db
+            self.update_state_db.assert_has_calls([
+                call('PROCESS_STATS|1234', 'UID', psutil._common.pids()),
+                call('PROCESS_STATS|1234', 'PPID', 5678),
+                call('PROCESS_STATS|1234', '%CPU', '20.5'),
+                call('PROCESS_STATS|1234', '%MEM', '10.5'),
+                call('PROCESS_STATS|1234', 'STIME', 1234567890),
+                call('PROCESS_STATS|1234', 'TT', None),
+                call('PROCESS_STATS|1234', 'TIME', 1234567890),
+                call('PROCESS_STATS|1234', 'CMD', 'python script.py'),
+                call('PROCESS_STATS|5678', 'UID', psutil._common.pids()),
+                call('PROCESS_STATS|5678', 'PPID', 0),
+                call('PROCESS_STATS|5678', '%CPU', '15.5'),
+                call('PROCESS_STATS|5678', '%MEM', '5.5'),
+                call('PROCESS_STATS|5678', 'STIME', 9876543210),
+                call('PROCESS_STATS|5678', 'TT', None),
+                call('PROCESS_STATS|5678', 'TIME', 9876543210),
+                call('PROCESS_STATS|5678', 'CMD', 'bash script.sh')
+            ])
+
     @patch('procdockerstatsd.getstatusoutput_noshell_pipe', return_value=([0, 0], ''))
     def test_update_fipsstats_command(self, mock_cmd):
         pdstatsd = procdockerstatsd.ProcDockerStats(procdockerstatsd.SYSLOG_IDENTIFIER)