Source code for mlthon.basics.rate_tracker

import abc


class RateTracker:
    @property
    def max_rate(self):
        return self._max_rate_

    @property
    def window(self):
        return self._window_ms_

    @abc.abstractmethod
    def is_valid_rate(self, ts: int) -> bool:
        """
        :param ts: epoch time
        :return:
        """
        raise NotImplementedError


class SlicingRateTracker(RateTracker):
    def __init__(self, max_rate: int, window_ms: int):
        """
        Track maximum rate allowed with in given time window
        :param max_rate: Maximum rate
        :param window_ms: Time window in milliseconds
        """
        self._max_rate_ = max_rate
        self._window_ms_ = window_ms
        self._curr_rate_ = 0
        self._end_ts_ = 0

    def is_valid_rate(self, ts: int) -> bool:
        if ts <= self._end_ts_:
            if self._curr_rate_ > self._max_rate_:
                return False
            self._curr_rate_ += 1
        elif ts > self._end_ts_:
            self._end_ts_ = ts + self.window
            self._curr_rate_ = 1
        return self._curr_rate_ <= self._max_rate_