[pkg/stanza/fileconsumer] Optimize indexed fingerprint matching (#47660)

This PR was implemented as part of #45203 but decided to split the
changes done to make them easier to review.

Updates
#27404

Changes fingerprint matching in `fileset`. Before this change, matching
was done with linear scans over the tracked files. After this change,
matching uses an internal index based on fingerprint length and
fingerprint content, so exact and prefix matches no longer depend on
scanning the full set every time.

As part of that implementation, `Fileset.Pop()` changes from removing
the first element to removing the last element. In other words, it
changes from FIFO behavior to LIFO behavior. This is an internal change
that makes removal O(1) and fits the indexed structure.

The benchmark was run before and after the change with `go test
./fileconsumer -run '^$' -bench BenchmarkPollManyFiles -benchmem
-count=5`.

  | Files watched | Before | After | Improvement |
  | --- | ---: | ---: | ---: |
  | 100 | 2.448 ms/op | 2.196 ms/op | 10.3% |
  | 500 | 12.27 ms/op | 10.77 ms/op | 12.3% |
  | 1000 | 28.20 ms/op | 21.91 ms/op | 22.3% |
  | 2000 | 64.59 ms/op | 44.90 ms/op | 30.5% |
  | 2500 | 87.65 ms/op | 57.22 ms/op | 34.7% |
  | 3000 | 112.98 ms/op | 69.61 ms/op | 38.4% |
  | Geomean | 28.55 ms/op | 21.27 ms/op | 25.5% |

---------

Signed-off-by: Israel Blancas <iblancasa@gmail.com>
Co-authored-by: Andrzej Stencel <andrzej.stencel@elastic.co>

Author: iblancasa

.chloggen/27404-filelog-indexed-fingerprint-matching.yaml

@@ -0,0 +1,23 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: enhancement
+
+# The name of the component, or a single word describing the area of concern, (e.g. receiver/filelog)
+component: receiver/file_log
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: "Improve polling performance when watching many files by indexing fingerprint matching."
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [27404]
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if the change is relevant to users of libraries or exported APIs.
+# Include 'developer' if the change is relevant only for developers of this repo.
+# Default: '[user]'
+change_logs: []

pkg/stanza/fileconsumer/benchmark_test.go

@@ -342,6 +342,50 @@ func BenchmarkConsumeFiles(b *testing.B) {
 	}
 }
 
