github.com/ava-labs/avalanchego@v1.11.11/network/throttling/inbound_conn_throttler.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package throttling
     5  
     6  import (
     7  	"context"
     8  	"net"
     9  
    10  	"golang.org/x/time/rate"
    11  )
    12  
    13  var _ net.Listener = (*throttledListener)(nil)
    14  
    15  // Wraps [listener] and returns a net.Listener that will accept at most
    16  // [maxConnsPerSec] connections per second.
    17  // [maxConnsPerSec] must be non-negative.
    18  func NewThrottledListener(listener net.Listener, maxConnsPerSec float64) net.Listener {
    19  	ctx, cancel := context.WithCancel(context.Background())
    20  	return &throttledListener{
    21  		ctx:           ctx,
    22  		ctxCancelFunc: cancel,
    23  		listener:      listener,
    24  		limiter:       rate.NewLimiter(rate.Limit(maxConnsPerSec), int(maxConnsPerSec)+1),
    25  	}
    26  }
    27  
    28  // [throttledListener] is a net.Listener that rate-limits
    29  // acceptance of incoming connections.
    30  // Note that InboundConnUpgradeThrottler rate-limits _upgrading_ of
    31  // inbound connections, whereas throttledListener rate-limits
    32  // _acceptance_ of inbound connections.
    33  type throttledListener struct {
    34  	// [ctx] is cancelled when Close() is called
    35  	ctx context.Context
    36  	// [ctxCancelFunc] cancels [ctx] when it's called
    37  	ctxCancelFunc func()
    38  	// The underlying listener
    39  	listener net.Listener
    40  	// Handles rate-limiting
    41  	limiter *rate.Limiter
    42  }
    43  
    44  func (l *throttledListener) Accept() (net.Conn, error) {
    45  	// Wait until the rate-limiter says to accept the
    46  	// next incoming connection. If l.Close() is called,
    47  	// Wait will return immediately.
    48  	if err := l.limiter.Wait(l.ctx); err != nil {
    49  		return nil, err
    50  	}
    51  	return l.listener.Accept()
    52  }
    53  
    54  func (l *throttledListener) Close() error {
    55  	// Cancel [l.ctx] so Accept() will return immediately
    56  	l.ctxCancelFunc()
    57  	return l.listener.Close()
    58  }
    59  
    60  func (l *throttledListener) Addr() net.Addr {
    61  	return l.listener.Addr()
    62  }