add redis lock

src/main/java/com/taoyuanx/securitydemo/security/lock/RedisLock.java

@@ -0,0 +1,102 @@
+package com.taoyuanx.securitydemo.security.lock;
+
+import com.taoyuanx.securitydemo.security.ratelimit.AbstractRateLimiter;
+import com.taoyuanx.securitydemo.utils.SpringContextUtil;
+import com.taoyuanx.securitydemo.utils.StringIntUtil;
+import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.core.io.ClassPathResource;
+import org.springframework.data.redis.core.StringRedisTemplate;
+import org.springframework.data.redis.core.script.DefaultRedisScript;
+import org.springframework.data.redis.core.script.RedisScript;
+import org.springframework.scripting.support.ResourceScriptSource;
+
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+
+/**
+ * @author dushitaoyuan
+ * @desc redis 限流实现
+ * 复制 spring-cloud-gateway,可与spring-boot,spring-cloud等分离使用
+ * @date 2019/9/5
+ */
+@Slf4j
+public class RedisLock {
+
+    @Getter
+    private String lockKey;
+    @Getter
+    private String lockValue;
+
+
+    public RedisLock(String lockKey) {
+        this(lockKey, UUID.randomUUID().toString());
+    }
+
+    public RedisLock(String lockKey, String lockValue) {
+        init();
+        this.lockKey = lockKey;
+        this.lockValue = lockValue;
+    }
+
+
+    public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
+        boolean lock = doTryLock();
+        if (!lock) {
+            unit.sleep(time);
+        }
+        return lock;
+    }
+
+    public boolean tryLock() {
+        return doTryLock();
+    }
+
+
+    public void unlock() {
+        String result = redisTemplate.execute(unLockScript, Arrays.asList(lockKey), lockValue);
+        if (FAILED.equals(result)) {
+            log.warn("释放锁失败,lock_value 不匹配 my_lock_value:[{}]", lockValue);
+        }
+    }
+
+    private boolean doTryLock() {
+        String result = redisTemplate.execute(lockScript, Arrays.asList(lockKey), lockValue, String.valueOf(LOCK_MAX_EXPIRE_SECONDS));
+        if (SUCCESSS.equals(result)) {
+            return true;
+        }
+        return false;
+    }
+
+    private static StringRedisTemplate redisTemplate;
+    private static RedisScript<String> lockScript;
+    private static RedisScript<String> unLockScript;
+    private static String SUCCESSS = "1", FAILED = "0";
+    /**
+     * 最大锁定5分钟
+     */
+    private static Long LOCK_MAX_EXPIRE_SECONDS = 300L;
+
+    private void init() {
+        if (Objects.isNull(lockScript)) {
+            DefaultRedisScript script = new DefaultRedisScript();
+            script.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/redis_lock.lua")));
+            script.setResultType(String.class);
+            lockScript = script;
+        }
+        if (Objects.isNull(unLockScript)) {
+            DefaultRedisScript script = new DefaultRedisScript();
+            script.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/redis_unlock.lua")));
+            script.setResultType(String.class);
+            unLockScript = script;
+        }
+        if (Objects.isNull(redisTemplate)) {
+            redisTemplate = SpringContextUtil.getBean(StringRedisTemplate.class);
+        }
+    }
+}

src/main/resources/META-INF/rate_limiter_count.lua

@@ -1,7 +1,7 @@
 --判断资源归0标记是否存在
 --标记实现,可通过 布隆过滤和bitmap实现 注意redis支持情况
 local bit_key_offset=tonumber(ARGV[3])
-redis.log(redis.LOG_WARNING, "bit_key_offset " .. ARGV[3])
+--redis.log(redis.LOG_WARNING, "bit_key_offset " .. ARGV[3])
 local zero_flag =redis.call("GETBIT", KEYS[2],bit_key_offset)
 if zero_flag == 1 then
     return -1

src/main/resources/META-INF/redis_lock.lua

