text: # Copyright 2019 Scalyr Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #   http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ------------------------------------------------------------------------
  #
  # author: Edward Chee <echee@scalyr.com>
  
  from __future__ import unicode_literals
  from __future__ import absolute_import
  
  __author__ = "echee@scalyr.com"
  
  if False:  # NOSONAR
      from typing import Dict
  
  import random
  import threading
  import time
  from collections import deque
  
  from six.moves import range
  
  import scalyr_agent.scalyr_logging as scalyr_logging
  
  
  MULTIPLY = "multiply"
  RESET_THE_MULTIPLY = "reset_then_multiply"
  
  
  class RateLimiterToken(object):
      def __init__(self, token_id):
          self._token_id = token_id
  
      @property
      def token_id(self):
          """Integer ID"""
          return self._token_id
  
      def __repr__(self):
          return "Token #%s" % self._token_id
  
  
  class BlockingRateLimiter(object):
      """A Rate Limiter that blocks for a random time interval such that over time, the average rate of acquire() calls
      is R where R varies between a max upper and min lower bound. R is a "cluster" rate where the semantics are
      each of num_agents Scalyr Agents instantiates one of these rate limiters such that the aggregate access rate over
      all agents is R.  (Therefore, the per-agent rate is R / num_agents).
  
      Whenever a certain number of consecutive successes are seen, the rate is increased up until max_cluster_rate.
      Likewise, whenever a failure is reported (on token release()), the rate is lowered up until a min_cluster_rate.
      Cluster rate is reduced by dividing current rate by backoff_factor.  Cluster rate is increased either by
      multiplying by increase_factor or resetting to initial rate, depending on the strategy.
      (Note: if cluster rate is already greater or equal to initial rate, it will be multiplied by the increase factor)
      """
  
      STRATEGY_MULTIPLY = MULTIPLY
      STRATEGY_RESET_THEN_MULTIPLY = RESET_THE_MULTIPLY
  
      # Hard limit to protect against user error.
      # With 5K agents, each agent will start off at 4rps.
      # With 1 agent, the agent will start off at 20K rps.
      # Note: we do not specify an upper bound to allow the rate limiter to push the boundaries.
      HARD_LIMIT_INITIAL_CLUSTER_RATE = 20000
  
      HARD_LIMIT_MIN_CLUSTER_RATE = 1
  
      registry = {}  # type: Dict[str, BlockingRateLimiter]
      registry_lock = threading.Lock()
  
      def __repr__(self):
          return "%s/%s" % (self.__class__.__name__, self._name)
  
      @classmethod
      def get_instance(cls, key, global_config, logger=None):
          """Returned a named instance of a rate limiter so that it can be shared by multiple threads/code areas
  
          If a named instance does not exist, create one, based on the global configuration settings.
          """
          cls.registry_lock.acquire()
          try:
              if key not in cls.registry:
                  cls.registry[key] = BlockingRateLimiter(
                      global_config.k8s_ratelimit_cluster_num_agents,
                      global_config.k8s_ratelimit_cluster_rps_init,
                      global_config.k8s_ratelimit_cluster_rps_max,
                      global_config.k8s_ratelimit_cluster_rps_min,
                      global_config.k8s_ratelimit_consecutive_increase_threshold,
                      global_config.k8s_ratelimit_strategy,
                      global_config.k8s_ratelimit_increase_factor,
                      global_config.k8s_ratelimit_backoff_factor,
                      global_config.k8s_ratelimit_max_concurrency,
                      logger=logger,
                      name=key,
                  )
              return cls.registry[key]
          finally:
              cls.registry_lock.release()
  
      def __init__(
          self,
          num_agents,
          initial_cluster_rate,
          max_cluster_rate,
          min_cluster_rate,
          consecutive_success_threshold,
          strategy=MULTIPLY,
          increase_factor=2.0,
          backoff_factor=0.5,
          max_concurrency=1,
          logger=None,
          fake_clock=None,
          name=None,
      ):
          """
          @param num_agents: Number of agents in the cluster
          @param initial_cluster_rate: initial cluster rate (requests/sec)
          @param max_cluster_rate: upper bound on cluster rate (requests/sec)
          @param min_cluster_rate: lower bound on cluster rate (requests/sec)
          @param consecutive_success_threshold: number of consecutive successes to trigger a rate increase
          @param strategy: strategy for changing rate (multiply or reset-then-multiply)
          @param increase_factor: multiplicative factor for increasing rate
          @param backoff_factor: multiplicative factor for decreasing rate
          @param max_concurrency: max number of tokens (for limiting concurrent requests) that can be acquired
          @param logger: If supplied, will log messages for debugging
          @parm name: Name string for identification purposes (optional)
  
          @type num_agents: int
          @type initial_cluster_rate: float
          @type max_cluster_rate: float
          @type min_cluster_rate: float
          @type consecutive_success_threshold: int
          @type strategy: six.text_type
          @type increase_factor: float
          @type backoff_factor: float
          @type max_concurrency: int
          @type logger: Logger
          @type name: six.text_type
          """
          # Validate input (Note: will raise exception and thus kill the agent process if invalid)
          strategies = [self.STRATEGY_MULTIPLY, self.STRATEGY_RESET_THEN_MULTIPLY]
          if strategy not in strategies:
              raise ValueError("Increase strategy must be one of %s" % strategies)
  
          if max_concurrency < 1:
              raise ValueError("max_concurrency must be greater than 0")
  
          if consecutive_success_threshold < 1:
              raise ValueError(
                  "consecutive_success_threshold must be a positive int. Value=%s"
                  % consecutive_success_threshold
              )
  
          if int(consecutive_success_threshold) != consecutive_success_threshold:
              raise ValueError(
                  "consecutive_success_threshold must be a positive int. Value=%s"
                  % consecutive_success_threshold
              )
  
          if (
              initial_cluster_rate < min_cluster_rate
              or initial_cluster_rate > max_cluster_rate
          ):
              raise ValueError(
                  "Initial cluster rate must be between lower and upper rates. Initial=%s, Lower=%s, Upper=%s."
                  % (initial_cluster_rate, min_cluster_rate, max_cluster_rate)
              )
  
          self._logger = logger
          self._num_agents = num_agents
          self._initial_cluster_rate = initial_cluster_rate
  
          self._max_cluster_rate = max_cluster_rate
          self._min_cluster_rate = min_cluster_rate
          # A flag that when set to True signifies that a one-time adjustment has been made to min/max
          # cluster rates.  This is needed to avoid repeat warning messages being logged every time
          # __init__ is called on a redundant instance of this class (e.g. when created by k8s_monitor_initialize()
          self._lazy_adjusted = False
  
          self._consecutive_success_threshold = consecutive_success_threshold
          self._strategy = strategy
          self._increase_factor = increase_factor
          self._backoff_factor = backoff_factor
  
          # target cluster rate
          self._current_cluster_rate = self._initial_cluster_rate
          self._consecutive_successes = 0
          # state for tracking actual rate
          self._first_action_time = None
          self._last_action_time = None
          self._num_actions = 0
          # lock for protecting all the above state for cluster & actual rates
          self._cluster_rate_lock = threading.RLock()
  
          self._max_concurrency = max_concurrency
  
          self._name = name
  
          # A queue of tokens
          # 2->TODO python3 does not allow None in sort. There is sort in 'self._initialize_token_queue()'
          self._ripe_time = 0.0
          self._token_queue = deque()
          # Condition variable to synchronize access to token queue
          self._token_queue_cv = threading.Condition()
  
          # fake clock for testing
          self._fake_clock = fake_clock
          self._test_mode_lock = threading.Lock()
  
          self._initialize_token_queue()
  
      def _lazy_adjust_min_max_rates(self):
          """Adjust min/max cluster rates to hard limits"""
  
          self._cluster_rate_lock.acquire()
          try:
              if self._lazy_adjusted:
                  return
  
              if self._initial_cluster_rate > self.HARD_LIMIT_INITIAL_CLUSTER_RATE:
                  old_initial_cluster_rate = self._initial_cluster_rate
                  self._initial_cluster_rate = self.HARD_LIMIT_INITIAL_CLUSTER_RATE
                  if self._logger:
                      self._logger.warn(
                          "RateLimiter: initial cluster rate of %.2f is too high.  Limiting to %.2f."
                          % (old_initial_cluster_rate, self._initial_cluster_rate)
                      )
              if self._min_cluster_rate < self.HARD_LIMIT_MIN_CLUSTER_RATE:
                  old_min_cluster_rate = self._min_cluster_rate
                  self._min_cluster_rate = self.HARD_LIMIT_MIN_CLUSTER_RATE
                  if self._logger:
                      self._logger.warn(
                          "RateLimiter: min cluster rate of %.2f is too low.  Increasing to %.2f."
                          % (old_min_cluster_rate, self._min_cluster_rate)
                      )
              if self._initial_cluster_rate < self._min_cluster_rate:
                  old_initial_cluster_rate = self._initial_cluster_rate
                  self._initial_cluster_rate = self._min_cluster_rate
                  self._current_cluster_rate = self._initial_cluster_rate
                  if self._logger:
                      self._logger.warn(
                          "RateLimiter: initial cluster rate of %.2f is too low.  Increasing to %.2f."
                          % (old_initial_cluster_rate, self._initial_cluster_rate)
                      )
              self._lazy_adjusted = True
          finally:
              self._cluster_rate_lock.release()
  
      def _get_actual_cluster_rate(self):
          """Calculate actual rate based on recorded actions.  If fewer than 2 actions, returns None"""
          self._cluster_rate_lock.acquire()
          try:
              if self._num_actions < 2:
                  return None
              rate = (
                  float(self._num_actions - 1)
                  * self._num_agents
                  / (self._last_action_time - self._first_action_time)
              )
              return rate
          finally:
              self._cluster_rate_lock.release()
  
      def _record_actual_rate(self):
          """Record a completed action in order to keep track of actual rate"""
          self._cluster_rate_lock.acquire()
          try:
              t = self._time()
              self._num_actions += 1
              if self._num_actions == 1:
                  self._first_action_time = t
              self._last_action_time = t
          finally:
              self._cluster_rate_lock.release()
  
      def _reset_actual_rate(self):
          """Reset actual rate state"""
          self._cluster_rate_lock.acquire()
          try:
              self._num_actions = 0
              self._first_action_time = None
              self._last_action_time = None
          finally:
              self._cluster_rate_lock.release()
  
      @property
      def name(self):
          return self._name
  
      @property
      def current_cluster_rate(self):
          return self._current_cluster_rate
  
      @property
      def current_agent_rate(self):
          return float(self._current_cluster_rate) / self._num_agents
  
      def _initialize_token_queue(self):
          """Initialize a queue with max_concurrency tokens."""
          self._token_queue_cv.acquire()
          for num in range(self._max_concurrency):
              token = RateLimiterToken(num)
              self._token_queue.append(token)
          self._ripe_time = self._get_next_ripe_time()
          self._token_queue_cv.release()
  
      def _increase_cluster_rate(self):
          """Increase cluster rate but no higher than configured max and no lower than current cluster rate.
  
          If actual rate was calculated, then the new rate is calculated off the lower of actual & current target rate.
          Otherwise, new rate is calculated off the existing target rate.
  
          Note: this method is not thread-safe. Callers should pre-acquire cluster_rate_lock.
          """
          current_target_rate = self._current_cluster_rate
  
          actual_rate = self._get_actual_cluster_rate()
          if actual_rate is not None:
              # new rate should be based off the lower of actual & target rates
              new_rate = min(actual_rate, current_target_rate) * self._increase_factor
          else:
              new_rate = current_target_rate * self._increase_factor
  
          # new rate cannot be higher than configured higher bound
          new_rate = min(new_rate, self._max_cluster_rate)
  
          # new rate cannot be lower than existing rate (which may happen if actual rate was really low)
          new_rate = max(new_rate, current_target_rate)
  
          self._current_cluster_rate = new_rate
  
          if self._logger:
              self._logger.log(
                  scalyr_logging.DEBUG_LEVEL_3,
                  "RateLimiter: increase cluster rate %s (%s) x %s -> %s"
                  % (current_target_rate, actual_rate, self._increase_factor, new_rate),
              )
  
          self._reset_actual_rate()
  
      def _decrease_cluster_rate(self):
          """Decrease cluster rate but no lower than configured min and no higher than current cluster rate.
  
          If actual rate was calculated, then the new rate is calculated off the higher of actual & current target rate.
          Otherwise, new rate is calculated off the existing target rate.
  
          Note: this method is not thread-safe. Callers should pre-acquire cluster_rate_lock.
          """
          current_target_rate = self._current_cluster_rate
  
          actual_rate = self._get_actual_cluster_rate()
          if actual_rate is not None:
              # new rate should be based off the higher of actual & target rates
              new_rate = max(actual_rate, current_target_rate) * self._backoff_factor
          else:
              new_rate = current_target_rate * self._backoff_factor
  
          # new rate cannot be lower than configured lower bound
          new_rate = max(new_rate, self._min_cluster_rate)
  
          # new rate cannot be higher than existing rate (which may happen if actual rate was really high)
          new_rate = min(new_rate, current_target_rate)
  
          self._current_cluster_rate = new_rate
  
          if self._logger:
              self._logger.log(
                  scalyr_logging.DEBUG_LEVEL_3,
                  "RateLimiter: decrease cluster rate %s (%s) x %s -> %s"
                  % (current_target_rate, actual_rate, self._backoff_factor, new_rate),
              )
          self._reset_actual_rate()
  
      def _adjust_cluster_rate(self, success):
          """Adjust cluster rate, backing off on failure or increasing on enough consecutive successes.
  
          Any failure will cause backoff.
          A success will accumulate until self._consecutive_success_threshold is reached, at which point the rate will be
              increased according to self._increase_strategy.  When this happens, the success accumulator resets to 0
  
          @param success: Whether an operation was successful
          @type success: bool
          """
          self._cluster_rate_lock.acquire()
          try:
              if success:
                  # Success case: If current rate can be increased, do so based on strategy
                  self._consecutive_successes += 1
  
                  # not enough consecutive successes
                  if self._consecutive_successes < self._consecutive_success_threshold:
                      return
  
                  self._consecutive_successes = 0
                  if self._current_cluster_rate < self._max_cluster_rate:
                      if self._strategy == self.STRATEGY_RESET_THEN_MULTIPLY:
                          if self._current_cluster_rate < self._initial_cluster_rate:
                              self._current_cluster_rate = self._initial_cluster_rate
                          else:
                              self._increase_cluster_rate()
                      else:
                          self._increase_cluster_rate()
              else:
                  # Failure case: backoff
                  self._consecutive_successes = 0
  
                  # currently, a single failure will cause backoff
                  if self._current_cluster_rate > self._min_cluster_rate:
                      if self._strategy == self.STRATEGY_RESET_THEN_MULTIPLY:
                          if self._current_cluster_rate > self._initial_cluster_rate:
                              self._current_cluster_rate = self._initial_cluster_rate
                          else:
                              self._decrease_cluster_rate()
                      else:
                          self._decrease_cluster_rate()
          finally:
              self._cluster_rate_lock.release()
  
      def _get_next_ripe_time(self):
          """Get the absolute time for when this token becomes next available
  
          A token contains state representing the last time it was acquired.
          In order to achieve a uniform cluster-wide average rate of R, the next time is calculated as:
  
          current_time + delta
  
          where delta = random number between (0, 2 * T)
          where T is the per-agent interval or num_agents / cluster rate
          """
          agent_interval = float(self._num_agents) / self._current_cluster_rate
          delta = random.uniform(0, 2 * agent_interval)
          next_ripe = max(self._ripe_time, self._time()) + delta
          return next_ripe
  
      def _time(self):
          """Returns absolute time in UTC epoch seconds"""
          if self._fake_clock:
              return self._fake_clock.time()
          else:
              return time.time()
  
      def _simulate_acquire_token(self):
          self._test_mode_lock.acquire()
          try:
              while len(self._token_queue) == 0 or self._time() < self._ripe_time:
                  self._test_mode_lock.release()
                  self._fake_clock.simulate_waiting()
                  self._test_mode_lock.acquire()
  
              # Head token is ripe
              token = self._token_queue.popleft()
              self._ripe_time = self._get_next_ripe_time()
              return token
          finally:
              self._test_mode_lock.release()
  
      def _simulate_release_token(self, token, success):
          self._adjust_cluster_rate(success)
          self._test_mode_lock.acquire()
          try:
              self._token_queue.append(token)
          finally:
              self._test_mode_lock.release()
          self._fake_clock.advance_time(
              increment_by=self._ripe_time - self._fake_clock.time()
          )
  
      def acquire_token(self):
          """Acquires a token, blocking until a token becomes available.
  
          The primary state used to determine token availability is self._ripe_time. Until that time is reached,
          all client threads are blocked.
  
          @return: a token object
          @rtype: RateLimiterToken
          """
          self._lazy_adjust_min_max_rates()
  
          if self._fake_clock:
              return self._simulate_acquire_token()
  
          # block until a token is available from the token heap
          if self._logger:
              self._logger.log(
                  scalyr_logging.DEBUG_LEVEL_5,
                  "[%s] RateLimiter.acquire_token : acquire()    "
                  % threading.current_thread().name,
              )
          self._token_queue_cv.acquire()
  
          try:
              while len(self._token_queue) == 0 or self._time() < self._ripe_time:
                  # no tokens available so sleep for a while
                  if len(self._token_queue) == 0:
                      # queue contained no tokens so wait indefinitely
                      if self._logger:
                          self._logger.log(
                              scalyr_logging.DEBUG_LEVEL_5,
                              "[%s] RateLimiter.acquire_token : wait()"
                              % threading.current_thread().name,
                          )
                      self._token_queue_cv.wait()
                  else:
                      # queue contained at least one token so sleep until the head token becomes ripe
                      sleep_time = max(0, self._ripe_time - self._time())
                      if sleep_time > 0:
                          if self._logger:
                              self._logger.log(
                                  scalyr_logging.DEBUG_LEVEL_5,
                                  "[%s] RateLimiter.acquire_token : wait(%s)"
                                  % (threading.current_thread().name, sleep_time),
                              )
                          self._token_queue_cv.wait(sleep_time)
  
              # Head token is ripe.
              token = self._token_queue.popleft()
              if self._logger:
                  self._logger.log(
                      scalyr_logging.DEBUG_LEVEL_5,
                      "[%s] RateLimiter grant token %s at %.2f"
                      % (threading.current_thread().name, token, self._time()),
                  )
              self._ripe_time = self._get_next_ripe_time()
  
              if self._logger:
                  self._logger.log(
                      scalyr_logging.DEBUG_LEVEL_5,
                      "[%s] RateLimiter.acquire_token : returning token %s"
                      % (threading.current_thread().name, token),
                  )
              return token
  
          finally:
              self._token_queue_cv.release()
  
      def release_token(self, token, success):
          """Release a token back to the rate limiter.  Adjusts the rate based on success/failure and then re-inserts
          token back into the queue
  
          @param token: Token to release
          @param success: Whether or not the client operation was successful
  
          @type token: RateLimiterToken
          @type success: bool
          """
          if not isinstance(token, RateLimiterToken):
              raise TypeError(
                  "Rate limiting token must be of type %s" % type(RateLimiterToken)
              )
  
          self._record_actual_rate()
  
          if self._fake_clock:
              return self._simulate_release_token(token, success)
  
          if self._logger:
              self._logger.log(
                  scalyr_logging.DEBUG_LEVEL_5,
                  "[%s] RateLimiter.release_token : before acquire"
                  % threading.current_thread().name,
              )
          self._token_queue_cv.acquire()
  
          try:
              # potentially adjust rate depending on reported outcome
              self._adjust_cluster_rate(success)
  
              # re-insert token
              self._token_queue.append(token)
  
              # awaken threads waiting for tokens
              if self._logger:
                  self._logger.log(
                      scalyr_logging.DEBUG_LEVEL_5,
                      "[%s] RateLimiter.release_token : notifyAll()"
                      % threading.current_thread().name,
                  )
              self._token_queue_cv.notifyAll()
  
          finally:
              self._token_queue_cv.release()
  