Merge pull request #2510 from enkore/issue/2473

Add --debug-profile option

3 files changed, 85 insertions, 3 deletions

diff --git a/docs/usage.rst b/docs/usage.rst
index 33c145350..924559b34 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -523,8 +523,8 @@ Miscellaneous Help
 .. include:: usage/help.rst.inc
 
 
-Debug Commands
---------------
+Debugging Facilities
+--------------------
 
 There is a ``borg debug`` command that has some subcommands which are all
 **not intended for normal use** and **potentially very dangerous** if used incorrectly.
@@ -543,6 +543,20 @@ They exist to improve debugging capabilities without direct system access, e.g.
 in case you ever run into some severe malfunction. Use them only if you know
 what you are doing or if a trusted |project_name| developer tells you what to do.
 
+Borg has a ``--debug-topic TOPIC`` option to enable specific debugging messages. Topics
+are generally not documented.
+
+A ``--debug-profile FILE`` option exists which writes a profile of the main program's
+execution to a file. The format of these files is not directly compatible with the
+Python profiling tools, since these use the "marshal" format, which is not intended
+to be secure (quoting the Python docs: "Never unmarshal data received from an untrusted
+or unauthenticated source.").
+
+The ``borg debug profile-convert`` command can be used to take a Borg profile and convert
+it to a profile file that is compatible with the Python tools.
+
+Additionally, if the filename specified for ``--debug-profile`` ends with ".pyprof" a
+Python compatible profile is generated. This is only intended for local use by developers.
 
 Additional Notes
 ----------------
diff --git a/src/borg/archiver.py b/src/borg/archiver.py
index 80a0dc4ba..7e316d07a 100644
--- a/src/borg/archiver.py
+++ b/src/borg/archiver.py
@@ -1791,6 +1791,13 @@ class Archiver:
                     print("object %s not found [info from chunks cache]." % hex_id)
         return EXIT_SUCCESS
 
+    def do_debug_convert_profile(self, args):
+        """convert Borg profile to Python profile"""
+        import marshal
+        with args.output, args.input:
+            marshal.dump(msgpack.unpack(args.input, use_list=False, encoding='utf-8'), args.output)
+        return EXIT_SUCCESS
+
     @with_repository(lock=False, manifest=False)
     def do_break_lock(self, args, repository):
         """Break the repository lock (e.g. in case it was left by a dead borg."""
@@ -2264,6 +2271,9 @@ class Archiver:
             add_common_option('--consider-part-files', dest='consider_part_files',
                               action='store_true', default=False,
                               help='treat part files like normal files (e.g. to list/extract them)')
+            add_common_option('--debug-profile', dest='debug_profile', default=None, metavar='FILE',
+                              help='Write execution profile in Borg format into FILE. For local use a Python-'
+                                   'compatible file can be generated by suffixing FILE with ".pyprof".')
 
         parser = argparse.ArgumentParser(prog=self.prog, description='Borg - Deduplicated Backups',
                                          add_help=False)
@@ -3668,6 +3678,20 @@ class Archiver:
         subparser.add_argument('ids', metavar='IDs', nargs='+', type=str,
                                help='hex object ID(s) to show refcounts for')
 
+        debug_convert_profile_epilog = process_epilog("""
+        Convert a Borg profile to a Python cProfile compatible profile.
+        """)
+        subparser = debug_parsers.add_parser('convert-profile', parents=[common_parser], add_help=False,
+                                          description=self.do_debug_convert_profile.__doc__,
+                                          epilog=debug_convert_profile_epilog,
+                                          formatter_class=argparse.RawDescriptionHelpFormatter,
+                                          help='convert Borg profile to Python profile (debug)')
+        subparser.set_defaults(func=self.do_debug_convert_profile)
+        subparser.add_argument('input', metavar='INPUT', type=argparse.FileType('rb'),
+                               help='Borg profile')
+        subparser.add_argument('output', metavar='OUTPUT', type=argparse.FileType('wb'),
+                               help='Output file')
+
         benchmark_epilog = process_epilog("These commands do various benchmarks.")
 
         subparser = subparsers.add_parser('benchmark', parents=[mid_common_parser], add_help=False,
@@ -3823,7 +3847,36 @@ class Archiver:
         self.prerun_checks(logger)
         if is_slow_msgpack():
             logger.warning("Using a pure-python msgpack! This will result in lower performance.")
-        return set_ec(func(args))
+        if args.debug_profile:
+            # Import only when needed - avoids a further increase in startup time
+            import cProfile
+            import marshal
+            logger.debug('Writing execution profile to %s', args.debug_profile)
+            # Open the file early, before running the main program, to avoid
+            # a very late crash in case the specified path is invalid.
+            with open(args.debug_profile, 'wb') as fd:
+                profiler = cProfile.Profile()
+                variables = dict(locals())
+                profiler.enable()
+                try:
+                    return set_ec(func(args))
+                finally:
+                    profiler.disable()
+                    profiler.snapshot_stats()
+                    if args.debug_profile.endswith('.pyprof'):
+                        marshal.dump(profiler.stats, fd)
+                    else:
+                        # We use msgpack here instead of the marshal module used by cProfile itself,
+                        # because the latter is insecure. Since these files may be shared over the
+                        # internet we don't want a format that is impossible to interpret outside
+                        # an insecure implementation.
+                        # See scripts/msgpack2marshal.py for a small script that turns a msgpack file
+                        # into a marshal file that can be read by e.g. pyprof2calltree.
+                        # For local use it's unnecessary hassle, though, that's why .pyprof makes
+                        # it compatible (see above).
+                        msgpack.pack(profiler.stats, fd, use_bin_type=True)
+        else:
+            return set_ec(func(args))
 
 
 def sig_info_handler(sig_no, stack):  # pragma: no cover
diff --git a/src/borg/testsuite/archiver.py b/src/borg/testsuite/archiver.py
index dd3285fd3..dfe6e75c5 100644
--- a/src/borg/testsuite/archiver.py
+++ b/src/borg/testsuite/archiver.py
@@ -3,6 +3,7 @@ import errno
 import json
 import logging
 import os
+import pstats
 import random
 import shutil
 import socket
@@ -1688,6 +1689,20 @@ class ArchiverTestCase(ArchiverTestCaseBase):
         assert log_message['levelname'] == 'DEBUG'  # there should only be DEBUG messages
         assert isinstance(log_message['message'], str)
 
+    def test_debug_profile(self):
+        self.create_test_files()
+        self.cmd('init', '--encryption=repokey', self.repository_location)
+        self.cmd('create', self.repository_location + '::test', 'input', '--debug-profile=create.prof')
+        self.cmd('debug', 'convert-profile', 'create.prof', 'create.pyprof')
+        stats = pstats.Stats('create.pyprof')
+        stats.strip_dirs()
+        stats.sort_stats('cumtime')
+
+        self.cmd('create', self.repository_location + '::test2', 'input', '--debug-profile=create.pyprof')
+        stats = pstats.Stats('create.pyprof')  # Only do this on trusted data!
+        stats.strip_dirs()
+        stats.sort_stats('cumtime')
+
     def test_common_options(self):
         self.create_test_files()
         self.cmd('init', '--encryption=repokey', self.repository_location)