perf: batch node persistence during builds

Author: tae2089

internal/app/ingest/workflow/build.go

@@ -107,26 +107,19 @@ func (b *buildPersistBatch) reset() {
 	b.bytes = 0
 }
 
-// bindAndReleaseNodeBatch upserts a parsed file's nodes and binds its comment annotations within a transaction-scoped store.
-// @intent persist nodes and their annotation bindings atomically per file before releasing comment buffers.
-// @sideEffect writes graph nodes and annotation rows via the transaction-scoped store.
-// @mutates graph nodes and annotations
-func (s *Service) bindAndReleaseNodeBatch(ctx context.Context, txStore ingestapp.GraphStore, batches []parsedBuildNodeBatch, idx int) error {
+// bindAndReleaseNodeBatch binds a parsed file's comments after its nodes have persisted, then releases comment buffers.
+// @intent preserve per-file annotation binding and release behavior after the enclosing flush persists all nodes together.
+// @sideEffect writes annotation rows via the transaction-scoped store.
+// @mutates graph annotations
+func (s *Service) bindAndReleaseNodeBatch(ctx context.Context, txStore ingestapp.GraphStore, storedNodesByFile map[string][]graph.Node, batches []parsedBuildNodeBatch, idx int) error {
 	parsed := &batches[idx]
 
-	if err := txStore.UpsertNodes(ctx, parsed.nodes); err != nil {
-		return trace.Wrap(err, "upsert nodes for "+parsed.relPath)
-	}
-
 	if len(parsed.tsComments) > 0 {
 		binderComments := toBinderComments(parsed.tsComments)
 		binder := binding.NewBinder()
 		bindings := binder.Bind(binderComments, parsed.nodes, parsed.language, parsed.sourceLines)
 
-		storedNodes, err := txStore.GetNodesByFile(ctx, parsed.relPath)
-		if err != nil {
-			return trace.Wrap(err, "get stored nodes for annotations")
-		}
+		storedNodes := storedNodesByFile[parsed.relPath]
 		storedMap := make(map[string]*graph.Node, len(storedNodes))
 		for i := range storedNodes {
 			key := storedNodes[i].QualifiedName + ":" + strconv.Itoa(storedNodes[i].StartLine)
@@ -602,19 +595,50 @@ func (s *Service) packageSemanticEdgeBatches(batches []parsedBuildNodeBatch) []p
 	return out
 }
 
-// flushBuildBatch persists the buffered nodes for the current batch.
-// @intent persist nodes before all edges so foreign-key style references can resolve.
+// flushBuildBatch persists the buffered nodes for the current bounded batch.
+// @intent persist all batch nodes before annotations and all edges so references can resolve with fewer store operations.
 // @sideEffect upserts graph nodes and annotations through the transaction-scoped store.
 // @mutates graph nodes and annotations
 func (s *Service) flushBuildBatch(ctx context.Context, txStore ingestapp.GraphStore, batch *buildPersistBatch) error {
 	if batch.files == 0 {
 		return nil
 	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+
+	nodeCount := 0
+	annotationFilePaths := make([]string, 0, len(batch.nodeBatches))
+	for _, parsed := range batch.nodeBatches {
+		nodeCount += len(parsed.nodes)
+		if len(parsed.tsComments) > 0 {
+			annotationFilePaths = append(annotationFilePaths, parsed.relPath)
+		}
+	}
+	nodes := make([]graph.Node, 0, nodeCount)
+	for _, parsed := range batch.nodeBatches {
+		nodes = append(nodes, parsed.nodes...)
+	}
+	if len(nodes) > 0 {
+		if err := txStore.UpsertNodes(ctx, nodes); err != nil {
+			return trace.Wrap(err, "upsert batch nodes")
+		}
+	}
+
+	var storedNodesByFile map[string][]graph.Node
+	if len(annotationFilePaths) > 0 {
+		stored, err := txStore.GetNodesByFiles(ctx, annotationFilePaths)
+		if err != nil {
+			return trace.Wrap(err, "get stored nodes for annotations")
+		}
+		storedNodesByFile = stored
+	}
+
 	for i := range batch.nodeBatches {
 		if err := ctx.Err(); err != nil {
 			return err
 		}
-		if err := s.bindAndReleaseNodeBatch(ctx, txStore, batch.nodeBatches, i); err != nil {
+		if err := s.bindAndReleaseNodeBatch(ctx, txStore, storedNodesByFile, batch.nodeBatches, i); err != nil {
 			return err
 		}
 	}

internal/app/ingest/workflow/indexer_test.go

@@ -116,6 +116,7 @@ type recordingGraphStore struct {
 	nextID                uint
 	nodesByFP             map[string][]graph.Node
 	edges                 []graph.Edge
+	upsertedNodeBatches   [][]graph.Node
 	upsertedEdges         [][]graph.Edge
 	fileSuffixLookupCalls int
 	importFileNodeCalls   int
@@ -137,6 +138,7 @@ func (r *recordingGraphStore) DeleteGraph(ctx context.Context) error {
 
 func (r *recordingGraphStore) UpsertNodes(ctx context.Context, nodes []graph.Node) error {
 	r.record("UpsertNodes")
+	r.upsertedNodeBatches = append(r.upsertedNodeBatches, append([]graph.Node(nil), nodes...))
 	for i := range nodes {
 		r.nextID++
 		nodes[i].ID = r.nextID
@@ -207,6 +209,7 @@ func (r *recordingGraphStore) GetNodesByQualifiedNames(ctx context.Context, name
 }
 
 func (r *recordingGraphStore) GetNodesByFiles(ctx context.Context, filePaths []string) (map[string][]graph.Node, error) {
+	r.record("GetNodesByFiles")
 	set := make(map[string]bool, len(filePaths))
 	for _, fp := range filePaths {
 		set[fp] = true
@@ -2334,7 +2337,7 @@ func Keep() {}
 		t.Fatalf("Build: %v", err)
 	}
 
-	want := []string{"DeleteGraph", "UpsertNodes", "GetNodesByFile", "UpsertAnnotation", "UpsertEdges"}
+	want := []string{"DeleteGraph", "UpsertNodes", "GetNodesByFiles", "UpsertAnnotation", "UpsertEdges"}
 	for i, op := range want[:4] {
 		if len(fakeStore.ops) <= i {
 			t.Fatalf("ops too short: got %v", fakeStore.ops)
@@ -2359,6 +2362,81 @@ func Keep() {}
 	}
 }
 
+func TestFlushBuildBatch_BatchesNodesAndAnnotationLookups(t *testing.T) {
+	fakeStore := newRecordingGraphStore(t)
+	released := 0
+	svc := &Service{
+		onBatchRelease: func([]parsedBuildNodeBatch, int) {
+			released++
+		},
+	}
+	batch := buildPersistBatch{
+		files: 2,
+		nodeBatches: []parsedBuildNodeBatch{
+			{
+				relPath: "first.go",
+				nodes: []graph.Node{{
+					QualifiedName: "sample.First",
+					FilePath:      "first.go",
+					Kind:          graph.NodeKindFunction,
+					StartLine:     2,
+					EndLine:       2,
+				}},
+				tsComments:  []ingest.CommentBlock{{StartLine: 1, EndLine: 1, Text: "@intent first"}},
+				language:    "go",
+				sourceLines: []string{"// @intent first", "func First() {}"},
+			},
+			{
+				relPath: "second.go",
+				nodes: []graph.Node{{
+					QualifiedName: "sample.Second",
+					FilePath:      "second.go",
+					Kind:          graph.NodeKindFunction,
+					StartLine:     2,
+					EndLine:       2,
+				}},
+				tsComments:  []ingest.CommentBlock{{StartLine: 1, EndLine: 1, Text: "@intent second"}},
+				language:    "go",
+				sourceLines: []string{"// @intent second", "func Second() {}"},
+			},
+		},
+	}
+
+	if err := svc.flushBuildBatch(context.Background(), fakeStore, &batch); err != nil {
+		t.Fatalf("flushBuildBatch: %v", err)
+	}
+	if got := len(fakeStore.upsertedNodeBatches); got != 1 {
+		t.Fatalf("node upsert calls = %d, want 1 (ops=%v)", got, fakeStore.ops)
+	}
+	if got := len(fakeStore.upsertedNodeBatches[0]); got != 2 {
+		t.Fatalf("batched node count = %d, want 2", got)
+	}
+	countOperation := func(want string) int {
+		count := 0
+		for _, op := range fakeStore.ops {
+			if op == want {
+				count++
+			}
+		}
+		return count
+	}
+	if got := countOperation("GetNodesByFiles"); got != 1 {
+		t.Fatalf("bulk annotation node lookups = %d, want 1 (ops=%v)", got, fakeStore.ops)
+	}
+	if got := countOperation("GetNodesByFile"); got != 0 {
+		t.Fatalf("per-file annotation node lookups = %d, want 0 (ops=%v)", got, fakeStore.ops)
+	}
+	if got := countOperation("UpsertAnnotation"); got != 2 {
+		t.Fatalf("annotation upserts = %d, want 2 (ops=%v)", got, fakeStore.ops)
+	}
+	if released != 2 {
+		t.Fatalf("released batches = %d, want 2", released)
+	}
+	if batch.files != 0 || len(batch.nodeBatches) != 0 {
+		t.Fatalf("batch was not reset: %+v", batch)
+	}
+}
+
 func TestBuild_FlushesLargeBuildInBoundedBatches(t *testing.T) {
 	fakeStore := newRecordingGraphStore(t)
 	svc := &Service{