search: add index batching (#104163)

* add basic search backend integration tests

* add search backend benchmark

* add benchmark indexServer

* fix

* lint

* add more tests

* lint

* do not use the poller

* batch write

* refactor and add tests

* improvements

* improvements

* cleanup

* only observe index success

* add monitorIndexEvents method

* nit use switch instead of if

* make newIndexQueueProcessor private

* simplify runProcessor

* go lint

pkg/storage/unified/testing/benchmark.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/storage/unified/resource"
 	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 )
 
 // BenchmarkOptions configures the benchmark parameters
@@ -43,6 +44,28 @@ type BenchmarkResult struct {
 	P99Latency    time.Duration
 }
 
+// initializeBackend sets up the backend with initial resources for each group and resource type combination
+func initializeBackend(ctx context.Context, backend resource.StorageBackend, opts *BenchmarkOptions) error {
+	for ns := 0; ns < opts.NumNamespaces; ns++ {
+		namespace := fmt.Sprintf("ns-%d", ns)
+		for g := 0; g < opts.NumGroups; g++ {
+			group := fmt.Sprintf("group-%d", g)
+			for r := 0; r < opts.NumResourceTypes; r++ {
+				resourceType := fmt.Sprintf("resource-%d", r)
+				_, err := writeEvent(ctx, backend, "init", resource.WatchEvent_ADDED,
+					WithNamespace(namespace),
+					WithGroup(group),
+					WithResource(resourceType),
+					WithValue([]byte("init")))
+				if err != nil {
+					return fmt.Errorf("failed to initialize backend: %w", err)
+				}
+			}
+		}
+	}
+	return nil
+}
+
 // runStorageBackendBenchmark runs a write throughput benchmark
 func runStorageBackendBenchmark(ctx context.Context, backend resource.StorageBackend, opts *BenchmarkOptions) (*BenchmarkResult, error) {
 	if opts == nil {
@@ -62,22 +85,6 @@ func runStorageBackendBenchmark(ctx context.Context, backend resource.StorageBac
 
 	var wg sync.WaitGroup
 
-	// Initialize each group and resource type combination in the init namespace
-	namespace := "ns-init"
-	for g := 0; g < opts.NumGroups; g++ {
-		group := fmt.Sprintf("group-%d", g)
-		for r := 0; r < opts.NumResourceTypes; r++ {
-			resourceType := fmt.Sprintf("resource-%d", r)
-			_, err := writeEvent(ctx, backend, "init", resource.WatchEvent_ADDED,
-				WithNamespace(namespace),
-				WithGroup(group),
-				WithResource(resourceType),
-				WithValue([]byte("init")))
-			if err != nil {
-				return nil, fmt.Errorf("failed to initialize backend: %w", err)
-			}
-		}
-	}
 	// Start workers
 	startTime := time.Now()
 	for workerID := 0; workerID < opts.Concurrency; workerID++ {
@@ -147,16 +154,24 @@ func runStorageBackendBenchmark(ctx context.Context, backend resource.StorageBac
 }
 
 // BenchmarkStorageBackend runs a benchmark test for a storage backend implementation
-func BenchmarkStorageBackend(b *testing.B, backend resource.StorageBackend, opts *BenchmarkOptions) {
+func BenchmarkStorageBackend(b testing.TB, backend resource.StorageBackend, opts *BenchmarkOptions) {
 	ctx := context.Background()
 
+	// Initialize the backend
+	err := initializeBackend(ctx, backend, opts)
+	require.NoError(b, err)
+
+	// Run the benchmark
 	result, err := runStorageBackendBenchmark(ctx, backend, opts)
 	require.NoError(b, err)
 
-	b.ReportMetric(result.Throughput, "writes/sec")
-	b.ReportMetric(float64(result.P50Latency.Milliseconds()), "p50-latency-ms")
-	b.ReportMetric(float64(result.P90Latency.Milliseconds()), "p90-latency-ms")
-	b.ReportMetric(float64(result.P99Latency.Milliseconds()), "p99-latency-ms")
+	// Only report metrics if we're running a benchmark
+	if bb, ok := b.(*testing.B); ok {
+		bb.ReportMetric(result.Throughput, "writes/sec")
+		bb.ReportMetric(float64(result.P50Latency.Milliseconds()), "p50-latency-ms")
+		bb.ReportMetric(float64(result.P90Latency.Milliseconds()), "p90-latency-ms")
+		bb.ReportMetric(float64(result.P99Latency.Milliseconds()), "p99-latency-ms")
+	}
 
 	// Also log the results for better visibility
 	b.Logf("Benchmark Configuration: Workers=%d, Resources=%d, Namespaces=%d, Groups=%d, Resource Types=%d", opts.Concurrency, opts.NumResources, opts.NumNamespaces, opts.NumGroups, opts.NumResourceTypes)
@@ -169,3 +184,268 @@ func BenchmarkStorageBackend(b *testing.B, backend resource.StorageBackend, opts
 	b.Logf("P90 Latency: %v", result.P90Latency)
 	b.Logf("P99 Latency: %v", result.P99Latency)
 }
+
+// runSearchBackendBenchmarkWriteThroughput runs a write throughput benchmark for search backend
+// This is a simple benchmark that writes a single resource/group/namespace because indices are per-tenant/group/resource.
+func runSearchBackendBenchmarkWriteThroughput(ctx context.Context, backend resource.SearchBackend, opts *BenchmarkOptions) (*BenchmarkResult, error) {
+	if opts == nil {
+		opts = DefaultBenchmarkOptions()
+	}
+
+	// Create channels for workers
+	jobs := make(chan int, opts.NumResources)
+	results := make(chan time.Duration, opts.NumResources)
+	errors := make(chan error, opts.NumResources)
+
+	// Fill the jobs channel
+	for i := 0; i < opts.NumResources; i++ {
+		jobs <- i
+	}
+	close(jobs)
+
+	var wg sync.WaitGroup
+
+	// Initialize namespace and resource type
+	nr := resource.NamespacedResource{
+		Namespace: "ns-init",
+		Group:     "group",
+		Resource:  "resource",
+	}
+
+	// Build initial index
+	size := int64(10000) // force the index to be on disk
+	index, err := backend.BuildIndex(ctx, nr, size, 0, nil, func(index resource.ResourceIndex) (int64, error) {
+		return 0, nil
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to initialize backend: %w", err)
+	}
+
+	// Start workers
+	startTime := time.Now()
+	for workerID := 0; workerID < opts.Concurrency; workerID++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			batch := make([]*resource.BulkIndexItem, 0, 1000)
+
+			for jobID := range jobs {
+				doc := &resource.IndexableDocument{
+					Key: &resource.ResourceKey{
+						Namespace: nr.Namespace,
+						Group:     nr.Group,
+						Resource:  nr.Resource,
+						Name:      fmt.Sprintf("item-%d", jobID),
+					},
+					Title: fmt.Sprintf("Document %d", jobID),
+					Tags:  []string{"tag1", "tag2"},
+					Fields: map[string]interface{}{
+						"field1": jobID,
+						"field2": fmt.Sprintf("value-%d", jobID),
+					},
+				}
+
+				batch = append(batch, &resource.BulkIndexItem{
+					Action: resource.ActionIndex,
+					Doc:    doc,
+				})
+
+				// If we've collected 100 items or this is the last job, process the batch
+				if len(batch) == 100 || jobID == opts.NumResources-1 {
+					writeStart := time.Now()
+					err := index.BulkIndex(&resource.BulkIndexRequest{
+						Items: batch,
+					})
+					if err != nil {
+						errors <- err
+						return
+					}
+
+					// Record the latency for each document in the batch
+					latency := time.Since(writeStart)
+					for i := 0; i < len(batch); i++ {
+						results <- latency
+					}
+
+					// Reset the batch
+					batch = batch[:0]
+				}
+			}
+		}()
+	}
+
+	// Wait for all workers to complete
+	wg.Wait()
+	close(results)
+	close(errors)
+
+	// Check for errors
+	if len(errors) > 0 {
+		return nil, <-errors // Return the first error encountered
+	}
+
+	// Collect all latencies
+	latencies := make([]time.Duration, 0, opts.NumResources)
+	for latency := range results {
+		latencies = append(latencies, latency)
+	}
+
+	// Sort latencies for percentile calculation
+	sort.Slice(latencies, func(i, j int) bool {
+		return latencies[i] < latencies[j]
+	})
+
+	totalDuration := time.Since(startTime)
+	throughput := float64(opts.NumResources) / totalDuration.Seconds()
+
+	return &BenchmarkResult{
+		TotalDuration: totalDuration,
+		WriteCount:    opts.NumResources,
+		Throughput:    throughput,
+		P50Latency:    latencies[len(latencies)*50/100],
+		P90Latency:    latencies[len(latencies)*90/100],
+		P99Latency:    latencies[len(latencies)*99/100],
+	}, nil
+}
+
+// BenchmarkSearchBackend runs a benchmark test for a search backend implementation
+func BenchmarkSearchBackend(tb testing.TB, backend resource.SearchBackend, opts *BenchmarkOptions) {
+	ctx := context.Background()
+
+	result, err := runSearchBackendBenchmarkWriteThroughput(ctx, backend, opts)
+	require.NoError(tb, err)
+
+	if b, ok := tb.(*testing.B); ok {
+		b.ReportMetric(result.Throughput, "writes/sec")
+		b.ReportMetric(float64(result.P50Latency.Milliseconds()), "p50-latency-ms")
+		b.ReportMetric(float64(result.P90Latency.Milliseconds()), "p90-latency-ms")
+		b.ReportMetric(float64(result.P99Latency.Milliseconds()), "p99-latency-ms")
+	}
+
+	// Also log the results for better visibility
+	tb.Logf("Benchmark Configuration: Workers=%d, Resources=%d, Namespaces=%d, Groups=%d, Resource Types=%d", opts.Concurrency, opts.NumResources, opts.NumNamespaces, opts.NumGroups, opts.NumResourceTypes)
+	tb.Logf("")
+	tb.Logf("Benchmark Results:")
+	tb.Logf("Total Duration: %v", result.TotalDuration)
+	tb.Logf("Write Count: %d", result.WriteCount)
+	tb.Logf("Throughput: %.2f writes/sec", result.Throughput)
+	tb.Logf("P50 Latency: %v", result.P50Latency)
+	tb.Logf("P90 Latency: %v", result.P90Latency)
+	tb.Logf("P99 Latency: %v", result.P99Latency)
+}
+
+func BenchmarkIndexServer(tb testing.TB, ctx context.Context, backend resource.StorageBackend, searchBackend resource.SearchBackend, opts *BenchmarkOptions) {
+	events := make(chan *resource.IndexEvent, opts.NumResources)
+	server, err := resource.NewResourceServer(resource.ResourceServerOptions{
+		Backend: backend,
+		Search: resource.SearchOptions{
+			Backend:         searchBackend,
+			IndexEventsChan: events,
+			Resources:       &testDocumentBuilderSupplier{opts: opts},
+		},
+	})
+	require.NoError(tb, err)
+	require.NotNil(tb, server)
+
+	// Initialize the backend
+	err = initializeBackend(ctx, backend, opts)
+	require.NoError(tb, err)
+
+	// Discard the latencies from the initial index build.
+	for i := 0; i < (opts.NumGroups * opts.NumResourceTypes * opts.NumNamespaces); i++ {
+		<-events
+	}
+
+	// Run the storage backend benchmark write throughput to create events
+	startTime := time.Now()
+	var result *BenchmarkResult
+	go func() {
+		result, err = runStorageBackendBenchmark(ctx, backend, opts)
+		require.NoError(tb, err)
+	}()
+
+	// Wait for all events to be processed
+	latencies := make([]float64, 0, opts.NumResources)
+	for i := 0; i < opts.NumResources; i++ {
+		evt := <-events
+		latencies = append(latencies, evt.Latency.Seconds())
+	}
+	totalDuration := time.Since(startTime)
+	// Calculate index latency percentiles
+	sort.Float64s(latencies)
+	var p50, p90, p99 float64
+	if len(latencies) > 0 {
+		p50 = latencies[len(latencies)*50/100]
+		p90 = latencies[len(latencies)*90/100]
+		p99 = latencies[len(latencies)*99/100]
+	}
+
+	// Report metrics if running a benchmark
+	if b, ok := tb.(*testing.B); ok {
+		b.ReportMetric(result.Throughput, "writes/sec")
+		b.ReportMetric(float64(result.P50Latency.Milliseconds()), "p50-latency-ms")
+		b.ReportMetric(float64(result.P90Latency.Milliseconds()), "p90-latency-ms")
+		b.ReportMetric(float64(result.P99Latency.Milliseconds()), "p99-latency-ms")
+		b.ReportMetric(p50, "p50-index-latency-s")
+		b.ReportMetric(p90, "p90-index-latency-s")
+		b.ReportMetric(p99, "p99-index-latency-s")
+	}
+
+	// Log results for better visibility
+	tb.Logf("Benchmark Configuration: Workers=%d, Resources=%d, Namespaces=%d, Groups=%d, Resource Types=%d",
+		opts.Concurrency, opts.NumResources, opts.NumNamespaces, opts.NumGroups, opts.NumResourceTypes)
+	tb.Logf("")
+	tb.Logf("Storage Benchmark Results:")
+	tb.Logf("Total Duration: %v", result.TotalDuration)
+	tb.Logf("Storage Write Count: %d", result.WriteCount)
+	tb.Logf("Storage Write Throughput: %.2f writes/sec", result.Throughput)
+	tb.Logf("P50 Write Latency: %v", result.P50Latency)
+	tb.Logf("P90 Write Latency: %v", result.P90Latency)
+	tb.Logf("P99 Write Latency: %v", result.P99Latency)
+	tb.Logf("")
+	tb.Logf("Index Latency Results:")
+	tb.Logf("Indexing Throughput: %.2f events/sec", float64(len(latencies))/totalDuration.Seconds())
+	tb.Logf("P50 Index Latency: %.3fs", p50)
+	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 *resource.ResourceKey, rv int64, value []byte) (*resource.IndexableDocument, error) {
+	return &resource.IndexableDocument{
+		Key:   key,
+		Title: fmt.Sprintf("Document %s", key.Name),
+		Tags:  []string{"test", "benchmark"},
+		Fields: map[string]interface{}{
+			"value": string(value),
+		},
+	}, nil
+}
+
+// testDocumentBuilderSupplier implements DocumentBuilderSupplier for testing
+type testDocumentBuilderSupplier struct {
+	opts *BenchmarkOptions
+}
+
+func (s *testDocumentBuilderSupplier) GetDocumentBuilders() ([]resource.DocumentBuilderInfo, error) {
+	builders := make([]resource.DocumentBuilderInfo, 0, s.opts.NumGroups*s.opts.NumResourceTypes)
+
+	// Add builders for all possible group/resource combinations
+	for g := 0; g < s.opts.NumGroups; g++ {
+		group := fmt.Sprintf("group-%d", g)
+		for r := 0; r < s.opts.NumResourceTypes; r++ {
+			resourceType := fmt.Sprintf("resource-%d", r)
+			builders = append(builders, resource.DocumentBuilderInfo{
+				GroupResource: schema.GroupResource{
+					Group:    group,
+					Resource: resourceType,
+				},
+				Builder: &testDocumentBuilder{},
+			})
+		}
+	}
+
+	return builders, nil
+}

pkg/storage/unified/testing/search_backend.go

@@ -0,0 +1,180 @@
+package test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
+	"github.com/grafana/grafana/pkg/util/testutil"
+)
+
+// Test names for the search backend test suite
+const (
+	TestBuildIndex    = "build index"
+	TestTotalDocs     = "total docs"
+	TestResourceIndex = "resource index"
+)
+
+// NewSearchBackendFunc is a function that creates a new SearchBackend instance
+type NewSearchBackendFunc func(ctx context.Context) resource.SearchBackend
+
+// RunSearchBackendTest runs the search backend test suite
+func RunSearchBackendTest(t *testing.T, newBackend NewSearchBackendFunc, opts *TestOptions) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	if opts == nil {
+		opts = &TestOptions{}
+	}
+
+	if opts.NSPrefix == "" {
+		opts.NSPrefix = "test-" + time.Now().Format("20060102150405")
+	}
+
+	t.Logf("Running tests with namespace prefix: %s", opts.NSPrefix)
+
+	cases := []struct {
+		name string
+		fn   func(*testing.T, resource.SearchBackend, string)
+	}{
+		{TestBuildIndex, runTestSearchBackendBuildIndex},
+		{TestTotalDocs, runTestSearchBackendTotalDocs},
+		{TestResourceIndex, runTestResourceIndex},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.fn(t, newBackend(context.Background()), opts.NSPrefix)
+		})
+	}
+}
+
+func runTestSearchBackendBuildIndex(t *testing.T, backend resource.SearchBackend, nsPrefix string) {
+	ctx := testutil.NewTestContext(t, time.Now().Add(5*time.Second))
+	ns := resource.NamespacedResource{
+		Namespace: nsPrefix + "-ns1",
+		Group:     "group",
+		Resource:  "resource",
+	}
+
+	// Get the index should return nil if the index does not exist
+	index, err := backend.GetIndex(ctx, ns)
+	require.NoError(t, err)
+	require.Nil(t, index)
+
+	// Build the index
+	index, err = backend.BuildIndex(ctx, ns, 0, 0, nil, func(index resource.ResourceIndex) (int64, error) {
+		// Write a test document
+		err := index.BulkIndex(&resource.BulkIndexRequest{
+			Items: []*resource.BulkIndexItem{
+				{
+					Action: resource.ActionIndex,
+					Doc: &resource.IndexableDocument{
+						Key: &resource.ResourceKey{
+							Namespace: ns.Namespace,
+							Group:     ns.Group,
+							Resource:  ns.Resource,
+							Name:      "doc1",
+						},
+						Title: "Document 1",
+					},
+				},
+			},
+		})
+		if err != nil {
+			return 0, err
+		}
+		return 1, nil
+	})
+	require.NoError(t, err)
+	require.NotNil(t, index)
+
+	// Get the index should now return the index
+	index, err = backend.GetIndex(ctx, ns)
+	require.NoError(t, err)
+	require.NotNil(t, index)
+}
+
+func runTestSearchBackendTotalDocs(t *testing.T, backend resource.SearchBackend, nsPrefix string) {
+	// Get total document count
+	count := backend.TotalDocs()
+	require.GreaterOrEqual(t, count, int64(0))
+}
+
+func runTestResourceIndex(t *testing.T, backend resource.SearchBackend, nsPrefix string) {
+	ctx := testutil.NewTestContext(t, time.Now().Add(5*time.Second))
+	ns := resource.NamespacedResource{
+		Namespace: nsPrefix + "-ns1",
+		Group:     "group",
+		Resource:  "resource",
+	}
+
+	// Build initial index with some test documents
+	index, err := backend.BuildIndex(ctx, ns, 3, 0, nil, func(index resource.ResourceIndex) (int64, error) {
+		err := index.BulkIndex(&resource.BulkIndexRequest{
+			Items: []*resource.BulkIndexItem{
+				{
+					Action: resource.ActionIndex,
+					Doc: &resource.IndexableDocument{
+						Key: &resource.ResourceKey{
+							Namespace: ns.Namespace,
+							Group:     ns.Group,
+							Resource:  ns.Resource,
+							Name:      "doc1",
+						},
+						Title: "Document 1",
+						Tags:  []string{"tag1", "tag2"},
+						Fields: map[string]interface{}{
+							"field1": 1,
+							"field2": "value1",
+						},
+					},
+				},
+				{
+					Action: resource.ActionIndex,
+					Doc: &resource.IndexableDocument{
+						Key: &resource.ResourceKey{
+							Namespace: ns.Namespace,
+							Group:     ns.Group,
+							Resource:  ns.Resource,
+							Name:      "doc2",
+						},
+						Title: "Document 2",
+						Tags:  []string{"tag2", "tag3"},
+						Fields: map[string]interface{}{
+							"field1": 2,
+							"field2": "value2",
+						},
+					},
+				},
+			},
+		})
+		require.NoError(t, err)
+		return int64(2), nil
+	})
+	require.NoError(t, err)
+	require.NotNil(t, index)
+
+	t.Run("Search", func(t *testing.T) {
+		req := &resource.ResourceSearchRequest{
+			Options: &resource.ListOptions{
+				Key: &resource.ResourceKey{
+					Namespace: ns.Namespace,
+					Group:     ns.Group,
+					Resource:  ns.Resource,
+				},
+			},
+			Fields: []string{"title", "folder", "tags"},
+			Query:  "tag3",
+			Limit:  10,
+		}
+		resp, err := index.Search(ctx, nil, req, nil)
+		require.NoError(t, err)
+		require.NotNil(t, resp)
+		require.Equal(t, int64(1), resp.TotalHits) // Only doc3 should have tag3 now
+	})
+}