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]
 	}
 	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
+}