multi_process_rss.py 3.6 KB

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