github.com/cozy/cozy-stack@v0.0.0-20240603063001-31110fa4cae1/pkg/limits/impl_inmemory.go (about)

     1  package limits
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  var counterCleanInterval = 1 * time.Second
     9  
    10  type memRef struct {
    11  	val int64
    12  	exp time.Time
    13  }
    14  
    15  // InMemory implementation ofr [Counter].
    16  type InMemory struct {
    17  	mu   sync.Mutex
    18  	vals map[string]*memRef
    19  }
    20  
    21  // NewInMemory returns a in-memory counter.
    22  func NewInMemory() *InMemory {
    23  	counter := &InMemory{vals: make(map[string]*memRef)}
    24  
    25  	go counter.cleaner()
    26  
    27  	return counter
    28  }
    29  
    30  func (i *InMemory) cleaner() {
    31  	for range time.Tick(counterCleanInterval) {
    32  		i.mu.Lock()
    33  
    34  		now := time.Now()
    35  		for k, v := range i.vals {
    36  			if now.After(v.exp) {
    37  				delete(i.vals, k)
    38  			}
    39  		}
    40  
    41  		i.mu.Unlock()
    42  	}
    43  }
    44  
    45  func (i *InMemory) Increment(key string, timeLimit time.Duration) (int64, error) {
    46  	i.mu.Lock()
    47  	defer i.mu.Unlock()
    48  
    49  	if _, ok := i.vals[key]; !ok {
    50  		i.vals[key] = &memRef{
    51  			val: 0,
    52  			exp: time.Now().Add(timeLimit),
    53  		}
    54  	}
    55  	i.vals[key].val++
    56  	return i.vals[key].val, nil
    57  }
    58  
    59  func (i *InMemory) Reset(key string) error {
    60  	i.mu.Lock()
    61  	defer i.mu.Unlock()
    62  
    63  	delete(i.vals, key)
    64  
    65  	return nil
    66  }