Error handling and refactoring in tools/linux/procfs.py, and rename multi-process-rss.py.

tools/multi-process-rss.py

-#!/usr/bin/env python
-# Copyright 2013 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.
-
-# Counts a resident set size (RSS) of multiple processes without double-counts.
-# If they share the same page frame, the page frame is counted only once.
-#
-# Usage:
-# ./multi-process-rss.py <pid>|<pid>r [...]
-#
-# If <pid> has 'r' at the end, all descendants of the process are accounted.
-#
-# Example:
-# ./multi-process-rss.py 12345 23456r
-#
-# The command line above counts the RSS of 1) process 12345, 2) process 23456
-# and 3) all descendant processes of process 23456.
-
-
-import collections
-import logging
-import os
-import psutil
-import sys
-
-
-if sys.platform.startswith('linux'):
-  _TOOLS_PATH = os.path.dirname(os.path.abspath(__file__))
-  _TOOLS_LINUX_PATH = os.path.join(_TOOLS_PATH, 'linux')
-  sys.path.append(_TOOLS_LINUX_PATH)
-  import procfs  # pylint: disable=F0401
-
-
-class _NullHandler(logging.Handler):
-  def emit(self, record):
-    pass
-
-
-_LOGGER = logging.getLogger('multi-process-rss')
-_LOGGER.addHandler(_NullHandler())
-
-
-def _recursive_get_children(pid):
-  try:
-    children = psutil.Process(pid).get_children()
-  except psutil.error.NoSuchProcess:
-    return []
-  descendant = []
-  for child in children:
-    descendant.append(child.pid)
-    descendant.extend(_recursive_get_children(child.pid))
-  return descendant
-
-
-def list_pids(argv):
-  pids = []
-  for arg in argv[1:]:
-    try:
-      if arg.endswith('r'):
-        recursive = True
-        pid = int(arg[:-1])
-      else:
-        recursive = False
-        pid = int(arg)
-    except ValueError:
-      raise SyntaxError("%s is not an integer." % arg)
-    else:
-      pids.append(pid)
-    if recursive:
-      children = _recursive_get_children(pid)
-      pids.extend(children)
-
-  pids = sorted(set(pids), key=pids.index)  # uniq: maybe slow, but simple.
-
-  return pids
-
-
-def count_pageframes(pids):
-  pageframes = collections.defaultdict(int)
-  pagemap_dct = {}
-  for pid in pids:
-    maps = procfs.ProcMaps.load(pid)
-    if not maps:
-      _LOGGER.warning('/proc/%d/maps not found.' % pid)
-      continue
-    pagemap = procfs.ProcPagemap.load(pid, maps)
-    if not pagemap:
-      _LOGGER.warning('/proc/%d/pagemap not found.' % pid)
-      continue
-    pagemap_dct[pid] = pagemap
-
-  for pid, pagemap in pagemap_dct.iteritems():
-    for vma in pagemap.vma_internals.itervalues():
-      for pageframe, number in vma.pageframes.iteritems():
-        pageframes[pageframe] += number
-
-  return pageframes
-
-
-def count_statm(pids):
-  resident = 0
-  shared = 0
-  private = 0
-
-  for pid in pids:
-    statm = procfs.ProcStatm.load(pid)
-    if not statm:
-      _LOGGER.warning('/proc/%d/statm not found.' % pid)
-      continue
-    resident += statm.resident
-    shared += statm.share
-    private += (statm.resident - statm.share)
-
-  return (resident, shared, private)
-
-
-def main(argv):
-  logging_handler = logging.StreamHandler()
-  logging_handler.setLevel(logging.WARNING)
-  logging_handler.setFormatter(logging.Formatter(
-      '%(asctime)s:%(name)s:%(levelname)s:%(message)s'))
-
-  _LOGGER.setLevel(logging.WARNING)
-  _LOGGER.addHandler(logging_handler)
-
-  if sys.platform.startswith('linux'):
-    logging.getLogger('procfs').setLevel(logging.WARNING)
-    logging.getLogger('procfs').addHandler(logging_handler)
-    pids = list_pids(argv)
-    pageframes = count_pageframes(pids)
-  else:
-    _LOGGER.error('%s is not supported.' % sys.platform)
-    return 1
-
-  # TODO(dmikurube): Classify this total RSS.
-  print len(pageframes) * 4096
-
-  return 0
-
-
-if __name__ == '__main__':
-  sys.exit(main(sys.argv))