From 07a2d042de507ed837295edb53a1fe3166d734a2 Mon Sep 17 00:00:00 2001
From: sergeyb <sergeyb@uber.com>
Date: Sat, 7 Mar 2026 21:41:54 +0000
Subject: [PATCH] feat(cancel): Graceful cancellation to GetChangedTargets

---
 controller/getchangedtargets.go | 40 +++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go
index 7c24178..95951ed 100644
--- a/controller/getchangedtargets.go
+++ b/controller/getchangedtargets.go
@@ -36,9 +36,10 @@ type job struct {
 	cancel            context.CancelFunc
 }
 
-// GetChangedTargets returns the changed targets between two revisions.
+// GetChangedTargets returns the changed targets between two revisions. If the client disconnects, the context will be cancelled, and the function returns with context.ErrCancelled.
 func (c *controller) GetChangedTargets(request *pb.GetChangedTargetsRequest, stream pb.TangoServiceGetChangedTargetsYARPCServer) error {
-	ctx := context.Background()
+	// Use the stream from the context as a parent context for the request. If the client disconnects, the context will be cancelled.
+	ctx := stream.Context()
 	if err := validateGetChangedTargetsRequest(request); err != nil {
 		c.logger.Error("GetChangedTargets: Invalid request", zap.Error(err))
 		return err
@@ -179,12 +180,35 @@ func (c *controller) compareTargetGraphs(ctx context.Context, firstGraph, second
 
 	// 1) Extract targets and metadata; index by canonical names
 	firstTargetsByID, firstMetadata := getTargetsAndMetadata(firstGraph)
+
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	secondTargetsByID, secondMetadata := getTargetsAndMetadata(secondGraph)
+
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	firstByName := buildNameIndex(firstTargetsByID, firstMetadata)
+
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	secondByName := buildNameIndex(secondTargetsByID, secondMetadata)
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	sourceFileRuleTypeID := detectSourceFileID(secondMetadata)
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	// 2) Build newTargetsMap and processing order (source files first if present)
 	newTargetsMap := make(map[string]*pb.OptimizedTarget, len(secondByName))
 	for name, t := range secondByName {
@@ -271,6 +295,10 @@ func (c *controller) compareTargetGraphs(ctx context.Context, firstGraph, second
 		}
 	}
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	// Iterate over the changed targets and check if any of them are DIRECT changes.
 	for name, ct := range changedByName {
 		if ct.GetChangeType() == pb.CHANGE_TYPE_DIRECT || ct.GetChangeType() == pb.CHANGE_TYPE_NEW {
@@ -305,11 +333,19 @@ func (c *controller) compareTargetGraphs(ctx context.Context, firstGraph, second
 		}
 	}
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	// Compute BFS distances from CHANGE_TYPE_DIRECT targets through the dependency graph.
 	if outputConfig.GetComputeDistances() {
 		computeDistances(c.logger, changedByName, secondByName, secondMetadata)
 	}
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	// TODO: https://github.com/uber/tango/issues/34
 	// only return changed targets changed within x distance from a direct target