魔の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 threading
import time
import uuid

# 互斥锁
class Mutex:
    def __init__(self, name, server="127.0.0.1"):
        self.name = name
        self.key_name = "MUTEX_" + name
        self.id = uuid.uuid4().hex
        self.redis = Redis(host=server)

    def acquire(self, blocking=True, ex=120):
        r = self.redis.set(self.key_name, self.id, ex=ex, nx=True)
        if blocking:
            while not r:
                time.sleep(0.01)
                r = self.redis.set(self.key_name, self.id, ex=ex, nx=True)
        return r

    def release(self):
        if self.acquired():
            self.redis.delete(self.key_name)

    def acquired(self):
        r = self.redis.get(self.key_name)
        return r != None and r.decode() == str(self.id)

    def __enter__(self):
        self.acquire()

    def __exit__(self, exc_type, exc_value, exc_trackback):
        self.release()
        if exc_value != None:
            raise exc_value


# 读写锁
class ReadWriteLock:
    MAX_READ_TIME = 120 # 最长加读锁时间（秒）
    '''
    正常不会用到MAX_READ_TIME，在加锁后异常并且没有对解锁进行正确处理时，会导致锁始终被一个已崩溃的进程持有无法释放，
    为了应对这种情况，通过设置MAX_READ_TIME，在加读锁时间超过MAX_READ_TIME没有解锁时自动解锁
    '''

    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
        self.id = uuid.uuid4().hex
        self.redis = Redis(host=server)
        self.lock_type = None

    def acquire_read_lock(self, blocking=True):
        # 这个mutex用来维护读写锁内部数据本身使用，在离开ReadWriteLock函数前必须释放
        mutex = Mutex(self.meta_lock_name, self.server)
        try:
            mutex.acquire()
            wlock_locked = self.redis.get(self.wlock_name)
            if wlock_locked:
                if blocking:
                    while wlock_locked:
                        mutex.release()
                        time.sleep(0.05)
                        mutex.acquire()
                        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-self.MAX_READ_TIME)*1000) )
            # 添加新的读锁
            pipeline.zadd(self.rlock_name, {self.id: int(now*1000)} )
            pipeline.expire(self.rlock_name, self.MAX_READ_TIME)
            pipeline.execute()
            self.lock_type = 'R'
            return True
        finally:
            mutex.release()
    
    def acquire_write_lock(self, blocking=True, ex=120):
        mutex = Mutex(self.meta_lock_name, self.server)
        try:
            mutex.acquire()
            # 移除一定时间前失效的读锁（主动处理异常未释放的读锁）
            self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-self.MAX_READ_TIME)*1000) )
            # 获取当前读锁数量
            r = self.redis.zcard(self.rlock_name)
            if r:
                if blocking:
                    while r:
                        mutex.release()
                        time.sleep(0.05)
                        mutex.acquire()
                        # 移除一定时间前失效的读锁（主动处理异常未释放的读锁）
                        self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-self.MAX_READ_TIME)*1000) )
                        # 获取当前读锁数量
                        r = self.redis.zcard(self.rlock_name)
                else:
                    return False
            # 尝试获取写锁
            r = self.redis.set(self.wlock_name, self.id, ex=ex, nx=True)
            if blocking:
                while not r:
                    mutex.release()
                    time.sleep(0.05)
                    mutex.acquire()
                    self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-self.MAX_READ_TIME)*1000) )
                    r = self.redis.zcard(self.rlock_name)
                    if r:
                        if blocking:
                            while r:
                                mutex.release()
                                time.sleep(0.05)
                                mutex.acquire()
                                self.redis.zremrangebyscore(self.rlock_name, 0, int((time.time()-self.MAX_READ_TIME)*1000) )
                                r = self.redis.zcard(self.rlock_name)
                        else:
                            return False
                    r = self.redis.set(self.wlock_name, self.id, ex=ex, nx=True)
            if r:
                self.lock_type = 'W'
            return r
        finally:
            mutex.release()

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

    def release(self):
        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 != None and r.decode() == str(self.id):
                self.redis.delete(self.wlock_name)
            self.lock_type = None

# -*- coding: utf-8 -*-
from sqlalchemy.ext.declarative import declarative_base, declared_attr
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager
from urllib import parse
import re
import datetime
import types
import decimal


