Load static symbol information lazily with some clean-ups.

tools/deep_memory_profiler/dmprof

 #!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """The deep heap profiler script for Chrome."""
 
 from datetime import datetime
 import json
 import optparse
 import os
 import re
 import shutil
 import subprocess
 import sys
 import tempfile
 
 FIND_RUNTIME_SYMBOLS_PATH = os.path.join(
     os.path.dirname(os.path.abspath(__file__)),
     os.pardir,
     'find_runtime_symbols')
 sys.path.append(FIND_RUNTIME_SYMBOLS_PATH)
 
+from find_runtime_symbols import find_runtime_symbols_list
 from prepare_symbol_info import prepare_symbol_info
-from find_runtime_symbols import find_runtime_symbols_list
+from static_symbols import StaticSymbols
 
 BUCKET_ID = 5
 VIRTUAL = 0
 COMMITTED = 1
 ALLOC_COUNT = 2
 FREE_COUNT = 3
 NULL_REGEX = re.compile('')
 
 POLICIES_JSON_PATH = os.path.join(
     os.path.dirname(os.path.abspath(__file__)),
@@ skipping 63 matching lines @@
     return "invalid heap profile dump: %s" % repr(self.value)
 
 
 class ObsoleteDumpVersionException(ParsingException):
   def __init__(self, value):
     self.value = value
   def __str__(self):
     return "obsolete heap profile dump version: %s" % repr(self.value)
 
 
+class DelayedStaticSymbols(object):
+  """Represents static symbol information loaded lazily."""
+
+  def __init__(self, prefix, keep=False):
+    self.maps_path = prefix + '.maps'
+    self.keep = keep
+    if keep:
+      self.prepared_data_dir = prefix + '.pre'
+    self.loaded_static_symbols = None
+
+  def get(self):
+    if not self.loaded_static_symbols:
+      if not self.keep:
+        self.prepared_data_dir = tempfile.mkdtemp()
+      try:
+        prepare_symbol_info(self.maps_path, self.prepared_data_dir)
+        self.loaded_static_symbols = StaticSymbols.load(self.prepared_data_dir)
+      finally:
+        if not self.keep:
+          shutil.rmtree(self.prepared_data_dir)
+    return self.loaded_static_symbols
+
+
 class Rule(object):
   """Represents one matching rule in a policy file."""
 
   def __init__(self, name, mmap, pattern):
     self.name = name
     self.mmap = mmap
     self.condition = re.compile(pattern + r'\Z')
 
 
 class Policy(object):
@@ skipping 66 matching lines @@
         symbols: A dict mapping runtime addresses to symbol names.
     """
     for line in self.stacktrace_lines:
       words = line.split()
       bucket = buckets.get(int(words[BUCKET_ID]))
       if not bucket:
         continue
       for i in range(0, BUCKET_ID - 1):
         sys.stdout.write(words[i] + ' ')
       for address in bucket.stacktrace:
-        sys.stdout.write((symbols.get(address) or address) + ' ')
+        sys.stdout.write((symbols.get(address) or ('0x%016x' % address)) + ' ')
       sys.stdout.write('\n')
 
   @staticmethod
   def accumulate_size_for_pprof(stacktrace_lines, rule_list, buckets,
                                 component_name, symbols):
     """Accumulates size of committed chunks and the number of allocated chunks.
 
     Args:
         stacktrace_lines: A list of strings which are valid as stacktraces.
         rule_list: A list of Rule objects.
@@ skipping 39 matching lines @@
           (component_name and
            component_name != get_component(rule_list, bucket, symbols))):
         continue
 
       sys.stdout.write('%6d: %8s [%6d: %8s] @' % (
           int(words[ALLOC_COUNT]) - int(words[FREE_COUNT]),
           words[COMMITTED],
           int(words[ALLOC_COUNT]) - int(words[FREE_COUNT]),
           words[COMMITTED]))
       for address in bucket.stacktrace:
