TRKJ/EdgeNode
2
0
Fork 0
mirror of https://github.com/TeaOSLab/EdgeNode.git synced 2025-11-06 18:10:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

ttlcache支持泛型

Browse Source
This commit is contained in:
GoEdgeLab
2023-10-05 08:28:16 +08:00
parent 3851a5c424
commit 6aba7fb295
10 changed files with 95 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
internal/caches/list_file.go
View File

@@ -27,7 +27,7 @@ type FileList struct {
onAdd func(item *Item) onAdd func(item *Item)
onRemove func(item *Item) onRemove func(item *Item)
memoryCache *ttlcache.Cache memoryCache *ttlcache.Cache[bool]
// 老数据库地址 // 老数据库地址
oldDir string oldDir string
@@ -36,7 +36,7 @@ type FileList struct {
func NewFileList(dir string) ListInterface { func NewFileList(dir string) ListInterface {
return &FileList{ return &FileList{
dir: dir, dir: dir,
memoryCache: ttlcache.NewCache(), memoryCache: ttlcache.NewCache[bool](),
} }
} }
@@ -120,7 +120,7 @@ func (this *FileList) Add(hash string, item *Item) error {
return err return err
} }
this.memoryCache.Write(hash, 1, this.maxExpiresAtForMemoryCache(item.ExpiredAt)) this.memoryCache.Write(hash, true, this.maxExpiresAtForMemoryCache(item.ExpiredAt))
if this.onAdd != nil { if this.onAdd != nil {
this.onAdd(item) this.onAdd(item)
@@ -158,7 +158,7 @@ func (this *FileList) Exist(hash string) (bool, error) {
} }
return false, err return false, err
} }
this.memoryCache.Write(hash, 1, this.maxExpiresAtForMemoryCache(expiredAt)) this.memoryCache.Write(hash, false, this.maxExpiresAtForMemoryCache(expiredAt))
return true, nil return true, nil
} }

2
internal/nodes/http_request_mismatch.go
View File

