实现请求连接数等限制

internal/utils/expires/id_key_map.go

@@ -0,0 +1,64 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package expires
+
+import "sync"
+
+type IdKeyMap struct {
+	idKeys map[int64]string // id => key
+	keyIds map[string]int64 // key => id
+
+	locker sync.Mutex
+}
+
+func NewIdKeyMap() *IdKeyMap {
+	return &IdKeyMap{
+		idKeys: map[int64]string{},
+		keyIds: map[string]int64{},
+	}
+}
+
+func (this *IdKeyMap) Add(id int64, key string) {
+	oldKey, ok := this.idKeys[id]
+	if ok {
+		delete(this.keyIds, oldKey)
+	}
+
+	oldId, ok := this.keyIds[key]
+	if ok {
+		delete(this.idKeys, oldId)
+	}
+
+	this.idKeys[id] = key
+	this.keyIds[key] = id
+}
+
+func (this *IdKeyMap) Key(id int64) (key string, ok bool) {
+	key, ok = this.idKeys[id]
+	return
+}
+
+func (this *IdKeyMap) Id(key string) (id int64, ok bool) {
+	id, ok = this.keyIds[key]
+	return
+}
+
+func (this *IdKeyMap) DeleteId(id int64) {
+	key, ok := this.idKeys[id]
+	if ok {
+		delete(this.keyIds, key)
+	}
+	delete(this.idKeys, id)
+}
+
+func (this *IdKeyMap) DeleteKey(key string) {
+	id, ok := this.keyIds[key]
+	if ok {
+		delete(this.idKeys, id)
+	}
+	delete(this.keyIds, key)
+}
+
+func (this *IdKeyMap) Len() int {
+	return len(this.idKeys)
+}

internal/utils/expires/id_key_map_test.go

@@ -0,0 +1,46 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package expires
+
+import (
+	"github.com/iwind/TeaGo/assert"
+	"github.com/iwind/TeaGo/logs"
+	"testing"
+)
+
+func TestNewIdKeyMap(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	var m = NewIdKeyMap()
+	m.Add(1, "1")
+	m.Add(1, "2")
+	m.Add(100, "100")
+	logs.PrintAsJSON(m.idKeys, t)
+	logs.PrintAsJSON(m.keyIds, t)
+
+	{
+		k, ok := m.Key(1)
+		a.IsTrue(ok)
+		a.IsTrue(k == "2")
+	}
+
+	{
+		_, ok := m.Key(2)
+		a.IsFalse(ok)
+	}
+
+	m.DeleteKey("2")
+
+	{
+		_, ok := m.Key(1)
+		a.IsFalse(ok)
+	}
+
+	logs.PrintAsJSON(m.idKeys, t)
+	logs.PrintAsJSON(m.keyIds, t)
+
+	m.DeleteId(100)
+
+	logs.PrintAsJSON(m.idKeys, t)
+	logs.PrintAsJSON(m.keyIds, t)
+}

internal/utils/expires/list.go

@@ -27,14 +27,19 @@ func NewList() *List {
 	return list
 }
 
+// Add 添加条目
+// 如果条目已经存在，则覆盖
 func (this *List) Add(itemId int64, expiresAt int64) {
 	this.locker.Lock()
 	defer this.locker.Unlock()
 
 	// 是否已经存在
-	_, ok := this.itemsMap[itemId]
+	oldExpiresAt, ok := this.itemsMap[itemId]
 	if ok {
-		this.removeItem(itemId)
+		if oldExpiresAt == expiresAt {
+			return
+		}
+		delete(this.expireMap, oldExpiresAt)
 	}
 
 	expireItemMap, ok := this.expireMap[expiresAt]
@@ -68,8 +73,9 @@ func (this *List) GC(timestamp int64, callback func(itemId int64)) {
 	}
 }
 
