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

pkg/setting/setting.go

@@ -563,6 +563,7 @@ type Cfg struct {
 	IndexMaxBatchSize                          int
 	IndexFileThreshold                         int
 	IndexMinCount                              int
+	IndexMaxCount                              int
 	IndexRebuildInterval                       time.Duration
 	IndexCacheTTL                              time.Duration
 	EnableSharding                             bool

pkg/setting/setting_unified_storage.go

@@ -64,6 +64,7 @@ func (cfg *Cfg) setUnifiedStorageConfig() {
 	cfg.InstanceID = section.Key("instance_id").String()
 	cfg.IndexFileThreshold = section.Key("index_file_threshold").MustInt(10)
 	cfg.IndexMinCount = section.Key("index_min_count").MustInt(1)
+	cfg.IndexMaxCount = section.Key("index_max_count").MustInt(0)
 	// default to 24 hours because usage insights summarizes the data every 24 hours
 	cfg.IndexRebuildInterval = section.Key("index_rebuild_interval").MustDuration(24 * time.Hour)
 	cfg.IndexCacheTTL = section.Key("index_cache_ttl").MustDuration(10 * time.Minute)

pkg/setting/setting_unified_storage_test.go

@@ -28,6 +28,16 @@ func TestCfg_setUnifiedStorageConfig(t *testing.T) {
 		_, err = s.NewKey("dataSyncerInterval", "10m")
 		assert.NoError(t, err)
 
+		// Add unified_storage section for index settings
+		unifiedStorageSection, err := cfg.Raw.NewSection("unified_storage")
+		assert.NoError(t, err)
+
+		_, err = unifiedStorageSection.NewKey("index_min_count", "5")
+		assert.NoError(t, err)
+
+		_, err = unifiedStorageSection.NewKey("index_max_count", "1000")
+		assert.NoError(t, err)
+
 		cfg.setUnifiedStorageConfig()
 
 		value, exists := cfg.UnifiedStorage["playlists.playlist.grafana.app"]
@@ -39,5 +49,22 @@ func TestCfg_setUnifiedStorageConfig(t *testing.T) {
 			DataSyncerRecordsLimit:               1001,
 			DataSyncerInterval:                   time.Minute * 10,
 		})
+
+		// Test that index settings are correctly parsed
+		assert.Equal(t, 5, cfg.IndexMinCount)
+		assert.Equal(t, 1000, cfg.IndexMaxCount)
+	})
+
+	t.Run("read unified_storage configs with defaults", func(t *testing.T) {
+		cfg := NewCfg()
+		err := cfg.Load(CommandLineArgs{HomePath: "../../", Config: "../../conf/defaults.ini"})
+		assert.NoError(t, err)
+
+		// Don't add any custom index settings, test defaults
+		cfg.setUnifiedStorageConfig()
+
+		// Test that default index settings are applied
+		assert.Equal(t, 1, cfg.IndexMinCount)
+		assert.Equal(t, 0, cfg.IndexMaxCount)
 	})
 }

pkg/storage/unified/resource/search.go

@@ -14,6 +14,7 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 
 	"github.com/grafana/authlib/types"
@@ -117,6 +118,7 @@ type searchSupport struct {
 	builders     *builderCache
 	initWorkers  int
 	initMinSize  int
+	initMaxSize  int
 
 	// Index queue processors
 	indexQueueProcessorsMutex sync.Mutex
@@ -156,6 +158,7 @@ func newSearchSupport(opts SearchOptions, storage StorageBackend, access types.A
 		log:                   slog.Default().With("logger", "resource-search"),
 		initWorkers:           opts.WorkerThreads,
 		initMinSize:           opts.InitMinCount,
+		initMaxSize:           opts.InitMaxCount,
 		indexMetrics:          indexMetrics,
 		clientIndexEventsChan: opts.IndexEventsChan,
 		indexEventsChan:       make(chan *IndexEvent),
@@ -406,8 +409,17 @@ func (s *searchSupport) buildIndexes(ctx context.Context, rebuild bool) (int, er
 				// we need to clear the cache to make sure we get the latest usage insights data
 				s.builders.clearNamespacedCache(info.NamespacedResource)
 			}
-			s.log.Debug("building index", "namespace", info.Namespace, "group", info.Group, "resource", info.Resource)
 			totalBatchesIndexed++
+
+			// If the count is too large, we need to set the index to empty.
+			// Only do this if the max size is set to a non-zero (default) value.
+			if s.initMaxSize > 0 && (info.Count > int64(s.initMaxSize)) {
+				s.log.Info("setting empty index for resource with count greater than max size", "namespace", info.Namespace, "group", info.Group, "resource", info.Resource, "count", info.Count, "maxSize", s.initMaxSize)
+				_, err := s.buildEmptyIndex(ctx, info.NamespacedResource, info.ResourceVersion)
+				return err
+			}
+
+			s.log.Debug("building index", "namespace", info.Namespace, "group", info.Group, "resource", info.Resource)
 			_, _, err := s.build(ctx, info.NamespacedResource, info.Count, info.ResourceVersion)
 			return err
 		})
