waf-platform/EdgeDNS/internal/nodes/manager_record.go

// 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()
}