type Counter struct {
	rate  int           // 計數(shù)周期內(nèi)最多允許的請求數(shù)
	begin time.Time     // 計數(shù)開始時間
	cycle time.Duration // 計數(shù)周期
	count int           // 計數(shù)周期內(nèi)累計收到的請求數(shù)
	lock  sync.Mutex
}
func (l *Counter) Allow() bool {
	l.lock.Lock()
	defer l.lock.Unlock()
	if l.count == l.rate-1 {
		now := time.Now()
		if now.Sub(l.begin) >= l.cycle {
			//速度允許范圍內(nèi)， 重置計數(shù)器
			l.Reset(now)
			return true
		} else {
			return false
		}
	} else {
		//沒有達到速率限制，計數(shù)加1
		l.count++
		return true
	}
}
func (l *Counter) Set(r int, cycle time.Duration) {
	l.rate = r
	l.begin = time.Now()
	l.cycle = cycle
	l.count = 0
}
func (l *Counter) Reset(t time.Time) {
	l.begin = t
	l.count = 0
}
func main() {
	var lr Counter
	lr.Set(3, time.Second) // 1s內(nèi)最多請求3次
	var wg sync.WaitGroup
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			if lr.Allow() {
				log.Println("ok:", i)
			} else {
				log.Println("fail:", i)
			}
			wg.Done()
		}(i)
		time.Sleep(200 * time.Millisecond)
	}
	wg.Wait()
}
/* 每秒三個
2021/11/20 16:22:39 ok: 0
2021/11/20 16:22:39 ok: 1
2021/11/20 16:22:39 ok: 2
2021/11/20 16:22:39 fail: 3
2021/11/20 16:22:40 fail: 4
2021/11/20 16:22:40 ok: 5
2021/11/20 16:22:40 ok: 6
2021/11/20 16:22:40 ok: 7
2021/11/20 16:22:41 ok: 8
2021/11/20 16:22:41 fail: 9
*/

var winMu map[string]*sync.RWMutex
func init() {
	winMu = make(map[string]*sync.RWMutex)
}
type timeSlot struct {
	timestamp time.Time // 這個timeSlot的時間起點
	count     int       // 落在這個timeSlot內(nèi)的請求數(shù)
}
func countReq(win []*timeSlot) int {
	var count int
	for _, ts := range win {
		count += ts.count
	}
	return count
}
type SlidingWindowLimiter struct {
	SlotDuration time.Duration // time slot的長度
	WinDuration  time.Duration // sliding window的長度
	numSlots     int           // window內(nèi)最多有多少個slot
	maxReq       int           // win duration內(nèi)允許的最大請求數(shù)
	windows      map[string][]*timeSlot
}
func NewSliding(slotDuration time.Duration, winDuration time.Duration, maxReq int) *SlidingWindowLimiter {
	return &SlidingWindowLimiter{
		SlotDuration: slotDuration,
		WinDuration:  winDuration,
		numSlots:     int(winDuration / slotDuration),
		windows:      make(map[string][]*timeSlot),
		maxReq:       maxReq,
	}
}
// 獲取user_id/ip的時間窗口
func (l *SlidingWindowLimiter) getWindow(uidOrIp string) []*timeSlot {
	win, ok := l.windows[uidOrIp]
	if !ok {
		win = make([]*timeSlot, 0, l.numSlots)
	}
	return win
}
func (l *SlidingWindowLimiter) storeWindow(uidOrIp string, win []*timeSlot) {
	l.windows[uidOrIp] = win
}
func (l *SlidingWindowLimiter) validate(uidOrIp string) bool {
	// 同一user_id/ip并發(fā)安全
	mu, ok := winMu[uidOrIp]
	if !ok {
		var m sync.RWMutex
		mu = &m
		winMu[uidOrIp] = mu
	}
	mu.Lock()
	defer mu.Unlock()
	win := l.getWindow(uidOrIp)
	now := time.Now()
	// 已經(jīng)過期的time slot移出時間窗
	timeoutOffset := -1
	for i, ts := range win {
		if ts.timestamp.Add(l.WinDuration).After(now) {
			break
		}
		timeoutOffset = i
	}
	if timeoutOffset > -1 {
		win = win[timeoutOffset+1:]
	}
	// 判斷請求是否超限
	var result bool
	if countReq(win) < l.maxReq {
		result = true
	}
	// 記錄這次的請求數(shù)
	var lastSlot *timeSlot
	if len(win) > 0 {
		lastSlot = win[len(win)-1]
		if lastSlot.timestamp.Add(l.SlotDuration).Before(now) {
			lastSlot = &timeSlot{timestamp: now, count: 1}
			win = append(win, lastSlot)
		} else {
			lastSlot.count++
		}
	} else {
		lastSlot = &timeSlot{timestamp: now, count: 1}
		win = append(win, lastSlot)
	}
	l.storeWindow(uidOrIp, win)
	return result
}
func (l *SlidingWindowLimiter) getUidOrIp() string {
	return "127.0.0.1"
}
func (l *SlidingWindowLimiter) IsLimited() bool {
	return !l.validate(l.getUidOrIp())
}
func main() {
	limiter := NewSliding(100*time.Millisecond, time.Second, 10)
	for i := 0; i < 5; i++ {
		fmt.Println(limiter.IsLimited())
	}
	time.Sleep(100 * time.Millisecond)
	for i := 0; i < 5; i++ {
		fmt.Println(limiter.IsLimited())
	}
	fmt.Println(limiter.IsLimited())
	for _, v := range limiter.windows[limiter.getUidOrIp()] {
		fmt.Println(v.timestamp, v.count)
	}
	fmt.Println("a thousand years later...")
	time.Sleep(time.Second)
	for i := 0; i < 7; i++ {
		fmt.Println(limiter.IsLimited())
	}
	for _, v := range limiter.windows[limiter.getUidOrIp()] {
		fmt.Println(v.timestamp, v.count)
	}
}

