rnsh/rnsh/retry.py

import asyncio
import logging
import threading
import time
import logging as __logging
from typing import Callable

module_logger = __logging.getLogger(__name__)


class RetryStatus:
    def __init__(self, tag: any, try_limit: int, wait_delay: float, retry_callback: Callable[[any, int], any],
                 timeout_callback: Callable[[any, int], None], tries: int = 1):
        self._log = module_logger.getChild(self.__class__.__name__)
        self._log.setLevel(logging.INFO)
        self.tag = tag
        self.try_limit = try_limit
        self.tries = tries
        self.wait_delay = wait_delay
        self.retry_callback = retry_callback
        self.timeout_callback = timeout_callback
        self.try_time = time.time()
        self.completed = False

    @property
    def ready(self):
        ready = time.time() > self.try_time + self.wait_delay
        # self._log.debug(f"ready check {self.tag} try_time {self.try_time} wait_delay {self.wait_delay} next_try {self.try_time + self.wait_delay} now {time.time()} exceeded {time.time() - self.try_time - self.wait_delay} ready {ready}")
        return ready

    @property
    def timed_out(self):
        return self.ready and self.tries >= self.try_limit

    def timeout(self):
        self.completed = True
        self.timeout_callback(self.tag, self.tries)

    def retry(self):
        self.tries += 1
        self.try_time = time.time()
        self.retry_callback(self.tag, self.tries)


class RetryThread:
    def __init__(self, loop_period: float = 0.25, name: str = "retry thread"):
        self._log = module_logger.getChild(self.__class__.__name__)
        self._loop_period = loop_period
        self._statuses: list[RetryStatus] = []
        self._tag_counter = 0
        self._lock = threading.RLock()
        self._run = True
        self._finished: asyncio.Future | None = None
        self._thread = threading.Thread(name=name, target=self._thread_run)
        self._thread.start()

    def close(self, loop: asyncio.AbstractEventLoop | None = None) -> asyncio.Future | None:
        self._log.debug("stopping timer thread")
        if loop is None:
            self._run = False
            self._thread.join()
            return None
        else:
            self._finished = loop.create_future()
            return self._finished

    def _thread_run(self):
        while self._run and self._finished is None:
            time.sleep(self._loop_period)
            ready: list[RetryStatus] = []
            prune: list[RetryStatus] = []
            with self._lock:
                ready.extend(list(filter(lambda s: s.ready, self._statuses)))
            for retry in ready:
                try:
                    if not retry.completed:
                        if retry.timed_out:
                            self._log.debug(f"timed out {retry.tag} after {retry.try_limit} tries")
                            retry.timeout()
                            prune.append(retry)
                        elif retry.ready:
                            self._log.debug(f"retrying {retry.tag}, try {retry.tries + 1}/{retry.try_limit}")
                            retry.retry()
                except Exception as e:
                    self._log.error(f"error processing retry id {retry.tag}: {e}")
                    prune.append(retry)

            with self._lock:
                for retry in prune:
                    self._log.debug(f"pruned retry {retry.tag}, retry count {retry.tries}/{retry.try_limit}")
                    try:
                        self._statuses.remove(retry)
                    except:
                        pass
        if self._finished is not None:
            self._finished.set_result(None)

    def _get_next_tag(self):
        self._tag_counter += 1
        return self._tag_counter

    def has_tag(self, tag: any) -> bool:
        with self._lock:
            return next(filter(lambda s: s.tag == tag, self._statuses), None) is not None

    def begin(self, try_limit: int, wait_delay: float, try_callback: Callable[[any, int], any],
              timeout_callback: Callable[[any, int], None], tag: int | None = None) -> any:
        self._log.debug(f"running first try")
        tag = try_callback(tag, 1)
        self._log.debug(f"first try got id {tag}")
        with self._lock:
            if tag is None:
                tag = self._get_next_tag()
            self.complete(tag)
            self._statuses.append(RetryStatus(tag=tag,
                                              tries=1,
                                              try_limit=try_limit,
                                              wait_delay=wait_delay,
                                              retry_callback=try_callback,
                                              timeout_callback=timeout_callback))
        self._log.debug(f"added retry timer for {tag}")

    def complete(self, tag: any):
        assert tag is not None
        with self._lock:
            status = next(filter(lambda l: l.tag == tag, self._statuses), None)
            if status is not None:
                status.completed = True
                self._statuses.remove(status)
                self._log.debug(f"completed {tag}")
                return

        self._log.debug(f"status not found to complete {tag}")

    def complete_all(self):
        with self._lock:
            for status in self._statuses:
                status.completed = True
                self._log.debug(f"completed {status.tag}")
            self._statuses.clear()