#!/usr/bin/env python

#
# Collect per-process memory use over time.
# Written by Lucas C. Villa Real <lucasvr@us.ibm.com>
# Released under the GNU GPL version 2.
#

import os
import time
import glob
import stat
import errno
import signal
import getopt
import subprocess

import PyDaemon

# Global keepalive flag. Reset by the SIGINT signal handler.
g_alive = True

# Directory where logs will be saved
g_datadir = os.getenv("PWD") + "/data"

# File containing the total memory information
g_dummyfile = "TotalMemory"

class OptionParser :
	def __init__(self) :
		self.appname = os.sys.argv[0]
		self.shortopts = "hvdr:p:"
		self.longopts = [ "help", "version", "debug", "report=", "plot=" ]
		self.toReport = 0
		self.toPlot = 0
		self.daemonize = True
		self.__optind = 0
	
	def ShowVersion(self, retval) :
		revision = "$Revision: 1 $"
		print "MemoryMonitor revision%s" %(revision.split("$")[1].split(":")[1])
		os.sys.exit(retval)

	def ShowUsage(self, retval) :
		print "Syntax: %s <options>\n" %self.appname
		print "Available options are:"
		print "  -d,   --debug       keep the daemon in foreground"
		print "  -v,   --version     show version"
		print "  -r N, --report=N    report the top N memory hog processes"
		print "  -p N, --plot=N      plot the memory consumption of the top N memory hog processes"
		print "  -h,   --help        this help\n"
		os.sys.exit(retval)

	def ParseOptions(self) :
		try :
			options, arguments = getopt.getopt(os.sys.argv[1:], self.shortopts, self.longopts)
		except getopt.GetoptError, detail :
			print "%s: %s" %(self.appname, str(detail))
			self.ShowUsage(1)

		for opt, arg in options :
			if opt in [ "-h", "--help" ] :
				self.ShowUsage(0)
			elif opt in [ "-v", "--version" ] :
				self.ShowVersion(0)
			elif opt in [ "-d", "--debug" ] :
				self.daemonize = False
			elif opt in [ "-r", "--report" ] :
				self.toReport = int(arg)
			elif opt in [ "-p", "--plot" ] :
				self.toPlot = int(arg)
			else :
				print "invalid option %s" %opt
				self.ShowUsage(1)

		return self


class ProcParser :
	def __init__(self) :
		self.contents = ""
		self.data = ""

	def parseStatusFile(self, pid, field) :
		contents = ""
		fp = open("/proc/%d/status" %pid)
		while True :
			buf = fp.readline(4096)
			if not buf :
				break
			idx = buf.find(field)
			if idx >= 0 :
				contents = buf[idx:].split()[1]
				break
		fp.close()
		return contents 

	def getVmSize(self, pid) :
		vmSize = self.parseStatusFile(pid, "VmSize:")
		if len(vmSize) :
			return int(vmSize)
		return -1

	def getVmSwap(self, pid) :
		vmSwap = self.parseStatusFile(pid, "VmSwap:")
		if len(vmSwap) :
			return int(vmSwap)
		return -1

	def getPPid(self, pid) :
		ppid = self.parseStatusFile(pid, "PPid:")
		if len(ppid) :
			return int(ppid)
		return -1

	def getProcessName(self, pid) :
		fp = open("/proc/%d/cmdline" %pid)
		buf = fp.readline(1024)
		fp.close()

		tokens = buf.split('\x00')
		if tokens[0].find('python') >= 0 or tokens[0].find('/bin/sh') >= 0 :
			processName = tokens[1]
		elif len(tokens[0]) > 0 :
			processName = tokens[0]
		else :
			fp = open("/proc/%d/stat" %pid)
			buf = fp.readline(1024)
			fp.close()

			idx_start = buf.find("(")
			idx_end = buf.find(")")
			if idx_start >= 0 and idx_end >= 0 :
				processName = buf[idx_start+1:idx_end]
			else :
				print "unknown --> ",
				print processName
				print "tokens  --> ",
				print tokens
				return ""

		if processName.find("/") >= 0  :
			processName = os.path.basename(processName)
		if processName.find(":") >= 0 :
			processName = processName.split(":")[0]
		if processName.find("-") == 0 :
			processName = processName[1:]

		return processName

	def getTotalMemory(self) :
		totalMemory = 0
		fp = open("/proc/meminfo")
		for line in fp.readlines() :
			if line.find("MemTotal:") == 0 :
				totalMemory = int(line.split()[1])
				break
		fp.close()
		return totalMemory


class ProcessInfo :
	def __init__(self) :
		self.pid = 0
		self.ppid = 0
		self.name = ""
		self.vmSize = 0
		self.vmSwap = 0

	def setPid(self, pid) :
		self.pid = int(pid)

	def setPPid(self, ppid) :
		self.ppid = int(ppid)

	def setName(self, name) :
		self.name = name

	def setVmSize(self, vmSize) :
		self.vmSize = int(vmSize)
	
	def setVmSwap(self, vmSwap) :
		self.vmSwap = int(vmSwap)

	def getPid(self) :
		return self.pid

	def getPPid(self) :
		return self.ppid

	def getName(self) :
		return self.name

	def getVmSize(self) :
		return self.vmSize

	def getVmSwap(self) :
		return self.vmSwap


