刘河 6 lat temu
rodzic
commit
72290275cc
3 zmienionych plików z 30 dodań i 13 usunięć
  1. 3 2
      utils/conn.go
  2. 23 0
      utils/pool.go
  3. 4 11
      utils/util.go

+ 3 - 2
utils/conn.go

@@ -17,8 +17,6 @@ import (
 )
 
 const cryptKey = "1234567812345678"
-const poolSize = 64 * 1024
-const poolSizeSmall = 10
 
 type CryptConn struct {
 	conn  net.Conn
@@ -116,6 +114,7 @@ func (s *SnappyConn) Read(b []byte) (n int, err error) {
 		}
 	}()
 	buf := bufPool.Get().([]byte)
+	defer bufPool.Put(buf)
 	if n, err = s.r.Read(buf); err != nil {
 		return
 	}
@@ -152,8 +151,10 @@ func (s *Conn) ReadLen(cLen int) ([]byte, error) {
 	var buf []byte
 	if cLen <= poolSizeSmall {
 		buf = bufPoolSmall.Get().([]byte)[:cLen]
+		defer bufPoolSmall.Put(buf)
 	} else {
 		buf = bufPool.Get().([]byte)[:cLen]
+		defer bufPool.Put(buf)
 	}
 	if n, err := io.ReadFull(s, buf); err != nil || n != cLen {
 		return buf, errors.New("读取指定长度错误" + err.Error())

+ 23 - 0
utils/pool.go

@@ -0,0 +1,23 @@
+package utils
+
+import "sync"
+
+const poolSize = 64 * 1024
+const poolSizeSmall = 10
+const poolSizeCopy = 32 * 1024
+
+var bufPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, poolSize)
+	},
+}
+var bufPoolSmall = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, poolSizeSmall)
+	},
+}
+var bufPoolCopy = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, poolSizeCopy)
+	},
+}

+ 4 - 11
utils/util.go

@@ -145,21 +145,13 @@ func GetIntNoerrByStr(str string) int {
 	return i
 }
 
-var bufPool = sync.Pool{
-	New: func() interface{} {
-		return make([]byte, poolSize)
-	},
-}
-var bufPoolSmall = sync.Pool{
-	New: func() interface{} {
-		return make([]byte, poolSizeSmall)
-	},
-}
+
 // io.copy的优化版,读取buffer长度原为32*1024,与snappy不同,导致读取出的内容存在差异,不利于解密,特此修改
 //废除
 func copyBuffer(dst io.Writer, src io.Reader) (written int64, err error) {
 	//TODO 回收问题
-	buf := bufPool.Get().([]byte)[:32*1024]
+	buf := bufPoolCopy.Get().([]byte)
+	defer bufPoolCopy.Put(buf)
 	for {
 		nr, er := src.Read(buf)
 		if nr > 0 {
@@ -190,6 +182,7 @@ func copyBuffer(dst io.Writer, src io.Reader) (written int64, err error) {
 func FlushConn(c net.Conn) {
 	c.SetReadDeadline(time.Now().Add(time.Second * 3))
 	buf := bufPool.Get().([]byte)
+	defer bufPool.Put(buf)
 	for {
 		if _, err := c.Read(buf); err != nil {
 			break