diff --git a/internal/caches/list_file.go b/internal/caches/list_file.go index 0fd9e13..7d19e38 100644 --- a/internal/caches/list_file.go +++ b/internal/caches/list_file.go @@ -27,7 +27,7 @@ type FileList struct { onAdd func(item *Item) onRemove func(item *Item) - memoryCache *ttlcache.Cache + memoryCache *ttlcache.Cache[bool] // 老数据库地址 oldDir string @@ -36,7 +36,7 @@ type FileList struct { func NewFileList(dir string) ListInterface { return &FileList{ dir: dir, - memoryCache: ttlcache.NewCache(), + memoryCache: ttlcache.NewCache[bool](), } } @@ -120,7 +120,7 @@ func (this *FileList) Add(hash string, item *Item) error { return err } - this.memoryCache.Write(hash, 1, this.maxExpiresAtForMemoryCache(item.ExpiredAt)) + this.memoryCache.Write(hash, true, this.maxExpiresAtForMemoryCache(item.ExpiredAt)) if this.onAdd != nil { this.onAdd(item) @@ -158,7 +158,7 @@ func (this *FileList) Exist(hash string) (bool, error) { } return false, err } - this.memoryCache.Write(hash, 1, this.maxExpiresAtForMemoryCache(expiredAt)) + this.memoryCache.Write(hash, false, this.maxExpiresAtForMemoryCache(expiredAt)) return true, nil } diff --git a/internal/nodes/http_request_mismatch.go b/internal/nodes/http_request_mismatch.go index acad224..54cd976 100644 --- a/internal/nodes/http_request_mismatch.go +++ b/internal/nodes/http_request_mismatch.go @@ -63,7 +63,7 @@ func (this *HTTPRequest) doMismatch() { // 要考虑到服务在切换集群时,域名未生效状态时,用户访问的仍然是老集群中的节点,就会产生找不到域名的情况 if len(remoteIP) > 0 { const maxAttempts = 100 - if ttlcache.SharedCache.IncreaseInt64("MISMATCH_DOMAIN:"+remoteIP, int64(1), time.Now().Unix()+60, false) > maxAttempts { + if ttlcache.SharedInt64Cache.IncreaseInt64("MISMATCH_DOMAIN:"+remoteIP, int64(1), time.Now().Unix()+60, false) > maxAttempts { // 在加入之前再次检查黑名单 if !waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteIP) { waf.SharedIPBlackList.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteIP, time.Now().Unix()+3600) diff --git a/internal/ttlcache/cache.go b/internal/ttlcache/cache.go index c709ae6..a5accb8 100644 --- a/internal/ttlcache/cache.go +++ b/internal/ttlcache/cache.go @@ -5,7 +5,7 @@ import ( "github.com/TeaOSLab/EdgeNode/internal/utils/fasttime" ) -var SharedCache = NewBigCache() +var SharedInt64Cache = NewBigCache[int64]() // Cache TTL缓存 // 最大的缓存时间为30 * 86400 @@ -13,24 +13,24 @@ var SharedCache = NewBigCache() // // Piece1 | Piece2 | Piece3 | ... // [ Item1, Item2, ... ] | ... -type Cache struct { +type Cache[T any] struct { isDestroyed bool - pieces []*Piece + pieces []*Piece[T] countPieces uint64 maxItems int gcPieceIndex int } -func NewBigCache() *Cache { +func NewBigCache[T any]() *Cache[T] { var delta = utils.SystemMemoryGB() / 2 if delta <= 0 { delta = 1 } - return NewCache(NewMaxItemsOption(delta * 1_000_000)) + return NewCache[T](NewMaxItemsOption(delta * 1_000_000)) } -func NewCache(opt ...OptionInterface) *Cache { +func NewCache[T any](opt ...OptionInterface) *Cache[T] { var countPieces = 256 var maxItems = 1_000_000 @@ -61,13 +61,13 @@ func NewCache(opt ...OptionInterface) *Cache { } } - var cache = &Cache{ + var cache = &Cache[T]{ countPieces: uint64(countPieces), maxItems: maxItems, } for i := 0; i < countPieces; i++ { - cache.pieces = append(cache.pieces, NewPiece(maxItems/countPieces)) + cache.pieces = append(cache.pieces, NewPiece[T](maxItems/countPieces)) } // Add to manager @@ -76,7 +76,7 @@ func NewCache(opt ...OptionInterface) *Cache { return cache } -func (this *Cache) Write(key string, value any, expiredAt int64) (ok bool) { +func (this *Cache[T]) Write(key string, value T, expiredAt int64) (ok bool) { if this.isDestroyed { return } @@ -92,20 +92,20 @@ func (this *Cache) Write(key string, value any, expiredAt int64) (ok bool) { } var uint64Key = HashKey([]byte(key)) var pieceIndex = uint64Key % this.countPieces - return this.pieces[pieceIndex].Add(uint64Key, &Item{ + return this.pieces[pieceIndex].Add(uint64Key, &Item[T]{ Value: value, expiredAt: expiredAt, }) } -func (this *Cache) IncreaseInt64(key string, delta int64, expiredAt int64, extend bool) int64 { +func (this *Cache[T]) IncreaseInt64(key string, delta T, expiredAt int64, extend bool) T { if this.isDestroyed { - return 0 + return any(0).(T) } var currentTimestamp = fasttime.Now().Unix() if expiredAt <= currentTimestamp { - return 0 + return any(0).(T) } var maxExpiredAt = currentTimestamp + 30*86400 @@ -117,24 +117,24 @@ func (this *Cache) IncreaseInt64(key string, delta int64, expiredAt int64, exten return this.pieces[pieceIndex].IncreaseInt64(uint64Key, delta, expiredAt, extend) } -func (this *Cache) Read(key string) (item *Item) { +func (this *Cache[T]) Read(key string) (item *Item[T]) { var uint64Key = HashKey([]byte(key)) return this.pieces[uint64Key%this.countPieces].Read(uint64Key) } -func (this *Cache) Delete(key string) { +func (this *Cache[T]) Delete(key string) { var uint64Key = HashKey([]byte(key)) this.pieces[uint64Key%this.countPieces].Delete(uint64Key) } -func (this *Cache) Count() (count int) { +func (this *Cache[T]) Count() (count int) { for _, piece := range this.pieces { count += piece.Count() } return } -func (this *Cache) GC() { +func (this *Cache[T]) GC() { var index = this.gcPieceIndex const maxPiecesPerGC = 4 for i := index; i < index+maxPiecesPerGC; i++ { @@ -151,13 +151,13 @@ func (this *Cache) GC() { this.gcPieceIndex = index } -func (this *Cache) Clean() { +func (this *Cache[T]) Clean() { for _, piece := range this.pieces { piece.Clean() } } -func (this *Cache) Destroy() { +func (this *Cache[T]) Destroy() { SharedManager.Remove(this) this.isDestroyed = true diff --git a/internal/ttlcache/cache_test.go b/internal/ttlcache/cache_test.go index f093b9c..7a83bdd 100644 --- a/internal/ttlcache/cache_test.go +++ b/internal/ttlcache/cache_test.go @@ -14,7 +14,7 @@ import ( ) func TestNewCache(t *testing.T) { - var cache = NewCache() + var cache = NewCache[int]() cache.Write("a", 1, time.Now().Unix()+3600) cache.Write("b", 2, time.Now().Unix()+1) cache.Write("c", 1, time.Now().Unix()+3602) @@ -48,7 +48,7 @@ func TestCache_Memory(t *testing.T) { testutils.StartMemoryStats(t) - var cache = NewCache() + var cache = NewCache[int]() var count = 20_000_000 for i := 0; i < count; i++ { cache.Write("a"+strconv.Itoa(i), 1, time.Now().Unix()+3600) @@ -73,21 +73,21 @@ func TestCache_Memory(t *testing.T) { func TestCache_IncreaseInt64(t *testing.T) { var a = assert.NewAssertion(t) - var cache = NewCache() + var cache = NewCache[int64]() var unixTime = time.Now().Unix() { cache.IncreaseInt64("a", 1, unixTime+3600, false) var item = cache.Read("a") t.Log(item) - a.IsTrue(item.Value == int64(1)) + a.IsTrue(item.Value == 1) a.IsTrue(item.expiredAt == unixTime+3600) } { cache.IncreaseInt64("a", 1, unixTime+3600+1, true) var item = cache.Read("a") t.Log(item) - a.IsTrue(item.Value == int64(2)) + a.IsTrue(item.Value == 2) a.IsTrue(item.expiredAt == unixTime+3600+1) } { @@ -103,7 +103,7 @@ func TestCache_IncreaseInt64(t *testing.T) { func TestCache_Read(t *testing.T) { runtime.GOMAXPROCS(1) - var cache = NewCache(PiecesOption{Count: 32}) + var cache = NewCache[int](PiecesOption{Count: 32}) for i := 0; i < 10_000_000; i++ { cache.Write("HELLO_WORLD_"+strconv.Itoa(i), i, time.Now().Unix()+int64(i%10240)+1) @@ -125,7 +125,7 @@ func TestCache_Read(t *testing.T) { } func TestCache_GC(t *testing.T) { - var cache = NewCache(&PiecesOption{Count: 5}) + var cache = NewCache[int](&PiecesOption{Count: 5}) cache.Write("a", 1, time.Now().Unix()+1) cache.Write("b", 2, time.Now().Unix()+2) cache.Write("c", 3, time.Now().Unix()+3) @@ -161,12 +161,12 @@ func TestCache_GC(t *testing.T) { func TestCache_GC2(t *testing.T) { runtime.GOMAXPROCS(1) - var cache1 = NewCache(NewPiecesOption(32)) + var cache1 = NewCache[int](NewPiecesOption(32)) for i := 0; i < 1_000_000; i++ { cache1.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 10))) } - var cache2 = NewCache(NewPiecesOption(5)) + var cache2 = NewCache[int](NewPiecesOption(5)) for i := 0; i < 1_000_000; i++ { cache2.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 10))) } @@ -178,7 +178,7 @@ func TestCache_GC2(t *testing.T) { } func TestCacheDestroy(t *testing.T) { - var cache = NewCache() + var cache = NewCache[int]() t.Log("count:", SharedManager.Count()) cache.Destroy() t.Log("count:", SharedManager.Count()) @@ -187,7 +187,7 @@ func TestCacheDestroy(t *testing.T) { func BenchmarkNewCache(b *testing.B) { runtime.GOMAXPROCS(1) - var cache = NewCache(NewPiecesOption(128)) + var cache = NewCache[int](NewPiecesOption(128)) for i := 0; i < 2_000_000; i++ { cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(10, 100))) } @@ -205,7 +205,7 @@ func BenchmarkNewCache(b *testing.B) { func BenchmarkCache_Add(b *testing.B) { runtime.GOMAXPROCS(1) - var cache = NewCache() + var cache = NewCache[int]() for i := 0; i < b.N; i++ { cache.Write(strconv.Itoa(i), i, fasttime.Now().Unix()+int64(i%1024)) } @@ -214,7 +214,7 @@ func BenchmarkCache_Add(b *testing.B) { func BenchmarkCache_Add_Parallel(b *testing.B) { runtime.GOMAXPROCS(1) - var cache = NewCache() + var cache = NewCache[int64]() var i int64 b.RunParallel(func(pb *testing.PB) { for pb.Next() { @@ -227,7 +227,7 @@ func BenchmarkCache_Add_Parallel(b *testing.B) { func BenchmarkNewCacheGC(b *testing.B) { runtime.GOMAXPROCS(1) - var cache = NewCache(NewPiecesOption(1024)) + var cache = NewCache[int](NewPiecesOption(1024)) for i := 0; i < 3_000_000; i++ { cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 100))) } @@ -244,7 +244,7 @@ func BenchmarkNewCacheGC(b *testing.B) { func BenchmarkNewCacheClean(b *testing.B) { runtime.GOMAXPROCS(1) - var cache = NewCache(NewPiecesOption(128)) + var cache = NewCache[int](NewPiecesOption(128)) for i := 0; i < 3_000_000; i++ { cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(10, 100))) } diff --git a/internal/ttlcache/item.go b/internal/ttlcache/item.go index 5b0c501..f6713d8 100644 --- a/internal/ttlcache/item.go +++ b/internal/ttlcache/item.go @@ -1,6 +1,6 @@ package ttlcache -type Item struct { - Value any +type Item[T any] struct { + Value T expiredAt int64 } diff --git a/internal/ttlcache/manager.go b/internal/ttlcache/manager.go index 53b15b1..9c631ad 100644 --- a/internal/ttlcache/manager.go +++ b/internal/ttlcache/manager.go @@ -11,17 +11,21 @@ import ( var SharedManager = NewManager() +type GCAble interface { + GC() +} + type Manager struct { ticker *time.Ticker locker sync.Mutex - cacheMap map[*Cache]zero.Zero + cacheMap map[GCAble]zero.Zero } func NewManager() *Manager { var manager = &Manager{ ticker: time.NewTicker(2 * time.Second), - cacheMap: map[*Cache]zero.Zero{}, + cacheMap: map[GCAble]zero.Zero{}, } goman.New(func() { @@ -41,13 +45,13 @@ func (this *Manager) init() { } } -func (this *Manager) Add(cache *Cache) { +func (this *Manager) Add(cache GCAble) { this.locker.Lock() this.cacheMap[cache] = zero.New() this.locker.Unlock() } -func (this *Manager) Remove(cache *Cache) { +func (this *Manager) Remove(cache GCAble) { this.locker.Lock() delete(this.cacheMap, cache) this.locker.Unlock() diff --git a/internal/ttlcache/piece.go b/internal/ttlcache/piece.go index 3e9d4df..40d48c0 100644 --- a/internal/ttlcache/piece.go +++ b/internal/ttlcache/piece.go @@ -3,12 +3,11 @@ package ttlcache import ( "github.com/TeaOSLab/EdgeNode/internal/utils/expires" "github.com/TeaOSLab/EdgeNode/internal/utils/fasttime" - "github.com/iwind/TeaGo/types" "sync" ) -type Piece struct { - m map[uint64]*Item +type Piece[T any] struct { + m map[uint64]*Item[T] expiresList *expires.List maxItems int lastGCTime int64 @@ -16,15 +15,15 @@ type Piece struct { locker sync.RWMutex } -func NewPiece(maxItems int) *Piece { - return &Piece{ - m: map[uint64]*Item{}, +func NewPiece[T any](maxItems int) *Piece[T] { + return &Piece[T]{ + m: map[uint64]*Item[T]{}, expiresList: expires.NewSingletonList(), maxItems: maxItems, } } -func (this *Piece) Add(key uint64, item *Item) (ok bool) { +func (this *Piece[T]) Add(key uint64, item *Item[T]) (ok bool) { this.locker.Lock() if this.maxItems > 0 && len(this.m) >= this.maxItems { this.locker.Unlock() @@ -38,11 +37,14 @@ func (this *Piece) Add(key uint64, item *Item) (ok bool) { return true } -func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64, extend bool) (result int64) { +func (this *Piece[T]) IncreaseInt64(key uint64, delta T, expiredAt int64, extend bool) (result T) { this.locker.Lock() item, ok := this.m[key] if ok && item.expiredAt > fasttime.Now().Unix() { - result = types.Int64(item.Value) + delta + int64Value, isInt64 := any(item.Value).(int64) + if isInt64 { + result = any(int64Value + any(delta).(int64)).(T) + } item.Value = result if extend { item.expiredAt = expiredAt @@ -51,7 +53,7 @@ func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64, exten } else { if len(this.m) < this.maxItems { result = delta - this.m[key] = &Item{ + this.m[key] = &Item[T]{ Value: delta, expiredAt: expiredAt, } @@ -63,7 +65,7 @@ func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64, exten return } -func (this *Piece) Delete(key uint64) { +func (this *Piece[T]) Delete(key uint64) { this.expiresList.Remove(key) this.locker.Lock() @@ -71,7 +73,7 @@ func (this *Piece) Delete(key uint64) { this.locker.Unlock() } -func (this *Piece) Read(key uint64) (item *Item) { +func (this *Piece[T]) Read(key uint64) (item *Item[T]) { this.locker.RLock() item = this.m[key] if item != nil && item.expiredAt < fasttime.Now().Unix() { @@ -82,14 +84,14 @@ func (this *Piece) Read(key uint64) (item *Item) { return } -func (this *Piece) Count() (count int) { +func (this *Piece[T]) Count() (count int) { this.locker.RLock() count = len(this.m) this.locker.RUnlock() return } -func (this *Piece) GC() { +func (this *Piece[T]) GC() { var currentTime = fasttime.Now().Unix() if this.lastGCTime == 0 { this.lastGCTime = currentTime - 3600 @@ -112,15 +114,15 @@ func (this *Piece) GC() { this.lastGCTime = currentTime } -func (this *Piece) Clean() { +func (this *Piece[T]) Clean() { this.locker.Lock() - this.m = map[uint64]*Item{} + this.m = map[uint64]*Item[T]{} this.locker.Unlock() this.expiresList.Clean() } -func (this *Piece) Destroy() { +func (this *Piece[T]) Destroy() { this.locker.Lock() this.m = nil this.locker.Unlock() @@ -128,7 +130,7 @@ func (this *Piece) Destroy() { this.expiresList.Clean() } -func (this *Piece) gcItemMap(itemMap expires.ItemMap) { +func (this *Piece[T]) gcItemMap(itemMap expires.ItemMap) { this.locker.Lock() for key := range itemMap { delete(this.m, key) diff --git a/internal/ttlcache/piece_test.go b/internal/ttlcache/piece_test.go index 0ec0ff1..4605f37 100644 --- a/internal/ttlcache/piece_test.go +++ b/internal/ttlcache/piece_test.go @@ -7,10 +7,10 @@ import ( ) func TestPiece_Add(t *testing.T) { - piece := NewPiece(10) - piece.Add(1, &Item{expiredAt: time.Now().Unix() + 3600}) - piece.Add(2, &Item{}) - piece.Add(3, &Item{}) + piece := NewPiece[int](10) + piece.Add(1, &Item[int]{expiredAt: time.Now().Unix() + 3600}) + piece.Add(2, &Item[int]{}) + piece.Add(3, &Item[int]{}) piece.Delete(3) for key, item := range piece.m { t.Log(key, item.Value) @@ -19,18 +19,18 @@ func TestPiece_Add(t *testing.T) { } func TestPiece_MaxItems(t *testing.T) { - piece := NewPiece(10) + piece := NewPiece[int](10) for i := 0; i < 1000; i++ { - piece.Add(uint64(i), &Item{expiredAt: time.Now().Unix() + 3600}) + piece.Add(uint64(i), &Item[int]{expiredAt: time.Now().Unix() + 3600}) } t.Log(len(piece.m)) } func TestPiece_GC(t *testing.T) { - piece := NewPiece(10) - piece.Add(1, &Item{Value: 1, expiredAt: time.Now().Unix() + 1}) - piece.Add(2, &Item{Value: 2, expiredAt: time.Now().Unix() + 1}) - piece.Add(3, &Item{Value: 3, expiredAt: time.Now().Unix() + 1}) + piece := NewPiece[int](10) + piece.Add(1, &Item[int]{Value: 1, expiredAt: time.Now().Unix() + 1}) + piece.Add(2, &Item[int]{Value: 2, expiredAt: time.Now().Unix() + 1}) + piece.Add(3, &Item[int]{Value: 3, expiredAt: time.Now().Unix() + 1}) t.Log("before gc ===") for key, item := range piece.m { t.Log(key, item.Value) @@ -46,9 +46,9 @@ func TestPiece_GC(t *testing.T) { } func TestPiece_GC2(t *testing.T) { - piece := NewPiece(10) + piece := NewPiece[int](10) for i := 0; i < 10_000; i++ { - piece.Add(uint64(i), &Item{Value: 1, expiredAt: time.Now().Unix() + int64(rands.Int(1, 10))}) + piece.Add(uint64(i), &Item[int]{Value: 1, expiredAt: time.Now().Unix() + int64(rands.Int(1, 10))}) } time.Sleep(1 * time.Second) diff --git a/internal/waf/utils/utils.go b/internal/waf/utils/utils.go index 7a7e0c7..9ab09d4 100644 --- a/internal/waf/utils/utils.go +++ b/internal/waf/utils/utils.go @@ -4,12 +4,11 @@ import ( "github.com/TeaOSLab/EdgeNode/internal/re" "github.com/TeaOSLab/EdgeNode/internal/ttlcache" "github.com/cespare/xxhash" - "github.com/iwind/TeaGo/types" "strconv" "time" ) -var cache = ttlcache.NewCache() +var cache = ttlcache.NewCache[int8]() // MatchStringCache 正则表达式匹配字符串,并缓存结果 func MatchStringCache(regex *re.Regexp, s string) bool { @@ -26,7 +25,7 @@ func MatchStringCache(regex *re.Regexp, s string) bool { var key = regex.IdString() + "@" + strconv.FormatUint(hash, 10) var item = cache.Read(key) if item != nil { - return types.Int8(item.Value) == 1 + return item.Value == 1 } var b = regex.MatchString(s) if b { @@ -52,10 +51,10 @@ func MatchBytesCache(regex *re.Regexp, byteSlice []byte) bool { var key = regex.IdString() + "@" + strconv.FormatUint(hash, 10) var item = cache.Read(key) if item != nil { - return types.Int8(item.Value) == 1 + return item.Value == 1 } if item != nil { - return types.Int8(item.Value) == 1 + return item.Value == 1 } var b = regex.Match(byteSlice) if b { diff --git a/internal/waf/utils/utils_test.go b/internal/waf/utils/utils_test.go index d5a0612..dabcc89 100644 --- a/internal/waf/utils/utils_test.go +++ b/internal/waf/utils/utils_test.go @@ -1,7 +1,8 @@ -package utils +package utils_test import ( "github.com/TeaOSLab/EdgeNode/internal/re" + "github.com/TeaOSLab/EdgeNode/internal/waf/utils" "net/http" "regexp" "runtime" @@ -13,16 +14,16 @@ import ( func TestMatchStringCache(t *testing.T) { regex := re.MustCompile(`\d+`) - t.Log(MatchStringCache(regex, "123")) - t.Log(MatchStringCache(regex, "123")) - t.Log(MatchStringCache(regex, "123")) + t.Log(utils.MatchStringCache(regex, "123")) + t.Log(utils.MatchStringCache(regex, "123")) + t.Log(utils.MatchStringCache(regex, "123")) } func TestMatchBytesCache(t *testing.T) { regex := re.MustCompile(`\d+`) - t.Log(MatchBytesCache(regex, []byte("123"))) - t.Log(MatchBytesCache(regex, []byte("123"))) - t.Log(MatchBytesCache(regex, []byte("123"))) + t.Log(utils.MatchBytesCache(regex, []byte("123"))) + t.Log(utils.MatchBytesCache(regex, []byte("123"))) + t.Log(utils.MatchBytesCache(regex, []byte("123"))) } func TestMatchRemoteCache(t *testing.T) { @@ -55,10 +56,10 @@ func BenchmarkMatchStringCache(b *testing.B) { var data = strings.Repeat("HELLO", 512) var regex = re.MustCompile(`(?iU)\b(eval|system|exec|execute|passthru|shell_exec|phpinfo)\b`) //b.Log(regex.Keywords()) - _ = MatchStringCache(regex, data) + _ = utils.MatchStringCache(regex, data) for i := 0; i < b.N; i++ { - _ = MatchStringCache(regex, data) + _ = utils.MatchStringCache(regex, data) } }