cleanup

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
2024-03-09 20:04:32 +03:00 · 2024-03-09 20:04:32 +03:00 · e430c15ae8
commit e430c15ae8
parent a39a2f398a
1 changed files with 2 additions and 1 deletions
--- 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