From 0a271c2d8201f14f9fc7a3c7eb3bbd3f6b5f962e Mon Sep 17 00:00:00 2001
From: lilydjwg <lilydjwg@gmail.com>
Date: Sat, 29 Jun 2024 23:43:22 +0800
Subject: [PATCH] use systemd's CPUUsageNSec and MemoryPeak properties if
 available

---
 lilac2/systemd.py | 91 +++++++++++++++++++++++++++++++++++------------
 1 file changed, 68 insertions(+), 23 deletions(-)

diff --git a/lilac2/systemd.py b/lilac2/systemd.py
index 0691ddc..09c9db9 100644
--- a/lilac2/systemd.py
+++ b/lilac2/systemd.py
@@ -1,6 +1,6 @@
 import os
 import subprocess
-from typing import Generator, Any
+from typing import Generator, Any, Optional
 import select
 import time
 import logging
@@ -13,7 +13,7 @@
 _available = None
 _check_lock = threading.Lock()
 
-def available() -> bool:
+def available() -> bool | dict[str, bool]:
   global _available
 
   with _check_lock:
@@ -22,14 +22,61 @@ def available() -> bool:
       logger.debug('systemd availability: %s', _available)
   return _available
 
-def _check_availability() -> bool:
+def _cgroup_memory_usage(cgroup: str) -> int:
+  mem_file = f'/sys/fs/cgroup{cgroup}/memory.peak'
+  with open(mem_file) as f:
+    return int(f.read().rstrip())
+
+def _cgroup_cpu_usage(cgroup: str) -> int:
+  cpu_file = f'/sys/fs/cgroup{cgroup}/cpu.stat'
+  with open(cpu_file) as f:
+    for l in f:
+      if l.startswith('usage_usec '):
+        return int(l.split()[1]) * 1000
+  return 0
+
+def _check_availability() -> bool | dict[str, bool]:
   if 'DBUS_SESSION_BUS_ADDRESS' not in os.environ:
     dbus = f'/run/user/{os.getuid()}/bus'
     if not os.path.exists(dbus):
       return False
     os.environ['DBUS_SESSION_BUS_ADDRESS'] = f'unix:path={dbus}'
-  p = subprocess.run(['systemd-run', '--quiet', '--user', '-u', 'lilac-check', 'true'])
-  return p.returncode == 0
+  p = subprocess.run([
+    'systemd-run', '--quiet', '--user',
+    '--remain-after-exit', '-u', 'lilac-check', 'true',
+  ])
+  if p.returncode != 0:
+    return False
+
+  try:
+    ps: dict[str, Optional[int]] = {
+      'CPUUsageNSec': None,
+      'MemoryPeak': None,
+    }
+    _read_service_int_properties('lilac-check', ps)
+
+    ret = {}
+    for k, v in ps.items():
+      ret[k] = v is not None
+
+    return ret
+  finally:
+    subprocess.run(['systemctl', '--user', 'stop', '--quiet', 'lilac-check'])
+
+def _read_service_int_properties(name: str, properties: dict[str, Optional[int]]) -> None:
+  cmd = [
+    'systemctl', '--user', 'show', f'{name}.service',
+  ] + [f'--property={k}' for k in properties]
+
+  out = subprocess.check_output(cmd, text=True)
+  for l in out.splitlines():
+    k, v = l.split('=', 1)
+    if k in properties:
+      try:
+        properties[k] = int(v)
+      except ValueError:
+        # [not set]
+        pass
 
 def start_cmd(
   name: str, cmd: Cmd,
@@ -117,32 +164,30 @@ def poll_rusage(name: str, deadline: float) -> tuple[RUsage, bool]:
       logger.warning('%s.service already finished: %s', name, state)
       return RUsage(0, 0), False
 
-    mem_file = f'/sys/fs/cgroup{cgroup}/memory.peak'
-
+    nsec = 0
     mem_max = 0
+    availability = available()
+    assert isinstance(availability, dict)
     for _ in _poll_cmd(pid):
-      with open(mem_file) as f:
-        mem_cur = int(f.read().rstrip())
-        mem_max = max(mem_cur, mem_max)
+      if not availability['CPUUsageNSec']:
+        nsec = _cgroup_cpu_usage(cgroup)
+      if not availability['MemoryPeak']:
+        mem_max = _cgroup_memory_usage(cgroup)
       if time.time() > deadline:
         timedout = True
         break
 
     # systemd will remove the cgroup as soon as the process exits
     # instead of racing with systemd, we just ask it for the data
-    nsec = 0
-    out = subprocess.check_output([
-      'systemctl', '--user', 'show', f'{name}.service',
-      '--property=CPUUsageNSec',
-    ], text=True)
-    for l in out.splitlines():
-      k, v = l.split('=', 1)
-      if k == 'CPUUsageNSec':
-        try:
-          nsec = int(v)
-        except ValueError:
-          # [not set]
-          pass
+    ps: dict[str, Optional[int]] = {
+      'CPUUsageNSec': None,
+      'MemoryPeak': None,
+    }
+    _read_service_int_properties(name, ps)
+    if n := ps['CPUUsageNSec']:
+      nsec = n
+    if n := ps['MemoryPeak']:
+      mem_max = n
 
   finally:
     logger.debug('stopping worker service')