From e430c15ae83e6e05f7d1a37894d4167d29709d30 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Sat, 9 Mar 2024 20:04:32 +0300
Subject: [PATCH] cleanup

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 sync/waitgroup.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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