github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/unicast/ratelimit/distributor.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/libp2p/go-libp2p/core/peer"
     7  
     8  	"github.com/onflow/flow-go/network/p2p"
     9  )
    10  
    11  // UnicastRateLimiterDistributor subscribes to rate limited peer events from RateLimiters.
    12  type UnicastRateLimiterDistributor struct {
    13  	consumers []p2p.RateLimiterConsumer
    14  	lock      sync.RWMutex
    15  }
    16  
    17  var _ p2p.RateLimiterConsumer = (*UnicastRateLimiterDistributor)(nil)
    18  
    19  // NewUnicastRateLimiterDistributor returns a new UnicastRateLimiterDistributor.
    20  func NewUnicastRateLimiterDistributor() *UnicastRateLimiterDistributor {
    21  	return &UnicastRateLimiterDistributor{
    22  		consumers: make([]p2p.RateLimiterConsumer, 0),
    23  	}
    24  }
    25  
    26  // AddConsumer adds a consumer to the consumers list.
    27  func (r *UnicastRateLimiterDistributor) AddConsumer(consumer p2p.RateLimiterConsumer) {
    28  	r.lock.Lock()
    29  	defer r.lock.Unlock()
    30  	r.consumers = append(r.consumers, consumer)
    31  }
    32  
    33  // OnRateLimitedPeer invokes each consumer callback with the rate limited peer info.
    34  func (r *UnicastRateLimiterDistributor) OnRateLimitedPeer(pid peer.ID, role, msgType, topic, reason string) {
    35  	r.lock.RLock()
    36  	defer r.lock.RUnlock()
    37  	for _, consumer := range r.consumers {
    38  		consumer.OnRateLimitedPeer(pid, role, msgType, topic, reason)
    39  	}
    40  }