限制各个内存缓存的条目数量，防止OOM

Author: iwind

teageo/utils.go

@@ -9,7 +9,7 @@ import (
 	"strings"
 )
 
-var ip2cityGrid = teamemory.NewGrid(32)
+var ip2cityGrid = teamemory.NewGrid(32, teamemory.NewLimitCountOpt(100_0000))
 
 // 加载Geo-City数据库
 func SetupDB() {

tealogs/accesslogs/access_log.go

@@ -20,7 +20,7 @@ import (
 	"time"
 )
 
-var userAgentGrid = teamemory.NewGrid(32)
+var userAgentGrid = teamemory.NewGrid(32, teamemory.NewLimitCountOpt(100_0000))
 var charsetReg = regexp.MustCompile("(?i)charset\\s*=\\s*([\\w-]+)")
 var headerReg = regexp.MustCompile("([A-Z])")
 

teamemory/cell.go

@@ -7,7 +7,8 @@ import (
 )
 
 type Cell struct {
-	LimitSize int64
+	LimitSize  int64
+	LimitCount int
 
 	mapping    map[uint64]*Item // key => item
 	list       *List            // { item1, item2, ... }
@@ -37,6 +38,12 @@ func (this *Cell) Write(hashKey uint64, item *Item) {
 		}
 	}
 
+	// limit count
+	if this.LimitCount > 0 && len(this.mapping) >= this.LimitCount {
+		this.locker.Unlock()
+		return
+	}
+
 	// trim memory
 	size := item.Size()
 	shouldTrim := false

teamemory/grid.go

@@ -24,7 +24,8 @@ type Grid struct {
 
 	gzipLevel int
 
-	limitSize int64
+	limitSize  int64
+	limitCount int
 }
 
 func NewGrid(countCells int, opt ...interface{}) *Grid {
@@ -38,6 +39,8 @@ func NewGrid(countCells int, opt ...interface{}) *Grid {
 			grid.gzipLevel = x.Level
 		case *LimitSizeOpt:
 			grid.limitSize = x.Size
+		case *LimitCountOpt:
+			grid.limitCount = x.Count
 		case *RecycleIntervalOpt:
 			grid.recycleInterval = x.Interval
 		}
@@ -52,6 +55,8 @@ func NewGrid(countCells int, opt ...interface{}) *Grid {
 	for i := 0; i < countCells; i++ {
 		cell := NewCell()
 		cell.LimitSize = int64(math.Floor(float64(grid.limitSize) / float64(countCells)))
+		cell.LimitCount = int(math.Floor(float64(grid.limitCount)) / float64(countCells))
+
 		cells = append(cells, cell)
 	}
 	grid.cells = cells

teamemory/grid_test.go

@@ -27,7 +27,7 @@ func TestMemoryGrid_Write(t *testing.T) {
 	grid.Delete([]byte("abc"))
 	t.Log(grid.Read([]byte("abc")))
 
-	for i := 0; i < 100; i ++ {
+	for i := 0; i < 100; i++ {
 		grid.WriteInt64([]byte(fmt.Sprintf("%d", i)), 123, 1)
 	}
 
@@ -45,6 +45,14 @@ func TestMemoryGrid_Write(t *testing.T) {
 	grid.Destroy()
 }
 
+func TestMemoryGrid_Write_LimitCount(t *testing.T) {
+	grid := NewGrid(2, NewLimitCountOpt(10))
+	for i := 0; i < 100; i++ {
+		grid.WriteInt64([]byte(strconv.Itoa(i)), int64(i), 30)
+	}
+	t.Log(grid.Stat().CountItems, "items")
+}
+
 func TestMemoryGrid_Compress(t *testing.T) {
 	grid := NewGrid(5, NewCompressOpt(1))
 	grid.WriteString([]byte("hello"), strings.Repeat("abcd", 10240), 30)
@@ -54,7 +62,7 @@ func TestMemoryGrid_Compress(t *testing.T) {
 
 func BenchmarkMemoryGrid_Performance(b *testing.B) {
 	grid := NewGrid(1024)
-	for i := 0; i < b.N; i ++ {
+	for i := 0; i < b.N; i++ {
 		grid.WriteInt64([]byte("key:"+strconv.Itoa(i)), int64(i), 3600)
 	}
 }
@@ -68,7 +76,7 @@ func TestMemoryGrid_Performance(t *testing.T) {
 
 	s := []byte(strings.Repeat("abcd", 10*1024))
 
-	for i := 0; i < 100000; i ++ {
+	for i := 0; i < 100000; i++ {
 		grid.WriteBytes([]byte(fmt.Sprintf("key:%d_%d", i, 1)), s, 3600)
 		item := grid.Read([]byte(fmt.Sprintf("key:%d_%d", i, 1)))
 		if item != nil {
@@ -97,10 +105,10 @@ func TestMemoryGrid_Performance_Concurrent(t *testing.T) {
 	wg := sync.WaitGroup{}
 	wg.Add(runtime.NumCPU())
 
-	for c := 0; c < runtime.NumCPU(); c ++ {
+	for c := 0; c < runtime.NumCPU(); c++ {
 		go func(c int) {
 			defer wg.Done()
-			for i := 0; i < 50000; i ++ {
+			for i := 0; i < 50000; i++ {
 				grid.WriteBytes([]byte(fmt.Sprintf("key:%d_%d", i, c)), s, 3600)
 				item := grid.Read([]byte(fmt.Sprintf("key:%d_%d", i, c)))
 				if item != nil {
@@ -128,7 +136,7 @@ func TestMemoryGrid_CompressPerformance(t *testing.T) {
 	now := time.Now()
 	data := []byte(strings.Repeat("abcd", 1024))
 
-	for i := 0; i < 100000; i ++ {
+	for i := 0; i < 100000; i++ {
 		grid.WriteBytes([]byte(fmt.Sprintf("key:%d", i)), data, 3600)
 		item := grid.Read([]byte(fmt.Sprintf("key:%d", i+100)))
 		if item != nil {

teamemory/item.go

@@ -9,7 +9,7 @@ import (
 	"unsafe"
 )
 
-type ItemType = int
+type ItemType = int8
 
 const (
 	ItemInt64     = 1
@@ -33,6 +33,8 @@ type Item struct {
 	// linked list
 	Prev *Item
 	Next *Item
+
+	size int64
 }
 
 func NewItem(key []byte, dataType ItemType) *Item {
@@ -79,5 +81,8 @@ func (this *Item) String() string {
 }
 
 func (this *Item) Size() int64 {
-	return int64(int(unsafe.Sizeof(*this)) + len(this.Key) + len(this.ValueBytes))
+	if this.size == 0 {
+		this.size = int64(int(unsafe.Sizeof(*this)) + len(this.Key) + len(this.ValueBytes))
+	}
+	return this.size
 }

teamemory/opt_limit_count.go

@@ -0,0 +1,11 @@
+package teamemory
+
+type LimitCountOpt struct {
+	Count int
+}
+
+func NewLimitCountOpt(count int) *LimitCountOpt {
+	return &LimitCountOpt{
+		Count: count,
+	}
+}

teawaf/checkpoints/cc.go

@@ -28,7 +28,7 @@ func (this *CCCheckpoint) Start() {
 	if this.grid != nil {
 		this.grid.Destroy()
 	}
-	this.grid = teamemory.NewGrid(32)
+	this.grid = teamemory.NewGrid(32, teamemory.NewLimitCountOpt(1000_0000))
 }
 
 func (this *CCCheckpoint) RequestValue(req *requests.Request, param string, options map[string]string) (value interface{}, sysErr error, userErr error) {

teawaf/utils/utils.go

@@ -9,7 +9,7 @@ import (
 	"strconv"
 )
 
-var grid = teamemory.NewGrid(32)
+var grid = teamemory.NewGrid(32, teamemory.NewLimitCountOpt(1000_0000))
 
 // 正则表达式匹配字符串，并缓存结果
 func MatchStringCache(regex *regexp.Regexp, s string) bool {
@@ -21,9 +21,9 @@ func MatchStringCache(regex *regexp.Regexp, s string) bool {
 	}
 	b := regex.MatchString(s)
 	if b {
-		grid.WriteInt64(key, 1, 3600)
+		grid.WriteInt64(key, 1, 1800)
 	} else {
-		grid.WriteInt64(key, 0, 3600)
+		grid.WriteInt64(key, 0, 1800)
 	}
 	return b
 }
@@ -38,9 +38,9 @@ func MatchBytesCache(regex *regexp.Regexp, byteSlice []byte) bool {
 	}
 	b := regex.Match(byteSlice)
 	if b {
-		grid.WriteInt64(key, 1, 3600)
+		grid.WriteInt64(key, 1, 1800)
 	} else {
-		grid.WriteInt64(key, 0, 3600)
+		grid.WriteInt64(key, 0, 1800)
 	}
 	return b
 }

teaweb/actions/default/agents/board/scripts/engine.go

@@ -22,7 +22,7 @@ import (
 	"time"
 )
 
-var engineCache = teamemory.NewGrid(1)
+var engineCache = teamemory.NewGrid(1, teamemory.NewLimitCountOpt(1000))
 var dayReg = regexp.MustCompile(`^(\d+)-(\d+)-(\d+)$`)
 
 // 脚本引擎