ext: Monkeypatch os.waitpid to extract CPU time from subprocess

Added utility class `TimedWaitPID` which monkey-patches os.waitpid()
with a functor that has the same signature, but calls os.wait4()
instead. This allows the process's user and system CPU time to be
obtained from the OS when using APIs (such as subprocess) which use
os.waitpid() internally.

The process CPU time is stored within the functor and can be read back
later by calling TimedWaitPID.get_time_for_pid().

JIRA: https://gem5.atlassian.net/browse/GEM5-548

Change-Id: I9ebe9ca1241a4f28c90ad31f672f32ac52786664
Reviewed-on: https://gem5-review.googlesource.com/c/public/gem5/+/32652
Maintainer: Bobby R. Bruce <bbruce@ucdavis.edu>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Hoa Nguyen <hoanguyen@ucdavis.edu>

diff --git a/ext/testlib/helper.py b/ext/testlib/helper.py
index ff83409866319ead26c84500819030494d9e3936..01ca539dc3ff907703d29e00621e9a3eb21984de 100644 (file)
--- a/ext/testlib/helper.py
+++ b/ext/testlib/helper.py
@@ -1,3 +1,15 @@
+# Copyright (c) 2020 ARM Limited
+# All rights reserved
+#
+# The license below extends only to copyright in the software and shall
+# not be construed as granting a license to any other intellectual
+# property including but not limited to intellectual property relating
+# to a hardware implementation of the functionality of the software
+# licensed hereunder.  You may use the software subject to the license
+# terms below provided that you ensure that this notice is replicated
+# unmodified and in its entirety in all distributions of the software,
+# modified or unmodified, in source code or in binary form.
+#
 # Copyright (c) 2017 Mark D. Hill and David A. Wood
 # All rights reserved.
 #
@@ -29,7 +41,7 @@
 '''
 Helper classes for writing tests with this test library.
 '''
-from collections import MutableSet
+from collections import MutableSet, namedtuple
 
 import difflib
 import errno
@@ -42,6 +54,83 @@ import tempfile
 import threading
 import time
 
+class TimedWaitPID(object):
+    """Utility to monkey-patch os.waitpid() with os.wait4().
+
+    This allows process usage time to be obtained directly from the OS
+    when used with APIs, such as `subprocess`, which use os.waitpid to
+    join child processes.
+
+    The resource usage data from os.wait4() is stored in a functor and
+    can be obtained using the get_time_for_pid() method.
+
+    To avoid unbounded memory usage, the time record is deleted after
+    it is read.
+
+    """
+    TimeRecord = namedtuple( "_TimeRecord", "user_time system_time" )
+
+    class Wrapper(object):
+        def __init__(self):
+            self._time_for_pid = {}
+            self._access_lock = threading.Lock()
+
+        def __call__(self, pid, options):
+            pid, status, resource_usage = os.wait4(pid, options)
+            with self._access_lock:
+                self._time_for_pid[pid] = (
+                    TimedWaitPID.TimeRecord(
+                        resource_usage.ru_utime,
+                        resource_usage.ru_stime
+                    )
+                )
+            return (pid, status)
+
+        def has_time_for_pid(self, pid):
+            with self._access_lock:
+                return pid in self._time_for_pid
+
+        def get_time_for_pid(self, pid):
+            with self._access_lock:
+                if pid not in self._time_for_pid:
+                    raise Exception("No resource usage for pid {}".format(pid))
+                time_for_pid = self._time_for_pid[pid]
+                del self._time_for_pid[pid]
+                return time_for_pid
+
+    _wrapper = None
+    _wrapper_lock = threading.Lock()
+    _original_os_waitpid = None
+
+    @staticmethod
+    def install():
+        with TimedWaitPID._wrapper_lock:
+            if TimedWaitPID._wrapper is None:
+                TimedWaitPID._wrapper = TimedWaitPID.Wrapper()
+            if TimedWaitPID._original_os_waitpid is None :
+                TimedWaitPID._original_os_waitpid = os.waitpid
+                os.waitpid = TimedWaitPID._wrapper
+
+    @staticmethod
+    def restore():
+        with TimedWaitPID._wrapper_lock:
+            if TimedWaitPID._original_os_waitpid is not None :
+                os.waitpid = TimedWaitPID._original_os_waitpid
+                TimedWaitPID._original_os_waitpid = None
+
+    @staticmethod
+    def has_time_for_pid(pid):
+        with TimedWaitPID._wrapper_lock:
+            return TimedWaitPID._wrapper.has_time_for_pid(pid)
+
+    @staticmethod
+    def get_time_for_pid(pid):
+        with TimedWaitPID._wrapper_lock:
+            return TimedWaitPID._wrapper.get_time_for_pid(pid)
+
+# Patch os.waitpid()
+TimedWaitPID.install()
+
 #TODO Tear out duplicate logic from the sandbox IOManager
 def log_call(logger, command, *popenargs, **kwargs):
     '''