基于权重分配算法实现_java如何根据权重分配流量算法

一、在日常工作中，经常有一些针对资源分配的策略需要按权重设置，比如：某个rpc服务调用等等。今天聊一聊按权重分配资源的实现。

具体代码如下：

1.WeightRandom.java

import lombok.Data;
 
import java.io.Serializable;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
 
public class WeightRandom {
    private Random random = new Random();
    private List<Resource> resources;
    private int factor;
 
    public void init() {
        this.resources = new ArrayList<Resource>();
        resources.add(new Resource("a", 30));
        resources.add(new Resource("b", 50));
        resources.add(new Resource("c", 1000));
        resources.add(new Resource("d", 100));
 
        // 初始化
        this.factor = 0;
        for (Resource resource:resources) {
            resource.setBegin(factor);
            factor += resource.getWeight();
        }
    }
 
    /**
     * 对比区间开始点
     * @return
     */
    public Resource random() {
        Resource selected = null;
        int rand = random.nextInt(factor);
        for (int i = resources.size() - 1; i >= 0; i--) {
            Resource resource = resources.get(i);
            if (rand >= resource.getBegin()) {
                selected = resource;
                break;
            }
        }
        return selected;
    }
 
    /**
     * 对比区间结束点
     * @return
     */
    public Resource randomX() {
        Resource selected = null;
        int rv = random.nextInt(factor);
        int endpoint = 0;//区间结束点
        for (Resource resource:resources) {
            endpoint+=resource.getWeight();
            if (rv<=endpoint) {
                selected = resource;
                break;
            }
        }
 
        return selected;
    }
 
    public static void main(String[] args) {
        Map<String, AtomicInteger> count = new HashMap<>(8);
        WeightRandom w = new WeightRandom();
        w.init();
        for (int i = 0; i < 10000; i++) {
            Resource resource = w.random();
            AtomicInteger ai = new AtomicInteger(0);
            AtomicInteger old = count.putIfAbsent(resource.getName(),ai);
            if(old==null){
                old=ai;
            }
            old.incrementAndGet();
        }
        System.out.println(count);
 
        System.out.println("=====chooseX=====");
        count.clear();
 
        for (int i = 0; i < 10000; i++) {
            Resource server = w.randomX();
            AtomicInteger ai = new AtomicInteger(0);
            AtomicInteger old = count.putIfAbsent(server.getName(),ai);
            if(old==null){
                old=ai;
            }
            old.incrementAndGet();
        }
        System.out.println(count);
    }
 
    @Data
    public static class Resource implements Serializable {
        private static final long serialVersionUID = -7447211367731992379L;
        private String name;
        private int weight;
        private int begin;
 
        public Resource(String name, int weigth) {
            this.name = name;
            this.weight = weigth;
        }
    }
}

2.WeightRandomStrategy.java

import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
 
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
 
/**
 * Java 基于权重按比例分配算法,对比区间结束点
 * @param <K>
 * @param <V>
 */
public class WeightRandomStrategy<K, V extends Number> {
    private TreeMap<Double, K> weightMap = new TreeMap<>();
 
    public WeightRandomStrategy(List<Pair<K, V>> list) {
        for (Pair<K, V> pair : list) {
            double lastWeight = this.weightMap.size() == 0 ? 0 : this.weightMap.lastKey();
            this.weightMap.put(pair.getValue().doubleValue() + lastWeight, pair.getKey());
        }
    }
 
    /**
     * 依据权重返回 K
     * @return
     */
    public K random() {
        double randomWeight = this.weightMap.lastKey() * Math.random();
        SortedMap<Double, K> tailMap = this.weightMap.tailMap(randomWeight, false);
        return this.weightMap.get(tailMap.firstKey());
    }
 
    public static void main(String[] args) {
        List<Pair<String, Integer>> list = new ArrayList<>();
        list.add(new ImmutablePair<>("A", 90));
        list.add(new ImmutablePair<>("B", 10));
        list.add(new ImmutablePair<>("C", 30));
        list.add(new ImmutablePair<>("D", 200));
        WeightRandomStrategy<String, Integer> strategy = new WeightRandomStrategy<>(list);
        Map<String, AtomicInteger> count = new HashMap<>(8);
        for (int i = 0; i < 10000; i++) {
            String key=strategy.random();
            AtomicInteger ai = new AtomicInteger(0);
            AtomicInteger old = count.putIfAbsent(key,ai);
            if(old==null){
                old=ai;
            }
            old.incrementAndGet();
        }
        System.out.println(count);
    }
}

3.WeightRandomStrategyX.java

import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
 
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;
 
/**
 * Java 基于权重按比例分配算法,对比区间开始点
 * @param <K>
 * @param <V>
 */
public class WeightRandomStrategyX<K, V extends Number> {
    private TreeMap<Double, K> weightMap = new TreeMap<>();
    private double factor=0D;
    public WeightRandomStrategyX(List<Pair<K, V>> list) {
        for (Pair<K, V> pair : list) {
            this.weightMap.put(factor, pair.getKey());
            this.factor+=pair.getValue().doubleValue();
        }
    }
 
    /**
     * 依据权重返回 K
     * @return
     */
    public K random() {
        double randomWeight = ThreadLocalRandom.current().nextDouble(factor);
        SortedMap<Double, K> tailMap = this.weightMap.headMap(randomWeight, false);
        return this.weightMap.get(tailMap.lastKey());
    }
 
    public static void main(String[] args) {
        List<Pair<String, Integer>> list = new ArrayList<>();
        list.add(new ImmutablePair<>("A", 90));
        list.add(new ImmutablePair<>("B", 10));
        list.add(new ImmutablePair<>("C", 30));
        list.add(new ImmutablePair<>("D", 200));
        WeightRandomStrategyX<String, Integer> strategy = new WeightRandomStrategyX<>(list);
 
        Map<String, AtomicInteger> count = new HashMap<>(8);
        for (int i = 0; i < 10000; i++) {
            String key=strategy.random();
            AtomicInteger ai = new AtomicInteger(0);
            AtomicInteger old = count.putIfAbsent(key,ai);
            if(old==null){
                old=ai;
            }
            old.incrementAndGet();
        }
        System.out.println(count);
    }
}

其中WeightRandomStrategy的实现摘自https://dorole.com/2014/，是基于treeMap，非常巧妙！

注：以上实现若用于生产环境中需要做充分测试。

二、参考资料

https://dorole.com/2014/