From 87885a3b546838318014bf379e0ad7f8f5f1f270 Mon Sep 17 00:00:00 2001
From: hein <hein.hpq@gmail.com>
Date: Mon, 23 Dec 2024 18:07:26 +0800
Subject: [PATCH] optimize: Caching the ConsistentHashSelector to avoid
 remapping when select every time.

---
 .../ConsistentHashLoadBalance.java            | 41 +++++++++++++++++-
 .../loadbalance/LoadBalanceTest.java          | 42 ++++++++++++++++++-
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/discovery/seata-discovery-core/src/main/java/org/apache/seata/discovery/loadbalance/ConsistentHashLoadBalance.java b/discovery/seata-discovery-core/src/main/java/org/apache/seata/discovery/loadbalance/ConsistentHashLoadBalance.java
index 1156bf00319..da0ff018593 100644
--- a/discovery/seata-discovery-core/src/main/java/org/apache/seata/discovery/loadbalance/ConsistentHashLoadBalance.java
+++ b/discovery/seata-discovery-core/src/main/java/org/apache/seata/discovery/loadbalance/ConsistentHashLoadBalance.java
@@ -45,9 +45,48 @@ public class ConsistentHashLoadBalance implements LoadBalance {
     private static final int VIRTUAL_NODES_NUM = ConfigurationFactory.getInstance().getInt(
         LOAD_BALANCE_CONSISTENT_HASH_VIRTUAL_NODES, VIRTUAL_NODES_DEFAULT);
 
+    /**
+     * The ConsistentHashSelectorWrapper that caches a {@link ConsistentHashSelector}.
+     */
+    private volatile ConsistentHashSelectorWrapper selectorWrapper;
+
     @Override
+    @SuppressWarnings("unchecked")
     public <T> T select(List<T> invokers, String xid) {
-        return new ConsistentHashSelector<>(invokers, VIRTUAL_NODES_NUM).select(xid);
+        if (selectorWrapper == null) {
+            synchronized (this) {
+                if (selectorWrapper == null) {
+                    selectorWrapper = new ConsistentHashSelectorWrapper(
+                            new ConsistentHashSelector<>(invokers, VIRTUAL_NODES_NUM), invokers.hashCode());
+                }
+            }
+        }
+        return (T) selectorWrapper.getSelector(invokers).select(xid);
+    }
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    private static final class ConsistentHashSelectorWrapper {
+
+        private volatile ConsistentHashSelector selector;
+        private volatile int invokersHashcode;
+
+        public ConsistentHashSelectorWrapper(ConsistentHashSelector selector, int invokersHashcode) {
+            this.selector = selector;
+            this.invokersHashcode = invokersHashcode;
+        }
+
+        public ConsistentHashSelector getSelector(List invokers) {
+            int hashCode;
+            if ((hashCode = invokers.hashCode()) != invokersHashcode) {
+                synchronized (this) {
+                    if (hashCode != invokersHashcode) {
+                        selector = new ConsistentHashSelector(invokers, VIRTUAL_NODES_NUM);
+                        invokersHashcode = hashCode;
+                    }
+                }
+            }
+            return selector;
+        }
     }
 
     private static final class ConsistentHashSelector<T> {
diff --git a/discovery/seata-discovery-core/src/test/java/org/apache/seata/discovery/loadbalance/LoadBalanceTest.java b/discovery/seata-discovery-core/src/test/java/org/apache/seata/discovery/loadbalance/LoadBalanceTest.java
index 8ca741799ab..d424a56b36a 100644
--- a/discovery/seata-discovery-core/src/test/java/org/apache/seata/discovery/loadbalance/LoadBalanceTest.java
+++ b/discovery/seata-discovery-core/src/test/java/org/apache/seata/discovery/loadbalance/LoadBalanceTest.java
@@ -21,11 +21,14 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 
+import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Stream;
 
@@ -85,7 +88,7 @@ public void testXIDLoadBalance_select(List<InetSocketAddress> addresses) throws
         Assertions.assertNotNull(inetSocketAddress);
         // test not found tc channel
         inetSocketAddress = loadBalance.select(addresses, "127.0.0.1:8199:123456");
-        Assertions.assertNotEquals(inetSocketAddress.getPort(),8199);
+        Assertions.assertNotEquals(inetSocketAddress.getPort(), 8199);
     }
 
     /**
@@ -108,6 +111,27 @@ public void testConsistentHashLoadBalance_select(List<InetSocketAddress> address
         Assertions.assertEquals(1, selected, "selected must be equal to 1");
     }
 
+    /**
+     * Test cached consistent hash load balance select.
+     *
+     * @param addresses the addresses
+     */
+    @ParameterizedTest
+    @MethodSource("addressProvider")
+    public void testCachedConsistentHashLoadBalance_select(List<InetSocketAddress> addresses) throws Exception {
+        ConsistentHashLoadBalance loadBalance = new ConsistentHashLoadBalance();
+        loadBalance.select(addresses, XID);
+        Object o1 = getConsistentHashSelectorByReflect(loadBalance);
+        loadBalance.select(addresses, XID);
+        Object o2 = getConsistentHashSelectorByReflect(loadBalance);
+        Assertions.assertEquals(o1, o2);
+        addresses = new ArrayList<>(addresses);
+        addresses.remove(ThreadLocalRandom.current().nextInt(addresses.size()));
+        loadBalance.select(addresses, XID);
+        Object o3 = getConsistentHashSelectorByReflect(loadBalance);
+        Assertions.assertNotEquals(o1, o3);
+    }
+
     /**
      * Test least active load balance select.
      *
@@ -166,6 +190,22 @@ public Map<InetSocketAddress, AtomicLong> getSelectedCounter(int runs, List<Inet
         return counter;
     }
 
+    /**
+     * Gets ConsistentHashSelector Instance By Reflect
+     *
+     * @param loadBalance the loadBalance
+     * @return the ConsistentHashSelector
+     */
+    public Object getConsistentHashSelectorByReflect(ConsistentHashLoadBalance loadBalance) throws Exception {
+        Field selectorWrapperField = ConsistentHashLoadBalance.class.getDeclaredField("selectorWrapper");
+        selectorWrapperField.setAccessible(true);
+        Object selectWrapper = selectorWrapperField.get(loadBalance);
+        Assertions.assertNotNull(selectWrapper);
+        Field selectorField = selectWrapper.getClass().getDeclaredField("selector");
+        selectorField.setAccessible(true);
+        return selectorField.get(selectWrapper);
+    }
+
     /**
      * Address provider object [ ] [ ].
      *