search: Handle index build errors gracefully (#108862)

* Close new index if we fail to build it.

* Respect context cancellation in getOrCreateIndex.

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