Skip to content

Commit e8e293f

Browse files
robberphexrobberphex
committed
introduce intermediate node to avoid ABA problem
1 parent 8d90fcf commit e8e293f

File tree

6 files changed

+99
-78
lines changed

6 files changed

+99
-78
lines changed

‎sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowChecker.java

+42-50
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 1999-2018 Alibaba Group Holding Ltd.
2+
* Copyright 1999-2024 Alibaba Group Holding Ltd.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -15,16 +15,6 @@
1515
*/
1616
package com.alibaba.csp.sentinel.slots.block.flow.param;
1717

18-
import java.lang.reflect.Array;
19-
import java.util.ArrayList;
20-
import java.util.Collection;
21-
import java.util.Collections;
22-
import java.util.List;
23-
import java.util.Set;
24-
import java.util.concurrent.TimeUnit;
25-
import java.util.concurrent.atomic.AtomicInteger;
26-
import java.util.concurrent.atomic.AtomicLong;
27-
2818
import com.alibaba.csp.sentinel.cluster.ClusterStateManager;
2919
import com.alibaba.csp.sentinel.cluster.TokenResult;
3020
import com.alibaba.csp.sentinel.cluster.TokenResultStatus;
@@ -37,6 +27,13 @@
3727
import com.alibaba.csp.sentinel.slots.statistic.cache.CacheMap;
3828
import com.alibaba.csp.sentinel.util.TimeUtil;
3929

