监听器
package listener
import (
"sync"
"time"
)
type Listener struct {
c chan Msg
bmut sync.Mutex
wmut sync.Mutex
ipBlackList map[string]string
ipWhitelist map[string]string
DeltaInvokeCount int
DeltaInvokeTime time.Duration
}
func NewListener(deltaInvokeCount int, deltaInvokeTime time.Duration, disabledIp ...string) *Listener {
l := &Listener{}
m := make(map[string]string)
l.ipBlackList = m
c := make(chan Msg)
l.c = c
l.DeltaInvokeCount = deltaInvokeCount
l.DeltaInvokeTime = deltaInvokeTime
l.DisableIp(disabledIp...)
wl := make(map[string]string)
l.ipWhitelist = wl
return l
}
type Msg struct {
Ip string
t time.Time
}
func (r *Listener) IsWork(ip string) bool {
r.wmut.Lock()
_, ok := r.ipWhitelist[ip]
r.wmut.Unlock()
if ok {
return true
}
r.bmut.Lock()
_, ok = r.ipBlackList[ip]
r.bmut.Unlock()
if ok {
return false
}
msg := Msg{
Ip: ip,
t: time.Now(),
}
r.c <- msg
return true
}
func (r *Listener) Start() {
type invoker struct {
Ip string
Count int
FirstInvokeTime time.Time
}
im := make(map[string]*invoker)
go func() {
for {
t := <-r.c
ip := t.Ip
iv, ok := im[ip]
if !ok {
i := &invoker{
Ip: ip,
Count: 0,
FirstInvokeTime: t.t,
}
im[ip] = i
iv = i
}
iv.Count++
if iv.Count >= r.DeltaInvokeCount {
if t.t.Sub(iv.FirstInvokeTime) <= r.DeltaInvokeTime {
r.DisableIp(ip)
}
iv.Count = 0
iv.FirstInvokeTime = t.t
}
}
}()
}
//Only temporarily removed from the blacklist,
//it is still possible to add to the blacklist again.
//To allow it completely, you must call AddToIpWhiteList.
func (r *Listener) EnableIp(ip ...string) {
for _, i := range ip {
r.bmut.Lock()
delete(r.ipBlackList, i)
r.bmut.Unlock()
}
}
func (r *Listener) DisableIp(ip ...string) {
for _, i := range ip {
r.bmut.Lock()
r.DeleteFromIpWhiteList(i)
r.ipBlackList[i] = i
r.bmut.Unlock()
}
}
func (r *Listener) AddToIpWhiteList(ip ...string) {
for _, i := range ip {
r.wmut.Lock()
r.EnableIp(i)
r.ipWhitelist[i] = i
r.wmut.Unlock()
}
}
func (r *Listener) DeleteFromIpWhiteList(ip ...string) {
for _, i := range ip {
r.wmut.Lock()
delete(r.ipWhitelist, i)
r.wmut.Unlock()
}
}
func (r *Listener) IpWhiteList() []string {
l := make([]string, 0)
r.wmut.Lock()
for _, i := range r.ipWhitelist {
l = append(l, i)
}
r.wmut.Unlock()
return l
}
func (r *Listener) IpBlackList() []string {
l := make([]string, 0)
r.bmut.Lock()
for _, i := range r.ipBlackList {
l = append(l, i)
}
r.bmut.Unlock()
return l
}
使用
func main() {
//单位时间内调用次数限制(也就是调用限制速度),初始禁用ip,可选
//示例,即在10秒内调用次数不可以超过20次,禁用ip1,ip2
lis := listener.NewListener(20, 10*time.Second, "ip1", "ip2")
lis.Start()
time.AfterFunc(10*time.Second, func() {
lis.AddToIpWhiteList("ip1")
fmt.Println("clear", lis.IpWhiteList(), lis.IpBlackList())
})
time.AfterFunc(20*time.Second, func() {
lis.DeleteFromIpWhiteList("ip1")
fmt.Println("clear", lis.IpWhiteList(), lis.IpBlackList())
})
go func() {
for i := 0; i < 100; i++ {
time.Sleep(400 * time.Millisecond)
work("ip1", lis)
}
}()
go func() {
for i := 0; i < 100; i++ {
time.Sleep(400 * time.Millisecond)
work("ip2", lis)
}
}()
//ip3开始小于限制的调用速度不会被禁用,后来速度改变被禁用
go func() {
for i := 0; i < 100; i++ {
if i > 60 {
time.Sleep(400 * time.Millisecond)
} else {
time.Sleep(540 * time.Millisecond)
}
work("ip3", lis)
}
}()
//只用于阻塞主线程,观察输出
time.Sleep(time.Hour)
}
func work(ip string, l *listener.Listener) {
if !l.IsWork(ip) {
fmt.Println("fail", ip)
return
}
fmt.Println("work", ip)
}