diff --git a/internal/cache/item.go b/internal/cache/item.go deleted file mode 100644 index 276da46..0000000 --- a/internal/cache/item.go +++ /dev/null @@ -1,6 +0,0 @@ -package cache - -type Item struct { - value interface{} - expiredAt int64 -} diff --git a/internal/cache/cache.go b/internal/ttlcache/cache.go similarity index 81% rename from internal/cache/cache.go rename to internal/ttlcache/cache.go index 4ec1fc2..f642f39 100644 --- a/internal/cache/cache.go +++ b/internal/ttlcache/cache.go @@ -1,4 +1,4 @@ -package cache +package ttlcache import ( "time" @@ -14,12 +14,14 @@ import ( type Cache struct { pieces []*Piece countPieces uint64 + maxItems int gcPieceIndex int } func NewCache(opt ...OptionInterface) *Cache { countPieces := 128 + maxItems := 1_000_000 for _, option := range opt { if option == nil { continue @@ -29,20 +31,25 @@ func NewCache(opt ...OptionInterface) *Cache { if o.Count > 0 { countPieces = o.Count } + case *MaxItemsOption: + if o.Count > 0 { + maxItems = o.Count + } } } cache := &Cache{ countPieces: uint64(countPieces), + maxItems: maxItems, } for i := 0; i < countPieces; i++ { - cache.pieces = append(cache.pieces, NewPiece()) + cache.pieces = append(cache.pieces, NewPiece(maxItems/countPieces)) } // start timer go func() { - ticker := time.NewTicker(1 * time.Second) + ticker := time.NewTicker(5 * time.Second) for range ticker.C { cache.GC() } @@ -51,7 +58,7 @@ func NewCache(opt ...OptionInterface) *Cache { return cache } -func (this *Cache) Add(key string, value interface{}, expiredAt int64) { +func (this *Cache) Write(key string, value interface{}, expiredAt int64) { currentTimestamp := time.Now().Unix() if expiredAt <= currentTimestamp { return @@ -64,12 +71,12 @@ func (this *Cache) Add(key string, value interface{}, expiredAt int64) { uint64Key := HashKey([]byte(key)) pieceIndex := uint64Key % this.countPieces this.pieces[pieceIndex].Add(uint64Key, &Item{ - value: value, + Value: value, expiredAt: expiredAt, }) } -func (this *Cache) Read(key string) (value *Item) { +func (this *Cache) Read(key string) (item *Item) { uint64Key := HashKey([]byte(key)) return this.pieces[uint64Key%this.countPieces].Read(uint64Key) } diff --git a/internal/cache/cache_test.go b/internal/ttlcache/cache_test.go similarity index 60% rename from internal/cache/cache_test.go rename to internal/ttlcache/cache_test.go index 0d88d69..ddeb34b 100644 --- a/internal/cache/cache_test.go +++ b/internal/ttlcache/cache_test.go @@ -1,4 +1,4 @@ -package cache +package ttlcache import ( "github.com/iwind/TeaGo/rands" @@ -10,15 +10,15 @@ import ( func TestNewCache(t *testing.T) { cache := NewCache() - cache.Add("a", 1, time.Now().Unix()+3600) - cache.Add("b", 2, time.Now().Unix()+3601) - cache.Add("a", 1, time.Now().Unix()+3602) - cache.Add("d", 1, time.Now().Unix()+1) + cache.Write("a", 1, time.Now().Unix()+3600) + cache.Write("b", 2, time.Now().Unix()+3601) + cache.Write("a", 1, time.Now().Unix()+3602) + cache.Write("d", 1, time.Now().Unix()+1) for _, piece := range cache.pieces { if len(piece.m) > 0 { for k, item := range piece.m { - t.Log(k, "=>", item.value, item.expiredAt) + t.Log(k, "=>", item.Value, item.expiredAt) } } } @@ -32,25 +32,25 @@ func BenchmarkCache_Add(b *testing.B) { cache := NewCache() for i := 0; i < b.N; i++ { - cache.Add(strconv.Itoa(i), i, time.Now().Unix()+int64(i%1024)) + cache.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(i%1024)) } } func TestCache_Read(t *testing.T) { runtime.GOMAXPROCS(1) - var cache = NewCache(PiecesOption{Count: 512}) + var cache = NewCache(PiecesOption{Count: 32}) for i := 0; i < 10_000_000; i++ { - cache.Add("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) } - /**total := 0 + total := 0 for _, piece := range cache.pieces { - t.Log(len(piece.m), "keys") + //t.Log(len(piece.m), "keys") total += len(piece.m) } - t.Log(total, "total keys")**/ + t.Log(total, "total keys") before := time.Now() for i := 0; i < 10_240; i++ { @@ -61,15 +61,15 @@ func TestCache_Read(t *testing.T) { func TestCache_GC(t *testing.T) { var cache = NewCache(&PiecesOption{Count: 5}) - cache.Add("a", 1, time.Now().Unix()+1) - cache.Add("b", 2, time.Now().Unix()+2) - cache.Add("c", 3, time.Now().Unix()+3) - cache.Add("d", 4, time.Now().Unix()+4) - cache.Add("e", 5, time.Now().Unix()+10) + cache.Write("a", 1, time.Now().Unix()+1) + cache.Write("b", 2, time.Now().Unix()+2) + cache.Write("c", 3, time.Now().Unix()+3) + cache.Write("d", 4, time.Now().Unix()+4) + cache.Write("e", 5, time.Now().Unix()+10) go func() { for i := 0; i < 1000; i++ { - cache.Add("f", 1, time.Now().Unix()+1) + cache.Write("f", 1, time.Now().Unix()+1) time.Sleep(10 * time.Millisecond) } }() @@ -83,7 +83,7 @@ func TestCache_GC(t *testing.T) { t.Log("now:", time.Now().Unix()) for _, p := range cache.pieces { for k, v := range p.m { - t.Log(k, v.value, v.expiredAt) + t.Log(k, v.Value, v.expiredAt) } } } @@ -93,7 +93,7 @@ func TestCache_GC2(t *testing.T) { cache := NewCache() for i := 0; i < 1_000_000; i++ { - cache.Add(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))) } for i := 0; i < 100; i++ { diff --git a/internal/ttlcache/item.go b/internal/ttlcache/item.go new file mode 100644 index 0000000..9720e29 --- /dev/null +++ b/internal/ttlcache/item.go @@ -0,0 +1,6 @@ +package ttlcache + +type Item struct { + Value interface{} + expiredAt int64 +} diff --git a/internal/cache/option.go b/internal/ttlcache/option.go similarity index 51% rename from internal/cache/option.go rename to internal/ttlcache/option.go index 286810d..09a4cf8 100644 --- a/internal/cache/option.go +++ b/internal/ttlcache/option.go @@ -1,4 +1,4 @@ -package cache +package ttlcache type OptionInterface interface { } @@ -10,3 +10,11 @@ type PiecesOption struct { func NewPiecesOption(count int) *PiecesOption { return &PiecesOption{Count: count} } + +type MaxItemsOption struct { + Count int +} + +func NewMaxItemsOption(count int) *MaxItemsOption { + return &MaxItemsOption{Count: count} +} diff --git a/internal/cache/piece.go b/internal/ttlcache/piece.go similarity index 77% rename from internal/cache/piece.go rename to internal/ttlcache/piece.go index 6f186e1..3e3ed13 100644 --- a/internal/cache/piece.go +++ b/internal/ttlcache/piece.go @@ -1,4 +1,4 @@ -package cache +package ttlcache import ( "github.com/TeaOSLab/EdgeNode/internal/utils" @@ -7,16 +7,21 @@ import ( ) type Piece struct { - m map[uint64]*Item - locker sync.RWMutex + m map[uint64]*Item + maxItems int + locker sync.RWMutex } -func NewPiece() *Piece { - return &Piece{m: map[uint64]*Item{}} +func NewPiece(maxItems int) *Piece { + return &Piece{m: map[uint64]*Item{}, maxItems: maxItems} } func (this *Piece) Add(key uint64, item *Item) () { this.locker.Lock() + if len(this.m) >= this.maxItems { + this.locker.Unlock() + return + } this.m[key] = item this.locker.Unlock() } diff --git a/internal/cache/piece_test.go b/internal/ttlcache/piece_test.go similarity index 54% rename from internal/cache/piece_test.go rename to internal/ttlcache/piece_test.go index 1f4d4e0..0ec0ff1 100644 --- a/internal/cache/piece_test.go +++ b/internal/ttlcache/piece_test.go @@ -1,4 +1,4 @@ -package cache +package ttlcache import ( "github.com/iwind/TeaGo/rands" @@ -7,25 +7,33 @@ import ( ) func TestPiece_Add(t *testing.T) { - piece := NewPiece() + piece := NewPiece(10) piece.Add(1, &Item{expiredAt: time.Now().Unix() + 3600}) piece.Add(2, &Item{}) piece.Add(3, &Item{}) piece.Delete(3) for key, item := range piece.m { - t.Log(key, item.value) + t.Log(key, item.Value) } t.Log(piece.Read(1)) } +func TestPiece_MaxItems(t *testing.T) { + piece := NewPiece(10) + for i := 0; i < 1000; i++ { + piece.Add(uint64(i), &Item{expiredAt: time.Now().Unix() + 3600}) + } + t.Log(len(piece.m)) +} + func TestPiece_GC(t *testing.T) { - piece := NewPiece() - 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(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}) t.Log("before gc ===") for key, item := range piece.m { - t.Log(key, item.value) + t.Log(key, item.Value) } time.Sleep(1 * time.Second) @@ -33,14 +41,14 @@ func TestPiece_GC(t *testing.T) { t.Log("after gc ===") for key, item := range piece.m { - t.Log(key, item.value) + t.Log(key, item.Value) } } func TestPiece_GC2(t *testing.T) { - piece := NewPiece() + piece := NewPiece(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{Value: 1, expiredAt: time.Now().Unix() + int64(rands.Int(1, 10))}) } time.Sleep(1 * time.Second) diff --git a/internal/cache/utils.go b/internal/ttlcache/utils.go similarity index 86% rename from internal/cache/utils.go rename to internal/ttlcache/utils.go index 1ea27b0..652814b 100644 --- a/internal/cache/utils.go +++ b/internal/ttlcache/utils.go @@ -1,4 +1,4 @@ -package cache +package ttlcache import "github.com/dchest/siphash" diff --git a/internal/cache/utils_test.go b/internal/ttlcache/utils_test.go similarity index 92% rename from internal/cache/utils_test.go rename to internal/ttlcache/utils_test.go index 83b2da5..5f9c88f 100644 --- a/internal/cache/utils_test.go +++ b/internal/ttlcache/utils_test.go @@ -1,4 +1,4 @@ -package cache +package ttlcache import ( "runtime" diff --git a/internal/waf/utils/utils.go b/internal/waf/utils/utils.go index f501475..482e5d4 100644 --- a/internal/waf/utils/utils.go +++ b/internal/waf/utils/utils.go @@ -2,14 +2,17 @@ package utils import ( "fmt" - "github.com/TeaOSLab/EdgeNode/internal/grids" + "github.com/TeaOSLab/EdgeNode/internal/ttlcache" "github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/dchest/siphash" + "github.com/iwind/TeaGo/types" "regexp" "strconv" + "time" ) -var grid = grids.NewGrid(32, grids.NewLimitCountOpt(1000_0000)) +//var grid = grids.NewGrid(32, grids.NewLimitCountOpt(1000_0000)) +var cache = ttlcache.NewCache() // 正则表达式匹配字符串,并缓存结果 func MatchStringCache(regex *regexp.Regexp, s string) bool { @@ -19,16 +22,16 @@ func MatchStringCache(regex *regexp.Regexp, s string) bool { } hash := siphash.Hash(0, 0, utils.UnsafeStringToBytes(s)) - key := []byte(fmt.Sprintf("%p_", regex) + strconv.FormatUint(hash, 10)) - item := grid.Read(key) + key := fmt.Sprintf("%p_", regex) + strconv.FormatUint(hash, 10) + item := cache.Read(key) if item != nil { - return item.ValueInt64 == 1 + return types.Int8(item.Value) == 1 } b := regex.MatchString(s) if b { - grid.WriteInt64(key, 1, 1800) + cache.Write(key, 1, time.Now().Unix()+1800) } else { - grid.WriteInt64(key, 0, 1800) + cache.Write(key, 0, time.Now().Unix()+1800) } return b } @@ -41,16 +44,19 @@ func MatchBytesCache(regex *regexp.Regexp, byteSlice []byte) bool { } hash := siphash.Hash(0, 0, byteSlice) - key := []byte(fmt.Sprintf("%p_", regex) + strconv.FormatUint(hash, 10)) - item := grid.Read(key) + key := fmt.Sprintf("%p_", regex) + strconv.FormatUint(hash, 10) + item := cache.Read(key) if item != nil { - return item.ValueInt64 == 1 + return types.Int8(item.Value) == 1 + } + if item != nil { + return types.Int8(item.Value) == 1 } b := regex.Match(byteSlice) if b { - grid.WriteInt64(key, 1, 1800) + cache.Write(key, 1, time.Now().Unix()+1800) } else { - grid.WriteInt64(key, 0, 1800) + cache.Write(key, 0, time.Now().Unix()+1800) } return b } diff --git a/internal/waf/utils/utils_test.go b/internal/waf/utils/utils_test.go index f9ae8e8..653b7ee 100644 --- a/internal/waf/utils/utils_test.go +++ b/internal/waf/utils/utils_test.go @@ -41,6 +41,7 @@ func TestMatchRemoteCache(t *testing.T) { func BenchmarkMatchStringCache(b *testing.B) { data := strings.Repeat("HELLO", 512) regex := regexp.MustCompile(`(?iU)\b(eval|system|exec|execute|passthru|shell_exec|phpinfo)\b`) + _ = MatchStringCache(regex, data) for i := 0; i < b.N; i++ { _ = MatchStringCache(regex, data)