chore: use WaitGroup in dualStackDialContext

This commit is contained in:
Larvan2 2023-08-30 17:07:49 +08:00
parent 86cf1dd54b
commit 414d8f2162

View file

@ -169,10 +169,15 @@ func dualStackDialContext(ctx context.Context, dialFn dialFunc, network string,
preferIPVersion := opt.prefer
fallbackTicker := time.NewTicker(fallbackTimeout)
defer fallbackTicker.Stop()
results := make(chan dialResult)
returned := make(chan struct{})
defer close(returned)
var wg sync.WaitGroup
racer := func(ips []netip.Addr, isPrimary bool) {
defer wg.Done()
result := dialResult{isPrimary: isPrimary}
defer func() {
select {
@ -186,27 +191,35 @@ func dualStackDialContext(ctx context.Context, dialFn dialFunc, network string,
result.Conn, result.error = dialFn(ctx, network, ips, port, opt)
}
var wait int
if len(ipv4s) != 0 {
wait++
wg.Add(1)
go racer(ipv4s, preferIPVersion != 6)
}
if len(ipv6s) != 0 {
wait++
wg.Add(1)
go racer(ipv6s, preferIPVersion != 4)
}
go func() {
wg.Wait()
close(results)
}()
var fallback dialResult
var errs []error
for i := 0; i < wait; {
loop:
for {
select {
case <-fallbackTicker.C:
if fallback.error == nil && fallback.Conn != nil {
return fallback.Conn, nil
}
case res := <-results:
i++
case res, ok := <-results:
if !ok {
break loop
}
if res.error == nil {
if res.isPrimary {
return res.Conn, nil
@ -221,6 +234,7 @@ func dualStackDialContext(ctx context.Context, dialFn dialFunc, network string,
}
}
}
if fallback.error == nil && fallback.Conn != nil {
return fallback.Conn, nil
}