#!/usr/bin/python 

## ######################################################
## Copyright (c) 2009, Nicholas Kidd 
## All rights reserved.
## 
## Redistribution and use in source and binary forms, with or without
## modification, are permitted provided that the following conditions are met:
##     * Redistributions of source code must retain the above copyright
##       notice, this list of conditions and the following disclaimer.
##     * Redistributions in binary form must reproduce the above copyright
##       notice, this list of conditions and the following disclaimer in the
##       documentation and/or other materials provided with the distribution.
##     * Neither the name of the <organization> nor the
##       names of its contributors may be used to endorse or promote products
##       derived from this software without specific prior written permission.
## 
## THIS SOFTWARE IS PROVIDED BY <copyright holder> ''AS IS'' AND ANY
## EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
## WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
## DISCLAIMED. IN NO EVENT SHALL <copyright holder> BE LIABLE FOR ANY
## DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
## (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
## LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
## ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
## (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
## SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
## ######################################################

## ######################################################
## Comments
##
## class JobMonitor:
##    Class to run a shell command and reports the time, memory usage,
##    and return code. All of the code is pieced together from various places online
##    (e.g., the python recipe "http://code.activestate.com/recipes/286222/"); 
##    however, I wanted something that was simple and aggregrated it all together.
##
##    Use cases
##       * Shell Job :
##         (cmd,rcode,time,mem) = jobmonitor.RunJob('[[ shell command ]]')
##       * Non-Shell :
##         (cmd,rcode,time,mem) = jobmonitor.RunJob([args,list],shell=False,timeout=XXX,stdout=o,stderr=e)


import os
import subprocess
import time
import sys
import time
import signal

## ######################################################
## Run the job [job] and return the timing and memory
## statistics. Simply a handy wrapper.
def RunJob(job,timeout=None,name=None,stdin=None,stdout=None,stderr=None,shell=True):
    jm = JobMonitor(job,timeout,name,stdin,stdout,stderr,shell)
    jm.Run()
    return jm.Result()

## #######################################################
## class MemMonitor
## Given a pid, read /proc/<pid>/status and returns
## memory, resident, or stacksize usage
##
## Taken from : http://code.activestate.com/recipes/286222/
class MemMonitor:
    _scale = {'kB': 1024.0, 'mB': 1024.0*1024.0,
              'KB': 1024.0, 'MB': 1024.0*1024.0}

    def __init__(self,pid):
        self.pid = pid
        self._proc_status = '/proc/%d/status' % pid

    def _VmB(self,VmKey):
        '''Private.
        '''
        # global _proc_status, _scale
        # get pseudo file  /proc/<pid>/status
        try:
            t = open(self._proc_status)
            v = t.read()
            t.close()
            # print v
            # get VmKey line e.g. 'VmRSS:  9999  kB\n ...'
            i = v.index(VmKey)
            v = v[i:].split(None, 3)  # whitespace
            if len(v) < 3:
                return 0.0  # invalid format?
            # convert Vm value to bytes
            return float(v[1]) * MemMonitor._scale[v[2]]
        except Exception,e:
            print '  Error ::> ',e
            return 0.0  # non-Linux?
    
    def memory(self,since=0.0):
        '''Return memory usage in bytes.
        '''
        return self._VmB('VmSize:') - since
    
    
    def resident(self,since=0.0):
        '''Return resident memory usage in bytes.
        '''
        return self._VmB('VmRSS:') - since
    
    
    def stacksize(self,since=0.0):
        '''Return stack size in bytes.
        '''
        return self._VmB('VmStk:') - since

## ########################################################
## Runs a shell command [cmd] and waits for completion.
## Records :
##   maximum memory used
##   user-execution time
##   return code
## @return (cmd, return code , time to execute)
class JobMonitor:
    Debug = False

    def __init__(self,cmd,
        timeout=None,
        name=None,
        stdin=None,
        stdout=None,
        stderr=None,
        shell=True):
        self.cmd     = cmd 
        if name == None:
          self.name = cmd
        else:
          self.name = name
        self.rc      = 0
        self.time    = 0
        self.max     = 0
        self.start   = 0
        self.timeout = timeout
        self.pgid    = 0
        self.stdin   = stdin
        self.stdout  = stdout
        self.stderr  = stderr
        self.shell   = shell

    def Result(self):
        return (self.cmd , self.rc , self.time , self.max)

    def Command(self):
        return self.cmd

    def ReturnCode(self):
        return self.rc

    def Time(self):
        return self.time
    
    def MaxMemory(self):
        return self.max

    def Run(self):
        if JobMonitor.Debug:
          print '--- Executing "%s"\n' % self.cmd
        try:
            self.start = time.time()
            p = subprocess.Popen( self.cmd ,
                stdin  = self.stdin,
                stdout = self.stdout,
                stderr = self.stderr,
                shell  = self.shell)
            if JobMonitor.Debug:
                print '  Pid :> ',p.pid
            monitor = MemMonitor(p.pid)
            self.rc = p.poll()
            while( self.rc == None ):
                mem = monitor.memory()
                if JobMonitor.Debug:
                    print '  Mem :> ',mem
                if (mem > self.max):
                    self.max = mem
                time.sleep(1)
                self.time = time.time() - self.start
                self.rc = p.poll()
                if (self.timeout != None and self.time >= self.timeout):
                  break
            # Check for timeout
            if (self.rc == None):
              os.kill(p.pid,signal.SIGKILL)
              killedpid,stat = os.waitpid(p.pid,0)#,os.WNOHANG)
              if killedpid == 0:
                print >> sys.stderr, "TODO: FIXME"
            if self.rc == None: #JobMonitor.Debug:
              print ' %s(%d) :> %f s for %s' % (('Completion','Timeout')[self.rc == None],p.pid,self.time,self.name)
        except OSError, e:
            print >>sys.stderr, "  >> Execution failed:",e
            self.rc = self.time = self.max = -60