github.com/ava-labs/avalanchego@v1.11.11/network/p2p/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 p2p
     5  
     6  import (
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/ava-labs/avalanchego/ids"
    11  	"github.com/ava-labs/avalanchego/utils/timer/mockable"
    12  )
    13  
    14  var _ Throttler = (*SlidingWindowThrottler)(nil)
    15  
    16  type Throttler interface {
    17  	// Handle returns true if a message from [nodeID] should be handled.
    18  	Handle(nodeID ids.NodeID) bool
    19  }
    20  
    21  // NewSlidingWindowThrottler returns a new instance of SlidingWindowThrottler.
    22  // Nodes are throttled if they exceed [limit] messages during an interval of
    23  // time over [period].
    24  // [period] and [limit] should both be > 0.
    25  func NewSlidingWindowThrottler(period time.Duration, limit int) *SlidingWindowThrottler {
    26  	now := time.Now()
    27  	return &SlidingWindowThrottler{
    28  		period: period,
    29  		limit:  float64(limit),
    30  		windows: [2]window{
    31  			{
    32  				start: now,
    33  				hits:  make(map[ids.NodeID]float64),
    34  			},
    35  			{
    36  				start: now.Add(-period),
    37  				hits:  make(map[ids.NodeID]float64),
    38  			},
    39  		},
    40  	}
    41  }
    42  
    43  // window is used internally by SlidingWindowThrottler to represent the amount
    44  // of hits from a node in the evaluation period beginning at [start]
    45  type window struct {
    46  	start time.Time
    47  	hits  map[ids.NodeID]float64
    48  }
    49  
    50  // SlidingWindowThrottler is an implementation of the sliding window throttling
    51  // algorithm.
    52  type SlidingWindowThrottler struct {
    53  	period time.Duration
    54  	limit  float64
    55  	clock  mockable.Clock
    56  
    57  	lock    sync.Mutex
    58  	current int
    59  	windows [2]window
    60  }
    61  
    62  // Handle returns true if the amount of calls received in the last [s.period]
    63  // time is less than [s.limit]
    64  //
    65  // This is calculated by adding the current period's count to a weighted count
    66  // of the previous period.
    67  func (s *SlidingWindowThrottler) Handle(nodeID ids.NodeID) bool {
    68  	s.lock.Lock()
    69  	defer s.lock.Unlock()
    70  
    71  	// The current window becomes the previous window if the current evaluation
    72  	// period is over
    73  	now := s.clock.Time()
    74  	sinceUpdate := now.Sub(s.windows[s.current].start)
    75  	if sinceUpdate >= 2*s.period {
    76  		s.rotate(now.Add(-s.period))
    77  	}
    78  	if sinceUpdate >= s.period {
    79  		s.rotate(now)
    80  		sinceUpdate = 0
    81  	}
    82  
    83  	currentHits := s.windows[s.current].hits
    84  	current := currentHits[nodeID]
    85  	previousFraction := float64(s.period-sinceUpdate) / float64(s.period)
    86  	previous := s.windows[1-s.current].hits[nodeID]
    87  	estimatedHits := current + previousFraction*previous
    88  	if estimatedHits >= s.limit {
    89  		// The peer has sent too many requests, drop this request.
    90  		return false
    91  	}
    92  
    93  	currentHits[nodeID]++
    94  	return true
    95  }
    96  
    97  func (s *SlidingWindowThrottler) rotate(t time.Time) {
    98  	s.current = 1 - s.current
    99  	s.windows[s.current] = window{
   100  		start: t,
   101  		hits:  make(map[ids.NodeID]float64),
   102  	}
   103  }