util: Pass limit to MergeSlice

Signed-off-by: 🌲 Harry 🌊 John 🏔 <johrry@amazon.com>

pkg/store/bucket.go

@@ -1854,7 +1854,7 @@ func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesReq
 					}
 				})
 
-				result = strutil.MergeSlices(res, extRes)
+				result = strutil.MergeSlices(int(req.Limit), res, extRes)
 			} else {
 				seriesReq := &storepb.SeriesRequest{
 					MinTime:              req.Start,
@@ -1949,10 +1949,7 @@ func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesReq
 		return nil, status.Error(codes.Unknown, errors.Wrap(err, "marshal label names response hints").Error())
 	}
 
-	names := strutil.MergeSlices(sets...)
-	if req.Limit > 0 && len(names) > int(req.Limit) {
-		names = names[:req.Limit]
-	}
+	names := strutil.MergeSlices(int(req.Limit), sets...)
 
 	return &storepb.LabelNamesResponse{
 		Names: names,
@@ -2071,7 +2068,7 @@ 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 = strutil.MergeSlices(int(req.Limit), res, []string{extLabelValue})
 				}
 				result = res
 			} else {
@@ -2169,10 +2166,7 @@ func (s *BucketStore) LabelValues(ctx context.Context, req *storepb.LabelValuesR
 		return nil, status.Error(codes.Unknown, errors.Wrap(err, "marshal label values response hints").Error())
 	}
 
-	vals := strutil.MergeSlices(sets...)
-	if req.Limit > 0 && len(vals) > int(req.Limit) {
-		vals = vals[:req.Limit]
-	}
+	vals := strutil.MergeSlices(int(req.Limit), sets...)
 
 	return &storepb.LabelValuesResponse{
 		Values: vals,

pkg/store/proxy.go

@@ -417,10 +417,7 @@ func (s *ProxyStore) LabelNames(ctx context.Context, originalRequest *storepb.La
 		return nil, err
 	}
 
-	result := strutil.MergeUnsortedSlices(names...)
-	if originalRequest.Limit > 0 && len(result) > int(originalRequest.Limit) {
-		result = result[:originalRequest.Limit]
-	}
+	result := strutil.MergeUnsortedSlices(int(originalRequest.Limit), names...)
 
 	return &storepb.LabelNamesResponse{
 		Names:    result,
@@ -525,10 +522,7 @@ func (s *ProxyStore) LabelValues(ctx context.Context, originalRequest *storepb.L
 		return nil, err
 	}
 
-	vals := strutil.MergeUnsortedSlices(all...)
-	if originalRequest.Limit > 0 && len(vals) > int(originalRequest.Limit) {
-		vals = vals[:originalRequest.Limit]
-	}
+	vals := strutil.MergeUnsortedSlices(int(originalRequest.Limit), all...)
 
 	return &storepb.LabelValuesResponse{
 		Values:   vals,

pkg/strutil/merge.go

@@ -10,33 +10,39 @@ import (
 
 // MergeSlices merges a set of sorted string slices into a single ones
 // while removing all duplicates.
-func MergeSlices(a ...[]string) []string {
+// If limit is set, only the first limit results will be returned. 0 to disable.
+func MergeSlices(limit int, a ...[]string) []string {
 	if len(a) == 0 {
 		return nil
 	}
 	if len(a) == 1 {
-		return a[0]
+		return truncateToLimit(limit, a[0])
 	}
 	l := len(a) / 2
-	return mergeTwoStringSlices(MergeSlices(a[:l]...), MergeSlices(a[l:]...))
+	return mergeTwoStringSlices(limit, MergeSlices(limit, a[:l]...), MergeSlices(limit, a[l:]...))
 }
 
 // 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 {
+// If limit is set, only the first limit results will be returned. 0 to disable.
+func MergeUnsortedSlices(limit int, a ...[]string) []string {
 	for _, s := range a {
 		if !sort.StringsAreSorted(s) {
 			sort.Strings(s)
 		}
 	}
-	return MergeSlices(a...)
+	return MergeSlices(limit, a...)
 }
 
-func mergeTwoStringSlices(a, b []string) []string {
+func mergeTwoStringSlices(limit int, a, b []string) []string {
+	a = truncateToLimit(limit, a)
+	b = truncateToLimit(limit, b)
+
 	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 {
@@ -56,5 +62,13 @@ func mergeTwoStringSlices(a, b []string) []string {
 	// Append all remaining elements.
 	res = append(res, a...)
 	res = append(res, b...)
+	res = truncateToLimit(limit, res)
 	return res
 }
+
+func truncateToLimit(limit int, a []string) []string {
+	if limit > 0 && len(a) > limit {
+		return a[:limit]
+	}
+	return a
+}

pkg/strutil/merge_test.go

@@ -0,0 +1,95 @@
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache License 2.0.
+
+package strutil
+
+import (
+	"testing"
+
+	"github.com/efficientgo/core/testutil"
+)
+
+func TestMergeSlices(t *testing.T) {
+	testCases := map[string]struct {
+		slices   [][]string
+		limit    int
+		expected []string
+	}{
+		"empty slice": {
+			slices: [][]string{
+				{},
+			},
+			expected: []string{},
+		},
+		"single slice with limit": {
+			slices: [][]string{
+				{"a", "b", "c", "d"},
+			},
+			limit:    2,
+			expected: []string{"a", "b"},
+		},
+		"multiple slices with limit": {
+			slices: [][]string{
+				{"a", "b", "d", "f"},
+				{"c", "e", "g"},
+				{"r", "s", "t"},
+			},
+			limit:    4,
+			expected: []string{"a", "b", "c", "d"},
+		},
+		"multiple slices without limit": {
+			slices: [][]string{
+				{"a", "b", "d", "f"},
+				{"c", "e", "g"},
+				{"r", "s"},
+			},
+			expected: []string{"a", "b", "c", "d", "e", "f", "g", "r", "s"},
+		},
+	}
+
+	for tcName, tc := range testCases {
+		t.Run(tcName, func(t *testing.T) {
+			res := MergeSlices(tc.limit, tc.slices...)
+			testutil.Equals(t, tc.expected, res)
+		})
+	}
+}
+
+func TestMergeUnsortedSlices(t *testing.T) {
+	testCases := map[string]struct {
+		slices   [][]string
+		limit    int
+		expected []string
+	}{
+		"empty slice": {
+			slices: [][]string{
+				{},
+			},
+			expected: []string{},
+		},
+		"multiple slices without limit": {
+			slices: [][]string{
+				{"d", "c", "b", "a"},
+				{"f", "g", "c"},
+				{"s", "r", "e"},
+			},
+			expected: []string{"a", "b", "c", "d", "e", "f", "g", "r", "s"},
+		},
+		"multiple slices with limit": {
+			slices: [][]string{
+				{"d", "c", "b", "a"},
+				{"f", "g", "c"},
+				{"s", "r", "e"},
+			},
+			limit:    5,
+			expected: []string{"a", "b", "c", "d", "e"},
+		},
+	}
+
+	for tcName, tc := range testCases {
+		t.Run(tcName, func(t *testing.T) {
+			res := MergeUnsortedSlices(tc.limit, tc.slices...)
+			testutil.Equals(t, tc.expected, res)
+		})
+	}
+}