codekingpro/portable-devtools

1"""Retry quota implementation.2 3 4"""5import threading6 7 8class RetryQuota:9    INITIAL_CAPACITY = 50010 11    def __init__(self, initial_capacity=INITIAL_CAPACITY, lock=None):12        self._max_capacity = initial_capacity13        self._available_capacity = initial_capacity14        if lock is None:15            lock = threading.Lock()16        self._lock = lock17 18    def acquire(self, capacity_amount):19        """Attempt to aquire a certain amount of capacity.20 21        If there's not sufficient amount of capacity available, ``False``22        is returned.  Otherwise, ``True`` is returned, which indicates that23        capacity was successfully allocated.24 25        """26        # The acquire() is only called when we encounter a retryable27        # response so we aren't worried about locking the entire method.28        with self._lock:29            if capacity_amount > self._available_capacity:30                return False31            self._available_capacity -= capacity_amount32            return True33 34    def release(self, capacity_amount):35        """Release capacity back to the retry quota.36 37        The capacity being released will be truncated if necessary38        to ensure the max capacity is never exceeded.39 40        """41        # Implementation note:  The release() method is called as part42        # of the "after-call" event, which means it gets invoked for43        # every API call.  In the common case where the request is44        # successful and we're at full capacity, we can avoid locking.45        # We can't exceed max capacity so there's no work we have to do.46        if self._max_capacity == self._available_capacity:47            return48        with self._lock:49            amount = min(50                self._max_capacity - self._available_capacity, capacity_amount51            )52            self._available_capacity += amount53 54    @property55    def available_capacity(self):56        return self._available_capacity57