Merge pull request #2 from lunasaw/dev_1.0.3

Dev 1.0.3 实现多种负载均衡算法，增加注入配置

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>io.github.lunasaw</groupId>
     <artifactId>zlm-spring-boot-starter</artifactId>
-    <version>1.0.2</version>
+    <version>1.0.3</version>
 
     <name>zlm-spring-boot-starter</name>
     <description>ZLMediaKit restful api and hook starter</description>
@@ -21,7 +21,7 @@
         <httpclient5.version>5.2.1</httpclient5.version>
         <github.username>lunasaw</github.username>
         <app.profiles>${project.name}</app.profiles>
-        <zlm-spring-boot-starter.version>1.0.2</zlm-spring-boot-starter.version>
+        <zlm-spring-boot-starter.version>1.0.3</zlm-spring-boot-starter.version>
     </properties>
 
     <developers>

src/main/java/io/github/lunasaw/zlm/config/ZlmAutoConfiguration.java

@@ -2,6 +2,9 @@
 
 import io.github.lunasaw.zlm.hook.service.ZlmHookService;
 import io.github.lunasaw.zlm.hook.service.impl.DefaultZlmHookServiceImpl;
+import io.github.lunasaw.zlm.node.LoadBalancer;
+import io.github.lunasaw.zlm.node.impl.*;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
@@ -21,9 +24,31 @@
 @ConditionalOnProperty(prefix = "zlm", name = "enable", havingValue = "true")
 public class ZlmAutoConfiguration {
 
+    @Autowired
+    ZlmProperties zlmProperties;
+
     @Bean
     @ConditionalOnMissingBean
     public ZlmHookService zlmHookService() {
         return new DefaultZlmHookServiceImpl();
     }
+
+    @Bean
+    @ConditionalOnMissingBean
+    public LoadBalancer loadBalancer() {
+        switch (zlmProperties.getBalance()) {
+            case RANDOM:
+                return new RandomLoadBalancer();
+            case ROUND_ROBIN:
+                return new RoundRobinLoadBalancer();
+            case CONSISTENT_HASHING:
+                return new ConsistentHashingLoadBalander();
+            case WEIGHT_RANDOM:
+                return new WeightRandomLoadBalancer();
+            case WEIGHT_ROUND_ROBIN:
+                return new WeightRoundRobinLoadBalancer();
+            default:
+                throw new RuntimeException("未找到负载均衡器");
+        }
+    }
 }

src/main/java/io/github/lunasaw/zlm/config/ZlmNodeConfig.java → src/main/java/io/github/lunasaw/zlm/config/ZlmNode.java

@@ -9,7 +9,7 @@
  * @description: zlm节点配置
  */
 @Data
-public class ZlmNodeConfig {
+public class ZlmNode {
 
     /**
      * The id of this node.
@@ -31,5 +31,13 @@ public class ZlmNodeConfig {
      */
     private boolean enabled = true;
 
+    /**
+     * Whether enable hook.
+     */
     private boolean hookEnabled = true;
+
+    /**
+     * The weight of this host.
+     */
+    private int weight = 100;
 }

src/main/java/io/github/lunasaw/zlm/config/ZlmProperties.java

@@ -1,13 +1,14 @@
 package io.github.lunasaw.zlm.config;
 
+import io.github.lunasaw.zlm.enums.LoadBalancerEnums;
 import lombok.Data;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.boot.context.properties.ConfigurationProperties;
 
-import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.stream.Collectors;
 
 /**
@@ -23,14 +24,31 @@ public class ZlmProperties implements InitializingBean {
     /**
      * 对外NodeMap
      */
-    private Map<String, ZlmNodeConfig> nodeMap = new HashMap<>();
+    public static Map<String, ZlmNode> nodeMap = new ConcurrentHashMap<>();
 
-    private List<ZlmNodeConfig> nodes = new ArrayList<>();
+    public static List<ZlmNode> nodes = new CopyOnWriteArrayList<>();
 
     private boolean enable = true;
 
+    private LoadBalancerEnums balance = LoadBalancerEnums.ROUND_ROBIN;
+
+    public static void addNode(ZlmNode zlmNode) {
+        if (zlmNode != null) {
+            nodes.add(zlmNode);
+            nodeMap.put(zlmNode.getServerId(), zlmNode);
+        }
+    }
+
+    public Map<String, ZlmNode> getNodeMap() {
+        return nodeMap;
+    }
+
+    public List<ZlmNode> getNodes() {
+        return nodes;
+    }
+
     @Override
     public void afterPropertiesSet() {
-        nodeMap = nodes.stream().filter(ZlmNodeConfig::isEnabled).collect(Collectors.toMap(ZlmNodeConfig::getServerId, node -> node));
+        nodeMap = nodes.stream().filter(ZlmNode::isEnabled).collect(Collectors.toMap(ZlmNode::getServerId, node -> node));
     }
 }

