feat(coroutine): Optimize Wait (#86)

* feat(coroutine): Optimize `Wait` Signed-off-by: Flc゛ <[email protected]> * feat(coroutine): Optimize `Wait` Signed-off-by: Flc゛ <[email protected]> * feat(coroutine): Optimize `Wait` Signed-off-by: Flc゛ <[email protected]> --------- Signed-off-by: Flc゛ <[email protected]>
go-fries · Jan 31, 2024 · f95070f · f95070f
1 parent 4c8992a
commit f95070f
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 12 deletions.
diff --git a/coroutine/README.md b/coroutine/README.md
@@ -41,5 +41,8 @@ func main() {
 
 	// Parallel Example2
 	coroutine.RunParallel(funcs...)
+
+	// Wait Example
+	coroutine.Wait(funcs...)
 }
 ```
diff --git a/coroutine/concurrent.go b/coroutine/concurrent.go
@@ -10,7 +10,6 @@ type Concurrent struct {
 
 func NewConcurrent(limit int) *Concurrent {
 	return &Concurrent{
-		wg: sync.WaitGroup{},
 		ch: make(chan struct{}, limit),
 		fs: make([]func(), 0),
 	}

diff --git a/coroutine/parallel.go b/coroutine/parallel.go
@@ -9,7 +9,6 @@ type Parallel struct {
 
 func NewParallel() *Parallel {
 	return &Parallel{
-		wg: sync.WaitGroup{},
 		fs: make([]func(), 0),
 	}
 }

diff --git a/coroutine/wait.go b/coroutine/wait.go
@@ -2,12 +2,14 @@ package coroutine
 
 import "sync"
 
-func Wait(fn func()) {
+func Wait(fs ...func()) {
 	var wg sync.WaitGroup
-	wg.Add(1)
-	go func() {
-		fn()
-		wg.Done()
-	}()
+	wg.Add(len(fs))
+	for _, f := range fs {
+		go func(f func()) {
+			defer wg.Done()
+			f()
+		}(f)
+	}
 	wg.Wait()
 }
diff --git a/coroutine/wait_test.go b/coroutine/wait_test.go
@@ -1,19 +1,31 @@
 package coroutine
 
 import (
+	"bytes"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestWait(t *testing.T) {
-	var msg string
+	var (
+		buffer bytes.Buffer
+		mu     sync.Mutex
+	)
 
 	Wait(func() {
-		time.Sleep(time.Second)
-		msg = "hello"
+		time.Sleep(100 * time.Millisecond)
+		mu.Lock()
+		defer mu.Unlock()
+		buffer.WriteString("hello")
+	}, func() {
+		time.Sleep(200 * time.Millisecond)
+		mu.Lock()
+		defer mu.Unlock()
+		buffer.WriteString(" world")
 	})
 
-	assert.Equal(t, "hello", msg)
+	assert.Equal(t, "hello world", buffer.String())
 }