content: /*
   * Copyright 1999-2024 Alibaba Group Holding Ltd.
   *
   * 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.
   */
  package com.alibaba.csp.sentinel.slots.block.flow.param;
  
  import com.alibaba.csp.sentinel.cluster.ClusterStateManager;
  import com.alibaba.csp.sentinel.cluster.TokenResult;
  import com.alibaba.csp.sentinel.cluster.TokenResultStatus;
  import com.alibaba.csp.sentinel.cluster.TokenService;
  import com.alibaba.csp.sentinel.cluster.client.TokenClientProvider;
  import com.alibaba.csp.sentinel.cluster.server.EmbeddedClusterTokenServerProvider;
  import com.alibaba.csp.sentinel.log.RecordLog;
  import com.alibaba.csp.sentinel.slotchain.ResourceWrapper;
  import com.alibaba.csp.sentinel.slots.block.RuleConstant;
  import com.alibaba.csp.sentinel.slots.statistic.cache.CacheMap;
  import com.alibaba.csp.sentinel.util.TimeUtil;
  
  import java.lang.reflect.Array;
  import java.time.format.DateTimeFormatter;
  import java.util.*;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.concurrent.atomic.AtomicReference;
  
  /**
   * Rule checker for parameter flow control.
   *
   * @author Eric Zhao
   * @since 0.2.0
   */
  public final class ParamFlowChecker {
  
      public static boolean passCheck(ResourceWrapper resourceWrapper, /*@Valid*/ ParamFlowRule rule, /*@Valid*/ int count,
                                      Object... args) {
          if (args == null) {
              return true;
          }
  
          int paramIdx = rule.getParamIdx();
          if (args.length <= paramIdx) {
              return true;
          }
  
          // Get parameter value.
          Object value = args[paramIdx];
  
          // Assign value with the result of paramFlowKey method
          if (value instanceof ParamFlowArgument) {
              value = ((ParamFlowArgument) value).paramFlowKey();
          }
          // If value is null, then pass
          if (value == null) {
              return true;
          }
  
          if (rule.isClusterMode() && rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {
              return passClusterCheck(resourceWrapper, rule, count, value);
          }
  
          return passLocalCheck(resourceWrapper, rule, count, value);
      }
  
      private static boolean passLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                            Object value) {
          try {
              if (Collection.class.isAssignableFrom(value.getClass())) {
                  for (Object param : ((Collection) value)) {
                      if (!passSingleValueCheck(resourceWrapper, rule, count, param)) {
                          return false;
                      }
                  }
              } else if (value.getClass().isArray()) {
                  int length = Array.getLength(value);
                  for (int i = 0; i < length; i++) {
                      Object param = Array.get(value, i);
                      if (!passSingleValueCheck(resourceWrapper, rule, count, param)) {
                          return false;
                      }
                  }
              } else {
                  return passSingleValueCheck(resourceWrapper, rule, count, value);
              }
          } catch (Throwable e) {
              RecordLog.warn("[ParamFlowChecker] Unexpected error", e);
          }
  
          return true;
      }
  
      static boolean passSingleValueCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                          Object value) {
          if (rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {
              if (rule.getControlBehavior() == RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER) {
                  return passThrottleLocalCheck(resourceWrapper, rule, acquireCount, value);
              } else {
                  return passDefaultLocalCheck(resourceWrapper, rule, acquireCount, value);
              }
          } else if (rule.getGrade() == RuleConstant.FLOW_GRADE_THREAD) {
              Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
              long threadCount = getParameterMetric(resourceWrapper).getThreadCount(rule.getParamIdx(), value);
              if (exclusionItems.contains(value)) {
                  int itemThreshold = rule.getParsedHotItems().get(value);
                  return ++threadCount <= itemThreshold;
              }
              long threshold = (long) rule.getCount();
              return ++threadCount <= threshold;
          }
  
          return true;
      }
  
      static boolean passDefaultLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                           Object value) {
          ParameterMetric metric = getParameterMetric(resourceWrapper);
          CacheMap<Object, AtomicReference<TokenUpdateStatus>> tokenCounters = metric == null ? null : metric.getRuleStampedTokenCounter(rule);
  
          DateTimeFormatter dtf = DateTimeFormatter.ISO_DATE_TIME;
          if (tokenCounters == null) {
              return true;
          }
  
          // Calculate max token count (threshold)
          Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
          long tokenCount = (long) rule.getCount();
          if (exclusionItems.contains(value)) {
              tokenCount = rule.getParsedHotItems().get(value);
          }
  
          if (tokenCount == 0) {
              return false;
          }
  
          long maxCount = tokenCount + rule.getBurstCount();
          if (acquireCount > maxCount) {
              return false;
          }
  
          while (true) {
              long currentTime = TimeUtil.currentTimeMillis();
  
              AtomicReference<TokenUpdateStatus> atomicLastStatus = tokenCounters.putIfAbsent(value, new AtomicReference<>(
                      new TokenUpdateStatus(currentTime, maxCount - acquireCount)
              ));
              if (atomicLastStatus == null) {
                  // Token never added, just replenish the tokens and consume {@code acquireCount} immediately.
                  return true;
              }
  
              // Calculate the time duration since last token was added.
              TokenUpdateStatus lastStatus = atomicLastStatus.get();
              long passTime = currentTime - lastStatus.getLastAddTokenTime();
              // A simplified token bucket algorithm that will replenish the tokens only when statistic window has passed.
              long newQps;
              if (passTime > rule.getDurationInSec() * 1000) {
                  long restQps = lastStatus.getRestQps();
                  long toAddCount = (passTime * tokenCount) / (rule.getDurationInSec() * 1000);
                  newQps = toAddCount + restQps > maxCount ? (maxCount - acquireCount)
                          : (restQps + toAddCount - acquireCount);
  
                  if (newQps < 0) {
                      return false;
                  }
                  TokenUpdateStatus newStatus = new TokenUpdateStatus(currentTime, newQps);
                  if (atomicLastStatus.compareAndSet(lastStatus, newStatus)) {
                      return true;
                  }
                  Thread.yield();
              } else {
                  newQps = lastStatus.getRestQps() - acquireCount;
                  if (newQps >= 0) {
                      TokenUpdateStatus newStatus = new TokenUpdateStatus(lastStatus.getLastAddTokenTime(), newQps);
                      if (atomicLastStatus.compareAndSet(lastStatus, newStatus)) {
                          return true;
                      }
                  } else {
                      return false;
                  }
  
                  Thread.yield();
              }
          }
      }
  
      static boolean passThrottleLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                            Object value) {
          ParameterMetric metric = getParameterMetric(resourceWrapper);
          CacheMap<Object, AtomicLong> timeRecorderMap = metric == null ? null : metric.getRuleTimeCounter(rule);
          if (timeRecorderMap == null) {
              return true;
          }
  
          // Calculate max token count (threshold)
          Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
          long tokenCount = (long) rule.getCount();
          if (exclusionItems.contains(value)) {
              tokenCount = rule.getParsedHotItems().get(value);
          }
  
          if (tokenCount == 0) {
              return false;
          }
  
          long costTime = Math.round(1.0 * 1000 * acquireCount * rule.getDurationInSec() / tokenCount);
          while (true) {
              long currentTime = TimeUtil.currentTimeMillis();
              AtomicLong timeRecorder = timeRecorderMap.putIfAbsent(value, new AtomicLong(currentTime));
              if (timeRecorder == null) {
                  return true;
              }
              //AtomicLong timeRecorder = timeRecorderMap.get(value);
              long lastPassTime = timeRecorder.get();
              long expectedTime = lastPassTime + costTime;
  
              if (expectedTime <= currentTime || expectedTime - currentTime < rule.getMaxQueueingTimeMs()) {
                  AtomicLong lastPastTimeRef = timeRecorderMap.get(value);
                  if (lastPastTimeRef.compareAndSet(lastPassTime, currentTime)) {
                      long waitTime = expectedTime - currentTime;
                      if (waitTime > 0) {
                          lastPastTimeRef.set(expectedTime);
                          try {
                              TimeUnit.MILLISECONDS.sleep(waitTime);
                          } catch (InterruptedException e) {
                              RecordLog.warn("passThrottleLocalCheck: wait interrupted", e);
                          }
                      }
                      return true;
                  } else {
                      Thread.yield();
                  }
              } else {
                  return false;
              }
          }
      }
  
      private static ParameterMetric getParameterMetric(ResourceWrapper resourceWrapper) {
          // Should not be null.
          return ParameterMetricStorage.getParamMetric(resourceWrapper);
      }
  
      @SuppressWarnings("unchecked")
      private static Collection<Object> toCollection(Object value) {
          if (value instanceof Collection) {
              return (Collection<Object>) value;
          } else if (value.getClass().isArray()) {
              List<Object> params = new ArrayList<Object>();
              int length = Array.getLength(value);
              for (int i = 0; i < length; i++) {
                  Object param = Array.get(value, i);
                  params.add(param);
              }
              return params;
          } else {
              return Collections.singletonList(value);
          }
      }
  
      private static boolean passClusterCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                              Object value) {
          try {
              Collection<Object> params = toCollection(value);
  
              TokenService clusterService = pickClusterService();
              if (clusterService == null) {
                  // No available cluster client or server, fallback to local or
                  // pass in need.
                  return fallbackToLocalOrPass(resourceWrapper, rule, count, params);
              }
  
              TokenResult result = clusterService.requestParamToken(rule.getClusterConfig().getFlowId(), count, params);
              switch (result.getStatus()) {
                  case TokenResultStatus.OK:
                      return true;
                  case TokenResultStatus.BLOCKED:
                      return false;
                  default:
                      return fallbackToLocalOrPass(resourceWrapper, rule, count, params);
              }
          } catch (Throwable ex) {
              RecordLog.warn("[ParamFlowChecker] Request cluster token for parameter unexpected failed", ex);
              return fallbackToLocalOrPass(resourceWrapper, rule, count, value);
          }
      }
  
      private static boolean fallbackToLocalOrPass(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                                   Object value) {
          if (rule.getClusterConfig().isFallbackToLocalWhenFail()) {
              return passLocalCheck(resourceWrapper, rule, count, value);
          } else {
              // The rule won't be activated, just pass.
              return true;
          }
      }
  
      private static TokenService pickClusterService() {
          if (ClusterStateManager.isClient()) {
              return TokenClientProvider.getClient();
          }
          if (ClusterStateManager.isServer()) {
              return EmbeddedClusterTokenServerProvider.getServer();
          }
          return null;
      }
  
      private ParamFlowChecker() {
      }
  }
  