fix case where context closes docs reader

mongodb-labs · Dec 7, 2024 · 28b8525 · 28b8525
1 parent 50708b7
commit 28b8525
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/internal/verifier/compare.go b/internal/verifier/compare.go
@@ -375,12 +375,17 @@ func iterateCursorToChannel(
 	cursor *mongo.Cursor,
 	writer chan<- bson.Raw,
 ) error {
+	defer close(writer)
+
 	for cursor.Next(ctx) {
 		state.NoteSuccess("received a document")
-		writer <- slices.Clone(cursor.Current)
-	}
 
-	close(writer)
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case writer <- slices.Clone(cursor.Current):
+		}
+	}
 
 	return errors.Wrap(cursor.Err(), "failed to iterate cursor")
 }