优化ttlcache

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"`

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

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

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

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

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

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

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

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

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