From a2f761a4ab17e53a45b8f1fa4773c4e0149709fe Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Mon, 25 Sep 2023 10:55:21 -0400
Subject: [PATCH] Get CPU usage from cgroup v2.

---
 benchmarks/__init__.py |  8 ++++++--
 benchmarks/conftest.py | 31 ++++++++++++++++++++-----------
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py
index 57f0be071..bcefa8d3a 100644
--- a/benchmarks/__init__.py
+++ b/benchmarks/__init__.py
@@ -1,8 +1,12 @@
-"""pytest-based end-to-end benchmarks of Tahoe-LAFS.
+"""
+pytest-based end-to-end benchmarks of Tahoe-LAFS.
 
 Usage:
 
-$ pytest benchmark --number-of-nodes=3
+$ systemd-run --user --scope pytest benchmark --number-of-nodes=3
 
 It's possible to pass --number-of-nodes multiple times.
+
+The systemd-run makes sure the tests run in their own cgroup so we get CPU
+accounting correct.
 """
diff --git a/benchmarks/conftest.py b/benchmarks/conftest.py
index 2b60d8cdb..1d522108b 100644
--- a/benchmarks/conftest.py
+++ b/benchmarks/conftest.py
@@ -5,6 +5,7 @@ The number of nodes is parameterized via a --number-of-nodes CLI option added
 to pytest.
 """
 
+import os
 from resource import getrusage, RUSAGE_CHILDREN
 from shutil import which, rmtree
 from tempfile import mkdtemp
@@ -108,6 +109,23 @@ def client_node(request, grid, storage_nodes, number_of_nodes) -> Client:
     print(f"Client node pid: {client_node.process.transport.pid}")
     return client_node
 
+def get_cpu_time_for_cgroup():
+    """
+    Get how many CPU seconds have been used in current cgroup so far.
+
+    Assumes we're running in a v2 cgroup.
+    """
+    with open("/proc/self/cgroup") as f:
+        cgroup = f.read().strip().split(":")[-1]
+        assert cgroup.startswith("/")
+        cgroup = cgroup[1:]
+    cpu_stat = os.path.join("/sys/fs/cgroup", cgroup, "cpu.stat")
+    with open(cpu_stat) as f:
+        for line in f.read().splitlines():
+            if line.startswith("usage_usec"):
+                return int(line.split()[1]) / 1_000_000
+    raise ValueError("Failed to find usage_usec")
+
 
 class Benchmarker:
     """Keep track of benchmarking results."""
@@ -115,20 +133,11 @@ class Benchmarker:
     @contextmanager
     def record(self, capsys: pytest.CaptureFixture[str], name, **parameters):
         """Record the timing of running some code, if it succeeds."""
-        process = Process()
-
-        def get_children_cpu_time():
-            cpu = 0
-            for subprocess in process.children():
-                usage = subprocess.cpu_times()
-                cpu += usage.system + usage.user
-            return cpu
-
-        start_cpu = get_children_cpu_time()
+        start_cpu = get_cpu_time_for_cgroup()
         start = time()
         yield
         elapsed = time() - start
-        end_cpu = get_children_cpu_time()
+        end_cpu = get_cpu_time_for_cgroup()
         elapsed_cpu = end_cpu - start_cpu
         # FOR now we just print the outcome:
         parameters = " ".join(f"{k}={v}" for (k, v) in parameters.items())