src/main/java/io/github/lunasaw/zlm/enums/LoadBalancerEnums.java

@@ -0,0 +1,29 @@
+package io.github.lunasaw.zlm.enums;
+
+import lombok.Getter;
+
+/**
+ * @author weidian
+ */
+
+@Getter
+public enum LoadBalancerEnums {
+    RANDOM(1, "Random"),
+
+    ROUND_ROBIN(2, "RoundRobin"),
+
+    CONSISTENT_HASHING(3, "ConsistentHashing"),
+
+    WEIGHT_ROUND_ROBIN(4, "WeightRoundRobin"),
+
+    WEIGHT_RANDOM(5, "WeightRandom");
+
+    private int code;
+    private String type;
+
+    private LoadBalancerEnums(int code, String type) {
+        this.code = code;
+        this.type = type;
+    }
+
+}

src/main/java/io/github/lunasaw/zlm/node/LoadBalancer.java

@@ -0,0 +1,14 @@
+package io.github.lunasaw.zlm.node;
+
+import io.github.lunasaw.zlm.config.ZlmNode;
+
+/**
+ * @author luna
+ * @date 2024/1/5
+ */
+public interface LoadBalancer {
+
+    ZlmNode selectNode(String key);
+
+    String getType();
+}

src/main/java/io/github/lunasaw/zlm/node/NodeService.java

@@ -1,19 +1,19 @@
 package io.github.lunasaw.zlm.node;
 
-import io.github.lunasaw.zlm.api.ZlmRestService;
-import io.github.lunasaw.zlm.config.ZlmNodeConfig;
 import io.github.lunasaw.zlm.config.ZlmProperties;
-import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.stereotype.Service;
 
-import java.util.Map;
-
 /**
  * @author luna
  * @date 2023/12/4
  */
 @Service
+@ConditionalOnBean(ZlmProperties.class)
 public class NodeService {
 
 
+    public void method() {
+
+    }
 }

src/main/java/io/github/lunasaw/zlm/node/impl/ConsistentHashingLoadBalander.java

@@ -0,0 +1,65 @@
+package io.github.lunasaw.zlm.node.impl;
+
+import io.github.lunasaw.zlm.config.ZlmNode;
+import io.github.lunasaw.zlm.config.ZlmProperties;
+import io.github.lunasaw.zlm.enums.LoadBalancerEnums;
+import io.github.lunasaw.zlm.node.LoadBalancer;
+
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * @author luna
+ * @date 2024/1/5
+ */
+public class ConsistentHashingLoadBalander implements LoadBalancer {
+
+    static final TreeMap<Integer, String> VIRTUAL_NODE_MAP = new TreeMap<>();
+
+    public ConsistentHashingLoadBalander() {
+        init();
+    }
+
+    public static void init() {
+        for (String nodeName : ZlmProperties.nodeMap.keySet()) {
+            ZlmNode nodeConfig = ZlmProperties.nodeMap.get(nodeName);
+            int weight = nodeConfig.getWeight();
+            for (int i = 0; i < weight * 10; i++) {
+                String virtualNodeName = nodeName + "#" + i;
+                int hash = getHash(virtualNodeName);
+                VIRTUAL_NODE_MAP.put(hash, virtualNodeName);
+            }
+        }
+    }
+
+    private static int getHash(String str) {
+        final int p = 16777619;
+        int hash = (int) 2166136261L;
+        for (int i = 0; i < str.length(); i++) {
+            hash = (hash ^ str.charAt(i)) * p;
+        }
+        hash += hash << 13;
+        hash ^= hash >> 7;
+        hash += hash << 3;
+        hash ^= hash >> 17;
+        hash += hash << 5;
+        return hash;
+    }
+
+    @Override
+    public ZlmNode selectNode(String key) {
+        int hash = getHash(key);
+        Map.Entry<Integer, String> entry = VIRTUAL_NODE_MAP.ceilingEntry(hash);
+        if (entry == null) {
+            entry = VIRTUAL_NODE_MAP.firstEntry();
+        }
+        String virtualNodeName = entry.getValue();
+        String nodeName = virtualNodeName.split("#")[0];
+        return ZlmProperties.nodeMap.get(nodeName);
+    }
+
+    @Override
+    public String getType() {
+        return LoadBalancerEnums.CONSISTENT_HASHING.getType();
+    }
+}

