Fix potential concurrency issue when updating flow rules (#1783)

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleManager.java

@@ -16,12 +16,14 @@
 package com.alibaba.csp.sentinel.slots.block.flow;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import com.alibaba.csp.sentinel.concurrent.NamedThreadFactory;
 import com.alibaba.csp.sentinel.log.RecordLog;
@@ -43,10 +45,11 @@ import com.alibaba.csp.sentinel.property.SentinelProperty;
  *
  * @author jialiang.linjl
  * @author Eric Zhao
+ * @author Weihua
  */
 public class FlowRuleManager {
 
-    private static final Map<String, List<FlowRule>> flowRules = new ConcurrentHashMap<String, List<FlowRule>>();
+    private static final AtomicReference<Map<String, List<FlowRule>>> flowRules = new AtomicReference<Map<String, List<FlowRule>>>();
 
     private static final FlowPropertyListener LISTENER = new FlowPropertyListener();
     private static SentinelProperty<List<FlowRule>> currentProperty = new DynamicSentinelProperty<List<FlowRule>>();
@@ -56,6 +59,7 @@ public class FlowRuleManager {
         new NamedThreadFactory("sentinel-metrics-record-task", true));
 
     static {
+        flowRules.set(Collections.<String, List<FlowRule>>emptyMap());
         currentProperty.addListener(LISTENER);
         SCHEDULER.scheduleAtFixedRate(new MetricTimerListener(), 0, 1, TimeUnit.SECONDS);
     }
@@ -83,7 +87,7 @@ public class FlowRuleManager {
      */
     public static List<FlowRule> getRules() {
         List<FlowRule> rules = new ArrayList<FlowRule>();
-        for (Map.Entry<String, List<FlowRule>> entry : flowRules.entrySet()) {
+        for (Map.Entry<String, List<FlowRule>> entry : flowRules.get().entrySet()) {
             rules.addAll(entry.getValue());
         }
         return rules;
@@ -99,11 +103,11 @@ public class FlowRuleManager {
     }
 
     static Map<String, List<FlowRule>> getFlowRuleMap() {
-        return flowRules;
+        return flowRules.get();
     }
 
     public static boolean hasConfig(String resource) {
-        return flowRules.containsKey(resource);
+        return flowRules.get().containsKey(resource);
     }
 
     public static boolean isOtherOrigin(String origin, String resourceName) {
@@ -111,7 +115,7 @@ public class FlowRuleManager {
             return false;
         }
 
-        List<FlowRule> rules = flowRules.get(resourceName);
+        List<FlowRule> rules = flowRules.get().get(resourceName);
 
         if (rules != null) {
             for (FlowRule rule : rules) {
@@ -129,21 +133,17 @@ public class FlowRuleManager {
         @Override
         public void configUpdate(List<FlowRule> value) {
             Map<String, List<FlowRule>> rules = FlowRuleUtil.buildFlowRuleMap(value);
-            if (rules != null) {
-                flowRules.clear();
-                flowRules.putAll(rules);
-            }
-            RecordLog.info("[FlowRuleManager] Flow rules received: {}", flowRules);
+            //the rules was always not null, it's no need to check nullable
+            //remove checking to avoid IDE warning
+            flowRules.set(rules);
+            RecordLog.info("[FlowRuleManager] Flow rules received: {}", rules);
         }
 
         @Override
         public void configLoad(List<FlowRule> conf) {
             Map<String, List<FlowRule>> rules = FlowRuleUtil.buildFlowRuleMap(conf);
-            if (rules != null) {
-                flowRules.clear();
-                flowRules.putAll(rules);
-            }
-            RecordLog.info("[FlowRuleManager] Flow rules loaded: {}", flowRules);
+            flowRules.set(rules);
+            RecordLog.info("[FlowRuleManager] Flow rules loaded: {}", rules);
         }
     }
 

sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleManagerTest.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 1999-2018 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author Weihua
+ */
+public class FlowRuleManagerTest {
+
+    public static final List<FlowRule> STATIC_RULES_1 = new ArrayList<FlowRule>();
+    public static final List<FlowRule> STATIC_RULES_2 = new ArrayList<FlowRule>();
+
+    static {
+        FlowRule first = new FlowRule();
+        first.setResource("/a/b/c");
+        first.setCount(100);
+        STATIC_RULES_1.add(first);
+
+        FlowRule second = new FlowRule();
+        second.setResource("/a/b/c");
+        second.setCount(200);
+        STATIC_RULES_2.add(second);
+
+    }
+
+    @Test
+    public void testLoadAndGetRules(){
+        FlowRuleManager.loadRules(STATIC_RULES_1);
+        assertEquals(1, FlowRuleManager.getRules().size()); // the initial size
+        new Thread(loader, "Loader").start();
+
+        for(int i = 0; i < 10000; i++){
+            //The initial size is 1, and the size after updating should also be 1,
+            //if the actual size is 0, that must be called after clear(),
+            // but before putAll() in FlowPropertyListener.configUpdate
+            assertEquals(1, FlowRuleManager.getRules().size());
+        }
+    }
+
+    public Runnable loader = new Runnable() {
+        @Override
+        public void run() {
+            for(int i = 0; i < 10000; i++){
+                //to guarantee that they're different and change happens
+                FlowRuleManager.loadRules(i % 2 == 0 ? STATIC_RULES_2 : STATIC_RULES_1);
+            }
+        }
+    };
+
+}