#!/usr/bin/env ruby
#
# author: komar@cock.li
#
# Usage: Place in /etc/munin/plugins/ (or link it there using ln -s)
#	Add this to your /etc/munin/plugin-conf.d/munin-node:
#		[membyuser]
#		user root # required if /proc can't be read by any user
#		env.USERS postgres squid nginx # All the users you don't want
#                                  # to be calculated as "others".
#                                  # User "root" is not allowed
#                                  # because munin sucks.

users = (ENV['USERS'] || "").split(/\s+/).sort.uniq
case ARGV.shift
when "autoconf"
  if users.empty?
    puts "$USERS not defined."
  else
    puts "yes"
  end
when "config"
	puts "graph_args --base 1024 --lower-limit 0"
	puts "graph_title Memory usage, by user"
	puts "graph_category system"
	puts "graph_info This graph shows memory usage, for monitored users."
	puts "graph_vlabel Bytes"
	puts "graph_period second"
	puts "graph_order #{users.join(' ')} others"
	(users + ["others"]).each do |user|
		puts "#{user}.label #{user}"
		puts "#{user}.info Memory used by user #{user}"
		puts "#{user}.type GAUGE"
		puts "#{user}.draw AREASTACK"
	end
else

  pss_per_uid = {}

  Dir.open("/proc").each do |dir|
    next unless dir =~ /[0-9]+/
    begin
      uid = File.stat(File.join("/proc", dir)).uid
      smaps_rollup = File.read(File.join("/proc", dir, "smaps_rollup"))
      pss      = smaps_rollup.match(    /^Pss\:\s+(\d+) kB/)[1].to_i * 1024
      swap_pss = smaps_rollup.match(/^SwapPss\:\s+(\d+) kB/)[1].to_i * 1024
      pss_per_uid[uid] = (pss_per_uid[uid] || 0) + pss + swap_pss
    rescue Errno::ESRCH, Errno::ENOENT
    end
  end

  require 'etc'

  uid_per_name = {}

  Etc::Passwd.each do |u|
    uid_per_name[u.name] = u.uid
  end

  users.each do |user|
    uid = uid_per_name[user]
    pss = pss_per_uid.delete(uid) || 0
    puts "#{user}.value #{pss}"
  end
  others = pss_per_uid.inject(0){|acc,p|acc + p[1]}
  puts "others.value #{others}"

end