From 7fbc950adc76cbc448d03d300efd634774cf5368 Mon Sep 17 00:00:00 2001
From: Filip Petkovski <filip.petkovsky@gmail.com>
Date: Mon, 28 Aug 2023 12:19:35 +0200
Subject: [PATCH] Fix mutable stringset memory usage (#6669)

This commit fixes the Insert function for the mutable stringset
to only insert unique labels instead of adding every label to the set.

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>
---
 pkg/store/bucket_e2e_test.go |  1 +
 pkg/stringset/set.go         | 18 +++++++++++++++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/pkg/store/bucket_e2e_test.go b/pkg/store/bucket_e2e_test.go
index 85b1d597d8..7262dee155 100644
--- a/pkg/store/bucket_e2e_test.go
+++ b/pkg/store/bucket_e2e_test.go
@@ -799,6 +799,7 @@ func TestBucketStore_LabelNamesSet_e2e(t *testing.T) {
 		for _, n := range []string{"a", "b", "c"} {
 			testutil.Assert(t, filter.Has(n), "expected filter to have %s", n)
 		}
+		testutil.Equals(t, 3, filter.Count())
 	})
 }
 
diff --git a/pkg/stringset/set.go b/pkg/stringset/set.go
index defe699353..080071570f 100644
--- a/pkg/stringset/set.go
+++ b/pkg/stringset/set.go
@@ -10,6 +10,10 @@ import (
 type Set interface {
 	Has(string) bool
 	HasAny([]string) bool
+	// Count returns the number of elements in the set.
+	// A value of -1 indicates infinite size and can be returned by a
+	// set representing all possible string values.
+	Count() int
 }
 
 type fixedSet struct {
@@ -38,6 +42,10 @@ func (f fixedSet) Has(s string) bool {
 	return f.cuckoo.Lookup([]byte(s))
 }
 
+func (f fixedSet) Count() int {
+	return int(f.cuckoo.Count())
+}
+
 type mutableSet struct {
 	cuckoo *cuckoo.ScalableCuckooFilter
 }
@@ -54,7 +62,7 @@ func New() MutableSet {
 }
 
 func (e mutableSet) Insert(s string) {
-	e.cuckoo.Insert([]byte(s))
+	e.cuckoo.InsertUnique([]byte(s))
 }
 
 func (e mutableSet) Has(s string) bool {
@@ -70,6 +78,10 @@ func (e mutableSet) HasAny(strings []string) bool {
 	return false
 }
 
+func (e mutableSet) Count() int {
+	return int(e.cuckoo.Count())
+}
+
 type allStringsSet struct{}
 
 func (e allStringsSet) HasAny(_ []string) bool {
@@ -83,3 +95,7 @@ func AllStrings() *allStringsSet {
 func (e allStringsSet) Has(_ string) bool {
 	return true
 }
+
+func (e allStringsSet) Count() int {
+	return -1
+}