1.4.5.2

EdgeDNS/internal/models/records_route.go

@@ -0,0 +1,134 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+//go:build plus
+
+package models
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
+	"github.com/iwind/TeaGo/types"
+)
+
+func RouteIdString(routeId int64) string {
+	return "id:" + types.String(routeId)
+}
+
+type RouteRecords struct {
+	routeRecordsMap map[string]*RecordIds // routeCode => { recordId1, recordId2, ... }
+	recordsMap      map[int64]*NSRecord   // recordId => *NSRecord
+}
+
+func NewRouteRecords() *RouteRecords {
+	return &RouteRecords{
+		routeRecordsMap: map[string]*RecordIds{},
+		recordsMap:      map[int64]*NSRecord{},
+	}
+}
+
+func (this *RouteRecords) Add(record *NSRecord) {
+	// 先删除
+	this.remove(record.Id)
+
+	// 添加记录
+	this.recordsMap[record.Id] = record
+
+	// 添加线路
+	var routeIds = record.RouteIds
+	if len(routeIds) == 0 || (len(routeIds) == 1 && routeIds[0] == "") {
+		routeIds = []string{"default"}
+	}
+
+	for _, routeId := range routeIds {
+		recordIds, ok := this.routeRecordsMap[routeId]
+		if !ok {
+			recordIds = NewRecordIds()
+			this.routeRecordsMap[routeId] = recordIds
+		}
+		recordIds.Add(record.Id, record.Weight)
+	}
+}
+
+// Find 查找与线路匹配的记录
+// strictMode 表示是否严格匹配线路
+func (this *RouteRecords) Find(routeCodes []string, config *dnsconfigs.NSAnswerConfig, strictMode bool) (records []*NSRecord, routeCode string) {
+	if config == nil {
+		config = dnsconfigs.DefaultNSAnswerConfig()
+	}
+
+	var maxSize = int(config.MaxSize)
+	if maxSize <= 0 {
+		maxSize = dnsconfigs.NSAnswerDefaultSize
+	}
+
+	// 查找匹配的线路
+	for _, routeId := range routeCodes {
+		recordIds, ok := this.routeRecordsMap[routeId]
+		if ok && !recordIds.IsEmpty() {
+			return this.recordsWithIds(recordIds, config.Mode, maxSize), routeId
+		}
+	}
+
+	// 查找默认线路
+	recordIds, ok := this.routeRecordsMap["default"]
+	if ok && !recordIds.IsEmpty() {
+		return this.recordsWithIds(recordIds, config.Mode, maxSize), "default"
+	}
+
+	// 随机一个
+	if !strictMode {
+		for _, record := range this.recordsMap {
+			return []*NSRecord{record}, "default"
+		}
+	}
+
+	return nil, ""
+}
+
+func (this *RouteRecords) Remove(recordId int64) {
+	this.remove(recordId)
+}
+
+func (this *RouteRecords) IsEmpty() bool {
+	return len(this.recordsMap) == 0
+}
+
+func (this *RouteRecords) remove(recordId int64) {
+	oldRecord, ok := this.recordsMap[recordId]
+	if !ok {
+		return
+	}
+
+	delete(this.recordsMap, recordId)
+
+	var oldRouteIds = oldRecord.RouteIds
+	if len(oldRouteIds) == 0 || (len(oldRouteIds) == 1 && oldRouteIds[0] == "") {
+		oldRouteIds = []string{"default"}
+	}
+
+	for _, routeId := range oldRouteIds {
+		recordIds, ok := this.routeRecordsMap[routeId]
+		if ok {
+			recordIds.Remove(recordId)
+			if recordIds.IsEmpty() {
+				delete(this.routeRecordsMap, routeId)
+			}
+		}
+	}
+}
+
+func (this *RouteRecords) recordsWithIds(recordIds *RecordIds, mode dnsconfigs.NSAnswerMode, maxSize int) (records []*NSRecord) {
+	// round-robin
+	if mode == dnsconfigs.NSAnswerModeRoundRobin {
+		var recordId = recordIds.NextId()
+		if recordId > 0 {
+			return []*NSRecord{this.recordsMap[recordId]}
+		}
+	}
+
+	// random
+	var randomIds = recordIds.RandomIds(maxSize)
+	for _, randomId := range randomIds {
+		records = append(records, this.recordsMap[randomId])
+	}
+
+	return
+}