Rework profiler to act as a context manager

7 files changed, 117 insertions, 146 deletions

diff --git a/buildstream/_context.py b/buildstream/_context.py
index afbf69726..a7d308a6f 100644
--- a/buildstream/_context.py
+++ b/buildstream/_context.py
@@ -30,7 +30,7 @@ from . import _site
 from . import _yaml
 from ._exceptions import LoadError, LoadErrorReason, BstError
 from ._message import Message, MessageType
-from ._profile import Topics, profile_start, profile_end
+from ._profile import Topics, PROFILER
 from ._artifactcache import ArtifactCache
 from ._sourcecache import SourceCache
 from ._cas import CASCache, CASQuota, CASCacheUsage
@@ -181,9 +181,8 @@ class Context():
     # override that configuration with the configuration file indicated
     # by *config*, if any was specified.
     #
+    @PROFILER.profile(Topics.LOAD_CONTEXT, "load")
     def load(self, config=None):
-        profile_start(Topics.LOAD_CONTEXT, 'load')
-
         # If a specific config file is not specified, default to trying
         # a $XDG_CONFIG_HOME/buildstream.conf file
         #
@@ -316,8 +315,6 @@ class Context():
                                  'strict', 'default-mirror',
                                  'remote-execution'])
 
-        profile_end(Topics.LOAD_CONTEXT, 'load')
-
     @property
     def artifactcache(self):
         if not self._artifactcache:
diff --git a/buildstream/_loader/loader.py b/buildstream/_loader/loader.py
index 4c2bfec77..965edbee8 100644
--- a/buildstream/_loader/loader.py
+++ b/buildstream/_loader/loader.py
@@ -25,7 +25,7 @@ from .._exceptions import LoadError, LoadErrorReason
 from .. import Consistency
 from .. import _yaml
 from ..element import Element
-from .._profile import Topics, profile_start, profile_end
+from .._profile import Topics, PROFILER
 from .._includes import Includes
 
 from .types import Symbol
@@ -108,12 +108,11 @@ class Loader():
         target_elements = []
 
         for target in targets:
-            profile_start(Topics.LOAD_PROJECT, target)
-            _junction, name, loader = self._parse_name(target, rewritable, ticker,
-                                                       fetch_subprojects=fetch_subprojects)
-            element = loader._load_file(name, rewritable, ticker, fetch_subprojects)
-            target_elements.append(element)
-            profile_end(Topics.LOAD_PROJECT, target)
+            with PROFILER.profile(Topics.LOAD_PROJECT, target):
+                _junction, name, loader = self._parse_name(target, rewritable, ticker,
+                                                           fetch_subprojects=fetch_subprojects)
+                element = loader._load_file(name, rewritable, ticker, fetch_subprojects)
+                target_elements.append(element)
 
         #
         # Now that we've resolve the dependencies, scan them for circular dependencies
@@ -127,10 +126,8 @@ class Loader():
             for element in target_elements
         )
 
-        profile_key = "_".join(t for t in targets)
-        profile_start(Topics.CIRCULAR_CHECK, profile_key)
-        self._check_circular_deps(dummy_target)
-        profile_end(Topics.CIRCULAR_CHECK, profile_key)
+        with PROFILER.profile(Topics.CIRCULAR_CHECK, "_".join(targets)):
+            self._check_circular_deps(dummy_target)
 
         ret = []
         #
@@ -138,9 +135,9 @@ class Loader():
         #
         for element in target_elements:
             loader = element._loader
-            profile_start(Topics.SORT_DEPENDENCIES, element.name)
-            loader._sort_dependencies(element)
-            profile_end(Topics.SORT_DEPENDENCIES, element.name)
+            with PROFILER.profile(Topics.SORT_DEPENDENCIES, element.name):
+                loader._sort_dependencies(element)
+
             # Finally, wrap what we have into LoadElements and return the target
             #
             ret.append(loader._collect_element(element))
