用户端可以添加WAF 黑白名单

internal/utils/expires/list.go

@@ -0,0 +1,114 @@
+package expires
+
+import (
+	"sync"
+	"time"
+)
+
+type ItemMap = map[int64]bool
+
+type List struct {
+	expireMap map[int64]ItemMap // expires timestamp => map[id]bool
+	itemsMap  map[int64]int64   // itemId => timestamp
+
+	locker sync.Mutex
+}
+
+func NewList() *List {
+	return &List{
+		expireMap: map[int64]ItemMap{},
+		itemsMap:  map[int64]int64{},
+	}
+}
+
+func (this *List) Add(itemId int64, expiredAt int64) {
+	if expiredAt <= time.Now().Unix() {
+		return
+	}
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
+	// 是否已经存在
+	_, ok := this.itemsMap[itemId]
+	if ok {
+		this.removeItem(itemId)
+	}
+
+	expireItemMap, ok := this.expireMap[expiredAt]
+	if ok {
+		expireItemMap[itemId] = true
+	} else {
+		expireItemMap = ItemMap{
+			itemId: true,
+		}
+		this.expireMap[expiredAt] = expireItemMap
+	}
+
+	this.itemsMap[itemId] = expiredAt
+}
+
+func (this *List) Remove(itemId int64) {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+	this.removeItem(itemId)
+}
+
+func (this *List) GC(timestamp int64, callback func(itemId int64)) {
+	this.locker.Lock()
+	itemMap := this.gcItems(timestamp)
+	this.locker.Unlock()
+
+	for itemId := range itemMap {
+		callback(itemId)
+	}
+}
+
+func (this *List) StartGC(callback func(itemId int64)) {
+	ticker := time.NewTicker(1 * time.Second)
+	lastTimestamp := int64(0)
+	for range ticker.C {
+		timestamp := time.Now().Unix()
+		if lastTimestamp == 0 {
+			lastTimestamp = timestamp - 3600
+		}
+
+		// 防止死循环
+		if lastTimestamp > timestamp {
+			continue
+		}
+
+		for i := lastTimestamp; i <= timestamp; i++ {
+			this.GC(timestamp, callback)
+		}
+
+		// 这样做是为了防止系统时钟突变
+		lastTimestamp = timestamp
+	}
+}
+
+func (this *List) removeItem(itemId int64) {
+	expiresAt, ok := this.itemsMap[itemId]
+	if !ok {
+		return
+	}
+	delete(this.itemsMap, itemId)
+
+	expireItemMap, ok := this.expireMap[expiresAt]
+	if ok {
+		delete(expireItemMap, itemId)
+		if len(expireItemMap) == 0 {
+			delete(this.expireMap, expiresAt)
+		}
+	}
+}
+
+func (this *List) gcItems(timestamp int64) ItemMap {
+	expireItemsMap, ok := this.expireMap[timestamp]
+	if ok {
+		for itemId := range expireItemsMap {
+			delete(this.itemsMap, itemId)
+		}
+		delete(this.expireMap, timestamp)
+	}
+	return expireItemsMap
+}

internal/utils/expires/list_test.go

@@ -0,0 +1,115 @@
+package expires
+
+import (
+	"github.com/iwind/TeaGo/logs"
+	timeutil "github.com/iwind/TeaGo/utils/time"
+	"math"
+	"testing"
+	"time"
+)
+
+func TestList_Add(t *testing.T) {
+	list := NewList()
+	list.Add(1, time.Now().Unix())
+	t.Log("===BEFORE===")
+	logs.PrintAsJSON(list.expireMap, t)
+	logs.PrintAsJSON(list.itemsMap, t)
+
+	list.Add(1, time.Now().Unix()+1)
+	list.Add(2, time.Now().Unix()+1)
+	list.Add(3, time.Now().Unix()+2)
+	t.Log("===AFTER===")
+	logs.PrintAsJSON(list.expireMap, t)
+	logs.PrintAsJSON(list.itemsMap, t)
+}
+
+func TestList_Add_Overwrite(t *testing.T) {
+	list := NewList()
+	list.Add(1, time.Now().Unix()+1)
+	list.Add(1, time.Now().Unix()+1)
+	list.Add(1, time.Now().Unix()+2)
+	logs.PrintAsJSON(list.expireMap, t)
+	logs.PrintAsJSON(list.itemsMap, t)
+}
+
+func TestList_Remove(t *testing.T) {
+	list := NewList()
+	list.Add(1, time.Now().Unix()+1)
+	list.Remove(1)
+	logs.PrintAsJSON(list.expireMap, t)
+	logs.PrintAsJSON(list.itemsMap, t)
+}
+
+func TestList_GC(t *testing.T) {
+	list := NewList()
+	list.Add(1, time.Now().Unix()+1)
+	list.Add(2, time.Now().Unix()+1)
+	list.Add(3, time.Now().Unix()+2)
+	list.GC(time.Now().Unix()+2, func(itemId int64) {
+		t.Log("gc:", itemId)
+	})
+	logs.PrintAsJSON(list.expireMap, t)
+	logs.PrintAsJSON(list.itemsMap, t)
+}
+
+func TestList_Start_GC(t *testing.T) {
+	list := NewList()
+	list.Add(1, time.Now().Unix()+1)
+	list.Add(2, time.Now().Unix()+1)
+	list.Add(3, time.Now().Unix()+2)
+	list.Add(4, time.Now().Unix()+5)
+
+	go func() {
+		list.StartGC(func(itemId int64) {
+			t.Log("gc:", itemId, timeutil.Format("H:i:s"))
+			time.Sleep(2 * time.Second)
+		})
+	}()
+
+	time.Sleep(10 * time.Second)
+}
+
+func TestList_ManyItems(t *testing.T) {
+	list := NewList()
+	for i := 0; i < 100_000; i++ {
+		list.Add(int64(i), time.Now().Unix()+1)
+	}
+
+	now := time.Now()
+	count := 0
+	list.GC(time.Now().Unix()+1, func(itemId int64) {
+		count++
+	})
+	t.Log("gc", count, "items")
+	t.Log(time.Since(now).Seconds()*1000, "ms")
+}
+
+func TestList_Map_Performance(t *testing.T) {
+	t.Log("max uint32", math.MaxUint32)
+
+	{
+		m := map[int64]int64{}
+		for i := 0; i < 1_000_000; i++ {
+			m[int64(i)] = time.Now().Unix()
+		}
+
+		now := time.Now()
+		for i := 0; i < 100_000; i++ {
+			delete(m, int64(i))
+		}
+		t.Log(time.Since(now).Seconds()*1000, "ms")
+	}
+
+	{
+		m := map[uint32]int64{}
+		for i := 0; i < 1_000_000; i++ {
+			m[uint32(i)] = time.Now().Unix()
+		}
+
+		now := time.Now()
+		for i := 0; i < 100_000; i++ {
+			delete(m, uint32(i))
+		}
+		t.Log(time.Since(now).Seconds()*1000, "ms")
+	}
+}