proxy/pkg/utils/chan.go

package utils

import (
	"context"
	"log/slog"
	"net"

	"github.com/pkg/errors"
)

func ChanConnAccept(ctx context.Context, ls net.Listener) chan net.Conn {
	ch := make(chan net.Conn)
	go func() {
		defer close(ch)
		for {
			conn, err := ls.Accept()
			if err != nil {
				if errors.Is(err, net.ErrClosed) {
					return
				}
				// 临时错误重试连接
				var ne net.Error
				if errors.As(err, &ne) && ne.Temporary() {
					slog.Debug("临时错误重试")
					continue
				}
				slog.Error("接受连接失败", err)
				return
			}
			// ctx 取消后退出
			select {
			case <-ctx.Done():
				Close(conn)
				return
			case ch <- conn:
			}
		}
	}()
	return ch
}

func ChanWgWait[T WaitGroup](ctx context.Context, wg T) chan struct{} {
	ch := make(chan struct{})
	go func() {
		defer close(ch)
		wg.Wait()
		select {
		case <-ctx.Done():
		case ch <- struct{}{}:
		}
	}()
	return ch
}