pubgo · kooksee · May 3, 2024 · May 3, 2024 · May 3, 2024 · May 3, 2024
diff --git a/syncutil/waitgroup.go b/syncutil/waitgroup.go
@@ -4,15 +4,14 @@ import (
 	"runtime"
 	"sync"
 	"sync/atomic"
-	_ "unsafe"
+	"time"
 
-	"github.com/pubgo/funk/async"
-	"github.com/pubgo/funk/fastrand"
+	"github.com/pubgo/funk/log"
+	"github.com/pubgo/funk/stack"
+	"github.com/pubgo/funk/try"
+	"github.com/rs/zerolog"
 )
 
-//go:linkname state sync.(*WaitGroup).state
-func state(*sync.WaitGroup) (*uint64, *uint32)
-
 var defaultConcurrent = uint32(runtime.NumCPU() * 2)
 
 func NewWaitGroup(maxConcurrent ...uint32) *WaitGroup {
@@ -24,50 +23,45 @@ func NewWaitGroup(maxConcurrent ...uint32) *WaitGroup {
 }
 
 type WaitGroup struct {
-	_             noCopy
 	wg            sync.WaitGroup
+	count         atomic.Uint32
 	maxConcurrent uint32
-	err           error
-}
-
-func (t *WaitGroup) Count() uint32 {
-	count, _ := state(&t.wg)
-	return uint32(atomic.LoadUint64(count) >> 32)
 }
 
-func (t *WaitGroup) check() {
-	// 阻塞, 等待任务处理完毕
-	// 采样率(1), 打印log
-	for t.Count() >= t.maxConcurrent {
-		if fastrand.Sampling(1) {
-			logs.Warn().
-				Uint32("current", t.Count()).
-				Uint32("maximum", t.maxConcurrent).
-				Msg("WaitGroup current concurrent number exceeds the maximum concurrent number of the system")
+func (t *WaitGroup) Count() uint32 { return t.count.Load() }
+func (t *WaitGroup) checkAndWait() {
+	for {
+		count := t.Count()
+		if count <= t.maxConcurrent {
+			break
 		}
 
-		runtime.Gosched()
+		if count < defaultConcurrent {
+			runtime.Gosched()
+		} else {
+			time.Sleep(time.Microsecond * 10)
+		}
 	}
 }
 
 func (t *WaitGroup) Go(fn func()) {
 	t.wg.Add(1)
-	t.check()
-	async.GoSafe(
-		func() error { fn(); return nil },
-		func(err error) {
-			t.err = err
-		},
-	)
+	t.checkAndWait()
+	go func() {
+		defer t.wg.Done()
+		err := try.Try(func() error {
+			fn()
+			return nil
+		})
+		if err != nil {
+			log.Err(err).
+				Func(func(e *zerolog.Event) {
+					e.Str("fn_stack", stack.CallerWithFunc(fn).String())
+				}).Msg("recovery func panic")
+		}
+	}()
 }
 
-func (t *WaitGroup) Wait() error {
+func (t *WaitGroup) Wait() {
 	t.wg.Wait()
-	return t.err
 }
-
-type noCopy struct{}
-
-// Lock is a no-op used by -copylocks checker from `go vet`.
-func (*noCopy) Lock()   {}
-func (*noCopy) Unlock() {}
diff --git a/syncutil/waitgroup.s b/syncutil/waitgroup.s
diff --git a/syncutil/waitgroup_test.go b/syncutil/waitgroup_test.go
@@ -0,0 +1,20 @@
+package syncutil
+
+import (
+	"testing"
+	"time"
+)
+
+func TestName(t *testing.T) {
+	now := time.Now()
+	defer func() {
+		t.Log(time.Since(now))
+	}()
+	wg := NewWaitGroup(10)
+	for i := 0; i < 10; i++ {
+		wg.Go(func() {
+			time.Sleep(time.Millisecond * 10)
+		})
+	}
+	wg.Wait()
+}