#!/usr/bin/env python
"""
Active CPU probe
"""
import os       # fork(), nice()
import optparse # parser
import resource # getrusage()
import struct   # pack/unpack()
import sys      # argv
import time     # time()

DSZ, ISZ = len(struct.pack('d', 0.0)), len(struct.pack('i', 0))

def os_readn(fd, nbytes):
    buf = ""
    while len(buf) < nbytes:
        s = os.read(fd, nbytes - len(buf))
        if s:
            buf += s
        else:
            time.sleep(0.001)
    return buf
       
class Probe:
    """Active probe.
    Start with a nice value and milliseconds
    Run spin(), then look at cpu_* attributes
    """
    who = resource.RUSAGE_SELF

    def __init__(self, ms=100, nice=10):
        self.last_ru, self.ru = None, None
        self.last_real, self.real = -1, -1
        self._count = 10000
        self.nice, self.ms = nice, ms
        self.cpu_avail = 0
        
    def _start(self):
        self.last_ru = resource.getrusage(self.who)
        self.last_real = time.time()
        
    def _stop(self):
        self.ru = resource.getrusage(self.who)
        self.real = time.time()

    def _avail(self):
        if self.last_ru is None or self.ru is None:
            return 0.
        cpu_time = self.ru.ru_utime + self.ru.ru_stime - \
                   self.last_ru.ru_utime - self.last_ru.ru_stime
        real_time = self.real - self.last_real
        return cpu_time / real_time
    
    def spin(self):
        """Spin for self.ms milliseconds.
        """
        if self.nice:
            """
            Since we can't ever reduce our nice value, we have to run any niced
            spinner in a forked process.  We use a communication channel to 
            report results from the forked process to the parent.
            """
            rfd, wfd = os.pipe()
            pid = os.fork()
            if pid: # parent
                os.close(wfd) # close write end
                # read results from child
                s = os_readn(rfd, DSZ + ISZ)
                self.cpu_avail, self._count = struct.unpack('di', s)
                os.wait()
            else: # child
                os.close(rfd) # close read end
                os.nice(self.nice) # do nice-ing
                # perform spin
                self._spin()
                # return result to parent
                os.write(wfd, struct.pack('d', self._avail()))
                # also return new loop count
                os.write(wfd, struct.pack('i', self._count))
                sys.exit(0)
        else:
            self._spin()
            
    def _spin(self):
        self._start()
        t_stop = self.last_real + self.ms/1000.
        x, n = 1, 0
        while time.time() < t_stop:
            for i in xrange(self._count):
                x *= 2
            n += 1
        self._stop()
        """
        Adjust the count if it falls outside range 10 to 100.
        This choice of range is arbitrary, and the lower limit is
        more significant (want to avoid over-spinning by more than 10%).
        """
        if n < 10:
            self._count = int(self._count * 0.8)
        elif n > 100:
            self._count = int(self._count * 1.2)
        
def main(args):
    # parse args
    usage = "%prog [options]"
    parser = optparse.OptionParser(usage=usage, version="1.0")
    parser.add_option('-m','--millis', action='store', dest='ms',
                      type="int", default=100,
                      help="number of milliseconds out of " +
                      "every second to run the probe (default=%default)")
    parser.add_option('-n','--nice', action='store', dest='nice',
                      type="int", default=0,
                      help="nice value to give to the process while "+
                      "probing (default=%default)")
    options, args = parser.parse_args(args)
    probe = Probe(ms=options.ms, nice=options.nice)
    while 1:
        probe.spin()
        #print probe.cpu_avail, probe._count
        print "%.1lf" % (probe.cpu_avail * 100.0,)
        time.sleep(1)

if __name__ == '__main__':
    main(sys.argv)
        
        
