优化代码

internal/utils/byte_pool.go

@@ -1,16 +1,25 @@
 package utils
 
-var BytePool1024 = NewBytePool(20480, 1024)
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/goman"
+	"github.com/iwind/TeaGo/Tea"
+	"time"
+)
 
-// pool for get byte slice
+var BytePool1k = NewBytePool(20480, 1024)
+var BytePool4k = NewBytePool(20480, 4*1024)
+var BytePool16k = NewBytePool(40960, 16*1024)
+var BytePool32k = NewBytePool(20480, 32*1024)
+
+// BytePool pool for get byte slice
 type BytePool struct {
-	c      chan []byte
-	length int
-
-	lastSize int
+	c       chan []byte
+	maxSize int
+	length  int
+	hasNew  bool
 }
 
-// 创建新对象
+// NewBytePool 创建新对象
 func NewBytePool(maxSize, length int) *BytePool {
 	if maxSize <= 0 {
 		maxSize = 1024
@@ -18,24 +27,47 @@ func NewBytePool(maxSize, length int) *BytePool {
 	if length <= 0 {
 		length = 128
 	}
-	pool := &BytePool{
-		c:      make(chan []byte, maxSize),
-		length: length,
+	var pool = &BytePool{
+		c:       make(chan []byte, maxSize),
+		maxSize: maxSize,
+		length:  length,
 	}
+
+	pool.init()
+
 	return pool
 }
 
-// 获取一个新的byte slice
+// 初始化
+func (this *BytePool) init() {
+	var ticker = time.NewTicker(2 * time.Minute)
+	if Tea.IsTesting() {
+		ticker = time.NewTicker(5 * time.Second)
+	}
+	goman.New(func() {
+		for range ticker.C {
+			if this.hasNew {
+				this.hasNew = false
+				continue
+			}
+
+			this.Purge()
+		}
+	})
+}
+
+// Get 获取一个新的byte slice
 func (this *BytePool) Get() (b []byte) {
 	select {
 	case b = <-this.c:
 	default:
 		b = make([]byte, this.length)
+		this.hasNew = true
 	}
 	return
 }
 
-// 放回一个使用过的byte slice
+// Put 放回一个使用过的byte slice
 func (this *BytePool) Put(b []byte) {
 	if cap(b) != this.length {
 		return
@@ -47,7 +79,30 @@ func (this *BytePool) Put(b []byte) {
 	}
 }
 
-// 当前的数量
+// Length 单个字节slice长度
+func (this *BytePool) Length() int {
+	return this.length
+}
+
+// Size 当前的数量
 func (this *BytePool) Size() int {
 	return len(this.c)
 }
+
+// Purge 清理
+func (this *BytePool) Purge() {
+	// 1%
+	var count = len(this.c) / 100
+	if count == 0 {
+		return
+	}
+
+Loop:
+	for i := 0; i < count; i++ {
+		select {
+		case <-this.c:
+		default:
+			break Loop
+		}
+	}
+}