diff --git a/internal/iplibrary/ip_item.go b/internal/iplibrary/ip_item.go index 344daa4..7c21bdc 100644 --- a/internal/iplibrary/ip_item.go +++ b/internal/iplibrary/ip_item.go @@ -13,7 +13,7 @@ const ( // IPItem IP条目 type IPItem struct { Type string `json:"type"` - Id int64 `json:"id"` + Id uint64 `json:"id"` IPFrom uint64 `json:"ipFrom"` IPTo uint64 `json:"ipTo"` ExpiredAt int64 `json:"expiredAt"` diff --git a/internal/iplibrary/ip_list.go b/internal/iplibrary/ip_list.go index d444a1c..5081dc4 100644 --- a/internal/iplibrary/ip_list.go +++ b/internal/iplibrary/ip_list.go @@ -13,9 +13,9 @@ var GlobalWhiteIPList = NewIPList() // IPList IP名单 // TODO IP名单可以分片关闭，这样让每一片的数据量减少，查询更快 type IPList struct { - itemsMap map[int64]*IPItem // id => item + itemsMap map[uint64]*IPItem // id => item sortedItems []*IPItem - allItemsMap map[int64]*IPItem // id => item + allItemsMap map[uint64]*IPItem // id => item expireList *expires.List @@ -24,12 +24,12 @@ type IPList struct { func NewIPList() *IPList { list := &IPList{ - itemsMap: map[int64]*IPItem{}, - allItemsMap: map[int64]*IPItem{}, + itemsMap: map[uint64]*IPItem{}, + allItemsMap: map[uint64]*IPItem{}, } expireList := expires.NewList() - expireList.OnGC(func(itemId int64) { + expireList.OnGC(func(itemId uint64) { list.Delete(itemId) }) list.expireList = expireList @@ -51,7 +51,7 @@ func (this *IPList) Sort() { this.locker.Unlock() } -func (this *IPList) Delete(itemId int64) { +func (this *IPList) Delete(itemId uint64) { this.locker.Lock() this.deleteItem(itemId) this.locker.Unlock() @@ -198,7 +198,7 @@ func (this *IPList) lookupIP(ip uint64) *IPItem { // 在不加锁的情况下删除某个Item // 将会被别的方法引用，切记不能加锁 -func (this *IPList) deleteItem(itemId int64) { +func (this *IPList) deleteItem(itemId uint64) { _, ok := this.itemsMap[itemId] if !ok { return diff --git a/internal/iplibrary/manager_ip_list.go b/internal/iplibrary/manager_ip_list.go index bfa3196..85c122d 100644 --- a/internal/iplibrary/manager_ip_list.go +++ b/internal/iplibrary/manager_ip_list.go @@ -213,7 +213,7 @@ func (this *IPListManager) processItems(items []*pb.IPItem, fromRemote bool) { changedLists[list] = zero.New() if item.IsDeleted { - list.Delete(item.Id) + list.Delete(uint64(item.Id)) // 从WAF名单中删除 waf.SharedIPBlackList.RemoveIP(item.IpFrom, item.ServerId, fromRemote) @@ -227,7 +227,7 @@ func (this *IPListManager) processItems(items []*pb.IPItem, fromRemote bool) { } list.AddDelay(&IPItem{ - Id: item.Id, + Id: uint64(item.Id), Type: item.Type, IPFrom: utils.IP2Long(item.IpFrom), IPTo: utils.IP2Long(item.IpTo), diff --git a/internal/ttlcache/cache.go b/internal/ttlcache/cache.go index 06db1c1..9d1fa35 100644 --- a/internal/ttlcache/cache.go +++ b/internal/ttlcache/cache.go @@ -24,7 +24,7 @@ type Cache struct { } func NewCache(opt ...OptionInterface) *Cache { - var countPieces = 128 + var countPieces = 256 var maxItems = 2_000_000 var totalMemory = utils.SystemMemoryGB() diff --git a/internal/ttlcache/cache_test.go b/internal/ttlcache/cache_test.go index ca2c999..9e2a54a 100644 --- a/internal/ttlcache/cache_test.go +++ b/internal/ttlcache/cache_test.go @@ -1,16 +1,20 @@ package ttlcache import ( + "github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/TeaOSLab/EdgeNode/internal/utils/testutils" + "github.com/iwind/TeaGo/assert" "github.com/iwind/TeaGo/rands" + "github.com/iwind/TeaGo/types" "runtime" "strconv" + "sync/atomic" "testing" "time" ) func TestNewCache(t *testing.T) { - cache := NewCache() + var cache = NewCache() cache.Write("a", 1, time.Now().Unix()+3600) cache.Write("b", 2, time.Now().Unix()+3601) cache.Write("a", 1, time.Now().Unix()+3602) @@ -23,10 +27,10 @@ func TestNewCache(t *testing.T) { } } } - t.Log(cache.Read("a")) + t.Log("a:", cache.Read("a")) time.Sleep(2 * time.Second) - t.Log(cache.Read("d")) - t.Log(cache.Count(), "items") + t.Log("d:", cache.Read("d")) // should be nil + t.Log("left:", cache.Count(), "items") } func TestCache_Memory(t *testing.T) { @@ -40,7 +44,7 @@ func TestCache_Memory(t *testing.T) { t.Log(cache.Count()) - time.Sleep(1 * time.Second) + time.Sleep(10 * time.Second) for i := 0; i < count; i++ { if i%2 == 0 { cache.Delete("a" + strconv.Itoa(i)) @@ -50,27 +54,29 @@ func TestCache_Memory(t *testing.T) { t.Log(cache.Count()) cache.Count() -} -func BenchmarkCache_Add(b *testing.B) { - runtime.GOMAXPROCS(1) - - cache := NewCache() - for i := 0; i < b.N; i++ { - cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(i%1024)) - } + time.Sleep(10 * time.Second) } func TestCache_IncreaseInt64(t *testing.T) { + var a = assert.NewAssertion(t) + var cache = NewCache() + var unixTime = time.Now().Unix() { - cache.IncreaseInt64("a", 1, time.Now().Unix()+3600) - t.Log(cache.Read("a")) + cache.IncreaseInt64("a", 1, unixTime+3600) + var item = cache.Read("a") + t.Log(item) + a.IsTrue(item.Value == int64(1)) + a.IsTrue(item.expiredAt == unixTime+3600) } { - cache.IncreaseInt64("a", 1, time.Now().Unix()+3600+1) - t.Log(cache.Read("a")) + cache.IncreaseInt64("a", 1, unixTime+3600+1) + var item = cache.Read("a") + t.Log(item) + a.IsTrue(item.Value == int64(2)) + a.IsTrue(item.expiredAt == unixTime+3600+1) } { cache.Write("b", 1, time.Now().Unix()+3600+2) @@ -124,11 +130,16 @@ func TestCache_GC(t *testing.T) { for i := 0; i < 20; i++ { cache.GC() t.Log("items:", cache.Count()) + + if cache.Count() == 0 { + break + } time.Sleep(1 * time.Second) } t.Log("now:", time.Now().Unix()) for _, p := range cache.pieces { + t.Log("expire list:", p.expiresList.Count(), p.expiresList) for k, v := range p.m { t.Log(k, v.Value, v.expiredAt) } @@ -171,3 +182,58 @@ func BenchmarkNewCache(b *testing.B) { } }) } + +func BenchmarkCache_Add(b *testing.B) { + runtime.GOMAXPROCS(1) + + var cache = NewCache() + for i := 0; i < b.N; i++ { + cache.Write(strconv.Itoa(i), i, utils.UnixTime()+int64(i%1024)) + } +} + +func BenchmarkCache_Add_Parallel(b *testing.B) { + runtime.GOMAXPROCS(1) + + var cache = NewCache() + var i int64 + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + var j = atomic.AddInt64(&i, 1) + cache.Write(types.String(j), j, utils.UnixTime()+i%1024) + } + }) +} + +func BenchmarkNewCacheGC(b *testing.B) { + runtime.GOMAXPROCS(1) + + var cache = NewCache(NewPiecesOption(1024)) + for i := 0; i < 3_000_000; i++ { + cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 100))) + } + //b.Log(cache.pieces[0].Count()) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + cache.GC() + } + }) +} + +func BenchmarkNewCacheClean(b *testing.B) { + runtime.GOMAXPROCS(1) + + var cache = NewCache(NewPiecesOption(128)) + for i := 0; i < 3_000_000; i++ { + cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(10, 100))) + } + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + cache.Clean() + } + }) +} diff --git a/internal/ttlcache/piece.go b/internal/ttlcache/piece.go index 1871adb..5110bb5 100644 --- a/internal/ttlcache/piece.go +++ b/internal/ttlcache/piece.go @@ -2,36 +2,40 @@ package ttlcache import ( "github.com/TeaOSLab/EdgeNode/internal/utils" + "github.com/TeaOSLab/EdgeNode/internal/utils/expires" "github.com/iwind/TeaGo/types" "sync" "time" ) type Piece struct { - m map[uint64]*Item - maxItems int + m map[uint64]*Item + expiresList *expires.List + maxItems int + lastGCTime int64 locker sync.RWMutex } func NewPiece(maxItems int) *Piece { - return &Piece{m: map[uint64]*Item{}, maxItems: maxItems} + return &Piece{ + m: map[uint64]*Item{}, + expiresList: expires.NewSingletonList(), + maxItems: maxItems, + } } func (this *Piece) Add(key uint64, item *Item) (ok bool) { this.locker.Lock() if len(this.m) >= this.maxItems { - // 尝试先删除过期的 - this.gcWithoutLocker() - - // 仍然是满的就跳过 - if len(this.m) >= this.maxItems { - this.locker.Unlock() - return - } + this.locker.Unlock() + return } this.m[key] = item this.locker.Unlock() + + this.expiresList.Add(key, item.expiredAt) + return true } @@ -42,6 +46,7 @@ func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64) (resu result = types.Int64(item.Value) + delta item.Value = result item.expiredAt = expiredAt + this.expiresList.Add(key, expiredAt) } else { if len(this.m) < this.maxItems { result = delta @@ -49,13 +54,17 @@ func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64) (resu Value: delta, expiredAt: expiredAt, } + this.expiresList.Add(key, expiredAt) } } this.locker.Unlock() + return } func (this *Piece) Delete(key uint64) { + this.expiresList.Remove(key) + this.locker.Lock() delete(this.m, key) this.locker.Unlock() @@ -80,29 +89,48 @@ func (this *Piece) Count() (count int) { } func (this *Piece) GC() { - this.locker.Lock() - this.gcWithoutLocker() - this.locker.Unlock() + var currentTime = time.Now().Unix() + if this.lastGCTime == 0 { + this.lastGCTime = currentTime - 3600 + } + + var min = this.lastGCTime + var max = currentTime + if min > max { + // 过去的时间比现在大，则从这一秒重新开始 + min = max + } + + for i := min; i <= max; i++ { + var itemMap = this.expiresList.GC(i) + if len(itemMap) > 0 { + this.gcItemMap(itemMap) + } + } + + this.lastGCTime = currentTime } func (this *Piece) Clean() { this.locker.Lock() this.m = map[uint64]*Item{} this.locker.Unlock() + + this.expiresList.Clean() } func (this *Piece) Destroy() { this.locker.Lock() this.m = nil this.locker.Unlock() + + this.expiresList.Clean() } -// 不加锁的gc -func (this *Piece) gcWithoutLocker() { - timestamp := time.Now().Unix() - for k, item := range this.m { - if item.expiredAt <= timestamp { - delete(this.m, k) - } +func (this *Piece) gcItemMap(itemMap expires.ItemMap) { + this.locker.Lock() + for key := range itemMap { + delete(this.m, key) } + this.locker.Unlock() } diff --git a/internal/utils/expires/list.go b/internal/utils/expires/list.go index c61caff..96ab4c9 100644 --- a/internal/utils/expires/list.go +++ b/internal/utils/expires/list.go @@ -5,21 +5,24 @@ import ( "sync" ) -type ItemMap = map[int64]zero.Zero +type ItemMap = map[uint64]zero.Zero type List struct { expireMap map[int64]ItemMap // expires timestamp => map[id]ItemMap - itemsMap map[int64]int64 // itemId => timestamp + itemsMap map[uint64]int64 // itemId => timestamp locker sync.Mutex - gcCallback func(itemId int64) + gcCallback func(itemId uint64) + gcBatchCallback func(itemIds ItemMap) + + lastTimestamp int64 } func NewList() *List { var list = &List{ expireMap: map[int64]ItemMap{}, - itemsMap: map[int64]int64{}, + itemsMap: map[uint64]int64{}, } SharedManager.Add(list) @@ -27,12 +30,25 @@ func NewList() *List { return list } +func NewSingletonList() *List { + var list = &List{ + expireMap: map[int64]ItemMap{}, + itemsMap: map[uint64]int64{}, + } + + return list +} + // Add 添加条目 // 如果条目已经存在，则覆盖 -func (this *List) Add(itemId int64, expiresAt int64) { +func (this *List) Add(itemId uint64, expiresAt int64) { this.locker.Lock() defer this.locker.Unlock() + if this.lastTimestamp == 0 || this.lastTimestamp > expiresAt { + this.lastTimestamp = expiresAt + } + // 是否已经存在 oldExpiresAt, ok := this.itemsMap[itemId] if ok { @@ -55,34 +71,61 @@ func (this *List) Add(itemId int64, expiresAt int64) { this.itemsMap[itemId] = expiresAt } -func (this *List) Remove(itemId int64) { +func (this *List) Remove(itemId uint64) { this.locker.Lock() defer this.locker.Unlock() this.removeItem(itemId) } -func (this *List) GC(timestamp int64, callback func(itemId int64)) { +func (this *List) GC(timestamp int64) ItemMap { + if this.lastTimestamp > timestamp+1 { + return nil + } this.locker.Lock() var itemMap = this.gcItems(timestamp) if len(itemMap) == 0 { this.locker.Unlock() - return + return itemMap } this.locker.Unlock() - if callback != nil { + if this.gcCallback != nil { for itemId := range itemMap { - callback(itemId) + this.gcCallback(itemId) } } + if this.gcBatchCallback != nil { + this.gcBatchCallback(itemMap) + } + + return itemMap } -func (this *List) OnGC(callback func(itemId int64)) *List { +func (this *List) Clean() { + this.locker.Lock() + this.itemsMap = map[uint64]int64{} + this.expireMap = map[int64]ItemMap{} + this.locker.Unlock() +} + +func (this *List) Count() int { + this.locker.Lock() + var count = len(this.itemsMap) + this.locker.Unlock() + return count +} + +func (this *List) OnGC(callback func(itemId uint64)) *List { this.gcCallback = callback return this } -func (this *List) removeItem(itemId int64) { +func (this *List) OnGCBatch(callback func(itemMap ItemMap)) *List { + this.gcBatchCallback = callback + return this +} + +func (this *List) removeItem(itemId uint64) { expiresAt, ok := this.itemsMap[itemId] if !ok { return diff --git a/internal/utils/expires/list_test.go b/internal/utils/expires/list_test.go index cb6a2b7..9c5c606 100644 --- a/internal/utils/expires/list_test.go +++ b/internal/utils/expires/list_test.go @@ -1,6 +1,7 @@ package expires import ( + "github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/iwind/TeaGo/assert" "github.com/iwind/TeaGo/logs" timeutil "github.com/iwind/TeaGo/utils/time" @@ -50,13 +51,34 @@ func TestList_Remove(t *testing.T) { } func TestList_GC(t *testing.T) { + var unixTime = time.Now().Unix() + t.Log("unixTime:", unixTime) + + var list = NewList() + list.Add(1, unixTime+1) + list.Add(2, unixTime+1) + list.Add(3, unixTime+2) + list.OnGC(func(itemId uint64) { + t.Log("gc:", itemId) + }) + t.Log("last unixTime:", list.lastTimestamp) + list.GC(time.Now().Unix() + 2) + logs.PrintAsJSON(list.expireMap, t) + logs.PrintAsJSON(list.itemsMap, t) + + t.Log(list.Count()) +} + +func TestList_GC_Batch(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) + list.Add(4, time.Now().Unix()+2) + list.OnGCBatch(func(itemMap ItemMap) { + t.Log("gc:", itemMap) }) + list.GC(time.Now().Unix() + 2) logs.PrintAsJSON(list.expireMap, t) logs.PrintAsJSON(list.itemsMap, t) } @@ -72,7 +94,7 @@ func TestList_Start_GC(t *testing.T) { list.Add(7, time.Now().Unix()+6) list.Add(8, time.Now().Unix()+6) - list.OnGC(func(itemId int64) { + list.OnGC(func(itemId uint64) { t.Log("gc:", itemId, timeutil.Format("H:i:s")) time.Sleep(2 * time.Second) }) @@ -87,17 +109,18 @@ func TestList_Start_GC(t *testing.T) { func TestList_ManyItems(t *testing.T) { list := NewList() for i := 0; i < 1_000; i++ { - list.Add(int64(i), time.Now().Unix()) + list.Add(uint64(i), time.Now().Unix()) } for i := 0; i < 1_000; i++ { - list.Add(int64(i), time.Now().Unix()+1) + list.Add(uint64(i), time.Now().Unix()+1) } now := time.Now() count := 0 - list.GC(time.Now().Unix()+1, func(itemId int64) { + list.OnGC(func(itemId uint64) { count++ }) + list.GC(time.Now().Unix() + 1) t.Log("gc", count, "items") t.Log(time.Now().Sub(now)) } @@ -171,15 +194,23 @@ func BenchmarkList_GC(b *testing.B) { var lists = []*List{} - for i := 0; i < 100; i++ { - lists = append(lists, NewList()) + for m := 0; m < 1_000; m++ { + var list = NewList() + for j := 0; j < 10_000; j++ { + list.Add(uint64(j), utils.UnixTime()+100) + } + lists = append(lists, list) } + b.ResetTimer() + var timestamp = time.Now().Unix() - for i := 0; i < b.N; i++ { - for _, list := range lists { - list.GC(timestamp, nil) + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + for _, list := range lists { + list.GC(timestamp) + } } - } + }) } diff --git a/internal/utils/expires/manager.go b/internal/utils/expires/manager.go index 8ea023d..cbbb360 100644 --- a/internal/utils/expires/manager.go +++ b/internal/utils/expires/manager.go @@ -31,31 +31,32 @@ func NewManager() *Manager { func (this *Manager) init() { var lastTimestamp = int64(0) for range this.ticker.C { - timestamp := time.Now().Unix() + var currentTime = time.Now().Unix() if lastTimestamp == 0 { - lastTimestamp = timestamp - 3600 + lastTimestamp = currentTime - 3600 } - if timestamp >= lastTimestamp { - for i := lastTimestamp; i <= timestamp; i++ { + if currentTime >= lastTimestamp { + for i := lastTimestamp; i <= currentTime; i++ { this.locker.Lock() for list := range this.listMap { - list.GC(i, list.gcCallback) + list.GC(i) } this.locker.Unlock() } } else { - for i := timestamp; i <= lastTimestamp; i++ { + // 如果过去的时间比现在大，则从这一秒重新开始 + for i := currentTime; i <= currentTime; i++ { this.locker.Lock() for list := range this.listMap { - list.GC(i, list.gcCallback) + list.GC(i) } this.locker.Unlock() } } // 这样做是为了防止系统时钟突变 - lastTimestamp = timestamp + lastTimestamp = currentTime } } diff --git a/internal/waf/ip_list.go b/internal/waf/ip_list.go index b2f529b..90c141d 100644 --- a/internal/waf/ip_list.go +++ b/internal/waf/ip_list.go @@ -27,26 +27,26 @@ const IPTypeAll = "*" // IPList IP列表管理 type IPList struct { expireList *expires.List - ipMap map[string]int64 // ip => id - idMap map[int64]string // id => ip + ipMap map[string]uint64 // ip => id + idMap map[uint64]string // id => ip listType IPListType - id int64 + id uint64 locker sync.RWMutex } // NewIPList 获取新对象 func NewIPList(listType IPListType) *IPList { var list = &IPList{ - ipMap: map[string]int64{}, - idMap: map[int64]string{}, + ipMap: map[string]uint64{}, + idMap: map[uint64]string{}, listType: listType, } e := expires.NewList() list.expireList = e - e.OnGC(func(itemId int64) { + e.OnGC(func(itemId uint64) { list.remove(itemId) }) @@ -150,7 +150,7 @@ func (this *IPList) RemoveIP(ip string, serverId int64, shouldExecute bool) { } } -func (this *IPList) remove(id int64) { +func (this *IPList) remove(id uint64) { this.locker.Lock() ip, ok := this.idMap[id] if ok { @@ -163,6 +163,6 @@ func (this *IPList) remove(id int64) { this.locker.Unlock() } -func (this *IPList) nextId() int64 { - return atomic.AddInt64(&this.id, 1) +func (this *IPList) nextId() uint64 { + return atomic.AddUint64(&this.id, 1) }