@@ -715,6 +727,21 @@ func (s *searchSupport) build(ctx context.Context, nsr NamespacedResource, size
 	return index, rv, err
 }
 
+// buildEmptyIndex creates an empty index without adding any documents
+func (s *searchSupport) buildEmptyIndex(ctx context.Context, nsr NamespacedResource, rv int64) (ResourceIndex, error) {
+	ctx, span := s.tracer.Start(ctx, tracingPrexfixSearch+"BuildEmptyIndex")
+	defer span.End()
+
+	fields := s.builders.GetFields(nsr)
+	s.log.Debug("Building empty index", "namespace", nsr.Namespace, "group", nsr.Group, "resource", nsr.Resource, "rv", rv)
+
+	// Build an empty index by passing a builder function that doesn't add any documents
+	return s.search.BuildIndex(ctx, nsr, 0, rv, fields, func(index ResourceIndex) (int64, error) {
+		// Return the resource version without adding any documents to the index
+		return rv, nil
+	})
+}
+
 type builderCache struct {
 	// The default builder
 	defaultBuilder DocumentBuilder
@@ -856,3 +883,75 @@ func (s *builderCache) clearNamespacedCache(key NamespacedResource) {
 	defer s.mu.Unlock()
 	s.ns.Remove(key)
 }
+
+// Test utilities for document building
+
+// testDocumentBuilder implements DocumentBuilder for testing
+type testDocumentBuilder struct{}
+
+func (b *testDocumentBuilder) BuildDocument(ctx context.Context, key *resourcepb.ResourceKey, rv int64, value []byte) (*IndexableDocument, error) {
+	// convert value to unstructured.Unstructured
+	var u unstructured.Unstructured
+	if err := u.UnmarshalJSON(value); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal value: %w", err)
+	}
+
+	title := ""
+	tags := []string{}
+	val := ""
+
+	spec, ok, _ := unstructured.NestedMap(u.Object, "spec")
+	if ok {
+		if v, ok := spec["title"]; ok {
+			title = v.(string)
+		}
+		if v, ok := spec["tags"]; ok {
+			if tagSlice, ok := v.([]interface{}); ok {
+				tags = make([]string, len(tagSlice))
+				for i, tag := range tagSlice {
+					if strTag, ok := tag.(string); ok {
+						tags[i] = strTag
+					}
+				}
+			}
+		}
+		if v, ok := spec["value"]; ok {
+			val = v.(string)
+		}
+	}
+	return &IndexableDocument{
+		Key: &resourcepb.ResourceKey{
+			Namespace: key.Namespace,
+			Group:     key.Group,
+			Resource:  key.Resource,
+			Name:      u.GetName(),
+		},
+		Title: title,
+		Tags:  tags,
+		Fields: map[string]interface{}{
+			"value": val,
+		},
+	}, nil
+}
+
+// TestDocumentBuilderSupplier implements DocumentBuilderSupplier for testing
+type TestDocumentBuilderSupplier struct {
+	GroupsResources map[string]string
+}
+
+func (s *TestDocumentBuilderSupplier) GetDocumentBuilders() ([]DocumentBuilderInfo, error) {
+	builders := make([]DocumentBuilderInfo, 0, len(s.GroupsResources))
+
+	// Add builders for all possible group/resource combinations
+	for group, resourceType := range s.GroupsResources {
+		builders = append(builders, DocumentBuilderInfo{
+			GroupResource: schema.GroupResource{
+				Group:    group,
+				Resource: resourceType,
+			},
+			Builder: &testDocumentBuilder{},
+		})
+	}
+
+	return builders, nil
+}

pkg/storage/unified/resource/search_test.go

@@ -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)
+		})
+	}
+}

pkg/storage/unified/resource/server.go

@@ -156,6 +156,10 @@ type SearchOptions struct {
 	// Skip building index on startup for small indexes
 	InitMinCount int
 
+	// Build empty index on startup for large indexes so that
+	// we don't re-attempt to build the index later.
+	InitMaxCount int
+
 	// Channel to watch for index events (for testing)
 	IndexEventsChan chan *IndexEvent
 

pkg/storage/unified/search/options.go

@@ -38,6 +38,7 @@ func NewSearchOptions(features featuremgmt.FeatureToggles, cfg *setting.Cfg, tra
 			Resources:       docs,
 			WorkerThreads:   cfg.IndexWorkers,
 			InitMinCount:    cfg.IndexMinCount,
+			InitMaxCount:    cfg.IndexMaxCount,
 			RebuildInterval: cfg.IndexRebuildInterval,
 		}, nil
 	}

