1.4.5.2

EdgeDNS/internal/models/record_ids.go

@@ -0,0 +1,172 @@
+// 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
+}