From 5281a6bc1b8e36b0c778f6b2a41119a18c185aa5 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Tue, 2 Apr 2024 15:08:58 +0100
Subject: [PATCH] TSDB: rebuild labels symbol-table on each compaction

Log begin/end for timing, plus some stats.

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>
---
 tsdb/db.go                |  3 +++
 tsdb/head_dedupelabels.go | 40 +++++++++++++++++++++++++++++++++++++++
 tsdb/head_other.go        |  7 +++++++
 3 files changed, 50 insertions(+)

diff --git a/tsdb/db.go b/tsdb/db.go
index b2175d4758..090d6fcf0c 100644
--- a/tsdb/db.go
+++ b/tsdb/db.go
@@ -1407,6 +1407,9 @@ func (db *DB) compactHead(head *RangeHead) error {
 	if err = db.head.truncateMemory(head.BlockMaxTime()); err != nil {
 		return fmt.Errorf("head memory truncate: %w", err)
 	}
+
+	db.head.RebuildSymbolTable(db.logger)
+
 	return nil
 }
 
diff --git a/tsdb/head_dedupelabels.go b/tsdb/head_dedupelabels.go
index 203f92e6a8..aaab7c25be 100644
--- a/tsdb/head_dedupelabels.go
+++ b/tsdb/head_dedupelabels.go
@@ -16,6 +16,9 @@
 package tsdb
 
 import (
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+
 	"github.com/prometheus/prometheus/model/labels"
 )
 
@@ -25,3 +28,40 @@ func (s *memSeries) labels() labels.Labels {
 	defer s.Unlock()
 	return s.lset
 }
+
+// RebuildSymbolTable goes through all the series in h, build a SymbolTable with all names and values,
+// replace each series' Labels with one using that SymbolTable.
+func (h *Head) RebuildSymbolTable(logger log.Logger) *labels.SymbolTable {
+	level.Info(logger).Log("msg", "RebuildSymbolTable starting")
+	st := labels.NewSymbolTable()
+	builder := labels.NewScratchBuilderWithSymbolTable(st, 0)
+	rebuildLabels := func(lbls labels.Labels) labels.Labels {
+		builder.Reset()
+		lbls.Range(func(l labels.Label) {
+			builder.Add(l.Name, l.Value)
+		})
+		return builder.Labels()
+	}
+
+	for i := 0; i < h.series.size; i++ {
+		h.series.locks[i].Lock()
+
+		for _, s := range h.series.hashes[i].unique {
+			s.Lock()
+			s.lset = rebuildLabels(s.lset)
+			s.Unlock()
+		}
+
+		for _, all := range h.series.hashes[i].conflicts {
+			for _, s := range all {
+				s.Lock()
+				s.lset = rebuildLabels(s.lset)
+				s.Unlock()
+			}
+		}
+
+		h.series.locks[i].Unlock()
+	}
+	level.Info(logger).Log("msg", "RebuildSymbolTable finished", "size", st.Len())
+	return st
+}
diff --git a/tsdb/head_other.go b/tsdb/head_other.go
index 9306913d8f..eb1b93a3e5 100644
--- a/tsdb/head_other.go
+++ b/tsdb/head_other.go
@@ -16,6 +16,8 @@
 package tsdb
 
 import (
+	"github.com/go-kit/log"
+
 	"github.com/prometheus/prometheus/model/labels"
 )
 
@@ -23,3 +25,8 @@ import (
 func (s *memSeries) labels() labels.Labels {
 	return s.lset
 }
+
+// No-op when not using dedupelabels.
+func (h *Head) RebuildSymbolTable(logger log.Logger) *labels.SymbolTable {
+	return nil
+}