-        sys.stdout.write(' ' + address)
+        sys.stdout.write(' 0x%016x' % address)
       sys.stdout.write('\n')
 
   def print_for_pprof(
       self, rule_list, buckets, maps_lines, component_name, symbols):
     """Converts the heap profile dump so it can be processed by pprof.
 
     Args:
         rule_list: A list of Rule objects.
         buckets: A dict mapping bucket ids to Bucket objects.
         maps_lines: A list of strings containing /proc/.../maps.
@@ skipping 303 matching lines @@
     sorted_sizes_list = sorted(
         sizes.iteritems(), key=(lambda x: x[1]), reverse=True)
     total = 0
     for size_pair in sorted_sizes_list:
       sys.stdout.write('%10d %s\n' % (size_pair[1], size_pair[0]))
       total += size_pair[1]
     sys.stderr.write('total: %d\n' % (total))
 
 
 def update_symbols(
-    symbol_path, maps_path, appeared_addresses, symbols):
+    symbol_path, delayed_static_symbols, appeared_addresses, symbols):
   """Updates address/symbol mapping on memory and in a .symbol cache file.
 
   It reads cached address/symbol mapping from a .symbol file if it exists.
   Then, it resolves unresolved addresses from a Chrome binary with pprof.
   Both mappings on memory and in a .symbol cache file are updated.
 
   Symbol files are formatted as follows:
     <Address> <Symbol>
     <Address> <Symbol>
     <Address> <Symbol>
     ...
 
   Args:
       symbol_path: A string representing a path for a .symbol file.
-      maps_path: A string of the path of /proc/.../maps.
+      delayed_static_symbols: A DelayedStaticSymbols object.
       appeared_addresses: A list of known addresses.
       symbols: A dict mapping runtime addresses to symbol names.
   """
   with open(symbol_path, mode='a+') as symbol_f:
     symbol_lines = symbol_f.readlines()
     if symbol_lines:
       for line in symbol_lines:
         items = line.split(None, 1)
         if len(items) == 1:
           items.append('??')
-        symbols[items[0]] = items[1].rstrip()
+        symbols[int(items[0], 16)] = items[1].rstrip()
     if symbols:
       sys.stderr.write('  Found %d symbols in cache.\n' % len(symbols))
     else:
       sys.stderr.write('  No symbols found in cache.\n')
 
     unresolved_addresses = sorted(
         a for a in appeared_addresses if a not in symbols)
 
     if not unresolved_addresses:
       sys.stderr.write('  No need to resolve any more addresses.\n')
     else:
-      sys.stderr.write('  %d addresses are unresolved.\n' %
+      sys.stderr.write('  %d addresses unresolved.\n' %
                        len(unresolved_addresses))
-      prepared_data_dir = tempfile.mkdtemp()
-      try:
-        prepare_symbol_info(maps_path, prepared_data_dir)
+      static_symbols = delayed_static_symbols.get()
+      symbol_list = find_runtime_symbols_list(
+          static_symbols, unresolved_addresses)
 
-        symbol_list = find_runtime_symbols_list(
-            prepared_data_dir, unresolved_addresses)
+      for address, symbol in zip(unresolved_addresses, symbol_list):
+        if not symbol:
+          symbol = '??'
+        stripped_symbol = symbol.strip()
+        symbols[address] = stripped_symbol
+        symbol_f.write('%x %s\n' % (address, stripped_symbol))
 
-        for address, symbol in zip(unresolved_addresses, symbol_list):
-          if not symbol:
-            symbol = '??'
-          stripped_symbol = symbol.strip()
-          symbols[address] = stripped_symbol
-          symbol_f.write('%s %s\n' % (address, stripped_symbol))
-      finally:
-        shutil.rmtree(prepared_data_dir)
+    sys.stderr.write('  All symbols resolved.\n')
 
 
 def parse_policy(policy_path):
   """Parses policy file.
 
   A policy file contains component's names and their
   stacktrace pattern written in regular expression.
   Those patterns are matched against each symbols of
   each stacktraces in the order written in the policy file
 
@@ skipping 53 matching lines @@
     buckets_path = '%s.%04d.buckets' % (prefix, n)
     if not os.path.exists(buckets_path):
       if n > 10:
         break
       n += 1
       continue
     sys.stderr.write('  %s\n' % buckets_path)
     with open(buckets_path, 'r') as buckets_f:
       for line in buckets_f:
         words = line.split()
