py-app-runner 0.5.49.dev0__py3-none-any.whl

py_app_runner/tick_service.py

@@ -0,0 +1,308 @@
+import asyncio
+import inspect
+import logging
+import random
+from collections.abc import Awaitable, Callable
+from signal import SIGHUP, SIGINT, SIGQUIT, SIGTERM, Signals
+
+from .timer import Timer
+
+CALLBACK_TYPE = Callable[[int, float], None] | Callable[[int, float], Awaitable[None]]
+SHUTDOWN_CALLBACK_TYPE = Callable[[], None] | Callable[[], Awaitable[None]] | None
+
+
+class ShutdownRequestedException(Exception):
+    """
+    This class can be used to signal parts of the service/app that a shutdown was requested.
+    When needToShutdown is called with withException=True, this exception will be raised.
+    This is useful for example when you want to stop a long running task, but don't want
+    shutdown exception to be logged.
+    """
+
+    pass
+
+
+class TickService:
+    # Intervals
+    default_tick_interval: float
+    tick_interval: float
+    tick_count: int
+    tick_time: float
+
+    # Callbacks
+    callbacks: list[CALLBACK_TYPE]
+    shut_down_callback: SHUTDOWN_CALLBACK_TYPE
+    back_off_callback: CALLBACK_TYPE | None
+
+    # Various
+    shutdown_event: asyncio.Event
+    current_loop: asyncio.AbstractEventLoop
+    back_off_mode: bool
+    back_off_count: int
+    logger_name: str
+    logger: logging.Logger
+    timer: Timer
+
+    skip_signal_handling: bool = False
+    """Skip signal handling. Be careful with this, as it can lead to unstoppable processes"""
+
+    #######################
+    ### Class Lifecycle ###
+    #######################
+
+    def __del__(self) -> None:
+        self.logger.debug("Dealloc")
+
+    def __init__(
+        self,
+        tick_interval: float,
+        callback: CALLBACK_TYPE | None = None,
+        shut_down_callback: SHUTDOWN_CALLBACK_TYPE = None,
+        back_off_callback: CALLBACK_TYPE | None = None,
+        instance_name: str | None = None,
+    ):
+        # Intervals
+        self.default_tick_interval = tick_interval
+        self.tick_interval = tick_interval
+        self.tick_count = 0
+        self.tick_time = 0
+
+        # Callbacks
+        self.callbacks = [] if callback is None else [callback]
+        self.shut_down_callback = shut_down_callback
+        self.back_off_callback = back_off_callback
+
+        # Various
+        self.shutdown_event = asyncio.Event()
+        self.back_off_mode = False
+        self.back_off_count = 0
+
+        instance_name = f".{instance_name}" if instance_name is not None else ""
+        self.logger_name = f"{self.__class__.__name__}{instance_name}"
+        self.logger = logging.getLogger(self.logger_name)
+        self.timer = Timer("TickService")
+
+    ###############
+    ### Helpers ###
+    ###############
+
+    def add_callback(self, callback: CALLBACK_TYPE) -> None:
+        self.callbacks.append(callback)
+
+    def set_interval(self, interval: float) -> None:
+        self.tick_interval = interval
+        self.logger.debug(f"Tick interval set to {self.tick_interval} seconds")
+
+    def increase_interval_by(self, amount: float) -> None:
+        """
+        Increases the tick interval by a specified amount.
+
+        :param amount: The amount in seconds to increase the tick interval by.
+        """
+        self.tick_interval += amount
+        self.logger.debug(f"Tick interval increased by {amount} seconds, now at {self.tick_interval} seconds")
+
+    def increase_interval_exp(self, factor: float = 1.1, max_interval: float = 300) -> None:
+        """
+        Increases the tick interval exponentially for the next tick.
+
+        :param factor: The exponential factor to multiply the current interval by. Default is 1.1.
+        :param max_interval: The maximum interval in seconds. Default is 300.
+        """
+        self.tick_interval = min(self.tick_interval * factor, max_interval)
+        self.tick_interval *= random.uniform(0.9, 1.1)  # Little jitter
+        self.logger.debug(f"Tick interval increased to {self.tick_interval} seconds")
+
+    def reset_interval(self) -> None:
+        """
+        Resets the tick interval to the default value.
+        """
+        self.tick_interval = self.default_tick_interval
+        self.logger.debug(f"Tick interval reset to {self.tick_interval} seconds")
+
+    def set_backoff(self, start: bool = True) -> None:
+        self.back_off_mode = start
+        self.back_off_count = 0
+
+    def format_time(self, seconds: float) -> str:
+        """
+        Formats seconds to human readable time.
+
+        :param seconds: The seconds to format.
+        :return: The formatted time.
+        """
+        if seconds < 60:
+            return f"{seconds:.2f}s"
+
+        if seconds < 3600:
+            m, s = divmod(seconds, 60)
+            return f"{int(m)}m {int(s)}s"
+
+        m, s = divmod(seconds, 60)
+        h, m = divmod(m, 60)
+        return f"{int(h)}h {int(m)}m {int(s)}s"
+
+    async def need_to_shutdown(self, with_exception: bool = False) -> bool:
+        """
+        Waits for 0.001 seconds and checks if the shutdown event is set.
+
+        Returns:
+            bool: True if the shutdown event is set, False otherwise.
+        """
+        try:
+            await asyncio.sleep(0.001)
+        except asyncio.CancelledError:
+            self.logger.debug("needToShutdown: Sleep was cancelled")
+
+        if with_exception:
+            if self.shutdown_event.is_set():
+                raise ShutdownRequestedException("Shutdown requested")
+
+        return self.shutdown_event.is_set()
+
+    #######################
+    ### Loop management ###
+    #######################
+
+    async def shutdown(self, signal: Signals) -> None:
+        self.logger.debug(f"Shutdown signal received: {signal.name} ({signal.value})")
+
+        # Lets first shutdown our event loop
+        self.shutdown_event.set()
+
+        # Then call external shutdown callbacks
+        if self.shut_down_callback is not None:
+            # Need to catch exceptions here, otherwise the loop will be stopped with a interval delay
+            try:
+                if inspect.iscoroutinefunction(self.shut_down_callback):
+                    await self.shut_down_callback()
+                else:
+                    self.shut_down_callback()
+            except Exception as e:
+                self.logger.exception(e)
+
+        # And finally we cancel all tasks so for example asyncio.sleep is terminated
+        tasks = [
+            t
+            for t in asyncio.all_tasks(loop=self.current_loop)
+            if t is not asyncio.current_task(loop=self.current_loop)
+        ]
+        [task.cancel() for task in tasks]
+
+        logging.debug(f"Cancelling {len(tasks)} outstanding tasks")
+        # return_exceptions: the CancelledError we just caused would otherwise
+        # propagate here and abort the rest of the shutdown.
+        await asyncio.gather(*tasks, return_exceptions=True)
+        logging.debug("All outstanding tasks are cancelled")
+
+        # Remove callback handlers to avoid circular references
+        self.callbacks = []
+        self.shut_down_callback = None
+
+    def shutdown_handler(self, signal: Signals) -> None:
+        asyncio.create_task(self.shutdown(signal))
+
+    def init_signals(self) -> None:
+        if self.skip_signal_handling:
+            self.logger.debug("Skipping signal handlers")
+            return
+
+        self.logger.debug("Initializing signal handlers")
+
+        # Attach signal handlers
+        signals = (SIGHUP, SIGTERM, SIGINT, SIGQUIT)
+        for s in signals:
+            self.current_loop.add_signal_handler(s, lambda s=s: self.shutdown_handler(s))
+
+    async def loop_wait(self, timeout: int | float | None = None) -> None:
+        if timeout is None:
+            timeout = self.tick_interval
+
+        self.logger.debug(f"loop_wait: Sleeping for {self.format_time(timeout)}")
+
+        try:
+            await asyncio.sleep(timeout)
+        except asyncio.CancelledError:
+            self.logger.debug("loop_wait: Sleep was cancelled")
+
+    async def run_loop(self) -> None:
+        self.current_loop = asyncio.get_running_loop()
+        self.init_signals()
+
+        # Start infinite loop
+        self.logger.debug("Starting process loop")
+        while not self.shutdown_event.is_set():
+            # Tick
+            self.logger.debug(f"Tick #{self.tick_count + 1} started")
+            async with self.timer.aenter("TickService.tick"):
+                await self.tick()
+
+            # Timer stats
+            self.timer.print_timer_stats()
+            tick_total = self.timer.total_times.get("TickService.tick", {}).get("total", 0)
+            self.logger.debug(f"Tick #{self.tick_count} finished in {self.format_time(tick_total)}")
+
+            # Per-run timings key on run_index; without a reset the runtimes dicts
+            # grow unbounded on long-running services.
+            self.timer.reset_timers()
+            self.timer.run_index += 1
+
+            # Sleep
+            if not self.shutdown_event.is_set():
+                await self.loop_wait()
+
+        self.logger.debug("Process loop is stopped")
+
+    async def stop_loop(self, timeout: float = 10.0) -> None:
+        """Ask the loop to stop and wait for the shutdown to finish.
+
+        Safe to call from another thread's event loop: the coroutine is scheduled on
+        the loop that owns this service and awaited from the caller's loop.
+        """
+        self.logger.debug("Stopping process loop")
+
+        current_loop = getattr(self, "current_loop", None)
+        if current_loop is None or current_loop.is_closed():
+            self.logger.debug("Loop was never started or is already closed")
+            self.shutdown_event.set()
+            return
+
+        future = asyncio.run_coroutine_threadsafe(self.shutdown(SIGTERM), current_loop)
+        try:
+            await asyncio.wait_for(asyncio.wrap_future(future), timeout=timeout)
+        except TimeoutError:
+            self.logger.warning(f"Shutdown did not finish within {timeout}s, continuing anyway")
+        except Exception as e:
+            self.logger.warning(f"Shutdown failed: {e}")
+
+    ####################
+    ### Tick methods ###
+    ####################
+
+    async def tick(self) -> None:
+        self.tick_count += 1
+        self.tick_time += self.tick_interval
+
+        # Backoff mode
+        if self.back_off_mode:
+            self.back_off_count += 1
+            if self.back_off_callback is not None:
+                if inspect.iscoroutinefunction(self.back_off_callback):
+                    await self.back_off_callback(self.back_off_count, self.tick_time)
+                else:
+                    self.back_off_callback(self.back_off_count, self.tick_time)
+
+        # Regular tick
+        else:
+            for callback in self.callbacks:
+                try:
+                    if inspect.iscoroutinefunction(callback):
+                        await callback(self.tick_count, self.tick_time)
+                    else:
+                        callback(self.tick_count, self.tick_time)
+                except asyncio.CancelledError:
+                    self.logger.debug("Tick callback cancelled")
+                    raise
+
+                if await self.need_to_shutdown():
+                    break

