实现自定义 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
+}