// 封裝業(yè)務(wù)邏輯的執(zhí)行結(jié)果
type Result struct {
	Msg string
}
// 執(zhí)行的業(yè)務(wù)邏輯函數(shù)
type Handler func() Result
// 每個請求來了，把需要執(zhí)行的業(yè)務(wù)邏輯封裝成Task，放入木桶，等待worker取出執(zhí)行
type Task struct {
	handler Handler     // worker從木桶中取出請求對象后要執(zhí)行的業(yè)務(wù)邏輯函數(shù)
	resChan chan Result // 等待worker執(zhí)行并返回結(jié)果的channel
	taskID  int
}
func NewTask(id int, handler Handler) Task {
	return Task{
		handler: handler,
		resChan: make(chan Result),
		taskID:  id,
	}
}
// 漏桶
type LeakyBucket struct {
	BucketSize int       // 木桶的大小
	WorkerNum  int       // 同時從木桶中獲取任務(wù)執(zhí)行的worker數(shù)量
	bucket     chan Task // 存方任務(wù)的木桶
}
func NewLeakyBucket(bucketSize int, workNum int) *LeakyBucket {
	return &LeakyBucket{
		BucketSize: bucketSize,
		WorkerNum:  workNum,
		bucket:     make(chan Task, bucketSize),
	}
}
func (b *LeakyBucket) AddTask(task Task) bool {
	// 如果木桶已經(jīng)滿了，返回false
	select {
	case b.bucket <- task:
	default:
		fmt.Printf("request[id=%d] is refused\n", task.taskID)
		return false
	}
	// 如果成功入桶，調(diào)用者會等待worker執(zhí)行結(jié)果
	resp := <-task.resChan
	fmt.Printf("request[id=%d] is run ok, resp[%v]\n", task.taskID, resp)
	return true
}
func (b *LeakyBucket) Start(ctx context.Context) {
	// 開啟worker從木桶拉取任務(wù)執(zhí)行
	for i := 0; i < b.WorkerNum; i++ {
		go func(ctx context.Context) {
			for {
				select {
				case <-ctx.Done():
					return
				default:
					task := <-b.bucket
					result := task.handler()
					task.resChan <- result
				}
			}
		}(ctx)
	}
}
func main() {
	bucket := NewLeakyBucket(10, 4)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bucket.Start(ctx) // 開啟消費者
  	// 模擬20個并發(fā)請求
	var wg sync.WaitGroup
	wg.Add(20)
	for i := 0; i < 20; i++ {
		go func(id int) {
			defer wg.Done()
			task := NewTask(id, func() Result {
				time.Sleep(300 * time.Millisecond)
				return Result{}
			})
			bucket.AddTask(task)
		}(i)
	}
	wg.Wait()
	time.Sleep(10 * time.Second)
}

type TokenBucket struct {
	lock sync.Mutex
	rate     time.Duration // 多長時間生成一個令牌
	capacity int           // 桶的容量
	tokens   int           // 桶中當前token數(shù)量
	last     time.Time     // 桶上次放token的時間戳 s
}
func NewTokenBucket(bucketSize int, tokenRate time.Duration) *TokenBucket {
	return &TokenBucket{
		capacity: bucketSize,
		rate:     tokenRate,
	}
}
// 驗證是否能獲取一個令牌 返回是否被限流
func (t *TokenBucket) Allow() bool {
	t.lock.Lock()
	defer t.lock.Unlock()
	now := time.Now()
	if t.last.IsZero() {
		// 第一次訪問初始化為最大令牌數(shù)
		t.last, t.tokens = now, t.capacity
	} else {
		if t.last.Add(t.rate).Before(now) {
			// 如果 now 與上次請求的間隔超過了 token rate
			// 則增加令牌，更新last
			t.tokens += int(now.Sub(t.last) / t.rate)
			if t.tokens > t.capacity {
				t.tokens = t.capacity
			}
			t.last = now
		}
	}
	if t.tokens > 0 {
		// 如果令牌數(shù)大于0，取走一個令牌
		t.tokens--
		return true
	}
	// 沒有令牌,則拒絕
	return false
}
func main() {
	tokenBucket := NewTokenBucket(5, 100*time.Millisecond)
	for i := 0; i < 10; i++ {
		fmt.Println(tokenBucket.Allow())
	}
	time.Sleep(100 * time.Millisecond)
	fmt.Println(tokenBucket.Allow())
}