Python中监视线程卡死并自动崩溃退出 WatchDog

监视主线程卡死，卡死时保错再主动崩溃退出，同时打印当时的调用堆栈。
还可以再外部结合监视崩溃自动重启就可以在卡死时实现自动重启。

import threading
import traceback
import time
import sys
import os
from functools import wraps

class WatchDog(threading.Thread):

    def __init__(self, timeout=10, echo=False):
        super(WatchDog, self).__init__()
        self.timeout = timeout
        self.echo = echo
        self.last_kicked_ts = time.time()
        self.lock = threading.Lock()
        self.thread_id = threading.currentThread().ident
        self.terminated = False
        self.setDaemon(True)
        self.start()

    def terminate(self):
        self.terminated = True
        self.join(self.timeout)

    def kick(self):
        self.lock.acquire()
        self.last_kicked_ts = time.time()
        self.lock.release()

    def bark(self):
        formated_frame_stack = self._get_formated_frame_stack()
        if self.echo:
            print("!!!!! WATCH DOG FAILURE TRIGGERED !!!!!\n" + formated_frame_stack, flush=True)
        pid = os.getpid()
        os.kill(pid, 2) # 通知进程退出
        time.sleep(5)   # 等待5秒
        os.kill(pid, 9) # 发送强制退出

    def run(self):
        while not self.terminated:
            ts = time.time()
            self.lock.acquire()
            is_timeout = ts - self.last_kicked_ts > self.timeout
            self.lock.release()
            if is_timeout:
                self.bark()
            n = int(max(self.timeout / 3, 1))
            for i in range(n*10):
                time.sleep(0.1)
                if self.terminated:
                    break

    @staticmethod
    def _get_thread(tid):
        for t in threading.enumerate():
            if t.ident == tid:
                return t
        return None

    @staticmethod
    def _get_frame_stack(tid):
        for thread_id, stack in sys._current_frames().items():
            if thread_id == tid:
                return stack
        return None

    def _get_formated_frame_stack(self):
        info = []
        th = self._get_thread(self.thread_id)
        stack = self._get_frame_stack(self.thread_id)
        info.append('%s thead_id=%d' % (th.name, self.thread_id))
        for filename, lineno, _, line in traceback.extract_stack(stack):
            info.append('    at %s(%s:%d)' % (line, filename[filename.rfind(os.path.sep) + 1:], lineno))
        return '\n'.join(info)


def watch_dog(timeout=10, echo=False):
    def inner(func):
        def wrapper(*args, **kw):
            dog = WatchDog(timeout=timeout, echo=echo)
            ret = func(*args, **kw)
            dog.terminate()
            return ret
        return wrapper
    return inner

用例1，监控函数超时

@watch_dog(timeout=3, echo=True)
def func():
    # do something
    time.sleep(5)

def main():
    func()

if __name__ == '__main__':
    main()

我们监视func，设置其3秒超时，然后在里面sleep 5秒来引起超时，可以看到输出如下

!!!!! WATCH DOG FAILURE TRIGGERED !!!!!
MainThread thead_id=2164
    at main()(watch_dog.py:97)
    at func()(watch_dog.py:94)
    at ret = func(*args, **kw)(watch_dog.py:81)
    at time.sleep(5)(watch_dog.py:91)
[Finished in 3.5s with exit code 2]

在程序启动3.5秒后就主动崩溃了

用例2，监控循环超时

def main():
    watch_dog = WatchDog(timeout=3, echo=True)

    # main loop
    # while True:
    for i in range(0, 10):
        time.sleep(i)
        print(f"i={i}")
        watch_dog.kick()
    watch_dog.terminated()


if __name__ == '__main__':
    main()

输出如下

i=0
i=1
i=2
i=3
!!!!! WATCH DOG FAILURE TRIGGERED !!!!!
MainThread thead_id=9060
    at main()(watch_dog.py:101)
    at time.sleep(i)(watch_dog.py:94)
[Finished in 10.1s with exit code 2]

用例3，超时自动重试

import multiprocessing

def main():
    ...

if __name__ == '__main__':
    RETRY_TIMES = 5
    for i in range(RETRY_TIMES): # 重试次数
        child_process = multiprocessing.Process(target=main)
        child_process.start()
        child_process.join()

        if child_process.exitcode == 0:
            print("子进程正常结束")
            exit(0)
        elif i < RETRY_TIMES - 1:
            print("子进程异常结束，即将重试")
        else:
            print("子进程异常结束，重试超次")
            exit(child_process.exitcode)