add context error check on merge of labels and label names

Signed-off-by: Erlan Zholdubai uulu <erlanz@amazon.com>

CHANGELOG.md

@@ -43,6 +43,7 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 - [#7289](https://github.com/thanos-io/thanos/pull/7289) Query Frontend: show warnings from downstream queries.
 - [#7308](https://github.com/thanos-io/thanos/pull/7308) Store: Batch TSDB Infos for blocks.
 - [#7301](https://github.com/thanos-io/thanos/pull/7301) Store Gateway: fix index header reader `PostingsOffsets` returning wrong values.
+- [#7277](https://github.com/thanos-io/thanos/pull/7277) Query: added context cancellation check on MergeSlices on get label and label values
 
 ### Added
 

pkg/store/bucket.go

@@ -1816,7 +1816,9 @@ func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesReq
 					}
 				})
 
-				result = strutil.MergeSlices(res, extRes)
+				if result, err = strutil.MergeSlices(ctx, res, extRes); err != nil {
+					return err
+				}
 			} else {
 				seriesReq := &storepb.SeriesRequest{
 					MinTime:              req.Start,
@@ -1910,9 +1912,13 @@ func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesReq
 	if err != nil {
 		return nil, status.Error(codes.Unknown, errors.Wrap(err, "marshal label names response hints").Error())
 	}
+	names, err := strutil.MergeSlices(ctx, sets...)
+	if err != nil {
+		return nil, err
+	}
 
 	return &storepb.LabelNamesResponse{
-		Names: strutil.MergeSlices(sets...),
+		Names: names,
 		Hints: anyHints,
 	}, nil
 }
@@ -2025,7 +2031,10 @@ func (s *BucketStore) LabelValues(ctx context.Context, req *storepb.LabelValuesR
 
 				// Add the external label value as well.
 				if extLabelValue := b.extLset.Get(req.Label); extLabelValue != "" {
-					res = strutil.MergeSlices(res, []string{extLabelValue})
+					res, err = strutil.MergeSlices(ctx, res, []string{extLabelValue})
+					if err != nil {
+						return err
+					}
 				}
 				result = res
 			} else {
@@ -2122,9 +2131,13 @@ func (s *BucketStore) LabelValues(ctx context.Context, req *storepb.LabelValuesR
 	if err != nil {
 		return nil, status.Error(codes.Unknown, errors.Wrap(err, "marshal label values response hints").Error())
 	}
+	values, err := strutil.MergeSlices(ctx, sets...)
+	if err != nil {
+		return nil, err
+	}
 
 	return &storepb.LabelValuesResponse{
-		Values: strutil.MergeSlices(sets...),
+		Values: values,
 		Hints:  anyHints,
 	}, nil
 }

pkg/store/proxy.go

@@ -559,8 +559,12 @@ func (s *ProxyStore) LabelNames(ctx context.Context, r *storepb.LabelNamesReques
 	}
 
 	level.Debug(s.logger).Log("msg", strings.Join(storeDebugMsgs, ";"))
+	res, err := strutil.MergeUnsortedSlices(ctx, names...)
+	if err != nil {
+		return nil, err
+	}
 	return &storepb.LabelNamesResponse{
-		Names:    strutil.MergeUnsortedSlices(names...),
+		Names:    res,
 		Warnings: warnings,
 	}, nil
 }
@@ -663,8 +667,12 @@ func (s *ProxyStore) LabelValues(ctx context.Context, r *storepb.LabelValuesRequ
 	}
 
 	level.Debug(s.logger).Log("msg", strings.Join(storeDebugMsgs, ";"))
+	values, err := strutil.MergeUnsortedSlices(ctx, all...)
+	if err != nil {
+		return nil, err
+	}
 	return &storepb.LabelValuesResponse{
-		Values:   strutil.MergeUnsortedSlices(all...),
+		Values:   values,
 		Warnings: warnings,
 	}, nil
 }

pkg/strutil/merge.go

