Skip to content

Commit bdfa411

Browse files
mazp99mazp99
committed
feat：将速率限制改为了服务端获取
1 parent 84c2911 commit bdfa411

File tree

7 files changed

+135
-30
lines changed

7 files changed

+135
-30
lines changed

‎dongtai-agent/src/main/java/io/dongtai/iast/agent/AgentLauncher.java‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import io.dongtai.iast.common.state.AgentState;
99
import io.dongtai.iast.common.state.State;
1010
import io.dongtai.iast.common.state.StateCause;
11-
import io.dongtai.iast.common.utils.limit.InterfaceRateLimiterUtil;
1211
import io.dongtai.log.DongTaiLog;
1312
import io.dongtai.log.ErrorCode;
1413

@@ -172,19 +171,6 @@ private static void install(final Instrumentation inst) {
172171
if (!AGENT_STATE.isException()) {
173172
AGENT_STATE.setState(State.RUNNING);
174173
}
175-
String rateCaps = System.getProperty("rate.caps");
176-
if (rateCaps != null && !rateCaps.isEmpty()){
177-
try {
178-
InterfaceRateLimiterUtil.initializeInstance(Long.parseLong(rateCaps));
179-
} catch (NumberFormatException e) {
180-
// 当出现异常时，降级为默认速率
181-
rateCaps = IastProperties.getInstance().cfg.getProperty("rate.caps");
182-
DongTaiLog.error("Interface Rate Limiter Initialization Failure Reason：{} \n " +
183-
"the default rate will be used {}",e.getMessage(),rateCaps);
184-
InterfaceRateLimiterUtil.initializeInstance(Long.parseLong(rateCaps));
185-
186-
}
187-
}
188174
} else {
189175
DongTaiLog.error(ErrorCode.AGENT_REGISTER_INFO_INVALID);
190176
AGENT_STATE.setState(State.EXCEPTION);

‎dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/ConfigMonitor.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import io.dongtai.iast.common.config.ConfigKey;
1010
import io.dongtai.iast.common.constants.AgentConstant;
1111
import io.dongtai.iast.common.constants.ApiPath;
12+
import io.dongtai.iast.common.utils.limit.InterfaceRateLimiterUtil;
1213
import io.dongtai.log.DongTaiLog;
1314
import io.dongtai.log.ErrorCode;
1415

@@ -48,6 +49,37 @@ private void updateConfig() {
4849
if (logLevel != null) {
4950
DongTaiLog.setLevel(DongTaiLog.parseLevel(logLevel));
5051
}
52+
53+
//获取是否开启qps限流
54+
Boolean enableQpsRate = ConfigBuilder.getInstance().get(ConfigKey.ENABLE_QPS_RATE_LIMIT);
55+
if (enableQpsRate) {
56+
int qpsRateLimit = ConfigBuilder.getInstance().get(ConfigKey.QPS_RATE_LIMIT);
57+
if (qpsRateLimit <= 0){
58+
DongTaiLog.error("qpsRateLimit the value cannot be less than 0");
59+
qpsRateLimit = 100;
60+
DongTaiLog.error("qpsRateLimit revert to 100");
61+
62+
}
63+
int tokenBucketPoolSize = ConfigBuilder.getInstance().get(ConfigKey.TOKEN_BUCKET_POOL_SIZE);
64+
if (tokenBucketPoolSize <= 0){
65+
DongTaiLog.error("tokenBucketPoolSize the value cannot be less than 0");
66+
tokenBucketPoolSize=5000;
67+
DongTaiLog.error("tokenBucketPoolSize revert to 5000");
68+
69+
}
70+
//判断是否已经开启，如果已经开启，则更新数据
71+
if (InterfaceRateLimiterUtil.getRateLimiterState()) {
72+
InterfaceRateLimiterUtil.updateTheData(qpsRateLimit, tokenBucketPoolSize);
73+
} else {
74+
//初始化令牌池工具，设置池大小和速率
75+
InterfaceRateLimiterUtil.initializeInstance(qpsRateLimit, tokenBucketPoolSize);
76+
}
77+
}else {
78+
if(InterfaceRateLimiterUtil.getRateLimiterState()){
79+
InterfaceRateLimiterUtil.turnOffTheRateLimiter();
80+
}
81+
}
82+
5183
}
5284

5385
@Override

‎dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigBuilder.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,12 @@ private ConfigBuilder() {
2828
Config.<String>create(ConfigKey.LOGGER_LEVEL));
2929
this.configMap.put(ConfigKey.VALIDATED_SINK,
3030
Config.<Boolean>create(ConfigKey.VALIDATED_SINK).setDefaultValue(false));
31+
this.configMap.put(ConfigKey.ENABLE_QPS_RATE_LIMIT,
32+
Config.<Boolean>create(ConfigKey.ENABLE_QPS_RATE_LIMIT).setDefaultValue(false));
33+
this.configMap.put(ConfigKey.QPS_RATE_LIMIT,
34+
Config.<Integer>create(ConfigKey.QPS_RATE_LIMIT).setDefaultValue(100));
35+
this.configMap.put(ConfigKey.TOKEN_BUCKET_POOL_SIZE,
36+
Config.<Integer>create(ConfigKey.TOKEN_BUCKET_POOL_SIZE).setDefaultValue(5000));
3137
}
3238

