package weightroundrobin

import (
    "fmt"
    "strings"
)

// 每一個后端服務(wù)定義
type BackendServer struct {
    // 實例權(quán)重
    Weight int
    // 當前的權(quán)重，初始為Weight
    currentWeight int
    // 后端服務(wù)名稱
    ServerName string
}

// 通過權(quán)重實現(xiàn)調(diào)用輪詢的定義
type WeightServerRoundRobin struct {
    // 所有有效的權(quán)重總和
    effectiveWeight int
    // 后端服務(wù)列表
    backendServerList []*BackendServer
}

// 創(chuàng)建一個負載輪詢器
func NewWeightServerRoundRobin() *WeightServerRoundRobin {
    return WeightServerRoundRobin{
        effectiveWeight: 0,
    }
}

// 增加后端服務(wù)名稱和權(quán)重
func (r *WeightServerRoundRobin) AddBackendServer(backendServer *BackendServer) {
    r.effectiveWeight += backendServer.Weight
    r.backendServerList = append(r.backendServerList, backendServer)
}

// 更具權(quán)重獲取一個后端服務(wù)名稱
func (r *WeightServerRoundRobin) GetBackendServer() *BackendServer {
    var expectBackendServer *BackendServer
    for _, backendServer := range r.backendServerList {
        // 給每個后端服務(wù)增加自身權(quán)重
        backendServer.currentWeight += backendServer.Weight
        if expectBackendServer == nil {
            expectBackendServer = backendServer
        }
        if backendServer.currentWeight > expectBackendServer.currentWeight {
            expectBackendServer = backendServer
        }
    }
    r.VisitBackendServerCurrentWeight()
    // 把選擇的后端服務(wù)權(quán)重減掉總權(quán)重
    expectBackendServer.currentWeight -= r.effectiveWeight
    return expectBackendServer
}

// 打印后端服務(wù)的當前權(quán)重變化
func (r *WeightServerRoundRobin) VisitBackendServerCurrentWeight() {
    var serverListForLog []string
    for _, backendServer := range r.backendServerList {
        serverListForLog = append(serverListForLog,
            fmt.Sprintf("%v", backendServer.currentWeight))
    }
    fmt.Printf("(%v)\n", strings.Join(serverListForLog, ", "))
}

package weightroundrobin

import (
    "fmt"
    "testing"
)

func TestNewWeightServerRoundRobin(t *testing.T) {
    weightServerRoundRobin := NewWeightServerRoundRobin()
    weightServerRoundRobin.AddBackendServer(BackendServer{
        ServerName: "ServerA",
        Weight: 5,
    })
    weightServerRoundRobin.AddBackendServer(BackendServer{
        ServerName: "ServerB",
        Weight: 3,
    })
    weightServerRoundRobin.AddBackendServer(BackendServer{
        ServerName: "ServerC",
        Weight: 1,
    })

    expectServerNameList := []string{
        "ServerA", "ServerB", "ServerA", "ServerC", "ServerA", "ServerB", "ServerA", "ServerB", "ServerA",
        //"ServerA", "ServerB", "ServerA", "ServerC", "ServerA", "ServerB", "ServerA", "ServerB", "ServerA",
    }
    fmt.Printf("(A, B, C)\n")
    for ii, expectServerName := range expectServerNameList {
        weightServerRoundRobin.VisitBackendServerCurrentWeight()
        backendServer := weightServerRoundRobin.GetBackendServer()
        if backendServer.ServerName != expectServerName {
            t.Errorf("%v.%v.expect:%v, actual:%v", t.Name(), ii, expectServerName, backendServer.ServerName)
            return
        }
    }
}