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
@@ -53,10 +54,56 @@ 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=lambda: psutil._common.pids(), pid=1234, ppid=lambda: 5678, memory_percent=lambda: 10.5, cpu_percent=lambda: 20.5, create_time=lambda: 1234567890, cmdline=lambda: ['python', 'script.py']),
+            MockProcess(uids=lambda: psutil._common.pids(), pid=5678, ppid=lambda: 0, memory_percent=lambda: 5.5, cpu_percent=lambda: 15.5, create_time=lambda: 9876543210, cmdline=lambda: ['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)