-        buckets[int(words[0])] = Bucket(words[2:], words[1] == 'mmap')
+        stacktrace = [int(address, 16) for address in words[2:]]
+        buckets[int(words[0])] = Bucket(stacktrace, words[1] == 'mmap')
     n += 1
 
   return buckets
 
 
 def determine_dump_path_list(dump_path, prefix):
   dump_path_list = [dump_path]
 
   # search for the sequence of files
   n = int(dump_path[len(dump_path) - 9 : len(dump_path) - 5])
@@ skipping 36 matching lines @@
   sys.stderr.write('Loading heap dump files.\n')
   appeared_addresses = set()
   dumps = []
   for path in dump_path_list:
     sys.stderr.write('  %s' % path)
     dumps.append(load_single_dump(path, buckets, appeared_addresses))
     sys.stderr.write('\n')
   return dumps, appeared_addresses
 
 
-def load_and_update_symbol_cache(prefix, appeared_addresses):
-  maps_path = prefix + '.maps'
+def load_and_update_symbol_cache(
+    prefix, appeared_addresses, delayed_static_symbols):
   symbol_path = prefix + '.symbols'
   sys.stderr.write('Loading and updating symbol cache: "%s".\n' % symbol_path)
   symbols = {}
-  update_symbols(symbol_path, maps_path, appeared_addresses, symbols)
+  update_symbols(
+      symbol_path, delayed_static_symbols, appeared_addresses, symbols)
   return symbols
 
 
 def load_default_policies():
   with open(POLICIES_JSON_PATH, mode='r') as policies_f:
     default_policies = json.load(policies_f)
   return default_policies
 
 
 def load_policy(policies_dict, policy_label):
@@ skipping 21 matching lines @@
     for specified_policy in policy_labels:
       if specified_policy in default_policies:
         specified_policies[specified_policy] = (
             default_policies[specified_policy])
     policies = load_policies_dict(specified_policies)
   else:
     policies = load_policies_dict(default_policies)
   return policies
 
 
+def load_basic_files_with_multiple_dumps(dump_path, keep):
+  prefix = find_prefix(dump_path)
+  buckets = load_buckets(prefix)
+  dumps, appeared_addresses = load_dumps(
+      determine_dump_path_list(dump_path, prefix), buckets)
+  delayed_static_symbols = DelayedStaticSymbols(prefix, keep)
+  symbols = load_and_update_symbol_cache(
+      prefix, appeared_addresses, delayed_static_symbols)
+  return buckets, dumps, appeared_addresses, delayed_static_symbols, symbols
+
+
+def load_basic_files_with_single_dump(dump_path, keep):
+  prefix = find_prefix(dump_path)
+  buckets = load_buckets(prefix)
+  dump, appeared_addresses = load_dump(dump_path, buckets)
+  delayed_static_symbols = DelayedStaticSymbols(prefix, keep)
+  symbols = load_and_update_symbol_cache(
+      prefix, appeared_addresses, delayed_static_symbols)
+  return buckets, dump, appeared_addresses, delayed_static_symbols, symbols
+
+
 def do_stacktrace(sys_argv):
-  parser = optparse.OptionParser(usage='Usage: %prog stacktrace <dump>')
+  parser = optparse.OptionParser(
+      'Usage: %prog stacktrace [--keep] <dump>')
+  parser.add_option('--keep', dest='keep', action='store_true')
   options, args = parser.parse_args(sys_argv)
 
   if len(args) != 2:
     parser.error('needs 1 argument.')
     return 1
 
   dump_path = args[1]
 
-  prefix = find_prefix(dump_path)
-  buckets = load_buckets(prefix)
-  dump, appeared_addresses = load_dump(dump_path, buckets)
-  symbols = load_and_update_symbol_cache(prefix, appeared_addresses)
+  buckets, dump, appeared_addresses, delayed_static_symbols, symbols = (
+      load_basic_files_with_single_dump(dump_path, options.keep))
 
   dump.print_stacktrace(buckets, symbols)
 
   return 0
 
 
 def do_csv(sys_argv):
-  parser = optparse.OptionParser('Usage: %prog csv [-p POLICY] <first-dump>')
+  parser = optparse.OptionParser(
+      'Usage: %prog csv [-p POLICY] [--keep] <first-dump>')
   parser.add_option('-p', '--policy', type='string', dest='policy',
                     help='profile with POLICY', metavar='POLICY')
