From de50c4dd5a5eddf4bf2d6aea5174727bc17c5ba2 Mon Sep 17 00:00:00 2001
From: tanghc <thc8775@163.com>
Date: Fri, 6 Sep 2019 19:46:14 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=88=86=E5=B8=83=E5=BC=8F?=
 =?UTF-8?q?=E9=99=90=E6=B5=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../gatewaycommon/bean/ConfigLimitDto.java    |  2 +-
 .../limit/RedisLimitManager.java              | 79 +++++++++++++++++++
 .../src/main/resources/sop/limit.lua          | 10 +++
 3 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/limit/RedisLimitManager.java
 create mode 100644 sop-common/sop-gateway-common/src/main/resources/sop/limit.lua

diff --git a/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/bean/ConfigLimitDto.java b/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/bean/ConfigLimitDto.java
index 92ed0866..c76e6ac2 100644
--- a/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/bean/ConfigLimitDto.java
+++ b/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/bean/ConfigLimitDto.java
@@ -37,7 +37,7 @@ public class ConfigLimitDto {
     /** 服务id, 数据库字段：service_id */
     private String serviceId;
 
-    /** 限流策略，1：漏桶策略，2：令牌桶策略, 数据库字段：limit_type */
+    /** 限流策略，1：窗口策略，2：令牌桶策略, 数据库字段：limit_type */
     private Byte limitType;
 
     /** 每秒可处理请求数, 数据库字段：exec_count_per_second */
diff --git a/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/limit/RedisLimitManager.java b/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/limit/RedisLimitManager.java
new file mode 100644
index 00000000..88a2972e
--- /dev/null
+++ b/sop-common/sop-gateway-common/src/main/java/com/gitee/sop/gatewaycommon/limit/RedisLimitManager.java
@@ -0,0 +1,79 @@
+package com.gitee.sop.gatewaycommon.limit;
+
+import com.gitee.sop.gatewaycommon.bean.ConfigLimitDto;
+import com.gitee.sop.gatewaycommon.bean.SopConstants;
+import org.apache.commons.codec.digest.DigestUtils;
+import org.apache.commons.io.IOUtils;
+import org.springframework.core.io.ClassPathResource;
+import org.springframework.data.redis.core.RedisTemplate;
+import org.springframework.data.redis.core.StringRedisTemplate;
+import org.springframework.data.redis.core.script.RedisScript;
+import org.springframework.util.Assert;
+
+import java.util.Collections;
+
+/**
+ * 基于redis限流管理
+ *
+ * @author tanghc
+ */
+public class RedisLimitManager extends DefaultLimitManager {
+
+    /**
+     * 限流脚本
+     */
+    private static final String DEFAULT_LIMIT_LUA_FILE_PATH = "/sop/limit.lua";
+
+    private static final Long REDIS_SUCCESS = 1L;
+
+    private StringRedisTemplate redisTemplate;
+    private String limitScript;
+    private String limitScriptSha1;
+
+    public RedisLimitManager(@SuppressWarnings("rawtypes") RedisTemplate redisTemplate) {
+        Assert.notNull(redisTemplate, "redisTemplate不能为null");
+        this.redisTemplate = new StringRedisTemplate(redisTemplate.getConnectionFactory());
+        ClassPathResource limitLua = new ClassPathResource(getLimitLuaFilePath());
+        try {
+            this.limitScript = IOUtils.toString(limitLua.getInputStream(), SopConstants.UTF8);
+            this.limitScriptSha1 = DigestUtils.sha1Hex(this.limitScript);
+        } catch (Exception e) {
+            throw new RuntimeException("读取脚本文件失败，脚本路径:" + getLimitLuaFilePath(), e);
+        }
+    }
+
+    public String getLimitLuaFilePath() {
+        return DEFAULT_LIMIT_LUA_FILE_PATH;
+    }
+
+    @Override
+    public boolean acquire(ConfigLimitDto routeConfig) {
+        String key = "sop:lmt:" + routeConfig.getRouteId();
+        int limitCount = routeConfig.getExecCountPerSecond();
+
+        Object result = redisTemplate.execute(
+                new RedisScript<Long>() {
+                    @Override
+                    public String getSha1() {
+                        return limitScriptSha1;
+                    }
+
+                    @Override
+                    public Class<Long> getResultType() {
+                        return Long.class;
+                    }
+
+                    @Override
+                    public String getScriptAsString() {
+                        return limitScript;
+                    }
+                },
+                // KEYS[1] key
+                Collections.singletonList(key),
+                // ARGV[1] limit
+                String.valueOf(limitCount)
+        );
+        return REDIS_SUCCESS.equals(result);
+    }
+
+}
diff --git a/sop-common/sop-gateway-common/src/main/resources/sop/limit.lua b/sop-common/sop-gateway-common/src/main/resources/sop/limit.lua
new file mode 100644
index 00000000..1a5be153
--- /dev/null
+++ b/sop-common/sop-gateway-common/src/main/resources/sop/limit.lua
@@ -0,0 +1,10 @@
+local key = KEYS[1] --限流KEY（一秒一个）
+local limit = tonumber(ARGV[1]) --限流大小
+local current = tonumber(redis.call("get", key)) or 0
+if current + 1 > limit then --如果超出限流大小
+   return 0
+else  --请求数+1，并设置2秒过期
+   redis.call("INCRBY", key,"1")
+   redis.call("expire", key,"2")
+   return 1
+end
\ No newline at end of file