From 3fca792a26df488bbb7130e9d2b66c9e6e147f45 Mon Sep 17 00:00:00 2001
From: Feng Pan <fenpan@microsoft.com>
Date: Thu, 29 Feb 2024 12:19:15 +0000
Subject: [PATCH] Update ProcessStats query by using API instead of parsing ps
 command.

---
 azure-pipelines.yml            |  1 +
 scripts/procdockerstatsd       | 31 +++++++++++++++++++++----------
 tests/procdockerstatsd_test.py | 17 +++++++++++++++--
 3 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index af214110..ca493785 100644
--- a/azure-pipelines.yml
+++ b/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
diff --git a/scripts/procdockerstatsd b/scripts/procdockerstatsd
index 19e579ae..5942b183 100755
--- a/scripts/procdockerstatsd
+++ b/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)
diff --git a/tests/procdockerstatsd_test.py b/tests/procdockerstatsd_test.py
index 40b222db..22925240 100644
--- a/tests/procdockerstatsd_test.py
+++ b/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,8 +54,20 @@ 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)
 
     @patch('procdockerstatsd.getstatusoutput_noshell_pipe', return_value=([0, 0], ''))