tools/multi_process_rss.py

   1 #!/usr/bin/env python
   2 # Copyright 2013 The Chromium Authors
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 # Counts a resident set size (RSS) of multiple processes without double-counts.
   7 # If they share the same page frame, the page frame is counted only once.
   8 #
   9 # Usage:
  10 # ./multi-process-rss.py <pid>|<pid>r [...]
  11 #
  12 # If <pid> has 'r' at the end, all descendants of the process are accounted.
  13 #
  14 # Example:
  15 # ./multi-process-rss.py 12345 23456r
  16 #
  17 # The command line above counts the RSS of 1) process 12345, 2) process 23456
  18 # and 3) all descendant processes of process 23456.
  19
  20 from __future__ import print_function
  21
  22 import collections
  23 import logging
  24 import os
  25 import psutil
  26 import sys
  27
  28
  29 if sys.platform.startswith('linux'):
  30   _TOOLS_PATH = os.path.dirname(os.path.abspath(__file__))
  31   _TOOLS_LINUX_PATH = os.path.join(_TOOLS_PATH, 'linux')
  32   sys.path.append(_TOOLS_LINUX_PATH)
  33   import procfs  # pylint: disable=F0401
  34
  35
  36 class _NullHandler(logging.Handler):
  37   def emit(self, record):
  38     pass
  39
  40
  41 _LOGGER = logging.getLogger('multi-process-rss')
  42 _LOGGER.addHandler(_NullHandler())
  43
  44
  45 def _recursive_get_children(pid):
  46   try:
  47     children = psutil.Process(pid).get_children()
  48   except psutil.error.NoSuchProcess:
  49     return []
  50   descendant = []
  51   for child in children:
  52     descendant.append(child.pid)
  53     descendant.extend(_recursive_get_children(child.pid))
  54   return descendant
  55
  56
  57 def list_pids(argv):
  58   pids = []
  59   for arg in argv[1:]:
  60     try:
  61       if arg.endswith('r'):
  62         recursive = True
  63         pid = int(arg[:-1])
  64       else:
  65         recursive = False
  66         pid = int(arg)
  67     except ValueError:
  68       raise SyntaxError("%s is not an integer." % arg)
  69     else:
  70       pids.append(pid)
  71     if recursive:
  72       children = _recursive_get_children(pid)
  73       pids.extend(children)
  74
  75   pids = sorted(set(pids), key=pids.index)  # uniq: maybe slow, but simple.
  76
  77   return pids
  78
  79
  80 def count_pageframes(pids):
  81   pageframes = collections.defaultdict(int)
  82   pagemap_dct = {}
  83   for pid in pids:
  84     maps = procfs.ProcMaps.load(pid)
  85     if not maps:
  86       _LOGGER.warning('/proc/%d/maps not found.' % pid)
  87       continue
  88     pagemap = procfs.ProcPagemap.load(pid, maps)
  89     if not pagemap:
  90       _LOGGER.warning('/proc/%d/pagemap not found.' % pid)
  91       continue
  92     pagemap_dct[pid] = pagemap
  93
  94   for pid, pagemap in pagemap_dct.iteritems():
  95     for vma in pagemap.vma_internals.itervalues():
  96       for pageframe, number in vma.pageframes.iteritems():
  97         pageframes[pageframe] += number
  98
  99   return pageframes
 100
 101
 102 def count_statm(pids):
 103   resident = 0
 104   shared = 0
 105   private = 0
 106
 107   for pid in pids:
 108     statm = procfs.ProcStatm.load(pid)
 109     if not statm:
 110       _LOGGER.warning('/proc/%d/statm not found.' % pid)
 111       continue
 112     resident += statm.resident
 113     shared += statm.share
 114     private += (statm.resident - statm.share)
 115
 116   return (resident, shared, private)
 117
 118
 119 def main(argv):
 120   logging_handler = logging.StreamHandler()
 121   logging_handler.setLevel(logging.WARNING)
 122   logging_handler.setFormatter(logging.Formatter(
 123       '%(asctime)s:%(name)s:%(levelname)s:%(message)s'))
 124
 125   _LOGGER.setLevel(logging.WARNING)
 126   _LOGGER.addHandler(logging_handler)
 127
 128   if sys.platform.startswith('linux'):
 129     logging.getLogger('procfs').setLevel(logging.WARNING)
 130     logging.getLogger('procfs').addHandler(logging_handler)
 131     pids = list_pids(argv)
 132     pageframes = count_pageframes(pids)
 133   else:
 134     _LOGGER.error('%s is not supported.' % sys.platform)
 135     return 1
 136
 137   # TODO(dmikurube): Classify this total RSS.
 138   print(len(pageframes) * 4096)
 139
 140   return 0
 141
 142
 143 if __name__ == '__main__':
 144   sys.exit(main(sys.argv))