gopkg.in/hpcloud/tail.v1@v1.0.0/ratelimiter/memory.go (about)

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