SpringBoot中基于AOP和Semaphore實現(xiàn)API限流

更新時間：2024年10月20日 09:04:05 作者：碼到三十五

調用速率限制是 Web API 中的常見要求,旨在防止濫用并確保公平使用資源,借助Spring Boot 中的 AOP,我們可以通過攔截方法調用并限制在特定時間范圍內允許的請求數(shù)量來實現(xiàn)速率限制,需要的朋友可以參考下

Spring Boot API 中的速率限制

可以使用各種技術在 Spring Boot API 中實現(xiàn)速率限制。一種常見的方法是使用 Spring AOP來攔截傳入的請求并實施速率限制。

步驟 1 - 定義速率限制配置

創(chuàng)建一個配置類，在其中定義速率限制參數(shù)，例如允許的請求數(shù)和時間段。

@Configuration
public class RateLimitConfig {
    @Value("${rate.limit.requests}")
    private int requests;

    @Value("${rate.limit.seconds}")
    private int seconds;

    // Getters and setters
}

步驟 2 - 創(chuàng)建速率限制方面

使用 Spring AOP 實現(xiàn)一個方面來攔截方法調用并強制執(zhí)行速率限制。

@Aspect
@Component
public class RateLimitAspect {
    @Autowired
    private RateLimitConfig rateLimitConfig;

    @Autowired
    private RateLimiter rateLimiter;

    @Around("@annotation(RateLimited)")
    public Object enforceRateLimit(ProceedingJoinPoint joinPoint) throws Throwable {
        String key = getKey(joinPoint);
        if (!rateLimiter.tryAcquire(key, rateLimitConfig.getRequests(), rateLimitConfig.getSeconds())) {
            throw new RateLimitExceededException("Rate limit exceeded");
        }
        return joinPoint.proceed();
    }

    private String getKey(ProceedingJoinPoint joinPoint) {
        //為正在調用的方法生成唯一密鑰
        //方法簽名、用戶ID、IP地址等。

    }
}

步驟 3 — 定義 RateLimited 注釋

創(chuàng)建自定義注釋來標記應受速率限制的方法。

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
  public @interface RateLimited {
}

步驟 4 - 實施速率限制器

創(chuàng)建速率限制器組件，使用令牌桶算法或任何其他合適的算法來管理速率限制。

@Component
public class RateLimiter {
    private final Map<String,RateLimitedSemaphore> semaphores = new ConcurrentHashMap<>();

    public boolean tryAcquire(String key, int requests, int seconds) {
        
        long currentTime = System.currentTimeMillis();

        // 計算時間窗口
        long startTime = currentTime - seconds * 1000;

        // 過期刪除
        cleanupExpiredEntries(startTime);

        // 獲取semaphore 
        RateLimitedSemaphore semaphore = semaphores.computeIfAbsent(key, k -> {
            RateLimitedSemaphore newSemaphore = new RateLimitedSemaphore(requests);
            newSemaphore.setLastAcquireTime(currentTime); // Set last acquire time
            return newSemaphore;
        });

        // 校驗 semaphore 
        boolean acquired = semaphore.tryAcquire();
        if (acquired) {
            semaphore.setLastAcquireTime(currentTime); 
            // 更新
        }
        return acquired;
    }

    private void cleanupExpiredEntries(long startTime) {
        Iterator<Map.Entry<String, RateLimitedSemaphore>> iterator = semaphores.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, RateLimitedSemaphore> entry = iterator.next();
            String key = entry.getKey();
            RateLimitedSemaphore semaphore = entry.getValue();
            if (semaphore.getLastAcquireTime() < startTime) {
                iterator.remove();
            }
        }
    }

    private class RateLimitedSemaphore extends Semaphore {
        private volatile long lastAcquireTime;

        public RateLimitedSemaphore(int permits) {
            super(permits);
        }

        public long getLastAcquireTime() {
            return lastAcquireTime;
        }

        public void setLastAcquireTime(long lastAcquireTime) {
            this.lastAcquireTime = lastAcquireTime;
        }
    }
}

步驟 5 - 注釋控制器方法

用注解來注釋應該進行速率限制的控制器方法 @RateLimited。

@RestController
public class MyController {
    @RateLimited
    @GetMapping("/api/resource")
    public ResponseEntity<String> getResource() {
        // Implementation
    }
}

步驟 6 - 配置速率限制屬性

application.properties在您的或中配置速率限制屬性 application.yml。

rate.limit.requests=10
rate.limit.seconds=60

要按 IP 地址限制請求，可以從傳入請求中提取 IP 地址并將其用作速率限制的密鑰：

private String getKey(HttpServletRequest request) {
    String ipAddress = request.getRemoteAddr();
    return ipAddress; //用ID做key
}

還需要修改enforceRateLimit 中的方法 RateLimitAspect 以將對象傳遞 HttpServletRequest 給 getKey 方法：

@Around("@annotation(RateLimited)")
public Object enforceRateLimit(ProceedingJoinPoint joinPoint) throws Throwable {
   
    ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
    HttpServletRequest request = requestAttributes.getRequest();

    String key = getKey(request);
    if (!rateLimiter.tryAcquire(key, rateLimitConfig.getRequests(), rateLimitConfig.getSeconds())) {
        throw new RateLimitExceededException("Rate limit exceeded");
    }
    return joinPoint.proceed();
}

以上就是SpringBoot中基于AOP和Semaphore實現(xiàn)API限流的詳細內容，更多關于SpringBoot實現(xiàn)API限流的資料請關注腳本之家其它相關文章！

您可能感興趣的文章: