Bug Report: VDiff leaks goroutines when canceled

[mhamza15]

Overview of the Issue

When a VDiff is canceled, StreamExecute can block indefinitely because it does not consider its context being canceled, and the result channel is not closed when the context is canceled:

vitess/go/vt/vttablet/tabletmanager/vdiff/table_differ.go

Lines 434 to 439 in b4d088d

	select {
	case <-ctx.Done():
	default:
	close(participant.result)
	close(gtidch)
	}

Reproduction Steps

Potential reproduction test case: mhamza15@c301eb3

❯ go test ./go/vt/vttablet/tabletmanager/vdiff/ -run TestShardStreamerNoLeakOnCancel -count=1 -timeout 120s -v
=== RUN   TestShardStreamerNoLeakOnCancel
--- FAIL: TestShardStreamerNoLeakOnCancel (0.00s)
panic: deadlock: main bubble goroutine has exited but blocked goroutines remain [recovered, repanicked]

goroutine 185 [running]:
testing.tRunner.func1.2({0x1077abc20, 0x4e0ddb18a2a0})
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/testing/testing.go:1974 +0x1a0
testing.tRunner.func1()
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/testing/testing.go:1977 +0x318
panic({0x1077abc20?, 0x4e0ddb18a2a0?})
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/runtime/panic.go:860 +0x12c
internal/synctest.Run(0x4e0ddade6120)
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/runtime/synctest.go:251 +0x2c4
testing/synctest.Test(0x4e0ddac1c248, 0x1079bbdc8)
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/testing/synctest/synctest.go:290 +0x88
vitess.io/vitess/go/vt/vttablet/tabletmanager/vdiff.TestShardStreamerNoLeakOnCancel(0x4e0ddac1c248?)
	/Users/mhamza/dev/vitess/go/vt/vttablet/tabletmanager/vdiff/primitive_executor_test.go:33 +0x24
testing.tRunner(0x4e0ddac1c248, 0x1079bbc40)
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/testing/testing.go:2036 +0xc4
created by testing.(*T).Run in goroutine 1
	/Users/mhamza/.local/share/mise/installs/go/1.26.0/src/testing/testing.go:2101 +0x3a8

goroutine 189 [chan receive (durable), synctest bubble 1]:
vitess.io/vitess/go/vt/vttablet/tabletmanager/vdiff.(*shardStreamer).StreamExecute(0x4e0ddaa0a240, {0x4e0ddaeb8480?, 0x0?}, {0x104825154?, 0x4e0ddaa0a2c0?}, 0x1079ec908?, 0x30?, 0x27?, 0x4e0ddade61a0)
	/Users/mhamza/dev/vitess/go/vt/vttablet/tabletmanager/vdiff/shard_streamer.go:45 +0x40
vitess.io/vitess/go/vt/vtgate/engine.runOneStream.func1()
	/Users/mhamza/dev/vitess/go/vt/vtgate/engine/merge_sort.go:232 +0x118
created by vitess.io/vitess/go/vt/vtgate/engine.runOneStream in goroutine 188
	/Users/mhamza/dev/vitess/go/vt/vtgate/engine/merge_sort.go:228 +0x134
FAIL	vitess.io/vitess/go/vt/vttablet/tabletmanager/vdiff	4.608s
FAIL

Binary Version

main

Operating System and Environment details

all

Log Fragments

goroutine 203 [chan receive (durable)]:
   vitess.io/vitess/go/vt/vttablet/tabletmanager/vdiff.(*shardStreamer).StreamExecute(...)
       shard_streamer.go:45
   vitess.io/vitess/go/vt/vtgate/engine.runOneStream.func1()
       merge_sort.go:232

[Devanshusharma2005]

Hi @mhamza15,
I tried running TestShardStreamerNoLeakOnCancel on the current main branch, but it looks like the test isn’t present there and nothing fails.
Is this issue still reproducible on main, or does it require the reproduction branch you linked?
Just want to confirm before digging deeper. Thanks!

[mhamza15]

Hi @Devanshusharma2005, thanks for taking a look! I already have a fix for it and will be putting up a PR soon. If you want to take a stab at it anyway for your own learning, copy the test case from the link above to your local setup, and that'll allow you to run it. Thanks!

[Devanshusharma2005]

Great!!
Could you please tag or mention me when you open the PR? I’d love to take a look at the fix and understand the approach. Thanks again!