From b19a330461972b601d1cd5dbbc66fac9442cb064 Mon Sep 17 00:00:00 2001
From: Feng Pan <fenpan@microsoft.com>
Date: Tue, 22 Oct 2024 12:35:40 +0000
Subject: [PATCH] [not4review] test

---
 scripts/procdockerstatsd       | 18 ++++++++++++++++--
 tests/procdockerstatsd_test.py |  5 +++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/scripts/procdockerstatsd b/scripts/procdockerstatsd
index bdc04a66..60e81cf2 100755
--- a/scripts/procdockerstatsd
+++ b/scripts/procdockerstatsd
@@ -24,6 +24,7 @@ REDIS_HOSTIP = "127.0.0.1"
 
 
 class ProcDockerStats(daemon_base.DaemonBase):
+    all_process_obj = {}
 
     def __init__(self, log_identifier):
         super(ProcDockerStats, self).__init__(log_identifier)
@@ -138,10 +139,18 @@ class ProcDockerStats(daemon_base.DaemonBase):
 
     def update_processstats_command(self):
         processdata = []
-        for process in psutil.process_iter(['pid', 'ppid', 'memory_percent', 'cpu_percent', 'create_time', 'cmdline']):
+        pid_set = set()
+        for process_obj in psutil.process_iter(['pid', 'ppid', 'memory_percent', 'cpu_percent', 'create_time', 'cmdline']):
             try:
+                pid = process_obj.pid
+                pid_set.add(pid)
+                # store process object and reuse for CPU utilization
+                if pid in self.all_process_obj:
+                    process = self.all_process_obj[pid]
+                else:
+                    self.all_process_obj[pid] = process_obj
+                    process = process_obj
                 uid = process.uids()[0]
-                pid = process.pid
                 ppid = process.ppid()
                 mem = process.memory_percent()
                 cpu = process.cpu_percent()
@@ -157,6 +166,11 @@ class ProcDockerStats(daemon_base.DaemonBase):
             except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                 pass
 
+        # erase dead process
+        for id in list(self.all_process_obj.keys()):
+            if id not in pid_set:
+                del self.all_process_obj[id]
+
         # wipe out all data before updating with new values
         self.state_db.delete_all_by_pattern('STATE_DB', 'PROCESS_STATS|*')
 
diff --git a/tests/procdockerstatsd_test.py b/tests/procdockerstatsd_test.py
index a15563bc..9dd022e9 100644
--- a/tests/procdockerstatsd_test.py
+++ b/tests/procdockerstatsd_test.py
@@ -120,8 +120,13 @@ def test_update_processstats_command(self):
         ]
 
         with patch("procdockerstatsd.psutil.process_iter", return_value=mocked_processes) as mock_process_iter:
+            pdstatsd.all_process_obj = {1234: MockProcess(uids=[1000], pid=1234, ppid=0, memory_percent=5.5, cpu_percent=99, create_time=valid_create_time2, cmdline=['bash', 'script.sh'], user_time=3.5, system_time=4.0),
+                                        6666: MockProcess(uids=[1000], pid=6666, ppid=0, memory_percent=5.5, cpu_percent=15.5, create_time=valid_create_time2, cmdline=['bash', 'script.sh'], user_time=3.5, system_time=4.0)}
             pdstatsd.update_processstats_command()
             mock_process_iter.assert_called_once()
+            assert pdstatsd.all_process_obj.get(1234).cpu_percent() == 99
+            assert pdstatsd.all_process_obj.get(5678).cpu_percent() == 15.5
+            assert pdstatsd.all_process_obj.get(6666) is None
 
     @patch('procdockerstatsd.getstatusoutput_noshell_pipe', return_value=([0, 0], ''))
     def test_update_fipsstats_command(self, mock_cmd):