From b667f6641ac031e26641cc75ed0956eaa64d5458 Mon Sep 17 00:00:00 2001 From: Feng Pan 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 | 83 +++++++++++++++++++++++++++++++++- 3 files changed, 103 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..0f7b1602 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 @@ -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)