go ip listener 限流设计（ip黑名单和白名单）

监听器

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)
}