graphlog.py

# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import sys
import time
from optparse import OptionParser
import memreport
import grapher
from subprocess import Popen, PIPE

def logger(iterations, wait, source="/proc", chrome_only=False):
  """ Logs smaps behavior of chrome and non-chrome processes.

  Required arguments:
  iterations -- number of memory data snapshots to take
  wait -- seconds between each data capture
  chrome -- bool decides if data from all or only chrome processes is acquired

  Keyword arguments:
  source -- the directory used for parsing (default: /proc)
  chrome_only -- If True, does not account for non-Chrome process

  Returns:
  snapshots_list -- list with ProcSnapshot objects
  """
  count = 0
  snapshots_list = []
  while count < iterations:
    snapshot = memreport.ProcSnapshot(source, chrome_only)
    #TODO: Decide what information we want to present instead of saving
    # all the data available
    snapshots_list.append(snapshot)
    count += 1
    if count < iterations:
      time.sleep(wait)

  return snapshots_list

def userArgs():
  """
  Function uses the OptionParser object to allow for voluntary user input from
  the command line.
  """
  parser = OptionParser()
  parser.add_option("-p", "--proc", type="string", metavar="SOURCE DIR",
                    dest="source", help="source directory [default: /proc]",
                    default="/proc")
  parser.add_option("-i", "--iterations", type='int', default=1,
                     help="number of snapshot iterations")
  parser.add_option("-t", "--wait_time", type='int', default=0,
                    help="time in seconds between each snapshot",
                    dest="interval")
  parser.add_option("-c", "--chrome_only", action="store_true", default=False,
                    help="show only chrome processes")
  parser.add_option("-g", "--graph", default="timelog",
                    help="choose graph type (timelog, snapshot, rss)")
  parser.add_option("-m", "--memory", default="Pss",
                    help="choose memory type to display")

  #TODO: add "save file to destination" option + generate it on the fly
  return parser.parse_args()


def setTimelogDataTable(snapshots_list, mem_type, show_gpu=False, uptime=0):
  '''
  Given a list of ProcSnapshots objects, return a list formatted properly for
  a line graph
  '''
  chrome, sys_pids_list = getUniquePidsList(snapshots_list)

  # Doing only for chrome right now. Figure out how to use this for
  # non-Chrome as well
  data_table = []
  data_table.append(["Time"])
  chrome_keys = sorted(chrome.keys(), key=int)
  for pid in chrome_keys:
    data_table[0].append(pid + ': ' + chrome[pid])
  if show_gpu:
    data_table[0].append("Total Graphics Memory")
    data_table[0].append("Current Graphics Memory in GTT")

  i = 0
  for snapshot in snapshots_list:
    i += 1
    data_table.append([int(snapshot.timestamp - snapshots_list[0].timestamp +
                           uptime)])
    # Iterate through all the pids in the first table
    for pid in chrome_keys:
      try:
        data_table[i].append(round(snapshot.chrome_smaps.pid[pid][mem_type] /
                                   1024.0, 1))
      except KeyError:
        data_table[i].append(0)
    if show_gpu:
      data_table[i].append(round(snapshot.gpu_mem.total / 1024.0, 1))
      data_table[i].append(round(snapshot.gpu_mem.current / 1024.0, 1))

  return data_table


def getUniquePidsList(snapshots_list):
  '''
  Return a sorted list with all the unique PIDs in a list of ProcSnapshot
  objects.
  '''
  # Create the two sets
  #chrome_set = set()
  sys_set = set()

  chrome = {}

  # Iterate through snapshots list
  for snapshot in snapshots_list:

    # Add chrome PIDs to set
    for pid in snapshot.chrome_pids_list:
      chrome[pid] = snapshot.chrome_types[pid]
      #chrome_set.add(pid)

    # Add non-Chrome PIDs to set
    for pid in snapshot.sys_pids_list:
      sys_set.add(pid)

  # Return sorted list of the sets
  return chrome, sorted(sys_set, key=int)


