From 9f9e408405cddbd429c0c9a6dc1c200431f935cc Mon Sep 17 00:00:00 2001 From: Carrie Edwards Date: Tue, 8 Jul 2025 13:24:12 -0700 Subject: [PATCH] Rename wrapper and use struct instead of map for type/unit tracking Signed-off-by: Carrie Edwards --- cmd/prometheus/main.go | 2 +- promql/engine_test.go | 2 +- ...orage.go => type_unit_mismatch_storage.go} | 76 +++++++++---------- 3 files changed, 39 insertions(+), 41 deletions(-) rename storage/{annotating_storage.go => type_unit_mismatch_storage.go} (51%) diff --git a/cmd/prometheus/main.go b/cmd/prometheus/main.go index cd686adea2..0bf685f6a4 100644 --- a/cmd/prometheus/main.go +++ b/cmd/prometheus/main.go @@ -785,7 +785,7 @@ func main() { var wrappedStorage storage.Storage = localStorage if cfg.scrape.EnableTypeAndUnitLabels { - wrappedStorage = storage.AnnotatingStorage(localStorage) + wrappedStorage = storage.TypeAndUnitMismatchStorage(localStorage) } var ( diff --git a/promql/engine_test.go b/promql/engine_test.go index f04bdfff9a..9bc1bd03bf 100644 --- a/promql/engine_test.go +++ b/promql/engine_test.go @@ -3722,7 +3722,7 @@ func TestTypeUnitMismatchAnnotations(t *testing.T) { store := promqltest.LoadedStorage(t, "load 1m\n"+strings.TrimSpace(testCase.data)) var wrappedStorage storage.Storage = store if testCase.typeAndUnitLabelsEnabled { - wrappedStorage = storage.AnnotatingStorage(store) + wrappedStorage = storage.TypeAndUnitMismatchStorage(store) } t.Cleanup(func() { _ = store.Close() }) diff --git a/storage/annotating_storage.go b/storage/type_unit_mismatch_storage.go similarity index 51% rename from storage/annotating_storage.go rename to storage/type_unit_mismatch_storage.go index 84bb5016ca..8c522b5904 100644 --- a/storage/annotating_storage.go +++ b/storage/type_unit_mismatch_storage.go @@ -21,63 +21,61 @@ import ( "github.com/prometheus/prometheus/util/annotations" ) -// AnnotatingStorage wraps given storage and tracks type and unit +// TypeAndUnitMismatchStorage wraps given storage and tracks type and unit // labels across series. It will produce an info annotation if there // is a mismatch between type or unit in series with the same name. -func AnnotatingStorage(s Storage) Storage { - return &annotatingStorage{ +func TypeAndUnitMismatchStorage(s Storage) Storage { + return &typeAndUnitMismatchStorage{ Storage: s, } } -type annotatingStorage struct { +type typeAndUnitMismatchStorage struct { Storage } -type annotatingQuerier struct { +type typeAndUnitMismatchQuerier struct { Querier } -func (s *annotatingStorage) Querier(mint, maxt int64) (Querier, error) { +func (s *typeAndUnitMismatchStorage) Querier(mint, maxt int64) (Querier, error) { q, err := s.Storage.Querier(mint, maxt) if err != nil { return nil, err } - return &annotatingQuerier{ + return &typeAndUnitMismatchQuerier{ Querier: q, }, nil } -func (q *annotatingQuerier) Select(ctx context.Context, sort bool, hints *SelectHints, matchers ...*labels.Matcher) SeriesSet { +func (q *typeAndUnitMismatchQuerier) Select(ctx context.Context, sort bool, hints *SelectHints, matchers ...*labels.Matcher) SeriesSet { ss := q.Querier.Select(ctx, sort, hints, matchers...) - return AnnotatingSeriesSet(ss) + return TypeAndUnitMismatchSeriesSet(ss) } -type annotatingSeriesSet struct { +type typeAndUnitInfo struct { + typ string + unit string +} + +type typeAndUnitMismatchSeriesSet struct { SeriesSet - seen map[string]struct { - Type string - Unit string - } + prev *typeAndUnitInfo annotations annotations.Annotations } -func AnnotatingSeriesSet(s SeriesSet) SeriesSet { - return &annotatingSeriesSet{ +func TypeAndUnitMismatchSeriesSet(s SeriesSet) SeriesSet { + return &typeAndUnitMismatchSeriesSet{ SeriesSet: s, - seen: make(map[string]struct { - Type string - Unit string - }), } } -func (s *annotatingSeriesSet) At() Series { +func (s *typeAndUnitMismatchSeriesSet) At() Series { return s.SeriesSet.At() } -func (s *annotatingSeriesSet) Next() bool { +func (s *typeAndUnitMismatchSeriesSet) Next() bool { if !s.SeriesSet.Next() { return false } @@ -87,29 +85,29 @@ func (s *annotatingSeriesSet) Next() bool { mType := series.Labels().Get("__type__") mUnit := series.Labels().Get("__unit__") - if prev, ok := s.seen[metric]; ok { - if prev.Type != mType || prev.Unit != mUnit { - if prev.Type == "unknown" || mType == "unknown" { - if prev.Type == "unknown" && mType != "unknown" { - s.seen[metric] = struct { - Type string - Unit string - }{Type: mType, Unit: mUnit} - } - } else { - s.annotations.Add(fmt.Errorf("%w %q", annotations.MismatchedTypeOrUnitInfo, metric)) - } + if s.prev == nil { + s.prev = &typeAndUnitInfo{ + typ: mType, + unit: mUnit, } } else { - s.seen[metric] = struct { - Type string - Unit string - }{Type: mType, Unit: mUnit} + if s.prev.typ != mType || s.prev.unit != mUnit { + // Skip annotation if either type is "unknown" - this allows unknown to coexist with known types + if s.prev.typ == "unknown" || mType == "unknown" { + // Update to the known type if we have one + if s.prev.typ == "unknown" && mType != "unknown" { + s.prev.typ = mType + s.prev.unit = mUnit + } + } else { + s.annotations.Add(fmt.Errorf("%w for metric %q", annotations.MismatchedTypeOrUnitInfo, metric)) + } + } } return true } -func (s *annotatingSeriesSet) Warnings() annotations.Annotations { +func (s *typeAndUnitMismatchSeriesSet) Warnings() annotations.Annotations { got := s.SeriesSet.Warnings() if got == nil { got = make(annotations.Annotations)