onestep 0.5.0__py3-none-any.whl

onestep/onestep.py

@@ -0,0 +1,281 @@
+import functools
+import collections
+import inspect
+import logging
+
+from inspect import isgenerator, iscoroutinefunction, isasyncgenfunction, isasyncgen
+from itertools import groupby
+from typing import Optional, List, Dict, Any, Callable, Union, Type
+
+from .broker.base import BaseBroker
+from .middleware import BaseMiddleware
+from .exception import StopMiddleware
+from .message import Message
+from .retry import TimesRetry
+from .signal import message_sent, started, stopped
+from .state import State
+from .worker import ThreadWorker, BaseWorker
+
+logger = logging.getLogger(__name__)
+
+MAX_WORKERS = 20
+DEFAULT_WORKERS = 1
+DEFAULT_WORKER_CLASS = ThreadWorker
+
+
+class BaseOneStep:
+    consumers: Dict[str, List[BaseWorker]] = collections.defaultdict(list)
+    state = State()  # 全局状态
+
+    def __init__(self, fn,
+                 group: str = "OneStep",
+                 name: Optional[str] = None,
+                 from_broker: Union[BaseBroker, List[BaseBroker], None] = None,
+                 to_broker: Union[BaseBroker, List[BaseBroker], None] = None,
+                 workers: Optional[int] = None,
+                 worker_class: Optional[Type[BaseWorker]] = None,
+                 middlewares: Optional[List[Any]] = None,
+                 retry: Union[Callable, object] = TimesRetry(),
+                 error_callback: Optional[Union[Callable, object]] = None):
+        self.group = group
+        self.fn = fn
+        self.name = name or fn.__name__
+        self.workers = workers or DEFAULT_WORKERS
+        self.worker_class = worker_class or DEFAULT_WORKER_CLASS
+        if self.workers > MAX_WORKERS:
+            logger.warning(f"workers[{self.workers}] greater than {MAX_WORKERS}")
+            self.workers = MAX_WORKERS
+        self.middlewares = middlewares or []
+        if self.middlewares:
+            for m in self.middlewares:
+                if not isinstance(m, BaseMiddleware):
+                    raise TypeError(f"middleware must be BaseMiddleware instance, not {type(m)}")
+
+        self.from_brokers = self._init_broker(from_broker)
+        self.to_brokers = self._init_broker(to_broker)
+        self.retry = retry
+        self.error_callback = error_callback
+
+        for broker in self.from_brokers:
+            self._add_consumer(broker)
+
+    @staticmethod
+    def _init_broker(broker: Union[BaseBroker, List[BaseBroker], None] = None):
+        if not broker:
+            return []
+
+        if isinstance(broker, BaseBroker):
+            return [broker]
+        if isinstance(broker, (list, tuple)):
+            return list(broker)
+        raise TypeError(
+            f"broker must be BaseBroker or list or tuple, not {type(broker)}")
+
+    def _add_consumer(self, broker):
+        """ 添加来源消费者 """
+        worker_class_params = inspect.signature(self.worker_class.__init__).parameters
+        if "workers" in worker_class_params:
+            self.consumers[self.group].append(
+                self.worker_class(onestep=self, broker=broker, workers=self.workers)
+            )
+        else:
+            for _ in range(self.workers):
+                self.consumers[self.group].append(
+                    self.worker_class(onestep=self, broker=broker)
+                )
+
+    @classmethod
+    def _find_consumers(cls, group: Optional[str] = None):
+        """按组查找消费者"""
+        if group is None:
+            consumers = [c for v in cls.consumers.values() for c in v]
+        else:
+            consumers = cls.consumers[group]
+        return consumers
+
+    @classmethod
+    def print_jobs(cls, group):
+        print("Jobs:")
+        prints = []
+        _consumers = cls._find_consumers(group)
+        group_instance = groupby(_consumers, key=lambda x: x.instance)
+        for instance, _ in group_instance:
+            prints.append([instance.name, instance.group, instance.workers, str(instance.from_brokers)])
+        print("{:<15} {:<10} {:<10} {:<20}".format("Job", "Group", "Workers", "From Brokers"))
+        for v in prints:
+            print("{:<15} {:<10} {:<10} {:<20}".format(*v))
+
+    @classmethod
+    def start(cls, group: Optional[str] = None, print_jobs: bool = False):
+        logger.debug(f"start group [{group or 'all'}]")
+        _consumers = cls._find_consumers(group)
+        if not _consumers:
+            logger.debug(f"no consumer found in group [{group or 'all'}]")
+            return
+        if print_jobs:
+            cls.print_jobs(group)
+        for consumer in _consumers:
+            consumer.start()
+            logger.debug(f"started: {consumer=}")
+
+    @classmethod
+    def shutdown(cls, group: Optional[str] = None):
+        logger.debug(f"stop group [{group or 'all'}]")
+        for consumer in cls._find_consumers(group):
+            consumer.shutdown()
+            logger.debug(f"stopped: {consumer=}")
+
+    def wraps(self, func):
+        @functools.wraps(func)
+        def wrapped_f(*args, **kwargs):
+            return self(*args, **kwargs)  # noqa
+
+        return wrapped_f
+
+    def send(self, result, broker=None):
+        """将返回的内容交给broker发送"""
+        brokers = self._init_broker(broker) or self.to_brokers
+        # 如果是Message类型，就不再封装
+        message = result if isinstance(result, Message) else Message(body=result)
+        self.before_emit("send", message=message)
+
+        if result and not brokers:
+            logger.debug("send(result): broker is empty")
+            return
+
+        if not result and brokers:
+            logger.debug("send(result): body is empty")
+            return
+
+        for broker in brokers:
+            message_sent.send(self, message=message, broker=broker)
+            broker.send(message, step=self)
+        self.after_emit("send", message=message)
+
+    def before_emit(self, signal, *args, **kwargs):
+        signal = "before_" + signal
+        self.emit(signal, *args, **kwargs)
+
+    def after_emit(self, signal, *args, **kwargs):
+        signal = "after_" + signal
+        self.emit(signal, *args, **kwargs)
+
+    def emit(self, signal, *args, **kwargs):
+        for middleware in self.middlewares:
+            if not hasattr(middleware, signal):
+                continue
+            try:
+                getattr(middleware, signal)(step=self, *args, **kwargs)
+            except StopMiddleware as e:
+                logger.debug(f"middleware<{middleware}> is stopped，reason: {e}")
+                break
+
+    @classmethod
+    def is_shutdown(cls, group):
+        # check all broker
+        _consumers = cls._find_consumers(group)
+        if not _consumers:
+            return True
+        return all(broker._shutdown for broker in _consumers)
+
+
+def decorator_func_proxy(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+class SyncOneStep(BaseOneStep):
+
+    def __call__(self, *args, **kwargs):
+        """同步执行原函数"""
+
+        result = self.fn(*args, **kwargs)
+
+        if isgenerator(result):
+            for item in result:
+                self.send(item)
+        else:
+            self.send(result)
+        return result
+
+
+class AsyncOneStep(BaseOneStep):
+
+    async def __call__(self, *args, **kwargs):
+        """"异步执行原函数"""
+
+        if iscoroutinefunction(self.fn):
+            result = await self.fn(*args, **kwargs)
+        else:
+            result = self.fn(*args, **kwargs)
+
+        if isasyncgen(result):
+            async for item in result:
+                self.send(item)
+        else:
+            self.send(result)
+        return result
+
+
+class step:
+
+    def __init__(self,
+                 *,
+                 group: str = "OneStep",
+                 name: Optional[str] = None,
+                 from_broker: Union[BaseBroker, List[BaseBroker], None] = None,
+                 to_broker: Union[BaseBroker, List[BaseBroker], None] = None,
+                 workers: Optional[int] = None,
+                 worker_class: Optional[Type[BaseWorker]] = None,
+                 middlewares: Optional[List[Any]] = None,
+                 retry: Union[Callable, object] = TimesRetry(),
+                 error_callback: Optional[Union[Callable, object]] = None):
+        self.params = {
+            "group": group,
+            "name": name,
+            "from_broker": from_broker,
+            "to_broker": to_broker,
+            "workers": workers,
+            "worker_class": worker_class,
+            "middlewares": middlewares,
+            "retry": retry,
+            "error_callback": error_callback
+        }
+
+    def __call__(self, func, *_args, **_kwargs):
+        func.__step_params__ = self.params
+        if iscoroutinefunction(func) or isasyncgenfunction(func):
+            os = AsyncOneStep(fn=func, **self.params)
+        else:
+            os = SyncOneStep(fn=func, **self.params)
+
+        return os.wraps(func)
+
+    @staticmethod
+    def start(group=None, block=None, print_jobs=False):
+        BaseOneStep.start(group=group, print_jobs=print_jobs)
+        started.send()
+        if block:
+            import time
+            while not BaseOneStep.is_shutdown(group):
+                time.sleep(1)
+
+    @staticmethod
+    def shutdown(group=None):
+        BaseOneStep.shutdown(group=group)
+        stopped.send()
+
+    @staticmethod
+    def set_debugging():
+
+        if not BaseOneStep.state.debug:
+            onestep_logger = logging.getLogger("onestep")
+            handler = logging.StreamHandler()
+            handler.setFormatter(logging.Formatter(
+                "[%(levelname)s] %(asctime)s %(name)s:%(message)s"))
+            onestep_logger.addHandler(handler)
+            onestep_logger.setLevel(logging.DEBUG)
+            BaseOneStep.state.debug = True

onestep/retry.py

@@ -0,0 +1,117 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Optional, Tuple, Union, Type
+
+from .exception import RetryInLocal, RetryInQueue
+from .message import Message
+
+
+class RetryStatus(Enum):
+    CONTINUE = 1  # 继续（执行重试）
+    END_WITH_CALLBACK = 2  # 结束（执行回调）
+    END_IGNORE_CALLBACK = 3  # 结束（忽略回调）
+
+
+class BaseRetry(ABC):
+
+    @abstractmethod
+    def __call__(self, *args, **kwargs) -> Optional[RetryStatus]:
+        pass
+
+
+class NeverRetry(BaseRetry):
+
+    def __call__(self, message) -> Optional[RetryStatus]:
+        return RetryStatus.END_WITH_CALLBACK
+
+
+class AlwaysRetry(BaseRetry):
+
+    def __call__(self, message) -> Optional[RetryStatus]:
+        return RetryStatus.CONTINUE
+
+
+class AllRetry(BaseRetry):
+    def __init__(self, *retries):
+        self.retries = retries
+
+    def __call__(self, message) -> Optional[RetryStatus]:
+        for retry in self.retries:
+            status = retry(message)
+            if status != RetryStatus.CONTINUE:
+                return status
+        return RetryStatus.CONTINUE
+
+
+class AnyRetry(BaseRetry):
+    def __init__(self, *retries):
+        self.retries = retries
+
+    def __call__(self, message) -> Optional[RetryStatus]:
+        for retry in self.retries:
+            status = retry(message)
+            if status == RetryStatus.CONTINUE:
+                return RetryStatus.CONTINUE
+        return RetryStatus.END_WITH_CALLBACK
+
+
+class TimesRetry(BaseRetry):
+
+    def __init__(self, times: int = 3):
+        self.times = times
+
+    def __call__(self, message) -> Optional[RetryStatus]:
+        if message.failure_count < self.times:
+            return RetryStatus.CONTINUE
+        return RetryStatus.END_WITH_CALLBACK
+
+
+class RetryIfException(BaseRetry):
+
+    def __init__(self, exceptions: Optional[Tuple[Union[Exception, Type]]] = Exception):
+        self.exceptions = exceptions
+
+    def __call__(self, message) -> Optional[RetryStatus]:
+        if isinstance(message.exception.exc_value, self.exceptions):
+            return RetryStatus.CONTINUE
+        return RetryStatus.END_WITH_CALLBACK
+
+
+class AdvancedRetry(TimesRetry):
+    """高级重试策略
+    
+    1. 本地重试：如果异常是 RetryInLocal 或 指定异常，且重试次数未达到上限，则本地重试，不回调
+    2. 队列重试：如果异常是 RetryInQueue，且重试次数未达到上限，则入队重试，不回调
+    3. 其他异常：如果异常不是 RetryInLocal 或 RetryInQueue 或 指定异常，则不重试，回调
+    注：待重试的异常若继承自 RetryException，则可单独指定重试次数，否则默认为 3 次
+    """
+
+    def __init__(self, times: int = 3, exceptions: Optional[Tuple[Union[Exception, Type]]] = None):
+        super().__init__(times=times)
+        self.exceptions = (RetryInLocal, RetryInQueue) + (exceptions or ())
+
+    def __call__(self, message: Message) -> Optional[RetryStatus]:
+        if isinstance(message.exception.exc_value, self.exceptions):
+            max_retry_times = getattr(message.exception.exc_value, "times", None) or self.times
+            if message.failure_count < max_retry_times:
+                if isinstance(message.exception.exc_value, RetryInQueue):
+                    return RetryStatus.END_IGNORE_CALLBACK
+                return RetryStatus.CONTINUE
+            else:
+                return RetryStatus.END_WITH_CALLBACK
+        else:
+            return RetryStatus.END_WITH_CALLBACK
+
+
+# error callback
+class BaseErrorCallback(ABC):
+
+    @abstractmethod
+    def __call__(self, *args, **kwargs):
+        pass
+
+
+class NackErrorCallBack(BaseErrorCallback):
+
+    def __call__(self, message):
+        message.reject()

onestep/signal.py

@@ -0,0 +1,11 @@
+from blinker import signal
+
+started = signal("started")
+stopped = signal("stopped")
+
+message_received = signal("message_received")
+message_sent = signal("message_sent")
+message_consumed = signal("message_consumed")
+message_error = signal("message_error")
+message_drop = signal("message_drop")
+message_requeue = signal("message_requeue")

onestep/state.py

@@ -0,0 +1,23 @@
+# -*- coding: utf-8 -*-
+import typing
+
+
+class State:
+    _state: typing.Dict[str, typing.Any]
+
+    def __init__(self, state: typing.Optional[typing.Dict[str, typing.Any]] = None):
+        if state is None:
+            state = {}
+        super().__setattr__("_state", state)
+
+    def __setattr__(self, key: typing.Any, value: typing.Any) -> None:
+        self._state[key] = value
+
+    def __getattr__(self, key: typing.Any) -> typing.Any:
+        try:
+            return self._state[key]
+        except KeyError:
+            return None
+
+    def __delattr__(self, key: typing.Any) -> None:
+        del self._state[key]

onestep/worker.py

@@ -0,0 +1,205 @@
+"""
+将指定的函数放入线程中运行
+"""
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, Iterable
+
+import logging
+import threading
+from asyncio import iscoroutinefunction
+from inspect import isasyncgenfunction
+
+from asgiref.sync import async_to_sync
+
+from .message import Message
+from .retry import RetryStatus
+from .broker import BaseBroker
+from . import exception
+from .signal import message_received, message_consumed, message_error, message_drop, message_requeue
+
+logger = logging.getLogger(__name__)
+
+
+class BaseWorker:
+    broker_exit: Dict[BaseBroker, bool] = {}
+    broker_exit_lock = threading.Lock()
+
+    def __init__(self, onestep, broker: BaseBroker, *args, **kwargs):
+        self.instance = onestep
+        self.retry = self.instance.retry
+        self.error_callback = self.instance.error_callback
+        self.broker = broker
+        self.args = args
+        self.kwargs = kwargs
+        self._shutdown = False
+
+    @property
+    def instance_name(self):
+        return self.instance.fn.__name__
+
+    def start(self):
+        """启动 Worker"""
+        raise NotImplementedError
+
+    def run(self):
+        """执行 Worker 的逻辑"""
+        raise NotImplementedError
+
+    def shutdown(self):
+        """关闭 Worker"""
+        raise NotImplementedError
+
+    def receive_messages(self) -> Iterable[Message]:
+        """ 从broker中获取消息 """
+        for result in self.broker.consume():
+            if self._shutdown:
+                break
+            if result is None:
+                continue
+            messages = result if isinstance(result, Iterable) else [result]
+            yield from messages
+            # when broker is once, it will shut down after receive a message
+            if self.broker.once:
+                self.shutdown()
+
+    def _run_real_instance(self, message: Message) -> None:
+        """ 执行实例的逻辑 """
+        if iscoroutinefunction(self.instance.fn) or isasyncgenfunction(self.instance.fn):
+            async_to_sync(self.instance)(message, *self.args, **self.kwargs)
+        else:
+            self.instance(message, *self.args, **self.kwargs)
+
+    def handle_message(self, message: Message):
+        """ 处理消息 """
+        message.broker = message.broker or self.broker
+        logger.debug(f"{self.instance.name} receive message<{message}> from {self.broker!r}")
+        message_received.send(self, message=message)
+
+        try:
+            self.broker.before_emit("consume", message=message, step=self.instance)
+            self.instance.before_emit("consume", message=message)
+            self._run_real_instance(message)
+            self.handle_success(message)
+            self.instance.after_emit("consume", message=message)
+            self.broker.after_emit("consume", message=message, step=self.instance)
+        except (exception.DropMessage, exception.RejectMessage) as e:
+            self.handle_drop(message, e)
+        except exception.RequeueMessage as e:
+            self.handle_requeue(message, e)
+        except Exception as e:
+            message_error.send(self, message=message, error=e)
+            self.handle_error(message, e)
+            self.handle_retry(message)
+        finally:
+            self.handle_cancel_consume(message)
+
+    def handle_success(self, message):
+        message_consumed.send(self, message=message)
+        message.confirm(step=self.instance)
+
+    def handle_drop(self, message, reason):
+        message_drop.send(self, message=message, reason=reason)
+        logger.warning(f"{self.instance.name} dropped <{type(reason).__name__}: {str(reason)}>")
+        message.reject(step=self.instance)
+
+    def handle_requeue(self, message, reason):
+        message_requeue.send(self, message=message, reason=reason)
+        logger.warning(f"{self.instance.name} requeue <{type(reason).__name__}: {str(reason)}>")
+        message.requeue(is_source=True, step=self.instance)
+
+    def handle_error(self, message, error):
+        if self.instance.state.debug:
+            logger.exception(f"{self.instance.name} run error <{type(error).__name__}: {str(error)}>")
+        else:
+            logger.error(f"{self.instance.name} run error <{type(error).__name__}: {str(error)}>")
+        message.set_exception()
+
+    def handle_cancel_consume(self, message):
+        if self.broker.cancel_consume and self.broker.cancel_consume(message):
+            self.shutdown()
+
+    def handle_retry(self, message):
+        retry_status = self.retry(message)
+        if retry_status is RetryStatus.END_WITH_CALLBACK:
+            if self.error_callback:
+                self.error_callback(message)
+            message.reject()
+        elif retry_status is RetryStatus.END_IGNORE_CALLBACK:
+            message.requeue()
+        elif retry_status is RetryStatus.CONTINUE:
+            message.requeue()
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} {self.instance.name}>"
+
+
+class ThreadWorker(BaseWorker):
+
+    def __init__(self, onestep, broker: BaseBroker, *args, **kwargs):
+        """
+        线程执行包装过的`onestep`函数
+        :param onestep: OneStep实例
+        :param broker: 监听的from broker
+        """
+        super().__init__(onestep, broker, *args, **kwargs)
+        self.thread = None
+
+    def start(self):
+        """启动单线程 Worker"""
+        self.thread = threading.Thread(target=self.run, daemon=True)
+        self.thread.start()
+
+    def run(self):
+        """线程执行包装过的`onestep`函数
+
+        `fn`为`onestep`函数，执行会调用`onestep`的`__call__`方法
+        :return:
+        """
+
+        while not self._shutdown:
+            with ThreadWorker.broker_exit_lock:
+                if ThreadWorker.broker_exit.get(self.broker, False):
+                    self.shutdown()
+                    break
+            for message in self.receive_messages():
+                self.handle_message(message)
+
+    def shutdown(self):
+        ThreadWorker.broker_exit[self.broker] = True
+        self.broker.shutdown()
+        self._shutdown = True
+
+
+class ThreadPoolWorker(BaseWorker):
+    broker_exit: Dict[BaseBroker, bool] = {}
+    broker_exit_lock = threading.Lock()
+
+    def __init__(self, onestep, broker: BaseBroker, workers=None, *args, **kwargs):
+        super().__init__(onestep, broker, *args, **kwargs)
+        self.executor = ThreadPoolExecutor(max_workers=workers)
+
+    def start(self):
+        """启动线程池 Worker"""
+        self.executor.submit(self.run)
+
+    def run(self):
+        """线程执行包装过的`onestep`函数
+
+        `fn`为`onestep`函数，执行会调用`onestep`的`__call__`方法
+        :return:
+        """
+
+        while not self._shutdown:
+            with ThreadPoolWorker.broker_exit_lock:
+                if ThreadPoolWorker.broker_exit.get(self.broker, False):
+                    self.shutdown()
+                    break
+            for message in self.receive_messages():
+                # 将消息处理提交到线程池中并发执行
+                self.executor.submit(self.handle_message, message)
+
+    def shutdown(self):
+        """关闭线程池 Worker"""
+        ThreadPoolWorker.broker_exit[self.broker] = True
+        self._shutdown = True
+        self.executor.shutdown()