memory_calc_linux.py

#!/usr/bin/env python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
import re
import string

PAGE_SIZE_IN_KB = 4096 // 1024

def eprint(*args, **kwargs):
  print(*args, file=sys.stderr, **kwargs)

def read_file(path):
  with open(path) as f:
    return f.read()

def make_proc_file_reader(name):
  def getter(pid):
    path = '/proc/{}/{}'.format(pid, name)
    return read_file(path)
  return getter

statm = make_proc_file_reader('statm')
status = make_proc_file_reader('status')
smaps = make_proc_file_reader('smaps')

def parse_heading(heading):
  parts = re.split(r'\s+', heading, maxsplit=6)
  if len(parts) == 5:
    parts.append('')
  if len(parts) == 6:
    parts.append('')
  return tuple(parts)

def slow_calc(pid):
  s = smaps(pid)
  lines = filter(None, s.split('\n'))
  data = {}
  i = 0
  while i < len(lines):
    fields = {}
    heading = lines[i]
    i += 1
    while i < len(lines) and (lines[i][0] not in string.hexdigits.lower()):
      m = re.match('^(\w+):\s+(\d+)\s+kB$', lines[i])
      if m:
        key, value = m.groups(1)
        fields[key] = int(value)
      i += 1
    data[parse_heading(heading)] = fields  
  
  annotated = {}
  for heading, fields in data.items():
    address, permissions, _, _, _, path, is_deleated = heading
    should_count = '3'
    if not path.startswith("/"):
      should_count = '1'
    elif path.startswith("/") and permissions == "rw-p":
      should_count = '2'
    elif path.startswith("/") and permissions == "r--p":
      should_count = '2'
    annotated[(should_count, path, permissions, heading)] = fields

  eprint()
  eprint('Slow calc for {}:'.format(pid))

  total = 0
  eprint('   path                                                                     permissions prv_d prv_c  swap shd_d shd_c  counting'.format(**locals()))
  for heading, fields in sorted(annotated.items()):
    should_count, path, permissions, _ = heading
    private_dirty = fields["Private_Dirty"]
    private_clean = fields["Private_Clean"]
    swap = fields["Swap"]
    shared_clean = fields["Shared_Clean"]
    shared_dirty = fields["Shared_Dirty"]
    if should_count == '1':
      reason = 'everything'
      total += private_dirty + private_clean + swap + shared_clean + shared_dirty
    elif should_count == '2':
      reason = 'prv_d + prv_c'
      total += private_dirty + private_clean
    elif should_count == '3':
      reason = 'nothing'
    else:
      assert False
    eprint('{should_count:1} {path:80} {permissions} {private_dirty:5} {private_clean:5} {swap:5} {shared_dirty:5} {shared_clean:5}  {reason}'.format(**locals()))

  return total

def get_vmswap(pid):
  text = status(pid)
  lines = text.split('\n')
  for line in lines:
    m = re.match('^VmSwap:\s+(\d+)\s+kB$', line)
    if m:
      return int(m.group(1))

def get_resident_and_shared(pid):
  s = statm(pid)
  nums = map(int, s.split(' '))
  size, resident, shared, text, lib, data, dt = nums
  return resident, shared

def fast_calc(pid):
  vmswap = get_vmswap(pid)
  resident, shared = get_resident_and_shared(pid)
  page_size = PAGE_SIZE_IN_KB
  total = vmswap + page_size * (resident - shared)
  eprint()
  eprint('Fast calc for {}:'.format(pid))
  eprint('vmswap + page_size * (resident - shared) = total'.format(**locals()))
  eprint('{vmswap:6} + {page_size:9} * ({resident:8} - {shared:6}) = {total:5}'.format(**locals()))
  return total

def get_result(pid):
  return pid, fast_calc(pid), slow_calc(pid)

def show_help():
  print('Usage: calc.py <pid> [<pid>...]')

def main(args):
  if '-h' in args or '--h' in args:
    show_help()
    exit()
  pids = [arg for arg in args[1:] if not arg.startswith('-')]
  results = map(lambda pid: get_result(pid), pids)
  eprint()
  print('pid   fast_total slow_total fast_total-slow_total')
  for pid, fast_total, slow_total  in results:
    print('{:8} {:7}    {:7}               {:7}'.format(pid, fast_total, slow_total, fast_total-slow_total))

if __name__ == '__main__':
  main(sys.argv)