pkg/storage/unified/sql/backend.go

@@ -300,6 +300,8 @@ func (b *backend) GetResourceStats(ctx context.Context, namespace string, minCou
 			}
 			if row.Count > int64(minCount) {
 				res = append(res, row)
+			} else {
+				b.log.Debug("skipping stats for resource with count less than min count", "namespace", row.Namespace, "group", row.Group, "resource", row.Resource, "count", row.Count, "minCount", minCount)
 			}
 		}
 		return err

pkg/storage/unified/testing/benchmark.go

@@ -13,8 +13,6 @@ import (
 	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
 
 	"github.com/stretchr/testify/require"
-	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
-	"k8s.io/apimachinery/pkg/runtime/schema"
 )
 
 // BenchmarkOptions configures the benchmark parameters
@@ -350,7 +348,7 @@ func BenchmarkIndexServer(tb testing.TB, ctx context.Context, backend resource.S
 		Search: resource.SearchOptions{
 			Backend:         searchBackend,
 			IndexEventsChan: events,
-			Resources:       &testDocumentBuilderSupplier{groupsResources: groupsResources},
+			Resources:       &resource.TestDocumentBuilderSupplier{GroupsResources: groupsResources},
 		},
 	})
 	require.NoError(tb, err)
@@ -424,73 +422,3 @@ func BenchmarkIndexServer(tb testing.TB, ctx context.Context, backend resource.S
 	tb.Logf("P90 Index Latency: %.3fs", p90)
 	tb.Logf("P99 Index Latency: %.3fs", p99)
 }
-
-// testDocumentBuilder implements DocumentBuilder for testing
-type testDocumentBuilder struct{}
-
-func (b *testDocumentBuilder) BuildDocument(ctx context.Context, key *resourcepb.ResourceKey, rv int64, value []byte) (*resource.IndexableDocument, error) {
-	// convert value to unstructured.Unstructured
-	var u unstructured.Unstructured
-	if err := u.UnmarshalJSON(value); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal value: %w", err)
-	}
-
-	title := ""
-	tags := []string{}
-	val := ""
-
-	spec, ok, _ := unstructured.NestedMap(u.Object, "spec")
-	if ok {
-		if v, ok := spec["title"]; ok {
-			title = v.(string)
-		}
-		if v, ok := spec["tags"]; ok {
-			if tagSlice, ok := v.([]interface{}); ok {
-				tags = make([]string, len(tagSlice))
-				for i, tag := range tagSlice {
-					if strTag, ok := tag.(string); ok {
-						tags[i] = strTag
-					}
-				}
-			}
-		}
-		if v, ok := spec["value"]; ok {
-			val = v.(string)
-		}
-	}
-	return &resource.IndexableDocument{
-		Key: &resourcepb.ResourceKey{
-			Namespace: key.Namespace,
-			Group:     key.Group,
-			Resource:  key.Resource,
-			Name:      u.GetName(),
-		},
-		Title: title,
-		Tags:  tags,
-		Fields: map[string]interface{}{
-			"value": val,
-		},
-	}, nil
-}
-
-// testDocumentBuilderSupplier implements DocumentBuilderSupplier for testing
-type testDocumentBuilderSupplier struct {
-	groupsResources map[string]string
-}
-
-func (s *testDocumentBuilderSupplier) GetDocumentBuilders() ([]resource.DocumentBuilderInfo, error) {
-	builders := make([]resource.DocumentBuilderInfo, 0, len(s.groupsResources))
-
-	// Add builders for all possible group/resource combinations
-	for group, resourceType := range s.groupsResources {
-		builders = append(builders, resource.DocumentBuilderInfo{
-			GroupResource: schema.GroupResource{
-				Group:    group,
-				Resource: resourceType,
-			},
-			Builder: &testDocumentBuilder{},
-		})
-	}
-
-	return builders, nil
-}

pkg/storage/unified/testing/search_and_storage.go

@@ -103,8 +103,8 @@ func RunTestSearchAndStorage(t *testing.T, ctx context.Context, backend resource
 			Backend: backend,
 			Search: resource.SearchOptions{
 				Backend: searchBackend,
-				Resources: &testDocumentBuilderSupplier{
+				Resources: &resource.TestDocumentBuilderSupplier{
-					groupsResources: map[string]string{
+					GroupsResources: map[string]string{
 						"test.grafana.app": "testresources",
 					},
 				},