github.com/koko1123/flow-go-1@v0.29.6/network/p2p/unicast/ratelimit/internal/limiter_map/rate_limiter_map.go (about)

     1  package limiter_map
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"golang.org/x/time/rate"
     8  
     9  	"github.com/libp2p/go-libp2p/core/peer"
    10  )
    11  
    12  type RateLimiterMetadata struct {
    13  	mu *sync.RWMutex
    14  	// limiter the rate limiter
    15  	limiter *rate.Limiter
    16  	// lastRateLimit the last timestamp this the peer was rate limited.
    17  	lastRateLimit time.Time
    18  	// lastAccessed the last timestamp when the limiter was used. This is used to Cleanup old limiter data
    19  	lastAccessed time.Time
    20  }
    21  
    22  // newRateLimiterMetadata returns a new RateLimiterMetadata
    23  func newRateLimiterMetadata(limiter *rate.Limiter) *RateLimiterMetadata {
    24  	return &RateLimiterMetadata{
    25  		mu:            &sync.RWMutex{},
    26  		limiter:       limiter,
    27  		lastAccessed:  time.Now(),
    28  		lastRateLimit: time.Time{},
    29  	}
    30  }
    31  
    32  // Limiter returns RateLimiterMetadata.limiter..
    33  func (m *RateLimiterMetadata) Limiter() *rate.Limiter {
    34  	m.mu.RLock()
    35  	defer m.mu.RUnlock()
    36  	return m.limiter
    37  }
    38  
    39  // LastRateLimit returns RateLimiterMetadata.lastRateLimit.
    40  func (m *RateLimiterMetadata) LastRateLimit() time.Time {
    41  	m.mu.RLock()
    42  	defer m.mu.RUnlock()
    43  	return m.lastRateLimit
    44  }
    45  
    46  // SetLastRateLimit sets RateLimiterMetadata.lastRateLimit.
    47  func (m *RateLimiterMetadata) SetLastRateLimit(lastRateLimit time.Time) {
    48  	m.mu.Lock()
    49  	defer m.mu.Unlock()
    50  	m.lastRateLimit = lastRateLimit
    51  }
    52  
    53  // LastAccessed returns RateLimiterMetadata.lastAccessed.
    54  func (m *RateLimiterMetadata) LastAccessed() time.Time {
    55  	m.mu.RLock()
    56  	defer m.mu.RUnlock()
    57  	return m.lastAccessed
    58  }
    59  
    60  // SetLastAccessed sets RateLimiterMetadata.lastAccessed.
    61  func (m *RateLimiterMetadata) SetLastAccessed(lastAccessed time.Time) {
    62  	m.mu.Lock()
    63  	defer m.mu.Unlock()
    64  	m.lastAccessed = lastAccessed
    65  }
    66  
    67  // RateLimiterMap stores a RateLimiterMetadata for each peer in an underlying map.
    68  type RateLimiterMap struct {
    69  	mu              sync.RWMutex
    70  	ttl             time.Duration
    71  	cleanupInterval time.Duration
    72  	limiters        map[peer.ID]*RateLimiterMetadata
    73  	done            chan struct{}
    74  }
    75  
    76  func NewLimiterMap(ttl, cleanupInterval time.Duration) *RateLimiterMap {
    77  	return &RateLimiterMap{
    78  		mu:              sync.RWMutex{},
    79  		limiters:        make(map[peer.ID]*RateLimiterMetadata),
    80  		ttl:             ttl,
    81  		cleanupInterval: cleanupInterval,
    82  		done:            make(chan struct{}),
    83  	}
    84  }
    85  
    86  // Get returns limiter in RateLimiterMap map
    87  func (r *RateLimiterMap) Get(peerID peer.ID) (*RateLimiterMetadata, bool) {
    88  	r.mu.RLock()
    89  	defer r.mu.RUnlock()
    90  	if lmtr, ok := r.limiters[peerID]; ok {
    91  		lmtr.SetLastAccessed(time.Now())
    92  		return lmtr, ok
    93  	}
    94  	return nil, false
    95  }
    96  
    97  // Store stores limiter in RateLimiterMap map
    98  func (r *RateLimiterMap) Store(peerID peer.ID, lmtr *rate.Limiter) {
    99  	r.mu.Lock()
   100  	defer r.mu.Unlock()
   101  	r.limiters[peerID] = newRateLimiterMetadata(lmtr)
   102  }
   103  
   104  // UpdateLastRateLimit sets the lastRateLimit field of the RateLimiterMetadata for a peer.
   105  func (r *RateLimiterMap) UpdateLastRateLimit(peerID peer.ID, lastRateLimit time.Time) {
   106  	r.mu.RLock()
   107  	defer r.mu.RUnlock()
   108  	r.limiters[peerID].SetLastRateLimit(lastRateLimit)
   109  }
   110  
   111  // Remove deletes peerID key from underlying map.
   112  func (r *RateLimiterMap) Remove(peerID peer.ID) {
   113  	r.mu.Lock()
   114  	defer r.mu.Unlock()
   115  	r.removeUnlocked(peerID)
   116  }
   117  
   118  // removeUnlocked removes peerID key from underlying map without acquiring a lock.
   119  func (r *RateLimiterMap) removeUnlocked(peerID peer.ID) {
   120  	delete(r.limiters, peerID)
   121  }
   122  
   123  // Cleanup check the TTL for all keys in map and Remove isExpired keys.
   124  func (r *RateLimiterMap) Cleanup() {
   125  	r.mu.Lock()
   126  	defer r.mu.Unlock()
   127  	for peerID, item := range r.limiters {
   128  		if time.Since(item.LastAccessed()) > r.ttl {
   129  			r.removeUnlocked(peerID)
   130  		}
   131  	}
   132  }
   133  
   134  // CleanupLoop starts a loop that periodically removes stale peers.
   135  func (r *RateLimiterMap) CleanupLoop() {
   136  	ticker := time.NewTicker(r.cleanupInterval)
   137  	defer ticker.Stop()
   138  	for {
   139  		select {
   140  		case <-ticker.C:
   141  			r.Cleanup()
   142  		case <-r.done:
   143  			return
   144  		}
   145  	}
   146  }
   147  
   148  // Close will Close the done channel starting the final full Cleanup and stopping the Cleanup loop.
   149  func (r *RateLimiterMap) Close() {
   150  	close(r.done)
   151  }