From 15b3de5893900c3d8fc4f7a87cad38d9bfdcceba Mon Sep 17 00:00:00 2001 From: Georges Chaudy Date: Fri, 9 May 2025 15:36:21 +0200 Subject: [PATCH] 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/resource/search.go | 194 +++++--- pkg/storage/unified/resource/search_queue.go | 143 ++++++ .../unified/resource/search_queue_test.go | 173 +++++++ pkg/storage/unified/resource/search_test.go | 54 +++ pkg/storage/unified/resource/server.go | 3 + pkg/storage/unified/search/bleve.go | 80 ++-- .../unified/search/bleve_integration_test.go | 56 +++ .../unified/search/bleve_performance_test.go | 58 ++- .../unified/search/bleve_search_test.go | 435 +++++++++++------- pkg/storage/unified/search/bleve_test.go | 263 ++++++----- .../unified/sql/test/benchmark_test.go | 68 ++- pkg/storage/unified/testing/benchmark.go | 322 ++++++++++++- pkg/storage/unified/testing/search_backend.go | 180 ++++++++ 13 files changed, 1562 insertions(+), 467 deletions(-) create mode 100644 pkg/storage/unified/resource/search_queue.go create mode 100644 pkg/storage/unified/resource/search_queue_test.go create mode 100644 pkg/storage/unified/resource/search_test.go create mode 100644 pkg/storage/unified/search/bleve_integration_test.go create mode 100644 pkg/storage/unified/testing/search_backend.go diff --git a/pkg/storage/unified/resource/search.go b/pkg/storage/unified/resource/search.go index 0f8da3f4a78..0a05ee97567 100644 --- a/pkg/storage/unified/resource/search.go +++ b/pkg/storage/unified/resource/search.go @@ -21,6 +21,8 @@ import ( folders "github.com/grafana/grafana/apps/folder/pkg/apis/folder/v1beta1" ) +const maxBatchSize = 1000 + type NamespacedResource struct { Namespace string Group string @@ -32,15 +34,28 @@ func (s *NamespacedResource) Valid() bool { return s.Namespace != "" && s.Group != "" && s.Resource != "" } +type IndexAction int + +const ( + ActionIndex IndexAction = iota + ActionDelete +) + +type BulkIndexItem struct { + Action IndexAction + Key *ResourceKey // Only used for delete actions + Doc *IndexableDocument // Only used for index actions +} + +type BulkIndexRequest struct { + Items []*BulkIndexItem + ResourceVersion int64 +} + type ResourceIndex interface { - // Add a document to the index. Note it may not be searchable until after flush is called - Write(doc *IndexableDocument) error - - // Mark a resource as deleted. Note it may not be searchable until after flush is called - Delete(key *ResourceKey) error - - // Make sure any changes to the index are flushed and available in the next search/origin calls - Flush() error + // BulkIndex allows for multiple index actions to be performed in a single call. + // The order of the items is guaranteed to be the same as the input + BulkIndex(req *BulkIndexRequest) error // Search within a namespaced resource // When working with federated queries, the additional indexes will be passed in explicitly @@ -96,6 +111,14 @@ type searchSupport struct { builders *builderCache initWorkers int initMinSize int + + // Index queue processors + indexQueueProcessorsMutex sync.Mutex + indexQueueProcessors map[string]*indexQueueProcessor + indexEventsChan chan *IndexEvent + + // testing + clientIndexEventsChan chan *IndexEvent } var ( @@ -117,14 +140,17 @@ func newSearchSupport(opts SearchOptions, storage StorageBackend, access types.A } support = &searchSupport{ - access: access, - tracer: tracer, - storage: storage, - search: opts.Backend, - log: slog.Default().With("logger", "resource-search"), - initWorkers: opts.WorkerThreads, - initMinSize: opts.InitMinCount, - indexMetrics: indexMetrics, + access: access, + tracer: tracer, + storage: storage, + search: opts.Backend, + log: slog.Default().With("logger", "resource-search"), + initWorkers: opts.WorkerThreads, + initMinSize: opts.InitMinCount, + indexMetrics: indexMetrics, + clientIndexEventsChan: opts.IndexEventsChan, + indexEventsChan: make(chan *IndexEvent), + indexQueueProcessors: make(map[string]*indexQueueProcessor), } info, err := opts.Resources.GetDocumentBuilders() @@ -394,10 +420,12 @@ func (s *searchSupport) init(ctx context.Context) error { continue } - s.handleEvent(watchctx, v) + s.dispatchEvent(watchctx, v) } }() + go s.monitorIndexEvents(ctx) + end := time.Now().Unix() s.log.Info("search index initialized", "duration_secs", end-start, "total_docs", s.search.TotalDocs()) if s.indexMetrics != nil { @@ -407,13 +435,11 @@ func (s *searchSupport) init(ctx context.Context) error { return nil } -// Async event -func (s *searchSupport) handleEvent(ctx context.Context, evt *WrittenEvent) { - ctx, span := s.tracer.Start(ctx, tracingPrexfixSearch+"HandleEvent") - if !slices.Contains([]WatchEvent_Type{WatchEvent_ADDED, WatchEvent_MODIFIED, WatchEvent_DELETED}, evt.Type) { - s.log.Info("ignoring watch event", "type", evt.Type) - return - } +// Async event dispatching +// This is called from the watch event loop +// It will dispatch the event to the appropriate index queue processor +func (s *searchSupport) dispatchEvent(ctx context.Context, evt *WrittenEvent) { + ctx, span := s.tracer.Start(ctx, tracingPrexfixSearch+"dispatchEvent") defer span.End() span.SetAttributes( attribute.String("event_type", evt.Type.String()), @@ -423,69 +449,60 @@ func (s *searchSupport) handleEvent(ctx context.Context, evt *WrittenEvent) { attribute.String("name", evt.Key.Name), ) + switch evt.Type { + case WatchEvent_ADDED, WatchEvent_MODIFIED, WatchEvent_DELETED: // OK + default: + s.log.Info("ignoring watch event", "type", evt.Type) + span.AddEvent("ignoring watch event", trace.WithAttributes(attribute.String("type", evt.Type.String()))) + } + nsr := NamespacedResource{ Namespace: evt.Key.Namespace, Group: evt.Key.Group, Resource: evt.Key.Resource, } - index, err := s.getOrCreateIndex(ctx, nsr) if err != nil { s.log.Warn("error getting index for watch event", "error", err) + span.RecordError(err) return } - - builder, err := s.builders.get(ctx, nsr) + // Get or create index queue processor for this index + indexQueueProcessor, err := s.getOrCreateIndexQueueProcessor(index, nsr) if err != nil { - s.log.Warn("error getting builder for watch event", "error", err) + s.log.Error("error getting index queue processor for watch event", "error", err) + span.RecordError(err) return } + indexQueueProcessor.Add(evt) +} - _, buildDocSpan := s.tracer.Start(ctx, tracingPrexfixSearch+"BuildDocument") - doc, err := builder.BuildDocument(ctx, evt.Key, evt.ResourceVersion, evt.Value) - if err != nil { - s.log.Warn("error building document watch event", "error", err) - return - } - buildDocSpan.End() - - switch evt.Type { - case WatchEvent_ADDED, WatchEvent_MODIFIED: - _, writeSpan := s.tracer.Start(ctx, tracingPrexfixSearch+"WriteDocument") - err = index.Write(doc) - writeSpan.End() - if err != nil { - s.log.Warn("error writing document watch event", "error", err) +func (s *searchSupport) monitorIndexEvents(ctx context.Context) { + var evt *IndexEvent + for { + select { + case <-ctx.Done(): return + case evt = <-s.indexEventsChan: } - if evt.Type == WatchEvent_ADDED && s.indexMetrics != nil { - s.indexMetrics.IndexedKinds.WithLabelValues(evt.Key.Resource).Inc() + if evt.Err != nil { + s.log.Error("error indexing watch event", "error", evt.Err) + continue } - case WatchEvent_DELETED: - _, deleteSpan := s.tracer.Start(ctx, tracingPrexfixSearch+"DeleteDocument") - err = index.Delete(evt.Key) - deleteSpan.End() - if err != nil { - s.log.Warn("error deleting document watch event", "error", err) - return + _, span := s.tracer.Start(ctx, tracingPrexfixSearch+"monitorIndexEvents") + defer span.End() + // record latency from when event was created to when it was indexed + span.AddEvent("index latency", trace.WithAttributes(attribute.Float64("latency_seconds", evt.Latency.Seconds()))) + s.log.Debug("indexed new object", "resource", evt.WrittenEvent.Key.Resource, "latency_seconds", evt.Latency.Seconds(), "name", evt.WrittenEvent.Key.Name, "namespace", evt.WrittenEvent.Key.Namespace, "rv", evt.WrittenEvent.ResourceVersion) + if evt.Latency.Seconds() > 1 { + s.log.Warn("high index latency object details", "resource", evt.WrittenEvent.Key.Resource, "latency_seconds", evt.Latency.Seconds(), "name", evt.WrittenEvent.Key.Name, "namespace", evt.WrittenEvent.Key.Namespace, "rv", evt.WrittenEvent.ResourceVersion) } if s.indexMetrics != nil { - s.indexMetrics.IndexedKinds.WithLabelValues(evt.Key.Resource).Dec() + s.indexMetrics.IndexLatency.WithLabelValues(evt.WrittenEvent.Key.Resource).Observe(evt.Latency.Seconds()) + } + if s.clientIndexEventsChan != nil { + s.clientIndexEventsChan <- evt } - default: - // do nothing - s.log.Warn("unknown watch event", "type", evt.Type) - } - - // record latency from when event was created to when it was indexed - latencySeconds := float64(time.Now().UnixMicro()-evt.ResourceVersion) / 1e6 - span.AddEvent("index latency", trace.WithAttributes(attribute.Float64("latency_seconds", latencySeconds))) - s.log.Debug("indexed new object", "resource", evt.Key.Resource, "latency_seconds", latencySeconds, "name", evt.Key.Name, "namespace", evt.Key.Namespace, "rv", evt.ResourceVersion) - if latencySeconds > 1 { - s.log.Warn("high index latency object details", "resource", evt.Key.Resource, "latency_seconds", latencySeconds, "name", evt.Key.Name, "namespace", evt.Key.Namespace, "rv", evt.ResourceVersion) - } - if s.indexMetrics != nil { - s.indexMetrics.IndexLatency.WithLabelValues(evt.Key.Resource).Observe(latencySeconds) } } @@ -534,6 +551,7 @@ func (s *searchSupport) build(ctx context.Context, nsr NamespacedResource, size Resource: nsr.Resource, Namespace: nsr.Namespace, } + index, err := s.search.BuildIndex(ctx, nsr, size, rv, fields, func(index ResourceIndex) (int64, error) { rv, err = s.storage.ListIterator(ctx, &ListRequest{ Limit: 1000000000000, // big number @@ -541,13 +559,15 @@ func (s *searchSupport) build(ctx context.Context, nsr NamespacedResource, size Key: key, }, }, func(iter ListIterator) error { + // Collect all documents in a single bulk request + items := make([]*BulkIndexItem, 0) + for iter.Next() { if err = iter.Error(); err != nil { return err } // Update the key name - // Or should we read it from the body? key.Name = iter.Name() // Convert it to an indexable document @@ -557,8 +577,18 @@ func (s *searchSupport) build(ctx context.Context, nsr NamespacedResource, size continue } - // And finally write it to the index - if err = index.Write(doc); err != nil { + // Add to bulk items + items = append(items, &BulkIndexItem{ + Action: ActionIndex, + Doc: doc, + }) + } + + // Perform single bulk index operation + if len(items) > 0 { + if err = index.BulkIndex(&BulkIndexRequest{ + Items: items, + }); err != nil { return err } } @@ -580,10 +610,6 @@ func (s *searchSupport) build(ctx context.Context, nsr NamespacedResource, size s.indexMetrics.IndexedKinds.WithLabelValues(key.Resource).Add(float64(docCount)) } - if err == nil { - err = index.Flush() - } - // rv is the last RV we read. when watching, we must add all events since that time return index, rv, err } @@ -703,3 +729,23 @@ func AsResourceKey(ns string, t string) (*ResourceKey, error) { return nil, fmt.Errorf("unknown resource type") } + +// getOrCreateIndexQueueProcessor returns an IndexQueueProcessor for the given index +func (s *searchSupport) getOrCreateIndexQueueProcessor(index ResourceIndex, nsr NamespacedResource) (*indexQueueProcessor, error) { + s.indexQueueProcessorsMutex.Lock() + defer s.indexQueueProcessorsMutex.Unlock() + + key := fmt.Sprintf("%s/%s/%s", nsr.Namespace, nsr.Group, nsr.Resource) + if indexQueueProcessor, ok := s.indexQueueProcessors[key]; ok { + return indexQueueProcessor, nil + } + + builder, err := s.builders.get(context.Background(), nsr) + if err != nil { + s.log.Error("error getting document builder", "error", err) + return nil, err + } + indexQueueProcessor := newIndexQueueProcessor(index, nsr, maxBatchSize, builder, s.indexEventsChan) + s.indexQueueProcessors[key] = indexQueueProcessor + return indexQueueProcessor, nil +} diff --git a/pkg/storage/unified/resource/search_queue.go b/pkg/storage/unified/resource/search_queue.go new file mode 100644 index 00000000000..71da4c04434 --- /dev/null +++ b/pkg/storage/unified/resource/search_queue.go @@ -0,0 +1,143 @@ +package resource + +import ( + "context" + "sync" + "time" +) + +// indexQueueProcessor manages queue-based operations for a specific index +// It is responsible for ingesting events for a single index +// It will batch events and send them to the index in a single bulk request +type indexQueueProcessor struct { + index ResourceIndex + nsr NamespacedResource + queue chan *WrittenEvent + batchSize int + builder DocumentBuilder + + resChan chan *IndexEvent // Channel to send results to the caller + + mu sync.Mutex + running bool +} + +type IndexEvent struct { + WrittenEvent *WrittenEvent + Action IndexAction + IndexableDocument *IndexableDocument // empty for delete actions + Timestamp time.Time + Latency time.Duration + Err error +} + +// newIndexQueueProcessor creates a new IndexQueueProcessor for the given index +func newIndexQueueProcessor(index ResourceIndex, nsr NamespacedResource, batchSize int, builder DocumentBuilder, resChan chan *IndexEvent) *indexQueueProcessor { + return &indexQueueProcessor{ + index: index, + nsr: nsr, + queue: make(chan *WrittenEvent, 1000), // Buffer size of 1000 events + batchSize: batchSize, + builder: builder, + resChan: resChan, + running: false, + } +} + +// Add adds an event to the queue and ensures the background processor is running +func (b *indexQueueProcessor) Add(evt *WrittenEvent) { + b.queue <- evt + + // Start the processor if it's not already running + b.mu.Lock() + defer b.mu.Unlock() + if !b.running { + b.running = true + go b.runProcessor() + } +} + +// runProcessor is the task processing the queue of written events +func (b *indexQueueProcessor) runProcessor() { + defer func() { + b.mu.Lock() + b.running = false + b.mu.Unlock() + }() + + for { + batch := make([]*WrittenEvent, 0, b.batchSize) + select { + case evt := <-b.queue: + batch = append(batch, evt) + case <-time.After(5 * time.Second): + // No events in the past few seconds, stop the processor + return + } + + prepare: + for len(batch) < b.batchSize { + select { + case evt := <-b.queue: + batch = append(batch, evt) + default: + break prepare + } + } + + b.process(batch) + } +} + +// process handles a batch of events +func (b *indexQueueProcessor) process(batch []*WrittenEvent) { + if len(batch) == 0 { + return + } + + // Create bulk request + req := &BulkIndexRequest{ + Items: make([]*BulkIndexItem, 0, len(batch)), + } + resp := make([]*IndexEvent, 0, len(batch)) + + for _, evt := range batch { + result := &IndexEvent{ + WrittenEvent: evt, + } + resp = append(resp, result) + + item := &BulkIndexItem{} + if evt.Type == WatchEvent_DELETED { + item.Action = ActionDelete + item.Key = evt.Key + } else { + item.Action = ActionIndex + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + doc, err := b.builder.BuildDocument(ctx, evt.Key, evt.ResourceVersion, evt.Value) + if err != nil { + result.Err = err + } else { + item.Doc = doc + result.IndexableDocument = doc + } + } + req.Items = append(req.Items, item) + } + + err := b.index.BulkIndex(req) + if err != nil { + for _, r := range resp { + r.Err = err + } + } + ts := time.Now() + if b.resChan != nil { + for _, r := range resp { + r.Timestamp = ts + r.Latency = time.Duration(ts.UnixMicro()-r.WrittenEvent.ResourceVersion) * time.Microsecond + b.resChan <- r + } + } +} diff --git a/pkg/storage/unified/resource/search_queue_test.go b/pkg/storage/unified/resource/search_queue_test.go new file mode 100644 index 00000000000..9e19ea501cf --- /dev/null +++ b/pkg/storage/unified/resource/search_queue_test.go @@ -0,0 +1,173 @@ +package resource + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +func TestNewIndexQueueProcessor(t *testing.T) { + mockIndex := &MockResourceIndex{} + mockBuilder := &MockDocumentBuilder{} + nsr := NamespacedResource{Resource: "test"} + + resChan := make(chan *IndexEvent) + + processor := newIndexQueueProcessor(mockIndex, nsr, 10, mockBuilder, resChan) + + assert.NotNil(t, processor) + assert.Equal(t, 10, processor.batchSize) + assert.NotNil(t, processor.queue) +} + +func TestIndexQueueProcessor_SingleEvent(t *testing.T) { + mockIndex := &MockResourceIndex{} + mockBuilder := &MockDocumentBuilder{} + nsr := NamespacedResource{Resource: "test"} + + resChan := make(chan *IndexEvent) + + processor := newIndexQueueProcessor(mockIndex, nsr, 10, mockBuilder, resChan) + + // Test data + key := ResourceKey{Resource: "test", Name: "obj1", Namespace: "default"} + evt := &WrittenEvent{ + Key: &key, + ResourceVersion: time.Now().UnixMicro(), + Type: WatchEvent_ADDED, + Value: []byte(`{"test": "data"}`), + } + + // Setup expectations + mockBuilder.On("BuildDocument", mock.Anything, &key, evt.ResourceVersion, evt.Value).Return(&IndexableDocument{Key: &key}, nil) + mockIndex.On("BulkIndex", mock.MatchedBy(func(req *BulkIndexRequest) bool { + return len(req.Items) == 1 && req.Items[0].Action == ActionIndex + })).Return(nil) + + // Start processor and wait for the document to be indexed + processor.Add(evt) + + resp := <-resChan + assert.NotNil(t, resp) + assert.Nil(t, resp.Err) + assert.Equal(t, &key, resp.IndexableDocument.Key) + + mockBuilder.AssertExpectations(t) + mockIndex.AssertExpectations(t) +} + +func TestIndexQueueProcessor_BatchProcessing(t *testing.T) { + mockIndex := &MockResourceIndex{} + mockBuilder := &MockDocumentBuilder{} + nsr := NamespacedResource{Namespace: "default", Resource: "test"} + + resChan := make(chan *IndexEvent) + + processor := newIndexQueueProcessor(mockIndex, nsr, 2, mockBuilder, resChan) + + // Test data for two events + events := []*WrittenEvent{ + { + Key: &ResourceKey{Resource: "test", Name: "obj1", Namespace: "default"}, + ResourceVersion: time.Now().UnixMicro(), + Type: WatchEvent_ADDED, + Value: []byte(`{"test": "data1"}`), + }, + { + Key: &ResourceKey{Resource: "test", Name: "obj2", Namespace: "default"}, + ResourceVersion: time.Now().UnixMicro(), + Type: WatchEvent_DELETED, + }, + } + + // Setup expectations + mockBuilder.On("BuildDocument", mock.Anything, events[0].Key, events[0].ResourceVersion, events[0].Value). + Return(&IndexableDocument{Key: events[0].Key}, nil) + mockIndex.On("BulkIndex", mock.MatchedBy(func(req *BulkIndexRequest) bool { + return len(req.Items) == 2 && + req.Items[0].Action == ActionIndex && + req.Items[1].Action == ActionDelete + })).Return(nil) + + // Start processor and add events + processor.Add(events[0]) + processor.Add(events[1]) + + r0 := <-resChan + assert.Nil(t, r0.Err) + assert.Equal(t, events[0].Key, r0.IndexableDocument.Key) + + r1 := <-resChan + assert.Nil(t, r1.Err) + assert.Nil(t, r1.IndexableDocument) // deleted event + + mockBuilder.AssertExpectations(t) + mockIndex.AssertExpectations(t) +} + +func TestIndexQueueProcessor_BuildDocumentError(t *testing.T) { + mockIndex := &MockResourceIndex{} + mockBuilder := &MockDocumentBuilder{} + nsr := NamespacedResource{Resource: "test"} + + resChan := make(chan *IndexEvent) + + processor := newIndexQueueProcessor(mockIndex, nsr, 10, mockBuilder, resChan) + + evt := &WrittenEvent{ + Key: &ResourceKey{Resource: "test", Name: "obj1", Namespace: "default"}, + ResourceVersion: time.Now().UnixMicro(), + Type: WatchEvent_ADDED, + Value: []byte(`invalid json`), + } + + // Setup expectations for error case + mockBuilder.On("BuildDocument", mock.Anything, evt.Key, evt.ResourceVersion, evt.Value). + Return(nil, assert.AnError) + + // The bulk index should not be called since document building failed + mockIndex.On("BulkIndex", mock.Anything).Return(nil).Maybe() + + processor.Add(evt) + + resp := <-resChan + assert.NotNil(t, resp) + assert.Error(t, resp.Err) + assert.Nil(t, resp.IndexableDocument) + + mockBuilder.AssertExpectations(t) + mockIndex.AssertExpectations(t) +} + +func TestIndexQueueProcessor_BulkIndexError(t *testing.T) { + mockIndex := &MockResourceIndex{} + mockBuilder := &MockDocumentBuilder{} + nsr := NamespacedResource{Resource: "test"} + + resChan := make(chan *IndexEvent) + + processor := newIndexQueueProcessor(mockIndex, nsr, 10, mockBuilder, resChan) + + evt := &WrittenEvent{ + Key: &ResourceKey{Resource: "test", Name: "obj1", Namespace: "default"}, + ResourceVersion: time.Now().UnixMicro(), + Type: WatchEvent_ADDED, + Value: []byte(`{"test": "data"}`), + } + + // Setup expectations + mockBuilder.On("BuildDocument", mock.Anything, evt.Key, evt.ResourceVersion, evt.Value). + Return(&IndexableDocument{Key: evt.Key}, nil) + mockIndex.On("BulkIndex", mock.Anything).Return(assert.AnError) + + processor.Add(evt) + + resp := <-resChan + assert.NotNil(t, resp) + assert.Error(t, resp.Err) + + mockBuilder.AssertExpectations(t) + mockIndex.AssertExpectations(t) +} diff --git a/pkg/storage/unified/resource/search_test.go b/pkg/storage/unified/resource/search_test.go new file mode 100644 index 00000000000..1f7d253002c --- /dev/null +++ b/pkg/storage/unified/resource/search_test.go @@ -0,0 +1,54 @@ +package resource + +import ( + "context" + + "github.com/grafana/authlib/types" + "github.com/stretchr/testify/mock" +) + +var _ ResourceIndex = &MockResourceIndex{} + +// Mock implementations +type MockResourceIndex struct { + mock.Mock +} + +func (m *MockResourceIndex) BulkIndex(req *BulkIndexRequest) error { + args := m.Called(req) + return args.Error(0) +} + +func (m *MockResourceIndex) Search(ctx context.Context, access types.AccessClient, req *ResourceSearchRequest, federate []ResourceIndex) (*ResourceSearchResponse, error) { + args := m.Called(ctx, access, req, federate) + return args.Get(0).(*ResourceSearchResponse), args.Error(1) +} + +func (m *MockResourceIndex) CountManagedObjects(ctx context.Context) ([]*CountManagedObjectsResponse_ResourceCount, error) { + args := m.Called(ctx) + return args.Get(0).([]*CountManagedObjectsResponse_ResourceCount), args.Error(1) +} + +func (m *MockResourceIndex) DocCount(ctx context.Context, folder string) (int64, error) { + args := m.Called(ctx, folder) + return args.Get(0).(int64), args.Error(1) +} + +func (m *MockResourceIndex) ListManagedObjects(ctx context.Context, req *ListManagedObjectsRequest) (*ListManagedObjectsResponse, error) { + args := m.Called(ctx, req) + return args.Get(0).(*ListManagedObjectsResponse), args.Error(1) +} + +var _ DocumentBuilder = &MockDocumentBuilder{} + +type MockDocumentBuilder struct { + mock.Mock +} + +func (m *MockDocumentBuilder) BuildDocument(ctx context.Context, key *ResourceKey, resourceVersion int64, value []byte) (*IndexableDocument, error) { + args := m.Called(ctx, key, resourceVersion, value) + if args.Get(0) == nil { + return nil, args.Error(1) + } + return args.Get(0).(*IndexableDocument), nil +} diff --git a/pkg/storage/unified/resource/server.go b/pkg/storage/unified/resource/server.go index b31d90558d7..3eb15c0e2d8 100644 --- a/pkg/storage/unified/resource/server.go +++ b/pkg/storage/unified/resource/server.go @@ -158,6 +158,9 @@ type SearchOptions struct { // Skip building index on startup for small indexes InitMinCount int + + // Channel to watch for index events (for testing) + IndexEventsChan chan *IndexEvent } type ResourceServerOptions struct { diff --git a/pkg/storage/unified/search/bleve.go b/pkg/storage/unified/search/bleve.go index 17f8fbeba23..83fa1d9fbb3 100644 --- a/pkg/storage/unified/search/bleve.go +++ b/pkg/storage/unified/search/bleve.go @@ -225,14 +225,12 @@ func (b *bleveBackend) BuildIndex(ctx context.Context, // Batch all the changes idx := &bleveIndex{ - key: key, - index: index, - batch: index.NewBatch(), - batchSize: b.opts.BatchSize, - fields: fields, - standard: resource.StandardSearchFields(), - features: b.features, - tracing: b.tracer, + key: key, + index: index, + fields: fields, + standard: resource.StandardSearchFields(), + features: b.features, + tracing: b.tracer, } idx.allFields, err = getAllFields(idx.standard, fields) @@ -245,12 +243,6 @@ func (b *bleveBackend) BuildIndex(ctx context.Context, if err != nil { return nil, err } - - // Flush the batch - err = idx.Flush() - if err != nil { - return nil, err - } } b.cacheMu.Lock() @@ -318,51 +310,33 @@ type bleveIndex struct { // The values returned with all allFields []*resource.ResourceTableColumnDefinition - - // only valid in single thread - batch *bleve.Batch - batchSize int // ??? not totally sure the units here - - features featuremgmt.FeatureToggles - tracing trace.Tracer + features featuremgmt.FeatureToggles + tracing trace.Tracer } -// Write implements resource.DocumentIndex. -func (b *bleveIndex) Write(v *resource.IndexableDocument) error { - v = v.UpdateCopyFields() +// BulkIndex implements resource.ResourceIndex. +func (b *bleveIndex) BulkIndex(req *resource.BulkIndexRequest) error { + if len(req.Items) == 0 { + return nil + } - // remove references (for now!) - v.References = nil - if b.batch != nil { - err := b.batch.Index(v.Key.SearchID(), v) - if err != nil { - return err + batch := b.index.NewBatch() + for _, item := range req.Items { + switch item.Action { + case resource.ActionIndex: + doc := item.Doc.UpdateCopyFields() + doc.References = nil // remove references (for now!) + + err := batch.Index(doc.Key.SearchID(), doc) + if err != nil { + return err + } + case resource.ActionDelete: + batch.Delete(item.Key.SearchID()) } - if b.batch.Size() > b.batchSize { - err = b.index.Batch(b.batch) - b.batch.Reset() // clear the batch - } - return err // nil } - return b.index.Index(v.Key.SearchID(), v) -} -// Delete implements resource.DocumentIndex. -func (b *bleveIndex) Delete(key *resource.ResourceKey) error { - if b.batch != nil { - return fmt.Errorf("unexpected delete while building batch") - } - return b.index.Delete(key.SearchID()) -} - -// Flush implements resource.DocumentIndex. -func (b *bleveIndex) Flush() (err error) { - if b.batch != nil { - err = b.index.Batch(b.batch) - b.batch.Reset() - b.batch = nil - } - return err + return b.index.Batch(batch) } func (b *bleveIndex) ListManagedObjects(ctx context.Context, req *resource.ListManagedObjectsRequest) (*resource.ListManagedObjectsResponse, error) { diff --git a/pkg/storage/unified/search/bleve_integration_test.go b/pkg/storage/unified/search/bleve_integration_test.go new file mode 100644 index 00000000000..6c28d560087 --- /dev/null +++ b/pkg/storage/unified/search/bleve_integration_test.go @@ -0,0 +1,56 @@ +package search + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/grafana/grafana/pkg/infra/tracing" + "github.com/grafana/grafana/pkg/services/featuremgmt" + "github.com/grafana/grafana/pkg/storage/unified/resource" + unitest "github.com/grafana/grafana/pkg/storage/unified/testing" +) + +func TestBleveSearchBackend(t *testing.T) { + if testing.Short() { + t.Skip("skipping integration test") + } + + // Run the search backend test suite + unitest.RunSearchBackendTest(t, func(ctx context.Context) resource.SearchBackend { + tempDir := t.TempDir() + + // Create a new bleve backend + backend, err := NewBleveBackend(BleveOptions{ + Root: tempDir, + FileThreshold: 5, + }, tracing.NewNoopTracerService(), featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchPermissionFiltering), nil) + require.NoError(t, err) + require.NotNil(t, backend) + + return backend + }, &unitest.TestOptions{ + NSPrefix: "bleve-test", + }) +} + +func TestSearchBackendBenchmark(t *testing.T) { + opts := &unitest.BenchmarkOptions{ + NumResources: 10000, + Concurrency: 1, // For now we only want to test the write throughput + NumNamespaces: 1, + NumGroups: 1, + NumResourceTypes: 1, + } + tempDir := t.TempDir() + + // Create a new bleve backend + backend, err := NewBleveBackend(BleveOptions{ + Root: tempDir, + }, tracing.NewNoopTracerService(), featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchPermissionFiltering), nil) + require.NoError(t, err) + require.NotNil(t, backend) + + unitest.BenchmarkSearchBackend(t, backend, opts) +} diff --git a/pkg/storage/unified/search/bleve_performance_test.go b/pkg/storage/unified/search/bleve_performance_test.go index 2d645235259..3598e579114 100644 --- a/pkg/storage/unified/search/bleve_performance_test.go +++ b/pkg/storage/unified/search/bleve_performance_test.go @@ -94,30 +94,50 @@ func newTestWriter(size int, batchSize int) IndexWriter { return func(index resource.ResourceIndex) (int64, error) { total := time.Now() start := time.Now() - for i := range size { + + // Create a batch of items + batch := make([]*resource.BulkIndexItem, 0, batchSize) + + for i := 0; i < size; i++ { name := fmt.Sprintf("name%d", i) - err := index.Write(&resource.IndexableDocument{ - RV: int64(i), - Name: name, - Key: &resource.ResourceKey{ - Name: name, - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + item := &resource.BulkIndexItem{ + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: int64(i), + Name: name, + Key: &resource.ResourceKey{ + Name: name, + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: name + "-title", }, - Title: name + "-title", - }) - if err != nil { - return 0, err } - // show progress for every batch - if i%batchSize == 0 && verbose { - fmt.Printf("Indexed %d documents\n", i) - end := time.Now() - fmt.Printf("Time taken for indexing batch: %s\n", end.Sub(start)) - start = time.Now() + + batch = append(batch, item) + + // When batch is full or this is the last item, process the batch + if len(batch) == batchSize || i == size-1 { + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: batch, + }) + if err != nil { + return 0, err + } + + if verbose { + end := time.Now() + fmt.Printf("Indexed %d documents\n", i+1) + fmt.Printf("Time taken for indexing batch: %s\n", end.Sub(start)) + start = time.Now() + } + + // Reset batch for next iteration + batch = make([]*resource.BulkIndexItem, 0, batchSize) } } + end := time.Now() logVerbose(fmt.Sprintf("Indexed %d documents in %s", size, end.Sub(total))) return 0, nil diff --git a/pkg/storage/unified/search/bleve_search_test.go b/pkg/storage/unified/search/bleve_search_test.go index 0d1598b3caa..ee1ee12f5fc 100644 --- a/pkg/storage/unified/search/bleve_search_test.go +++ b/pkg/storage/unified/search/bleve_search_test.go @@ -30,28 +30,37 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("when query is empty, sort documents by title instead of search score", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "name1", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "name1", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "bbb", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name2", + Key: &resource.ResourceKey{ + Name: "name2", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "aaa", + }, + }, }, - Title: "bbb", - }) - require.NoError(t, err) - err = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name2", - Key: &resource.ResourceKey{ - Name: "name2", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "aaa", }) require.NoError(t, err) @@ -65,28 +74,37 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("will boost phrase match query over match query results", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "name1", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "name1", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "I want to say a hello", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name2", + Key: &resource.ResourceKey{ + Name: "name2", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "we want hello", + }, + }, }, - Title: "I want to say a hello", - }) - require.NoError(t, err) - err = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name2", - Key: &resource.ResourceKey{ - Name: "name2", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "we want hello", }) require.NoError(t, err) @@ -100,28 +118,37 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("will prioritize matches", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "name1", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "name1", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "Asserts Dashboards", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name2", + Key: &resource.ResourceKey{ + Name: "name2", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "New dashboard 10", + }, + }, }, - Title: "Asserts Dashboards", - }) - require.NoError(t, err) - err = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name2", - Key: &resource.ResourceKey{ - Name: "name2", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "New dashboard 10", }) require.NoError(t, err) @@ -134,28 +161,37 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("will boost exact match query over match phrase query results", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "name1", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "name1", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "we want hello pls", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name2", + Key: &resource.ResourceKey{ + Name: "name2", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "we want hello", + }, + }, }, - Title: "we want hello pls", - }) - require.NoError(t, err) - err = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name2", - Key: &resource.ResourceKey{ - Name: "name2", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "we want hello", }) require.NoError(t, err) @@ -169,16 +205,23 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("title with numbers will match document", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "aaa", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "aaa", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "A123456", + }, + }, }, - Title: "A123456", }) require.NoError(t, err) @@ -203,29 +246,37 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("title will match escaped characters", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "aaa", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "aaa", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "what\"s up", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 2, + Name: "name2", + Key: &resource.ResourceKey{ + Name: "name2", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "what\"s that", + }, + }, }, - Title: "what\"s up", - }) - require.NoError(t, err) - - err = index.Write(&resource.IndexableDocument{ - RV: 2, - Name: "name2", - Key: &resource.ResourceKey{ - Name: "name2", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "what\"s that", }) require.NoError(t, err) @@ -242,16 +293,23 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("title search will match document", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "aaa", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "aaa", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "I want to say a wonderfully Hello to the WORLD! Hello-world", + }, + }, }, - Title: "I want to say a wonderfully Hello to the WORLD! Hello-world", }) require.NoError(t, err) @@ -300,40 +358,51 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("title search will NOT match documents", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "name1", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "name1", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "I want to say a wonderful Hello to the WORLD! Hello-world", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name2", + Key: &resource.ResourceKey{ + Name: "name2", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "A0456", + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name3", + Key: &resource.ResourceKey{ + Name: "name3", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "mash-A02382-10", + }, + }, }, - Title: "I want to say a wonderful Hello to the WORLD! Hello-world", - }) - require.NoError(t, err) - err = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name2", - Key: &resource.ResourceKey{ - Name: "name2", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "A0456", - }) - require.NoError(t, err) - err = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name3", - Key: &resource.ResourceKey{ - Name: "name3", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, - }, - Title: "mash-A02382-10", }) require.NoError(t, err) @@ -358,35 +427,47 @@ func TestCanSearchByTitle(t *testing.T) { t.Run("title search with character will match one document", func(t *testing.T) { index, _ := newTestDashboardsIndex(t, threshold, 2, 2, noop) - err := index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "name1", - Key: &resource.ResourceKey{ - Name: "aaa", - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "name1", + Key: &resource.ResourceKey{ + Name: "aaa", + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: "foo", + }, + }, }, - Title: "foo", }) require.NoError(t, err) for i, v := range search.TermCharacters { - err = index.Write(&resource.IndexableDocument{ - RV: int64(i), - Name: fmt.Sprintf("name%d", i), - Key: &resource.ResourceKey{ - Name: fmt.Sprintf("name%d", i), - Namespace: key.Namespace, - Group: key.Group, - Resource: key.Resource, + err = index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: int64(i), + Name: fmt.Sprintf("name%d", i), + Key: &resource.ResourceKey{ + Name: fmt.Sprintf("name%d", i), + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, + Title: fmt.Sprintf(`test foo%d%sbar`, i, v), + }, + }, }, - Title: fmt.Sprintf(`test foo%d%sbar`, i, v), }) require.NoError(t, err) - } - for i, v := range search.TermCharacters { title := fmt.Sprintf(`test foo%d%sbar`, i, v) query := newQueryByTitle(title) res, err := index.Search(context.Background(), nil, query, nil) diff --git a/pkg/storage/unified/search/bleve_test.go b/pkg/storage/unified/search/bleve_test.go index 0ac8d9c0c5c..f3a1473c2f0 100644 --- a/pkg/storage/unified/search/bleve_test.go +++ b/pkg/storage/unified/search/bleve_test.go @@ -65,91 +65,107 @@ func TestBleveBackend(t *testing.T) { Group: key.Group, Resource: key.Resource, }, 2, rv, info.Fields, func(index resource.ResourceIndex) (int64, error) { - _ = index.Write(&resource.IndexableDocument{ - RV: 1, - Name: "aaa", - Key: &resource.ResourceKey{ - Name: "aaa", - Namespace: "ns", - Group: "dashboard.grafana.app", - Resource: "dashboards", - }, - Title: "aaa (dash)", - Folder: "xxx", - Fields: map[string]any{ - DASHBOARD_PANEL_TYPES: []string{"timeseries", "table"}, - DASHBOARD_ERRORS_TODAY: 25, - DASHBOARD_VIEWS_LAST_1_DAYS: 50, - }, - Labels: map[string]string{ - utils.LabelKeyDeprecatedInternalID: "10", // nolint:staticcheck - }, - Tags: []string{"aa", "bb"}, - Manager: &utils.ManagerProperties{ - Kind: utils.ManagerKindRepo, - Identity: "repo-1", - }, - Source: &utils.SourceProperties{ - Path: "path/to/aaa.json", - Checksum: "xyz", - TimestampMillis: 1609462800000, // 2021 - }, - }) - _ = index.Write(&resource.IndexableDocument{ - RV: 2, - Name: "bbb", - Key: &resource.ResourceKey{ - Name: "bbb", - Namespace: "ns", - Group: "dashboard.grafana.app", - Resource: "dashboards", - }, - Title: "bbb (dash)", - Folder: "xxx", - Fields: map[string]any{ - DASHBOARD_PANEL_TYPES: []string{"timeseries"}, - DASHBOARD_ERRORS_TODAY: 40, - DASHBOARD_VIEWS_LAST_1_DAYS: 100, - }, - Tags: []string{"aa"}, - Labels: map[string]string{ - "region": "east", - utils.LabelKeyDeprecatedInternalID: "11", // nolint:staticcheck - }, - Manager: &utils.ManagerProperties{ - Kind: utils.ManagerKindRepo, - Identity: "repo-1", - }, - Source: &utils.SourceProperties{ - Path: "path/to/bbb.json", - Checksum: "hijk", - TimestampMillis: 1640998800000, // 2022 - }, - }) - _ = index.Write(&resource.IndexableDocument{ - RV: 3, - Key: &resource.ResourceKey{ - Name: "ccc", - Namespace: "ns", - Group: "dashboard.grafana.app", - Resource: "dashboards", - }, - Name: "ccc", - Title: "ccc (dash)", - Folder: "zzz", - Manager: &utils.ManagerProperties{ - Kind: utils.ManagerKindRepo, - Identity: "repo2", - }, - Source: &utils.SourceProperties{ - Path: "path/in/repo2.yaml", - }, - Fields: map[string]any{}, - Tags: []string{"aa"}, - Labels: map[string]string{ - "region": "west", + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Name: "aaa", + Key: &resource.ResourceKey{ + Name: "aaa", + Namespace: "ns", + Group: "dashboard.grafana.app", + Resource: "dashboards", + }, + Title: "aaa (dash)", + Folder: "xxx", + Fields: map[string]any{ + DASHBOARD_PANEL_TYPES: []string{"timeseries", "table"}, + DASHBOARD_ERRORS_TODAY: 25, + DASHBOARD_VIEWS_LAST_1_DAYS: 50, + }, + Labels: map[string]string{ + utils.LabelKeyDeprecatedInternalID: "10", // nolint:staticcheck + }, + Tags: []string{"aa", "bb"}, + Manager: &utils.ManagerProperties{ + Kind: utils.ManagerKindRepo, + Identity: "repo-1", + }, + Source: &utils.SourceProperties{ + Path: "path/to/aaa.json", + Checksum: "xyz", + TimestampMillis: 1609462800000, // 2021 + }, + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 2, + Name: "bbb", + Key: &resource.ResourceKey{ + Name: "bbb", + Namespace: "ns", + Group: "dashboard.grafana.app", + Resource: "dashboards", + }, + Title: "bbb (dash)", + Folder: "xxx", + Fields: map[string]any{ + DASHBOARD_PANEL_TYPES: []string{"timeseries"}, + DASHBOARD_ERRORS_TODAY: 40, + DASHBOARD_VIEWS_LAST_1_DAYS: 100, + }, + Tags: []string{"aa"}, + Labels: map[string]string{ + "region": "east", + utils.LabelKeyDeprecatedInternalID: "11", // nolint:staticcheck + }, + Manager: &utils.ManagerProperties{ + Kind: utils.ManagerKindRepo, + Identity: "repo-1", + }, + Source: &utils.SourceProperties{ + Path: "path/to/bbb.json", + Checksum: "hijk", + TimestampMillis: 1640998800000, // 2022 + }, + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 3, + Key: &resource.ResourceKey{ + Name: "ccc", + Namespace: "ns", + Group: "dashboard.grafana.app", + Resource: "dashboards", + }, + Name: "ccc", + Title: "ccc (dash)", + Folder: "zzz", + Manager: &utils.ManagerProperties{ + Kind: utils.ManagerKindRepo, + Identity: "repo2", + }, + Source: &utils.SourceProperties{ + Path: "path/in/repo2.yaml", + }, + Fields: map[string]any{}, + Tags: []string{"aa"}, + Labels: map[string]string{ + "region": "west", + }, + }, + }, }, }) + if err != nil { + return 0, err + } return rv, nil }) require.NoError(t, err) @@ -330,42 +346,55 @@ func TestBleveBackend(t *testing.T) { Group: key.Group, Resource: key.Resource, }, 2, rv, fields, func(index resource.ResourceIndex) (int64, error) { - _ = index.Write(&resource.IndexableDocument{ - RV: 1, - Key: &resource.ResourceKey{ - Name: "zzz", - Namespace: "ns", - Group: "folder.grafana.app", - Resource: "folders", - }, - Title: "zzz (folder)", - Manager: &utils.ManagerProperties{ - Kind: utils.ManagerKindRepo, - Identity: "repo-1", - }, - Source: &utils.SourceProperties{ - Path: "path/to/folder.json", - Checksum: "xxxx", - TimestampMillis: 300, - }, - Labels: map[string]string{ - utils.LabelKeyDeprecatedInternalID: "123", - }, - }) - _ = index.Write(&resource.IndexableDocument{ - RV: 2, - Key: &resource.ResourceKey{ - Name: "yyy", - Namespace: "ns", - Group: "folder.grafana.app", - Resource: "folders", - }, - Title: "yyy (folder)", - Labels: map[string]string{ - "region": "west", - utils.LabelKeyDeprecatedInternalID: "321", + err := index.BulkIndex(&resource.BulkIndexRequest{ + Items: []*resource.BulkIndexItem{ + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 1, + Key: &resource.ResourceKey{ + Name: "zzz", + Namespace: "ns", + Group: "folder.grafana.app", + Resource: "folders", + }, + Title: "zzz (folder)", + Manager: &utils.ManagerProperties{ + Kind: utils.ManagerKindRepo, + Identity: "repo-1", + }, + Source: &utils.SourceProperties{ + Path: "path/to/folder.json", + Checksum: "xxxx", + TimestampMillis: 300, + }, + Labels: map[string]string{ + utils.LabelKeyDeprecatedInternalID: "123", + }, + }, + }, + { + Action: resource.ActionIndex, + Doc: &resource.IndexableDocument{ + RV: 2, + Key: &resource.ResourceKey{ + Name: "yyy", + Namespace: "ns", + Group: "folder.grafana.app", + Resource: "folders", + }, + Title: "yyy (folder)", + Labels: map[string]string{ + "region": "west", + utils.LabelKeyDeprecatedInternalID: "321", + }, + }, + }, }, }) + if err != nil { + return 0, err + } return rv, nil }) require.NoError(t, err) diff --git a/pkg/storage/unified/sql/test/benchmark_test.go b/pkg/storage/unified/sql/test/benchmark_test.go index b91bcb003ab..b29975fc61d 100644 --- a/pkg/storage/unified/sql/test/benchmark_test.go +++ b/pkg/storage/unified/sql/test/benchmark_test.go @@ -2,20 +2,24 @@ package test import ( "context" + "os" "testing" "time" - infraDB "github.com/grafana/grafana/pkg/infra/db" + "github.com/grafana/grafana/pkg/infra/db" + "github.com/grafana/grafana/pkg/infra/tracing" + "github.com/grafana/grafana/pkg/services/featuremgmt" "github.com/grafana/grafana/pkg/setting" "github.com/grafana/grafana/pkg/storage/unified/resource" + "github.com/grafana/grafana/pkg/storage/unified/search" "github.com/grafana/grafana/pkg/storage/unified/sql" "github.com/grafana/grafana/pkg/storage/unified/sql/db/dbimpl" test "github.com/grafana/grafana/pkg/storage/unified/testing" "github.com/stretchr/testify/require" ) -func newTestBackend(b *testing.B) resource.StorageBackend { - dbstore := infraDB.InitTestDB(b) +func newTestBackend(b testing.TB) resource.StorageBackend { + dbstore := db.InitTestDB(b) eDB, err := dbimpl.ProvideResourceDB(dbstore, setting.NewCfg(), nil) require.NoError(b, err) require.NotNil(b, eDB) @@ -32,10 +36,62 @@ func newTestBackend(b *testing.B) resource.StorageBackend { return backend } -func BenchmarkSQLStorageBackend(b *testing.B) { +func TestIntegrationBenchmarkSQLStorageBackend(t *testing.T) { + if testing.Short() { + t.Skip("skipping integration test in short mode") + } + if db.IsTestDBSpanner() { + t.Skip("Skipping benchmark on Spanner") + } opts := test.DefaultBenchmarkOptions() - if infraDB.IsTestDbSQLite() { + if db.IsTestDbSQLite() { opts.Concurrency = 1 // to avoid SQLite database is locked error } - test.BenchmarkStorageBackend(b, newTestBackend(b), opts) + test.BenchmarkStorageBackend(t, newTestBackend(t), opts) +} + +func TestIntegrationBenchmarkResourceServer(t *testing.T) { + if testing.Short() { + t.Skip("skipping integration test in short mode") + } + if db.IsTestDBSpanner() { + t.Skip("Skipping benchmark on Spanner") + } + + ctx := context.Background() + opts := &test.BenchmarkOptions{ + NumResources: 1000, + Concurrency: 10, + NumNamespaces: 1, + NumGroups: 1, + NumResourceTypes: 1, + } + tempDir := t.TempDir() + t.Cleanup(func() { + _ = os.RemoveAll(tempDir) + }) + // Create a new bleve backend + search, err := search.NewBleveBackend(search.BleveOptions{ + Root: tempDir, + }, tracing.NewNoopTracerService(), featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchPermissionFiltering), nil) + require.NoError(t, err) + require.NotNil(t, search) + + // Create a new resource backend + dbstore := db.InitTestDB(t) + eDB, err := dbimpl.ProvideResourceDB(dbstore, setting.NewCfg(), nil) + require.NoError(t, err) + require.NotNil(t, eDB) + + storage, err := sql.NewBackend(sql.BackendOptions{ + DBProvider: eDB, + IsHA: false, + }) + require.NoError(t, err) + require.NotNil(t, storage) + + err = storage.Init(ctx) + require.NoError(t, err) + + test.BenchmarkIndexServer(t, ctx, storage, search, opts) } diff --git a/pkg/storage/unified/testing/benchmark.go b/pkg/storage/unified/testing/benchmark.go index 2f81f1a077b..11c577625ce 100644 --- a/pkg/storage/unified/testing/benchmark.go +++ b/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 +} diff --git a/pkg/storage/unified/testing/search_backend.go b/pkg/storage/unified/testing/search_backend.go new file mode 100644 index 00000000000..edb03073273 --- /dev/null +++ b/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 + }) +}