diff --git a/buildstream/_pipeline.py b/buildstream/_pipeline.py
index 5dec9001d..c176b82f6 100644
--- a/buildstream/_pipeline.py
+++ b/buildstream/_pipeline.py
@@ -28,7 +28,7 @@ from pyroaring import BitMap  # pylint: disable=no-name-in-module
 
 from ._exceptions import PipelineError
 from ._message import Message, MessageType
-from ._profile import Topics, profile_start, profile_end
+from ._profile import Topics, PROFILER
 from . import Scope, Consistency
 from ._project import ProjectRefStorage
 
@@ -107,22 +107,19 @@ class Pipeline():
         # First concatenate all the lists for the loader's sake
         targets = list(itertools.chain(*target_groups))
 
-        profile_start(Topics.LOAD_PIPELINE, "_".join(t.replace(os.sep, '-') for t in targets))
+        with PROFILER.profile(Topics.LOAD_PIPELINE, "_".join(t.replace(os.sep, "-") for t in targets)):
+            elements = self._project.load_elements(targets,
+                                                   rewritable=rewritable,
+                                                   fetch_subprojects=fetch_subprojects)
 
-        elements = self._project.load_elements(targets,
-                                               rewritable=rewritable,
-                                               fetch_subprojects=fetch_subprojects)
+            # Now create element groups to match the input target groups
+            elt_iter = iter(elements)
+            element_groups = [
+                [next(elt_iter) for i in range(len(group))]
+                for group in target_groups
+            ]
 
-        # Now create element groups to match the input target groups
-        elt_iter = iter(elements)
-        element_groups = [
-            [next(elt_iter) for i in range(len(group))]
-            for group in target_groups
-        ]
-
-        profile_end(Topics.LOAD_PIPELINE, "_".join(t.replace(os.sep, '-') for t in targets))
-
-        return tuple(element_groups)
+            return tuple(element_groups)
 
     # resolve_elements()
     #
diff --git a/buildstream/_profile.py b/buildstream/_profile.py
index f29e070c4..3931a4cf0 100644
--- a/buildstream/_profile.py
+++ b/buildstream/_profile.py
@@ -18,18 +18,16 @@
 #  Authors:
 #        Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
 #        James Ennis <james.ennis@codethink.co.uk>
+#        Benjamin Schubert <bschubert15@bloomberg.net>
 
+
+import contextlib
 import cProfile
 import pstats
 import os
 import datetime
 import time
 
-# Track what profile topics are active
-active_topics = set()
-active_profiles = {}
-initialized = False
-
 
 # Use the topic values here to decide what to profile
 # by setting them in the BST_PROFILE environment variable.
@@ -55,99 +53,89 @@ class Topics():
     ALL = 'all'
 
 
-class Profile():
-    def __init__(self, topic, key, message):
-        self.message = message
-        self.key = topic + '-' + key
-        self.start = time.time()
+class _Profile:
+    def __init__(self, key, message):
         self.profiler = cProfile.Profile()
+
+        self.key = key
+        self.message = message
+
+        self.start_time = time.time()
+        filename_template = os.path.join(
+            os.getcwd(),
+            "profile-{}-{}".format(
+                datetime.datetime.fromtimestamp(self.start_time).strftime("%Y%m%dT%H%M%S"),
+                self.key.replace("/", "-").replace(".", "-")
+            )
+        )
+        self.log_filename = "{}.log".format(filename_template)
+        self.cprofile_filename = "{}.cprofile".format(filename_template)
+
+    def __enter__(self):
+        self.start()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.stop()
+        self.save()
+
+    def start(self):
         self.profiler.enable()
 
-    def end(self):
+    def stop(self):
         self.profiler.disable()
 