class ModelBase(object):

    @declared_attr
    def __tablename__(cls):
        # return cls.__name__.lower()
        return re.sub(r'([A-Z])', r'_\1', cls.__name__[0].lower()+cls.__name__[1:]).lower()

    @classmethod
    def props(cls):
        if cls.__base__.__name__ == "Base":
            return [c for c in cls.__table__.columns]
        elif cls.__base__.__base__.__name__ == "Base":
            super_column_names = [c for c in cls.__base__.__table__.columns]
            column_names = [c for c in cls.__table__.columns if c.name != 'id']
            return super_column_names + column_names
        else:
            assert(False and "多层继承后的props获取暂未实现")

    @classmethod
    def prop_names(cls):
        if cls.__base__.__name__ == "Base":
            return [c.name for c in cls.__table__.columns]
        elif cls.__base__.__base__.__name__ == "Base":
            super_column_names = [c.name for c in cls.__base__.__table__.columns]
            column_names = [c.name for c in cls.__table__.columns if c.name != 'id']
            return super_column_names + column_names
        else:
            assert(False and "多层继承后的prop_names获取暂未实现")

    @classmethod
    def prop(cls, prop_name):
        return cls.__table__.columns[prop_name]

    # 对象在输出时如何序列化成字符串
    def __repr__(self):
        attrs = []
        # for c in self.__table__.columns:
        for name in self.prop_names():
            attr = getattr(self, name)
            if type(attr) in (str, datetime.date, datetime.time, datetime.datetime):
                attrs.append(f"{name}='{attr}'")
            else:
                attrs.append(f"{name}={attr}")
        return f"{self.__class__.__name__}({', '.join(attrs)})"

    # 转换成dict时的键列表
    def keys(self):
        return self.prop_names()

    def __getitem__(self, item):
        return getattr(self, item)

    def __setitem__(self, item, value):
        return setattr(self, item, value)

    def to_dict(self):
        return {name: to_jsonable(getattr(self, name)) for name in self.keys()}


def to_dict(db_obj):
    if isinstance(db_obj, ModelBase):
        return db_obj.to_dict()
    else:
        return db_obj

def to_list(db_objs):
    return [to_dict(db_obj) for db_obj in db_objs]


# Base = declarative_base(cls=ModelBase)


# 把ORM对象转成可序列化成JSON的对象，对于ORM对象的list转换为dict的list，对于ORM对象转换成dict
def to_jsonable(o):
    if isinstance(o, list):
        return [to_jsonable(e) for e in o]
    if isinstance(o, dict):
        return {k: to_jsonable(v) for (k,v) in o.items()}
    if isinstance(o, ModelBase):
        return o.to_dict()
    if isinstance(o, datetime.datetime):
        return o.strftime('%Y-%m-%d %H:%M:%S')
    if isinstance(o, datetime.date):
        return o.strftime('%Y-%m-%d')
    if isinstance(o, datetime.time):
        return o.strftime('%H:%M:%S')
    if isinstance(o, decimal.Decimal):
        return float(o)
    return o


engines = {}

r"""
不同数据库URL样例
r"sqlite:///:memory:"  sqlite内存数据库
r"sqlite:///C:\path\foo.db" sqlite在windows上
r"sqlite:////root/data/foo.db" sqlite在linux上
r"mysql+pymysql://username:password@server/dbname" MySQL
r"postgresql+psycopg2://username:password@server/dbname" Postgresql
r"mssql+pymssql://username:password@server/dbname?charset=utf8" Microsoft SQL Server
"""
def init_engine(
        url=None, name="main",
        dialect=None, username="", password="", server="", dbname=None,
        **kwargs):
    if dialect:
        if dialect.lower() in ("mysql", ):
            dialect = "mysql+pymysql"
        elif dialect.lower() in ("postgresql", "postgres", "pgsql"):
            dialect = "postgresql+psycopg2"
        elif dialect.lower() in ("mssql", "sqlserver"):
            dialect = "mssql+pymssql"
    if url is None:
        if password:
            url = '{}://{}:{}@{}/{}'.format(dialect, username, parse.quote_plus(password), server, dbname)
        elif username:
            url = '{}://{}@{}/{}'.format(dialect, username, server, dbname)
        else:
            url = '{}://{}/{}'.format(dialect, server, dbname)

    # # 设置连接池默认大小
    # if "pool_size" not in kwargs:
    #     kwargs["pool_size"] = 5

    # # 设置连接默认回收时间(MySQL默认8小时)
    # if "pool_recycle" not in kwargs:
    #     kwargs["pool_recycle"] = 28000

    # 默认启用pool_pre_ping
    if "pool_pre_ping" not in kwargs:
        kwargs["pool_pre_ping"] = True

    engine = create_engine(url, **kwargs)
    session_maker = sessionmaker(expire_on_commit=False)
    session_maker.configure(bind=engine)

    @contextmanager
    def _session_scope(maker):
        session = maker()
        mutexes = []
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()
    engine.session_scope = types.MethodType(_session_scope, session_maker)
    engines[name] = engine
    return engine