-func (this *List) OnGC(callback func(itemId int64)) {
+func (this *List) OnGC(callback func(itemId int64)) *List {
 	this.gcCallback = callback
+	return this
 }
 
 func (this *List) removeItem(itemId int64) {

internal/utils/expires/list_test.go

@@ -1,9 +1,11 @@
 package expires
 
 import (
+	"github.com/iwind/TeaGo/assert"
 	"github.com/iwind/TeaGo/logs"
 	timeutil "github.com/iwind/TeaGo/utils/time"
 	"math"
+	"runtime"
 	"testing"
 	"time"
 )
@@ -24,12 +26,19 @@ func TestList_Add(t *testing.T) {
 }
 
 func TestList_Add_Overwrite(t *testing.T) {
+	var timestamp = time.Now().Unix()
+
 	list := NewList()
-	list.Add(1, time.Now().Unix()+1)
-	list.Add(1, time.Now().Unix()+1)
-	list.Add(1, time.Now().Unix()+2)
+	list.Add(1, timestamp+1)
+	list.Add(1, timestamp+1)
+	list.Add(1, timestamp+2)
 	logs.PrintAsJSON(list.expireMap, t)
 	logs.PrintAsJSON(list.itemsMap, t)
+
+	var a = assert.NewAssertion(t)
+	a.IsTrue(len(list.itemsMap) == 1)
+	a.IsTrue(len(list.expireMap) == 1)
+	a.IsTrue(list.itemsMap[1] == timestamp+2)
 }
 
 func TestList_Remove(t *testing.T) {
@@ -77,7 +86,10 @@ func TestList_Start_GC(t *testing.T) {
 
 func TestList_ManyItems(t *testing.T) {
 	list := NewList()
-	for i := 0; i < 100_000; i++ {
+	for i := 0; i < 1_000; i++ {
+		list.Add(int64(i), time.Now().Unix())
+	}
+	for i := 0; i < 1_000; i++ {
 		list.Add(int64(i), time.Now().Unix()+1)
 	}
 
@@ -87,35 +99,69 @@ func TestList_ManyItems(t *testing.T) {
 		count++
 	})
 	t.Log("gc", count, "items")
-	t.Log(time.Since(now).Seconds()*1000, "ms")
+	t.Log(time.Now().Sub(now))
 }
 
 func TestList_Map_Performance(t *testing.T) {
 	t.Log("max uint32", math.MaxUint32)
 
+	var timestamp = time.Now().Unix()
+
 	{
 		m := map[int64]int64{}
 		for i := 0; i < 1_000_000; i++ {
-			m[int64(i)] = time.Now().Unix()
+			m[int64(i)] = timestamp
 		}
 
 		now := time.Now()
 		for i := 0; i < 100_000; i++ {
 			delete(m, int64(i))
 		}
-		t.Log(time.Since(now).Seconds()*1000, "ms")
+		t.Log(time.Now().Sub(now))
+	}
+
+	{
+		m := map[uint64]int64{}
+		for i := 0; i < 1_000_000; i++ {
+			m[uint64(i)] = timestamp
+		}
+
+		now := time.Now()
+		for i := 0; i < 100_000; i++ {
+			delete(m, uint64(i))
+		}
+		t.Log(time.Now().Sub(now))
 	}
 
 	{
 		m := map[uint32]int64{}
 		for i := 0; i < 1_000_000; i++ {
-			m[uint32(i)] = time.Now().Unix()
+			m[uint32(i)] = timestamp
 		}
 
 		now := time.Now()
 		for i := 0; i < 100_000; i++ {
 			delete(m, uint32(i))
 		}
-		t.Log(time.Since(now).Seconds()*1000, "ms")
+		t.Log(time.Now().Sub(now))
+	}
+}
+
+func Benchmark_Map_Uint64(b *testing.B) {
+	runtime.GOMAXPROCS(1)
+	var timestamp = uint64(time.Now().Unix())
+
+	var i uint64
+	var count uint64 = 1_000_000
+
+	m := map[uint64]uint64{}
+	for i = 0; i < count; i++ {
+		m[i] = timestamp
+	}
+
+	for n := 0; n < b.N; n++ {
+		for i = 0; i < count; i++ {
+			_ = m[i]
+		}
 	}
 }