subprocess.Popen execute Command / Invoke function timeout control.

command-func-timeout.py

#coding: utf-8

import os
import time
import signal
# from datetime import datetime
from multiprocessing import Process, Queue
import subprocess
import threading

import functools


class TimeoutError(Exception): pass


class SnoozeAlarm(threading.Thread):
  def __init__(self, zzz):
    threading.Thread.__init__(self)
    self.setDaemon(True)
    self.zzz = zzz

  def run(self):
    time.sleep(self.zzz)
    os.kill(os.getpid(), signal.SIGALRM)
    

def timeout(seconds, error_message = 'Function call timed out'):
    def decorated(func):
        def _handle_timeout(signum, frame):
            raise TimeoutError(error_message)

        def wrapper(*args, **kwargs):
            signal.signal(signal.SIGALRM, _handle_timeout)
            SnoozeAlarm(seconds).start()
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)
            return result

        return functools.wraps(func)(wrapper)

    return decorated


def timeout_run(func, seconds, msg):
    d_func = timeout(seconds, msg)(func)
    try:
        rv = d_func()
        ok = True
    except TimeoutError:
        ok, rv = False, None
        
    return ok, rv
    

# !!!Unused, may unsafe!    
class Func(object):
    def __init__(self, _worker):
        
        def worker(q):
            q.put(_worker())

        self._worker = _worker
        self.worker = worker
        self.process = None
        self.rv = None
        self.queue = Queue()
        self.ok = False


    def run(self, timeout):
        
        def target():
            self.process = Process(target=self.worker, args=(self.queue,))
            self.process.start()
            self.process.join()
        
        thread = threading.Thread(target=target)
        thread.start()
        thread.join(timeout)
        
        if thread.is_alive():
            if self.process:
                self.process.terminate()
                os.kill(self.process.pid, signal.SIGKILL)
                self.process.join()
                self.queue.close()
            thread.join()
        else:
            self.rv = self.queue.get()
            self.ok = True
            
        return self.ok, self.rv


class Command(object):
    
    KILL = -9
    TERM = -15
    
    def __init__(self, cmd):
        self.cmd = cmd.replace('\n', ' ')

        self.process = None     # Timeout too fast
        self.returncode = None
        self.output = None
        self.error = None


    def run(self, timeout, sig=KILL):
        def target():
            self.process = subprocess.Popen(self.cmd, shell=True,
                                            stdout=subprocess.PIPE,
                                            stderr=subprocess.PIPE)
            self.process.wait()
            self.returncode = self.process.returncode

        thread = threading.Thread(target=target)
        thread.start()
        thread.join(timeout)

        if thread.is_alive():
            if self.process:
                try:
                    if sig == Command.KILL:
                        self.process.kill()
                    elif sig == Command.TERM:
                        self.process.terminate()
                except OSError:
                    pass
            thread.join(0.3)
        else:
            self.output, self.error = self.process.communicate() # target() block here
        
        return self.returncode, self.output, self.error