49 lines
811 B
Go
49 lines
811 B
Go
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
|
|
}
|