魔のkyo的BLOG

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

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

def main():
    func()

if __name__ == '__main__':
    main()

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()

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)

# -*- coding: utf-8 -*-
import threading
import time
import queue
from dataclasses import dataclass, field
from typing import Any
import signal
import traceback
import os


@dataclass(order=True)
class ScheduledItem:
    time_sec: float
    cb: Any = field(compare=False)


class Driver:
    def __init__(self):
        self._scheduled = queue.PriorityQueue()
        self._scheduled_every_lock = threading.Lock()
        self._scheduled_every = []
        self._callbacks_lock = threading.Lock()
        self._callbacks = {}
        self._async_queue = queue.Queue()
        self._epoch_sec = time.time()
        self._debug_epoch = None
        self._last_epoch = None
        self._sigterm_received = False

    # 得到driver内的当前时间
    def get_epoch(self):
        return self._epoch_sec

    # 设置调试时间，设置了调试时间driver内将不更新当前时间而采用调试时间
    def set_debug_epoch(self, epoch):
        self._debug_epoch = epoch

    # 主循环
    # wait_sync_interval 每趟循环的时间间隔，如果为0相当于占用一个core
    # intercept_signals  拦截并处理信号，默认处理CTRL+C和kill，收到时退出主循环
    def run(self, wait_sync_interval=0, intercept_signals=(signal.SIGINT, signal.SIGTERM)):
        if intercept_signals:
            self.intercept_sigterm(intercept_signals)
        while not self.sigterm_received():
            self.run_step(wait_sync_interval)

    # 执行一趟主逻辑，一般放在主循环中执行
    def run_step(self, wait_sync_interval=0):
        if self._debug_epoch:
            self._epoch_sec = self._debug_epoch
        else:
            self._epoch_sec = time.time()
            if self._last_epoch is not None:
                if self._epoch_sec - self._last_epoch < wait_sync_interval:
                    t = wait_sync_interval - (self._epoch_sec - self._last_epoch)
                    time.sleep(t)
                    self._epoch_sec = time.time()
        self._last_epoch = self._epoch_sec

        self._do_async()
        self._do_schedule()
        self._do_schedule_every()
        
    # 计划单次定时任务
    def schedule(self, cb, time_sec):
        self._scheduled.put_nowait( ScheduledItem(time_sec, cb) )

    # 计划重复任务
    def schedule_every(self, cb, interval_sec):
        self._scheduled_every_lock.acquire()
        self._scheduled_every.append( { "next_sec":self._epoch_sec+interval_sec, "interval":interval_sec, "cb":cb } )
        self._scheduled_every_lock.release()

    # 增加消息接收者
    def add_receiver(self, topic_or_type, cb):
        self._callbacks_lock.acquire()
        if topic_or_type not in self._callbacks:
            self._callbacks[topic_or_type] = set()
        self._callbacks[topic_or_type].add(cb)
        self._callbacks_lock.release()
        return cb

    # 删除消息接收者
    def remove_receiver(self, topic_or_type, cb):
        self._callbacks_lock.acquire()
        if topic_or_type in self._callbacks:
            if cb in self._callbacks[topic_or_type]:
                self._callbacks[topic_or_type].remove(cb)
        self._callbacks_lock.release()

    # 同步发送消息
    def send(self, obj, topic=None):
        if topic == None:
            topic = type(obj)
        cbs = []
        self._callbacks_lock.acquire()
        if topic in self._callbacks.keys():
            cbs = list(self._callbacks[topic])
        self._callbacks_lock.release()
        for cb in cbs:
            cb(obj)

    # 异步发送消息
    def send_async(self, obj, topic=None):
        self._async_queue.put_nowait( (obj, topic) )

    def _do_async(self):
        while not self._async_queue.empty():
            self.send(*self._async_queue.get_nowait())

    def _do_schedule(self):
        i = 0
        while not self._scheduled.empty():
            item = self._scheduled.get_nowait()
            if item.time_sec > self._epoch_sec:
                self._scheduled.put_nowait(item)
                break
            item.cb(self._epoch_sec)

    def _do_schedule_every(self):
        cbs = []
        self._scheduled_every_lock.acquire()
        for o in self._scheduled_every:
            while self._epoch_sec >= o["next_sec"]:
                cbs.append(o["cb"])
                o["next_sec"] += o["interval"]
        self._scheduled_every_lock.release()
        for cb in cbs:
            cb(self._epoch_sec)

    # 拦截终止信号
    def intercept_sigterm(self, intercept_signals):
        for sig in intercept_signals:
            signal.signal(sig, bind(Driver._on_signal, self))

    # 是否收到了终止信号
    def sigterm_received(self):
        return self._sigterm_received

    def _on_signal(self, signum, stack):
        print(f"Signal #{signum} received, Traceback:")
        for filename, lineno, _, line in traceback.extract_stack(stack):
            print('    at %s(%s:%d)' % (line, filename[filename.rfind(os.path.sep) + 1:], lineno))
        self._sigterm_received = True