@@ -63,7 +63,7 @@ func (this *HTTPRequest) doMismatch() {
// 要考虑到服务在切换集群时，域名未生效状态时，用户访问的仍然是老集群中的节点，就会产生找不到域名的情况 // 要考虑到服务在切换集群时，域名未生效状态时，用户访问的仍然是老集群中的节点，就会产生找不到域名的情况
if len(remoteIP) > 0 { if len(remoteIP) > 0 {
const maxAttempts = 100 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) { if !waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteIP) {
waf.SharedIPBlackList.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteIP, time.Now().Unix()+3600) waf.SharedIPBlackList.Add(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteIP, time.Now().Unix()+3600)

38
internal/ttlcache/cache.go
View File

@@ -5,7 +5,7 @@ import (
"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime" "github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
) )
var SharedCache = NewBigCache() var SharedInt64Cache = NewBigCache[int64]()
// Cache TTL缓存 // Cache TTL缓存
// 最大的缓存时间为30 * 86400 // 最大的缓存时间为30 * 86400
@@ -13,24 +13,24 @@ var SharedCache = NewBigCache()
// //
// Piece1 | Piece2 | Piece3 | ... // Piece1 | Piece2 | Piece3 | ...
// [ Item1, Item2, ... ] | ... // [ Item1, Item2, ... ] | ...
type Cache struct { type Cache[T any] struct {
isDestroyed bool isDestroyed bool
pieces []*Piece pieces []*Piece[T]
countPieces uint64 countPieces uint64
maxItems int maxItems int
gcPieceIndex int gcPieceIndex int
} }
func NewBigCache() *Cache { func NewBigCache[T any]() *Cache[T] {
var delta = utils.SystemMemoryGB() / 2 var delta = utils.SystemMemoryGB() / 2
if delta <= 0 { if delta <= 0 {
delta = 1 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 countPieces = 256
var maxItems = 1_000_000 var maxItems = 1_000_000
@@ -61,13 +61,13 @@ func NewCache(opt ...OptionInterface) *Cache {
} }
} }
var cache = &Cache{ var cache = &Cache[T]{
countPieces: uint64(countPieces), countPieces: uint64(countPieces),
maxItems: maxItems, maxItems: maxItems,
} }
for i := 0; i < countPieces; i++ { 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 // Add to manager
@@ -76,7 +76,7 @@ func NewCache(opt ...OptionInterface) *Cache {
return 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 { if this.isDestroyed {
return return
} }
@@ -92,20 +92,20 @@ func (this *Cache) Write(key string, value any, expiredAt int64) (ok bool) {
} }
var uint64Key = HashKey([]byte(key)) var uint64Key = HashKey([]byte(key))
var pieceIndex = uint64Key % this.countPieces var pieceIndex = uint64Key % this.countPieces
return this.pieces[pieceIndex].Add(uint64Key, &Item{ return this.pieces[pieceIndex].Add(uint64Key, &Item[T]{
Value: value, Value: value,
expiredAt: expiredAt, 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 { if this.isDestroyed {
return 0 return any(0).(T)
} }
var currentTimestamp = fasttime.Now().Unix() var currentTimestamp = fasttime.Now().Unix()
if expiredAt <= currentTimestamp { if expiredAt <= currentTimestamp {
return 0 return any(0).(T)
} }
var maxExpiredAt = currentTimestamp + 30*86400 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) 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)) var uint64Key = HashKey([]byte(key))
return this.pieces[uint64Key%this.countPieces].Read(uint64Key) 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)) var uint64Key = HashKey([]byte(key))
this.pieces[uint64Key%this.countPieces].Delete(uint64Key) 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 { for _, piece := range this.pieces {
count += piece.Count() count += piece.Count()
} }
return return
} }
func (this *Cache) GC() { func (this *Cache[T]) GC() {
var index = this.gcPieceIndex var index = this.gcPieceIndex
const maxPiecesPerGC = 4 const maxPiecesPerGC = 4
for i := index; i < index+maxPiecesPerGC; i++ { for i := index; i < index+maxPiecesPerGC; i++ {
@@ -151,13 +151,13 @@ func (this *Cache) GC() {
this.gcPieceIndex = index this.gcPieceIndex = index
} }
func (this *Cache) Clean() { func (this *Cache[T]) Clean() {
for _, piece := range this.pieces { for _, piece := range this.pieces {
piece.Clean() piece.Clean()
} }
} }
func (this *Cache) Destroy() { func (this *Cache[T]) Destroy() {
SharedManager.Remove(this) SharedManager.Remove(this)
this.isDestroyed = true this.isDestroyed = true

30
internal/ttlcache/cache_test.go
View File

@@ -14,7 +14,7 @@ import (
) )
func TestNewCache(t *testing.T) { func TestNewCache(t *testing.T) {
var cache = NewCache() var cache = NewCache[int]()
cache.Write("a", 1, time.Now().Unix()+3600) cache.Write("a", 1, time.Now().Unix()+3600)
cache.Write("b", 2, time.Now().Unix()+1) cache.Write("b", 2, time.Now().Unix()+1)
cache.Write("c", 1, time.Now().Unix()+3602) cache.Write("c", 1, time.Now().Unix()+3602)
@@ -48,7 +48,7 @@ func TestCache_Memory(t *testing.T) {
testutils.StartMemoryStats(t) testutils.StartMemoryStats(t)
var cache = NewCache() var cache = NewCache[int]()
var count = 20_000_000 var count = 20_000_000
for i := 0; i < count; i++ { for i := 0; i < count; i++ {
cache.Write("a"+strconv.Itoa(i), 1, time.Now().Unix()+3600) 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) { func TestCache_IncreaseInt64(t *testing.T) {
var a = assert.NewAssertion(t) var a = assert.NewAssertion(t)
var cache = NewCache() var cache = NewCache[int64]()
var unixTime = time.Now().Unix() var unixTime = time.Now().Unix()
{ {
cache.IncreaseInt64("a", 1, unixTime+3600, false) cache.IncreaseInt64("a", 1, unixTime+3600, false)
var item = cache.Read("a") var item = cache.Read("a")
t.Log(item) t.Log(item)
a.IsTrue(item.Value == int64(1)) a.IsTrue(item.Value == 1)
a.IsTrue(item.expiredAt == unixTime+3600) a.IsTrue(item.expiredAt == unixTime+3600)
} }
{ {
cache.IncreaseInt64("a", 1, unixTime+3600+1, true) cache.IncreaseInt64("a", 1, unixTime+3600+1, true)
var item = cache.Read("a") var item = cache.Read("a")
t.Log(item) t.Log(item)
a.IsTrue(item.Value == int64(2)) a.IsTrue(item.Value == 2)
a.IsTrue(item.expiredAt == unixTime+3600+1) a.IsTrue(item.expiredAt == unixTime+3600+1)
} }
{ {
@@ -103,7 +103,7 @@ func TestCache_IncreaseInt64(t *testing.T) {
func TestCache_Read(t *testing.T) { func TestCache_Read(t *testing.T) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache = NewCache(PiecesOption{Count: 32}) var cache = NewCache[int](PiecesOption{Count: 32})
for i := 0; i < 10_000_000; i++ { for i := 0; i < 10_000_000; i++ {
cache.Write("HELLO_WORLD_"+strconv.Itoa(i), i, time.Now().Unix()+int64(i%10240)+1) 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) { 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("a", 1, time.Now().Unix()+1)
cache.Write("b", 2, time.Now().Unix()+2) cache.Write("b", 2, time.Now().Unix()+2)
cache.Write("c", 3, time.Now().Unix()+3) cache.Write("c", 3, time.Now().Unix()+3)
@@ -161,12 +161,12 @@ func TestCache_GC(t *testing.T) {
func TestCache_GC2(t *testing.T) { func TestCache_GC2(t *testing.T) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache1 = NewCache(NewPiecesOption(32)) var cache1 = NewCache[int](NewPiecesOption(32))
for i := 0; i < 1_000_000; i++ { for i := 0; i < 1_000_000; i++ {
cache1.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 10))) 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++ { for i := 0; i < 1_000_000; i++ {
cache2.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 10))) 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) { func TestCacheDestroy(t *testing.T) {
var cache = NewCache() var cache = NewCache[int]()
t.Log("count:", SharedManager.Count()) t.Log("count:", SharedManager.Count())
cache.Destroy() cache.Destroy()
t.Log("count:", SharedManager.Count()) t.Log("count:", SharedManager.Count())
@@ -187,7 +187,7 @@ func TestCacheDestroy(t *testing.T) {
func BenchmarkNewCache(b *testing.B) { func BenchmarkNewCache(b *testing.B) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache = NewCache(NewPiecesOption(128)) var cache = NewCache[int](NewPiecesOption(128))
for i := 0; i < 2_000_000; i++ { for i := 0; i < 2_000_000; i++ {
cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(10, 100))) 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) { func BenchmarkCache_Add(b *testing.B) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache = NewCache() var cache = NewCache[int]()
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
cache.Write(strconv.Itoa(i), i, fasttime.Now().Unix()+int64(i%1024)) 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) { func BenchmarkCache_Add_Parallel(b *testing.B) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache = NewCache() var cache = NewCache[int64]()
var i int64 var i int64
b.RunParallel(func(pb *testing.PB) { b.RunParallel(func(pb *testing.PB) {
for pb.Next() { for pb.Next() {
@@ -227,7 +227,7 @@ func BenchmarkCache_Add_Parallel(b *testing.B) {
func BenchmarkNewCacheGC(b *testing.B) { func BenchmarkNewCacheGC(b *testing.B) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache = NewCache(NewPiecesOption(1024)) var cache = NewCache[int](NewPiecesOption(1024))
for i := 0; i < 3_000_000; i++ { for i := 0; i < 3_000_000; i++ {
cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 100))) 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) { func BenchmarkNewCacheClean(b *testing.B) {
runtime.GOMAXPROCS(1) runtime.GOMAXPROCS(1)
var cache = NewCache(NewPiecesOption(128)) var cache = NewCache[int](NewPiecesOption(128))
for i := 0; i < 3_000_000; i++ { for i := 0; i < 3_000_000; i++ {
cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(10, 100))) cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(10, 100)))
} }

4
internal/ttlcache/item.go
View File

@@ -1,6 +1,6 @@
package ttlcache package ttlcache
type Item struct { type Item[T any] struct {
Value any Value T
expiredAt int64 expiredAt int64
} }

12
internal/ttlcache/manager.go
View File

@@ -11,17 +11,21 @@ import (
var SharedManager = NewManager() var SharedManager = NewManager()
type GCAble interface {
GC()
}
type Manager struct { type Manager struct {
ticker *time.Ticker ticker *time.Ticker
locker sync.Mutex locker sync.Mutex
cacheMap map[*Cache]zero.Zero cacheMap map[GCAble]zero.Zero
} }
func NewManager() *Manager { func NewManager() *Manager {
var manager = &Manager{ var manager = &Manager{
ticker: time.NewTicker(2 * time.Second), ticker: time.NewTicker(2 * time.Second),
cacheMap: map[*Cache]zero.Zero{}, cacheMap: map[GCAble]zero.Zero{},
} }
goman.New(func() { 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.locker.Lock()
this.cacheMap[cache] = zero.New() this.cacheMap[cache] = zero.New()
this.locker.Unlock() this.locker.Unlock()
} }
func (this *Manager) Remove(cache *Cache) { func (this *Manager) Remove(cache GCAble) {
this.locker.Lock() this.locker.Lock()
delete(this.cacheMap, cache) delete(this.cacheMap, cache)
this.locker.Unlock() this.locker.Unlock()

38
internal/ttlcache/piece.go
View File

@@ -3,12 +3,11 @@ package ttlcache
import ( import (
"github.com/TeaOSLab/EdgeNode/internal/utils/expires" "github.com/TeaOSLab/EdgeNode/internal/utils/expires"
"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime" "github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
"github.com/iwind/TeaGo/types"
"sync" "sync"
) )
type Piece struct { type Piece[T any] struct {
m map[uint64]*Item m map[uint64]*Item[T]
expiresList *expires.List expiresList *expires.List
maxItems int maxItems int
lastGCTime int64 lastGCTime int64
@@ -16,15 +15,15 @@ type Piece struct {
locker sync.RWMutex locker sync.RWMutex
} }
func NewPiece(maxItems int) *Piece { func NewPiece[T any](maxItems int) *Piece[T] {
return &Piece{ return &Piece[T]{
m: map[uint64]*Item{}, m: map[uint64]*Item[T]{},
expiresList: expires.NewSingletonList(), expiresList: expires.NewSingletonList(),
maxItems: maxItems, 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() this.locker.Lock()
if this.maxItems > 0 && len(this.m) >= this.maxItems { if this.maxItems > 0 && len(this.m) >= this.maxItems {
this.locker.Unlock() this.locker.Unlock()
@@ -38,11 +37,14 @@ func (this *Piece) Add(key uint64, item *Item) (ok bool) {
return true 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() this.locker.Lock()
item, ok := this.m[key] item, ok := this.m[key]
if ok && item.expiredAt > fasttime.Now().Unix() { 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 item.Value = result
if extend { if extend {
item.expiredAt = expiredAt item.expiredAt = expiredAt
@@ -51,7 +53,7 @@ func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64, exten
} else { } else {
if len(this.m) < this.maxItems { if len(this.m) < this.maxItems {
result = delta result = delta
this.m[key] = &Item{ this.m[key] = &Item[T]{
Value: delta, Value: delta,
expiredAt: expiredAt, expiredAt: expiredAt,
} }
@@ -63,7 +65,7 @@ func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64, exten
return return
} }
func (this *Piece) Delete(key uint64) { func (this *Piece[T]) Delete(key uint64) {
this.expiresList.Remove(key) this.expiresList.Remove(key)
this.locker.Lock() this.locker.Lock()
@@ -71,7 +73,7 @@ func (this *Piece) Delete(key uint64) {
this.locker.Unlock() this.locker.Unlock()
} }
func (this *Piece) Read(key uint64) (item *Item) { func (this *Piece[T]) Read(key uint64) (item *Item[T]) {
this.locker.RLock() this.locker.RLock()
item = this.m[key] item = this.m[key]
if item != nil && item.expiredAt < fasttime.Now().Unix() { if item != nil && item.expiredAt < fasttime.Now().Unix() {
@@ -82,14 +84,14 @@ func (this *Piece) Read(key uint64) (item *Item) {
return return
} }
func (this *Piece) Count() (count int) { func (this *Piece[T]) Count() (count int) {
this.locker.RLock() this.locker.RLock()
count = len(this.m) count = len(this.m)
this.locker.RUnlock() this.locker.RUnlock()
return return
} }
func (this *Piece) GC() { func (this *Piece[T]) GC() {
var currentTime = fasttime.Now().Unix() var currentTime = fasttime.Now().Unix()
if this.lastGCTime == 0 { if this.lastGCTime == 0 {
this.lastGCTime = currentTime - 3600 this.lastGCTime = currentTime - 3600
@@ -112,15 +114,15 @@ func (this *Piece) GC() {
this.lastGCTime = currentTime this.lastGCTime = currentTime
} }
func (this *Piece) Clean() { func (this *Piece[T]) Clean() {
this.locker.Lock() this.locker.Lock()
this.m = map[uint64]*Item{} this.m = map[uint64]*Item[T]{}
this.locker.Unlock() this.locker.Unlock()
this.expiresList.Clean() this.expiresList.Clean()
} }
func (this *Piece) Destroy() { func (this *Piece[T]) Destroy() {
this.locker.Lock() this.locker.Lock()
this.m = nil this.m = nil
this.locker.Unlock() this.locker.Unlock()
@@ -128,7 +130,7 @@ func (this *Piece) Destroy() {
this.expiresList.Clean() this.expiresList.Clean()
} }
func (this *Piece) gcItemMap(itemMap expires.ItemMap) { func (this *Piece[T]) gcItemMap(itemMap expires.ItemMap) {
this.locker.Lock() this.locker.Lock()
for key := range itemMap { for key := range itemMap {
delete(this.m, key) delete(this.m, key)

24
internal/ttlcache/piece_test.go
View File

@@ -7,10 +7,10 @@ import (
) )
func TestPiece_Add(t *testing.T) { func TestPiece_Add(t *testing.T) {
piece := NewPiece(10) piece := NewPiece[int](10)
piece.Add(1, &Item{expiredAt: time.Now().Unix() + 3600}) piece.Add(1, &Item[int]{expiredAt: time.Now().Unix() + 3600})
piece.Add(2, &Item{}) piece.Add(2, &Item[int]{})
piece.Add(3, &Item{}) piece.Add(3, &Item[int]{})
piece.Delete(3) piece.Delete(3)
for key, item := range piece.m { for key, item := range piece.m {
t.Log(key, item.Value) t.Log(key, item.Value)
@@ -19,18 +19,18 @@ func TestPiece_Add(t *testing.T) {
} }
func TestPiece_MaxItems(t *testing.T) { func TestPiece_MaxItems(t *testing.T) {
piece := NewPiece(10) piece := NewPiece[int](10)
for i := 0; i < 1000; i++ { 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)) t.Log(len(piece.m))
} }
func TestPiece_GC(t *testing.T) { func TestPiece_GC(t *testing.T) {
piece := NewPiece(10) piece := NewPiece[int](10)
piece.Add(1, &Item{Value: 1, expiredAt: time.Now().Unix() + 1}) piece.Add(1, &Item[int]{Value: 1, expiredAt: time.Now().Unix() + 1})
piece.Add(2, &Item{Value: 2, expiredAt: time.Now().Unix() + 1}) piece.Add(2, &Item[int]{Value: 2, expiredAt: time.Now().Unix() + 1})
piece.Add(3, &Item{Value: 3, expiredAt: time.Now().Unix() + 1}) piece.Add(3, &Item[int]{Value: 3, expiredAt: time.Now().Unix() + 1})
t.Log("before gc ===") t.Log("before gc ===")
for key, item := range piece.m { for key, item := range piece.m {
t.Log(key, item.Value) t.Log(key, item.Value)
@@ -46,9 +46,9 @@ func TestPiece_GC(t *testing.T) {
} }
func TestPiece_GC2(t *testing.T) { func TestPiece_GC2(t *testing.T) {
piece := NewPiece(10) piece := NewPiece[int](10)
for i := 0; i < 10_000; i++ { 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) time.Sleep(1 * time.Second)

9
internal/waf/utils/utils.go
View File

@@ -4,12 +4,11 @@ import (
"github.com/TeaOSLab/EdgeNode/internal/re" "github.com/TeaOSLab/EdgeNode/internal/re"
"github.com/TeaOSLab/EdgeNode/internal/ttlcache" "github.com/TeaOSLab/EdgeNode/internal/ttlcache"
"github.com/cespare/xxhash" "github.com/cespare/xxhash"
"github.com/iwind/TeaGo/types"
"strconv" "strconv"
"time" "time"
) )
var cache = ttlcache.NewCache() var cache = ttlcache.NewCache[int8]()
// MatchStringCache 正则表达式匹配字符串，并缓存结果 // MatchStringCache 正则表达式匹配字符串，并缓存结果
func MatchStringCache(regex *re.Regexp, s string) bool { 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 key = regex.IdString() + "@" + strconv.FormatUint(hash, 10)
var item = cache.Read(key) var item = cache.Read(key)
if item != nil { if item != nil {
return types.Int8(item.Value) == 1 return item.Value == 1
} }
var b = regex.MatchString(s) var b = regex.MatchString(s)
if b { if b {
@@ -52,10 +51,10 @@ func MatchBytesCache(regex *re.Regexp, byteSlice []byte) bool {
var key = regex.IdString() + "@" + strconv.FormatUint(hash, 10) var key = regex.IdString() + "@" + strconv.FormatUint(hash, 10)
var item = cache.Read(key) var item = cache.Read(key)
if item != nil { if item != nil {
return types.Int8(item.Value) == 1 return item.Value == 1
} }
if item != nil { if item != nil {
return types.Int8(item.Value) == 1 return item.Value == 1
} }
var b = regex.Match(byteSlice) var b = regex.Match(byteSlice)
if b { if b {

19
internal/waf/utils/utils_test.go
View File

@@ -1,7 +1,8 @@
package utils package utils_test
import ( import (
"github.com/TeaOSLab/EdgeNode/internal/re" "github.com/TeaOSLab/EdgeNode/internal/re"
"github.com/TeaOSLab/EdgeNode/internal/waf/utils"
"net/http" "net/http"
"regexp" "regexp"
"runtime" "runtime"
@@ -13,16 +14,16 @@ import (
func TestMatchStringCache(t *testing.T) { func TestMatchStringCache(t *testing.T) {
regex := re.MustCompile(`\d+`) regex := re.MustCompile(`\d+`)
t.Log(MatchStringCache(regex, "123")) t.Log(utils.MatchStringCache(regex, "123"))
t.Log(MatchStringCache(regex, "123")) t.Log(utils.MatchStringCache(regex, "123"))
t.Log(MatchStringCache(regex, "123")) t.Log(utils.MatchStringCache(regex, "123"))
} }
func TestMatchBytesCache(t *testing.T) { func TestMatchBytesCache(t *testing.T) {
regex := re.MustCompile(`\d+`) regex := re.MustCompile(`\d+`)
t.Log(MatchBytesCache(regex, []byte("123"))) t.Log(utils.MatchBytesCache(regex, []byte("123")))
t.Log(MatchBytesCache(regex, []byte("123"))) t.Log(utils.MatchBytesCache(regex, []byte("123")))
t.Log(MatchBytesCache(regex, []byte("123"))) t.Log(utils.MatchBytesCache(regex, []byte("123")))
} }
func TestMatchRemoteCache(t *testing.T) { func TestMatchRemoteCache(t *testing.T) {
@@ -55,10 +56,10 @@ func BenchmarkMatchStringCache(b *testing.B) {
var data = strings.Repeat("HELLO", 512) var data = strings.Repeat("HELLO", 512)
var regex = re.MustCompile(`(?iU)\b(eval|system|exec|execute|passthru|shell_exec|phpinfo)\b`) var regex = re.MustCompile(`(?iU)\b(eval|system|exec|execute|passthru|shell_exec|phpinfo)\b`)
//b.Log(regex.Keywords()) //b.Log(regex.Keywords())
_ = MatchStringCache(regex, data) _ = utils.MatchStringCache(regex, data)
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
_ = MatchStringCache(regex, data) _ = utils.MatchStringCache(regex, data)
} }
} }