fix(embedding): not-found metric classification + backfill pagination (#91)

* fix: record embedding worker not-found as skipped, not failed_final

When GetFeedbackRecord returns not-found at the start of Work — a benign race where
the record was deleted or its tenant purged between enqueue and processing — the
worker recorded terminal-failure metrics (a get_record_failed worker error plus a
failed_final outcome) before the not-found short-circuit. That inflates the failure
metrics and can trip false alerts, and it was inconsistent with the same worker's
not-found-on-write path, which already records skipped.

Check not-found first and record it as skipped (no worker error); record the failure
metrics only for genuine errors.

* perf: stream the embedding backfill in keyset pages

ListFeedbackRecordIDsForBackfill and BackfillEmbeddings materialized every eligible
record id in one slice; on a large deployment that is a memory spike when the CLI runs.
They now keyset-paginate (id > cursor LIMIT n), enqueuing page by page so the full set
is never held in memory at once.

Author: xernobyl

internal/repository/embeddings_repository.go

@@ -95,17 +95,24 @@ func (r *EmbeddingsRepository) DeleteByFeedbackRecordAndModel(
 	})
 }
 
-// ListFeedbackRecordIDsForBackfill returns IDs of feedback records that have non-empty value_text
-// and no row in embeddings for the given model (so they need an embedding for that model).
-func (r *EmbeddingsRepository) ListFeedbackRecordIDsForBackfill(ctx context.Context, model string) ([]uuid.UUID, error) {
+// ListFeedbackRecordIDsForBackfill returns one keyset page (fr.id > afterID, ordered by id,
+// at most limit rows) of feedback-record IDs that have non-empty value_text and no row in
+// embeddings for the given model (so they need an embedding for that model). Pass uuid.Nil
+// as afterID for the first page.
+func (r *EmbeddingsRepository) ListFeedbackRecordIDsForBackfill(
+	ctx context.Context, model string, afterID uuid.UUID, limit int,
+) ([]uuid.UUID, error) {
 	rows, err := r.db.Query(ctx, `
 		SELECT fr.id FROM feedback_records fr
 		WHERE fr.value_text IS NOT NULL AND trim(fr.value_text) != ''
+		  AND fr.id > $2
 		  AND NOT EXISTS (
 		    SELECT 1 FROM embeddings e
 		    WHERE e.feedback_record_id = fr.id AND e.model = $1
 		  )
-	`, model)
+		ORDER BY fr.id
+		LIMIT $3
+	`, model, afterID, limit)
 	if err != nil {
 		return nil, fmt.Errorf("list feedback record ids for backfill: %w", err)
 	}

internal/service/feedback_records_service.go

@@ -44,7 +44,9 @@ type FeedbackRecordsRepository interface {
 type EmbeddingsRepository interface {
 	Upsert(ctx context.Context, feedbackRecordID uuid.UUID, model string, embedding []float32) error
 	DeleteByFeedbackRecordAndModel(ctx context.Context, feedbackRecordID uuid.UUID, model string) error
-	ListFeedbackRecordIDsForBackfill(ctx context.Context, model string) ([]uuid.UUID, error)
+	ListFeedbackRecordIDsForBackfill(
+		ctx context.Context, model string, afterID uuid.UUID, limit int,
+	) ([]uuid.UUID, error)
 }
 
 // FeedbackRecordsService handles business logic for feedback records.
@@ -293,38 +295,58 @@ func (s *FeedbackRecordsService) SetEmbedding(
 	return nil
 }
 
+// embeddingBackfillPageSize bounds how many record ids the embedding backfill lists and
+// enqueues per keyset page, so a large deployment is never fully materialized in memory.
+const embeddingBackfillPageSize = 500
+
 // BackfillEmbeddings enqueues embedding jobs for the given model for all feedback records that have
 // non-empty value_text and no embedding row for that model (existing rows are replaced by upsert when the job runs).
-// Returns the number of jobs enqueued. Requires embeddingInserter and embeddingQueueName to be set.
+// It streams the records in keyset pages. Returns the number of jobs enqueued. Requires embeddingInserter
+// and embeddingQueueName to be set.
 func (s *FeedbackRecordsService) BackfillEmbeddings(ctx context.Context, model string) (int, error) {
 	if s.embeddingInserter == nil || s.embeddingQueueName == "" {
 		return 0, ErrEmbeddingBackfillNotConfigured
 	}
 
-	ids, err := s.embeddingsRepo.ListFeedbackRecordIDsForBackfill(ctx, model)
-	if err != nil {
-		return 0, fmt.Errorf("list ids for embedding backfill: %w", err)
-	}
-
 	opts := &river.InsertOpts{
 		Queue:       s.embeddingQueueName,
 		MaxAttempts: s.embeddingMaxAttempts,
 		UniqueOpts:  river.UniqueOpts{ByArgs: true, ByPeriod: uniqueByPeriodEmbedding},
 	}
 
 	enqueued := 0
+	afterID := uuid.Nil
 
-	for _, id := range ids {
-		_, err := s.embeddingInserter.Insert(ctx, FeedbackEmbeddingArgs{
-			FeedbackRecordID: id,
-			Model:            model,
-			ValueTextHash:    "backfill",
-		}, opts)
+	for {
+		ids, err := s.embeddingsRepo.ListFeedbackRecordIDsForBackfill(ctx, model, afterID, embeddingBackfillPageSize)
 		if err != nil {
-			return enqueued, fmt.Errorf("enqueue embedding job for %s: %w", id, err)
+			return enqueued, fmt.Errorf("list ids for embedding backfill: %w", err)
+		}
+
+		if len(ids) == 0 {
+			break
+		}
+
+		for _, id := range ids {
+			_, err := s.embeddingInserter.Insert(ctx, FeedbackEmbeddingArgs{
+				FeedbackRecordID: id,
+				Model:            model,
+				ValueTextHash:    "backfill",
+			}, opts)
+			if err != nil {
+				return enqueued, fmt.Errorf("enqueue embedding job for %s: %w", id, err)
+			}
+
+			enqueued++
 		}
 
-		enqueued++
+		// Advance the keyset cursor past the last id seen; the query excludes
+		// already-embedded records, so the cursor always moves forward.
+		afterID = ids[len(ids)-1]
+
+		if len(ids) < embeddingBackfillPageSize {
+			break
+		}
 	}
 
 	return enqueued, nil

internal/workers/feedback_embedding.go

@@ -64,6 +64,22 @@ func (w *FeedbackEmbeddingWorker) Work(ctx context.Context, job *river.Job[servi
 
 	record, err := w.embeddingService.GetFeedbackRecord(ctx, args.FeedbackRecordID)
 	if err != nil {
+		// Not-found means the record was deleted or its tenant purged between enqueue and
+		// now: a benign race, not a terminal failure. Record it as skipped (consistent with
+		// the not-found-on-write path) so it does not trip failure alerts.
+		if errors.Is(err, huberrors.ErrNotFound) {
+			if w.metrics != nil {
+				w.metrics.RecordEmbeddingOutcome(ctx, "skipped")
+				w.metrics.RecordEmbeddingDuration(ctx, time.Since(start), "skipped")
+			}
+
+			slog.Info("embedding: record gone before embed, skipping",
+				"feedback_record_id", args.FeedbackRecordID,
+			)
+
+			return nil
+		}
+
 		if w.metrics != nil {
 			w.metrics.RecordWorkerError(ctx, "get_record_failed")
 			w.metrics.RecordEmbeddingOutcome(ctx, "failed_final")
@@ -75,11 +91,6 @@ func (w *FeedbackEmbeddingWorker) Work(ctx context.Context, job *river.Job[servi
 			"error", err,
 		)
 
-		// Only suppress retries for not-found; transient DB/network errors should retry.
-		if errors.Is(err, huberrors.ErrNotFound) {
-			return nil
-		}
-
 		return fmt.Errorf("get feedback record: %w", err)
 	}
 

internal/workers/feedback_embedding_test.go

@@ -4,16 +4,43 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/riverqueue/river"
 	"github.com/riverqueue/river/rivertype"
 
 	"github.com/formbricks/hub/internal/huberrors"
 	"github.com/formbricks/hub/internal/models"
+	"github.com/formbricks/hub/internal/observability"
 	"github.com/formbricks/hub/internal/service"
 )
 
+// countingEmbeddingMetrics records outcome/worker-error counts for assertions.
+type countingEmbeddingMetrics struct {
+	outcomes  map[string]int
+	workerErr map[string]int
+}
+
+func newCountingEmbeddingMetrics() *countingEmbeddingMetrics {
+	return &countingEmbeddingMetrics{outcomes: map[string]int{}, workerErr: map[string]int{}}
+}
+
+func (m *countingEmbeddingMetrics) RecordJobsEnqueued(context.Context, int64)   {}
+func (m *countingEmbeddingMetrics) RecordProviderError(context.Context, string) {}
+
+func (m *countingEmbeddingMetrics) RecordEmbeddingOutcome(_ context.Context, status string) {
+	m.outcomes[status]++
+}
+
+func (m *countingEmbeddingMetrics) RecordWorkerError(_ context.Context, reason string) {
+	m.workerErr[reason]++
+}
+
+func (m *countingEmbeddingMetrics) RecordEmbeddingDuration(context.Context, time.Duration, string) {}
+
+var _ observability.EmbeddingMetrics = (*countingEmbeddingMetrics)(nil)
+
 type mockEmbeddingService struct {
 	record          *models.FeedbackRecord
 	getErr          error
@@ -69,6 +96,26 @@ func textRecord(valueText string) *models.FeedbackRecord {
 	return record
 }
 
+func TestFeedbackEmbeddingWorker_GetNotFoundRecordsSkipped(t *testing.T) {
+	metrics := newCountingEmbeddingMetrics()
+	svc := &mockEmbeddingService{getErr: huberrors.NewNotFoundError("feedback record", "gone")}
+	worker := NewFeedbackEmbeddingWorker(svc, &mockEmbeddingClient{}, "", metrics)
+
+	if err := worker.Work(context.Background(), embeddingJob()); err != nil {
+		t.Fatalf("Work() error = %v, want nil (not-found completes)", err)
+	}
+
+	// A not-found GET is a benign delete/purge race: record it as skipped, never
+	// failed_final (which would trip failure alerts) and not as a worker error.
+	if metrics.outcomes["skipped"] != 1 || metrics.outcomes["failed_final"] != 0 {
+		t.Fatalf("skipped=%d failed_final=%d, want 1/0", metrics.outcomes["skipped"], metrics.outcomes["failed_final"])
+	}
+
+	if metrics.workerErr["get_record_failed"] != 0 {
+		t.Fatalf("get_record_failed=%d, want 0 (not-found is not a worker error)", metrics.workerErr["get_record_failed"])
+	}
+}
+
 func TestFeedbackEmbeddingWorker_Work_SetEmbeddingConflict(t *testing.T) {
 	ctx := context.Background()
 

tests/embedding_backfill_test.go

@@ -0,0 +1,162 @@
+package tests
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/riverqueue/river"
+	"github.com/riverqueue/river/rivertype"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/formbricks/hub/internal/config"
+	"github.com/formbricks/hub/internal/models"
+	"github.com/formbricks/hub/internal/repository"
+	"github.com/formbricks/hub/internal/service"
+	"github.com/formbricks/hub/pkg/database"
+)
+
+// countingEmbeddingInserter records the FeedbackEmbeddingArgs jobs enqueued.
+type countingEmbeddingInserter struct {
+	ids []uuid.UUID
+}
+
+func (c *countingEmbeddingInserter) Insert(
+	_ context.Context, args river.JobArgs, _ *river.InsertOpts,
+) (*rivertype.JobInsertResult, error) {
+	if a, ok := args.(service.FeedbackEmbeddingArgs); ok {
+		c.ids = append(c.ids, a.FeedbackRecordID)
+	}
+
+	return &rivertype.JobInsertResult{}, nil
+}
+
+func embeddingBackfillRepos(t *testing.T) (*repository.FeedbackRecordsRepository, *repository.EmbeddingsRepository) {
+	t.Helper()
+
+	cfg, err := config.Load()
+	require.NoError(t, err)
+
+	db, err := database.NewPostgresPool(
+		context.Background(), cfg.Database.URL, database.WithPoolConfig(cfg.Database.PoolConfig()))
+	require.NoError(t, err)
+
+	t.Cleanup(db.Close)
+
+	return repository.NewFeedbackRecordsRepository(db), repository.NewEmbeddingsRepository(db)
+}
+
+// TestListFeedbackRecordIDsForBackfill_KeysetPagination locks the new keyset behavior of the
+// embedding backfill query: each page is bounded by the limit, ids are returned strictly
+// ascending with no id on two pages, and paging from uuid.Nil eventually returns every record
+// that needs an embedding. A fresh model is used so the created records are all eligible.
+func TestListFeedbackRecordIDsForBackfill_KeysetPagination(t *testing.T) {
+	ctx := context.Background()
+	feedbackRepo, embeddingsRepo := embeddingBackfillRepos(t)
+
+	model := "backfill-keyset-" + uuid.NewString()
+	tenant := uuid.NewString()
+
+	makeText := func(value string) uuid.UUID {
+		rec, err := feedbackRepo.Create(ctx, &models.CreateFeedbackRecordRequest{
+			SourceType:   "formbricks",
+			SubmissionID: uuid.NewString(),
+			TenantID:     tenant,
+			FieldID:      "q1",
+			FieldType:    models.FieldTypeText,
+			ValueText:    &value,
+		})
+		require.NoError(t, err)
+
+		return rec.ID
+	}
+
+	mine := map[uuid.UUID]bool{}
+	for _, value := range []string{"one", "two", "three", "four", "five"} {
+		mine[makeText(value)] = true
+	}
+
+	seen := map[uuid.UUID]bool{}
+	afterID := uuid.Nil
+
+	var prev uuid.UUID
+
+	hasPrev := false
+
+	for {
+		page, err := embeddingsRepo.ListFeedbackRecordIDsForBackfill(ctx, model, afterID, 2)
+		require.NoError(t, err)
+		require.LessOrEqual(t, len(page), 2, "LIMIT bounds the page size")
+
+		if len(page) == 0 {
+			break
+		}
+
+		for _, id := range page {
+			require.False(t, seen[id], "an id must not appear on two pages")
+			seen[id] = true
+
+			if hasPrev {
+				require.Negativef(t, bytes.Compare(prev[:], id[:]), "ids are returned strictly ascending")
+			}
+
+			prev = id
+			hasPrev = true
+		}
+
+		afterID = page[len(page)-1]
+
+		if len(page) < 2 {
+			break
+		}
+	}
+
+	for id := range mine {
+		assert.Truef(t, seen[id], "record %s needing an embedding must be returned across pages", id)
+	}
+}
+
+// TestBackfillEmbeddings_StreamsAllEligible drives the service backfill against the real DB
+// with a recording inserter and asserts it enqueues an embedding job for every record that
+// needs one (a fresh model makes the created records eligible).
+func TestBackfillEmbeddings_StreamsAllEligible(t *testing.T) {
+	ctx := context.Background()
+	feedbackRepo, embeddingsRepo := embeddingBackfillRepos(t)
+
+	model := "backfill-stream-" + uuid.NewString()
+	tenant := uuid.NewString()
+
+	makeText := func(value string) uuid.UUID {
+		rec, err := feedbackRepo.Create(ctx, &models.CreateFeedbackRecordRequest{
+			SourceType:   "formbricks",
+			SubmissionID: uuid.NewString(),
+			TenantID:     tenant,
+			FieldID:      "q1",
+			FieldType:    models.FieldTypeText,
+			ValueText:    &value,
+		})
+		require.NoError(t, err)
+
+		return rec.ID
+	}
+
+	mine := []uuid.UUID{makeText("one"), makeText("two"), makeText("three")}
+
+	inserter := &countingEmbeddingInserter{}
+	svc := service.NewFeedbackRecordsService(feedbackRepo, embeddingsRepo, model, nil, inserter, "embeddings", 3)
+
+	enqueued, err := svc.BackfillEmbeddings(ctx, model)
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, enqueued, len(mine))
+
+	got := map[uuid.UUID]bool{}
+	for _, id := range inserter.ids {
+		got[id] = true
+	}
+
+	for _, id := range mine {
+		assert.Truef(t, got[id], "record %s needing an embedding must be enqueued", id)
+	}
+}