py_app_runner/timer.py

@@ -0,0 +1,289 @@
+import logging
+import time
+from collections.abc import AsyncGenerator, Generator
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, TypedDict
+
+
+class TimerProperties(TypedDict):
+    level: int
+    total: float
+    runtimes: dict[int, float]
+
+
+class Timer:
+    name: str
+    logger: logging.Logger
+
+    start_times: dict[str, float]
+    total_times: dict[str, TimerProperties]
+
+    timers: list[tuple[int, str]]
+    level: int
+    run_index: int
+
+    def __init__(self, name: str):
+        self.name = name
+        self.logger = logging.getLogger(f"{self.__class__.__name__}")
+
+        self.start_times = {}
+        self.total_times = {}
+        self.timers = []
+        self.level = 0
+        self.run_index = 0
+
+    # * Class generator
+    @contextmanager
+    def enter(self, name: str) -> Generator["Timer", None, None]:
+        self.level += 1
+        self.start(name=name)
+
+        try:
+            yield self
+        finally:
+            self.stop(name=name)
+            self.level -= 1
+
+    @asynccontextmanager
+    async def aenter(self, name: str) -> AsyncGenerator["Timer", None]:
+        self.level += 1
+        self.start(name=name)
+
+        try:
+            yield self
+        finally:
+            self.stop(name=name)
+            self.level -= 1
+
+    # * Class context manager
+    def __enter__(self) -> "Timer":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.stop()
+
+    async def __aenter__(self) -> "Timer":
+        self.start()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.stop()
+
+    # * Start / stop
+    def start(self, name: str | None = None, skip_append: bool = False) -> None:
+        if name is None:
+            name = self.name
+
+        if len(self.timers) > 0:
+            (
+                level,
+                prev_timer_name,
+            ) = self.timers[-1]
+            if prev_timer_name != name and level != self.level:
+                self.stop(name=prev_timer_name, skip_remove=True)
+
+        self.start_times[name] = time.perf_counter()
+
+        if not skip_append:
+            self.timers.append(
+                (
+                    self.level,
+                    name,
+                )
+            )
+
+    def stop(
+        self,
+        name: str | None = None,
+        skip_remove: bool = False,
+    ) -> None:
+        if name is None:
+            name = self.name
+        end_time = time.perf_counter()
+
+        if name not in self.start_times:
+            self.logger.warning(f"Timer {name} stopped before it was started")
+            return None
+
+        elapsed = end_time - self.start_times[name]
+        if name not in self.total_times:
+            self.total_times[name] = TimerProperties(level=self.level, total=0, runtimes={})
+
+        self.total_times[name]["runtimes"][self.run_index] = elapsed
+        self.total_times[name]["total"] += elapsed
+
+        if not skip_remove:
+            # Remove current timer
+            self.timers.pop()
+
+            # Continue previous timer
+            if len(self.timers) > 0:
+                (
+                    level,
+                    prev_timer_name,
+                ) = self.timers[-1]
+                if prev_timer_name != name and level != self.level:
+                    self.start(name=prev_timer_name, skip_append=True)
+
+    # * Stats
+    def reset_timers(self) -> None:
+        """Resets the current timer. Call this at the start of an iteration."""
+        self.start_times = {}
+        self.total_times = {}
+
+    # * Output
+    def print_timer_stats_table(self, return_data: bool = False) -> str | None:
+        """Prints the current timing information into a table."""
+        timer_stats = "\n\n"
+
+        all_fn_names = [k for k in self.total_times.keys()]
+        if len(all_fn_names) == 0:
+            return None
+
+        # Max width of level column
+        max_level_width = max([self.total_times[k].get("level", 0) for k in all_fn_names])
+        if max_level_width % 2 == 1:
+            max_level_width += 1
+        max_level_width = max(max_level_width, 6)
+
+        # Max width of name column
+        max_name_width = max([len(k) for k in all_fn_names] + [4])
+        if max_name_width % 2 == 1:
+            max_name_width += 1
+
+        # Format string
+        format_str = (
+            f"{{:>{max_level_width}}} | {{:>{max_name_width}}} | "
+            f"{{:>10.4f}} | {{:>10.4f}} | {{:>10.4f}} | {{:>10.4f}} | {{:>5}}"
+        )
+
+        header = (
+            f"{{:>{max_level_width}}} | {{:>{max_name_width}}} | {{:^10}} | {{:^10}} | {{:^10}} | {{:^10}} | {{:^5}}"
+        ).format("Level", "Name", "Total (ms)", "Min (ms)", "Max (ms)", "Avg (ms)", "Count")
+        timer_stats += f"{header}\n"
+
+        sep_idx = header.find("|")
+        sep_text = ("-" * sep_idx) + "+" + ("-" * (len(header) - sep_idx - 1))
+        timer_stats += f"{sep_text}\n"
+
+        for name in all_fn_names:
+            runtimes = self.total_times[name].get("runtimes", {}).values()
+            total_time = self.total_times[name].get("total", 0) * 1000
+            min_time = min(runtimes) * 1000 if runtimes else 0
+            max_time = max(runtimes) * 1000 if runtimes else 0
+            avg_time = (sum(runtimes) / len(runtimes)) * 1000 if runtimes else 0
+            run_count = len(runtimes)
+
+            timer_stats += format_str.format(
+                self.total_times[name].get("level"),
+                name,
+                total_time,
+                min_time,
+                max_time,
+                avg_time,
+                run_count,
+            )
+            timer_stats += "\n"
+
+        totals_format_str = (
+            f"{{:>{max_level_width}}} | {{:>{max_name_width}}} | "
+            f"{{:>10.4f}} | {{:>10.4}} | {{:>10.4}} | {{:>10.4}} | {{:>10.4}}"
+        )
+        timer_stats += f"{sep_text}\n"
+        timer_stats += totals_format_str.format(
+            "",
+            "Total",
+            self.total_time(self.total_times) * 1000,
+            "",
+            "",
+            "",
+            "",
+        )
+        timer_stats += "\n\n"
+
+        if return_data:
+            return timer_stats
+
+        self.logger.debug(timer_stats)
+        return None
+
+    def print_timer_stats_json(self, return_data: bool = False) -> str | None:
+        """Prints the current timing information into a JSON string."""
+        pass
+
+    def print_timer_stats_csv(self, return_data: bool = False) -> str | None:
+        """Prints the current timing information into a CSV string."""
+        all_fn_names = list(self.total_times.keys())
+        if not all_fn_names:
+            return None
+
+        # Prepare header
+        header = all_fn_names + ["Total"]
+
+        # Get all runtime indices and find the maximum index
+        all_indices: set[int] = set()
+        for name in all_fn_names:
+            all_indices.update(self.total_times[name].get("runtimes", {}).keys())
+        max_index = max(all_indices) if all_indices else 0
+
+        # Prepare data rows
+        data_rows: list[list[str]] = []
+        for i in range(max_index + 1):
+            row: list[str] = []
+            row_total = 0.0
+            for name in all_fn_names:
+                runtime = self.total_times[name].get("runtimes", {}).get(i)
+                if runtime is not None:
+                    value = runtime * 1000
+                    row.append(f"{value:.4f}")
+                    row_total += value
+                else:
+                    row.append("")
+            row.append(f"{row_total:.4f}")
+            data_rows.append(row)
+
+        # Prepare total row
+        total_row: list[str] = []
+        grand_total = 0.0
+        for name in all_fn_names:
+            total = sum(self.total_times[name].get("runtimes", {}).values()) * 1000
+            total_row.append(f"{total:.4f}")
+            grand_total += total
+        total_row.append(f"{grand_total:.4f}")
+
+        # Combine all rows
+        all_rows = [header] + data_rows + [total_row]
+
+        # Convert to CSV string
+        output = "\n".join([",".join(str(cell) for cell in row) for row in all_rows])
+
+        if return_data:
+            return output
+
+        self.logger.info(f"\n{output}")
+        return None
+
+    def print_timer_stats(
+        self,
+        output_format: str = "table",
+        return_data: bool = False,
+    ) -> str | None:
+        """
+        Prints the current timing information into a table.
+
+        output_format: str - choices: ["table", "json", "csv"]
+        """
+
+        if output_format == "table":
+            return self.print_timer_stats_table(return_data=return_data)
+        elif output_format == "json":
+            return self.print_timer_stats_json(return_data=return_data)
+        elif output_format == "csv":
+            return self.print_timer_stats_csv(return_data=return_data)
+
+        return None
+
+    def total_time(self, total_times: dict[str, TimerProperties]) -> float:
+        """Returns the total amount accumulated across all functions in seconds."""
+        return sum([timer.get("total", 0) for _name, timer in total_times.items()])