github.com/godaddy-x/freego@v1.0.156/cache/limiter/rate_limiter_local.go (about)

     1  package rate
     2  
     3  import (
     4  	"github.com/godaddy-x/freego/cache"
     5  	"github.com/godaddy-x/freego/zlog"
     6  	"sync"
     7  )
     8  
     9  type RateLimiter interface {
    10  	Allow(resource string) bool // true=接受请求 false=拒绝请求
    11  }
    12  
    13  type LocalRateLimiter struct {
    14  	mu     sync.Mutex
    15  	cache  cache.Cache
    16  	option Option
    17  }
    18  
    19  type Option struct {
    20  	Limit       float64
    21  	Bucket      int
    22  	Expire      int
    23  	Distributed bool
    24  }
    25  
    26  func NewRateLimiter(option Option) RateLimiter {
    27  	if option.Distributed {
    28  		return &RedisRateLimiter{option: option}
    29  	}
    30  	return &LocalRateLimiter{cache: new(cache.LocalMapManager).NewCache(30, 3), option: option}
    31  }
    32  
    33  // key=过滤关键词 limit=速率 bucket=容量 expire=过期时间/秒
    34  func (self *LocalRateLimiter) getLimiter(resource string) *Limiter {
    35  	if len(resource) == 0 {
    36  		return nil
    37  	}
    38  	var limiter *Limiter
    39  	if v, b, _ := self.cache.Get(resource, nil); b {
    40  		limiter = v.(*Limiter)
    41  	}
    42  	if limiter == nil {
    43  		self.mu.Lock()
    44  		if v, b, _ := self.cache.Get(resource, nil); b {
    45  			limiter = v.(*Limiter)
    46  		}
    47  		if limiter == nil {
    48  			limiter = NewLimiter(Limit(self.option.Limit), self.option.Bucket)
    49  			if err := self.cache.Put(resource, limiter, self.option.Expire); err != nil {
    50  				zlog.Error("cache put failed", 0, zlog.AddError(err))
    51  			}
    52  		}
    53  		self.mu.Unlock()
    54  	}
    55  	return limiter
    56  }
    57  
    58  func (self *LocalRateLimiter) Allow(resource string) bool {
    59  	limiter := self.getLimiter(resource)
    60  	if limiter == nil {
    61  		return false
    62  	}
    63  	return limiter.Allow()
    64  }