package cn.dijia478.util;

import java.time.LocalTime;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 滑動時間窗口限流工具
 * 本限流工具只適用于單機(jī)版，如果想要做全局限流，可以按本程序的思想，用redis的List結(jié)構(gòu)去實現(xiàn)
 *
 * @author dijia478
 * @date 2020-10-13 10:53
 */
public class SlideWindow {

  /** 隊列id和隊列的映射關(guān)系，隊列里面存儲的是每一次通過時候的時間戳，這樣可以使得程序里有多個限流隊列 */
  private volatile static Map<String, List<Long>> MAP = new ConcurrentHashMap<>();

  private SlideWindow() {}

  public static void main(String[] args) throws InterruptedException {
    while (true) {
      // 任意10秒內(nèi)，只允許2次通過
      System.out.println(LocalTime.now().toString() + SlideWindow.isGo("ListId", 2, 10000L));
      // 睡眠0-10秒
      Thread.sleep(1000 * new Random().nextInt(10));
    }
  }

  /**
   * 滑動時間窗口限流算法
   * 在指定時間窗口，指定限制次數(shù)內(nèi)，是否允許通過
   *
   * @param listId   隊列id
   * @param count   限制次數(shù)
   * @param timeWindow 時間窗口大小
   * @return 是否允許通過
   */
  public static synchronized boolean isGo(String listId, int count, long timeWindow) {
    // 獲取當(dāng)前時間
    long nowTime = System.currentTimeMillis();
    // 根據(jù)隊列id，取出對應(yīng)的限流隊列，若沒有則創(chuàng)建
    List<Long> list = MAP.computeIfAbsent(listId, k -> new LinkedList<>());
    // 如果隊列還沒滿，則允許通過，并添加當(dāng)前時間戳到隊列開始位置
    if (list.size() < count) {
      list.add(0, nowTime);
      return true;
    }

    // 隊列已滿（達(dá)到限制次數(shù)），則獲取隊列中最早添加的時間戳
    Long farTime = list.get(count - 1);
    // 用當(dāng)前時間戳 減去 最早添加的時間戳
    if (nowTime - farTime <= timeWindow) {
      // 若結(jié)果小于等于timeWindow，則說明在timeWindow內(nèi)，通過的次數(shù)大于count
      // 不允許通過
      return false;
    } else {
      // 若結(jié)果大于timeWindow，則說明在timeWindow內(nèi)，通過的次數(shù)小于等于count
      // 允許通過，并刪除最早添加的時間戳，將當(dāng)前時間添加到隊列開始位置
      list.remove(count - 1);
      list.add(0, nowTime);
      return true;
    }
  }

}