1"""
2Decorator and functions to profile Salt using cProfile
3"""
4
5
6import datetime
7import logging
8import os
9import pstats
10import subprocess
11
12import salt.utils.files
13import salt.utils.hashutils
14import salt.utils.path
15import salt.utils.stringutils
16
17log = logging.getLogger(__name__)
18
19try:
20    import cProfile
21
22    HAS_CPROFILE = True
23except ImportError:
24    HAS_CPROFILE = False
25
26
27def profile_func(filename=None):
28    """
29    Decorator for adding profiling to a nested function in Salt
30    """
31
32    def proffunc(fun):
33        def profiled_func(*args, **kwargs):
34            logging.info("Profiling function %s", fun.__name__)
35            try:
36                profiler = cProfile.Profile()
37                retval = profiler.runcall(fun, *args, **kwargs)
38                profiler.dump_stats(filename or "{}_func.profile".format(fun.__name__))
39            except OSError:
40                logging.exception("Could not open profile file %s", filename)
41
42            return retval
43
44        return profiled_func
45
46    return proffunc
47
48
49def activate_profile(test=True):
50    pr = None
51    if test:
52        if HAS_CPROFILE:
53            pr = cProfile.Profile()
54            pr.enable()
55        else:
56            log.error("cProfile is not available on your platform")
57    return pr
58
59
60def output_profile(pr, stats_path="/tmp/stats", stop=False, id_=None):
61    if pr is not None and HAS_CPROFILE:
62        try:
63            pr.disable()
64            if not os.path.isdir(stats_path):
65                os.makedirs(stats_path)
66            date = datetime.datetime.now().isoformat()
67            if id_ is None:
68                id_ = salt.utils.hashutils.random_hash(size=32)
69            ficp = os.path.join(stats_path, "{}.{}.pstats".format(id_, date))
70            fico = os.path.join(stats_path, "{}.{}.dot".format(id_, date))
71            ficn = os.path.join(stats_path, "{}.{}.stats".format(id_, date))
72            if not os.path.exists(ficp):
73                pr.dump_stats(ficp)
74                with salt.utils.files.fopen(ficn, "w") as fic:
75                    pstats.Stats(pr, stream=fic).sort_stats("cumulative")
76            log.info("PROFILING: %s generated", ficp)
77            log.info("PROFILING (cumulative): %s generated", ficn)
78            pyprof = salt.utils.path.which("pyprof2calltree")
79            cmd = [pyprof, "-i", ficp, "-o", fico]
80            if pyprof:
81                failed = False
82                try:
83                    pro = subprocess.Popen(
84                        cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE
85                    )
86                except OSError:
87                    failed = True
88                if pro.returncode:
89                    failed = True
90                if failed:
91                    log.error("PROFILING (dot problem")
92                else:
93                    log.info("PROFILING (dot): %s generated", fico)
94                log.trace("pyprof2calltree output:")
95                log.trace(
96                    salt.utils.stringutils.to_str(pro.stdout.read()).strip()
97                    + salt.utils.stringutils.to_str(pro.stderr.read()).strip()
98                )
99            else:
100                log.info("You can run %s for additional stats.", cmd)
101        finally:
102            if not stop:
103                pr.enable()
104    return pr
105