Embeddings: add progress updates during indexing (#53202)

This adds incremental stats reporting to our embedding indexing jobs and
displays it in the UI while the job is processing.

Author: camdencheek

client/web/src/enterprise/site-admin/cody/RepoEmbeddingJobNode.tsx

@@ -36,6 +36,7 @@ export const RepoEmbeddingJobNode: FC<RepoEmbeddingJobNodeProps> = ({
     queuedAt,
     startedAt,
     failureMessage,
+    stats,
     onCancel,
 }) => (
     <li className="list-group-item p-2">
@@ -60,6 +61,7 @@ export const RepoEmbeddingJobNode: FC<RepoEmbeddingJobNodeProps> = ({
                             queuedAt={queuedAt}
                             startedAt={startedAt}
                             failureMessage={failureMessage}
+                            stats={stats}
                         />
                     </div>
                 </div>
@@ -84,14 +86,25 @@ export const RepoEmbeddingJobNode: FC<RepoEmbeddingJobNodeProps> = ({
 )
 
 const RepoEmbeddingJobExecutionInfo: FC<
-    Pick<RepoEmbeddingJobFields, 'state' | 'cancel' | 'finishedAt' | 'failureMessage' | 'queuedAt' | 'startedAt'>
-> = ({ state, cancel, finishedAt, queuedAt, startedAt, failureMessage }) => {
+    Pick<
+        RepoEmbeddingJobFields,
+        'state' | 'cancel' | 'finishedAt' | 'failureMessage' | 'queuedAt' | 'startedAt' | 'stats'
+    >
+> = ({ state, cancel, finishedAt, queuedAt, startedAt, failureMessage, stats }) => {
     const [isPopoverOpen, setIsPopoverOpen] = useState(false)
+    const estimatedFinish = calculateEstimatedFinish(
+        startedAt,
+        stats.filesScheduled,
+        stats.filesEmbedded,
+        stats.filesSkipped
+    )
+
     return (
         <>
             {state === RepoEmbeddingJobState.COMPLETED && finishedAt && (
                 <small>
-                    Completed <Timestamp date={finishedAt} />
+                    Completed embedding {stats.filesEmbedded} files ({stats.filesSkipped} skipped){' '}
+                    <Timestamp date={finishedAt} />
                 </small>
             )}
             {state === RepoEmbeddingJobState.CANCELED && finishedAt && (
@@ -114,6 +127,11 @@ const RepoEmbeddingJobExecutionInfo: FC<
                 <small>
                     {cancel ? (
                         'Cancelling ...'
+                    ) : estimatedFinish ? (
+                        <>
+                            Expected to finish <Timestamp date={estimatedFinish} /> (
+                            {stats.filesSkipped + stats.filesEmbedded}/{stats.filesScheduled} files)
+                        </>
                     ) : (
                         <>
                             Started processing <Timestamp date={startedAt} />
@@ -140,6 +158,31 @@ const RepoEmbeddingJobExecutionInfo: FC<
     )
 }
 
+function calculateEstimatedFinish(
+    startedAt: string | null,
+    filesScheduled: number,
+    filesEmbedded: number,
+    filesSkipped: number,
+    now?: number
+): Date | null {
+    const currentTime = now ?? Date.now()
+    if (!startedAt) {
+        return null
+    }
+    const startTime = Date.parse(startedAt)
+    if (filesScheduled === 0) {
+        // There is a period between when the job starts processing and when
+        // we know how many files need to be processed. In the case where
+        // we do not have an update with the number of files scheduled,
+        // we cannot calculate a meaningful ETA.
+        return null
+    }
+    const proportionFinished = (filesEmbedded + filesSkipped) / filesScheduled
+    const timeElapsed = currentTime - startTime
+    const estimatedTotalTime = timeElapsed / proportionFinished
+    return new Date(startTime + estimatedTotalTime)
+}
+
 function getRepoEmbeddingJobStateBadgeVariant(state: RepoEmbeddingJobState): BadgeVariantType {
     switch (state) {
         case RepoEmbeddingJobState.COMPLETED:

client/web/src/enterprise/site-admin/cody/backend.ts

@@ -33,6 +33,11 @@ const REPO_EMBEDDING_JOB_FRAGMENT = gql`
             oid
             abbreviatedOID
         }
+        stats {
+            filesScheduled
+            filesEmbedded
+            filesSkipped
+        }
     }
 `
 
@@ -64,6 +69,9 @@ export const useRepoEmbeddingJobsConnection = (
             const { repoEmbeddingJobs } = dataOrThrowErrors(result)
             return repoEmbeddingJobs
         },
+        options: {
+            pollInterval: 5000,
+        },
     })
 
 export const SCHEDULE_REPO_EMBEDDING_JOBS = gql`

cmd/frontend/graphqlbackend/embeddings.go

@@ -81,4 +81,11 @@ type RepoEmbeddingJobResolver interface {
 	Cancel() bool
 	Repo(ctx context.Context) (*RepositoryResolver, error)
 	Revision(ctx context.Context) (*GitCommitResolver, error)
+	Stats(context.Context) (RepoEmbeddingJobStatsResolver, error)
+}
+
+type RepoEmbeddingJobStatsResolver interface {
+	FilesEmbedded() int32
+	FilesScheduled() int32
+	FilesSkipped() int32
 }

cmd/frontend/graphqlbackend/embeddings.graphql

@@ -217,6 +217,35 @@ type RepoEmbeddingJob implements Node {
     The revision at which the repo was embedded.
     """
     revision: GitCommit
+
+    """
+    Statistics about the embeddings index job.
+    This will be updated periodically while the embeddings job is processing.
+    """
+    stats: RepoEmbeddingsStats!
+}
+
+"""
+Statistics about an embeddings index job.
+This will be updated periodically while the embeddings job is processing.
+"""
+type RepoEmbeddingsStats {
+    """
+    The number of files scheduled to be embedded.
+    """
+    filesScheduled: Int!
+
+    """
+    The number of files we generated embeddings for.
+    This will be updated periodically while the embeddings job is processing.
+    """
+    filesEmbedded: Int!
+
+    """
+    The number of files skipped.
+    This will be updated periodically while the embeddings job is processing.
+    """
+    filesSkipped: Int!
 }
 
 """

enterprise/cmd/frontend/internal/embeddings/resolvers/repo_embedding_jobs.go

@@ -57,7 +57,11 @@ func (s *repoEmbeddingJobsConnectionStore) ComputeNodes(ctx context.Context, arg
 	}
 	resolvers := make([]graphqlbackend.RepoEmbeddingJobResolver, 0, len(jobs))
 	for _, job := range jobs {
-		resolvers = append(resolvers, &repoEmbeddingJobResolver{db: s.db, gitserverClient: s.gitserverClient, job: job})
+		resolvers = append(resolvers, &repoEmbeddingJobResolver{
+			db:              s.db,
+			gitserverClient: s.gitserverClient,
+			job:             job,
+		})
 	}
 	return resolvers, nil
 }
@@ -149,6 +153,15 @@ func (r *repoEmbeddingJobResolver) Cancel() bool {
 	return r.job.Cancel
 }
 
+func (r *repoEmbeddingJobResolver) Stats(ctx context.Context) (graphqlbackend.RepoEmbeddingJobStatsResolver, error) {
+	store := repobg.NewRepoEmbeddingJobsStore(r.db)
+	stats, err := store.GetRepoEmbeddingJobStats(ctx, r.job.ID)
+	if err != nil {
+		return nil, err
+	}
+	return &repoEmbeddingJobStatsResolver{stats}, nil
+}
+
 func (r *repoEmbeddingJobResolver) compute(ctx context.Context) (*graphqlbackend.RepositoryResolver, error) {
 	r.once.Do(func() {
 		repo, err := r.db.Repos().Get(ctx, r.job.RepoID)
@@ -189,3 +202,26 @@ func unmarshalRepoEmbeddingJobID(id graphql.ID) (jobID int, err error) {
 	err = relay.UnmarshalSpec(id, &jobID)
 	return
 }
+
+type repoEmbeddingJobStatsResolver struct {
+	stats repobg.EmbedRepoStats
+}
+
+func (r *repoEmbeddingJobStatsResolver) FilesScheduled() int32 {
+	return int32(r.stats.CodeIndexStats.FilesScheduled + r.stats.TextIndexStats.FilesScheduled)
+}
+
+func (r *repoEmbeddingJobStatsResolver) FilesEmbedded() int32 {
+	return int32(r.stats.CodeIndexStats.FilesEmbedded + r.stats.TextIndexStats.FilesEmbedded)
+}
+
+func (r *repoEmbeddingJobStatsResolver) FilesSkipped() int32 {
+	skipped := 0
+	for _, count := range r.stats.CodeIndexStats.FilesSkipped {
+		skipped += count
+	}
+	for _, count := range r.stats.TextIndexStats.FilesSkipped {
+		skipped += count
+	}
+	return int32(skipped)
+}

enterprise/cmd/worker/internal/embeddings/repo/handler.go

@@ -9,7 +9,7 @@ import (
 	codeintelContext "github.com/sourcegraph/sourcegraph/enterprise/internal/codeintel/context"
 	edb "github.com/sourcegraph/sourcegraph/enterprise/internal/database"
 	"github.com/sourcegraph/sourcegraph/enterprise/internal/embeddings"
-	repoembeddingsbg "github.com/sourcegraph/sourcegraph/enterprise/internal/embeddings/background/repo"
+	bgrepo "github.com/sourcegraph/sourcegraph/enterprise/internal/embeddings/background/repo"
 	"github.com/sourcegraph/sourcegraph/enterprise/internal/embeddings/embed"
 	"github.com/sourcegraph/sourcegraph/internal/actor"
 	"github.com/sourcegraph/sourcegraph/internal/api"
@@ -27,10 +27,10 @@ type handler struct {
 	uploadStore            uploadstore.Store
 	gitserverClient        gitserver.Client
 	contextService         embed.ContextService
-	repoEmbeddingJobsStore repoembeddingsbg.RepoEmbeddingJobsStore
+	repoEmbeddingJobsStore bgrepo.RepoEmbeddingJobsStore
 }
 
-var _ workerutil.Handler[*repoembeddingsbg.RepoEmbeddingJob] = &handler{}
+var _ workerutil.Handler[*bgrepo.RepoEmbeddingJob] = &handler{}
 
 // The threshold to embed the entire file is slightly larger than the chunk threshold to
 // avoid splitting small files unnecessarily.
@@ -49,7 +49,7 @@ var splitOptions = codeintelContext.SplitOptions{
 	ChunkEarlySplitTokensThreshold: embeddingChunkEarlySplitTokensThreshold,
 }
 
-func (h *handler) Handle(ctx context.Context, logger log.Logger, record *repoembeddingsbg.RepoEmbeddingJob) error {
+func (h *handler) Handle(ctx context.Context, logger log.Logger, record *bgrepo.RepoEmbeddingJob) error {
 	if !conf.EmbeddingsEnabled() {
 		return errors.New("embeddings are not configured or disabled")
 	}
@@ -99,6 +99,12 @@ func (h *handler) Handle(ctx context.Context, logger log.Logger, record *repoemb
 		return err
 	}
 
+	reportStats := func(stats *bgrepo.EmbedRepoStats) {
+		if err := h.repoEmbeddingJobsStore.UpdateRepoEmbeddingJobStats(ctx, record.ID, stats); err != nil {
+			logger.Error("failed to update embedding stats", log.Error(err))
+		}
+	}
+
 	repoEmbeddingIndex, toRemove, stats, err := embed.EmbedRepo(
 		ctx,
 		embeddingsClient,
@@ -107,11 +113,14 @@ func (h *handler) Handle(ctx context.Context, logger log.Logger, record *repoemb
 		ranks,
 		opts,
 		logger,
+		reportStats,
 	)
 	if err != nil {
 		return err
 	}
 
+	reportStats(stats) // final, complete report
+
 	logger.Info(
 		"finished generating repo embeddings",
 		log.String("repoName", string(repo.Name)),