content: /*
   * Copyright 1999-2018 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.controller;
  
  import com.alibaba.csp.sentinel.node.Node;
  import com.alibaba.csp.sentinel.node.OccupyTimeoutProperty;
  import com.alibaba.csp.sentinel.slots.block.RuleConstant;
  import com.alibaba.csp.sentinel.slots.block.flow.PriorityWaitException;
  import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingController;
  import com.alibaba.csp.sentinel.util.TimeUtil;
  
  /**
   * Default throttling controller (immediately reject strategy).
   *
   * @author jialiang.linjl
   * @author Eric Zhao
   */
  public class DefaultController implements TrafficShapingController {
  
      private static final int DEFAULT_AVG_USED_TOKENS = 0;
  
      private double count;
      private int grade;
  
      public DefaultController(double count, int grade) {
          this.count = count;
          this.grade = grade;
      }
  
      @Override
      public boolean canPass(Node node, int acquireCount) {
          return canPass(node, acquireCount, false);
      }
  
      @Override
      public boolean canPass(Node node, int acquireCount, boolean prioritized) {
          int curCount = avgUsedTokens(node);
          if (curCount + acquireCount > count) {
              if (prioritized && grade == RuleConstant.FLOW_GRADE_QPS) {
                  long currentTime;
                  long waitInMs;
                  currentTime = TimeUtil.currentTimeMillis();
                  waitInMs = node.tryOccupyNext(currentTime, acquireCount, count);
                  if (waitInMs < OccupyTimeoutProperty.getOccupyTimeout()) {
                      node.addWaitingRequest(currentTime + waitInMs, acquireCount);
                      node.addOccupiedPass(acquireCount);
                      sleep(waitInMs);
  
                      // PriorityWaitException indicates that the request will pass after waiting for {@link @waitInMs}.
                      throw new PriorityWaitException(waitInMs);
                  }
              }
              return false;
          }
          return true;
      }
  
      private int avgUsedTokens(Node node) {
          if (node == null) {
              return DEFAULT_AVG_USED_TOKENS;
          }
          return grade == RuleConstant.FLOW_GRADE_THREAD ? node.curThreadNum() : (int)(node.passQps());
      }
  
      private void sleep(long timeMillis) {
          try {
              Thread.sleep(timeMillis);
          } catch (InterruptedException e) {
              // Ignore.
          }
      }
  }
  