算法之权重轮询算法

Posted 2022-11-29 浅弋、璃鱼

实现:

package transform

import (
	"fmt"
	"strings"
	"sync"

	"project/pkg/base/log"
)

// 权重轮询调度算法

// 每一个服务定义
type RecvServer struct  // nolint
	Weight        int
	currentWeight int
	RecvAddr      string


// 通过权重实现调用轮询的定义
type WeightServerRoundRobin struct  // nolint
	effectiveWeight int
	recvServerList  []*RecvServer
	sync.Mutex


// NewWeightServerRoundRobin 创建一个负载轮询器
func NewWeightServerRoundRobin() *WeightServerRoundRobin 
	return &WeightServerRoundRobin
		effectiveWeight: 0,
	


// AddRecvServer 新增一个负载节点
func (r *WeightServerRoundRobin) AddRecvServer(recvServer *RecvServer) 
	r.Lock()
	defer r.Unlock()
	r.effectiveWeight += recvServer.Weight
	r.recvServerList = append(r.recvServerList, recvServer)


// GetRecvServer 获取可用的负载的节点
func (r *WeightServerRoundRobin) GetRecvServer() *RecvServer 
	var expectRecvServer *RecvServer
	if len(r.recvServerList) == 0 
		// 没有可用的
		return expectRecvServer
	
	r.Lock()
	defer r.Unlock()
	for _, recvServer := range r.recvServerList 
		// 给每个服务增加自身权重
		recvServer.currentWeight += recvServer.Weight
		if expectRecvServer == nil 
			expectRecvServer = recvServer
		
		if recvServer.currentWeight > expectRecvServer.currentWeight 
			expectRecvServer = recvServer
		
	
	// r.VisitRecvServerCurrentWeight()
	// 把选择的服务权重减掉总权重
	expectRecvServer.currentWeight -= r.effectiveWeight
	return expectRecvServer


// UpdateWeighted 更新recvServer节点的权重newWeighted
func (r *WeightServerRoundRobin) UpdateWeighted(recvServer *RecvServer, newWeighted int) 
	r.Lock()
	defer r.Unlock()
	subWeight := recvServer.Weight - newWeighted
	recvServer.Weight = newWeighted
	r.effectiveWeight -= subWeight


// VisitRecvServerCurrentWeight 打印服务的当前权重变化
func (r *WeightServerRoundRobin) VisitRecvServerCurrentWeight() 
	serverListForLog := []string
	for _, recvServer := range r.recvServerList 
		serverListForLog = append(serverListForLog,
			fmt.Sprintf("%v", recvServer.currentWeight))
	
	log.Infof("(%v)\\n", strings.Join(serverListForLog, ", "))

测试:

package transform

import (
	"fmt"
	"testing"
)

func TestRecvServerRoundRobin1(t *testing.T) 
	weightServerRoundRobin := NewWeightServerRoundRobin()
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerA",
		Weight:   5,
	)
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerB",
		Weight:   3,
	)
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerC",
		Weight:   2,
	)

	for i := 0; i < 10; i++ 
		weightServerRoundRobin.VisitRecvServerCurrentWeight()
		recvServer := weightServerRoundRobin.GetRecvServer()
		if recvServer == nil 
			t.Error("获取不到可用的负载")
			return
		
		fmt.Printf("%d[%v], actual:%v\\n", i, t.Name(), recvServer.RecvAddr)
	


func TestRecvServerRoundRobin2(t *testing.T) 
	weightServerRoundRobin := NewWeightServerRoundRobin()
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerA",
		Weight:   5,
	)
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerB",
		Weight:   3,
	)
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerC",
		Weight:   2,
	)

	fmt.Println("(A, B, C)")
	for i := 0; i <= 10; i++ 
		weightServerRoundRobin.VisitRecvServerCurrentWeight()
		recvServer := weightServerRoundRobin.GetRecvServer()
		if recvServer == nil 
			t.Error("获取不到可用的负载")
			return
		
		// rn := rand.Intn(10)
		rn := 2
		if i == 0 
			weightServerRoundRobin.UpdateWeighted(recvServer, rn)
		
		fmt.Printf("%d[%v], actual:%v; effectiveWeight: %d\\n", rn, t.Name(), recvServer.RecvAddr, weightServerRoundRobin.effectiveWeight)
	


func TestRecvServerRoundRobin3(t *testing.T) 
	weightServerRoundRobin := NewWeightServerRoundRobin()
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerA",
		Weight:   5,
	)
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerB",
		Weight:   3,
	)
	weightServerRoundRobin.AddRecvServer(&RecvServer
		RecvAddr: "ServerC",
		Weight:   1,
	)

	expectServerNameList := []string
		"ServerA", "ServerB", "ServerA", "ServerC", "ServerA", "ServerB", "ServerA", "ServerB", "ServerA",
		"ServerA", "ServerB", "ServerA", "ServerC", "ServerA", "ServerB", "ServerA", "ServerB", "ServerA",
	
	fmt.Println("(A, B, C)")
	for ii, expectServerName := range expectServerNameList 
		weightServerRoundRobin.VisitRecvServerCurrentWeight()
		recvServer := weightServerRoundRobin.GetRecvServer()
		if recvServer == nil 
			t.Error("获取不到可用的负载")
			return
		
		if recvServer.RecvAddr != expectServerName 
			t.Errorf("%v.%v.expect:%v, actual:%v", t.Name(), ii, expectServerName, recvServer.RecvAddr)
			return