search: Handle index build errors gracefully (#108862)

* Close new index if we fail to build it.

* Respect context cancellation in getOrCreateIndex.

pkg/storage/unified/resource/search.go

@@ -647,7 +647,7 @@ func (s *searchSupport) getOrCreateIndex(ctx context.Context, key NamespacedReso
 		return idx, nil
 	}
 
-	idxInt, err, _ := s.buildIndex.Do(key.String(), func() (interface{}, error) {
+	ch := s.buildIndex.DoChan(key.String(), func() (interface{}, error) {
 		// Recheck if some other goroutine managed to build an index in the meantime.
 		// (That is, it finished running this function and stored the index into the cache)
 		idx, err := s.search.GetIndex(ctx, key)
@@ -681,10 +681,16 @@ func (s *searchSupport) getOrCreateIndex(ctx context.Context, key NamespacedReso
 		}
 		return idx, nil
 	})
-	if err != nil {
-		return nil, err
+
+	select {
+	case res := <-ch:
+		if res.Err != nil {
+			return nil, res.Err
+		}
+		return res.Val.(ResourceIndex), nil
+	case <-ctx.Done():
+		return nil, fmt.Errorf("failed to get index: %w", ctx.Err())
 	}
-	return idxInt.(ResourceIndex), nil
 }
 
 func (s *searchSupport) build(ctx context.Context, nsr NamespacedResource, size int64, rv int64) (ResourceIndex, int64, error) {

pkg/storage/unified/resource/search_test.go

@@ -332,3 +332,57 @@ func TestSearchGetOrCreateIndex(t *testing.T) {
 	require.Equal(t, int64(50), search.buildIndexCalls[0].size)
 	require.Equal(t, int64(11111111), search.buildIndexCalls[0].resourceVersion)
 }
+
+func TestSearchGetOrCreateIndexWithCancellation(t *testing.T) {
+	// Setup mock implementations
+	storage := &mockStorageBackend{
+		resourceStats: []ResourceStats{
+			{NamespacedResource: NamespacedResource{Namespace: "ns", Group: "group", Resource: "resource"}, Count: 50, ResourceVersion: 11111111},
+		},
+	}
+	search := &slowSearchBackend{
+		mockSearchBackend: mockSearchBackend{},
+	}
+	supplier := &TestDocumentBuilderSupplier{
+		GroupsResources: map[string]string{
+			"group": "resource",
+		},
+	}
+
+	// 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:  0,
+	}
+
+	support, err := newSearchSupport(opts, storage, nil, nil, noop.NewTracerProvider().Tracer("test"), nil, nil, nil)
+	require.NoError(t, err)
+	require.NotNil(t, support)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
+	defer cancel()
+
+	_, err = support.getOrCreateIndex(ctx, NamespacedResource{Namespace: "ns", Group: "group", Resource: "resource"})
+	// Make sure we get context deadline error
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+
+	// Wait until indexing is finished.
+	search.wg.Wait()
+
+	require.NotEmpty(t, search.buildIndexCalls)
+}
+
+type slowSearchBackend struct {
+	mockSearchBackend
+	wg sync.WaitGroup
+}
+
+func (m *slowSearchBackend) BuildIndex(ctx context.Context, key NamespacedResource, size int64, resourceVersion int64, fields SearchableDocumentFields, builder func(index ResourceIndex) (int64, error)) (ResourceIndex, error) {
+	m.wg.Add(1)
+	defer m.wg.Done()
+	time.Sleep(1 * time.Second)
+	return m.mockSearchBackend.BuildIndex(ctx, key, size, resourceVersion, fields, builder)
+}

pkg/storage/unified/search/bleve.go

@@ -202,22 +202,47 @@ func (b *bleveBackend) BuildIndex(
 	_, span := b.tracer.Start(ctx, tracingPrexfixBleve+"BuildIndex")
 	defer span.End()
 
-	var index bleve.Index
-	fileIndexName := "" // Name of the file-based index, or empty for in-memory indexes.
-
-	build := true
 	mapper, err := GetBleveMappings(fields)
 	if err != nil {
 		return nil, err
 	}
 
-	cachedIndex := b.getCachedIndex(key)
+	// Prepare fields before opening/creating indexes, so that we don't need to deal with closing them in case of errors.
+	standardSearchFields := resource.StandardSearchFields()
+	allFields, err := getAllFields(standardSearchFields, fields)
+	if err != nil {
+		return nil, err
+	}
 
 	logWithDetails := b.log.With("namespace", key.Namespace, "group", key.Group, "resource", key.Resource, "size", size, "rv", resourceVersion)
 
+	// Close the newly created/opened index by default.
+	closeIndex := true
+	// This function is added via defer after new index has been created/opened, to make sure we close it properly when needed.
+	// Whether index needs closing or not is controlled by closeIndex.
+	closeIndexOnExit := func(index bleve.Index, indexDir string) {
+		if !closeIndex {
+			return
+		}
+
+		if closeErr := index.Close(); closeErr != nil {
+			logWithDetails.Error("Failed to close index after index build failure", "err", closeErr)
+		}
+		if indexDir != "" {
+			if removeErr := os.RemoveAll(indexDir); removeErr != nil {
+				logWithDetails.Error("Failed to remove index directory after index build failure", "err", removeErr)
+			}
+		}
+	}
+
 	resourceDir := filepath.Join(b.opts.Root, cleanFileSegment(key.Namespace), cleanFileSegment(fmt.Sprintf("%s.%s", key.Resource, key.Group)))
 
+	var index bleve.Index
+	cachedIndex := b.getCachedIndex(key)
+	fileIndexName := "" // Name of the file-based index, or empty for in-memory indexes.
 	newIndexType := indexStorageMemory
+	build := true
+
 	if size > b.opts.FileThreshold {
 		newIndexType = indexStorageFile
 
@@ -231,6 +256,7 @@ func (b *bleveBackend) BuildIndex(
 		if index != nil {
 			build = false
 			logWithDetails.Debug("Existing index found on filesystem", "directory", filepath.Join(resourceDir, fileIndexName))
+			defer closeIndexOnExit(index, "") // Close index, but don't delete directory.
 		} else {
 			// Building index from scratch. Index name has a time component in it to be unique, but if
 			// we happen to create non-unique name, we bump the time and try again.
@@ -256,6 +282,7 @@ func (b *bleveBackend) BuildIndex(
 			}
 
 			logWithDetails.Info("Building index using filesystem", "directory", indexDir)
+			defer closeIndexOnExit(index, indexDir) // Close index, and delete new index directory.
 		}
 	} else {
 		index, err = bleve.NewMemOnly(mapper)
@@ -263,6 +290,7 @@ func (b *bleveBackend) BuildIndex(
 			return nil, fmt.Errorf("error creating new in-memory bleve index: %w", err)
 		}
 		logWithDetails.Info("Building index using memory")
+		defer closeIndexOnExit(index, "") // Close index, don't cleanup directory.
 	}
 
 	// Batch all the changes
@@ -271,21 +299,18 @@ func (b *bleveBackend) BuildIndex(
 		index:        index,
 		indexStorage: newIndexType,
 		fields:       fields,
-		standard:     resource.StandardSearchFields(),
+		allFields:    allFields,
+		standard:     standardSearchFields,
 		features:     b.features,
 		tracing:      b.tracer,
 	}
 
-	idx.allFields, err = getAllFields(idx.standard, fields)
-	if err != nil {
-		return nil, err
-	}
-
 	if build {
 		start := time.Now()
 		_, err = builder(idx)
 		if err != nil {
-			return nil, err
+			logWithDetails.Error("Failed to build index", "err", err)
+			return nil, fmt.Errorf("failed to build index: %w", err)
 		}
 		elapsed := time.Since(start)
 		logWithDetails.Info("Finished building index", "elapsed", elapsed)
@@ -303,6 +328,9 @@ func (b *bleveBackend) BuildIndex(
 		logWithDetails.Info("Storing index in cache", "key", key, "expiration", idx.expiration)
 	}
 
+	// We're storing index in the cache, so we can't close it.
+	closeIndex = false
+
 	b.cacheMx.Lock()
 	prev := b.cache[key]
 	b.cache[key] = idx

pkg/storage/unified/search/bleve_test.go

@@ -1026,3 +1026,36 @@ func TestCleanOldIndexes(t *testing.T) {
 		require.Len(t, files, 0)
 	})
 }
+
+func TestBleveIndexWithFailures(t *testing.T) {
+	t.Run("in-memory index", func(t *testing.T) {
+		testBleveIndexWithFailures(t, false)
+	})
+	t.Run("file-based index", func(t *testing.T) {
+		testBleveIndexWithFailures(t, true)
+	})
+}
+
+func testBleveIndexWithFailures(t *testing.T, fileBased bool) {
+	backend, _ := setupBleveBackend(t, 5, time.Nanosecond, "")
+
+	ns := resource.NamespacedResource{
+		Namespace: "test",
+		Group:     "group",
+		Resource:  "resource",
+	}
+
+	size := int64(1)
+	if fileBased {
+		// size=100 is above FileThreshold (5), make it a file-based index.
+		size = 100
+	}
+	_, err := backend.BuildIndex(context.Background(), ns, size, 100, nil, func(index resource.ResourceIndex) (int64, error) {
+		return 0, fmt.Errorf("fail")
+	})
+	require.Error(t, err)
+
+	// Even though previous build of the index failed, new building of the index should work.
+	_, err = backend.BuildIndex(context.Background(), ns, size, 100, nil, indexTestDocs(ns, int(size)))
+	require.NoError(t, err)
+}