nilmrun/nilmrun/processmanager.py

#!/usr/bin/python

from nilmdb.utils.printf import *

import threading
import multiprocessing
import cStringIO
import sys
import os
import time
import uuid

class LogReceiver(object):
    """Spawn a thread that listens to a pipe for log messages,
    and stores them locally."""
    def __init__(self, pipe):
        self.pipe = pipe
        self.log = cStringIO.StringIO()
        self.thread = threading.Thread(target = self.run)
        self.thread.start()

    def run(self):
        try:
            while True:
                data = self.pipe.recv_bytes()
                self.log.write(data)
        except EOFError:
            self.pipe.close()
            return

    def getvalue(self):
        return self.log.getvalue()

    def clear(self):
        self.log = cStringIO.StringIO()

class LogSender(object): # pragma: no cover (runs in a different process)
    """File-like object that writes output to a pipe"""
    def __init__(self, pipe):
        self.pipe = pipe

    def close(self):
        if self.pipe:
            self.pipe.close()
        self.pipe = None

    def write(self, data):
        if self.pipe:
            self.pipe.send_bytes(data)

    def flush(self):
        pass

    def isatty(self):
        return False

class Process(object):
    """Spawn and manage a running process"""
    def __init__(self, name, function, parameters):
        self.parameters = parameters
        self.start_time = None
        self.name = name
        pipes = multiprocessing.Pipe(duplex = False)
        self._log = LogReceiver(pipes[0])
        self._process = multiprocessing.Process(
            target = self._tramp, name = name,
            args = (function, pipes, parameters))
        self._process.daemon = True
        self._process.start()
        pipes[1].close()
        self.start_time = time.time()
        self.pid = str(uuid.uuid1(self._process.pid or 0))

    def _tramp(self, function, pipes, parameters): # pragma: no cover
        # Remap stdio in the child before calling function
        pipes[0].close()
        logfile = LogSender(pipes[1])
        sys.stdin = open(os.devnull, 'r')
        sys.stdout = logfile
        sys.stderr = logfile
        function(parameters)

    def terminate(self, timeout = 1):
        self._process.join(timeout)
        if self.alive:
            self._process.terminate()
        self._process.join(timeout)
        if self.alive:
            return False
        return True

    def clear_log(self):
        self._log.clear()

    @property
    def log(self):
        return self._log.getvalue()

    @property
    def alive(self):
        return self._process.is_alive()

    @property
    def exitcode(self):
        return self._process.exitcode

class ProcessManager(object):
    def __init__(self):
        self.processes = {}

    def __iter__(self):
        return iter(self.processes.keys())

    def __getitem__(self, key):
        return self.processes[key]

    def run(self, name, function, parameters):
        new = Process(name, function, parameters)
        self.processes[new.pid] = new
        return new.pid

    def terminate(self, pid):
        return self.processes[pid].terminate()

    def remove(self, pid):
        del self.processes[pid]