@@ -4,42 +4,54 @@
 package strutil
 
 import (
+	"context"
 	"sort"
 	"strings"
 )
 
 // MergeSlices merges a set of sorted string slices into a single ones
 // while removing all duplicates.
-func MergeSlices(a ...[]string) []string {
+func MergeSlices(ctx context.Context, a ...[]string) ([]string, error) {
 	if len(a) == 0 {
-		return nil
+		return nil, nil
 	}
 	if len(a) == 1 {
-		return a[0]
+		return a[0], nil
 	}
 	l := len(a) / 2
-	return mergeTwoStringSlices(MergeSlices(a[:l]...), MergeSlices(a[l:]...))
+	a1, err := MergeSlices(ctx, a[:l]...)
+	if err != nil {
+		return nil, err
+	}
+	a2, err := MergeSlices(ctx, a[l:]...)
+	if err != nil {
+		return nil, err
+	}
+	return mergeTwoStringSlices(ctx, a1, a2)
 }
 
 // MergeUnsortedSlices behaves like StringSlices but input slices are validated
 // for sortedness and are sorted if they are not ordered yet.
-func MergeUnsortedSlices(a ...[]string) []string {
+func MergeUnsortedSlices(ctx context.Context, a ...[]string) ([]string, error) {
 	for _, s := range a {
 		if !sort.StringsAreSorted(s) {
 			sort.Strings(s)
 		}
 	}
-	return MergeSlices(a...)
+	return MergeSlices(ctx, a...)
 }
 
-func mergeTwoStringSlices(a, b []string) []string {
+func mergeTwoStringSlices(ctx context.Context, a, b []string) ([]string, error) {
 	maxl := len(a)
 	if len(b) > len(a) {
 		maxl = len(b)
 	}
 	res := make([]string, 0, maxl*10/9)
 
 	for len(a) > 0 && len(b) > 0 {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 		d := strings.Compare(a[0], b[0])
 
 		if d == 0 {
@@ -56,5 +68,5 @@ func mergeTwoStringSlices(a, b []string) []string {
 	// Append all remaining elements.
 	res = append(res, a...)
 	res = append(res, b...)
-	return res
+	return res, nil
 }

pkg/strutil/merge_test.go

@@ -0,0 +1,29 @@
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache License 2.0.
+
+package strutil
+
+import (
+	"context"
+	"testing"
+
+	"github.com/efficientgo/core/testutil"
+)
+
+func TestMergeSlices(t *testing.T) {
+	var (
+		a1  = []string{"Alice", "Benjamin", "Charlotte", "David", "Eleanor", "Frederick", "Grace", "Henry", "Isla", "James"}
+		a2  = []string{"Ava", "Emma", "Ethan", "Isabella", "Jacob", "Liam", "Mason", "Noah", "Olivia", "Sophia"}
+		a3  = []string{"Elizabeth", "George", "Katherine", "Margaret", "Michael", "Robert", "Sarah", "Thomas", "Victoria", "William"}
+		a4  = []string{"Amelie", "Carlos", "Dmitri", "Elena", "Fiona", "Gustavo", "Hiro", "Ingrid", "Jorge", "Klara"}
+		res = []string{"Alice", "Amelie", "Ava", "Benjamin", "Carlos", "Charlotte", "David", "Dmitri", "Eleanor", "Elena", "Elizabeth", "Emma", "Ethan", "Fiona", "Frederick", "George", "Grace", "Gustavo", "Henry", "Hiro", "Ingrid", "Isabella", "Isla", "Jacob", "James", "Jorge", "Katherine", "Klara", "Liam", "Margaret", "Mason", "Michael", "Noah", "Olivia", "Robert", "Sarah", "Sophia", "Thomas", "Victoria", "William"}
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	result, err := MergeSlices(ctx, a1, a2, a3, a4)
+
+	testutil.Ok(t, err)
+	testutil.Equals(t, result, res)
+}