135 lines
3.2 KiB
Go
135 lines
3.2 KiB
Go
// 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
|
|
}
|