-        dt = datetime.datetime.fromtimestamp(self.start)
-        timestamp = dt.strftime('%Y%m%dT%H%M%S')
+    def save(self):
+        self._save_log()
+        self.profiler.dump_stats(self.cprofile_filename)
+
+    def _save_log(self):
+        heading = "\n".join([
+            "-" * 64,
+            "Profile for key: {}".format(self.key),
+            "Started at: {}".format(self.start_time),
+            "\n\t{}".format(self.message) if self.message else "",
+            "-" * 64,
+            ""  # for a final new line
+        ])
+
+        with open(self.log_filename, "a") as fp:
+            fp.write(heading)
+            ps = pstats.Stats(self.profiler, stream=fp).sort_stats("cumulative")
+            ps.print_stats()
 
-        filename = self.key.replace('/', '-')
-        filename = filename.replace('.', '-')
-        filename = os.path.join(os.getcwd(), 'profile-' + timestamp + '-' + filename)
 
-        time_ = dt.strftime('%Y-%m-%d %H:%M:%S')  # Human friendly format
-        self.__write_log(filename + '.log', time_)
+class _Profiler:
+    def __init__(self, settings):
+        self.active_topics = set()
+        self.enabled_topics = set()
 
-        self.__write_binary(filename + '.cprofile')
+        if settings:
+            self.enabled_topics = {
+                topic
+                for topic in settings.split(":")
+            }
 
-    ########################################
-    #            Private Methods           #
-    ########################################
+    @contextlib.contextmanager
+    def profile(self, topic, key, message=None):
+        if not self._is_profile_enabled(topic):
+            yield
+            return
 
-    def __write_log(self, filename, time_):
-        with open(filename, "a", encoding="utf-8") as f:
-            heading = '================================================================\n'
-            heading += 'Profile for key: {}\n'.format(self.key)
-            heading += 'Started at: {}\n'.format(time_)
-            if self.message:
-                heading += '\n    {}'.format(self.message)
-            heading += '================================================================\n'
-            f.write(heading)
-            ps = pstats.Stats(self.profiler, stream=f).sort_stats('cumulative')
-            ps.print_stats()
+        key = "{}-{}".format(topic, key)
 
-    def __write_binary(self, filename):
-        self.profiler.dump_stats(filename)
+        assert key not in self.active_topics
+        self.active_topics.add(key)
 
+        profiler = _Profile(key, message)
 
-# profile_start()
-#
-# Start profiling for a given topic.
-#
-# Args:
-#    topic (str): A topic name
-#    key (str): A key for this profile run
-#    message (str): An optional message to print in profile results
-#
-def profile_start(topic, key, message=None):
-    if not profile_enabled(topic):
-        return
+        with profiler:
+            yield
 
-    # Start profiling and hold on to the key
-    profile = Profile(topic, key, message)
-    assert active_profiles.get(profile.key) is None
-    active_profiles[profile.key] = profile
+        self.active_topics.remove(key)
 
+    def _is_profile_enabled(self, topic):
+        return topic in self.enabled_topics or Topics.ALL in self.enabled_topics
 
-# profile_end()
-#
-# Ends a profiling session previously
-# started with profile_start()
-#
-# Args:
-#    topic (str): A topic name
-#    key (str): A key for this profile run
-#
-def profile_end(topic, key):
-    if not profile_enabled(topic):
-        return
-
-    topic_key = topic + '-' + key
-    profile = active_profiles.get(topic_key)
-    assert profile
-    profile.end()
-    del active_profiles[topic_key]
-
-
-def profile_init():
-    global initialized  # pylint: disable=global-statement
-    if not initialized:
-        setting = os.getenv('BST_PROFILE')
-        if setting:
-            topics = setting.split(':')
-            for topic in topics:
-                active_topics.add(topic)
-        initialized = True
-
-
-def profile_enabled(topic):
-    profile_init()
-    return topic in active_topics or Topics.ALL in active_topics
+
+# Export a profiler to be used by BuildStream
+PROFILER = _Profiler(os.getenv("BST_PROFILE"))
diff --git a/buildstream/_project.py b/buildstream/_project.py
index f5da960cf..7a1734a25 100644
--- a/buildstream/_project.py
+++ b/buildstream/_project.py
@@ -30,7 +30,7 @@ from . import _cachekey
 from . import _site
 from . import _yaml
 from ._artifactelement import ArtifactElement
