Initial commit (code only without large binaries)

EdgeDNS/internal/nodes/manager_record.go

@@ -0,0 +1,254 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeDNS/internal/dbs"
+	"github.com/TeaOSLab/EdgeDNS/internal/models"
+	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeDNS/internal/rpc"
+	"sync"
+	"time"
+)
+
+// RecordManager 记录管理器
+type RecordManager struct {
+	recordsMap map[int64]*models.DomainRecords // domainId => RecordsMap
+
+	db      *dbs.DB
+	locker  sync.RWMutex
+	version int64
+
+	notifier chan bool
+}
+
+// NewRecordManager 获取新记录管理器对象
+func NewRecordManager(db *dbs.DB) *RecordManager {
+	return &RecordManager{
+		db:         db,
+		recordsMap: map[int64]*models.DomainRecords{},
+		notifier:   make(chan bool, 8),
+	}
+}
+
+// Start 启动自动任务
+func (this *RecordManager) Start() {
+	remotelogs.Println("RECORD_MANAGER", "starting ...")
+
+	// 从本地数据库中加载数据
+	err := this.Load()
+	if err != nil {
+		if rpc.IsConnError(err) {
+			remotelogs.Debug("RECORD_MANAGER", "load failed: "+err.Error())
+		} else {
+			remotelogs.Error("RECORD_MANAGER", "load failed: "+err.Error())
+		}
+	}
+
+	// 初始化运行
+	err = this.LoopAll()
+	if err != nil {
+		if rpc.IsConnError(err) {
+			remotelogs.Debug("RECORD_MANAGER", "loop failed: "+err.Error())
+		} else {
+			remotelogs.Error("RECORD_MANAGER", "loop failed: "+err.Error())
+		}
+	}
+
+	// 更新
+	var ticker = time.NewTicker(30 * time.Second)
+	for {
+		select {
+		case <-ticker.C:
+		case <-this.notifier:
+		}
+
+		err := this.LoopAll()
+		if err != nil {
+			if rpc.IsConnError(err) {
+				remotelogs.Debug("RECORD_MANAGER", "loop failed: "+err.Error())
+			} else {
+				remotelogs.Error("RECORD_MANAGER", "loop failed: "+err.Error())
+			}
+		}
+	}
+}
+
+// Load 从数据库中加载数据
+func (this *RecordManager) Load() error {
+	var offset = 0
+	var size = 10000
+	for {
+		records, err := this.db.ListRecords(offset, size)
+		if err != nil {
+			return err
+		}
+		if len(records) == 0 {
+			break
+		}
+
+		this.locker.Lock()
+		for _, record := range records {
+			domainRecords, ok := this.recordsMap[record.DomainId]
+			if !ok {
+				domainRecords = models.NewDomainRecords()
+				this.recordsMap[record.DomainId] = domainRecords
+			}
+			domainRecords.Add(record)
+
+			if record.Version > this.version {
+				this.version = record.Version
+			}
+		}
+		this.locker.Unlock()
+
+		offset += size
+	}
+
+	if this.version > 0 {
+		this.version++
+	}
+
+	return nil
+}
+
+func (this *RecordManager) LoopAll() error {
+	for {
+		hasNext, err := this.Loop()
+		if err != nil {
+			return err
+		}
+		if !hasNext {
+			break
+		}
+	}
+	return nil
+}
+
+// Loop 单次循环任务
+func (this *RecordManager) Loop() (hasNext bool, err error) {
+	client, err := rpc.SharedRPC()
+	if err != nil {
+		return false, err
+	}
+	resp, err := client.NSRecordRPC.ListNSRecordsAfterVersion(client.Context(), &pb.ListNSRecordsAfterVersionRequest{
+		Version: this.version,
+		Size:    20000,
+	})
+	if err != nil {
+		return false, err
+	}
+	var records = resp.NsRecords
+	if len(records) == 0 {
+		return false, nil
+	}
+	for _, record := range records {
+		this.processRecord(record)
+		if record.Version > this.version {
+			this.version = record.Version
+		}
+	}
+	this.version++
+	return true, nil
+}
+
+func (this *RecordManager) FindRecords(domainId int64, routeCodes []string, recordName string, recordType dnsconfigs.RecordType, strictMode bool) (records []*models.NSRecord, routeCode string) {
+	this.locker.RLock()
+	domainRecords, ok := this.recordsMap[domainId]
+	if ok {
+		records, routeCode = domainRecords.Find(routeCodes, recordName, recordType, sharedNodeConfig.Answer, strictMode)
+	}
+	this.locker.RUnlock()
+	return
+}
+
+// NotifyUpdate 通知更新
+func (this *RecordManager) NotifyUpdate() {
+	select {
+	case this.notifier <- true:
+	default:
+	}
+}
+
+// 处理单条记录
+func (this *RecordManager) processRecord(record *pb.NSRecord) {
+	if record.NsDomain == nil {
+		return
+	}
+
+	if !record.IsOn || record.IsDeleted {
+		this.locker.Lock()
+		domainRecords, ok := this.recordsMap[record.NsDomain.Id]
+		if ok {
+			domainRecords.Remove(record.Id)
+		}
+		this.locker.Unlock()
+
+		// 从数据库中删除
+		if this.db != nil {
+			err := this.db.DeleteRecord(record.Id)
+			if err != nil {
+				remotelogs.Error("RECORD_MANAGER", "delete record from db failed: "+err.Error())
+			}
+		}
+
+		return
+	}
+
+	// 存入数据库
+	if this.db != nil {
+		exists, err := this.db.ExistsRecord(record.Id)
+		if err != nil {
+			remotelogs.Error("RECORD_MANAGER", "query failed: "+err.Error())
+		} else {
+			var routeIds = []string{}
+			for _, route := range record.NsRoutes {
+				routeIds = append(routeIds, route.Code)
+			}
+
+			if exists {
+				err = this.db.UpdateRecord(record.Id, record.NsDomain.Id, record.Name, record.Type, record.Value, record.MxPriority, record.SrvPriority, record.SrvWeight, record.SrvPort, record.CaaFlag, record.CaaTag, record.Ttl, record.Weight, routeIds, record.Version)
+				if err != nil {
+					remotelogs.Error("RECORD_MANAGER", "update failed: "+err.Error())
+				}
+			} else {
+				err = this.db.InsertRecord(record.Id, record.NsDomain.Id, record.Name, record.Type, record.Value, record.MxPriority, record.SrvPriority, record.SrvWeight, record.SrvPort, record.CaaFlag, record.CaaTag, record.Ttl, record.Weight, routeIds, record.Version)
+				if err != nil {
+					remotelogs.Error("RECORD_MANAGER", "insert failed: "+err.Error())
+				}
+			}
+		}
+	}
+
+	// 加入缓存Map
+	this.locker.Lock()
+	domainRecords, ok := this.recordsMap[record.NsDomain.Id]
+	if !ok {
+		domainRecords = models.NewDomainRecords()
+		this.recordsMap[record.NsDomain.Id] = domainRecords
+	}
+	var routeIds = []string{}
+	for _, r := range record.NsRoutes {
+		routeIds = append(routeIds, r.Code)
+	}
+	domainRecords.Add(&models.NSRecord{
+		Id:          record.Id,
+		Name:        record.Name,
+		Type:        record.Type,
+		Value:       record.Value,
+		MXPriority:  record.MxPriority,
+		SRVPriority: record.SrvPriority,
+		SRVWeight:   record.SrvWeight,
+		SRVPort:     record.SrvPort,
+		CAAFlag:     record.CaaFlag,
+		CAATag:      record.CaaTag,
+		Ttl:         record.Ttl,
+		Weight:      record.Weight,
+		Version:     record.Version,
+		RouteIds:    routeIds,
+		DomainId:    record.NsDomain.Id,
+	})
+	this.locker.Unlock()
+}