def exit_engine(name="main"):
    if name in engines:
        del engines[name]


def get_engine(name="main"):
    if name in engines:
        return engines[name]
    else:
        return None


def session_scope(name="main"):
    if name in engines:
        return engines[name].session_scope()
    else:
        raise Exception("engine未初始化")

# -*- coding: utf-8 -*-
from sqlalchemy import Table, Column, Integer, Float, Numeric, String, DateTime, Date, Time, Boolean
from sqlalchemy import PrimaryKeyConstraint, ForeignKey, Index, UniqueConstraint
from sqlalchemy import text, func, and_, or_, not_, asc, inspect, desc, distinct
from sqlalchemy.orm import relationship, backref
from sqlalchemy.ext.declarative import declarative_base

from utils.db_ import ModelBase, init_engine, exit_engine, get_engine, session_scope, to_jsonable

Base = declarative_base(cls=ModelBase)


class Clazz(Base):
    id = Column(Integer(), primary_key=True, autoincrement=True)
    enroll_year = Column(Integer(), comment="入学年份")


class Student(Base):
    __table_args__ = (
        UniqueConstraint('clazz_id', 'name'),  # 创建多字段唯一约束
        Index('ix_clazz_id_stu_no', 'clazz_id', 'stu_no'),  # 创建多字段索引
        )
    id = Column(Integer(), primary_key=True, autoincrement=True) # 自增ID作为单字段主键，多字段主键也可以在多字段上都写上primary_key=True
    stu_no = Column(String(), index=True, unique=True, comment="学号") # 创建单字段索引，创建单字段唯一约束
    name = Column(String(), nullable=True, comment="姓名") # 创建非空约束
    clazz_id = Column(Integer(), ForeignKey('clazz.id', ondelete='RESTRICT'), nullable=True, comment="班级ID") # 创建外键

    clazz = relationship('Clazz')  # 多对一（一对多）关系映射


class Course(Base):
    id = Column(Integer(), primary_key=True, autoincrement=True)
    name = Column(String(), comment="课程名")

    students = relationship('Student', secondary='student_course', backref=backref('courses'))  # 多对多映射关系,backref相当于给Student对象加了courses属性

    def keys(self):
        return super().keys() + ["students",]


# 选课（多对多关系表）
student_course = Table('student_course', Base.metadata,
    Column('student_id', Integer(), ForeignKey('student.id'), primary_key=True),
    Column('course_id', Integer(), ForeignKey('course.id'), primary_key=True))


def main():
    engine = init_engine(url=r"sqlite:///:memory:")
    Base.metadata.create_all(engine)  # 根据定义创建出所有表
    with session_scope() as session:
        clazz = Clazz(enroll_year=2022) # 创建一个clazz对象
        session.add(clazz) # 用clazz对象创建一行记录
        session.flush() # 立即执行
        session.refresh(clazz) # 更新clazz，把刚刚创建的一行记录的id取出来更新到对象中
        print(clazz)

        course = Course(name="数学")
        session.add(course)
        session.flush()
        course = session.query(Course).filter(Course.name=="数学").first() # 把刚插入的记录查询出来
        print(course)

        student = Student(stu_no="007", name="张三", clazz_id=clazz.id)
        student.courses.append(course) # 测试backref给student加了courses属性，并且可以通过其维护关系表
        session.add(course)
        session.flush()

        session.refresh(course)
        print(course.students) 
        print(to_jsonable(course)) #to_jsonable有students属性需要覆盖keys方法来支持

        student.name = "李四"
        session.flush()  # 这里相当于执行了update


        # 执行原生SQL查询
        r = session.execute("select id, name from student").fetchall()
        for e in r:
            print(e) # 可以看到已经变成了李四



if __name__ == '__main__':
    main()