perf(orchestrator): store treehash mapping before graph computation

orchestrator/metrics.go

@@ -16,3 +16,6 @@ package orchestrator
 
 // opGetTargetGraph is snake_cased after the Orchestrator.GetTargetGraph method.
 const _opGetTargetGraph = "get_target_graph"
+
+// _opTreehashCacheWrite tracks treehash mapping upload that runs before graph computation.
+const _opTreehashCacheWrite = "treehash_cache_write"

orchestrator/native_orchestrator.go

@@ -186,6 +186,26 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT
 	} else {
 		logger.Infow("GetTargetGraph: bypass_cache=true, computing target graph")
 	}
+	// Store the treehash mapping in the background before the (potentially
+	// slow) graph computation so concurrent or subsequent requests can
+	// resolve it without waiting for the graph to finish.
+	go func() {
+		bgOp := metrics.Begin(e, _opTreehashCacheWrite, metrics.FastDurationBuckets)
+		thCachePath := cachekey.GetTreehashCachePath(build)
+		putErr := b.storage.Put(b.appCtx, storage.UploadRequest{
+			Key:    thCachePath,
+			Reader: bytes.NewReader([]byte(treehash)),
+		})
+		bgOp.Complete(putErr)
+		if putErr != nil {
+			logger.Warnw("GetTargetGraph: Failed to eagerly store treehash mapping",
+				zap.String("path", thCachePath), zap.Error(putErr))
+		} else {
+			logger.Infow("GetTargetGraph: Eagerly stored treehash mapping",
+				zap.String("path", thCachePath), zap.String("treehash", treehash))
+		}
+	}()
+
 	// Compute the target graph and store it in storage.
 	runner := b.graphRunner
 	if runner == nil {
@@ -229,12 +249,6 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT
 	if err != nil {
 		return nil, fmt.Errorf("write graph to storage at %s: %w", treehashPath, err)
 	}
-	treehashCachePath := cachekey.GetTreehashCachePath(build)
-	treehashReader := bytes.NewReader([]byte(treehash))
-	err = b.storage.Put(ctx, storage.UploadRequest{Key: treehashCachePath, Reader: treehashReader})
-	if err != nil {
-		return nil, fmt.Errorf("store treehash mapping at %s: %w", treehashCachePath, err)
-	}
 	recordStep(e, "cache_write_duration", cacheWriteStart, metrics.FastDurationBuckets)
 	graphReader, err := storage.NewGraphReader(ctx, b.storage, treehashPath)
 	if err != nil {

orchestrator/native_orchestrator_test.go

@@ -89,14 +89,16 @@ func TestNative_GetTargetGraph_Success(t *testing.T) {
 func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
+	appCtx, appCancel := context.WithCancel(context.Background())
+	defer appCancel()
 	st := storagemock.NewMockStorage(ctrl)
 	// First attempt returns NotFound to trigger compute path.
 	st.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, storage.NewNotFoundError("missing"))
-	// Expect writes (graph list and treehash cache mapping)
+	// Expect writes (graph stream and background treehash-mapping upload).
 	st.EXPECT().Put(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req storage.UploadRequest) error {
 		_, err := io.Copy(io.Discard, req.Reader)
 		return err
-	}).MinTimes(2)
+	}).MinTimes(1).MaxTimes(3)
 	// After compute, second read returns a valid delimited stream with one message
 	var buf bytes.Buffer
 	_ = json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}})
@@ -119,7 +121,7 @@ func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) {
 			RuleType: "go_library",
 		},
 	}}, nil)
-	o, err := NewNativeOrchestrator(context.Background(), Params{
+	o, err := NewNativeOrchestrator(appCtx, Params{
 		Storage:     st,
 		RepoManager: rm,
 		Logger:      zaptest.NewLogger(t).Sugar(),
@@ -140,19 +142,23 @@ func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) {
 func TestNative_GetTargetGraph_UnknownStrategy_UserError(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
+	appCtx, appCancel := context.WithCancel(context.Background())
+	defer appCancel()
 
 	st := storagemock.NewMockStorage(ctrl)
+	// The background treehash-mapping goroutine fires before the strategy check.
+	st.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	g := gitmock.NewMockInterface(ctrl)
 	g.EXPECT().RevParse(gomock.Any(), "HEAD^{tree}").Return("th", nil)
 	ws := workspacemock.NewMockWorkspace(ctrl)
-	ws.EXPECT().Path().Return("/tmp/ws")
+	ws.EXPECT().Path().Return("/tmp/ws").AnyTimes()
 	ws.EXPECT().Checkout(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)
 	ws.EXPECT().ApplyRequests(gomock.Any(), gomock.Any()).Return(nil)
 	ws.EXPECT().Release().Return(nil)
 	rm := repomanagermock.NewMockRepoManager(ctrl)
 	rm.EXPECT().Lease(gomock.Any(), gomock.Any()).Return(ws, nil)
 
-	o, err := NewNativeOrchestrator(context.Background(), Params{
+	o, err := NewNativeOrchestrator(appCtx, Params{
 		Storage:     st,
 		RepoManager: rm,
 		Logger:      zaptest.NewLogger(t).Sugar(),