Unified Storage/Search: Add max count config for indexing (#107255)
* Add max count config for indexing * Build empty index when max count is exceeded * Address linting * Refactor buildIndexes * Add test for max count threshold * Update test doc comments * Refactor TestBuildIndexes_MaxCountThreshold to not use mock framework * Rename mocks used in TestBuildIndexes_MaxCountThreshold * Refactor mockResourceIndex * Test setting of indexing threshold configs * Tweak comments, log * Fix logging in buildEmptyIndex * Export and reuse TestDocumentBuilderSupplier * Reuse MockResourceIndex
This commit is contained in:
Arati R.
@@ -2,9 +2,12 @@ package resource
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/grafana/authlib/types"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.opentelemetry.io/otel/trace/noop"
|
||||
|
||||
"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
|
||||
)
|
||||
@@ -54,3 +57,217 @@ func (m *MockDocumentBuilder) BuildDocument(ctx context.Context, key *resourcepb
|
||||
}
|
||||
return args.Get(0).(*IndexableDocument), nil
|
||||
}
|
||||
|
||||
// mockStorageBackend implements StorageBackend for testing
|
||||
type mockStorageBackend struct {
|
||||
resourceStats []ResourceStats
|
||||
}
|
||||
|
||||
func (m *mockStorageBackend) GetResourceStats(ctx context.Context, namespace string, minCount int) ([]ResourceStats, error) {
|
||||
var result []ResourceStats
|
||||
for _, stat := range m.resourceStats {
|
||||
// Apply the minCount filter like the real implementation does
|
||||
if stat.Count > int64(minCount) {
|
||||
result = append(result, stat)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (m *mockStorageBackend) WriteEvent(ctx context.Context, event WriteEvent) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (m *mockStorageBackend) ReadResource(ctx context.Context, req *resourcepb.ReadRequest) *BackendReadResponse {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockStorageBackend) WatchWriteEvents(ctx context.Context) (<-chan *WrittenEvent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockStorageBackend) ListIterator(ctx context.Context, req *resourcepb.ListRequest, callback func(ListIterator) error) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (m *mockStorageBackend) ListHistory(ctx context.Context, req *resourcepb.ListRequest, callback func(ListIterator) error) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// mockSearchBackend implements SearchBackend for testing with tracking capabilities
|
||||
type mockSearchBackend struct {
|
||||
buildIndexCalls []buildIndexCall
|
||||
buildEmptyIndexCalls []buildEmptyIndexCall
|
||||
}
|
||||
|
||||
type buildIndexCall struct {
|
||||
key NamespacedResource
|
||||
size int64
|
||||
resourceVersion int64
|
||||
fields SearchableDocumentFields
|
||||
}
|
||||
|
||||
type buildEmptyIndexCall struct {
|
||||
key NamespacedResource
|
||||
size int64 // should be 0 for empty indexes
|
||||
resourceVersion int64
|
||||
fields SearchableDocumentFields
|
||||
}
|
||||
|
||||
func (m *mockSearchBackend) GetIndex(ctx context.Context, key NamespacedResource) (ResourceIndex, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockSearchBackend) BuildIndex(ctx context.Context, key NamespacedResource, size int64, resourceVersion int64, fields SearchableDocumentFields, builder func(index ResourceIndex) (int64, error)) (ResourceIndex, error) {
|
||||
index := &MockResourceIndex{}
|
||||
index.On("BulkIndex", mock.Anything).Return(nil).Maybe()
|
||||
index.On("DocCount", mock.Anything, mock.Anything).Return(int64(0), nil).Maybe()
|
||||
|
||||
// Call the builder function (required by the contract)
|
||||
_, err := builder(index)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Determine if this is an empty index based on size
|
||||
// Empty indexes are characterized by size == 0
|
||||
if size == 0 {
|
||||
// This is an empty index (buildEmptyIndex was called)
|
||||
m.buildEmptyIndexCalls = append(m.buildEmptyIndexCalls, buildEmptyIndexCall{
|
||||
key: key,
|
||||
size: size,
|
||||
resourceVersion: resourceVersion,
|
||||
fields: fields,
|
||||
})
|
||||
} else {
|
||||
// This is a normal index (build was called)
|
||||
m.buildIndexCalls = append(m.buildIndexCalls, buildIndexCall{
|
||||
key: key,
|
||||
size: size,
|
||||
resourceVersion: resourceVersion,
|
||||
fields: fields,
|
||||
})
|
||||
}
|
||||
|
||||
return index, nil
|
||||
}
|
||||
|
||||
func (m *mockSearchBackend) TotalDocs() int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
func TestBuildIndexes_MaxCountThreshold(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initMaxSize int
|
||||
resourceStats []ResourceStats
|
||||
expectedNormalBuilds []string // expected NamespacedResource strings that should be built normally
|
||||
expectedEmptyBuilds []string // expected NamespacedResource strings that should be built as empty
|
||||
}{
|
||||
{
|
||||
name: "max count disabled (0) - all resources built normally",
|
||||
initMaxSize: 0,
|
||||
resourceStats: []ResourceStats{
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group1", Resource: "resource1"}, Count: 50},
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group1", Resource: "resource2"}, Count: 150},
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group2", Resource: "resource1"}, Count: 250},
|
||||
},
|
||||
expectedNormalBuilds: []string{
|
||||
"ns1/group1/resource1",
|
||||
"ns1/group1/resource2",
|
||||
"ns1/group2/resource1",
|
||||
},
|
||||
expectedEmptyBuilds: []string{},
|
||||
},
|
||||
{
|
||||
name: "max count 100 - resources above threshold get empty indexes",
|
||||
initMaxSize: 100,
|
||||
resourceStats: []ResourceStats{
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group1", Resource: "resource1"}, Count: 50}, // normal build
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group1", Resource: "resource2"}, Count: 150}, // empty build
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group2", Resource: "resource1"}, Count: 250}, // empty build
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group2", Resource: "resource2"}, Count: 80}, // normal build
|
||||
},
|
||||
expectedNormalBuilds: []string{
|
||||
"ns1/group1/resource1",
|
||||
"ns1/group2/resource2",
|
||||
},
|
||||
expectedEmptyBuilds: []string{
|
||||
"ns1/group1/resource2",
|
||||
"ns1/group2/resource1",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "max count 300 - no resources exceed threshold",
|
||||
initMaxSize: 300,
|
||||
resourceStats: []ResourceStats{
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group1", Resource: "resource1"}, Count: 50}, // normal build
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group1", Resource: "resource2"}, Count: 150}, // normal build
|
||||
{NamespacedResource: NamespacedResource{Namespace: "ns1", Group: "group2", Resource: "resource1"}, Count: 250}, // normal build
|
||||
},
|
||||
expectedNormalBuilds: []string{
|
||||
"ns1/group1/resource1",
|
||||
"ns1/group1/resource2",
|
||||
"ns1/group2/resource1",
|
||||
},
|
||||
expectedEmptyBuilds: []string{},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Setup mock implementations
|
||||
storage := &mockStorageBackend{
|
||||
resourceStats: tt.resourceStats,
|
||||
}
|
||||
search := &mockSearchBackend{
|
||||
buildIndexCalls: []buildIndexCall{},
|
||||
buildEmptyIndexCalls: []buildEmptyIndexCall{},
|
||||
}
|
||||
supplier := &TestDocumentBuilderSupplier{
|
||||
GroupsResources: map[string]string{
|
||||
"group1": "resource1",
|
||||
"group2": "resource2",
|
||||
},
|
||||
}
|
||||
|
||||
// Create search support with the specified initMaxSize
|
||||
opts := SearchOptions{
|
||||
Backend: search,
|
||||
Resources: supplier,
|
||||
WorkerThreads: 1,
|
||||
InitMinCount: 1, // set min count to default for this test
|
||||
InitMaxCount: tt.initMaxSize,
|
||||
}
|
||||
|
||||
support, err := newSearchSupport(opts, storage, nil, nil, noop.NewTracerProvider().Tracer("test"), nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, support)
|
||||
|
||||
// Call buildIndexes
|
||||
ctx := context.Background()
|
||||
indexesBuilt, err := support.buildIndexes(ctx, false)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify the correct number of indexes were built (normal + empty)
|
||||
expectedTotal := len(tt.expectedNormalBuilds) + len(tt.expectedEmptyBuilds)
|
||||
require.Equal(t, expectedTotal, indexesBuilt)
|
||||
|
||||
// Verify the correct resources were built normally
|
||||
actualNormalBuilds := make([]string, len(search.buildIndexCalls))
|
||||
for i, call := range search.buildIndexCalls {
|
||||
actualNormalBuilds[i] = call.key.String()
|
||||
}
|
||||
require.ElementsMatch(t, tt.expectedNormalBuilds, actualNormalBuilds)
|
||||
|
||||
// Verify the correct resources were built as empty indexes
|
||||
actualEmptyBuilds := make([]string, len(search.buildEmptyIndexCalls))
|
||||
for i, call := range search.buildEmptyIndexCalls {
|
||||
actualEmptyBuilds[i] = call.key.String()
|
||||
// Verify that empty indexes are built with size 0
|
||||
require.Equal(t, int64(0), call.size, "Empty index should be built with size 0")
|
||||
}
|
||||
require.ElementsMatch(t, tt.expectedEmptyBuilds, actualEmptyBuilds)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user