improve(memory): label, label_value memory

Author: JinzhouChen-xmu

server/controller/prometheus/cache/label.go

@@ -114,13 +114,18 @@ func (l *label) Add(batch []*controller.PrometheusLabel) {
 }
 
 func (l *label) refresh(args ...interface{}) error {
+	var count int64
+	if err := l.org.DB.Model(&metadbmodel.PrometheusLabel{}).Count(&count).Error; err != nil {
+		return err
+	}
+
 	rows, err := l.org.DB.Model(&metadbmodel.PrometheusLabel{}).Select("id", "name", "value").Rows()
 	if err != nil {
 		return err
 	}
 	defer rows.Close()
 
-	newActive := make(map[LabelKey]int)
+	newActive := make(map[LabelKey]int, count)
 	for rows.Next() {
 		var id int
 		var name, value string

server/controller/prometheus/cache/label_value.go

@@ -98,13 +98,18 @@ func (lv *labelValue) Add(batch []*controller.PrometheusLabelValue) {
 }
 
 func (lv *labelValue) refresh(args ...interface{}) error {
+	var count int64
+	if err := lv.org.DB.Model(&metadbmodel.PrometheusLabelValue{}).Count(&count).Error; err != nil {
+		return err
+	}
+
 	rows, err := lv.org.DB.Model(&metadbmodel.PrometheusLabelValue{}).Select("id", "value").Rows()
 	if err != nil {
 		return err
 	}
 	defer rows.Close()
 
-	newActive := make(map[string]int)
+	newActive := make(map[string]int, count)
 	for rows.Next() {
 		var id int
 		var value string

server/controller/prometheus/encoder/label.go

@@ -32,6 +32,9 @@ type label struct {
 	lock         sync.Mutex
 	resourceType string
 	labelKeyToID map[cache.LabelKey]int
+
+	isRefreshing bool
+	pendingKeys  map[cache.LabelKey]int
 }
 
 func newLabel(org *common.ORG) *label {
@@ -43,32 +46,49 @@ func newLabel(org *common.ORG) *label {
 }
 
 func (l *label) store(item *metadbmodel.PrometheusLabel) {
-	l.labelKeyToID[cache.NewLabelKey(item.Name, item.Value)] = item.ID
+	key := cache.NewLabelKey(item.Name, item.Value)
+	l.labelKeyToID[key] = item.ID
+
+	if l.isRefreshing {
+		l.pendingKeys[key] = item.ID
+	}
 }
 
 func (l *label) getID(key cache.LabelKey) (int, bool) {
 	id, ok := l.labelKeyToID[key]
 	return id, ok
 }
 
-func (l *label) refresh(args ...interface{}) error {
-	// Snapshot existing keys before querying DB, to identify entries added
-	// by encode() during the query window (those must be preserved even if
-	// not yet visible in the DB snapshot).
+func (l *label) MarkRefresh() {
 	l.lock.Lock()
-	preKeys := make(map[cache.LabelKey]struct{}, len(l.labelKeyToID))
-	for k := range l.labelKeyToID {
-		preKeys[k] = struct{}{}
+	defer l.lock.Unlock()
+	l.isRefreshing = true
+	l.pendingKeys = make(map[cache.LabelKey]int)
+}
+
+func (l *label) MarkRefreshDone() {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+	l.isRefreshing = false
+	l.pendingKeys = nil
+}
+
+func (l *label) refresh(args ...interface{}) error {
+	l.MarkRefresh()
+	defer l.MarkRefreshDone()
+
+	var count int64
+	if err := l.org.DB.Model(&metadbmodel.PrometheusLabel{}).Count(&count).Error; err != nil {
+		return err
 	}
-	l.lock.Unlock()
 
 	rows, err := l.org.DB.Model(&metadbmodel.PrometheusLabel{}).Select("id", "name", "value").Rows()
 	if err != nil {
 		return err
 	}
 	defer rows.Close()
 
-	newMap := make(map[cache.LabelKey]int)
+	newMap := make(map[cache.LabelKey]int, count)
 	for rows.Next() {
 		var id int
 		var name, value string
@@ -84,15 +104,12 @@ func (l *label) refresh(args ...interface{}) error {
 	}
 
 	l.lock.Lock()
-	for k, v := range l.labelKeyToID {
-		if _, wasInSnapshot := preKeys[k]; !wasInSnapshot {
-			// Written by encode() after the snapshot; may not be in the DB
-			// snapshot yet, so preserve it to avoid a spurious cache miss.
-			newMap[k] = v
-		}
+	for k, v := range l.pendingKeys {
+		newMap[k] = v
 	}
 	l.labelKeyToID = newMap
 	l.lock.Unlock()
+
 	return nil
 }
 

server/controller/prometheus/encoder/label_value.go

@@ -31,6 +31,9 @@ type labelValue struct {
 	lock         sync.Mutex
 	resourceType string
 	strToID      map[string]int
+
+	isRefreshing bool
+	pendingKeys  map[string]int
 }
 
 func newLabelValue(org *common.ORG) *labelValue {
@@ -41,16 +44,29 @@ func newLabelValue(org *common.ORG) *labelValue {
 	}
 }
 
-func (lv *labelValue) refresh(args ...interface{}) error {
-	// Snapshot existing keys before querying DB, to identify entries added
-	// by encode() during the query window (those must be preserved even if
-	// not yet visible in the DB snapshot).
+func (lv *labelValue) MarkRefresh() {
 	lv.lock.Lock()
-	preKeys := make(map[string]struct{}, len(lv.strToID))
-	for k := range lv.strToID {
-		preKeys[k] = struct{}{}
+	defer lv.lock.Unlock()
+	lv.isRefreshing = true
+	lv.pendingKeys = make(map[string]int)
+}
+
+func (lv *labelValue) MarkRefreshDone() {
+	lv.lock.Lock()
+	defer lv.lock.Unlock()
+	lv.isRefreshing = false
+	lv.pendingKeys = nil
+}
+
+func (lv *labelValue) refresh(args ...interface{}) error {
+	lv.MarkRefresh()
+	defer lv.MarkRefreshDone()
+
+	var count int64
+	if err := lv.org.DB.Model(&metadbmodel.PrometheusLabelValue{}).Count(&count).Error; err != nil {
+		log.Errorf("db query %s failed: %v", lv.resourceType, err, lv.org.LogPrefix)
+		return err
 	}
-	lv.lock.Unlock()
 
 	rows, err := lv.org.DB.Model(&metadbmodel.PrometheusLabelValue{}).Select("id", "value").Rows()
 	if err != nil {
@@ -59,7 +75,7 @@ func (lv *labelValue) refresh(args ...interface{}) error {
 	}
 	defer rows.Close()
 
-	newMap := make(map[string]int)
+	newMap := make(map[string]int, count)
 	for rows.Next() {
 		var id int
 		var value string
@@ -75,15 +91,12 @@ func (lv *labelValue) refresh(args ...interface{}) error {
 	}
 
 	lv.lock.Lock()
-	for k, v := range lv.strToID {
-		if _, wasInSnapshot := preKeys[k]; !wasInSnapshot {
-			// Written by encode() after the snapshot; may not be in the DB
-			// snapshot yet, so preserve it to avoid a spurious cache miss.
-			newMap[k] = v
-		}
+	for k, v := range lv.pendingKeys {
+		newMap[k] = v
 	}
 	lv.strToID = newMap
 	lv.lock.Unlock()
+
 	return nil
 }
 
@@ -124,4 +137,8 @@ func (lv *labelValue) getID(str string) (int, bool) {
 
 func (lv *labelValue) store(item *metadbmodel.PrometheusLabelValue) {
 	lv.strToID[item.Value] = item.ID
+
+	if lv.isRefreshing {
+		lv.pendingKeys[item.Value] = item.ID
+	}
 }