github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/collection/ingest/rate_limiter.go (about)

     1  package ingest
     2  
     3  import (
     4  	"strings"
     5  	"sync"
     6  
     7  	"golang.org/x/time/rate"
     8  
     9  	"github.com/onflow/flow-go/model/flow"
    10  )
    11  
    12  // AddressRateLimiter limits the rate of ingested transactions with a given payer address.
    13  type AddressRateLimiter struct {
    14  	mu       sync.RWMutex
    15  	limiters map[flow.Address]*rate.Limiter
    16  	limit    rate.Limit // X messages allowed per second
    17  	burst    int        // X messages allowed at one time
    18  }
    19  
    20  // AddressRateLimiter limits the rate of ingested transactions with a given payer address.
    21  // It allows the given "limit" amount messages per second with a "burst" amount of messages to be sent at once
    22  //
    23  // for example,
    24  // To config 1 message per 100 milliseconds, convert to per second first, which is 10 message per second,
    25  // so limit is 10 ( rate.Limit(10) ), and burst is 1.
    26  // Note: rate.Limit(0.1), burst = 1 means 1 message per 10 seconds, instead of 1 message per 100 milliseconds.
    27  //
    28  // To config 3 message per minute, the per-second-basis is 0.05 (3/60), so the limit should be rate.Limit(0.05),
    29  // and burst is 3.
    30  //
    31  // Note: The rate limit configured for each node may differ from the effective network-wide rate limit
    32  // for a given payer. In particular, the number of clusters and the message propagation factor will
    33  // influence how the individual rate limit translates to a network-wide rate limit.
    34  // For example, suppose we have 5 collection clusters and configure each Collection Node with a rate
    35  // limit of 1 message per second. Then, the effective network-wide rate limit for a payer address would
    36  // be *at least* 5 messages per second.
    37  func NewAddressRateLimiter(limit rate.Limit, burst int) *AddressRateLimiter {
    38  	return &AddressRateLimiter{
    39  		limiters: make(map[flow.Address]*rate.Limiter),
    40  		limit:    limit,
    41  		burst:    burst,
    42  	}
    43  }
    44  
    45  // Allow returns whether the given address should be allowed (not rate limited)
    46  func (r *AddressRateLimiter) Allow(address flow.Address) bool {
    47  	return !r.IsRateLimited(address)
    48  }
    49  
    50  // IsRateLimited returns whether the given address should be rate limited
    51  func (r *AddressRateLimiter) IsRateLimited(address flow.Address) bool {
    52  	r.mu.RLock()
    53  	limiter, ok := r.limiters[address]
    54  	r.mu.RUnlock()
    55  
    56  	if !ok {
    57  		return false
    58  	}
    59  
    60  	rateLimited := !limiter.Allow()
    61  	return rateLimited
    62  }
    63  
    64  // AddAddress add an address to be rate limited
    65  func (r *AddressRateLimiter) AddAddress(address flow.Address) {
    66  	r.mu.Lock()
    67  	defer r.mu.Unlock()
    68  
    69  	_, ok := r.limiters[address]
    70  	if ok {
    71  		return
    72  	}
    73  
    74  	r.limiters[address] = rate.NewLimiter(r.limit, r.burst)
    75  }
    76  
    77  // RemoveAddress remove an address for being rate limited
    78  func (r *AddressRateLimiter) RemoveAddress(address flow.Address) {
    79  	r.mu.Lock()
    80  	defer r.mu.Unlock()
    81  
    82  	delete(r.limiters, address)
    83  }
    84  
    85  // GetAddresses get the list of rate limited address
    86  func (r *AddressRateLimiter) GetAddresses() []flow.Address {
    87  	r.mu.RLock()
    88  	defer r.mu.RUnlock()
    89  
    90  	addresses := make([]flow.Address, 0, len(r.limiters))
    91  	for address := range r.limiters {
    92  		addresses = append(addresses, address)
    93  	}
    94  
    95  	return addresses
    96  }
    97  
    98  // GetLimitConfig get the limit config
    99  func (r *AddressRateLimiter) GetLimitConfig() (rate.Limit, int) {
   100  	r.mu.RLock()
   101  	defer r.mu.RUnlock()
   102  	return r.limit, r.burst
   103  }
   104  
   105  // SetLimitConfig update the limit config
   106  // Note all the existing limiters will be updated, and reset
   107  func (r *AddressRateLimiter) SetLimitConfig(limit rate.Limit, burst int) {
   108  	r.mu.Lock()
   109  	defer r.mu.Unlock()
   110  
   111  	for address := range r.limiters {
   112  		r.limiters[address] = rate.NewLimiter(limit, burst)
   113  	}
   114  
   115  	r.limit = limit
   116  	r.burst = burst
   117  }
   118  
   119  // Util functions
   120  func AddAddresses(r *AddressRateLimiter, addresses []flow.Address) {
   121  	for _, address := range addresses {
   122  		r.AddAddress(address)
   123  	}
   124  }
   125  
   126  func RemoveAddresses(r *AddressRateLimiter, addresses []flow.Address) {
   127  	for _, address := range addresses {
   128  		r.RemoveAddress(address)
   129  	}
   130  }
   131  
   132  // parse addresses string into a list of flow addresses
   133  func ParseAddresses(addresses string) ([]flow.Address, error) {
   134  	addressList := make([]flow.Address, 0)
   135  	for _, addr := range strings.Split(addresses, ",") {
   136  		addr = strings.TrimSpace(addr)
   137  		if addr == "" {
   138  			continue
   139  		}
   140  		flowAddr, err := flow.StringToAddress(addr)
   141  		if err != nil {
   142  			return nil, err
   143  		}
   144  		addressList = append(addressList, flowAddr)
   145  	}
   146  	return addressList, nil
   147  }