173 lines
3.0 KiB
Go
173 lines
3.0 KiB
Go
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
|
||
//go:build plus
|
||
|
||
package models
|
||
|
||
import (
|
||
"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
|
||
"github.com/iwind/TeaGo/rands"
|
||
"math/rand"
|
||
"time"
|
||
)
|
||
|
||
func init() {
|
||
rand.Seed(time.Now().UnixNano())
|
||
}
|
||
|
||
type recordIdInfo struct {
|
||
Id int64
|
||
Weight int32
|
||
}
|
||
|
||
type RecordIds struct {
|
||
IdList []*recordIdInfo
|
||
IdBucket []int64
|
||
|
||
RoundIndex int
|
||
|
||
totalWeight int64
|
||
}
|
||
|
||
func NewRecordIds() *RecordIds {
|
||
return &RecordIds{}
|
||
}
|
||
|
||
func (this *RecordIds) IsEmpty() bool {
|
||
return len(this.IdList) == 0
|
||
}
|
||
|
||
func (this *RecordIds) Add(newId int64, weight int32) {
|
||
if weight <= 0 {
|
||
weight = 10
|
||
}
|
||
|
||
const maxWeight = 999999
|
||
if weight > maxWeight {
|
||
weight = maxWeight
|
||
}
|
||
|
||
// 检查是否存在
|
||
for _, idInfo := range this.IdList {
|
||
if idInfo.Id == newId {
|
||
return
|
||
}
|
||
}
|
||
|
||
// 添加
|
||
this.IdList = append(this.IdList, &recordIdInfo{
|
||
Id: newId,
|
||
Weight: weight,
|
||
})
|
||
|
||
// 重置数据
|
||
this.resetData()
|
||
}
|
||
|
||
func (this *RecordIds) Remove(oldId int64) {
|
||
defer this.resetData()
|
||
|
||
var newIdList = []*recordIdInfo{}
|
||
for _, idInfo := range this.IdList {
|
||
if idInfo.Id == oldId {
|
||
continue
|
||
}
|
||
newIdList = append(newIdList, idInfo)
|
||
}
|
||
this.IdList = newIdList
|
||
}
|
||
|
||
// NextId for round-robin
|
||
func (this *RecordIds) NextId() int64 {
|
||
var l = len(this.IdList)
|
||
|
||
if l == 0 {
|
||
return 0
|
||
}
|
||
|
||
if l == 1 {
|
||
return this.IdList[0].Id
|
||
}
|
||
|
||
if this.RoundIndex > l-1 {
|
||
this.RoundIndex = 0
|
||
}
|
||
|
||
var id = this.IdList[this.RoundIndex].Id
|
||
|
||
this.RoundIndex++
|
||
return id
|
||
}
|
||
|
||
func (this *RecordIds) RandomIds(count int) []int64 {
|
||
if count <= 0 {
|
||
count = dnsconfigs.NSAnswerDefaultSize
|
||
}
|
||
|
||
var totalRecords = len(this.IdList)
|
||
|
||
if totalRecords == 0 {
|
||
return nil
|
||
}
|
||
|
||
if totalRecords == 1 {
|
||
return []int64{this.IdList[0].Id} // duplicate
|
||
}
|
||
|
||
if totalRecords < count {
|
||
count = totalRecords
|
||
}
|
||
|
||
var totalIds = len(this.IdBucket)
|
||
var startIndex = rands.Int(0, totalIds-1)
|
||
var endIndex = startIndex + count - 1
|
||
if endIndex <= totalIds-1 {
|
||
return this.IdBucket[startIndex : endIndex+1]
|
||
}
|
||
return append(this.IdBucket[startIndex:totalIds], this.IdBucket[0:endIndex-totalIds+1]...)
|
||
}
|
||
|
||
func (this *RecordIds) resetData() {
|
||
this.resetWeight()
|
||
}
|
||
|
||
func (this *RecordIds) resetWeight() {
|
||
var totalWeight int64
|
||
|
||
var weightMap = map[int32]bool{} // weight => bool
|
||
var hasUniqueWeights = false
|
||
var ids []int64
|
||
|
||
for _, idInfo := range this.IdList {
|
||
totalWeight += int64(idInfo.Weight)
|
||
|
||
// 检查是否有不同的权重
|
||
if len(weightMap) > 0 && !weightMap[idInfo.Weight] {
|
||
hasUniqueWeights = true
|
||
}
|
||
weightMap[idInfo.Weight] = true
|
||
|
||
ids = append(ids, idInfo.Id)
|
||
}
|
||
|
||
// 根据权重，重新组织IDs
|
||
if hasUniqueWeights {
|
||
var newIds = []int64{}
|
||
for _, idInfo := range this.IdList {
|
||
for i := int32(0); i < idInfo.Weight; i++ {
|
||
newIds = append(newIds, idInfo.Id)
|
||
}
|
||
}
|
||
ids = newIds
|
||
}
|
||
|
||
var countIds = len(ids)
|
||
if countIds > 0 {
|
||
rand.Shuffle(countIds, func(i, j int) {
|
||
ids[i], ids[j] = ids[j], ids[i]
|
||
})
|
||
}
|
||
|
||
this.totalWeight = totalWeight
|
||
this.IdBucket = ids
|
||
}
|