class ProcessList :
	def __init__(self) :
		self.processList = []
		self.myPid = 0

	def getList(self) :
		self.myPid = os.readlink("/proc/self")
		self.processList = []
		for i in glob.glob("/proc/*/stat") :
			try :
				info = os.stat(i)
			except OSError, e :
				# ignore errors if we failed to open the stat file
				continue

			if stat.S_ISDIR(info.st_mode) :
				# skip directories
				continue
		
			pid = i.split("/")[2]
			if pid == "self" or pid == "1" or pid == self.myPid :
				# skip "self" and the init task
				continue
			pid = int(pid)

			procParser = ProcParser()
			vmSize = procParser.getVmSize(pid)
			if vmSize == -1 :
				# skip kernel threads
				continue

			p = ProcessInfo()
			p.setPid(pid)
			p.setVmSize(vmSize)
			p.setPPid(procParser.getPPid(pid))
			p.setName(procParser.getProcessName(pid))
			p.setVmSwap(procParser.getVmSwap(pid))

			self.processList.append(p)
		return self.processList


class MemoryMonitor :
	def __init__(self) :
		self.captureInterval = 30
		self.availableMemory = 0
		self.referenceTime = int(time.time())

	def createWorkDirectory(self) :
		try :
			os.mkdir(g_datadir)
		except OSError, e :
			if e[0] != errno.EEXIST :
				return False
		return True

	def openFile(self, pObject) :
		return open("%s/%s:%d" %(g_datadir, pObject.getName(), pObject.getPid()), "a")

	def start(self) :
		global g_alive
		print "Monitoring memory usage. Press CTRL+C to stop."

		self.createWorkDirectory()

		self.availableMemory = ProcParser().getTotalMemory()
		totalMemoryFP = open("%s/%s" %(g_datadir, g_dummyfile), "w+")

		while g_alive :
			currentTime = int(time.time())
			processes = ProcessList().getList()
			for p in processes :
				f = self.openFile(p)
				if f :
					f.write("%s %d\n" %(currentTime - self.referenceTime, p.getVmSize()))
					f.close()
			totalMemoryFP.write("%d %d\n" %((currentTime - self.referenceTime), self.availableMemory))
			time.sleep(self.captureInterval)

		totalMemoryFP.close()


class Report :
	def __init__(self, toReport, toPlot) :
		self.gnuplot = None
		self.topList = []
		self.toReport = toReport
		self.toPlot = toPlot

	def getPeak(self, file) :
		peak = 0
		f = open(file)
		for line in f.readlines() :
			memoryUsage = int(line.split()[1])
			if memoryUsage > peak :
				peak = memoryUsage
		f.close()
		return peak

	def scan(self) :
		numElements = max(self.toReport, self.toPlot)
		for file in glob.glob(g_datadir + "/*") :
			if file == "%s/%s" %(g_datadir, g_dummyfile) :
				# skip the dummy file
				continue
			peak = self.getPeak(file)
			if len(self.topList) < numElements :
				self.topList.append((file, peak))
			else :
				idx = 0
				idxToReplace = -1
				for element in self.topList :
					if peak > element[1] :
						idxToReplace = idx
					idx += 1
				if idxToReplace != -1 :
					self.topList.pop(idxToReplace)
					self.topList.append((file, peak))
		self.topList.sort(key=lambda element: element[1], reverse=True)

	def printReport(self) :
		print "Report:"
		for element in self.topList :
			processName = element[0].strip(g_datadir + "/").split(":")[0]
			processId = int(element[0].strip(g_datadir + "/").split(":")[1])
			vmPeak = element[1] / 1024
			print "Process: %s, PID: %d, VmSize Peak: %d MB" %(processName, processId, vmPeak)

	def plotReport(self) :
		print "Plotting report"
		try :
			self.gnuplot = subprocess.Popen(['gnuplot'], stdin=subprocess.PIPE)
		except OSError, e :
			print "Error executing gnuplot: %s" %e[1]
			os.sys.exit(1)
		
		self.gnuplot.stdin.write("plot '%s/%s' title 'Total Memory' with lines\n" %(g_datadir, g_dummyfile))
		for element in self.topList :
			processName = element[0].strip(g_datadir + "/").split(":")[0]
			processId = int(element[0].strip(g_datadir + "/").split(":")[1])
			self.gnuplot.stdin.write("replot '%s' title '%s (%d)' with lines\n" %(element[0], processName, processId))
		
		print "Press Enter to exit."
		os.read(0, 1024)
		

# Signal handler which deals with SIGINT
def cleanup(signum, stack_frame) :
	global g_alive
	g_alive = False

def main() :
	userOptions = OptionParser()
	opts = userOptions.ParseOptions()

	if opts.toReport or opts.toPlot :
		r = Report(opts.toReport, opts.toPlot)
		r.scan()
		if opts.toReport :
			r.printReport()
		if opts.toPlot :
			r.plotReport()
		os.sys.exit(0)

	if opts.daemonize :
		PyDaemon.createDaemon()

	signal.signal(signal.SIGINT, cleanup)
	mm = MemoryMonitor()
	mm.start()
	os.sys.exit(0)

if __name__ == "__main__" :
	main()