+func BenchmarkPollManyFiles(b *testing.B) {
+	fileCounts := []int{100, 500, 1000, 2000, 2500, 3000}
+	for _, numFiles := range fileCounts {
+		b.Run(fmt.Sprintf("Files_%d", numFiles), func(b *testing.B) {
+			rootDir := b.TempDir()
+
+			for i := range numFiles {
+				path := filepath.Join(rootDir, fmt.Sprintf("file%04d.log", i))
+				f := filetest.OpenFile(b, path)
+				_, err := f.WriteString(path + "\n")
+				require.NoError(b, err)
+				require.NoError(b, f.Close())
+			}
+
+			cfg := NewConfig()
+			cfg.Include = []string{filepath.Join(rootDir, "*.log")}
+			cfg.StartAt = "beginning"
+			cfg.MaxConcurrentFiles = numFiles * 2
+			cfg.PollInterval = time.Microsecond
+
+			set := componenttest.NewNopTelemetrySettings()
+			callback := func(context.Context, [][]byte, map[string]any, int64, []int64) error {
+				return nil
+			}
+
+			op, err := cfg.Build(set, callback)
+			require.NoError(b, err)
+
+			ctx := b.Context()
+			persister := testutil.NewUnscopedMockPersister()
+			op.instantiateTracker(ctx, persister)
+			op.persister = persister
+
+			op.poll(ctx)
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for b.Loop() {
+				op.poll(ctx)
+			}
+		})
+	}
+}
+
 // BenchmarkFingerprintComparison benchmarks fingerprint comparison operations
 // This isolates the cost of fingerprint matching from file I/O
 func BenchmarkFingerprintComparison(b *testing.B) {
@@ -411,8 +455,8 @@ func BenchmarkFilesetMatch(b *testing.B) {
 				fs := fileset.New[*reader.Metadata](size)
 				fs.Add(metadatas...)
 
-				// Measure the Match operation
-				result := fs.Match(targetFp, fileset.Equal)
+				// Measure the indexed exact-match operation.
+				result := fs.MatchEqual(targetFp)
 				if result == nil {
 					b.Fatal("expected to find match")
 				}

pkg/stanza/fileconsumer/internal/archive/archive.go

@@ -93,7 +93,7 @@ func (a *archive) FindFiles(ctx context.Context, fps []*fingerprint.Fingerprint)
 				// we've already found a match for this index, continue
 				continue
 			}
-			if md := data.Match(fp, fileset.StartsWith); md != nil {
+			if md := data.MatchStartsWith(fp); md != nil {
 				// update the matched metada for the index
 				matchedMetadata[j] = md
 				archiveModified = true

pkg/stanza/fileconsumer/internal/fileset/fileset.go

@@ -22,12 +22,29 @@ type Matchable interface {
 	GetFingerprint() *fingerprint.Fingerprint
 }
 
+type bucket struct {
+	indices []int
+}
+
+type bucketPos struct {
+	bucket *bucket
+	idx    int
+}
+
 type Fileset[T Matchable] struct {
-	readers []T
+	readers       []T
+	positions     []bucketPos
+	buckets       map[int]map[string]*bucket
+	bucketLengths []int
 }
 
 func New[T Matchable](capacity int) *Fileset[T] {
-	return &Fileset[T]{readers: make([]T, 0, capacity)}
+	return &Fileset[T]{
+		readers:       make([]T, 0, capacity),
+		positions:     make([]bucketPos, 0, capacity),
+		buckets:       make(map[int]map[string]*bucket),
+		bucketLengths: make([]int, 0),
+	}
 }
 
 func (set *Fileset[T]) Len() int {
@@ -38,33 +55,188 @@ func (set *Fileset[T]) Get() []T {
 	return set.readers
 }
 
+func (set *Fileset[T]) Reset() {
+	var zero T
+	for i := range set.readers {
+		set.readers[i] = zero
+	}
+	set.readers = set.readers[:0]
+	for i := range set.positions {
+		set.positions[i] = bucketPos{}
+	}
+	set.positions = set.positions[:0]
+	clear(set.buckets)
+	set.bucketLengths = set.bucketLengths[:0]
+}
+
+func (set *Fileset[T]) Reindex() {
+	if len(set.readers) == 0 {
+		clear(set.buckets)
+		set.bucketLengths = set.bucketLengths[:0]
+		return
+	}
+
+	clear(set.buckets)
+	set.bucketLengths = set.bucketLengths[:0]
+
+	for idx, reader := range set.readers {
+		pos := bucketPos{}
+		if fp := reader.GetFingerprint(); fp != nil && fp.Len() > 0 {
+			pos = set.insertIntoBucket(idx, fp)
+		}
+		set.positions[idx] = pos
+	}
+}
+
 func (set *Fileset[T]) Pop() (T, error) {
-	// return first element from the array and remove it
 	var val T
 	if len(set.readers) == 0 {
 		return val, errFilesetEmpty
 	}
-	r := set.readers[0]
-	set.readers = slices.Delete(set.readers, 0, 1)
-	return r, nil
+	return set.removeAt(len(set.readers) - 1), nil
 }
 
 func (set *Fileset[T]) Add(readers ...T) {
-	// add open readers
-	set.readers = append(set.readers, readers...)
+	for _, reader := range readers {
+		set.readers = append(set.readers, reader)
+		idx := len(set.readers) - 1
+		pos := bucketPos{}
+		if fp := reader.GetFingerprint(); fp != nil && fp.Len() > 0 {
+			pos = set.insertIntoBucket(idx, fp)
+		}
+		set.positions = append(set.positions, pos)
+	}
+}
+
+func (set *Fileset[T]) MatchEqual(fp *fingerprint.Fingerprint) T {
+	var zero T
+	if fp == nil || fp.Len() == 0 {
+		return zero
+	}
+	bucketsByLength, ok := set.buckets[fp.Len()]
+	if !ok {
+		return zero
+	}
+	bucket := bucketsByLength[fp.Key()]
+	if bucket == nil {
+		return zero
+	}
+	for _, idx := range slices.Backward(bucket.indices) {
+		if fp.Equal(set.readers[idx].GetFingerprint()) {
+			return set.removeAt(idx)
+		}
+	}
+	return zero
+}
+
+func (set *Fileset[T]) MatchStartsWith(fp *fingerprint.Fingerprint) T {
+	var zero T
+	if fp == nil || fp.Len() == 0 {
+		return zero
+	}
+	fpKey := fp.Key()
+	for _, length := range set.bucketLengths {
+		if length > len(fpKey) {
+			continue
+		}
+		bucketsByLength, ok := set.buckets[length]
+		if !ok {
+			continue
+		}
+		bucket := bucketsByLength[fpKey[:length]]
+		if bucket == nil {
+			continue
+		}
+		for _, idx := range slices.Backward(bucket.indices) {
+			if fp.StartsWith(set.readers[idx].GetFingerprint()) {
+				return set.removeAt(idx)
+			}
+		}
+	}
+	return zero
 }
 
 func (set *Fileset[T]) Match(fp *fingerprint.Fingerprint, cmp func(a, b *fingerprint.Fingerprint) bool) T {
 	var val T
 	for idx, r := range set.readers {
 		if cmp(fp, r.GetFingerprint()) {
-			set.readers = append(set.readers[:idx], set.readers[idx+1:]...)
-			return r
+			return set.removeAt(idx)
 		}
 	}
 	return val
 }
 
+func (set *Fileset[T]) insertIntoBucket(idx int, fp *fingerprint.Fingerprint) bucketPos {
+	keyLen := fp.Len()
+	key := fp.Key()
+
+	bucketsByLength := set.buckets[keyLen]
+	if bucketsByLength == nil {
+		bucketsByLength = make(map[string]*bucket)
+		set.buckets[keyLen] = bucketsByLength
+		set.bucketLengths = append(set.bucketLengths, keyLen)
+		slices.SortFunc(set.bucketLengths, func(a, b int) int {
+			return b - a
+		})
+	}
+
+	b := bucketsByLength[key]
+	if b == nil {
+		b = &bucket{}
+		bucketsByLength[key] = b
+	}
+	b.indices = append(b.indices, idx)
+
+	return bucketPos{
+		bucket: b,
+		idx:    len(b.indices) - 1,
+	}
+}
+
+func (set *Fileset[T]) removeAt(idx int) T {
+	val := set.readers[idx]
+	set.removeFromBucket(idx)
+
+	lastIdx := len(set.readers) - 1
+	if idx != lastIdx {
+		set.readers[idx] = set.readers[lastIdx]
+		set.positions[idx] = set.positions[lastIdx]
+		set.updateBucketPosition(idx)
+	}
+
+	var zero T
+	set.readers[lastIdx] = zero
+	set.readers = set.readers[:lastIdx]
+	set.positions[lastIdx] = bucketPos{}
+	set.positions = set.positions[:lastIdx]
+
+	return val
+}
+
+func (set *Fileset[T]) removeFromBucket(idx int) {
+	pos := set.positions[idx]
+	if pos.bucket == nil {
+		return
+	}
+
+	last := len(pos.bucket.indices) - 1
+	swapIdx := pos.bucket.indices[last]
+	pos.bucket.indices[pos.idx] = swapIdx
+	pos.bucket.indices = pos.bucket.indices[:last]
+	if pos.idx < last {
+		set.positions[swapIdx].idx = pos.idx
+	}
+	set.positions[idx] = bucketPos{}
+}
+
+func (set *Fileset[T]) updateBucketPosition(idx int) {
+	pos := &set.positions[idx]
+	if pos.bucket == nil {
+		return
+	}
+	pos.bucket.indices[pos.idx] = idx
+}
+
 // comparators
 func StartsWith(a, b *fingerprint.Fingerprint) bool {
 	return a.StartsWith(b)