增加限流功能

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/pom.xml

@@ -0,0 +1,23 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>hsweb-concurrent-rate-limiter</artifactId>
+        <groupId>org.hswebframework.web</groupId>
+        <version>3.0.4-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>hsweb-concurrent-rate-limiter-api</artifactId>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+            <version>20.0</version>
+            <optional>true</optional>
+        </dependency>
+    </dependencies>
+
+</project>

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/src/main/java/org/hswebframework/web/concurrent/AbstractRateLimiterManager.java

@@ -0,0 +1,26 @@
+package org.hswebframework.web.concurrent;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+public abstract class AbstractRateLimiterManager implements RateLimiterManager {
+    private final Map<String, RateLimiter> counterStore = new HashMap<>(128);
+
+    protected abstract RateLimiter createRateLimiter(String key, double permits, TimeUnit timeUnit);
+
+    @Override
+    public RateLimiter getRateLimiter(String key, double permits, TimeUnit timeUnit) {
+        RateLimiter counter = counterStore.get(key);
+        if (counter != null) {
+            return counter;
+        }
+        synchronized (counterStore) {
+            return counterStore.computeIfAbsent(key, k -> createRateLimiter(key, permits, timeUnit));
+        }
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/src/main/java/org/hswebframework/web/concurrent/GuavaRateLimiterManager.java

@@ -0,0 +1,27 @@
+package org.hswebframework.web.concurrent;
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+public class GuavaRateLimiterManager extends AbstractRateLimiterManager {
+    @Override
+    protected RateLimiter createRateLimiter(String key, double permits, TimeUnit timeUnit) {
+        long seconds = timeUnit.toSeconds(1);
+        double permitsPerSecond = permits;
+
+        if (seconds > 0) {
+            permitsPerSecond = permits / timeUnit.toSeconds(1);
+        } else {
+            if (timeUnit == TimeUnit.MILLISECONDS) {
+                permitsPerSecond = permits / 1000D;
+            }
+        }
+
+        com.google.common.util.concurrent.RateLimiter rateLimiter =
+                com.google.common.util.concurrent.RateLimiter.create(permitsPerSecond);
+        return rateLimiter::tryAcquire;
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/src/main/java/org/hswebframework/web/concurrent/RateLimiter.java

@@ -0,0 +1,17 @@
+package org.hswebframework.web.concurrent;
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+public interface RateLimiter {
+
+    boolean tryAcquire(int permits, long timeout, TimeUnit timeUnit);
+
+    default boolean tryAcquire(long timeout, TimeUnit unit) {
+        return tryAcquire(1, timeout, unit);
+    }
+
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/src/main/java/org/hswebframework/web/concurrent/RateLimiterManager.java

@@ -0,0 +1,12 @@
+package org.hswebframework.web.concurrent;
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+public interface RateLimiterManager {
+
+    RateLimiter getRateLimiter(String key, double permits, TimeUnit timeUnit);
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/src/main/java/org/hswebframework/web/concurrent/annotation/RateLimiter.java

@@ -0,0 +1,39 @@
+package org.hswebframework.web.concurrent.annotation;
+
+import java.lang.annotation.*;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * 开启限流
+ *
+ * @author zhouhao
+ * @see org.hswebframework.web.concurrent.RateLimiter
+ * @since 3.0.4
+ */
+@Target({ElementType.TYPE, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@Inherited
+@Documented
+public @interface RateLimiter {
+
+    /**
+     * key,支持spel表达式: ${#username+'login'} .默认以当前注解的方法为key
+     *
+     * @return 限流Key
+     */
+    String[] key() default {};
+
+    /**
+     * @return 时间单位内允许访问次数, 如:每秒100次
+     */
+    double permits() default 100D;
+
+    /**
+     * @return 时间单位, 支持毫秒及以上的时间单位
+     */
+    TimeUnit timeUnit() default TimeUnit.SECONDS;
+
+    long acquire() default 10;
+
+    TimeUnit acquireTimeUnit() default TimeUnit.SECONDS;
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-api/src/test/java/org/hswebframework/web/concurrent/GuavaRateLimiterManagerTest.java

@@ -0,0 +1,28 @@
+package org.hswebframework.web.concurrent;
+
+import lombok.SneakyThrows;
+import org.junit.Test;
+
+import java.util.concurrent.TimeUnit;
+
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+public class GuavaRateLimiterManagerTest {
+    GuavaRateLimiterManager manager = new GuavaRateLimiterManager();
+
+    @SneakyThrows
+    @Test
+    public void testRateLimiter() {
+        RateLimiter limiter = manager.getRateLimiter("test", 1, TimeUnit.SECONDS);
+        for (int i = 0; i < 100; i++) {
+            if (!limiter.tryAcquire(10, TimeUnit.SECONDS)) {
+                throw new UnsupportedOperationException();
+            }
+            System.out.println(i + ":" + System.currentTimeMillis());
+        }
+
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/pom.xml

@@ -0,0 +1,52 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>hsweb-concurrent-rate-limiter</artifactId>
+        <groupId>org.hswebframework.web</groupId>
+        <version>3.0.4-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>hsweb-concurrent-rate-limiter-starter</artifactId>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.hswebframework.web</groupId>
+            <artifactId>hsweb-concurrent-rate-limiter-api</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.hswebframework.web</groupId>
+            <artifactId>hsweb-commons-utils</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.hswebframework.web</groupId>
+            <artifactId>hsweb-authorization-api</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+            <version>20.0</version>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.hswebframework.web</groupId>
+            <artifactId>hsweb-tests</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework</groupId>
+            <artifactId>spring-aspects</artifactId>
+        </dependency>
+    </dependencies>
+</project>

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/src/main/java/org/hswebframework/web/concurent/RateLimiterAopAdvisor.java

@@ -0,0 +1,69 @@
+package org.hswebframework.web.concurent;
+
+import lombok.extern.slf4j.Slf4j;
+import org.hswebframework.web.ExpressionUtils;
+import org.hswebframework.web.authorization.Authentication;
+import org.hswebframework.web.concurrent.RateLimiterManager;
+import org.hswebframework.web.concurrent.annotation.RateLimiter;
+import org.springframework.aop.MethodBeforeAdvice;
+import org.springframework.aop.support.StaticMethodMatcherPointcutAdvisor;
+import org.springframework.core.LocalVariableTableParameterNameDiscoverer;
+import org.springframework.core.ParameterNameDiscoverer;
+import org.springframework.core.annotation.AnnotationUtils;
+import org.springframework.util.ClassUtils;
+
+import java.lang.reflect.Method;
+import java.util.*;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+@Slf4j
+public class RateLimiterAopAdvisor extends StaticMethodMatcherPointcutAdvisor {
+
+    private static final long                    serialVersionUID = -1076122956392948260L;
+    private static final ParameterNameDiscoverer nameDiscoverer   = new LocalVariableTableParameterNameDiscoverer();
+
+    public RateLimiterAopAdvisor(RateLimiterManager rateLimiterManager) {
+        setAdvice((MethodBeforeAdvice) (method, args, target) -> {
+            String[] names = nameDiscoverer.getParameterNames(method);
+            RateLimiter limiter = Optional.ofNullable(AnnotationUtils.findAnnotation(method, RateLimiter.class))
+                    .orElseGet(() -> AnnotationUtils.findAnnotation(ClassUtils.getUserClass(target), RateLimiter.class));
+            if (limiter != null) {
+                List<String> keyExpressionList = new ArrayList<>(Arrays.asList(limiter.key()));
+                if (keyExpressionList.isEmpty()) {
+                    keyExpressionList.add(method.toString());
+                }
+                for (String keyExpress : keyExpressionList) {
+                    if (keyExpress.contains("${")) {
+                        Map<String, Object> params = new HashMap<>();
+                        params.put("user", Authentication.current().map(Authentication::getUser).orElse(null));
+                        for (int i = 0; i < args.length; i++) {
+                            params.put(names.length > i ? names[i] : "arg" + i, args[i]);
+                            params.put("arg" + i, args[i]);
+
+                        }
+                        keyExpress = ExpressionUtils.analytical(keyExpress, params, "spel");
+                    }
+                    log.debug("do rate limiter:[{}]. ", keyExpress);
+                    boolean success = rateLimiterManager
+                            .getRateLimiter(keyExpress, limiter.permits(), limiter.timeUnit())
+                            .tryAcquire(limiter.acquire(), limiter.acquireTimeUnit());
+                    if (!success) {
+                        throw new TimeoutException("请求超时");
+                    }
+                }
+            }
+        });
+    }
+
+
+    @Override
+    public boolean matches(Method method, Class<?> targetClass) {
+
+        return AnnotationUtils.findAnnotation(method, RateLimiter.class) != null
+                || AnnotationUtils.findAnnotation(targetClass, RateLimiter.class) != null;
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/src/main/java/org/hswebframework/web/concurent/RateLimiterAutoConfiguration.java

@@ -0,0 +1,30 @@
+package org.hswebframework.web.concurent;
+
+import org.hswebframework.web.concurrent.GuavaRateLimiterManager;
+import org.hswebframework.web.concurrent.RateLimiterManager;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+@Configuration
+public class RateLimiterAutoConfiguration {
+
+    @Bean
+    @ConditionalOnMissingBean(RateLimiterManager.class)
+    @ConditionalOnClass(name = "com.google.common.util.concurrent.RateLimiter")
+    public GuavaRateLimiterManager guavaRateLimiterManager() {
+        return new GuavaRateLimiterManager();
+    }
+
+    @Bean
+    @ConditionalOnBean(RateLimiterManager.class)
+    public RateLimiterAopAdvisor rateLimiterAopAdvisor(RateLimiterManager rateLimiterManager) {
+        return new RateLimiterAopAdvisor(rateLimiterManager);
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/src/main/resources/META-INF/spring.factories

@@ -0,0 +1,3 @@
+# Auto Configure
+org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
+org.hswebframework.web.concurent.RateLimiterAutoConfiguration

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/src/test/groovy/org/hswebframework/web/concurent/RateLimiterAopAdvisorTest.groovy

@@ -0,0 +1,48 @@
+package org.hswebframework.web.concurent
+
+import org.springframework.beans.factory.annotation.Autowired
+import org.springframework.boot.test.context.SpringBootTest
+import org.springframework.test.context.ContextConfiguration
+import spock.lang.Specification
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+@SpringBootTest(classes = TestApplication.class)
+@ContextConfiguration
+class RateLimiterAopAdvisorTest extends Specification {
+
+    @Autowired
+    public TestService testService;
+
+    def "测试限流"() {
+        TestService.counter.set(0);
+        given:
+        testService.test();
+        def timeoutException;
+        try {
+            testService.test();
+        } catch (Exception e) {
+            timeoutException = e;
+        }
+        expect:
+        TestService.counter.get() == 1
+        timeoutException != null
+    }
+
+    def "测试指定key限流"() {
+        TestService.counter.set(0);
+        given:
+        testService.test("test");
+        def timeoutException;
+        try {
+            testService.test("test");
+        } catch (Exception e) {
+            timeoutException = e;
+        }
+        expect:
+        TestService.counter.get() == 1
+        timeoutException != null
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/src/test/groovy/org/hswebframework/web/concurent/TestApplication.java

@@ -0,0 +1,11 @@
+package org.hswebframework.web.concurent;
+
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+@SpringBootApplication
+public class TestApplication {
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/hsweb-concurrent-rate-limiter-starter/src/test/groovy/org/hswebframework/web/concurent/TestService.java

@@ -0,0 +1,30 @@
+package org.hswebframework.web.concurent;
+
+import org.hswebframework.web.concurrent.annotation.RateLimiter;
+import org.springframework.stereotype.Service;
+
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * @author zhouhao
+ * @since 3.0.4
+ */
+@Service
+public class TestService {
+
+    public static AtomicLong counter = new AtomicLong();
+
+
+    @RateLimiter(permits = 1, acquire = 500, acquireTimeUnit = TimeUnit.MILLISECONDS) //一秒一次
+    public String test() {
+        counter.incrementAndGet();
+        return "ok";
+    }
+
+    @RateLimiter(key = "${#name}", permits = 1, acquire = 500, acquireTimeUnit = TimeUnit.MILLISECONDS)
+    public String test(String name) {
+        counter.incrementAndGet();
+        return name;
+    }
+}

hsweb-concurrent/hsweb-concurrent-rate-limiter/pom.xml

@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>hsweb-concurrent</artifactId>
+        <groupId>org.hswebframework.web</groupId>
+        <version>3.0.4-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>hsweb-concurrent-rate-limiter</artifactId>
+    <packaging>pom</packaging>
+    <modules>
+        <module>hsweb-concurrent-rate-limiter-api</module>
+        <module>hsweb-concurrent-rate-limiter-starter</module>
+    </modules>
+
+
+</project>

hsweb-concurrent/pom.xml

@@ -35,6 +35,7 @@
         <module>hsweb-concurrent-lock</module>
         <module>hsweb-concurrent-counter</module>
         <module>hsweb-concurrent-async-job</module>
+        <module>hsweb-concurrent-rate-limiter</module>
     </modules>
 
 

hsweb-starter/hsweb-spring-boot-starter/src/main/java/org/hswebframework/web/starter/RestControllerExceptionTranslator.java

@@ -54,6 +54,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 
 @RestControllerAdvice
@@ -148,6 +149,15 @@ public class RestControllerExceptionTranslator {
         return ResponseMessage.error(400, results.getResults().isEmpty() ? e.getMessage() : results.getResults().get(0).getMessage()).result(results.getResults());
     }
 
+    @ExceptionHandler(TimeoutException.class)
+    @ResponseStatus(HttpStatus.GATEWAY_TIMEOUT)
+    ResponseMessage handleException(TimeoutException exception) {
+        String msg = Optional.ofNullable(exception.getMessage())
+                .orElse("访问服务超时");
+        logger.warn(exception.getMessage(), exception);
+        return ResponseMessage.error(504, msg);
+    }
+
     @ExceptionHandler(RuntimeException.class)
     @ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR)
     ResponseMessage handleException(RuntimeException exception) {