From a833224c95c859724fb949bff8bbfb2bb80ac367 Mon Sep 17 00:00:00 2001
From: Michael Droettboom <mdboom@gmail.com>
Date: Tue, 27 Aug 2024 07:37:54 -0400
Subject: [PATCH] Make hooks work in bench_command subprocesses (#197)

* Make hooks work in bench_command subprocesses

* Fix tests
---
 pyperf/_command.py      | 11 ++++++++---
 pyperf/_hooks.py        | 13 +++++++++++++
 pyperf/_process_time.py | 38 +++++++++++++++++++++++++++++++++++---
 pyperf/_worker.py       | 11 ++---------
 4 files changed, 58 insertions(+), 15 deletions(-)

diff --git a/pyperf/_command.py b/pyperf/_command.py
index 465d971a..592bf5c0 100644
--- a/pyperf/_command.py
+++ b/pyperf/_command.py
@@ -1,4 +1,5 @@
 import functools
+import json
 import os.path
 import subprocess
 import sys
@@ -23,18 +24,22 @@ def bench_command(command, task, loops):
                         % proc.returncode)
 
     rss = None
+    metadata = {}
     try:
         lines = output.splitlines()
         timing = float(lines[0])
-        if len(lines) >= 2:
-            rss = int(lines[1])
+        rss = int(lines[1])
+        metadata = json.loads(lines[2])
     except ValueError:
         raise ValueError("failed to parse script output: %r" % output)
 
-    if rss:
+    if rss and rss > 0:
         # store the maximum
         max_rss = task.metadata.get('command_max_rss', 0)
         task.metadata['command_max_rss'] = max(max_rss, rss)
+
+    task.metadata.update(metadata)
+
     return timing
 
 
diff --git a/pyperf/_hooks.py b/pyperf/_hooks.py
index cb621047..226e150f 100644
--- a/pyperf/_hooks.py
+++ b/pyperf/_hooks.py
@@ -31,6 +31,19 @@ def get_selected_hooks(hook_names):
         yield hook_mapping[hook_name]
 
 
+def instantiate_selected_hooks(hook_names):
+    hook_managers = {}
+    for hook in get_selected_hooks(hook_names):
+        try:
+            hook_managers[hook.name] = hook.load()()
+        except HookError as e:
+            print(f"ERROR setting up hook '{hook.name}':", file=sys.stderr)
+            print(str(e), file=sys.stderr)
+            sys.exit(1)
+
+    return hook_managers
+
+
 class HookError(Exception):
     pass
 
diff --git a/pyperf/_process_time.py b/pyperf/_process_time.py
index 88bceb4e..199dc885 100644
--- a/pyperf/_process_time.py
+++ b/pyperf/_process_time.py
@@ -14,6 +14,8 @@
 If resource.getrusage() is available: compute the maximum RSS memory in bytes
 per process and writes it into stdout as a second line.
 """
+import contextlib
+import json
 import os
 import subprocess
 import sys
@@ -91,6 +93,27 @@ def bench_process(loops, args, kw, profile_filename=None):
     return (dt, max_rss)
 
 
+def load_hooks(metadata):
+    hook_names = []
+    while "--hook" in sys.argv:
+        hook_idx = sys.argv.index("--hook")
+        hook_name = sys.argv[hook_idx + 1]
+        hook_names.append(hook_name)
+        del sys.argv[hook_idx]
+        del sys.argv[hook_idx]
+
+    if len(hook_names):
+        # Only import pyperf if we know we have hooks
+        import pyperf._hooks
+
+        hook_managers = pyperf._hooks.instantiate_selected_hooks(hook_names)
+        metadata["hooks"] = ", ".join(hook_managers.values())
+    else:
+        hook_managers = {}
+
+    return hook_managers
+
+
 def main():
     # Make sure that the pyperf module wasn't imported
     if 'pyperf' in sys.modules:
@@ -111,6 +134,9 @@ def main():
     else:
         profile_filename = None
 
+    metadata = {}
+    hook_managers = load_hooks(metadata)
+
     loops = int(sys.argv[1])
     args = sys.argv[2:]
 
@@ -125,15 +151,21 @@ def main():
         kw['stdout'] = devnull
     kw['stderr'] = subprocess.STDOUT
 
-    dt, max_rss = bench_process(loops, args, kw, profile_filename)
+    with contextlib.ExitStack() as stack:
+        for hook in hook_managers.values():
+            stack.enter_context(hook)
+        dt, max_rss = bench_process(loops, args, kw, profile_filename)
 
     if devnull is not None:
         devnull.close()
 
+    for hook in hook_managers.values():
+        hook.teardown(metadata)
+
     # Write timing in seconds into stdout
     print(dt)
-    if max_rss:
-        print(max_rss)
+    print(max_rss or -1)
+    print(json.dumps(metadata))
 
 
 if __name__ == "__main__":
diff --git a/pyperf/_worker.py b/pyperf/_worker.py
index 5d423e8b..1c2b31dd 100644
--- a/pyperf/_worker.py
+++ b/pyperf/_worker.py
@@ -6,7 +6,7 @@
 import pyperf
 from pyperf._formatter import (format_number, format_value, format_values,
                                format_timedelta)
-from pyperf._hooks import get_selected_hooks, HookError
+from pyperf._hooks import instantiate_selected_hooks
 from pyperf._utils import MS_WINDOWS, MAC_OS, percentile, median_abs_dev
 
 
@@ -60,14 +60,7 @@ def _compute_values(self, values, nvalue,
 
         task_func = self.task_func
 
-        hook_managers = {}
-        for hook in get_selected_hooks(args.hook):
-            try:
-                hook_managers[hook.name] = hook.load()()
-            except HookError as e:
-                print(f"ERROR setting up hook '{hook.__name__}:'", file=sys.stderr)
-                print(str(e), file=sys.stderr)
-                sys.exit(1)
+        hook_managers = instantiate_selected_hooks(args.hook)
         if len(hook_managers):
             self.metadata["hooks"] = ", ".join(hook_managers.keys())