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

使用泛型优化计数器内存

Browse Source
This commit is contained in:
刘祥超
2023-11-15 15:57:41 +08:00
parent 768384dcf0
commit 59f27215d3
6 changed files with 59 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/nodes/client_conn.go
View File

@@ -305,7 +305,7 @@ func (this *ClientConn) increaseSYNFlood(synFloodConfig *firewallconfigs.SYNFloo
// 非TLS，设置为两倍，防止误封 // 非TLS，设置为两倍，防止误封
minAttempts = 2 * minAttempts minAttempts = 2 * minAttempts
} }
if result >= types.Uint64(minAttempts) { if result >= types.Uint32(minAttempts) {
var timeout = synFloodConfig.TimeoutSeconds var timeout = synFloodConfig.TimeoutSeconds
if timeout <= 0 { if timeout <= 0 {
timeout = 600 timeout = 600

42
internal/utils/counters/counter.go
View File

@@ -13,12 +13,16 @@ import (
const maxItemsPerGroup = 50_000 const maxItemsPerGroup = 50_000
var SharedCounter = NewCounter().WithGC() var SharedCounter = NewCounter[uint32]().WithGC()
type Counter struct { type SupportedUIntType interface {
uint32 | uint64
}
type Counter[T SupportedUIntType] struct {
countMaps uint64 countMaps uint64
locker *syncutils.RWMutex locker *syncutils.RWMutex
itemMaps []map[uint64]*Item itemMaps []map[uint64]*Item[T]
gcTicker *time.Ticker gcTicker *time.Ticker
gcIndex int gcIndex int
@@ -26,18 +30,18 @@ type Counter struct {
} }
// NewCounter create new counter // NewCounter create new counter
func NewCounter() *Counter { func NewCounter[T SupportedUIntType]() *Counter[T] {
var count = utils.SystemMemoryGB() * 8 var count = utils.SystemMemoryGB() * 8
if count < 8 { if count < 8 {
count = 8 count = 8
} }
var itemMaps = []map[uint64]*Item{} var itemMaps = []map[uint64]*Item[T]{}
for i := 0; i < count; i++ { for i := 0; i < count; i++ {
itemMaps = append(itemMaps, map[uint64]*Item{}) itemMaps = append(itemMaps, map[uint64]*Item[T]{})
} }
var counter = &Counter{ var counter = &Counter[T]{
countMaps: uint64(count), countMaps: uint64(count),
locker: syncutils.NewRWMutex(count), locker: syncutils.NewRWMutex(count),
itemMaps: itemMaps, itemMaps: itemMaps,
@@ -47,7 +51,7 @@ func NewCounter() *Counter {
} }
// WithGC start the counter with gc automatically // WithGC start the counter with gc automatically
func (this *Counter) WithGC() *Counter { func (this *Counter[T]) WithGC() *Counter[T] {
if this.gcTicker != nil { if this.gcTicker != nil {
return this return this
} }
@@ -62,7 +66,7 @@ func (this *Counter) WithGC() *Counter {
} }
// Increase key // Increase key
func (this *Counter) Increase(key uint64, lifeSeconds int) uint64 { func (this *Counter[T]) Increase(key uint64, lifeSeconds int) T {
var index = int(key % this.countMaps) var index = int(key % this.countMaps)
this.locker.RLock(index) this.locker.RLock(index)
var item = this.itemMaps[index][key] var item = this.itemMaps[index][key]
@@ -70,7 +74,7 @@ func (this *Counter) Increase(key uint64, lifeSeconds int) uint64 {
if item == nil { if item == nil {
// no need to care about duplication // no need to care about duplication
// always insert new item even when itemMap is full // always insert new item even when itemMap is full
item = NewItem(lifeSeconds) item = NewItem[T](lifeSeconds)
this.locker.Lock(index) this.locker.Lock(index)
this.itemMaps[index][key] = item this.itemMaps[index][key] = item
this.locker.Unlock(index) this.locker.Unlock(index)
@@ -83,12 +87,12 @@ func (this *Counter) Increase(key uint64, lifeSeconds int) uint64 {
} }
// IncreaseKey increase string key // IncreaseKey increase string key
func (this *Counter) IncreaseKey(key string, lifeSeconds int) uint64 { func (this *Counter[T]) IncreaseKey(key string, lifeSeconds int) T {
return this.Increase(this.hash(key), lifeSeconds) return this.Increase(this.hash(key), lifeSeconds)
} }
// Get value of key // Get value of key
func (this *Counter) Get(key uint64) uint64 { func (this *Counter[T]) Get(key uint64) T {
var index = int(key % this.countMaps) var index = int(key % this.countMaps)
this.locker.RLock(index) this.locker.RLock(index)
defer this.locker.RUnlock(index) defer this.locker.RUnlock(index)
@@ -100,12 +104,12 @@ func (this *Counter) Get(key uint64) uint64 {
} }
// GetKey get value of string key // GetKey get value of string key
func (this *Counter) GetKey(key string) uint64 { func (this *Counter[T]) GetKey(key string) T {
return this.Get(this.hash(key)) return this.Get(this.hash(key))
} }
// Reset key // Reset key
func (this *Counter) Reset(key uint64) { func (this *Counter[T]) Reset(key uint64) {
var index = int(key % this.countMaps) var index = int(key % this.countMaps)
this.locker.RLock(index) this.locker.RLock(index)
var item = this.itemMaps[index][key] var item = this.itemMaps[index][key]
@@ -119,12 +123,12 @@ func (this *Counter) Reset(key uint64) {
} }
// ResetKey string key // ResetKey string key
func (this *Counter) ResetKey(key string) { func (this *Counter[T]) ResetKey(key string) {
this.Reset(this.hash(key)) this.Reset(this.hash(key))
} }
// TotalItems get items count // TotalItems get items count
func (this *Counter) TotalItems() int { func (this *Counter[T]) TotalItems() int {
var total = 0 var total = 0
for i := 0; i < int(this.countMaps); i++ { for i := 0; i < int(this.countMaps); i++ {
@@ -137,7 +141,7 @@ func (this *Counter) TotalItems() int {
} }
// GC garbage expired items // GC garbage expired items
func (this *Counter) GC() { func (this *Counter[T]) GC() {
this.gcLocker.Lock() this.gcLocker.Lock()
var gcIndex = this.gcIndex var gcIndex = this.gcIndex
@@ -186,11 +190,11 @@ func (this *Counter) GC() {
} }
} }
func (this *Counter) CountMaps() int { func (this *Counter[T]) CountMaps() int {
return int(this.countMaps) return int(this.countMaps)
} }
// calculate hash of the key // calculate hash of the key
func (this *Counter) hash(key string) uint64 { func (this *Counter[T]) hash(key string) uint64 {
return xxhash.Sum64String(key) return xxhash.Sum64String(key)
} }

18
internal/utils/counters/counter_test.go
View File

@@ -19,7 +19,7 @@ import (
func TestCounter_Increase(t *testing.T) { func TestCounter_Increase(t *testing.T) {
var a = assert.NewAssertion(t) var a = assert.NewAssertion(t)
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
a.IsTrue(counter.Increase(1, 10) == 1) a.IsTrue(counter.Increase(1, 10) == 1)
a.IsTrue(counter.Increase(1, 10) == 2) a.IsTrue(counter.Increase(1, 10) == 2)
a.IsTrue(counter.Increase(2, 10) == 1) a.IsTrue(counter.Increase(2, 10) == 1)
@@ -32,7 +32,7 @@ func TestCounter_Increase(t *testing.T) {
func TestCounter_IncreaseKey(t *testing.T) { func TestCounter_IncreaseKey(t *testing.T) {
var a = assert.NewAssertion(t) var a = assert.NewAssertion(t)
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
a.IsTrue(counter.IncreaseKey("1", 10) == 1) a.IsTrue(counter.IncreaseKey("1", 10) == 1)
a.IsTrue(counter.IncreaseKey("1", 10) == 2) a.IsTrue(counter.IncreaseKey("1", 10) == 2)
a.IsTrue(counter.IncreaseKey("2", 10) == 1) a.IsTrue(counter.IncreaseKey("2", 10) == 1)
@@ -47,7 +47,7 @@ func TestCounter_GC(t *testing.T) {
return return
} }
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
counter.Increase(1, 20) counter.Increase(1, 20)
time.Sleep(1 * time.Second) time.Sleep(1 * time.Second)
counter.Increase(1, 20) counter.Increase(1, 20)
@@ -61,7 +61,7 @@ func TestCounter_GC2(t *testing.T) {
return return
} }
var counter = counters.NewCounter().WithGC() var counter = counters.NewCounter[uint32]().WithGC()
for i := 0; i < 1e5; i++ { for i := 0; i < 1e5; i++ {
counter.Increase(uint64(i), rands.Int(10, 300)) counter.Increase(uint64(i), rands.Int(10, 300))
} }
@@ -79,7 +79,7 @@ func TestCounterMemory(t *testing.T) {
var stat = &runtime.MemStats{} var stat = &runtime.MemStats{}
runtime.ReadMemStats(stat) runtime.ReadMemStats(stat)
var counter = counters.NewCounter().WithGC() var counter = counters.NewCounter[uint32]()
for i := 0; i < 1_000_000; i++ { for i := 0; i < 1_000_000; i++ {
counter.Increase(uint64(i), rands.Int(10, 300)) counter.Increase(uint64(i), rands.Int(10, 300))
} }
@@ -98,7 +98,7 @@ func TestCounterMemory(t *testing.T) {
func BenchmarkCounter_Increase(b *testing.B) { func BenchmarkCounter_Increase(b *testing.B) {
runtime.GOMAXPROCS(4) runtime.GOMAXPROCS(4)
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
b.ResetTimer() b.ResetTimer()
var i uint64 var i uint64
@@ -114,7 +114,7 @@ func BenchmarkCounter_Increase(b *testing.B) {
func BenchmarkCounter_IncreaseKey(b *testing.B) { func BenchmarkCounter_IncreaseKey(b *testing.B) {
runtime.GOMAXPROCS(4) runtime.GOMAXPROCS(4)
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
go func() { go func() {
var ticker = time.NewTicker(100 * time.Millisecond) var ticker = time.NewTicker(100 * time.Millisecond)
@@ -138,7 +138,7 @@ func BenchmarkCounter_IncreaseKey(b *testing.B) {
func BenchmarkCounter_IncreaseKey2(b *testing.B) { func BenchmarkCounter_IncreaseKey2(b *testing.B) {
runtime.GOMAXPROCS(4) runtime.GOMAXPROCS(4)
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
go func() { go func() {
var ticker = time.NewTicker(1 * time.Millisecond) var ticker = time.NewTicker(1 * time.Millisecond)
@@ -162,7 +162,7 @@ func BenchmarkCounter_IncreaseKey2(b *testing.B) {
func BenchmarkCounter_GC(b *testing.B) { func BenchmarkCounter_GC(b *testing.B) {
runtime.GOMAXPROCS(4) runtime.GOMAXPROCS(4)
var counter = counters.NewCounter() var counter = counters.NewCounter[uint32]()
for i := uint64(0); i < 1e5; i++ { for i := uint64(0); i < 1e5; i++ {
counter.IncreaseKey(types.String(i), 20) counter.IncreaseKey(types.String(i), 20)

30
internal/utils/counters/item.go
View File

@@ -6,16 +6,16 @@ import (
"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime" "github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
) )
type Item struct { const spanMaxValue = 10_000_000
lifeSeconds int64
spanSeconds int64
spans []uint64
type Item[T SupportedUIntType] struct {
spans []T
lastUpdateTime int64 lastUpdateTime int64
lifeSeconds int64
spanSeconds int64
} }
func NewItem(lifeSeconds int) *Item { func NewItem[T SupportedUIntType](lifeSeconds int) *Item[T] {
if lifeSeconds <= 0 { if lifeSeconds <= 0 {
lifeSeconds = 60 lifeSeconds = 60
} }
@@ -25,21 +25,23 @@ func NewItem(lifeSeconds int) *Item {
} }
var countSpans = lifeSeconds/spanSeconds + 1 /** prevent index out of bounds **/ var countSpans = lifeSeconds/spanSeconds + 1 /** prevent index out of bounds **/
return &Item{ return &Item[T]{
lifeSeconds: int64(lifeSeconds), lifeSeconds: int64(lifeSeconds),
spanSeconds: int64(spanSeconds), spanSeconds: int64(spanSeconds),
spans: make([]uint64, countSpans), spans: make([]T, countSpans),
lastUpdateTime: fasttime.Now().Unix(), lastUpdateTime: fasttime.Now().Unix(),
} }
} }
func (this *Item) Increase() (result uint64) { func (this *Item[T]) Increase() (result T) {
var currentTime = fasttime.Now().Unix() var currentTime = fasttime.Now().Unix()
var currentSpanIndex = this.calculateSpanIndex(currentTime) var currentSpanIndex = this.calculateSpanIndex(currentTime)
// return quickly // return quickly
if this.lastUpdateTime == currentTime { if this.lastUpdateTime == currentTime {
if this.spans[currentSpanIndex] < spanMaxValue {
this.spans[currentSpanIndex]++ this.spans[currentSpanIndex]++
}
for _, count := range this.spans { for _, count := range this.spans {
result += count result += count
} }
@@ -69,7 +71,9 @@ func (this *Item) Increase() (result uint64) {
} }
} }
if this.spans[currentSpanIndex] < spanMaxValue {
this.spans[currentSpanIndex]++ this.spans[currentSpanIndex]++
}
this.lastUpdateTime = currentTime this.lastUpdateTime = currentTime
for _, count := range this.spans { for _, count := range this.spans {
@@ -79,7 +83,7 @@ func (this *Item) Increase() (result uint64) {
return return
} }
func (this *Item) Sum() (result uint64) { func (this *Item[T]) Sum() (result T) {
if this.lastUpdateTime == 0 { if this.lastUpdateTime == 0 {
return 0 return 0
} }
@@ -104,16 +108,16 @@ func (this *Item) Sum() (result uint64) {
return result return result
} }
func (this *Item) Reset() { func (this *Item[T]) Reset() {
for index := range this.spans { for index := range this.spans {
this.spans[index] = 0 this.spans[index] = 0
} }
} }
func (this *Item) IsExpired(currentTime int64) bool { func (this *Item[T]) IsExpired(currentTime int64) bool {
return this.lastUpdateTime < currentTime-this.lifeSeconds-this.spanSeconds return this.lastUpdateTime < currentTime-this.lifeSeconds-this.spanSeconds
} }
func (this *Item) calculateSpanIndex(timestamp int64) int { func (this *Item[T]) calculateSpanIndex(timestamp int64) int {
return int(timestamp % this.lifeSeconds / this.spanSeconds) return int(timestamp % this.lifeSeconds / this.spanSeconds)
} }

8
internal/utils/counters/item_test.go
View File

@@ -17,7 +17,7 @@ func TestItem_Increase(t *testing.T) {
return return
} }
var item = counters.NewItem(10) var item = counters.NewItem[uint32](10)
t.Log(item.Increase(), item.Sum()) t.Log(item.Increase(), item.Sum())
time.Sleep(1 * time.Second) time.Sleep(1 * time.Second)
t.Log(item.Increase(), item.Sum()) t.Log(item.Increase(), item.Sum())
@@ -41,7 +41,7 @@ func TestItem_Increase2(t *testing.T) {
var a = assert.NewAssertion(t) var a = assert.NewAssertion(t)
var item = counters.NewItem(20) var item = counters.NewItem[uint32](20)
for i := 0; i < 100; i++ { for i := 0; i < 100; i++ {
t.Log(item.Increase(), item.Sum(), timeutil.Format("H:i:s")) t.Log(item.Increase(), item.Sum(), timeutil.Format("H:i:s"))
time.Sleep(2 * time.Second) time.Sleep(2 * time.Second)
@@ -58,7 +58,7 @@ func TestItem_IsExpired(t *testing.T) {
var currentTime = time.Now().Unix() var currentTime = time.Now().Unix()
var item = counters.NewItem(10) var item = counters.NewItem[uint32](10)
t.Log(item.IsExpired(currentTime)) t.Log(item.IsExpired(currentTime))
time.Sleep(10 * time.Second) time.Sleep(10 * time.Second)
t.Log(item.IsExpired(currentTime)) t.Log(item.IsExpired(currentTime))
@@ -73,7 +73,7 @@ func BenchmarkItem_Increase(b *testing.B) {
b.RunParallel(func(pb *testing.PB) { b.RunParallel(func(pb *testing.PB) {
for pb.Next() { for pb.Next() {
var item = counters.NewItem(60) var item = counters.NewItem[uint32](60)
item.Increase() item.Increase()
item.Sum() item.Sum()
} }

5
internal/waf/checkpoints/cc2.go
View File

@@ -76,7 +76,8 @@ func (this *CC2Checkpoint) RequestValue(req requests.Request, param string, opti
} }
var ccKey = "WAF-CC-" + types.String(ruleId) + "-" + strings.Join(keyValues, "@") var ccKey = "WAF-CC-" + types.String(ruleId) + "-" + strings.Join(keyValues, "@")
value = counters.SharedCounter.IncreaseKey(ccKey, period) var ccValue = counters.SharedCounter.IncreaseKey(ccKey, period)
value = ccValue
// 基于指纹统计 // 基于指纹统计
var enableFingerprint = true var enableFingerprint = true
@@ -96,7 +97,7 @@ func (this *CC2Checkpoint) RequestValue(req requests.Request, param string, opti
} }
var fpCCKey = "WAF-CC-" + types.String(ruleId) + "-" + strings.Join(fpKeyValues, "@") var fpCCKey = "WAF-CC-" + types.String(ruleId) + "-" + strings.Join(fpKeyValues, "@")
var fpValue = counters.SharedCounter.IncreaseKey(fpCCKey, period) var fpValue = counters.SharedCounter.IncreaseKey(fpCCKey, period)
if fpValue > value.(uint64) { if fpValue > ccValue {
value = fpValue value = fpValue
} }
} }