Merge pull request #24103 from makinacorpus/profiling

Profiling

salt/cli/caller.py

@@ -6,6 +6,7 @@
 
 # Import python libs
 from __future__ import absolute_import, print_function
+
 import os
 import sys
 import time
@@ -27,6 +28,8 @@
 from salt.utils import is_windows
 from salt.utils import print_cli
 from salt.utils import kinds
+from salt.utils import activate_profile
+from salt.utils import output_profile
 from salt.cli import daemons
 
 try:
@@ -123,8 +126,16 @@ def run(self):
         '''
         Execute the salt call logic
         '''
+        profiling_enabled = self.opts.get('profiling_enabled', False)
         try:
-            ret = self.call()
+            pr = activate_profile(profiling_enabled)
+            try:
+                ret = self.call()
+            finally:
+                output_profile(pr,
+                               stats_path=self.opts.get('profiling_path',
+                                                        '/tmp/stats'),
+                               stop=True)
             out = ret.get('out', 'nested')
             if self.opts['metadata']:
                 print_ret = ret

salt/cli/run.py

@@ -5,6 +5,8 @@
 import os
 
 from salt.utils import parsers
+from salt.utils import activate_profile
+from salt.utils import output_profile
 from salt.utils.verify import check_user
 from salt.exceptions import SaltClientError
 
@@ -23,6 +25,7 @@ def run(self):
 
         # Setup file logging!
         self.setup_logfile_logger()
+        profiling_enabled = self.options.profiling_enabled
 
         runner = salt.runner.Runner(self.config)
         if self.options.doc:
@@ -33,6 +36,13 @@ def run(self):
         # someone tries to use the runners via the python API
         try:
             if check_user(self.config['user']):
-                runner.run()
+                pr = activate_profile(profiling_enabled)
+                try:
+                    runner.run()
+                finally:
+                    output_profile(
+                        pr,
+                        stats_path=self.options.profiling_path,
+                        stop=True)
         except SaltClientError as exc:
             raise SystemExit(str(exc))

salt/modules/test.py

@@ -9,11 +9,11 @@
 import sys
 import time
 import traceback
-import hashlib
 import random
 
 # Import Salt libs
 import salt
+import salt.utils
 import salt.version
 import salt.loader
 import salt.ext.six as six
@@ -469,18 +469,26 @@ def opts_pkg():
     return ret
 
 
-def rand_str(size=9999999999):
+def rand_str(size=9999999999, hash_type=None):
     '''
     Return a random string
 
+        size
+            size of the string to generate
+        hash_type
+            hash type to use
+
+            .. versionadded:: 2015.5.2
+
     CLI Example:
 
     .. code-block:: bash
 
         salt '*' test.rand_str
     '''
-    hasher = getattr(hashlib, __opts__.get('hash_type', 'md5'))
-    return hasher(str(random.SystemRandom().randint(0, size))).hexdigest()
+    if not hash_type:
+        hash_type = __opts__.get('hash_type', 'md5')
+    return salt.utils.rand_str(hash_type=hash_type, size=size)
 
 
 def exception(message='Test Exception'):

salt/utils/__init__.py

@@ -25,11 +25,13 @@
 import socket
 import stat
 import sys
+import pstats
 import tempfile
 import time
 import types
 import warnings
 import string
+import subprocess
 
 # Import 3rd-party libs
 import salt.ext.six as six
@@ -42,6 +44,13 @@
 from stat import S_IMODE
 # pylint: enable=import-error,redefined-builtin
 
+
+try:
+    import cProfile
+    HAS_CPROFILE = True
+except ImportError:
+    HAS_CPROFILE = False
+
 # Try to load pwd, fallback to getpass if unsuccessful
 # Import 3rd-party libs
 try:
@@ -420,7 +429,6 @@ def profile_func(filename=None):
     '''
     def proffunc(fun):
         def profiled_func(*args, **kwargs):
-            import cProfile
             logging.info('Profiling function {0}'.format(fun.__name__))
             try:
                 profiler = cProfile.Profile()
@@ -437,6 +445,16 @@ def profiled_func(*args, **kwargs):
     return proffunc
 
 
+def rand_str(size=9999999999, hash_type=None):
+    '''
+    Return a random string
+    '''
+    if not hash_type:
+        hash_type = 'md5'
+    hasher = getattr(hashlib, hash_type)
+    return hasher(str(random.SystemRandom().randint(0, size))).hexdigest()
+
+
 def which(exe=None):
     '''
     Python clone of /usr/bin/which
