github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/framework/ratelimit/memrate/memorystore.go (about)

     1  package memrate
     2  
     3  import (
     4  	"context"
     5  	lru "github.com/hashicorp/golang-lru"
     6  	"github.com/hashicorp/golang-lru/simplelru"
     7  	logger "github.com/unionj-cloud/go-doudou/toolkit/zlogger"
     8  	"github.com/unionj-cloud/go-doudou/framework/ratelimit"
     9  	"sync"
    10  )
    11  
    12  const defaultMaxKeys = 256
    13  
    14  type LimiterFn func(ctx context.Context, store *MemoryStore, key string) ratelimit.Limiter
    15  
    16  type MemoryStore struct {
    17  	keys      *lru.Cache
    18  	maxKeys   int
    19  	onEvicted simplelru.EvictCallback
    20  	limiterFn LimiterFn
    21  	mu        sync.RWMutex
    22  }
    23  
    24  type MemoryStoreOption func(*MemoryStore)
    25  
    26  // WithMaxKeys set maxKeys
    27  func WithMaxKeys(maxKeys int) MemoryStoreOption {
    28  	return func(ls *MemoryStore) {
    29  		ls.maxKeys = maxKeys
    30  	}
    31  }
    32  
    33  // WithOnEvicted set onEvicted
    34  func WithOnEvicted(onEvicted func(key interface{}, value interface{})) MemoryStoreOption {
    35  	return func(ls *MemoryStore) {
    36  		ls.onEvicted = onEvicted
    37  	}
    38  }
    39  
    40  func NewMemoryStore(fn LimiterFn, opts ...MemoryStoreOption) *MemoryStore {
    41  	store := &MemoryStore{
    42  		maxKeys:   defaultMaxKeys,
    43  		limiterFn: fn,
    44  	}
    45  
    46  	for _, opt := range opts {
    47  		opt(store)
    48  	}
    49  
    50  	if store.onEvicted != nil {
    51  		store.keys, _ = lru.NewWithEvict(store.maxKeys, store.onEvicted)
    52  	} else {
    53  		store.keys, _ = lru.New(store.maxKeys)
    54  	}
    55  
    56  	return store
    57  }
    58  
    59  // GetLimiter returns the rate limiter for the provided key if it exists,
    60  // otherwise calls addKey to add key to the map
    61  func (store *MemoryStore) GetLimiter(key string) ratelimit.Limiter {
    62  	return store.GetLimiterCtx(context.Background(), key)
    63  }
    64  
    65  func (store *MemoryStore) addKeyCtx(ctx context.Context, key string) ratelimit.Limiter {
    66  	store.mu.Lock()
    67  	defer store.mu.Unlock()
    68  
    69  	limiter, exists := store.keys.Get(key)
    70  	if exists {
    71  		return limiter.(ratelimit.Limiter)
    72  	}
    73  
    74  	limiter = store.limiterFn(ctx, store, key)
    75  	store.keys.Add(key, limiter)
    76  
    77  	return limiter.(ratelimit.Limiter)
    78  }
    79  
    80  // GetLimiterCtx returns the rate limiter for the provided key if it exists,
    81  // otherwise calls addKey to add key to the map
    82  func (store *MemoryStore) GetLimiterCtx(ctx context.Context, key string) ratelimit.Limiter {
    83  	store.mu.RLock()
    84  
    85  	limiter, exists := store.keys.Get(key)
    86  	if !exists {
    87  		store.mu.RUnlock()
    88  		return store.addKeyCtx(ctx, key)
    89  	}
    90  
    91  	store.mu.RUnlock()
    92  	return limiter.(ratelimit.Limiter)
    93  }
    94  
    95  func (store *MemoryStore) DeleteKey(key string) {
    96  	store.mu.Lock()
    97  	defer store.mu.Unlock()
    98  
    99  	store.keys.Remove(key)
   100  	logger.Debug().Msgf("[go-doudou] key %s is deleted from store", key)
   101  }