#!/usr/bin/python
#!@PYTHON@
# proc-time.py -- Get resource usage from /proc for a command (GNU/Linux).
#
# awaiting PC's linux-mm patch for getrusage
#
# source file of the GNU LilyPond music typesetter
# Licence: GNU GPL
# 
# (c) 2000 Jan Nieuwenhuizen <janneke@gnu.org>
#
# GNU time doesn't report memory stuff on FreeBSD either,
# but I couldn't find an easy way like to get this info (like /proc).
#
# C translation also available

name = 'proc-time'
version = '@TOPLEVEL_VERSION@'
if version == '@' + 'TOPLEVEL_VERSION' + '@':
	version = '(unknown version)'		# uGUHGUHGHGUGH
  
import string
import getopt
import sys
import os
import time
import posix
import signal
import select

# Attempt to fix problems with limited stack size set by Python!
# Sets unlimited stack size. Note that the resource module only
# is available on UNIX.
try:
       import resource
       resource.setrlimit (resource.RLIMIT_STACK, (-1, -1))
except:
       pass

def identify ():
	sys.stderr.write ("%s from LilyPond %s\n" % (name, version))

def print_usage ():
	print (r"""
Usage: %s [OPTION]... COMMAND

Get resource usage from /proc for COMMAND (GNU/Linux).

Options:
  -b, --heartbeat     show info at every probe
  -h, --help          this help
  -i, --interval=TIME set heartbeat to TIME seconds
  -n, --no-fork       don't fork
  -t, --test          test mode
  -V, --verbose       be verbose
  -v, --version       version information
""" % name)

def print_version ():
	sys.stderr.write (r"""%s (GNU LilyPond) %s\n""" % (name, version))

(options, files) = getopt.getopt (sys.argv[1:], 'bhi:ntvV', ['heartbeat', 'help', 'interval=', 'no-fork', 'test', 'verbose' 'version'])

heartbeat = 0
interval = 0.5
nofork = 0
test = 0
verbose = 0

for opt in options:
	o = opt[0]
	a = opt[1]
	if o == '--heartbeat' or o == '-b':
		heartbeat = 1
	elif o == '--help' or o == '-h':
		print_usage ()
		sys.exit (0)
	elif o == '--interval' or o == '-i':
		interval = string.atof (a)
	elif o == '--no-fork' or o == '-n':
		nofork = 1
	elif o == '--test' or o == '-t':
		test = 1
        elif o == '--verbose' or o == '-V':
	      verbose = 1
	elif o == '--version' or o == '-v':
		print_version ()
		sys.exit (0)
	else:
		print o
		raise getopt.error

if len (files) < 1:
	help ()
	sys.exit (2)

status = 0
PAGE_SIZE = 4096.0
JIF_TO_SEC = 0.01
PAGE_TO_MEG = PAGE_SIZE / (1024 * 1024)

def head (name):
	try:
		f = open (name, "r")
		if f:
			str = f.readline ()
			f.close ()
		else:
			str = 0
	except:
		str = 0
	return str

class process_info:
	def __init__ (self, pid):
		self.t = 0
		self.id = pid
		self.cum_size = 0
		self.max_size = 0
		self.cum_rss = 0
		self.max_rss = 0
		self.stop = 0
		self.stat_name = "/proc/%d/stat" % self.id
		self.statm_name = "/proc/%d/statm" % self.id
		self.user_jiffies = 0
		self.system_jiffies = 0
		self.start = time.time ()

	def update_mem_stats (self):
		stat_str = head (self.statm_name)
		if stat_str:
			(size, rss, shared, code, data, lib, dirty) = string.split (stat_str, " ")
			self.t = self.t + 1
			size = string.atoi (size)
			self.max_size = max (size, self.max_size)
			self.cum_size = self.cum_size + size
			rss = string.atoi (rss)
			self.max_rss = max (rss, self.max_rss)
			self.cum_rss = self.cum_rss + rss


	def print_mem_stats (self):
		if self.t:
			  avg_size = self.cum_size / self.t
			  avg_rss = self.cum_rss / self.t
			  sys.stderr.write ("MAXSIZE: %6.3fM(%d), MAXRSS: %6.3fM(%d)\n" % (self.max_size * PAGE_TO_MEG, self.max_size, self.max_rss * PAGE_TO_MEG, self.max_rss))
			  sys.stderr.write ("AVGSIZE: %6.3fM(%d), AVGRSS: %6.3fM(%d)\n" % (avg_size * PAGE_TO_MEG, avg_size, avg_rss * PAGE_TO_MEG, avg_rss))

	def update_time_stats (self):
		stat_str = head (self.stat_name)
		if stat_str:
			stat_array = string.split (stat_str, " ")
			self.user_jiffies = string.atoi (stat_array[13])
			self.system_jiffies = string.atoi (stat_array[14])
		self.stop = time.time ()

	def print_time_stats (self):
		sys.stderr.write ("user: %6.2f(%d) system: %6.2f(%d)\n" % 
		(self.user_jiffies * JIF_TO_SEC, self.user_jiffies, self.system_jiffies * JIF_TO_SEC, self.system_jiffies))
		sys.stderr.write ("elapsed: %6.3fM\n" % (self.stop - self.start))



command = files[0]
args = files

def handler (*args):
	global child_pi
	sys.stderr.write ("\n")
	child_pi.update_time_stats ()
	child_pi.update_mem_stats ()
	child_pi.print_time_stats ()
	child_pi.print_mem_stats ()
	if test:
		sys.stderr.write ("handlerOUT")

signal.signal (signal.SIGINT, handler)
signal.signal (signal.SIGTERM, handler)
signal.signal (signal.SIGCHLD, handler)

# why doesn't it work when put in a function?
#def run ():
#	global child_pi
#	if nofork:
#		child_pi = process_info (os.fork ())
#	else
#		child_pi = process_info (1)
#	if child_pi.id:
#		#Parent
#		child_pi.update_mem_stats ()
#		while 1:
#			try:

if nofork:
	child_pi = process_info (1)
else:
	child_pi = process_info (os.fork ())
if child_pi.id:
	#Parent
	child_pi.update_mem_stats ()
	if interval <= 0:
		  (pid, status) = os.waitpid (child_pi.id, 0)
		  sys.exit (status);
	while 1:
		try:
			(pid, status) = os.waitpid (child_pi.id, os.WNOHANG)
			if pid:
				break
			time.sleep (interval)
			child_pi.update_mem_stats ()

			# hit enter to get mem stats at interval
			(i, o, e) =  select.select ([sys.stdin], [], [], 0)
			if i:
				sys.stdin.read (1)
				sys.stderr.write ("\n")
				child_pi.print_mem_stats ()
			elif heartbeat:
				sys.stderr.write ("\n")
				child_pi.print_mem_stats ()
		except:
			break

else:
	# Child
	os.execvp (command, args)

if not status:
	try:
		(pid, status) = os.waitpid (child_pi.id, 0)
	except:
		pass

if status:
    sys.stderr.write ("Command exited with non-zero exit status: %d\n" % status)

sys.exit (status)