def bind(func, *args, **kw):
    return lambda *_args, **_kw: func(*args, *_args, **kw, **_kw)


driver = Driver()


if __name__ == '__main__':
    print("initialize")

    driver.schedule(lambda epoch: print(f"Do sth. one time"), driver.get_epoch()+3)
    driver.schedule_every(lambda epoch: print(f"Do sth. regularly"), 1)
    driver.schedule(lambda _: os.system(f"kill {os.getpid()}"), driver.get_epoch()+10)
    driver.run(0.1)

    print("finalize")

# -*- coding: utf-8 -*-
from redis import Redis
import socket
import os
import threading
import time
import uuid
from contextlib import ExitStack, contextmanager

# 加锁时间片
# 对于超过一个加锁时间片的锁应该重复使用expire设置过期时间
# 为的是在程序异常结束时锁能被尽快自动释放，所以时间片不应太长，
# 也不应设置过短，不然一是会频繁的访问redis续过期时长导致资源浪费，
# 二是在很卡的时候负责续时的线程可能无法正常续时。
# 建议设置在5~20秒
LOCK_TIME_SLICE = 10


class KeepLockThread(threading.Thread):

    def __init__(self, redis, key_name):
        super().__init__()
        self.redis = redis
        self.key_name = key_name

        self.terminated = False
        self.lock = threading.Lock()

    def run(self):
        while True:
            time.sleep(LOCK_TIME_SLICE//3+1)

            with self.lock:
                if not self.terminated:
                    self.redis.expire(self.key_name, LOCK_TIME_SLICE)
                else:
                    return

    def stop(self):
        with self.lock:
            self.terminated = True


# 互斥锁
class Mutex:
    def __init__(self, name, server="127.0.0.1"):
        self.name = name
        self.key_name = "MUTEX_" + name
        hostname = socket.gethostname()
        pid = os.getpid()
        tid = threading.get_ident()
        # 得到唯一ID，主机名 PID TID 是为了调试用
        self.id = f"{hostname}_{pid}_{tid}_{uuid.uuid4().hex}"
        self.redis = Redis(host=server)
        self.thread = None

    # blocking: 是否阻塞
    # shortlived: 是否短暂的加锁，如果是短暂的则不会创建额外的线程维护过期时间，所谓短暂就是确定加锁时间小于LOCK_TIME_SLICE
    #             加锁时间是从获得锁到释放锁的时间，不包括阻塞等待的时间，可以阻塞等待很长时间仍然是短暂的加锁。
    def acquire(self, blocking=True, shortlived=False):
        r = self.redis.set(self.key_name, self.id, ex=LOCK_TIME_SLICE, nx=True)
        if blocking:
            while not r:
                time.sleep(0.05)
                r = self.redis.set(self.key_name, self.id, ex=LOCK_TIME_SLICE, nx=True)
        if r:
            if not shortlived:
                # 如果加锁成功了，那么应该创建线程维护过期时间
                self.thread = KeepLockThread(self.redis, self.key_name)
                self.thread.daemon = True
                self.thread.start()
        return r

    def release(self):
        if self.acquired():
            if self.thread:
                self.thread.stop()
                self.thread = None
            self.redis.delete(self.key_name)

    def acquired(self):
        r = self.redis.get(self.key_name)
        if r is not None and r.decode() == str(self.id):
            return True
        else:
            return False

    def __enter__(self):
        self.acquire()

    def __exit__(self, exc_type, exc_value, exc_trackback):
        self.release()
        if exc_value is not None:
            raise exc_value


def mutex(names, server="127.0.0.1"):
    def medium(func):
        def wrapper(*args, **kw):
            if type(names) == list:
                with ExitStack() as stack:
                    for name in names:
                        stack.enter_context(Mutex(name))
                    return func(*args, **kw)
            else:
                with Mutex(names, server):
                    return func(*args, **kw)
        wrapper.__name__ = func.__name__
        wrapper.__doc__ = func.__doc__
        return wrapper
    return medium


class KeepReadLockThread(threading.Thread):

    def __init__(self, redis, lock_id, rlock_name):
        super().__init__()
        self.redis = redis
        self.id = lock_id
        self.rlock_name = rlock_name

        self.terminated = False
        self.lock = threading.Lock()

    def run(self):
        while True:
            time.sleep(LOCK_TIME_SLICE//3+1)

            with self.lock:
                if not self.terminated:
                    pipeline = self.redis.pipeline()
                    now = time.time()
                    pipeline.zadd(self.rlock_name, {self.id: int(now*1000)} )
                    pipeline.expire(self.rlock_name, LOCK_TIME_SLICE)
                    pipeline.execute()
                else:
                    return

    def stop(self):
        with self.lock:
            self.terminated = True


# 读写锁
class ReadWriteLock:

    def __init__(self, name, server="127.0.0.1"):
        self.name = name
        self.server = server
        self.rlock_name = "RLOCK_" + name
        self.wlock_name = "WLOCK_" + name
        self.meta_lock_name = "META_" + name

        hostname = socket.gethostname()
        pid = os.getpid()
        tid = threading.get_ident()
        # 得到唯一ID，主机名 PID TID 是为了调试用
        self.id = f"{hostname}_{pid}_{tid}_{uuid.uuid4().hex}"

        self.redis = Redis(host=server)
        self.lock_type = None
        self.thread = None

    def acquire_read_lock(self, blocking=True, shortlived=False):
        # 这个mutex用来维护读写锁内部数据本身使用，在离开ReadWriteLock函数前必须释放
        mutex = Mutex(self.meta_lock_name, self.server)
        try:
            mutex.acquire(shortlived=True)
            wlock_locked = self.redis.get(self.wlock_name)
            if wlock_locked:
                if blocking:
                    while wlock_locked:
                        mutex.release()
                        time.sleep(0.05)
                        mutex.acquire(shortlived=True)
                        wlock_locked = self.redis.get(self.wlock_name)
                else:
                    return False
            
            pipeline = self.redis.pipeline()
            now = time.time()
            # 移除一定时间前失效的读锁（主动处理异常未释放的读锁）
            pipeline.zremrangebyscore(self.rlock_name, 0, int((now-LOCK_TIME_SLICE)*1000) )
            # 添加新的读锁
            pipeline.zadd(self.rlock_name, {self.id: int(now*1000)} )
            pipeline.expire(self.rlock_name, LOCK_TIME_SLICE)
            pipeline.execute()
            self.lock_type = 'R'

            if not shortlived:
                # 创建线程维护读锁过期时间
                self.thread = KeepReadLockThread(self.redis, self.id, self.rlock_name)
                self.thread.daemon = True
                self.thread.start()

            return True
        finally:
            mutex.release()
    
    def acquire_write_lock(self, blocking=True, shortlived=False):
        mutex = Mutex(self.meta_lock_name, self.server)
        try:
            # 注意这里使用mutex时shortlived始终为True，它不是由外部使用读写锁是否短暂决定
            # 而是加读写锁时内部需要用到mutex进行短暂加锁，以保证获取到读锁为0后到加写锁前读锁不会新增
            mutex.acquire(shortlived=True)
            # 移除一定时间前失效的读锁（主动处理异常未释放的读锁）
            self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-LOCK_TIME_SLICE)*1000) )
            # 获取当前读锁数量
            r = self.redis.zcard(self.rlock_name)
            if r:
                if blocking:
                    while r:
                        mutex.release()
                        time.sleep(0.05)
                        mutex.acquire(shortlived=True)
                        # 移除一定时间前失效的读锁（主动处理异常未释放的读锁）
                        self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-LOCK_TIME_SLICE)*1000) )
                        # 获取当前读锁数量
                        r = self.redis.zcard(self.rlock_name)
                else:
                    return False
            # 尝试获取写锁
            r = self.redis.set(self.wlock_name, self.id, ex=LOCK_TIME_SLICE, nx=True)
            if blocking:
                while not r:
                    mutex.release()
                    time.sleep(0.05)
                    # 没有获取到写锁，那么可能读锁又会被其他地方获取，所以要重新先等到读锁为0
                    mutex.acquire(shortlived=True)
                    self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-LOCK_TIME_SLICE)*1000) )
                    r = self.redis.zcard(self.rlock_name)
                    if r:
                        if blocking:
                            while r:
                                mutex.release()
                                time.sleep(0.05)
                                mutex.acquire(shortlived=True)
                                self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-LOCK_TIME_SLICE)*1000) )
                                r = self.redis.zcard(self.rlock_name)
                        else:
                            return False
                    r = self.redis.set(self.wlock_name, self.id, ex=LOCK_TIME_SLICE, nx=True)
            if r:
                if not shortlived:
                    self.thread = KeepLockThread(self.redis, self.wlock_name)
                    self.thread.daemon = True
                    self.thread.start()

                self.lock_type = 'W'
            return r
        finally:
            mutex.release()

    def acquired(self, lock_type = "any"):
        if lock_type == "any":
            if self.lock_type is not None:
                return True
            else:
                return False
        else:
            if self.lock_type == lock_type:
                return True
            else:
                return False

    def release(self):
        if self.thread:
            self.thread.stop()
            self.thread = None
        if self.lock_type == 'R':
            self.redis.zrem(self.rlock_name, self.id)
            self.lock_type = None
        elif self.lock_type == 'W':
            r = self.redis.get(self.wlock_name)
            if r is not None and r.decode() == str(self.id):
                self.redis.delete(self.wlock_name)
            self.lock_type = None


@contextmanager
def read_lock(name, server="127.0.0.1"):
    rwlock = ReadWriteLock(name=name)
    try:
        rwlock.acquire_read_lock()
        yield rwlock
    finally:
        rwlock.release()


@contextmanager
def write_lock(name, server="127.0.0.1"):
    rwlock = ReadWriteLock(name=name)
    try:
        rwlock.acquire_write_lock()
        yield rwlock
    finally:
        rwlock.release()