wally/sensors/main.py

import os
import sys
import time
import json
import glob
import signal
import os.path
import argparse

from .sensors.utils import SensorInfo
from .daemonize import Daemonize
from .discover import all_sensors
from .protocol import create_protocol


# load all sensors
from . import sensors
sensors_dir = os.path.dirname(sensors.__file__)
for fname in glob.glob(os.path.join(sensors_dir, "*.py")):
    mod_name = os.path.basename(fname[:-3])
    __import__("sensors.sensors." + mod_name)


def get_values(required_sensors):
    result = {}
    for sensor_name, params in required_sensors:
        if sensor_name in all_sensors:
            result.update(all_sensors[sensor_name](**params))
        else:
            msg = "Sensor {0!r} isn't available".format(sensor_name)
            raise ValueError(msg)
    return time.time(), result


def parse_args(args):
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--daemon',
                        choices=('start', 'stop', 'status'),
                        default=None)

    parser.add_argument('-u', '--url', default='stdout://')
    parser.add_argument('-t', '--timeout', type=float, default=1)
    parser.add_argument('-l', '--list-sensors', action='store_true')
    parser.add_argument('sensors_config', type=argparse.FileType('r'),
                        default=None, nargs='?')
    return parser.parse_args(args[1:])


def daemon_main(required_sensors, opts):
    try:
        source_id = str(required_sensors.pop('source_id'))
    except KeyError:
        source_id = None

    sender = create_protocol(opts.url)
    prev = {}

    while True:
        gtime, data = get_values(required_sensors.items())
        curr = {'time': SensorInfo(gtime, True)}
        for name, val in data.items():
            if val.is_accumulated:
                if name in prev:
                    curr[name] = SensorInfo(val.value - prev[name], True)
                prev[name] = val.value
            else:
                curr[name] = SensorInfo(val.value, False)

        if source_id is not None:
            curr['source_id'] = source_id

        sender.send(curr)

        time.sleep(opts.timeout)


def pid_running(pid):
    return os.path.exists("/proc/" + str(pid))


def main(argv):
    opts = parse_args(argv)

    if opts.list_sensors:
        print "\n".join(sorted(all_sensors))
        return 0

    if opts.daemon is not None:
        pid_file = "/tmp/sensors.pid"
        if opts.daemon == 'start':
            required_sensors = json.loads(opts.sensors_config.read())

            if "protocol_data" not in required_sensors:
                raise ValueError("No protocol data provided in config")

            def root_func():
                daemon_main(required_sensors, opts)

            daemon = Daemonize(app="perfcollect_app",
                               pid=pid_file,
                               action=root_func)
            daemon.start()
        elif opts.daemon == 'stop':
            if os.path.isfile(pid_file):
                pid = int(open(pid_file).read())
                if pid_running(pid):
                    os.kill(pid, signal.SIGTERM)

                time.sleep(0.1)

                if pid_running(pid):
                    os.kill(pid, signal.SIGKILL)

                if os.path.isfile(pid_file):
                    os.unlink(pid_file)
        elif opts.daemon == 'status':
            if os.path.isfile(pid_file):
                pid = int(open(pid_file).read())
                if pid_running(pid):
                    print "running"
                    return
            print "stopped"
        else:
            raise ValueError("Unknown daemon operation {}".format(opts.daemon))
    else:
        required_sensors = json.loads(opts.sensors_config.read())
        daemon_main(required_sensors, opts)
    return 0

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