perf: improve iter_process speed by directly access psutil internal functions

deploy/Windows/utils.py

@@ -1,6 +1,5 @@
 import os
 import re
-from dataclasses import dataclass
 from typing import Callable, Generic, Iterable, TypeVar
 
 from deploy.Windows.atomic import atomic_read, atomic_write
@@ -67,7 +66,6 @@ def poor_yaml_read(file):
         dict:
     """
     content = atomic_read(file)
-
     data = {}
     regex = re.compile(r'^(.*?):(.*?)$')
     for line in content.splitlines():
@@ -113,38 +111,12 @@ def poor_yaml_write(data, file, template_file=DEPLOY_TEMPLATE):
     atomic_write(file, text)
 
 
-@dataclass
-class DataProcessInfo:
-    proc: object  # psutil.Process or psutil._pswindows.Process
-    pid: int
-
-    @cached_property
-    def name(self):
-        try:
-            name = self.proc.name()
-        except:
-            name = ''
-        return name
-
-    @cached_property
-    def cmdline(self):
-        try:
-            cmdline = self.proc.cmdline()
-        except:
-            # psutil.AccessDenied
-            # NoSuchProcess: process no longer exists (pid=xxx)
-            cmdline = []
-        cmdline = ' '.join(cmdline).replace(r'\\', '/').replace('\\', '/')
-        return cmdline
-
-    def __str__(self):
-        # Don't print `proc`, it will take some time to get process properties
-        return f'DataProcessInfo(name="{self.name}", pid={self.pid}, cmdline="{self.cmdline}")'
-
-    __repr__ = __str__
-
-
-def iter_process() -> Iterable[DataProcessInfo]:
+def iter_process() -> "Iterable[tuple[int, list[str]]]":
+    """
+    Yields:
+        int: pid
+        list[str]: cmdline, and it's guaranteed to have at least one element
+    """
     try:
         import psutil
     except ModuleNotFoundError:
@@ -154,16 +126,48 @@ def iter_process() -> Iterable[DataProcessInfo]:
         # Since this is a one-time-usage, we access psutil._psplatform.Process directly
         # to bypass the call of psutil.Process.is_running().
         # This only costs about 0.017s.
+        # If you do psutil.process_iter(['pid', 'cmdline']) it will take over 1s
+        import psutil._psutil_windows as cetx
+        for pid in psutil.pids():
+            # 0 and 4 are always represented in taskmgr and process-hacker
+            if pid == 0 or pid == 4:
+                continue
+            try:
+                # This would be fast on psutil<=5.9.8 taking overall time 0.027s
+                # but taking 0.39s on psutil>=6.0.0
+                cmdline = cetx.proc_cmdline(pid, use_peb=True)
+            except (psutil.AccessDenied, psutil.NoSuchProcess, OSError):
+                # psutil.AccessDenied
+                # NoSuchProcess: process no longer exists (pid=xxx)
+                # ProcessLookupError: [Errno 3] assume no such process (originated from psutil_pid_is_running -> 0)
+                # OSError: [WinError 87] 参数错误。: '(originated from ReadProcessMemory)'
+                continue
+
+            # Validate cmdline
+            if not cmdline:
+                continue
+            try:
+                exe = cmdline[0]
+            except IndexError:
+                continue
+            # \??\C:\Windows\system32\conhost.exe
+            if exe.startswith(r'\??'):
+                continue
+            yield pid, cmdline
+    else:
+        # No optimizations yet
         for pid in psutil.pids():
             proc = psutil._psplatform.Process(pid)
-            yield DataProcessInfo(
-                proc=proc,
-                pid=proc.pid,
-            )
-    else:
-        # This will cost about 0.45s, even `attr` is given.
-        for proc in psutil.process_iter():
-            yield DataProcessInfo(
-                proc=proc,
-                pid=proc.pid,
-            )
+            try:
+                cmdline = proc.cmdline()
+            except (psutil.AccessDenied, psutil.NoSuchProcess, OSError):
+                continue
+
+            # Validate cmdline
+            if not cmdline:
+                continue
+            try:
+                cmdline[0]
+            except IndexError:
+                continue
+            yield pid, cmdline