com.github.edgar615.util.vertx.redis.ratelimit.MultSlidingWIndowRateLimit Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vertx-util-redis Show documentation
The newest version!
package com.github.edgar615.util.vertx.redis.ratelimit;

import com.github.edgar615.util.vertx.redis.AbstractLuaEvaluator;
import io.vertx.core.AsyncResult;
import io.vertx.core.Future;
import io.vertx.core.Handler;
import io.vertx.core.Vertx;
import io.vertx.core.json.JsonArray;
import io.vertx.redis.RedisClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * Created by edgar on 17-5-29.
 */
public class MultSlidingWIndowRateLimit extends AbstractLuaEvaluator {
  private static final Logger LOGGER = LoggerFactory.getLogger(MultSlidingWIndowRateLimit.class);

  public MultSlidingWIndowRateLimit(Vertx vertx, RedisClient redisClient, Future completed) {
    super(vertx, redisClient, "com/github/edgar615/util/redis/lua/multi_sliding_window_ratelimit.lua", completed);
  }

  /**
   * 限流
   *
   * @param rules  限流集合，必须包含三个元素:subject，limit,interval
   * @param handler 　回调
   */
  public void rateLimit(List rules,
                        Handler>          handler) {
    JsonArray limitArray;
    try {
      limitArray = checkArgument(rules);
    } catch (Exception e) {
      handler.handle(Future.failedFuture(e));
      return;
    }
    List keys = new ArrayList<>();
    List args = new ArrayList<>();
    args.add(limitArray.encode());
    args.add(Instant.now().getEpochSecond() + "");
    evaluate(keys, args, ar -> {
      if (ar.failed()) {
        LOGGER.error("rateLimit failed", ar.cause());
        handler.handle(Future.failedFuture("rateLimit failed"));
        return;
      }
      List subjects = rules.stream()
              .map(l -> l.getSubject())
              .collect(Collectors.toList());
      RateLimitUtils.createResult(ar.result(), subjects, handler);
    });
  }

  private JsonArray checkArgument(List limits) {
    if (limits.size() == 0) {
      throw new IllegalArgumentException("limits cannot empty");
    }
    JsonArray limitArray = new JsonArray();
    for (int i = 0; i < limits.size(); i++) {
      SlidingWindowRateLimitRule limit = limits.get(i);
      try {
        limitArray.add(new JsonArray().add(limit.getSubject())
            .add(limit.getLimit())
            .add(limit.getInterval())
            .add(limit.getPrecision()));
      } catch (Exception e) {
        throw new IllegalArgumentException(e);
      }
    }
    return limitArray;
  }

}