#!/usr/bin/env ruby
$LOAD_PATH.unshift("#{File.dirname(__FILE__)}/../lib")
require 'passenger/platform_info'
 
# Container for tabular data.
class Table
  def initialize(column_names)
    @column_names = column_names
    @rows = []
  end
  
  def add_row(values)
    @rows << values.to_a
  end
  
  def add_rows(list_of_rows)
    list_of_rows.each do |row|
      add_row(row)
    end
  end
  
  def remove_column(name)
    i = @column_names.index(name)
    @column_names.delete_at(i)
    @rows.each do |row|
      row.delete_at(i)
    end
  end
  
  def to_s(title = nil)
    max_column_widths = [1] * @column_names.size
    (@rows + [@column_names]).each do |row|
      row.each_with_index do |value, i|
        max_column_widths[i] = [value.to_s.size, max_column_widths[i]].max
      end
    end
    
    format_string = max_column_widths.map{ |i| "%-#{i}s" }.join("  ")
    header = sprintf(format_string, *@column_names).rstrip << "\n"
    if title
      free_space = header.size - title.size - 2
      if free_space <= 0
        left_bar_size = 3
        right_bar_size = 3
      else
        left_bar_size = free_space / 2
        right_bar_size = free_space - left_bar_size
      end
      result = "#{"-" * left_bar_size} #{title} #{"-" * right_bar_size}\n"
      result << header
    else
      result = header.dup
    end
    result << ("-" * header.size) << "\n"
    @rows.each do |row|
      result << sprintf(format_string, *row).rstrip << "\n"
    end
    result
  end
end
 
class MemoryStats
  class Process
    attr_accessor :pid
    attr_accessor :ppid
    attr_accessor :threads
    attr_accessor :vm_size              # in KB
    attr_accessor :name
    attr_accessor :private_dirty_rss    # in KB
    
    def vm_size_in_mb
      return sprintf("%.1f MB", vm_size / 1024.0)
    end
    
    def private_dirty_rss_in_mb
      if private_dirty_rss.is_a?(Numeric)
        return sprintf("%.1f MB", private_dirty_rss / 1024.0)
      else
        return "?"
      end
    end
    
    def to_a
      return [pid, ppid, threads, vm_size_in_mb, private_dirty_rss_in_mb, name]
    end
  end
  
  def start
    apache_processes = list_processes(:exe => PlatformInfo::HTTPD)
    print_process_list("Apache processes", apache_processes)
    
    puts
    passenger_processes = list_processes(:match => /(^Passenger |^Rails:|ApplicationPoolServerExecutable)/)
    print_process_list("Passenger processes", passenger_processes, :show_ppid => false)
    
    if RUBY_PLATFORM !~ /linux/
      puts
      puts "*** WARNING: The private dirty RSS can only be displayed " <<
        "on Linux. You're currently using '#{RUBY_PLATFORM}'."
    elsif ::Process.uid != 0 && (apache_processes + passenger_processes).any?{ |p| p.private_dirty_rss.nil? }
      puts
      puts "*** WARNING: Please run this tool as root. Otherwise the " <<
        "private dirty RSS of processes cannot be determined."
    end
  end
  
  # Returns a list of Process objects that match the given search criteria.
  #
  #  # Search by executable path.
  #  list_processes(:exe => '/usr/sbin/apache2')
  #  
  #  # Search by executable name.
  #  list_processes(:name => 'ruby1.8')
  #  
  #  # Search by process name.
  #  list_processes(:match => 'Passenger FrameworkSpawner')
  def list_processes(options)
    if options[:exe]
      name = options[:exe].sub(/.*\/(.*)/, '\1')
      ps = "ps -C '#{name}'"
    elsif options[:name]
      ps = "ps -C '#{options[:name]}'"
    elsif options[:match]
      ps = "ps -A"
    else
      raise ArgumentError, "Invalid options."
    end
    
    processes = []
    list = `#{ps} -o pid,ppid,nlwp,vsz,command`.split("\n")
    list.shift
    list.each do |line|
      line.gsub!(/^ */, '')
      line.gsub!(/ *$/, '')
      
      p = Process.new
      p.pid, p.ppid, p.threads, p.vm_size, p.name = line.split(/ +/, 5)
      if p.name !~ /^ps/ && (!options[:match] || p.name.match(options[:match]))
        [:pid, :ppid, :threads, :vm_size].each do |attr|
          p.send("#{attr}=", p.send(attr).to_i)
        end
        p.private_dirty_rss = determine_private_dirty_rss(p.pid)
        processes << p
      end
    end
    return processes
  end
 
private
  # Returns the private dirty RSS for the given process, in KB.
  def determine_private_dirty_rss(pid)
    total = 0
    File.read("/proc/#{pid}/smaps").split("\n").each do |line|
      line =~ /^(Private)_Dirty: +(\d+)/
      if $2
        total += $2.to_i
      end
    end
    return total
  rescue Errno::EACCES, Errno::ENOENT
    return nil
  end
  
  def print_process_list(title, processes, options = {})
    table = Table.new(%w{PID PPID Threads VMSize Private Name})
    table.add_rows(processes)
    if options.has_key?(:show_ppid) && !options[:show_ppid]
      table.remove_column('PPID')
    end
    puts table.to_s(title)
    
    total_private_dirty_rss = 0
    some_private_dirty_rss_cannot_be_determined = false
    processes.each do |p|
      if p.private_dirty_rss.is_a?(Numeric)
        total_private_dirty_rss += p.private_dirty_rss
      else
        some_private_dirty_rss_cannot_be_determined = true
      end
    end
    puts   "### Processes: #{processes.size}"
    printf "### Total private dirty RSS: %.2f MB", total_private_dirty_rss / 1024.0
    if some_private_dirty_rss_cannot_be_determined
      puts " (?)"
    else
      puts
    end
  end
end
 
MemoryStats.new.start