aws_xray_sdk/core/sampling/reservoir.py

import threading
from enum import Enum


class Reservoir:
    """
    Centralized thread-safe reservoir which holds fixed sampling
    quota, borrowed count and TTL.
    """
    def __init__(self):
        self._lock = threading.Lock()

        self._quota = None
        self._TTL = None

        self._this_sec = 0
        self._taken_this_sec = 0
        self._borrowed_this_sec = 0

        self._report_interval = 1
        self._report_elapsed = 0

    def borrow_or_take(self, now, can_borrow):
        """
        Decide whether to borrow or take one quota from
        the reservoir. Return ``False`` if it can neither
        borrow nor take. This method is thread-safe.
        """
        with self._lock:
            return self._borrow_or_take(now, can_borrow)

    def load_quota(self, quota, TTL, interval):
        """
        Load new quota with a TTL. If the input is None,
        the reservoir will continue using old quota until it
        expires or has a non-None quota/TTL in a future load.
        """
        if quota is not None:
            self._quota = quota
        if TTL is not None:
            self._TTL = TTL
        if interval is not None:
            self._report_interval = interval / 10

    @property
    def quota(self):
        return self._quota

    @property
    def TTL(self):
        return self._TTL

    def _time_to_report(self):
        if self._report_elapsed + 1 >= self._report_interval:
            self._report_elapsed = 0
            return True
        else:
            self._report_elapsed += 1

    def _borrow_or_take(self, now, can_borrow):
        self._adjust_this_sec(now)
        # Don't borrow if the quota is available and fresh.
        if (self._quota is not None and self._quota >= 0 and
                self._TTL is not None and self._TTL >= now):
            if(self._taken_this_sec >= self._quota):
                return ReservoirDecision.NO

            self._taken_this_sec = self._taken_this_sec + 1
            return ReservoirDecision.TAKE

        # Otherwise try to borrow if the quota is not present or expired.
        if can_borrow:
            if self._borrowed_this_sec >= 1:
                return ReservoirDecision.NO

            self._borrowed_this_sec = self._borrowed_this_sec + 1
            return ReservoirDecision.BORROW

    def _adjust_this_sec(self, now):
        if now != self._this_sec:
            self._taken_this_sec = 0
            self._borrowed_this_sec = 0
            self._this_sec = now


class ReservoirDecision(Enum):
    """
    An Enum of decisions the reservoir could make based on
    assigned quota with TTL and the current timestamp/usage.
    """
    TAKE = 'take'
    BORROW = 'borrow'
    NO = 'no'