1# Copyright 2019, David Wilson
2#
3# Redistribution and use in source and binary forms, with or without
4# modification, are permitted provided that the following conditions are met:
5#
6# 1. Redistributions of source code must retain the above copyright notice,
7# this list of conditions and the following disclaimer.
8#
9# 2. Redistributions in binary form must reproduce the above copyright notice,
10# this list of conditions and the following disclaimer in the documentation
11# and/or other materials provided with the distribution.
12#
13# 3. Neither the name of the copyright holder nor the names of its contributors
14# may be used to endorse or promote products derived from this software without
15# specific prior written permission.
16#
17# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
18# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
21# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
22# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
23# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
24# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
25# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
26# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
27# POSSIBILITY OF SUCH DAMAGE.
28
29# !mitogen: minify_safe
30
31"""
32mitogen.profiler
33    Record and report cProfile statistics from a run. Creates one aggregated
34    output file, one aggregate containing only workers, and one for the
35    top-level process.
36
37Usage:
38    mitogen.profiler record <dest_path> <tool> [args ..]
39    mitogen.profiler report <dest_path> [sort_mode]
40    mitogen.profiler stat <sort_mode> <tool> [args ..]
41
42Mode:
43    record: Record a trace.
44    report: Report on a previously recorded trace.
45    stat: Record and report in a single step.
46
47Where:
48    dest_path: Filesystem prefix to write .pstats files to.
49    sort_mode: Sorting mode; defaults to "cumulative". See:
50        https://docs.python.org/2/library/profile.html#pstats.Stats.sort_stats
51
52Example:
53    mitogen.profiler record /tmp/mypatch ansible-playbook foo.yml
54    mitogen.profiler dump /tmp/mypatch-worker.pstats
55"""
56
57from __future__ import print_function
58import os
59import pstats
60import shutil
61import subprocess
62import sys
63import tempfile
64import time
65
66
67def try_merge(stats, path):
68    try:
69        stats.add(path)
70        return True
71    except Exception as e:
72        print('%s failed. Will retry. %s' % (path, e))
73        return False
74
75
76def merge_stats(outpath, inpaths):
77    first, rest = inpaths[0], inpaths[1:]
78    for x in range(1):
79        try:
80            stats = pstats.Stats(first)
81        except EOFError:
82            time.sleep(0.2)
83            continue
84
85        print("Writing %r..." % (outpath,))
86        for path in rest:
87            #print("Merging %r into %r.." % (os.path.basename(path), outpath))
88            for x in range(5):
89                if try_merge(stats, path):
90                    break
91                time.sleep(0.2)
92
93    stats.dump_stats(outpath)
94
95
96def generate_stats(outpath, tmpdir):
97    print('Generating stats..')
98    all_paths = []
99    paths_by_ident = {}
100
101    for name in os.listdir(tmpdir):
102        if name.endswith('-dump.pstats'):
103            ident, _, pid = name.partition('-')
104            path = os.path.join(tmpdir, name)
105            all_paths.append(path)
106            paths_by_ident.setdefault(ident, []).append(path)
107
108    merge_stats('%s-all.pstat' % (outpath,), all_paths)
109    for ident, paths in paths_by_ident.items():
110        merge_stats('%s-%s.pstat' % (outpath, ident), paths)
111
112
113def do_record(tmpdir, path, *args):
114    env = os.environ.copy()
115    fmt = '%(identity)s-%(pid)s.%(now)s-dump.%(ext)s'
116    env['MITOGEN_PROFILING'] = '1'
117    env['MITOGEN_PROFILE_FMT'] = os.path.join(tmpdir, fmt)
118    rc = subprocess.call(args, env=env)
119    generate_stats(path, tmpdir)
120    return rc
121
122
123def do_report(tmpdir, path, sort='cumulative'):
124    stats = pstats.Stats(path).sort_stats(sort)
125    stats.print_stats(100)
126
127
128def do_stat(tmpdir, sort, *args):
129    valid_sorts = pstats.Stats.sort_arg_dict_default
130    if sort not in valid_sorts:
131        sys.stderr.write('Invalid sort %r, must be one of %s\n' %
132                         (sort, ', '.join(sorted(valid_sorts))))
133        sys.exit(1)
134
135    outfile = os.path.join(tmpdir, 'combined')
136    do_record(tmpdir, outfile, *args)
137    aggs = ('app.main', 'mitogen.broker', 'mitogen.child_main',
138            'mitogen.service.pool', 'Strategy', 'WorkerProcess',
139            'all')
140    for agg in aggs:
141        path = '%s-%s.pstat' % (outfile, agg)
142        if os.path.exists(path):
143            print()
144            print()
145            print('------ Aggregation %r ------' % (agg,))
146            print()
147            do_report(tmpdir, path, sort)
148            print()
149
150
151def main():
152    if len(sys.argv) < 2 or sys.argv[1] not in ('record', 'report', 'stat'):
153        sys.stderr.write(__doc__.lstrip())
154        sys.exit(1)
155
156    func = globals()['do_' + sys.argv[1]]
157    tmpdir = tempfile.mkdtemp(prefix='mitogen.profiler')
158    try:
159        sys.exit(func(tmpdir, *sys.argv[2:]) or 0)
160    finally:
161        shutil.rmtree(tmpdir)
162
163if __name__ == '__main__':
164    main()
165