feat: improve code (#193)

* feat: improve code

* feat: improve code

* feat: improve code

Author: renyijiu

internal/api/api.go

@@ -403,14 +403,15 @@ func (r *RequestClient) CreateRecord(parent string, rc *openDpsV1alpha1Resource.
 
 	// record's labels may have not been ensured
 	recordLabels := rc.GetLabels()
+	ensuredLabels := make([]*openDpsV1alpha1Resource.Label, 0, len(recordLabels))
 	for _, label := range recordLabels {
 		l, err := r.ensureLabel(parent, label.GetDisplayName())
 		if err != nil {
 			return nil, err
 		}
-		recordLabels = append(recordLabels, l)
+		ensuredLabels = append(ensuredLabels, l)
 	}
-	rc.SetLabels(recordLabels)
+	rc.SetLabels(ensuredLabels)
 
 	req := openDpsV1alpha1Service.CreateRecordRequest{
 		Parent: parent,

internal/collector/collector.go

@@ -234,10 +234,18 @@ func handleRecordCaches(uploadChan chan string, reqClient *api.RequestClient, co
 		if rcName, ok := rc.Record["name"].(string); ok && rcName != "" {
 			select {
 			case uploadChan <- rc.GetRecordCachePath():
-				log.Infof("Record cache %s is ready to upload", rcName)
+				log.WithFields(log.Fields{
+					"recordCache": rc.GetRecordCachePath(),
+					"recordName":  rcName,
+					"project":     rc.ProjectName,
+				}).Infof("record cache is ready to upload")
 			default:
 				recordSet.Release(rc.GetRecordCachePath())
-				log.Infof("Upload channel is full, skip uploading record cache %s", rcName)
+				log.WithFields(log.Fields{
+					"recordCache": rc.GetRecordCachePath(),
+					"recordName":  rcName,
+					"project":     rc.ProjectName,
+				}).Infof("upload channel is full, skip enqueueing record cache")
 			}
 		} else {
 			recordSet.Release(rc.GetRecordCachePath())
@@ -552,6 +560,17 @@ func createRecord(deviceInfo *openDpsV1alpha1Resource.Device, recordCache *model
 
 	title := getRecordTitle(recordCache)
 	description := getRecordDescription(recordCache)
+	ruleName, _ := recordCache.DiagnosisTask["rule_name"].(string)
+	ruleDisplayName, _ := recordCache.DiagnosisTask["rule_display_name"].(string)
+	recordLog := log.WithFields(log.Fields{
+		"recordCache":     recordCache.GetRecordCachePath(),
+		"project":         recordCache.ProjectName,
+		"title":           title,
+		"ruleName":        ruleName,
+		"ruleDisplayName": ruleDisplayName,
+		"device":          deviceInfo.GetName(),
+		"totalFiles":      len(recordCache.OriginalFiles),
+	})
 
 	labels := make([]*openDpsV1alpha1Resource.Label, 0)
 	for _, label := range recordCache.Labels {
@@ -573,18 +592,18 @@ func createRecord(deviceInfo *openDpsV1alpha1Resource.Device, recordCache *model
 		Device:      deviceInfo,
 		Metadata:    metadata,
 	}
-	ruleName, ok := recordCache.DiagnosisTask["rule_name"].(string)
-	if ok {
+	if ruleName != "" {
 		record.Rules = []*openDpsV1alpha1Resource.DiagnosisRule{
 			{
 				Name: ruleName,
 			},
 		}
 	}
 
+	recordLog.Infof("creating cloud record")
 	record, err := reqClient.CreateRecord(recordCache.ProjectName, record)
 	if err != nil {
-		log.Errorf("create record failed: %v", err)
+		recordLog.Errorf("create record failed: %v", err)
 		return
 	}
 	recordCache.Record = map[string]interface{}{
@@ -604,7 +623,9 @@ func createRecord(deviceInfo *openDpsV1alpha1Resource.Device, recordCache *model
 		latest.Record = recordCache.Record
 		return nil
 	}); err != nil {
-		log.Errorf("save record cache failed: %v", err)
+		recordLog.WithField("recordName", record.GetName()).Errorf("save record cache failed after record creation: %v", err)
+	} else {
+		recordLog.WithField("recordName", record.GetName()).Infof("created cloud record and saved record cache")
 	}
 }
 

internal/mod/http/server/rules.go

@@ -19,13 +19,13 @@ import (
 	"encoding/json"
 	"net/http"
 	"slices"
+	"sync/atomic"
 
 	"github.com/ThreeDotsLabs/watermill"
 	gcmessage "github.com/ThreeDotsLabs/watermill/message"
 	"github.com/ThreeDotsLabs/watermill/pubsub/gochannel"
 	"github.com/coscene-io/coscout/pkg/constant"
 	"github.com/coscene-io/coscout/pkg/rule_engine"
-	mapset "github.com/deckarep/golang-set/v2"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/time/rate"
 )
@@ -106,7 +106,9 @@ func RulesHandler(pubSub *gochannel.GoChannel) func(w http.ResponseWriter, r *ht
 }
 
 func ActiveTopicsHandler(ctx context.Context, pubSub *gochannel.GoChannel) func(w http.ResponseWriter, r *http.Request) {
-	activeTopics := mapset.NewSet[string]()
+	var activeTopics atomic.Value
+	activeTopics.Store([]string{})
+
 	go func() {
 		msg, err := pubSub.Subscribe(ctx, constant.TopicConfigTopicsMsg)
 		if err != nil {
@@ -131,14 +133,19 @@ func ActiveTopicsHandler(ctx context.Context, pubSub *gochannel.GoChannel) func(
 					continue
 				}
 				log.Infof("Received active topics: %v", topics)
-				activeTopics = mapset.NewSet(topics...)
+				activeTopics.Store(topics)
 			}
 		}
 	}()
 
 	return func(w http.ResponseWriter, r *http.Request) {
+		topics, ok := activeTopics.Load().([]string)
+		if !ok {
+			log.Errorf("Unexpected active topics value type")
+			topics = []string{}
+		}
 		res := ActiveTopicsResponse{
-			Topics: activeTopics.ToSlice(),
+			Topics: topics,
 		}
 		bytes, err := json.Marshal(res)
 		if err != nil {

internal/mod/rule/engine.go

@@ -17,6 +17,7 @@ package rule
 import (
 	"encoding/json"
 	"strconv"
+	"sync"
 	"time"
 
 	"buf.build/gen/go/coscene-io/coscene-openapi/protocolbuffers/go/coscene/openapi/dataplatform/v1alpha1/resources"
@@ -42,9 +43,10 @@ const (
 
 // Engine represents the rule engine that processes messages against rules.
 type Engine struct {
-	reqClient  api.RequestClient
-	deviceName string
+	reqClient api.RequestClient
 
+	mu           sync.RWMutex
+	deviceName   string
 	rules        []*rule_engine.Rule
 	activeTopics mapset.Set[string]
 
@@ -120,28 +122,62 @@ func (e *Engine) UpdateRules(apiRules []*resources.DiagnosisRule, configTopics [
 			}
 		}
 	}
-	e.rules = rules
-
 	configTopicSet := mapset.NewSet(configTopics...)
+	var newActiveTopics mapset.Set[string]
 	if hasWildcardTopicRule {
-		e.activeTopics = mapset.NewSet[string]()
+		newActiveTopics = mapset.NewSet[string]()
 	} else {
-		e.activeTopics = activeTopics.Intersect(configTopicSet)
+		newActiveTopics = activeTopics.Intersect(configTopicSet)
 	}
 
+	e.mu.Lock()
+	e.rules = rules
+	e.activeTopics = newActiveTopics
+	e.mu.Unlock()
+
 	log.Infof("Updated %d valid rules", len(rules))
 }
 
 // ActiveTopics returns all the active topics in the rule engine.
-func (e *Engine) ActiveTopics() mapset.Set[string] {
-	return e.activeTopics
+func (e *Engine) ActiveTopics() []string {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+	if e.activeTopics == nil {
+		return []string{}
+	}
+	return e.activeTopics.ToSlice()
+}
+
+func (e *Engine) activeTopicSet() mapset.Set[string] {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+	if e.activeTopics == nil {
+		return mapset.NewSet[string]()
+	}
+	return e.activeTopics.Clone()
+}
+
+func (e *Engine) SetDeviceName(name string) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	e.deviceName = name
+}
+
+func (e *Engine) getDeviceName() string {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+	return e.deviceName
 }
 
 // ConsumeNext shows how to process a message through the rule engine.
 func (e *Engine) ConsumeNext(item rule_engine.RuleItem) {
 	log.Debugf("consuming message: %+v", item)
 
-	for _, rule := range e.rules {
+	e.mu.RLock()
+	rules := e.rules
+	e.mu.RUnlock()
+
+	for _, rule := range rules {
 		if rule.Topics.Cardinality() > 0 && !rule.Topics.Contains(item.Topic) {
 			log.Debugf("rule %s topics %v does not match topic %s, skipping", rule.Metadata["rule_display_name"], rule.Topics, item.Topic)
 			continue
@@ -169,9 +205,17 @@ func (e *Engine) ConsumeNext(item rule_engine.RuleItem) {
 
 		log.Debugf("rule %s: isActive: %v, preActivationTime: %v, msgTs: %v", rule.Metadata["rule_display_name"], isActive, prevActivationTime, msgTs)
 		if isActive {
-			log.Infof("rule %s met, save collect info", rule.Metadata["rule_display_name"])
-
 			collectInfoId := uuid.New().String()
+			ruleLog := log.WithFields(log.Fields{
+				"collectID":       collectInfoId,
+				"ruleName":        rule.Metadata["rule_name"],
+				"ruleDisplayName": rule.Metadata["rule_display_name"],
+				"topic":           item.Topic,
+				"source":          item.Source,
+				"ts":              item.Ts,
+			})
+			ruleLog.Infof("rule met, save collect info")
+
 			additionalArgs := map[string]interface{}{}
 			for k, v := range rule.Metadata {
 				additionalArgs[k] = v
@@ -180,11 +224,13 @@ func (e *Engine) ConsumeNext(item rule_engine.RuleItem) {
 
 			for _, action := range rule.Actions {
 				if err := action.Run(curActivation, additionalArgs); err != nil {
-					log.Errorf("failed to run action %s: %v", action.Name, err)
+					ruleLog.WithField("action", action.Name).Errorf("failed to run action: %v", err)
 				}
 			}
 			if err := model.PublishCollectInfo(collectInfoId); err != nil {
-				log.Errorf("failed to publish collect info %s: %v", collectInfoId, err)
+				ruleLog.Errorf("failed to publish collect info: %v", err)
+			} else {
+				ruleLog.Infof("published collect info")
 			}
 		}
 	}
@@ -239,8 +285,6 @@ func (e *Engine) cleanupDebounceTime() {
 // uploadActionImpl is the implementation of "upload" function,
 // it stores the collect info in the database.
 func (e *Engine) uploadActionImpl(kwargs map[string]interface{}) error {
-	log.Infof("triggered saving collect info")
-
 	// Extract required fields with type assertions
 	beforeStr, ok := kwargs["before"].(string)
 	if !ok {
@@ -325,9 +369,20 @@ func (e *Engine) uploadActionImpl(kwargs map[string]interface{}) error {
 		return errors.Errorf("rule must be a DiagnosisRule object")
 	}
 
-	canUpload := e.canUpload(uploadLimit, rule)
-	if err := e.reqClient.HitDiagnosisRule(rule, e.deviceName, canUpload); err != nil {
-		return errors.Wrap(err, "failed to hit diagnosis rule")
+	deviceName := e.getDeviceName()
+	actionLog := log.WithFields(log.Fields{
+		"collectID":       collectInfoId,
+		"ruleName":        ruleName,
+		"ruleDisplayName": ruleDisplayName,
+		"project":         projectName,
+		"device":          deviceName,
+		"triggerTs":       triggerTsFloat,
+	})
+	actionLog.Infof("triggered saving collect info")
+
+	canUpload := e.canUpload(uploadLimit, rule, deviceName)
+	if err := e.reqClient.HitDiagnosisRule(rule, deviceName, canUpload); err != nil {
+		actionLog.Errorf("failed to hit diagnosis rule, continuing with save: %v", err)
 	}
 
 	// Calculate start and end times
@@ -360,43 +415,86 @@ func (e *Engine) uploadActionImpl(kwargs map[string]interface{}) error {
 		WhiteList:  whiteList,
 	}
 	collectInfo.Skip = !canUpload
+	actionLog.WithFields(log.Fields{
+		"canUpload": canUpload,
+		"start":     collectInfo.Cut.Start,
+		"end":       collectInfo.Cut.End,
+		"labels":    labels,
+	}).Infof("prepared collect info")
 
 	// Save the collect info
 	if err = collectInfo.SaveDraft(); err != nil {
 		return errors.Wrap(err, "failed to save collect info")
 	}
-	log.Infof("saved collect info %s", collectInfoId)
+	actionLog.Infof("saved collect info draft")
 
 	return nil
 }
 
-func (e *Engine) canUpload(uploadLimit *resources.UploadLimit, rule *resources.DiagnosisRule) bool {
+func (e *Engine) canUpload(uploadLimit *resources.UploadLimit, rule *resources.DiagnosisRule, deviceName string) bool {
+	ruleName := rule.GetName()
+	ruleDisplayName := rule.GetDisplayName()
 	if uploadLimit.HasDevice() {
-		count, err := e.reqClient.CountDiagnosisRuleHits(rule, e.deviceName)
+		count, err := e.reqClient.CountDiagnosisRuleHits(rule, deviceName)
 		if err != nil {
-			log.Errorf("failed to count diagnosis rule hits: %v, skipping", err)
-			return false
+			log.WithFields(log.Fields{
+				"ruleName":        ruleName,
+				"ruleDisplayName": ruleDisplayName,
+				"device":          deviceName,
+				"scope":           "device",
+			}).Errorf("failed to count diagnosis rule hits, defaulting to allow upload: %v", err)
+			return true
 		}
 
-		log.Infof("device count: %d, limit: %d", count, uploadLimit.GetDevice().GetTimes())
+		log.WithFields(log.Fields{
+			"ruleName":        ruleName,
+			"ruleDisplayName": ruleDisplayName,
+			"device":          deviceName,
+			"scope":           "device",
+			"count":           count,
+			"limit":           uploadLimit.GetDevice().GetTimes(),
+		}).Infof("checked diagnosis rule upload limit")
 
 		if count >= uploadLimit.GetDevice().GetTimes() {
-			log.Infof("device count %d exceeds limit %d, skipping", count, uploadLimit.GetDevice().GetTimes())
+			log.WithFields(log.Fields{
+				"ruleName":        ruleName,
+				"ruleDisplayName": ruleDisplayName,
+				"device":          deviceName,
+				"scope":           "device",
+				"count":           count,
+				"limit":           uploadLimit.GetDevice().GetTimes(),
+			}).Infof("diagnosis rule upload limit reached, skipping upload")
 			return false
 		}
 	}
 
 	if uploadLimit.HasGlobal() {
 		count, err := e.reqClient.CountDiagnosisRuleHits(rule, "")
 		if err != nil {
-			log.Errorf("failed to count diagnosis rule hits: %v, skipping", err)
-			return false
+			log.WithFields(log.Fields{
+				"ruleName":        ruleName,
+				"ruleDisplayName": ruleDisplayName,
+				"scope":           "global",
+			}).Errorf("failed to count diagnosis rule hits, defaulting to allow upload: %v", err)
+			return true
 		}
 
-		log.Infof("global count: %d, limit: %d", count, uploadLimit.GetGlobal().GetTimes())
+		log.WithFields(log.Fields{
+			"ruleName":        ruleName,
+			"ruleDisplayName": ruleDisplayName,
+			"scope":           "global",
+			"count":           count,
+			"limit":           uploadLimit.GetGlobal().GetTimes(),
+		}).Infof("checked diagnosis rule upload limit")
 
 		if count >= uploadLimit.GetGlobal().GetTimes() {
-			log.Infof("global count %d exceeds limit %d, skipping", count, uploadLimit.GetGlobal().GetTimes())
+			log.WithFields(log.Fields{
+				"ruleName":        ruleName,
+				"ruleDisplayName": ruleDisplayName,
+				"scope":           "global",
+				"count":           count,
+				"limit":           uploadLimit.GetGlobal().GetTimes(),
+			}).Infof("diagnosis rule upload limit reached, skipping upload")
 			return false
 		}
 	}