实现自定义 wg 以统计协程数量

pkg/utils/chan.go

@@ -0,0 +1,48 @@
+package utils
+
+import (
+	"context"
+	"log/slog"
+	"net"
+
+	"github.com/pkg/errors"
+)
+
+func ConnChan(ctx context.Context, ls net.Listener) chan net.Conn {
+	connCh := make(chan net.Conn)
+	go func() {
+		for {
+			conn, err := ls.Accept()
+			if err != nil {
+				slog.Error("接受连接失败", err)
+				// 临时错误重试连接
+				var ne net.Error
+				if errors.As(err, &ne) && ne.Temporary() {
+					slog.Debug("临时错误重试")
+					continue
+				}
+				return
+			}
+			// ctx 取消后退出
+			select {
+			case <-ctx.Done():
+				Close(conn)
+				return
+			case connCh <- conn:
+			}
+		}
+	}()
+	return connCh
+}
+
+func WaitChan(ctx context.Context, wg *CountWaitGroup) chan struct{} {
+	ch := make(chan struct{})
+	go func() {
+		wg.Wait()
+		select {
+		case <-ctx.Done():
+		case ch <- struct{}{}:
+		}
+	}()
+	return ch
+}

pkg/utils/sync.go

@@ -0,0 +1,29 @@
+package utils
+
+import (
+	"sync"
+	"sync/atomic"
+)
+
+type CountWaitGroup struct {
+	wg  sync.WaitGroup
+	num atomic.Uint64
+}
+
+func (c *CountWaitGroup) Add(delta uint64) {
+	c.wg.Add(int(delta))
+	c.num.Add(delta)
+}
+
+func (c *CountWaitGroup) Done() {
+	c.wg.Done()
+	c.num.Add(-1)
+}
+
+func (c *CountWaitGroup) Wait() {
+	c.wg.Wait()
+}
+
+func (c *CountWaitGroup) Count() uint64 {
+	return c.num.Load()
+}

pkg/utils/utils.go

@@ -1,13 +1,8 @@
 package utils
 
 import (
-	"context"
 	"io"
 	"log/slog"
-	"net"
-	"sync"
-
-	"github.com/pkg/errors"
 )
 
 func ReadByte(reader io.Reader) (byte, error) {
@@ -36,42 +31,3 @@ func Close[T io.Closer](v T) {
 		slog.Warn("对象关闭失败", "err", err)
 	}
 }
-
-func ConnChan(ctx context.Context, ls net.Listener) chan net.Conn {
-	connCh := make(chan net.Conn)
-	go func() {
-		for {
-			conn, err := ls.Accept()
-			if err != nil {
-				slog.Error("接受连接失败", err)
-				// 临时错误重试连接
-				var ne net.Error
-				if errors.As(err, &ne) && ne.Temporary() {
-					slog.Debug("临时错误重试")
-					continue
-				}
-				return
-			}
-			// ctx 取消后退出
-			select {
-			case <-ctx.Done():
-				Close(conn)
-				return
-			case connCh <- conn:
-			}
-		}
-	}()
-	return connCh
-}
-
-func WaitChan(ctx context.Context, wg *sync.WaitGroup) chan struct{} {
-	ch := make(chan struct{})
-	go func() {
-		wg.Wait()
-		select {
-		case <-ctx.Done():
-		case ch <- struct{}{}:
-		}
-	}()
-	return ch
-}