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

改进WAF缓存性能

Browse Source
This commit is contained in:
GoEdgeLab
2020-11-21 21:43:03 +08:00
parent c32959a6c8
commit 5485a5f663
11 changed files with 98 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
internal/cache/item.go vendored
View File

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

19
internal/cache/cache.go → internal/ttlcache/cache.go
View File

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

40
internal/cache/cache_test.go → internal/ttlcache/cache_test.go
View File

@@ -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++ {

6
internal/ttlcache/item.go Normal file
View File

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

10
internal/cache/option.go → internal/ttlcache/option.go
View File

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

11
internal/cache/piece.go → internal/ttlcache/piece.go
View File

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

30
internal/cache/piece_test.go → internal/ttlcache/piece_test.go
View File

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

2
internal/cache/utils.go → internal/ttlcache/utils.go
View File

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

2
internal/cache/utils_test.go → internal/ttlcache/utils_test.go
View File

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

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

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

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

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