web_huakui_server/web/service/srv_search/es.go

package srv_search

import (
	"context"
	"errors"
	"fmt"
	"hkdj-web-server/pub/common/log"
	"hkdj-web-server/pub/common/srvdef"
	"hkdj-web-server/pub/common/utils"
	"hkdj-web-server/pub/esmanager"
	"hkdj-web-server/pub/models/esmod"
	"hkdj-web-server/pub/models/videoinformationmod"
	"strconv"
	"time"

	"go.mongodb.org/mongo-driver/bson/primitive"
)

var syncIsRuining = false

func StartES() {
	log.Info("srv_search:Start...")
	defer log.Info("srv_search:Start!!!")
	// 初始化 ES 索引 mapping（已存在时报错可忽略）
	//InitArticleESIndex()
	//InitToolESIndex()
	startSyncEs()
	_, _ = skdCron.AddFunc("0 0/17 * * * ?", func() {
		startSyncEs()
	})
	skdCron.Start()
}

func startSyncEs() {
	ctx := context.Background()
	updateEsData(ctx, esmod.GetVideoIndex(), getVideoEsData)
}

func StopES() {
	log.Info("srv_search:Stop...")
	defer log.Info("srv_search:Stop!!!")
}

func updateEsData(ctx context.Context, esIndex string, fn func(ctx context.Context, startTime, endTime *time.Time, esIndex string) (esmanager.M, *srvdef.SrvErr)) {
	log.Info("srv_search:do...", log.Any("ESIndex", esIndex))
	defer log.Info("srv_search:do end", log.Any("EsIndex", esIndex))
	if fn == nil {
		return
	}
	esRecord, serr := esmod.GetRecordTimeByESIndex(ctx, esIndex)
	if serr != nil {
		log.Error(fmt.Sprintf("srv_search:startEngine GetRecordByESIndex failed! Index:%v,err:%+v:", esIndex, serr))
		return
	}
	isNew := false
	if esRecord == nil {
		isNew = true
		esRecord = &esmod.ESRecord{
			Id:         primitive.NewObjectID(),
			EsIndex:    esIndex,
			RecordTime: time.Time{},
			Timestamp:  esmod.ESTimestamp{},
			UpdatedAt:  time.Time{},
			CreatedAt:  time.Now(),
		}
	}
	var now = time.Now()
	var esData esmanager.M
	// 跨天凌晨5点重置所有数据
	if isNew || ((now.Hour() >= 5 && esRecord.UpdatedAt.Local().Hour() < 5) || esRecord.UpdatedAt.Local().Day() < now.Day()) {
		esData, serr = fn(ctx, nil, nil, esIndex)
		if serr != nil {
			log.Warn("updateEsData fn return err", log.Any("index", esIndex), log.E(serr))
			return
		}
		_ = delAllFromES(esIndex)
		isUpdate, err := commit(esData, esIndex, time.Now())
		if err != nil {
			log.Error("Elastic sync video data RunSync commit err! ", log.E(err))
			return
		}

		if isUpdate {
			// 有成功过就算成功
			if err := esmod.SetOneEsRecordTime(nil, esIndex, now); err != nil {
				return
			}
		}
	} else {
		timestamp := esRecord.RecordTime
		var recentMinute = utils.RecentMinute(now, esmod.FiveMinuteScale)
		for recentMinute.After(timestamp) {
			var end = timestamp.Add(esmod.FiveMinuteScale * time.Minute)
			if end.After(recentMinute) {
				end = recentMinute
			}
			esData, serr = fn(ctx, &timestamp, &end, esIndex)
			if serr != nil {
				log.Warn("updateEsData fn return err", log.Any("index", esIndex), log.E(serr))
				return
			}
			if esData != nil && len(esData) > 0 {
				isUpdate, err := commit(esData, esIndex, time.Now())
				if err != nil {
					log.Error("Elastic sync video data RunSync commit err! ", log.E(err))
					return
				}
				if isUpdate {
					// 有成功过就算成功
					if err := esmod.SetOneEsRecordTime(nil, esIndex, now); err != nil {
						return
					}
				}
			}
			timestamp = end
		}
	}
}

func commit(data esmanager.M, esIndex string, syncTime time.Time) (bool, error) {
	bulkData := handleBulkCountLimit(data)
	isHaveSuccess := false
	if bulkData != nil && len(bulkData) > 0 {
		for _, doc := range bulkData {
			if err := esmanager.Instance.Bulk(esIndex, doc); err != nil {
				log.Warn(fmt.Sprintf("Elastic sync commit video data SyncDataToES failed!, err:%+v\n", err))
				// 出错继续
			} else {
				isHaveSuccess = true
			}
		}
	}
	if isHaveSuccess {
		return isHaveSuccess, nil
	} else {
		return isHaveSuccess, errors.New("Elastic sync commit video data UpsertOneTrans failed!")
	}
}

func handleBulkCountLimit(data esmanager.M) []esmanager.M {
	if data == nil || len(data) == 0 {
		return make([]esmanager.M, 0)
	}
	rsl := make([]esmanager.M, 0)
	rsl = append(rsl, make(esmanager.M, batchSize))
	index := 0
	for k, v := range data {
		rsl[index][k] = v
		if len(rsl[index]) >= batchSize {
			index++
			rsl = append(rsl, make(esmanager.M, batchSize))
		}
	}
	return rsl
}

// DelVideoFromES 从 ES 实时删除单条视频（视频删除/下架时调用）
func DelVideoFromES(id int64) {
	if esmanager.Instance == nil {
		return
	}
	strId := strconv.FormatInt(id, 10)
	if err := batchDelFromES(esmod.GetVideoIndex(), []string{strId}); err != nil {
		log.Warn("srv_search:DelVideoFromES failed", log.Any("id", id), log.E(err))
	}
}

// UpsertVideoToES 将单条视频实时写入/更新 ES（视频上架时调用）
func UpsertVideoToES(info *videoinformationmod.VideoInformation) {
	if esmanager.Instance == nil {
		return
	}
	esData := videoToEsData([]*videoinformationmod.VideoInformation{info})
	if err := AddMultiDataToES(esmod.GetVideoIndex(), esData); err != nil {
		log.Warn("srv_search:UpsertVideoToES failed", log.Any("id", info.Id), log.E(err))
	}
}