diff --git a/util/iter/iter.go b/util/iter/iter.go index f5596a6..e9cafc8 100644 --- a/util/iter/iter.go +++ b/util/iter/iter.go @@ -252,6 +252,17 @@ func (iter Iter[T]) CollectSlice() []T { return output } +func (iter Iter[T]) Count() uint64 { + output := uint64(0) + + iter(func(_ T) Flow { + output++ + return Continue + }) + + return output +} + func FromMap[K comparable, V any](m map[K]V) Iter[Pair[K, V]] { return func(yield Yield[Pair[K, V]]) Flow { for k, v := range m { diff --git a/util/labelindex/selectors_bench_test.go b/util/labelindex/selectors_bench_test.go index 6b836b6..1ac93c6 100644 --- a/util/labelindex/selectors_bench_test.go +++ b/util/labelindex/selectors_bench_test.go @@ -16,71 +16,100 @@ package labelindex_test import ( "fmt" + "math/rand/v2" "testing" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "github.com/kubewharf/podseidon/util/iter" "github.com/kubewharf/podseidon/util/labelindex" ) -const numSelectors = 10000 - func BenchmarkSelectorsInsert(b *testing.B) { index := labelindex.NewSelectors[string]() + rng := rand.New(rand.NewChaCha8([32]byte{})) + + items := make([]iter.Pair[string, metav1.LabelSelector], b.N/2*2) + appId := 0 + subAppId := 0 + + for itemId := range b.N / 2 { + if rng.IntN(2) == 0 { + appId++ + subAppId = 0 + } else { + subAppId++ + } - for appId := range b.N { - err := index.Track(fmt.Sprint(appId), metav1.LabelSelector{ - MatchLabels: map[string]string{ - "aaa": "aaa", - "appId": fmt.Sprint(-appId), - "zzz": "zzz", - }, - }) - require.NoError(b, err) + items[itemId] = iter.NewPair(fmt.Sprintf("%d-%d", appId, subAppId), metav1.LabelSelector{MatchLabels: map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(appId), + "subApp": fmt.Sprint(subAppId), + "zzz": "zzz", + }}) + items[itemId+b.N/2] = iter.NewPair(fmt.Sprintf("%d-%d", appId, subAppId), metav1.LabelSelector{MatchLabels: map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(-appId), + "subApp": fmt.Sprint(subAppId), + "zzz": "zzz", + }}) } - for appId := range b.N { - err := index.Track(fmt.Sprint(appId), metav1.LabelSelector{ - MatchLabels: map[string]string{ - "aaa": "aaa", - "appId": fmt.Sprint(appId), - "zzz": "zzz", - }, - }) - require.NoError(b, err) + b.ResetTimer() + + for _, item := range items { + err := index.Track(item.Left, item.Right) + if err != nil { + b.Log(err) + b.FailNow() + } } } -func BenchmarkSelectorsQueryBroadExact(b *testing.B) { +func BenchmarkSelectorsQueryBroad(b *testing.B) { // Generate data index := labelindex.NewSelectors[string]() - for appId := range numSelectors { - err := index.Track(fmt.Sprint(appId), metav1.LabelSelector{ - MatchLabels: map[string]string{ - "aaa": "aaa", - "appId": fmt.Sprint(appId), - "zzz": "zzz", - }, - }) - require.NoError(b, err) - } + rng := rand.New(rand.NewChaCha8([32]byte{})) - b.ResetTimer() + appId := 0 + subAppId := 0 + + for appId < b.N { + err := index.Track(fmt.Sprintf("%d-%d", appId, subAppId), metav1.LabelSelector{MatchLabels: map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(appId), + "subApp": fmt.Sprint(subAppId), + "zzz": "zzz", + }}) + if err != nil { + b.Log(err) + b.FailNow() + } - for appId := range b.N { - appIdString := fmt.Sprint(appId % numSelectors) + if rng.IntN(2) == 0 { + appId++ + subAppId = 0 + } else { + subAppId++ + } + } - resultIter, _ := index.Query(map[string]string{ - "aaa": "aaa", - "appId": appIdString, - "zzz": "zzz", - }) - result := resultIter.CollectSlice() + queries := iter.Map(iter.Range(0, b.N), func(appId int) map[string]string { + return map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(appId), + "subApp": "0", + "zzz": "zzz", + } + }).CollectSlice() + + b.ResetTimer() - assert.Len(b, result, 1) - assert.Equal(b, []string{appIdString}, result) + for _, query := range queries { + resultIter, _ := index.Query(query) + result := resultIter.Count() + assert.Equal(b, uint64(1), result) } } diff --git a/util/labelindex/sets_bench_test.go b/util/labelindex/sets_bench_test.go index e8af308..8079586 100644 --- a/util/labelindex/sets_bench_test.go +++ b/util/labelindex/sets_bench_test.go @@ -16,66 +16,98 @@ package labelindex_test import ( "fmt" + "math/rand/v2" "testing" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "github.com/kubewharf/podseidon/util/iter" "github.com/kubewharf/podseidon/util/labelindex" ) -const numSets = 1000000 - func BenchmarkSetsInsert(b *testing.B) { index := labelindex.NewSets[string]() + rng := rand.New(rand.NewChaCha8([32]byte{})) + + items := make([]iter.Pair[string, map[string]string], b.N/2*2) + appId := 0 + subAppId := 0 - for appId := range b.N { - index.Track(fmt.Sprint(appId), map[string]string{ - "aaa": "aaa", - "appId": fmt.Sprint(-appId), - "zzz": "zzz", + for itemId := range b.N / 2 { + if rng.IntN(2) == 0 { + appId++ + subAppId = 0 + } else { + subAppId++ + } + + items[itemId] = iter.NewPair(fmt.Sprintf("%d-%d", appId, subAppId), map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(appId), + "subApp": fmt.Sprint(subAppId), + "zzz": "zzz", + }) + items[itemId+b.N/2] = iter.NewPair(fmt.Sprintf("%d-%d", appId, subAppId), map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(-appId), + "subApp": fmt.Sprint(subAppId), + "zzz": "zzz", }) } - for appId := range b.N { - index.Track(fmt.Sprint(appId), map[string]string{ - "aaa": "aaa", - "appId": fmt.Sprint(appId), - "zzz": "zzz", - }) + b.ResetTimer() + + for _, item := range items { + index.Track(item.Left, item.Right) } } -func BenchmarkSetsQueryBroadExact(b *testing.B) { +func BenchmarkSetsQueryBroad(b *testing.B) { // Generate data index := labelindex.NewSets[string]() + rng := rand.New(rand.NewChaCha8([32]byte{})) - for appId := range numSets { - index.Track(fmt.Sprint(appId), map[string]string{ - "aaa": "aaa", - "appId": fmt.Sprint(appId), - "zzz": "zzz", - }) - } + appId := 0 + subAppId := 0 - b.ResetTimer() + for appId < b.N { + if rng.IntN(2) == 0 { + appId++ + subAppId = 0 + } else { + subAppId++ + } - for appId := range b.N { - appIdString := fmt.Sprint(appId % numSets) + index.Track(fmt.Sprintf("%d-%d", appId, subAppId), map[string]string{ + "aaa": "aaa", + "appId": fmt.Sprint(appId), + "subApp": fmt.Sprint(subAppId), + "zzz": "zzz", + }) + } - resultIter, err := index.Query(metav1.LabelSelector{ + queries := iter.Map(iter.Range(0, b.N), func(appId int) metav1.LabelSelector { + return metav1.LabelSelector{ MatchLabels: map[string]string{ "aaa": "aaa", - "appId": appIdString, + "appId": fmt.Sprint(appId), "zzz": "zzz", }, - }) + } + }).CollectSlice() + + b.ResetTimer() + + for _, query := range queries { + resultIter, err := index.Query(query) require.NoError(b, err) - result := resultIter.CollectSlice() + resultCount := resultIter.Count() - assert.Len(b, result, 1) - assert.Equal(b, []string{appIdString}, result) + if resultCount == 0 { + b.Fail() + b.Logf("query %v yields no results", query) + } } }