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

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

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