src/main/java/io/github/lunasaw/zlm/node/impl/RandomLoadBalancer.java

@@ -0,0 +1,27 @@
+package io.github.lunasaw.zlm.node.impl;
+
+import io.github.lunasaw.zlm.config.ZlmNode;
+import io.github.lunasaw.zlm.config.ZlmProperties;
+import io.github.lunasaw.zlm.enums.LoadBalancerEnums;
+import io.github.lunasaw.zlm.node.LoadBalancer;
+
+import java.util.Random;
+
+/**
+ * @author luna
+ * @date 2024/1/5
+ */
+public class RandomLoadBalancer implements LoadBalancer {
+    static final Random RANDOM = new Random();
+
+    @Override
+    public ZlmNode selectNode(String key) {
+        int index = RANDOM.nextInt(ZlmProperties.nodes.size());
+        return ZlmProperties.nodes.get(index);
+    }
+
+    @Override
+    public String getType() {
+        return LoadBalancerEnums.RANDOM.getType();
+    }
+}

src/main/java/io/github/lunasaw/zlm/node/impl/RoundRobinLoadBalancer.java

@@ -0,0 +1,28 @@
+package io.github.lunasaw.zlm.node.impl;
+
+import io.github.lunasaw.zlm.config.ZlmNode;
+import io.github.lunasaw.zlm.config.ZlmProperties;
+import io.github.lunasaw.zlm.enums.LoadBalancerEnums;
+import io.github.lunasaw.zlm.node.LoadBalancer;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * @author luna
+ * @date 2024/1/5
+ */
+public class RoundRobinLoadBalancer implements LoadBalancer {
+
+    static final AtomicInteger SEQUENCE = new AtomicInteger(0);
+
+    @Override
+    public ZlmNode selectNode(String key) {
+        int index = SEQUENCE.getAndIncrement() % ZlmProperties.nodes.size();
+        return ZlmProperties.nodes.get(index);
+    }
+
+    @Override
+    public String getType() {
+        return LoadBalancerEnums.RANDOM.getType();
+    }
+}

src/main/java/io/github/lunasaw/zlm/node/impl/WeightRandomLoadBalancer.java

@@ -0,0 +1,60 @@
+package io.github.lunasaw.zlm.node.impl;
+
+import io.github.lunasaw.zlm.config.ZlmNode;
+import io.github.lunasaw.zlm.config.ZlmProperties;
+import io.github.lunasaw.zlm.enums.LoadBalancerEnums;
+import io.github.lunasaw.zlm.node.LoadBalancer;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * @author luna
+ * @date 2024/1/5
+ */
+public class WeightRandomLoadBalancer implements LoadBalancer {
+
+    static final AtomicInteger TOTAL_WEIGHT = new AtomicInteger(0);
+    static final Map<String, Integer> WEIGHT_MAP = new ConcurrentHashMap<>();
+    static final Random RANDOM = new Random();
+
+
+    public WeightRandomLoadBalancer() {
+        init();
+    }
+
+    public static void init() {
+        List<String> nodeList = new ArrayList<>(ZlmProperties.nodeMap.keySet());
+        int size = nodeList.size();
+        if (size > 0) {
+            for (String nodeName : nodeList) {
+                ZlmNode nodeConfig = ZlmProperties.nodeMap.get(nodeName);
+                int weight = nodeConfig.getWeight();
+                WEIGHT_MAP.put(nodeName, weight);
+                TOTAL_WEIGHT.addAndGet(weight);
+            }
+        }
+    }
+
+    @Override
+    public ZlmNode selectNode(String key) {
+        int randomNum = RANDOM.nextInt(TOTAL_WEIGHT.get());
+        int currentWeight = 0;
+        for (String nodeName : ZlmProperties.nodeMap.keySet()) {
+            currentWeight += WEIGHT_MAP.get(nodeName);
+            if (currentWeight > randomNum) {
+                return ZlmProperties.nodeMap.get(nodeName);
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public String getType() {
+        return LoadBalancerEnums.WEIGHT_RANDOM.getType();
+    }
+}