#!/usr/bin/env python
"""
Measure CPU utilization with an active CPU probe.
"""
import os       # fork(), nice()
import resource # getrusage()
import struct   # pack/unpack()
import sys      # argv
import time     # time()
#
from netlogger import util
from netlogger.nllog import OptionParser, get_logger, DoesLogging

## Globals

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

g_stop = 0

## Functions

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

def on_kill(signo, frame):
    global g_stop
    get_logger(__file__).info("killed")
    g_stop = 1

## Classes

class Probe(DoesLogging):
    """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):
        DoesLogging.__init__(self)
        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
        self.log.debug("avail", cpu_time=cpu_time, real_time=real_time)
        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()
            self.cpu_avail = self._avail()

    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)

## Program entry point

def main(args):
    # parse args
    desc = ' '.join(__doc__.split())
    parser = OptionParser(description=desc)
    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)
    log = get_logger(__file__)  # Should be first done, just after parsing args
    log.info("start")
    util.handleSignals(
        (on_kill, ('SIGTERM', 'SIGINT', 'SIGUSR2', 'SIGUSR1', 'SIGHUP')))
    probe = Probe(ms=options.ms, nice=options.nice)
    status = 0
    while not g_stop:
        log.debug("iterate.start")
        probe.spin()
        print "%.1lf" % (probe.cpu_avail * 100.0,)
        time.sleep(1)
        log.debug("iterate.end", status=0)
    log.info("end", status=status)
    return status

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