79 lines
1.1 KiB
Go
79 lines
1.1 KiB
Go
// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type Bucket struct {
|
|
mu sync.Mutex
|
|
FillInterval time.Duration
|
|
Burst int
|
|
v int
|
|
quitCh chan struct{}
|
|
started bool
|
|
closed bool
|
|
}
|
|
|
|
func (b *Bucket) startLocked() {
|
|
b.v = b.Burst
|
|
b.quitCh = make(chan struct{})
|
|
b.started = true
|
|
|
|
t := time.NewTicker(b.FillInterval)
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-b.quitCh:
|
|
return
|
|
case <-t.C:
|
|
b.tick()
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (b *Bucket) tick() {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
|
|
if b.v < b.Burst {
|
|
b.v++
|
|
}
|
|
}
|
|
|
|
func (b *Bucket) Close() {
|
|
b.mu.Lock()
|
|
if !b.started {
|
|
b.closed = true
|
|
b.mu.Unlock()
|
|
return
|
|
}
|
|
if b.closed {
|
|
b.mu.Unlock()
|
|
return
|
|
}
|
|
b.closed = true
|
|
b.mu.Unlock()
|
|
|
|
b.quitCh <- struct{}{}
|
|
}
|
|
|
|
func (b *Bucket) TryGet() int {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
|
|
if !b.started {
|
|
b.startLocked()
|
|
}
|
|
if b.v > 0 {
|
|
b.v--
|
|
return b.v + 1
|
|
}
|
|
return 0
|
|
}
|