diff --git a/sync/waitgroup.go b/sync/waitgroup.go
index d7fb7140..be6872f2 100644
--- a/sync/waitgroup.go
+++ b/sync/waitgroup.go
@@ -57,8 +57,9 @@ func (g *WaitGroup) WaitContext(ctx context.Context) {
 
 	select {
 	case <-ctx.Done():
+		g.drain.Store(true)
 		g.wg.Add(-int(g.c.Load()))
-		<-done
+		g.drain.Store(false)
 		return
 	case <-done:
 		return