fix(worker): Cancel the context when opening connection to leader for…

… streaming snapshot (#6033) (#6046) Co-authored-by: Daniel Mai <[email protected]> Co-authored-by: Manish R Jain <[email protected]>
hypermodeinc · Jul 21, 2020 · 0cfab06 · 0cfab06
1 parent 0b8cf73
commit 0cfab06
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/worker/snapshot.go b/worker/snapshot.go
@@ -49,8 +49,12 @@ func (n *node) populateSnapshot(snap pb.Snapshot, pl *conn.Pool) (int, error) {
 	con := pl.Get()
 	c := pb.NewWorkerClient(con)
 
+	// We should absolutely cancel the context when we return from this function, that way, the
+	// leader who is sending the snapshot would stop sending.
+	ctx, cancel := context.WithCancel(n.ctx)
+	defer cancel()
+
 	// Set my RaftContext on the snapshot, so it's easier to locate me.
-	ctx := n.ctx
 	snap.Context = n.RaftContext
 	stream, err := c.StreamSnapshot(ctx)
 	if err != nil {