From 35f7065d05e71747c5bd130b1e42246609d6cf11 Mon Sep 17 00:00:00 2001
From: $MY_NAME <komuw05@gmail.com>
Date: Wed, 20 Dec 2023 16:47:42 +0300
Subject: [PATCH] ok

---
 go.mod            |  1 +
 go.sum            |  2 ++
 sync/sync.go      |  4 +++
 sync/sync_test.go | 76 ++++++++++++++++++++++++++---------------------
 4 files changed, 49 insertions(+), 34 deletions(-)

diff --git a/go.mod b/go.mod
index 0f4e6a9a..5297ad58 100644
--- a/go.mod
+++ b/go.mod
@@ -16,6 +16,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/rs/dnscache v0.0.0-20230804202142-fc85eb664529 // indirect
+	github.com/sourcegraph/conc v0.3.0 // indirect
 	github.com/tmthrgd/httputils v0.0.0-20190904060602-27fdf7d93acd // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/sync v0.5.0 // indirect
diff --git a/go.sum b/go.sum
index 103610e4..f372051d 100644
--- a/go.sum
+++ b/go.sum
@@ -39,6 +39,8 @@ github.com/rs/zerolog v1.31.0 h1:FcTR3NnLWW+NnTwwhFWiJSZr4ECLpqCm6QsEnyvbV4A=
 github.com/rs/zerolog v1.31.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
+github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d h1:X4+kt6zM/OVO6gbJdAfJR60MGPsqCzbtXNnjoGqdfAs=
 github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
diff --git a/sync/sync.go b/sync/sync.go
index 6f4681a0..79d6bfe6 100644
--- a/sync/sync.go
+++ b/sync/sync.go
@@ -17,6 +17,7 @@ import (
 //
 // Use [New] to get a valid Waitgroup
 type WaitGroup struct {
+	mu     sync.Mutex
 	wg     sync.WaitGroup
 	cancel context.CancelCauseFunc
 
@@ -121,7 +122,10 @@ func (w *WaitGroup) Go(funcs ...func() error) error {
 					}
 				}(f)
 			}
+
+			w.mu.Lock()
 			w.wg.Wait()
+			w.mu.Unlock()
 		}
 
 		if w.cancel != nil {
diff --git a/sync/sync_test.go b/sync/sync_test.go
index 1289b084..da7e8c86 100644
--- a/sync/sync_test.go
+++ b/sync/sync_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/sourcegraph/conc"
 	"golang.org/x/sync/errgroup"
 
 	"go.akshayshah.org/attest"
@@ -109,42 +110,25 @@ func TestErrGroup(t *testing.T) {
 	t.Run("concurrency", func(t *testing.T) {
 		t.Parallel()
 
-		// wgLimited, _ := errgroup.WithContext(context.Background())
-		// wgLimited.SetLimit(1)
-		wgUnlimited, _ := errgroup.WithContext(context.Background())
-
-		// {
-		// 	funcs := []func() error{}
-		// 	for i := 0; i <= 4; i++ {
-		// 		funcs = append(funcs,
-		// 			func() error {
-		// 				return nil
-		// 			},
-		// 		)
-		// 	}
-
-		// 	go func() {
-		// 		wgLimited.Go(func() error {
-		// 			return nil
-		// 		})
-		// 	}()
-		// 	wgLimited.Go(func() error {
-		// 		return nil
-		// 	})
-
-		// 	err := wgLimited.Wait()
-		// 	attest.Ok(t, err)
-		// }
+		{
+			wgLimited, _ := errgroup.WithContext(context.Background())
+			wgLimited.SetLimit(1)
+
+			go func() {
+				wgLimited.Go(func() error {
+					return nil
+				})
+			}()
+			wgLimited.Go(func() error {
+				return nil
+			})
+
+			err := wgLimited.Wait()
+			attest.Ok(t, err)
+		}
 
 		{
-			funcs := []func() error{}
-			for i := 0; i <= 4; i++ {
-				funcs = append(funcs,
-					func() error {
-						return nil
-					},
-				)
-			}
+			wgUnlimited, _ := errgroup.WithContext(context.Background())
 
 			go func() {
 				wgUnlimited.Go(func() error {
@@ -160,3 +144,27 @@ func TestErrGroup(t *testing.T) {
 		}
 	})
 }
+
+func TestConc(t *testing.T) {
+	t.Parallel()
+
+	t.Run("concurrency", func(t *testing.T) {
+		t.Parallel()
+
+		{
+			wgLimited := conc.NewWaitGroup()
+
+			go func() {
+				wgLimited.Go(func() {
+					return
+				})
+				wgLimited.Wait()
+			}()
+			wgLimited.Go(func() {
+				return
+			})
+
+			wgLimited.Wait()
+		}
+	})
+}