Timeout on a function call

后端 未结 18 1845
挽巷
挽巷 2020-11-21 04:53

I\'m calling a function in Python which I know may stall and force me to restart the script.

How do I call the function or what do I wrap it in so that if it takes

18条回答
  •  佛祖请我去吃肉
    2020-11-21 05:08

    Here is a POSIX version that combines many of the previous answers to deliver following features:

    1. Subprocesses blocking the execution.
    2. Usage of the timeout function on class member functions.
    3. Strict requirement on time-to-terminate.

    Here is the code and some test cases:

    import threading
    import signal
    import os
    import time
    
    class TerminateExecution(Exception):
        """
        Exception to indicate that execution has exceeded the preset running time.
        """
    
    
    def quit_function(pid):
        # Killing all subprocesses
        os.setpgrp()
        os.killpg(0, signal.SIGTERM)
    
        # Killing the main thread
        os.kill(pid, signal.SIGTERM)
    
    
    def handle_term(signum, frame):
        raise TerminateExecution()
    
    
    def invoke_with_timeout(timeout, fn, *args, **kwargs):
        # Setting a sigterm handler and initiating a timer
        old_handler = signal.signal(signal.SIGTERM, handle_term)
        timer = threading.Timer(timeout, quit_function, args=[os.getpid()])
        terminate = False
    
        # Executing the function
        timer.start()
        try:
            result = fn(*args, **kwargs)
        except TerminateExecution:
            terminate = True
        finally:
            # Restoring original handler and cancel timer
            signal.signal(signal.SIGTERM, old_handler)
            timer.cancel()
    
        if terminate:
            raise BaseException("xxx")
    
        return result
    
    ### Test cases
    def countdown(n):
        print('countdown started', flush=True)
        for i in range(n, -1, -1):
            print(i, end=', ', flush=True)
            time.sleep(1)
        print('countdown finished')
        return 1337
    
    
    def really_long_function():
        time.sleep(10)
    
    
    def really_long_function2():
        os.system("sleep 787")
    
    
    # Checking that we can run a function as expected.
    assert invoke_with_timeout(3, countdown, 1) == 1337
    
    # Testing various scenarios
    t1 = time.time()
    try:
        print(invoke_with_timeout(1, countdown, 3))
        assert(False)
    except BaseException:
        assert(time.time() - t1 < 1.1)
        print("All good", time.time() - t1)
    
    t1 = time.time()
    try:
        print(invoke_with_timeout(1, really_long_function2))
        assert(False)
    except BaseException:
        assert(time.time() - t1 < 1.1)
        print("All good", time.time() - t1)
    
    
    t1 = time.time()
    try:
        print(invoke_with_timeout(1, really_long_function))
        assert(False)
    except BaseException:
        assert(time.time() - t1 < 1.1)
        print("All good", time.time() - t1)
    
    # Checking that classes are referenced and not
    # copied (as would be the case with multiprocessing)
    
    
    class X:
        def __init__(self):
            self.value = 0
    
        def set(self, v):
            self.value = v
    
    
    x = X()
    invoke_with_timeout(2, x.set, 9)
    assert x.value == 9
    

提交回复
热议问题