30+
import java.lang.reflect.Array;
31+
import java.time.format.DateTimeFormatter;
32+
import java.util.*;
33+
import java.util.concurrent.TimeUnit;
34+
import java.util.concurrent.atomic.AtomicLong;
35+
import java.util.concurrent.atomic.AtomicReference;
36+
4037
/**
4138
* Rule checker for parameter flow control.
4239
*
@@ -46,7 +43,7 @@
4643
public final class ParamFlowChecker {
4744

4845
public static boolean passCheck(ResourceWrapper resourceWrapper, /*@Valid*/ ParamFlowRule rule, /*@Valid*/ int count,
49-
Object... args) {
46+
Object... args) {
5047
if (args == null) {
5148
return true;
5249
}
@@ -79,7 +76,7 @@ private static boolean passLocalCheck(ResourceWrapper resourceWrapper, ParamFlow
7976
Object value) {
8077
try {
8178
if (Collection.class.isAssignableFrom(value.getClass())) {
82-
for (Object param : ((Collection)value)) {
79+
for (Object param : ((Collection) value)) {
8380
if (!passSingleValueCheck(resourceWrapper, rule, count, param)) {
8481
return false;
8582
}
@@ -117,7 +114,7 @@ static boolean passSingleValueCheck(ResourceWrapper resourceWrapper, ParamFlowRu
117114
int itemThreshold = rule.getParsedHotItems().get(value);
118115
return ++threadCount <= itemThreshold;
119116
}
120-
long threshold = (long)rule.getCount();
117+
long threshold = (long) rule.getCount();
121118
return ++threadCount <= threshold;
122119
}
123120

@@ -127,16 +124,16 @@ static boolean passSingleValueCheck(ResourceWrapper resourceWrapper, ParamFlowRu
127124
static boolean passDefaultLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
128125
Object value) {
129126
ParameterMetric metric = getParameterMetric(resourceWrapper);
130-
CacheMap<Object, AtomicLong> tokenCounters = metric == null ? null : metric.getRuleTokenCounter(rule);
131-
CacheMap<Object, AtomicLong> timeCounters = metric == null ? null : metric.getRuleTimeCounter(rule);
127+
CacheMap<Object, AtomicReference<TokenUpdateStatus>> tokenCounters = metric == null ? null : metric.getRuleStampedTokenCounter(rule);
132128

133-
if (tokenCounters == null || timeCounters == null) {
129+
DateTimeFormatter dtf = DateTimeFormatter.ISO_DATE_TIME;
130+
if (tokenCounters == null) {
134131
return true;
135132
}
136133

137134
// Calculate max token count (threshold)
138135
Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
139-
long tokenCount = (long)rule.getCount();
136+
long tokenCount = (long) rule.getCount();
140137
if (exclusionItems.contains(value)) {
141138
tokenCount = rule.getParsedHotItems().get(value);
142139
}
@@ -153,49 +150,44 @@ static boolean passDefaultLocalCheck(ResourceWrapper resourceWrapper, ParamFlowR
153150
while (true) {
154151
long currentTime = TimeUtil.currentTimeMillis();
155152

156-
AtomicLong lastAddTokenTime = timeCounters.putIfAbsent(value, new AtomicLong(currentTime));
157-
if (lastAddTokenTime == null) {
153+
AtomicReference<TokenUpdateStatus> atomicLastStatus = tokenCounters.putIfAbsent(value, new AtomicReference<>(
154+
new TokenUpdateStatus(currentTime, maxCount - acquireCount)
155+
));
156+
if (atomicLastStatus == null) {
158157
// Token never added, just replenish the tokens and consume {@code acquireCount} immediately.
159-
tokenCounters.putIfAbsent(value, new AtomicLong(maxCount - acquireCount));
160158
return true;
161159
}
162160

163161
// Calculate the time duration since last token was added.
164-
long passTime = currentTime - lastAddTokenTime.get();
162+
TokenUpdateStatus lastStatus = atomicLastStatus.get();
163+
long passTime = currentTime - lastStatus.getLastAddTokenTime();
165164
// A simplified token bucket algorithm that will replenish the tokens only when statistic window has passed.
165+
long newQps;
166166
if (passTime > rule.getDurationInSec() * 1000) {
167-
AtomicLong oldQps = tokenCounters.putIfAbsent(value, new AtomicLong(maxCount - acquireCount));
168-
if (oldQps == null) {
169-
// Might not be accurate here.
170-
lastAddTokenTime.set(currentTime);
171-
return true;
172-
} else {
173-
long restQps = oldQps.get();
174-
long toAddCount = (passTime * tokenCount) / (rule.getDurationInSec() * 1000);
175-
long newQps = toAddCount + restQps > maxCount ? (maxCount - acquireCount)
167+
long restQps = lastStatus.getRestQps();
168+
long toAddCount = (passTime * tokenCount) / (rule.getDurationInSec() * 1000);
169+
newQps = toAddCount + restQps > maxCount ? (maxCount - acquireCount)
176170
: (restQps + toAddCount - acquireCount);
177171

178-
if (newQps < 0) {
179-
return false;
180-
}
181-
if (oldQps.compareAndSet(restQps, newQps)) {
182-
lastAddTokenTime.set(currentTime);
183-
return true;
184-
}
185-
Thread.yield();
172+
if (newQps < 0) {
173+
return false;
186174
}
175+
TokenUpdateStatus newStatus = new TokenUpdateStatus(currentTime, newQps);
176+
if (atomicLastStatus.compareAndSet(lastStatus, newStatus)) {
177+
return true;
178+
}
179+
Thread.yield();
187180
} else {
188-
AtomicLong oldQps = tokenCounters.get(value);
189-
if (oldQps != null) {
190-
long oldQpsValue = oldQps.get();
191-
if (oldQpsValue - acquireCount >= 0) {
192-
if (oldQps.compareAndSet(oldQpsValue, oldQpsValue - acquireCount)) {
193-
return true;
194-
}
195-
} else {
196-
return false;
181+
newQps = lastStatus.getRestQps() - acquireCount;
182+
if (newQps >= 0) {
183+
TokenUpdateStatus newStatus = new TokenUpdateStatus(lastStatus.getLastAddTokenTime(), newQps);
184+
if (atomicLastStatus.compareAndSet(lastStatus, newStatus)) {
185+
return true;
197186
}
187+
} else {
188+
return false;
198189
}
190+
199191
Thread.yield();
200192
}
201193
}
@@ -211,7 +203,7 @@ static boolean passThrottleLocalCheck(ResourceWrapper resourceWrapper, ParamFlow
211203

212204
// Calculate max token count (threshold)
213205
Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
214-
long tokenCount = (long)rule.getCount();
206+
long tokenCount = (long) rule.getCount();
215207
if (exclusionItems.contains(value)) {
216208
tokenCount = rule.getParsedHotItems().get(value);
217209
}
@@ -261,7 +253,7 @@ private static ParameterMetric getParameterMetric(ResourceWrapper resourceWrappe
261253
@SuppressWarnings("unchecked")
262254
private static Collection<Object> toCollection(Object value) {
263255
if (value instanceof Collection) {
264-
return (Collection<Object>)value;
256+
return (Collection<Object>) value;
265257
} else if (value.getClass().isArray()) {
266258
List<Object> params = new ArrayList<Object>();
267259
int length = Array.getLength(value);

‎sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParameterMetric.java

+16-13
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Map;
2222
import java.util.concurrent.atomic.AtomicInteger;
2323
import java.util.concurrent.atomic.AtomicLong;
24+
import java.util.concurrent.atomic.AtomicReference;
2425

2526
import com.alibaba.csp.sentinel.log.RecordLog;
2627
import com.alibaba.csp.sentinel.slots.statistic.cache.CacheMap;
@@ -46,25 +47,35 @@ public class ParameterMetric {
4647
* @since 1.6.0
4748
*/
4849
private final Map<ParamFlowRule, CacheMap<Object, AtomicLong>> ruleTimeCounters = new HashMap<>();
50+
4951
/**
5052
* Format: (rule, (value, tokenCounter))
5153
*
5254
* @since 1.6.0
5355
*/
54-
private final Map<ParamFlowRule, CacheMap<Object, AtomicLong>> ruleTokenCounter = new HashMap<>();
56+
private final Map<ParamFlowRule, CacheMap<Object, AtomicReference<TokenUpdateStatus>>> ruleTokenCounter = new HashMap<>();
57+
5558
private final Map<Integer, CacheMap<Object, AtomicInteger>> threadCountMap = new HashMap<>();
5659

5760
/**
5861
* Get the token counter for given parameter rule.
5962
*
6063
* @param rule valid parameter rule
6164
* @return the associated token counter
62-
* @since 1.6.0
65+
* @since 1.8.8
6366
*/
64-
public CacheMap<Object, AtomicLong> getRuleTokenCounter(ParamFlowRule rule) {
67+
CacheMap<Object, AtomicReference<TokenUpdateStatus>> getRuleStampedTokenCounter(ParamFlowRule rule) {
6568
return ruleTokenCounter.get(rule);
6669
}
6770

71+
public void clear() {
72+
synchronized (lock) {
73+
ruleTimeCounters.clear();
74+
ruleTokenCounter.clear();
75+
threadCountMap.clear();
76+
}
77+
}
78+
6879
/**
6980
* Get the time record counter for given parameter rule.
7081
*
@@ -76,14 +87,6 @@ public CacheMap<Object, AtomicLong> getRuleTimeCounter(ParamFlowRule rule) {
7687
return ruleTimeCounters.get(rule);
7788
}
7889

79-
public void clear() {
80-
synchronized (lock) {
81-
threadCountMap.clear();
82-
ruleTimeCounters.clear();
83-
ruleTokenCounter.clear();
84-
}
85-
}
86-
8790
public void clearForRule(ParamFlowRule rule) {
8891
synchronized (lock) {
8992
ruleTimeCounters.remove(rule);
@@ -106,7 +109,7 @@ public void initialize(ParamFlowRule rule) {
106109
synchronized (lock) {
107110
if (ruleTokenCounter.get(rule) == null) {
108111
long size = Math.min(BASE_PARAM_MAX_CAPACITY * rule.getDurationInSec(), TOTAL_MAX_CAPACITY);
109-
ruleTokenCounter.put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(size));
112+
ruleTokenCounter.put(rule, new ConcurrentLinkedHashMapWrapper<>(size));
110113
}
111114
}
112115
}
@@ -253,7 +256,7 @@ public long getThreadCount(int index, Object value) {
253256
*
254257
* @return the token counter map
255258
*/
256-
Map<ParamFlowRule, CacheMap<Object, AtomicLong>> getRuleTokenCounterMap() {
259+
Map<ParamFlowRule, CacheMap<Object, AtomicReference<TokenUpdateStatus>>> getRuleTokenCounterMap() {
257260
return ruleTokenCounter;
258261
}
259262

‎sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/TokenUpdateStatus.java

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package com.alibaba.csp.sentinel.slots.block.flow.param;
2+
3+
class TokenUpdateStatus {
4+
5+
private final long lastAddTokenTime;
6+
7+
private final long restQps;
8+
9+
public TokenUpdateStatus(long lastAddTokenTime, long restQps) {
10+
this.lastAddTokenTime = lastAddTokenTime;
11+
this.restQps = restQps;
12+
}
13+
14+
public long getLastAddTokenTime() {
15+
return lastAddTokenTime;
16+
}
17+
18+
public long getRestQps() {
19+
return restQps;
20+
}
21+
22+
@Override
23+
public String toString() {
24+
return "TokenUpdateStatus{" +
25+
"hash=" + System.identityHashCode(this) +
26+
", lastAddTokenTime=" + lastAddTokenTime +
27+
", requestCount=" + restQps +
28+
'}';
29+
}
30+
}

‎sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowCheckerTest.java

+3-2
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.Map;
3232
import java.util.concurrent.TimeUnit;
3333
import java.util.concurrent.atomic.AtomicLong;
34+
import java.util.concurrent.atomic.AtomicReference;
3435

3536
import static org.junit.Assert.assertFalse;
3637
import static org.junit.Assert.assertTrue;
@@ -158,7 +159,7 @@ public void testPassLocalCheckForCollection() throws InterruptedException {
158159
ParameterMetric metric = new ParameterMetric();
159160
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
160161
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
161-
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
162+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
162163

163164
assertTrue(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, list));
164165
assertFalse(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, list));
@@ -215,7 +216,7 @@ public Object paramFlowKey() {
215216
ParameterMetric metric = new ParameterMetric();
216217
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
217218
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
218-
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
219+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
219220

220221
assertTrue(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, args));
221222
assertFalse(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, args));

‎sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowDefaultCheckerTest.java

+5-10
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,7 @@ public void testCheckQpsWithLongIntervalAndHighThreshold() {
6262
ParameterMetric metric = new ParameterMetric();
6363
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
6464
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
65-
metric.getRuleTokenCounterMap().put(rule,
66-
new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
65+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
6766

6867
// We mock the time directly to avoid unstable behaviour.
6968
setCurrentMillis(mocked, System.currentTimeMillis());
@@ -102,8 +101,7 @@ public void testParamFlowDefaultCheckSingleQps() {
102101
ParameterMetric metric = new ParameterMetric();
103102
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
104103
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
105-
metric.getRuleTokenCounterMap().put(rule,
106-
new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
104+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
107105

108106
// We mock the time directly to avoid unstable behaviour.
109107
setCurrentMillis(mocked, System.currentTimeMillis());
@@ -144,8 +142,7 @@ public void testParamFlowDefaultCheckSingleQpsWithBurst() throws InterruptedExce
144142
ParameterMetric metric = new ParameterMetric();
145143
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
146144
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
147-
metric.getRuleTokenCounterMap().put(rule,
148-
new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
145+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
149146

150147
// We mock the time directly to avoid unstable behaviour.
151148
setCurrentMillis(mocked, System.currentTimeMillis());
@@ -216,8 +213,7 @@ public void testParamFlowDefaultCheckQpsInDifferentDuration() throws Interrupted
216213
ParameterMetric metric = new ParameterMetric();
217214
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
218215
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
219-
metric.getRuleTokenCounterMap().put(rule,
220-
new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
216+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
221217

222218
// We mock the time directly to avoid unstable behaviour.
223219
setCurrentMillis(mocked, System.currentTimeMillis());
@@ -268,8 +264,7 @@ public void testParamFlowDefaultCheckSingleValueCheckQpsMultipleThreads() throws
268264
ParameterMetric metric = new ParameterMetric();
269265
ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
270266
metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
271-
metric.getRuleTokenCounterMap().put(rule,
272-
new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));
267+
metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<>(4000));
273268
int threadCount = 40;
274269

275270
final CountDownLatch waitLatch = new CountDownLatch(threadCount);

‎sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowSlotTest.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424
import static org.mockito.Mockito.when;
2525

2626
import java.util.Collections;
27-
import java.util.concurrent.atomic.AtomicInteger;
2827
import java.util.concurrent.atomic.AtomicLong;
28+
import java.util.concurrent.atomic.AtomicReference;
2929

3030
import org.junit.After;
3131
import org.junit.Before;
@@ -100,9 +100,9 @@ public void testEntryWhenParamFlowExists() throws Throwable {
100100
ParameterMetric metric = mock(ParameterMetric.class);
101101

102102
CacheMap<Object, AtomicLong> map = new ConcurrentLinkedHashMapWrapper<>(4000);
103-
CacheMap<Object, AtomicLong> map2 = new ConcurrentLinkedHashMapWrapper<>(4000);
103+
CacheMap<Object, AtomicReference<TokenUpdateStatus>> map2 = new ConcurrentLinkedHashMapWrapper<>(4000);
104104
when(metric.getRuleTimeCounter(rule)).thenReturn(map);
105-
when(metric.getRuleTokenCounter(rule)).thenReturn(map2);
105+
when(metric.getRuleStampedTokenCounter(rule)).thenReturn(map2);
106106
map.put(argToGo, new AtomicLong(TimeUtil.currentTimeMillis()));
107107

108108
// Insert the mock metric to control pass or block.

0 commit comments

Comments
 (0)