def setSnapshotDataTable(snapshot, mem_type, chrome_only):
  '''
  Given a list of ProcSnapshots objects, return a list formatted properly for
  a pie graph
  '''
  data_table = []
  data_table.append(["PID", "Size"])
  mem_type = mem_type.capitalize()
  sorted_chrome = snapshot.chrome_smaps.sortByMemType(mem_type)
  for item in sorted_chrome:
    if item[0] != "Total":
      data_table.append([item[0], item[1][mem_type]])

  if not chrome_only:
    sorted_sys = snapshot.sys_smaps.sortByMemType(mem_type)
    for item in sorted_sys:
      if item[0] != "Total":
        data_table.append([item[0], item[1][mem_type]])

  return data_table

def setRssDataTable(snapshot, chrome_only):
  '''
  Given a list of ProcSnapshots objects, return a list formatted properly for
  a bar graph
  '''
  data_table = []
  data_table.append(["PID", "Private", "Shared"])
  for pid in (snapshot.chrome_pids_list + ["Total"]):
    private = snapshot.chrome_smaps.pid[pid]["Private"]
    shared = snapshot.chrome_smaps.pid[pid]["Shared"]
    data_table.append([pid, private, shared])

  if not chrome_only:
    for pid in (snapshot.sys_pids_list + ["Total"]):
      private = snapshot.sys_smaps.pid[pid]["Private"]
      shared = snapshot.sys_smaps.pid[pid]["Shared"]
      data_table.append([pid, private, shared])

  return data_table

def getSystemParameters():
  sys_info = {}

  # first, get uname
  sys_info["uname"] = Popen("uname -a", shell=True,
                            stdout=PIPE).stdout.readline()

  # now get lsb info
  lsb_f = open("/etc/lsb-release", 'r')
  lsb = lsb_f.readlines()
  lsb_f.close()
  for line in lsb:
    line = line.split('=')
    try:
      sys_info[line[0]] = line[1]
    except IndexError:
      pass

  return sys_info

# Main Function
#TODO(thiagog): Add command line support to input all required fields for
# graph creation.
if __name__ == "__main__":
  options, args = userArgs()

  if options.graph not in ("rss", "snapshot", "timelog"):
    sys.exit("ERROR: graph argument invalid")

  uptime_file = open("/proc/uptime", 'r')
  uptime = float(uptime_file.readline().split()[0])
  uptime_file.close()

  sys_info = getSystemParameters()

  timestamp = time.asctime(time.localtime())
  snapshots = logger(options.iterations, options.interval,
                     options.source, options.chrome_only)


  gpu_mem = ("Pss", "Rss")

  # For testing right now
  mem_type = "Pss"


  # Testing on linux machine (no gpu memory information)
  show_gpu = False

  html_args = {}

  if options.graph == "timelog":
    graphs = ("Pss", "Rss", "Private_Clean", "Private_Dirty", "Shared_Clean",
              "Shared_Dirty", "Private", "Shared", "Referenced", "Anonymous")
    for mem_type in graphs:
      show_gpu = mem_type in gpu_mem
      html_args[mem_type] = setTimelogDataTable(snapshots, mem_type, show_gpu,
                                                uptime)

  if options.graph == "snapshot":
    snapshot = snapshots[0]
    html_args["title"] = mem_type + " [kBs] per PID"
    html_args["data"] = setSnapshotDataTable(snapshot, mem_type,
                                        options.chrome_only)

  if options.graph == "rss":
    snapshot = snapshots[0]
    html_args["title"] = "Private and Shared memory usage [kBs] per PID"
    html_args["haxis"] = "PIDs"
    html_args["data"] = setRssDataTable(snapshot, True)

  html_args["version"] = sys_info["uname"]
  html_args["release"] = sys_info["CHROMEOS_RELEASE_DESCRIPTION"]
  html_args["timestamp"] = timestamp
  #print html_args["data"]
  print grapher.chooseGraph(options.graph, html_args)