-from ._profile import Topics, profile_start, profile_end
+from ._profile import Topics, PROFILER
 from ._exceptions import LoadError, LoadErrorReason
 from ._options import OptionPool
 from ._artifactcache import ArtifactCache
@@ -156,9 +156,8 @@ class Project():
         self._fully_loaded = False
         self._project_includes = None
 
-        profile_start(Topics.LOAD_PROJECT, self.directory.replace(os.sep, '-'))
-        self._load(parent_loader=parent_loader)
-        profile_end(Topics.LOAD_PROJECT, self.directory.replace(os.sep, '-'))
+        with PROFILER.profile(Topics.LOAD_PROJECT, self.directory.replace(os.sep, '-')):
+            self._load(parent_loader=parent_loader)
 
         self._partially_loaded = True
 
diff --git a/buildstream/_scheduler/scheduler.py b/buildstream/_scheduler/scheduler.py
index 176900b33..50ad7f07a 100644
--- a/buildstream/_scheduler/scheduler.py
+++ b/buildstream/_scheduler/scheduler.py
@@ -29,7 +29,7 @@ from contextlib import contextmanager
 # Local imports
 from .resources import Resources, ResourceType
 from .jobs import JobStatus, CacheSizeJob, CleanupJob
-from .._profile import Topics, profile_start, profile_end
+from .._profile import Topics, PROFILER
 
 
 # A decent return code for Scheduler.run()
@@ -156,14 +156,11 @@ class Scheduler():
         self._check_cache_management()
 
         # Start the profiler
-        profile_start(Topics.SCHEDULER, "_".join(queue.action_name for queue in self.queues))
-
-        # Run the queues
-        self._sched()
-        self.loop.run_forever()
-        self.loop.close()
-
-        profile_end(Topics.SCHEDULER, "_".join(queue.action_name for queue in self.queues))
+        with PROFILER.profile(Topics.SCHEDULER, "_".join(queue.action_name for queue in self.queues)):
+            # Run the queues
+            self._sched()
+            self.loop.run_forever()
+            self.loop.close()
 
         # Stop handling unix signals
         self._disconnect_signals()
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index d2ece163d..73bc4badd 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -37,7 +37,7 @@ from ._message import Message, MessageType
 from ._scheduler import Scheduler, SchedStatus, TrackQueue, FetchQueue, \
     SourcePushQueue, BuildQueue, PullQueue, ArtifactPushQueue
 from ._pipeline import Pipeline, PipelineSelection
-from ._profile import Topics, profile_start, profile_end
+from ._profile import Topics, PROFILER
 from .types import _KeyStrength
 from . import utils, _yaml, _site
 from . import Scope, Consistency
@@ -117,19 +117,15 @@ class Stream():
                        except_targets=(),
                        use_artifact_config=False,
                        load_refs=False):
-
-        profile_start(Topics.LOAD_SELECTION, "_".join(t.replace(os.sep, '-') for t in targets))
-
-        target_objects, _ = self._load(targets, (),
-                                       selection=selection,
-                                       except_targets=except_targets,
-                                       fetch_subprojects=False,
-                                       use_artifact_config=use_artifact_config,
-                                       load_refs=load_refs)
-
-        profile_end(Topics.LOAD_SELECTION, "_".join(t.replace(os.sep, '-') for t in targets))
-
-        return target_objects
+        with PROFILER.profile(Topics.LOAD_SELECTION, "_".join(t.replace(os.sep, "-") for t in targets)):
+            target_objects, _ = self._load(targets, (),
+                                           selection=selection,
+                                           except_targets=except_targets,
+                                           fetch_subprojects=False,
+                                           use_artifact_config=use_artifact_config,
+                                           load_refs=load_refs)
+
+            return target_objects
 
     # shell()
     #