summaryrefslogtreecommitdiff
path: root/buildstream/_profile.py
blob: f29e070c43633eeb3727c6ea700e9d62b2368177 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
#
#  Copyright (C) 2017 Codethink Limited
#  Copyright (C) 2019 Bloomberg Finance LP
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public
#  License as published by the Free Software Foundation; either
#  version 2 of the License, or (at your option) any later version.
#
#  This library is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	 See the GNU
#  Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with this library. If not, see <http://www.gnu.org/licenses/>.
#
#  Authors:
#        Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
#        James Ennis <james.ennis@codethink.co.uk>

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.
#
# Multiple topics can be set with the ':' separator.
#
# E.g.:
#
#   BST_PROFILE=circ-dep-check:sort-deps bst <command> <args>
#
# The special 'all' value will enable all profiles.
class Topics():
    CIRCULAR_CHECK = 'circ-dep-check'
    SORT_DEPENDENCIES = 'sort-deps'
    LOAD_LOADER = 'load-loader'
    LOAD_CONTEXT = 'load-context'
    LOAD_PROJECT = 'load-project'
    LOAD_PIPELINE = 'load-pipeline'
    LOAD_SELECTION = 'load-selection'
    SCHEDULER = 'scheduler'
    SHOW = 'show'
    ARTIFACT_RECEIVE = 'artifact-receive'
    ALL = 'all'


class Profile():
    def __init__(self, topic, key, message):
        self.message = message
        self.key = topic + '-' + key
        self.start = time.time()
        self.profiler = cProfile.Profile()
        self.profiler.enable()

    def end(self):
        self.profiler.disable()

        dt = datetime.datetime.fromtimestamp(self.start)
        timestamp = dt.strftime('%Y%m%dT%H%M%S')

        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_)

        self.__write_binary(filename + '.cprofile')

    ########################################
    #            Private Methods           #
    ########################################

    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()

    def __write_binary(self, filename):
        self.profiler.dump_stats(filename)


# 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

    # 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


# 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