From 8cd1470bb835f4d264fb7bcf4ce0fdfbfa4d245b Mon Sep 17 00:00:00 2001
From: Valentin Schneider <valentin.schneider@arm.com>
Date: Fri, 24 Aug 2018 10:27:47 +0100
Subject: [PATCH] module/cpufreq: Add a contextmanager for temporary governor
 changes

We may sometime want to temporarily use another governor, and then
restore whatever governor was used previously - for instance, RT-app
calibration ([1]) needs to use the performance governor, but we don't
want this to interfere with e.g. our current experiment.

[1]: https://github.com/ARM-software/lisa/blob/master/libs/wlgen/wlgen/rta.py#L118
---
 devlib/module/cpufreq.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/devlib/module/cpufreq.py b/devlib/module/cpufreq.py
index 1f502da..9da0c8e 100644
--- a/devlib/module/cpufreq.py
+++ b/devlib/module/cpufreq.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from contextlib import contextmanager
+
 from devlib.module import Module
 from devlib.exception import TargetStableError
 from devlib.utils.misc import memoized
@@ -95,6 +97,39 @@ class CpufreqModule(Module):
         self.target.write_value(sysfile, governor)
         self.set_governor_tunables(cpu, governor, **kwargs)
 
+    @contextmanager
+    def use_governor(self, governor, cpus=None, **kwargs):
+        """
+        Use a given governor, then restore previous governor(s)
+
+        :param governor: Governor to use on all targeted CPUs (see :meth:`set_governor`)
+        :type governor: str
+
+        :param cpus: CPUs affected by the governor change (all by default)
+        :type cpus: list
+
+        :Keyword Arguments: Governor tunables, See :meth:`set_governor_tunables`
+        """
+        if not cpus:
+            cpus = range(self.target.number_of_cpus)
+
+        # Setting a governor & tunables for a cpu will set them for all cpus
+        # in the same clock domain, so only manipulating one cpu per domain
+        # is enough
+        domains = set(self.get_affected_cpus(cpu)[0] for cpu in cpus)
+        prev_governors = {cpu : (self.get_governor(cpu), self.get_governor_tunables(cpu))
+                          for cpu in domains}
+
+        for cpu in domains:
+            self.set_governor(cpu, governor, **kwargs)
+
+        try:
+            yield
+
+        finally:
+            for cpu, (governor, tunables) in prev_governors.items():
+                self.set_governor(cpu, governor, **tunables)
+
     def list_governor_tunables(self, cpu):
         """Returns a list of tunables available for the governor on the specified CPU."""
         if isinstance(cpu, int):