From b52862badfad26e0f4006f761b7807913714a37a Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Sun, 29 Sep 2024 10:50:25 +0200
Subject: [PATCH] lib/logstorage: return the expected `hits` results from
 `uniq` pipe when the number of unique values reaches the specified limit

Previously `uniq` pipe could return zero `hits` if the number of found unique values equals the specified limit.
This wasn't expected in most cases.
---
 lib/logstorage/pipe_uniq.go           |  4 ++--
 lib/logstorage/pipe_uniq_test.go      | 32 +++++++++++++++++++++++++--
 lib/logstorage/stats_count_uniq.go    |  2 +-
 lib/logstorage/stats_uniq_values.go   |  2 +-
 lib/logstorage/stats_values.go        |  2 +-
 lib/logstorage/storage_search_test.go |  6 ++---
 6 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/lib/logstorage/pipe_uniq.go b/lib/logstorage/pipe_uniq.go
index 2292fe127..c33cfc7f0 100644
--- a/lib/logstorage/pipe_uniq.go
+++ b/lib/logstorage/pipe_uniq.go
@@ -135,7 +135,7 @@ type pipeUniqProcessorShardNopad struct {
 //
 // It returns false if the block cannot be written because of the exceeded limit.
 func (shard *pipeUniqProcessorShard) writeBlock(br *blockResult) bool {
-	if limit := shard.pu.limit; limit > 0 && uint64(len(shard.m)) >= limit {
+	if limit := shard.pu.limit; limit > 0 && uint64(len(shard.m)) > limit {
 		return false
 	}
 
@@ -301,7 +301,7 @@ func (pup *pipeUniqProcessor) flush() error {
 
 	// There is little sense in returning partial hits when the limit on the number of unique entries is reached.
 	// It is better from UX experience is to return zero hits instead.
-	resetHits := pup.pu.limit > 0 && uint64(len(m)) >= pup.pu.limit
+	resetHits := pup.pu.limit > 0 && uint64(len(m)) > pup.pu.limit
 
 	// write result
 	wctx := &pipeUniqWriteContext{
diff --git a/lib/logstorage/pipe_uniq_test.go b/lib/logstorage/pipe_uniq_test.go
index 3c1cb7372..f3f61effd 100644
--- a/lib/logstorage/pipe_uniq_test.go
+++ b/lib/logstorage/pipe_uniq_test.go
@@ -114,13 +114,41 @@ func TestPipeUniq(t *testing.T) {
 		{
 			{"a", "2"},
 			{"b", "3"},
-			{"hits", "0"},
+			{"hits", "2"},
 		},
 		{
 			{"a", `2`},
 			{"b", `54`},
 			{"c", "d"},
-			{"hits", "0"},
+			{"hits", "1"},
+		},
+	})
+
+	f("uniq hits limit 3", [][]Field{
+		{
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"a", "2"},
+			{"b", "3"},
+		},
+		{
+			{"a", `2`},
+			{"b", `54`},
+			{"c", "d"},
+		},
+	}, [][]Field{
+		{
+			{"a", "2"},
+			{"b", "3"},
+			{"hits", "2"},
+		},
+		{
+			{"a", `2`},
+			{"b", `54`},
+			{"c", "d"},
+			{"hits", "1"},
 		},
 	})
 
diff --git a/lib/logstorage/stats_count_uniq.go b/lib/logstorage/stats_count_uniq.go
index b7aa87e4e..42a71a03f 100644
--- a/lib/logstorage/stats_count_uniq.go
+++ b/lib/logstorage/stats_count_uniq.go
@@ -355,7 +355,7 @@ func (sup *statsCountUniqProcessor) updateState(v []byte) int {
 
 func (sup *statsCountUniqProcessor) limitReached() bool {
 	limit := sup.su.limit
-	return limit > 0 && uint64(len(sup.m)) >= limit
+	return limit > 0 && uint64(len(sup.m)) > limit
 }
 
 func parseStatsCountUniq(lex *lexer) (*statsCountUniq, error) {
diff --git a/lib/logstorage/stats_uniq_values.go b/lib/logstorage/stats_uniq_values.go
index a266b2c02..25417efc9 100644
--- a/lib/logstorage/stats_uniq_values.go
+++ b/lib/logstorage/stats_uniq_values.go
@@ -215,7 +215,7 @@ func (sup *statsUniqValuesProcessor) updateState(v string) int {
 
 func (sup *statsUniqValuesProcessor) limitReached() bool {
 	limit := sup.su.limit
-	return limit > 0 && uint64(len(sup.m)) >= limit
+	return limit > 0 && uint64(len(sup.m)) > limit
 }
 
 func marshalJSONArray(items []string) string {
diff --git a/lib/logstorage/stats_values.go b/lib/logstorage/stats_values.go
index aab5475d5..eb9c100de 100644
--- a/lib/logstorage/stats_values.go
+++ b/lib/logstorage/stats_values.go
@@ -184,7 +184,7 @@ func (svp *statsValuesProcessor) finalizeStats() string {
 
 func (svp *statsValuesProcessor) limitReached() bool {
 	limit := svp.sv.limit
-	return limit > 0 && uint64(len(svp.values)) >= limit
+	return limit > 0 && uint64(len(svp.values)) > limit
 }
 
 func parseStatsValues(lex *lexer) (*statsValues, error) {
diff --git a/lib/logstorage/storage_search_test.go b/lib/logstorage/storage_search_test.go
index 4fd7ed0ac..35530309e 100644
--- a/lib/logstorage/storage_search_test.go
+++ b/lib/logstorage/storage_search_test.go
@@ -384,9 +384,9 @@ func TestStorageRunQuery(t *testing.T) {
 		}
 
 		resultsExpected := []ValueWithHits{
-			{`{instance="host-0:234",job="foobar"}`, 0},
-			{`{instance="host-1:234",job="foobar"}`, 0},
-			{`{instance="host-2:234",job="foobar"}`, 0},
+			{`{instance="host-0:234",job="foobar"}`, 385},
+			{`{instance="host-1:234",job="foobar"}`, 385},
+			{`{instance="host-2:234",job="foobar"}`, 385},
 		}
 		if !reflect.DeepEqual(results, resultsExpected) {
 			t.Fatalf("unexpected result; got\n%v\nwant\n%v", results, resultsExpected)