3339
public static ConfigBuilder getInstance() {
@@ -68,6 +74,9 @@ public void update(JSONObject config) {
6874
updateBool(config, ConfigKey.JsonKey.JSON_ENABLE_LOGGER);
6975
updateString(config, ConfigKey.JsonKey.JSON_LOGGER_LEVEL);
7076
updateBool(config, ConfigKey.JsonKey.JSON_VALIDATED_SINK);
77+
updateBool(config, ConfigKey.JsonKey.JSON_ENABLE_QPS_RATE_LIMIT);
78+
updateInt(config, ConfigKey.JsonKey.JSON_QPS_RATE_LIMIT);
79+
updateInt(config, ConfigKey.JsonKey.JSON_TOKEN_BUCKET_POOL_SIZE);
7180
updateRequestDenyList(config);
7281
}
7382

‎dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigKey.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ public enum ConfigKey {
99
ENABLE_LOGGER,
1010
LOGGER_LEVEL,
1111
VALIDATED_SINK,
12+
ENABLE_QPS_RATE_LIMIT,
13+
QPS_RATE_LIMIT,
14+
TOKEN_BUCKET_POOL_SIZE,
1215
;
1316

1417
public enum JsonKey {
@@ -20,6 +23,9 @@ public enum JsonKey {
2023
JSON_ENABLE_LOGGER("enable_log", ENABLE_LOGGER),
2124
JSON_LOGGER_LEVEL("log_level", LOGGER_LEVEL),
2225
JSON_VALIDATED_SINK("report_validated_sink", VALIDATED_SINK),
26+
JSON_ENABLE_QPS_RATE_LIMIT("enable_qps_rate_limit", ENABLE_QPS_RATE_LIMIT),
27+
JSON_QPS_RATE_LIMIT("qps_rate_limit", QPS_RATE_LIMIT),
28+
JSON_TOKEN_BUCKET_POOL_SIZE("token_bucket_pool_size", TOKEN_BUCKET_POOL_SIZE),
2329
;
2430

2531
private final String key;

‎dongtai-common/src/main/java/io/dongtai/iast/common/utils/limit/InterfaceRateLimiter.java‎

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,22 @@
1414
*/
1515
public class InterfaceRateLimiter {
1616

17-
private final InterfaceRateLimiterSoftReferenceHashMap<String, Bucket> buckets = new InterfaceRateLimiterSoftReferenceHashMap<>();
17+
private final InterfaceRateLimiterSoftReferenceHashMap<String> buckets = new InterfaceRateLimiterSoftReferenceHashMap<>();
18+
19+
private int theNumberOfTokenBucketPools;
1820

1921

2022
//最高速率
21-
private final long rateCaps;
23+
private long rateCaps;
2224

2325

24-
private InterfaceRateLimiter(long rateCaps) {
26+
private InterfaceRateLimiter(long rateCaps, int theNumberOfTokenBucketPools) {
2527
this.rateCaps = rateCaps;
28+
this.theNumberOfTokenBucketPools = theNumberOfTokenBucketPools;
2629
}
2730

28-
public static InterfaceRateLimiter getInstance(long rateCaps) {
29-
return new InterfaceRateLimiter(rateCaps);
31+
public static InterfaceRateLimiter getInstance(long rateCaps, int theNumberOfTokenBucketPools) {
32+
return new InterfaceRateLimiter(rateCaps, theNumberOfTokenBucketPools);
3033
}
3134

3235
/**
@@ -43,8 +46,8 @@ public boolean whetherItPassesOrNot(String interfaceName) {
4346
if (bucket == null) {
4447
return true;
4548
}
46-
}else {
47-
if (buckets.size() >= 5000){
49+
} else {
50+
if (buckets.size() >= theNumberOfTokenBucketPools) {
4851
return true;
4952
}
5053
bucket = Bucket.builder().addLimit(Bandwidth.simple(rateCaps, Duration.ofSeconds(1))).build();
@@ -55,4 +58,24 @@ public boolean whetherItPassesOrNot(String interfaceName) {
5558
return probe.isConsumed();
5659
}
5760

61+
/**
62+
* 更新速率限制器的配置
63+
* @param rateCaps 速率
64+
* @param theNumberOfTokenBucketPools 桶池大小
65+
* 只有当值发生改变时，才会开启更新，否则不执行操作
66+
*/
67+
public void updateTheData(long rateCaps, int theNumberOfTokenBucketPools) {
68+
if (this.rateCaps != rateCaps){
69+
this.rateCaps = rateCaps;
70+
buckets.updateTheData(this.rateCaps);
71+
}
72+
if (this.theNumberOfTokenBucketPools != theNumberOfTokenBucketPools) {
73+
this.theNumberOfTokenBucketPools = theNumberOfTokenBucketPools;
74+
//需要更新桶池大小时，判断已有的是否超过最大限制，如果超过，则删除到指定大小
75+
if (theNumberOfTokenBucketPools - buckets.size() < 0){
76+
int i = buckets.size() - theNumberOfTokenBucketPools;
77+
buckets.deleteTheElement(i);
78+
}
79+
}
80+
}
5881
}

‎dongtai-common/src/main/java/io/dongtai/iast/common/utils/limit/InterfaceRateLimiterSoftReferenceHashMap.java‎

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,59 @@
11
package io.dongtai.iast.common.utils.limit;
22

3+
import io.github.bucket4j.Bandwidth;
4+
import io.github.bucket4j.Bucket;
5+
36
import java.lang.ref.SoftReference;
7+
import java.time.Duration;
8+
import java.util.Iterator;
9+
import java.util.Map;
410
import java.util.concurrent.ConcurrentHashMap;
11+
import java.util.concurrent.atomic.AtomicInteger;
512

613
/**
714
* @author mazepeng
815
* @date 2023/12/7 16:56
916
*/
10-
public class InterfaceRateLimiterSoftReferenceHashMap<K,V> {
17+
public class InterfaceRateLimiterSoftReferenceHashMap<K> {
1118

12-
private final ConcurrentHashMap<K, SoftReference<V>> map = new ConcurrentHashMap<>();
19+
private final ConcurrentHashMap<K, SoftReference<Bucket>> map = new ConcurrentHashMap<>();
1320

14-
public void put(K key, V value) {
21+
public void put(K key, Bucket value) {
1522
map.put(key, new SoftReference<>(value));
1623
}
1724

18-
public V get(K key) {
19-
SoftReference<V> softReference = map.get(key);
25+
/**
26+
* 更新令牌桶的速率限制
27+
* @param rateCaps 速率
28+
*/
29+
public void updateTheData(long rateCaps){
30+
map.replaceAll((key, value) -> {
31+
Bucket build = Bucket.builder().addLimit(Bandwidth.simple(rateCaps, Duration.ofSeconds(1))).build();
32+
return new SoftReference<>(build);
33+
});
34+
}
35+
36+
/**
37+
* 删除指定个桶
38+
* @param theNumberOfTokenBucketPools 删除的数量
39+
*/
40+
public void deleteTheElement(int theNumberOfTokenBucketPools){
41+
AtomicInteger atomicInteger = new AtomicInteger(theNumberOfTokenBucketPools);
42+
for (Iterator<Map.Entry<K, SoftReference<Bucket>>> iterator = map.entrySet().iterator(); ; iterator.hasNext()) {
43+
Map.Entry<K, SoftReference<Bucket>> next = iterator.next();
44+
45+
int i = atomicInteger.decrementAndGet();
46+
if (i >= 0) {
47+
iterator.remove();
48+
} else {
49+
// 满足终止条件时，跳出循环
50+
break;
51+
}
52+
}
53+
}
54+
55+
public Bucket get(K key) {
56+
SoftReference<Bucket> softReference = map.get(key);
2057
if (softReference != null) {
2158
return softReference.get();
2259
}

‎dongtai-common/src/main/java/io/dongtai/iast/common/utils/limit/InterfaceRateLimiterUtil.java‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,18 +24,22 @@ public static boolean getRateLimiterState(){
2424
return turnOnTheRateLimiter;
2525
}
2626

27-
public static void turnOffTheRateLimiter(boolean state){
28-
turnOnTheRateLimiter = true;
27+
/**
28+
* 关闭速率限制器，并将限制器置为空释放空间
29+
*/
30+
public static void turnOffTheRateLimiter(){
31+
turnOnTheRateLimiter = false;
32+
instance = null;
2933
}
3034

3135
/**
3236
* 初始化速率限制
3337
*/
34-
public static void initializeInstance(long rateCaps) {
38+
public static void initializeInstance(long rateCaps,int theNumberOfTokenBucketPools) {
3539
if (instance == null) {
3640
synchronized (InterfaceRateLimiterUtil.class) {
3741
if (instance == null) {
38-
instance = InterfaceRateLimiter.getInstance(rateCaps);
42+
instance = InterfaceRateLimiter.getInstance(rateCaps,theNumberOfTokenBucketPools);
3943
turnOnTheRateLimiter = true;
4044
}
4145
}
@@ -56,5 +60,13 @@ public static boolean whetherItPassesOrNot(String interfaceName){
5660
return instance.whetherItPassesOrNot(interfaceName);
5761
}
5862

63+
/**
64+
* 更新令牌桶的设置
65+
* @param rateCaps 速率
66+
* @param theNumberOfTokenBucketPools 令牌桶池上限
67+
*/
68+
public static void updateTheData(long rateCaps,int theNumberOfTokenBucketPools) {
69+
instance.updateTheData(rateCaps,theNumberOfTokenBucketPools);
70+
}
5971
}
6072

0 commit comments

Comments
 (0)