github.com/MetalBlockchain/metalgo@v1.11.9/network/dialer/dialer.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package dialer
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"net"
    10  	"net/netip"
    11  	"time"
    12  
    13  	"go.uber.org/zap"
    14  
    15  	"github.com/MetalBlockchain/metalgo/network/throttling"
    16  	"github.com/MetalBlockchain/metalgo/utils/logging"
    17  )
    18  
    19  var _ Dialer = (*dialer)(nil)
    20  
    21  // Dialer attempts to create a connection with the provided IP/port pair
    22  type Dialer interface {
    23  	// If [ctx] is canceled, gives up trying to connect to [ip]
    24  	// and returns an error.
    25  	Dial(ctx context.Context, ip netip.AddrPort) (net.Conn, error)
    26  }
    27  
    28  type dialer struct {
    29  	dialer    net.Dialer
    30  	log       logging.Logger
    31  	network   string
    32  	throttler throttling.DialThrottler
    33  }
    34  
    35  type Config struct {
    36  	ThrottleRps       uint32        `json:"throttleRps"`
    37  	ConnectionTimeout time.Duration `json:"connectionTimeout"`
    38  }
    39  
    40  // NewDialer returns a new Dialer that calls net.Dial with the provided network.
    41  // [network] is the network passed into Dial. Should probably be "TCP".
    42  // [dialerConfig.connectionTimeout] gives the timeout when dialing an IP.
    43  // [dialerConfig.throttleRps] gives the max number of outgoing connection attempts/second.
    44  // If [dialerConfig.throttleRps] == 0, outgoing connections aren't rate-limited.
    45  func NewDialer(network string, dialerConfig Config, log logging.Logger) Dialer {
    46  	var throttler throttling.DialThrottler
    47  	if dialerConfig.ThrottleRps <= 0 {
    48  		throttler = throttling.NewNoDialThrottler()
    49  	} else {
    50  		throttler = throttling.NewDialThrottler(int(dialerConfig.ThrottleRps))
    51  	}
    52  	log.Debug(
    53  		"creating dialer",
    54  		zap.Uint32("throttleRPS", dialerConfig.ThrottleRps),
    55  		zap.Duration("dialTimeout", dialerConfig.ConnectionTimeout),
    56  	)
    57  	return &dialer{
    58  		dialer:    net.Dialer{Timeout: dialerConfig.ConnectionTimeout},
    59  		log:       log,
    60  		network:   network,
    61  		throttler: throttler,
    62  	}
    63  }
    64  
    65  func (d *dialer) Dial(ctx context.Context, ip netip.AddrPort) (net.Conn, error) {
    66  	if err := d.throttler.Acquire(ctx); err != nil {
    67  		return nil, err
    68  	}
    69  	d.log.Verbo("dialing",
    70  		zap.Stringer("ip", ip),
    71  	)
    72  	conn, err := d.dialer.DialContext(ctx, d.network, ip.String())
    73  	if err != nil {
    74  		return nil, fmt.Errorf("error while dialing %s: %w", ip, err)
    75  	}
    76  	return conn, nil
    77  }