#!/usr/bin/env ruby
# from: https://gist.github.com/kenn/5105061/raw/ac7ebc6be7008c35b72560cc4e05b7cc14eb4919/memstats.rb

#------------------------------------------------------------------------------
# Aggregate Print useful information from /proc/[pid]/smaps
#
# pss  - Roughly the amount of memory that is "really" being used by the pid
# swap - Amount of swap this process is currently using
#
# Reference:
#  http://www.mjmwired.net/kernel/Documentation/filesystems/proc.txt#361
#
# Example:
#   # ./memstats.rb 4386
#   Process:             4386
#   Command Line:        /usr/bin/mongod -f /etc/mongo/mongod.conf
#   Memory Summary:
#     private_clean             107,132 kB
#     private_dirty           2,020,676 kB
#     pss                     2,127,860 kB
#     rss                     2,128,536 kB
#     shared_clean                  728 kB
#     shared_dirty                    0 kB
#     size                  149,281,668 kB
#     swap                    1,719,792 kB
#------------------------------------------------------------------------------

class Mapping
  FIELDS = %w[ size rss shared_clean shared_dirty private_clean private_dirty swap pss ]
  attr_reader :address_start
  attr_reader :address_end
  attr_reader :perms
  attr_reader :offset
  attr_reader :device_major
  attr_reader :device_minor
  attr_reader :inode
  attr_reader :region

  attr_accessor :size
  attr_accessor :rss
  attr_accessor :shared_clean
  attr_accessor :shared_dirty
  attr_accessor :private_dirty
  attr_accessor :private_clean
  attr_accessor :swap
  attr_accessor :pss

  def initialize( lines )

    FIELDS.each do |field|
      self.send("#{field}=", 0)
    end

    parse_first_line( lines.shift )
    lines.each do |l|
      parse_field_line(l)
    end
  end

  def parse_first_line( line )
    parts = line.strip.split
    @address_start, @address_end = parts[0].split("-")
    @perms = parts[1]
    @offset = parts[2]
    @device_major, @device_minor = parts[3].split(":")
    @inode = parts[4]
    @region = parts[5] || "anonymous"
  end

  def parse_field_line( line )
    parts = line.strip.split
    field = parts[0].downcase.sub(':','')
    if respond_to? "#{field}="
      value = Float(parts[1]).to_i
      self.send( "#{field}=", value )
    end
  end
end

def consume_mapping( map_lines, totals )
  m = Mapping.new( map_lines )

  Mapping::FIELDS.each do |field|
    totals[field] += m.send( field )
  end
  return m
end

abort 'usage: memstats [pid]' unless ARGV.first
pid = ARGV.shift.to_i
totals = Hash.new(0)
mappings = []

File.open( "/proc/#{pid}/smaps" ) do |smaps|

  map_lines = []

  loop do
    break if smaps.eof?
    line = smaps.readline.strip
    case line
    when /\w+:\s+/
      map_lines << line
    when /[0-9a-f]+:[0-9a-f]+\s+/
      if map_lines.size > 0 then
        mappings << consume_mapping( map_lines, totals )
      end
      map_lines.clear
      map_lines << line
    else
      break
    end
  end
end

# http://rubyforge.org/snippet/download.php?type=snippet&id=511
def format_number( n )
  n.to_s.gsub(/(\d)(?=\d{3}+(?:\.|$))(\d{3}\..*)?/,'\1,\2')
end

def get_commandline( pid )
  commandline = IO.read( "/proc/#{pid}/cmdline" ).split("\0")
  if commandline.first =~ /java$/ then
    loop { break if commandline.shift == "-jar" }
    return "[java] #{commandline.shift}"
  end
  return commandline.join(' ')
end

if ARGV.include? '--yaml'
  require 'yaml'
  puts Hash[*totals.map do |k,v|
    [k + '_kb', v]
  end.flatten].to_yaml
else
  puts "#{"Process:".ljust(20)} #{pid}"
  puts "#{"Command Line:".ljust(20)} #{get_commandline(pid)}"
  puts "Memory Summary:"
  totals.keys.sort.each do |k|
    puts "  #{k.ljust(20)} #{format_number( totals[k] ).rjust(12)} kB"
  end
end


