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
 	}