+  parser.add_option('--keep', dest='keep', action='store_true')
   options, args = parser.parse_args(sys_argv)
 
   if len(args) != 2:
     parser.error('needs 1 argument.')
     return 1
 
   dump_path = args[1]
 
-  prefix = find_prefix(dump_path)
-  buckets = load_buckets(prefix)
-  dumps, appeared_addresses = load_dumps(
-      determine_dump_path_list(dump_path, prefix), buckets)
-  symbols = load_and_update_symbol_cache(prefix, appeared_addresses)
+  buckets, dumps, appeared_addresses, delayed_static_symbols, symbols = (
+      load_basic_files_with_multiple_dumps(dump_path, options.keep))
   policies = load_policies(options.policy)
 
   max_components = 0
   for policy in policies:
     max_components = max(max_components, len(policies[policy].components))
 
   for policy in sorted(policies):
     rule_list = policies[policy].rules
     components = policies[policy].components
 
@@ skipping 14 matching lines @@
       sys.stdout.write('%s%s\n' % (
             ','.join(s), ',' * (max_components - len(components))))
 
     for bucket in buckets.itervalues():
       bucket.clear_component_cache()
 
   return 0
 
 
 def do_json(sys_argv):
-  parser = optparse.OptionParser('Usage: %prog json [-p POLICY] <first-dump>')
+  parser = optparse.OptionParser(
+      'Usage: %prog json [-p POLICY] [--keep] <first-dump>')
   parser.add_option('-p', '--policy', type='string', dest='policy',
                     help='profile with POLICY', metavar='POLICY')
+  parser.add_option('--keep', dest='keep', action='store_true')
   options, args = parser.parse_args(sys_argv)
 
   if len(args) != 2:
     parser.error('needs 1 argument.')
     return 1
 
   dump_path = args[1]
 
-  prefix = find_prefix(dump_path)
-  buckets = load_buckets(prefix)
-  dumps, appeared_addresses = load_dumps(
-      determine_dump_path_list(dump_path, prefix), buckets)
-  symbols = load_and_update_symbol_cache(prefix, appeared_addresses)
+  buckets, dumps, appeared_addresses, delayed_static_symbols, symbols = (
+      load_basic_files_with_multiple_dumps(dump_path, options.keep))
   policies = load_policies(options.policy)
 
   json_base = {
     'version': 'JSON_DEEP_2',
     'policies': {},
   }
 
   for policy in sorted(policies):
     rule_list = policies[policy].rules
     components = policies[policy].components
@@ skipping 13 matching lines @@
 
     for bucket in buckets.itervalues():
       bucket.clear_component_cache()
 
   json.dump(json_base, sys.stdout, indent=2, sort_keys=True)
 
   return 0
 
 
 def do_list(sys_argv):
-  parser = optparse.OptionParser('Usage: %prog [-p POLICY] list <first-dump>')
+  parser = optparse.OptionParser(
+      'Usage: %prog [-p POLICY] [--keep] list <first-dump>')
   parser.add_option('-p', '--policy', type='string', dest='policy',
                     help='profile with POLICY', metavar='POLICY')
+  parser.add_option('--keep', dest='keep', action='store_true')
   options, args = parser.parse_args(sys_argv)
 
   if len(args) != 2:
     parser.error('needs 1 argument.')
     return 1
 
   dump_path = args[1]
 
-  prefix = find_prefix(dump_path)
-  buckets = load_buckets(prefix)
-  dumps, appeared_addresses = load_dumps(
-      determine_dump_path_list(dump_path, prefix), buckets)
-  symbols = load_and_update_symbol_cache(prefix, appeared_addresses)
+  buckets, dumps, appeared_addresses, delayed_static_symbols, symbols = (
+      load_basic_files_with_multiple_dumps(dump_path, options.keep))
   policies = load_policies(options.policy)
 
   for policy in sorted(policies):
     rule_list = policies[policy].rules
     components = policies[policy].components
 
     component_sizes = dumps[0].apply_policy(
         rule_list, buckets, dumps[0].dump_time, components, symbols)
     sys.stdout.write('%s:\n' % policy)
     for c in components:
       if c in ['hour', 'minute', 'second']:
         sys.stdout.write('%30s %10.3f\n' % (c, component_sizes[c]))
       else:
         sys.stdout.write('%30s %10.3f\n' % (
             c, component_sizes[c] / 1024.0 / 1024.0))
 
     for bucket in buckets.itervalues():
       bucket.clear_component_cache()
 
   return 0
 
 
 def do_expand(sys_argv):
   parser = optparse.OptionParser(
-      'Usage: %prog expand <dump> <policy> <component> <depth>')
+      'Usage: %prog expand [--keep] <dump> <policy> <component> <depth>')
+  parser.add_option('--keep', dest='keep', action='store_true')
   options, args = parser.parse_args(sys_argv)
 
   if len(args) != 5:
     parser.error('needs 4 arguments.')
     return 1
 
   dump_path = args[1]
   target_policy = args[2]
   component_name = args[3]
   depth = args[4]
 
