48 lines
		
	
	
		
			1.2 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
		
		
			
		
	
	
			48 lines
		
	
	
		
			1.2 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
| 
								 | 
							
								package ratecheck
							 | 
						|||
| 
								 | 
							
								
							 | 
						|||
| 
								 | 
							
								import (
							 | 
						|||
| 
								 | 
							
									"github.com/juju/ratelimit"
							 | 
						|||
| 
								 | 
							
									gocache "github.com/patrickmn/go-cache"
							 | 
						|||
| 
								 | 
							
									"go-admin/pkg/utility"
							 | 
						|||
| 
								 | 
							
									"time"
							 | 
						|||
| 
								 | 
							
								)
							 | 
						|||
| 
								 | 
							
								
							 | 
						|||
| 
								 | 
							
								var (
							 | 
						|||
| 
								 | 
							
									// Map of limiters with TTL
							 | 
						|||
| 
								 | 
							
									tokenBuckets       = gocache.New(120*time.Minute, 1*time.Minute)
							 | 
						|||
| 
								 | 
							
									userDur            = 1 * time.Second
							 | 
						|||
| 
								 | 
							
									userSize     int64 = 1
							 | 
						|||
| 
								 | 
							
								
							 | 
						|||
| 
								 | 
							
									orderDur = 20 * time.Second
							 | 
						|||
| 
								 | 
							
								)
							 | 
						|||
| 
								 | 
							
								
							 | 
						|||
| 
								 | 
							
								// CheckRateLimit 根据key检测,在规定时间内,是否超过访问次数,限流
							 | 
						|||
| 
								 | 
							
								func CheckRateLimit(key string, duration time.Duration, size int64) bool {
							 | 
						|||
| 
								 | 
							
									if _, found := tokenBuckets.Get(key); !found {
							 | 
						|||
| 
								 | 
							
										tokenBuckets.Set(
							 | 
						|||
| 
								 | 
							
											key,
							 | 
						|||
| 
								 | 
							
											ratelimit.NewBucket(duration, size),
							 | 
						|||
| 
								 | 
							
											duration)
							 | 
						|||
| 
								 | 
							
									}
							 | 
						|||
| 
								 | 
							
									expiringMap, found := tokenBuckets.Get(key)
							 | 
						|||
| 
								 | 
							
									if !found {
							 | 
						|||
| 
								 | 
							
										return false
							 | 
						|||
| 
								 | 
							
									}
							 | 
						|||
| 
								 | 
							
									return expiringMap.(*ratelimit.Bucket).TakeAvailable(1) > 0
							 | 
						|||
| 
								 | 
							
								}
							 | 
						|||
| 
								 | 
							
								
							 | 
						|||
| 
								 | 
							
								// CheckUserRateLimit 根据key检测,在规定时间内,单个用户是否超过访问次数,限流,默认5秒1次请求
							 | 
						|||
| 
								 | 
							
								func CheckUserRateLimit(userid int, methodName string) bool {
							 | 
						|||
| 
								 | 
							
									key := methodName + "-" + utility.IntTostring(userid)
							 | 
						|||
| 
								 | 
							
									return CheckRateLimit(key, userDur, userSize)
							 | 
						|||
| 
								 | 
							
								}
							 | 
						|||
| 
								 | 
							
								
							 | 
						|||
| 
								 | 
							
								// 检测订单成交是否重复推送
							 | 
						|||
| 
								 | 
							
								func CheckOrderIdIsExist(tradeId string) bool {
							 | 
						|||
| 
								 | 
							
									_, found := tokenBuckets.Get(tradeId)
							 | 
						|||
| 
								 | 
							
									if !found {
							 | 
						|||
| 
								 | 
							
										tokenBuckets.Set(tradeId, true, orderDur)
							 | 
						|||
| 
								 | 
							
									}
							 | 
						|||
| 
								 | 
							
									return found
							 | 
						|||
| 
								 | 
							
								}
							 |