github.com/BGrewell/tail@v1.0.1-0.20210309152823-689d25348e0e/ratelimiter/memory.go (about)

     1  package ratelimiter
     2  
     3  import (
     4  	"errors"
     5  	"time"
     6  )
     7  
     8  const (
     9  	GC_SIZE   int           = 100
    10  	GC_PERIOD time.Duration = 60 * time.Second
    11  )
    12  
    13  type Memory struct {
    14  	store           map[string]LeakyBucket
    15  	lastGCCollected time.Time
    16  }
    17  
    18  func NewMemory() *Memory {
    19  	m := new(Memory)
    20  	m.store = make(map[string]LeakyBucket)
    21  	m.lastGCCollected = time.Now()
    22  	return m
    23  }
    24  
    25  func (m *Memory) GetBucketFor(key string) (*LeakyBucket, error) {
    26  
    27  	bucket, ok := m.store[key]
    28  	if !ok {
    29  		return nil, errors.New("miss")
    30  	}
    31  
    32  	return &bucket, nil
    33  }
    34  
    35  func (m *Memory) SetBucketFor(key string, bucket LeakyBucket) error {
    36  
    37  	if len(m.store) > GC_SIZE {
    38  		m.GarbageCollect()
    39  	}
    40  
    41  	m.store[key] = bucket
    42  
    43  	return nil
    44  }
    45  
    46  func (m *Memory) GarbageCollect() {
    47  	now := time.Now()
    48  
    49  	// rate limit GC to once per minute
    50  	if now.Unix() >= m.lastGCCollected.Add(GC_PERIOD).Unix() {
    51  		for key, bucket := range m.store {
    52  			// if the bucket is drained, then GC
    53  			if bucket.DrainedAt().Unix() < now.Unix() {
    54  				delete(m.store, key)
    55  			}
    56  		}
    57  
    58  		m.lastGCCollected = now
    59  	}
    60  }