@@ -525,6 +543,62 @@ def which_bin(exes):
     return None
 
 
+def activate_profile(test=True):
+    pr = None
+    if test:
+        if HAS_CPROFILE:
+            pr = cProfile.Profile()
+            pr.enable()
+        else:
+            log.error('cProfile is not available on your platform')
+    return pr
+
+
+def output_profile(pr, stats_path='/tmp/stats', stop=False, id_=None):
+    if pr is not None and HAS_CPROFILE:
+        try:
+            pr.disable()
+            if not os.path.isdir(stats_path):
+                os.makedirs(stats_path)
+            date = datetime.datetime.now().isoformat()
+            if id_ is None:
+                id_ = rand_str(size=32)
+            ficp = os.path.join(stats_path, '{0}.{1}.pstats'.format(id_, date))
+            fico = os.path.join(stats_path, '{0}.{1}.dot'.format(id_, date))
+            ficn = os.path.join(stats_path, '{0}.{1}.stats'.format(id_, date))
+            if not os.path.exists(ficp):
+                pr.dump_stats(ficp)
+                with open(ficn, 'w') as fic:
+                    pstats.Stats(pr, stream=fic).sort_stats('cumulative')
+            log.info('PROFILING: {0} generated'.format(ficp))
+            log.info('PROFILING (cumulative): {0} generated'.format(ficn))
+            pyprof = which('pyprof2calltree')
+            cmd = [pyprof, '-i', ficp, '-o', fico]
+            if pyprof:
+                failed = False
+                try:
+                    pro = subprocess.Popen(
+                        cmd, shell=False,
+                        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+                except OSError:
+                    failed = True
+                if pro.returncode:
+                    failed = True
+                if failed:
+                    log.error('PROFILING (dot problem')
+                else:
+                    log.info('PROFILING (dot): {0} generated'.format(fico))
+                log.trace('pyprof2calltree output:')
+                log.trace(pro.stdout.read().strip() +
+                          pro.stderr.read().strip())
+            else:
+                log.info('You can run {0} for additional stats.'.format(cmd))
+        finally:
+            if not stop:
+                pr.enable()
+    return pr
+
+
 def list_files(directory):
     '''
     Return a list of all files found under directory

salt/utils/parsers.py

@@ -1377,6 +1377,34 @@ def _mixin_after_parsed(self):
             )
 
 
+class ProfilingPMixIn(six.with_metaclass(MixInMeta, object)):
+    _mixin_prio_ = 130
+
+    def _mixin_setup(self):
+        group = self.profiling_group = optparse.OptionGroup(
+            self,
+            'Profiling support',
+            # Include description here as a string
+        )
+
+        group.add_option(
+            '--profiling-path',
+            dest='profiling_path',
+            default='/tmp/stats',
+            help=('Folder that will hold all'
+                  ' Stats generations path (/tmp/stats)')
+        )
+        group.add_option(
+            '--enable-profiling',
+            dest='profiling_enabled',
+            default=False,
+            action='store_true',
+            help=('Enable generating profiling stats'
+                  ' in /tmp/stats (--profiling-path)')
+        )
+        self.add_option_group(group)
+
+
 class CloudCredentialsMixIn(six.with_metaclass(MixInMeta, object)):
     _mixin_prio_ = 30
 
@@ -2053,7 +2081,8 @@ class SaltCallOptionParser(six.with_metaclass(OptionParserMeta,
                                               OutputOptionsMixIn,
                                               HardCrashMixin,
                                               SaltfileMixIn,
-                                              ArgsStdinMixIn)):
+                                              ArgsStdinMixIn,
+                                              ProfilingPMixIn)):
 
     description = ('Salt call is used to execute module functions locally '
                    'on a minion')
@@ -2256,7 +2285,8 @@ class SaltRunOptionParser(six.with_metaclass(OptionParserMeta,
                                              HardCrashMixin,
                                              SaltfileMixIn,
                                              OutputOptionsMixIn,
-                                             ArgsStdinMixIn)):
+                                             ArgsStdinMixIn,
+                                             ProfilingPMixIn)):
 
     default_timeout = 1