改进WAF缓存性能

internal/cache/item.go

@@ -1,6 +0,0 @@
-package cache
-
-type Item struct {
-	value     interface{}
-	expiredAt int64
-}

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

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.Write("a", 1, time.Now().Unix()+3600)
-	cache.Add("b", 2, time.Now().Unix()+3601)
+	cache.Write("b", 2, time.Now().Unix()+3601)
-	cache.Add("a", 1, time.Now().Unix()+3602)
+	cache.Write("a", 1, time.Now().Unix()+3602)
-	cache.Add("d", 1, time.Now().Unix()+1)
+	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.Write("a", 1, time.Now().Unix()+1)
-	cache.Add("b", 2, time.Now().Unix()+2)
+	cache.Write("b", 2, time.Now().Unix()+2)
-	cache.Add("c", 3, time.Now().Unix()+3)
+	cache.Write("c", 3, time.Now().Unix()+3)
-	cache.Add("d", 4, time.Now().Unix()+4)
+	cache.Write("d", 4, time.Now().Unix()+4)
-	cache.Add("e", 5, time.Now().Unix()+10)
+	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++ {

internal/ttlcache/item.go

@@ -0,0 +1,6 @@
+package ttlcache
+
+type Item struct {
+	Value     interface{}
+	expiredAt int64
+}

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

internal/ttlcache/piece.go

@@ -1,4 +1,4 @@
-package cache
+package ttlcache
 
 import (
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
@@ -8,15 +8,20 @@ import (
 
 type Piece struct {
 	m        map[uint64]*Item
+	maxItems int
 	locker   sync.RWMutex
 }
 
-func NewPiece() *Piece {
+func NewPiece(maxItems int) *Piece {
-	return &Piece{m: map[uint64]*Item{}}
+	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()
 }

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 := NewPiece(10)
-	piece.Add(1, &Item{value: 1, expiredAt: time.Now().Unix() + 1})
+	piece.Add(1, &Item{Value: 1, expiredAt: time.Now().Unix() + 1})
-	piece.Add(2, &Item{value: 2, 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.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)

internal/ttlcache/utils.go

@@ -1,4 +1,4 @@
-package cache
+package ttlcache
 
 import "github.com/dchest/siphash"
 

internal/ttlcache/utils_test.go

@@ -1,4 +1,4 @@
-package cache
+package ttlcache
 
 import (
 	"runtime"

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))
+	key := fmt.Sprintf("%p_", regex) + strconv.FormatUint(hash, 10)
-	item := grid.Read(key)
+	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))
+	key := fmt.Sprintf("%p_", regex) + strconv.FormatUint(hash, 10)
-	item := grid.Read(key)
+	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
 }

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)