@@ -0,0 +1,23 @@
+-- get a lock
+-- 获取锁成功，则返回 1
+
+--redis.log(redis.LOG_WARNING, "key " .. KEYS[1])
+--redis.log(redis.LOG_WARNING, "ARGV1 " .. ARGV[1])
+--redis.log(redis.LOG_WARNING, "ARGV " .. ARGV[2])
+local lock_key = KEYS[1]
+local lock_value = ARGV[1]
+local ttl = tonumber(ARGV[2])
+local result = redis.call('setnx', lock_key, lock_value)
+--redis.log(redis.LOG_WARNING, "result " .. result)
+if result == 1 then
+    redis.call('expire', lock_key, ttl)
+    return '1'
+else
+    local value = redis.call('get', lock_key)
+    -- lock_value 一致可重入
+    if (value == lock_value) then
+        result = '1';
+        redis.call('expire', lock_key, ttl)
+    end
+end
+return tostring(result)

src/main/resources/META-INF/redis_unlock.lua

@@ -0,0 +1,12 @@
+-- unlock
+local lock_key = KEYS[1]
+local lock_value = ARGV[1]
+--[[--redis.log(redis.LOG_WARNING, "key " .. KEYS[1])
+--redis.log(redis.LOG_WARNING, "ARGV1 " .. ARGV[1])]]
+local result = redis.call('get', lock_key)
+--redis.log(redis.LOG_WARNING, "result " .. result)
+if result == lock_value then
+    redis.call('del', lock_key)
+    return '1';
+end
+return '0'

src/main/resources/application.properties

@@ -1,9 +1,9 @@
 spring.profiles.active=@profiles.active@
 
-spring.redis.database=2
-spring.redis.host=192.168.30.211
+spring.redis.database=0
+spring.redis.host=172.16.0.32
 spring.redis.port=6379
-spring.redis.password=
+spring.redis.password=guoruiredis
 
 server.servlet.session.timeout=1800
 

src/main/resources/logback-spring.xml

@@ -2,7 +2,7 @@
 
 <configuration scan="true" scanPeriod="10 seconds">
     <springProperty name="LOG_HOME" source="logging.path" defaultValue="./logs"/>
-    <property name="LOG_PREFIX" value="seed-boot"/>
+    <property name="LOG_PREFIX" value="demo"/>
 
     <conversionRule conversionWord="clr" converterClass="org.springframework.boot.logging.logback.ColorConverter"/>
     <conversionRule conversionWord="wex" converterClass="org.springframework.boot.logging.logback.WhitespaceThrowableProxyConverter"/>
@@ -114,7 +114,7 @@
 
 
         <springProfile name="dev">
-                <root level="DEBUG">
+                <root level="INFO">
                     <appender-ref ref="CONSOLE"/>
                 </root>
         </springProfile>

src/test/java/com/taoyuanx/securitydemo/RedisLockTest.java

@@ -0,0 +1,62 @@
+package com.taoyuanx.securitydemo;
+
+import com.taoyuanx.securitydemo.controller.BussinessController;
+import com.taoyuanx.securitydemo.security.lock.RedisLock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.test.context.junit4.SpringRunner;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+@RunWith(SpringRunner.class)
+@SpringBootTest
+public class RedisLockTest {
+    @Test
+    public void lockTestOne() throws InterruptedException {
+        RedisLock redisLock = new RedisLock("lock:xxx","11");
+        if (redisLock.tryLock()) {
+            System.out.println("获取到锁了");
+            redisLock.unlock();
+
+        }
+    }
+    @Test
+    public void lockTest() throws InterruptedException {
+        RedisLock redisLock = new RedisLock("lock:xxx");
+        for (int i = 0; i < 100; i++) {
+            executorService.submit(() -> {
+                try {
+                    if (redisLock.tryLock(3,TimeUnit.SECONDS)) {
+                        count++;
+                        System.out.println("执行");
+                        redisLock.unlock();
+                    }
+                } catch (InterruptedException e) {
+                    e.printStackTrace();
+                }
+                System.out.println("未执行");
+            });
+        }
+        System.out.println(count);
+        executorService.shutdown();
+        while (!executorService.awaitTermination(3, TimeUnit.SECONDS)){
+            System.out.println("休眠3秒");
+        }
+        System.out.println("结束");
+    }
+
+    private ExecutorService executorService;
+    private Integer count = 0;
+
+    @Before
+    public void before() throws Exception {
+        executorService = Executors.newFixedThreadPool(10);
+    }
+
+
+}