github.com/qiniu/go-sdk/v7@v7.27.0/internal/dialer/dialer.go (about)

     1  package dialer
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"net"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  type (
    12  	DialOptions struct {
    13  		Timeout   time.Duration
    14  		KeepAlive time.Duration
    15  	}
    16  
    17  	eitherConnOrError struct {
    18  		conn net.Conn
    19  		err  error
    20  	}
    21  
    22  	dialerErrs struct {
    23  		errs []error
    24  	}
    25  )
    26  
    27  func DialContext(ctx context.Context, network string, ips []net.IP, port string, dialOptions DialOptions) (net.Conn, error) {
    28  	var wg sync.WaitGroup
    29  	resultsChan := make(chan eitherConnOrError, len(ips))
    30  	cancels := make([]context.CancelFunc, 0, len(ips))
    31  	err := &dialerErrs{errs: make([]error, 0, len(ips))}
    32  	interval := dialOptions.Timeout / time.Duration(len(ips))
    33  
    34  	defer func() {
    35  		for _, cancel := range cancels {
    36  			cancel()
    37  		}
    38  		wg.Wait()
    39  		close(resultsChan)
    40  	}()
    41  
    42  	if len(ips) > 0 {
    43  		ip := ips[0]
    44  		ips = ips[1:]
    45  		newCtx, newCancel := context.WithCancel(ctx)
    46  		cancels = append(cancels, newCancel)
    47  		wg.Add(1)
    48  		dialContextAsync(newCtx, &wg, network, ip, port, DialOptions{Timeout: dialOptions.Timeout, KeepAlive: dialOptions.KeepAlive}, resultsChan)
    49  	} else {
    50  		return nil, errors.New("no ip could be dialed")
    51  	}
    52  
    53  	ticker := time.NewTicker(interval)
    54  	defer ticker.Stop()
    55  	for {
    56  		select {
    57  		case <-ticker.C:
    58  			if len(ips) > 0 {
    59  				ip := ips[0]
    60  				ips = ips[1:]
    61  				newCtx, newCancel := context.WithCancel(ctx)
    62  				cancels = append(cancels, newCancel)
    63  				wg.Add(1)
    64  				dialContextAsync(newCtx, &wg, network, ip, port, DialOptions{Timeout: dialOptions.Timeout - interval*time.Duration(len(cancels)-1), KeepAlive: dialOptions.KeepAlive}, resultsChan)
    65  			} else {
    66  				return nil, err
    67  			}
    68  		case connOrErr := <-resultsChan:
    69  			if connOrErr.err != nil {
    70  				err.errs = append(err.errs, connOrErr.err)
    71  			} else if connOrErr.conn != nil {
    72  				return connOrErr.conn, nil
    73  			}
    74  		}
    75  	}
    76  }
    77  
    78  func dialContextSync(ctx context.Context, network string, ip net.IP, port string, dialOptions DialOptions) (net.Conn, error) {
    79  	dialer := net.Dialer{Timeout: dialOptions.Timeout, KeepAlive: dialOptions.KeepAlive}
    80  	newAddr := ip.String()
    81  	if port != "" {
    82  		newAddr = net.JoinHostPort(newAddr, port)
    83  	}
    84  	return dialer.DialContext(ctx, network, newAddr)
    85  }
    86  
    87  func dialContextAsync(ctx context.Context, wg *sync.WaitGroup, network string, ip net.IP, port string, dialOptions DialOptions, c chan<- eitherConnOrError) {
    88  	go func() {
    89  		defer wg.Done()
    90  		conn, err := dialContextSync(ctx, network, ip, port, dialOptions)
    91  		if err != nil {
    92  			c <- eitherConnOrError{err: err}
    93  		} else {
    94  			c <- eitherConnOrError{conn: conn}
    95  		}
    96  	}()
    97  }
    98  
    99  func (e *dialerErrs) Error() string {
   100  	if len(e.errs) > 0 {
   101  		return e.errs[0].Error()
   102  	} else {
   103  		return context.DeadlineExceeded.Error()
   104  	}
   105  }
   106  
   107  func (e *dialerErrs) Unwrap() error {
   108  	if len(e.errs) > 0 {
   109  		return e.errs[0]
   110  	} else {
   111  		return context.DeadlineExceeded
   112  	}
   113  }
   114  
   115  func (e *dialerErrs) Timeout() bool {
   116  	if len(e.errs) > 0 {
   117  		if te, ok := e.errs[0].(interface{ Timeout() bool }); ok {
   118  			return te.Timeout()
   119  		} else {
   120  			return false
   121  		}
   122  	} else {
   123  		return true
   124  	}
   125  }