-  prefix = find_prefix(dump_path)
-  buckets = load_buckets(prefix)
-  dump, appeared_addresses = load_dump(dump_path, buckets)
-  symbols = load_and_update_symbol_cache(prefix, appeared_addresses)
+  buckets, dump, appeared_addresses, delayed_static_symbols, symbols = (
+      load_basic_files_with_single_dump(dump_path, options.keep))
   policies = load_policies(target_policy)
 
   rule_list = policies[target_policy].rules
 
   dump.expand(rule_list, buckets, component_name, int(depth), symbols)
 
   return 0
 
 
 def do_pprof(sys_argv):
   parser = optparse.OptionParser(
-      'Usage: %prog pprof [-c COMPONENT] <dump> <policy>')
+      'Usage: %prog pprof [-c COMPONENT] [--keep] <dump> <policy>')
   parser.add_option('-c', '--component', type='string', dest='component',
                     help='restrict to COMPONENT', metavar='COMPONENT')
+  parser.add_option('--keep', dest='keep', action='store_true')
   options, args = parser.parse_args(sys_argv)
 
   if len(args) != 3:
     parser.error('needs 2 arguments.')
     return 1
 
   dump_path = args[1]
   target_policy = args[2]
   component = options.component
 
-  prefix = find_prefix(dump_path)
-  buckets = load_buckets(prefix)
-  dump, appeared_addresses = load_dump(dump_path, buckets)
-  symbols = load_and_update_symbol_cache(prefix, appeared_addresses)
+  buckets, dump, appeared_addresses, delayed_static_symbols, symbols = (
+      load_basic_files_with_single_dump(dump_path, options.keep))
   policies = load_policies(target_policy)
 
   rule_list = policies[target_policy].rules
 
-  with open(prefix + '.maps', 'r') as maps_f:
+  with open(find_prefix(dump_path) + '.maps', 'r') as maps_f:
     maps_lines = maps_f.readlines()
   dump.print_for_pprof(rule_list, buckets, maps_lines, component, symbols)
 
   return 0
 
 
 def main():
   COMMANDS = {
     'csv': do_csv,
     'expand': do_expand,
@@ skipping 18 matching lines @@
 
 Commands:
    csv          Classify memory usage in CSV
    expand       Show all stacktraces contained in the specified component
    json         Classify memory usage in JSON
    list         Classify memory usage in simple listing format
    pprof        Format the profile dump so that it can be processed by pprof
    stacktrace   Convert runtime addresses to symbol names
 
 Quick Reference:
-   dmprof csv [-p POLICY] <first-dump>
-   dmprof expand <dump> <policy> <component> <depth>
-   dmprof json [-p POLICY] <first-dump>
-   dmprof list [-p POLICY] <first-dump>
-   dmprof pprof [-c COMPONENT] <dump> <policy>
-   dmprof stacktrace <dump>
+   dmprof csv [-p POLICY] [--keep] <first-dump>
+   dmprof expand [--keep] <dump> <policy> <component> <depth>
+   dmprof json [-p POLICY] [--keep] <first-dump>
+   dmprof list [-p POLICY] [--keep] <first-dump>
+   dmprof pprof [-c COMPONENT] [--keep] <dump> <policy>
+   dmprof stacktrace [--keep] <dump>
 """ % (sys.argv[0]))
     sys.exit(1)
   action = sys.argv.pop(1)
 
   return COMMANDS[action](sys.argv)
 
 
 if __name__ == '__main__':
   sys.exit(main())