Cancel delta stream on delta send (#218)

castai · Mar 6, 2024 · 80c13ad · 80c13ad
1 parent 2a6e399
commit 80c13ad
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/cmd/controller/state/delta/controller.go b/cmd/controller/state/delta/controller.go
@@ -96,13 +96,17 @@ func (c *Controller) sendDeltas(ctx context.Context, firstDeltaReport bool) {
 		return
 	}
 
+	// Cancel context to close stream after deltas are sent.
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	if firstDeltaReport {
 		ctx = metadata.AppendToOutgoingContext(ctx,
 			"x-delta-full-snapshot", "true",
 		)
 	}
 	deltaStream, err := c.castaiClient.KubernetesDeltaIngest(ctx, grpc.UseCompressor(gzip.Name))
-	if err != nil {
+	if err != nil && !errors.Is(err, context.Canceled) {
 		c.log.Warnf("creating delta upload stream: %v", err)
 		return
 	}