add memory pool

conn/tcp2/tcp2.go

@@ -8,12 +8,14 @@ import (
 	"io"
 	"log"
 	"net"
+	"slices"
 	"strings"
 	"time"
 
 	"git.me9.top/git/tinymq/config"
 	"git.me9.top/git/tinymq/conn"
 	"git.me9.top/git/tinymq/util"
+	"git.me9.top/git/tinymq/util/pool"
 )
 
 const PROTO string = "tcp"
@@ -80,7 +82,8 @@ func Server(cf *config.Config, bind string, hash string, fn conn.ServerConnectFu
 				var div []byte
 				if ci.IvLen > 0 {
 					// 服务端 IV
-					eiv = make([]byte, ci.IvLen)
+					eiv = pool.Get(ci.IvLen)
+					defer pool.Put(eiv)
 					_, err = rand.Read(eiv)
 					if err != nil {
 						log.Println("[tcp2 Server rand.Read ERROR]", err)
@@ -99,7 +102,8 @@ func Server(cf *config.Config, bind string, hash string, fn conn.ServerConnectFu
 						log.Println("[tcp2 Server SetReadDeadline ERROR]", err)
 						return
 					}
-					div = make([]byte, ci.IvLen)
+					div = pool.Get(ci.IvLen)
+					defer pool.Put(div)
 					_, err := io.ReadFull(conn, div)
 					if err != nil {
 						log.Println("[tcp2 Server ReadFull ERROR]", err)
@@ -156,7 +160,8 @@ func Dial(cf *config.Config, addr string, hash string) (conn.Connect, error) {
 	var div []byte
 	if ci.IvLen > 0 {
 		// 客户端 IV
-		eiv = make([]byte, ci.IvLen)
+		eiv = pool.Get(ci.IvLen)
+		defer pool.Put(eiv)
 		_, err = rand.Read(eiv)
 		if err != nil {
 			log.Println("[tcp2 Client rand.Read ERROR]", err)
@@ -175,7 +180,8 @@ func Dial(cf *config.Config, addr string, hash string) (conn.Connect, error) {
 			log.Println("[tcp2 Client SetReadDeadline ERROR]", err)
 			return nil, err
 		}
-		div = make([]byte, ci.IvLen)
+		div = pool.Get(ci.IvLen)
+		defer pool.Put(div)
 		_, err := io.ReadFull(conn, div)
 		if err != nil {
 			log.Println("[tcp2 Client ReadFull ERROR]", err)
@@ -197,9 +203,31 @@ func Dial(cf *config.Config, addr string, hash string) (conn.Connect, error) {
 	return c, nil
 }
 
-// 发送数据到网络
-// 如果有加密函数的话会直接修改源数据
-func (c *Tcp2) WriteRaw(buf []byte) (err error) {
+// 发送数据包到网络
+// recycle 指示是否需要将raw放回内存池
+func (c *Tcp2) WriteRawPackage(raw []byte, recycle bool) (err error) {
+	var buf []byte
+	var index int
+	dlen := len(raw)
+	if dlen >= MAX_LENGTH {
+		buf = pool.Get(dlen + 2 + 4 + 1)
+		index = 2 + 4
+		binary.BigEndian.PutUint16(buf[:2], MAX_LENGTH)
+		binary.BigEndian.PutUint32(buf[2:6], uint32(dlen))
+	} else {
+		buf = pool.Get(dlen + 2 + 1)
+		index = 2
+		binary.BigEndian.PutUint16(buf[:2], uint16(dlen))
+	}
+	copy(buf[index:index+dlen], raw)
+	index += dlen
+	buf[index] = util.CRC8(raw)
+
+	if recycle {
+		pool.Put(raw)
+	}
+	defer pool.Put(buf)
+
 	if c.cipher != nil {
 		c.cipher.Encrypt(buf, buf)
 	}
@@ -217,20 +245,62 @@ func (c *Tcp2) WriteRaw(buf []byte) (err error) {
 	}
 }
 
-// 申请内存并写入数据长度信息
-// 还多申请一个字节用于保存crc
-func (c *Tcp2) writeDataLen(dlen int) (buf []byte, start int) {
+// 从连接中读取数据包，已经剥离了头部长度和尾部crc
+// recycle 指示是否需要手动将缓存放回内存池
+func (c *Tcp2) ReadRawPackage(deadline int) (buf []byte, recycle bool, err error) {
+	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
+	if err != nil {
+		return nil, false, err
+	}
+	buf = pool.Get(2)
+	// 读取数据流长度
+	_, err = io.ReadFull(c.conn, buf)
+	if err != nil {
+		pool.Put(buf)
+		return nil, false, err
+	}
+	// 将读出来的数据进行解密
+	if c.cipher != nil {
+		c.cipher.Decrypt(buf, buf)
+	}
+	dlen := uint32(binary.BigEndian.Uint16(buf))
+	pool.Put(buf)
+	if dlen < 2 {
+		return nil, false, errors.New("length is less to 2")
+	}
 	if dlen >= MAX_LENGTH {
-		buf = make([]byte, dlen+2+4+1)
-		start = 2 + 4
-		binary.BigEndian.PutUint16(buf[:2], MAX_LENGTH)
-		binary.BigEndian.PutUint32(buf[2:6], uint32(dlen))
-	} else {
-		buf = make([]byte, dlen+2+1)
-		start = 2
-		binary.BigEndian.PutUint16(buf[:2], uint16(dlen))
+		// 数据包比较大，通过后面的4位长度来表示实际长度
+		buf := pool.Get(4)
+		_, err := io.ReadFull(c.conn, buf)
+		if err != nil {
+			pool.Put(buf)
+			return nil, false, err
+		}
+		if c.cipher != nil {
+			c.cipher.Decrypt(buf, buf)
+		}
+		dlen = binary.BigEndian.Uint32(buf)
+		pool.Put(buf)
+		if dlen < MAX_LENGTH {
+			return nil, false, errors.New("wrong length in read message")
+		}
 	}
-	return
+	// 读取指定长度的数据
+	buf = pool.Get(int(dlen) + 1) // 最后一个是crc的值
+	_, err = io.ReadFull(c.conn, buf)
+	if err != nil {
+		pool.Put(buf)
+		return nil, false, err
+	}
+	if c.cipher != nil {
+		c.cipher.Decrypt(buf, buf)
+	}
+	// 检查CRC8
+	if util.CRC8(buf[:dlen]) != buf[dlen] {
+		pool.Put(buf)
+		return nil, false, errors.New("CRC error")
+	}
+	return buf[:dlen], true, nil
 }
 
 // 发送Auth信息
@@ -247,9 +317,8 @@ func (c *Tcp2) WriteAuthInfo(channel string, auth []byte) (err error) {
 	}
 	// id(65502)+proto(string)+version(uint8)+option(byte)+channel(string)+auth([]byte)
 	dlen := 2 + 1 + protoLen + 1 + 1 + 2 + channelLen + len(auth)
-
-	buf, start := c.writeDataLen(dlen)
-	index := start
+	buf := pool.Get(dlen)
+	index := 0
 	binary.BigEndian.PutUint16(buf[index:index+2], config.ID_AUTH)
 	index += 2
 
@@ -274,59 +343,7 @@ func (c *Tcp2) WriteAuthInfo(channel string, auth []byte) (err error) {
 	copy(buf[index:index+channelLen], []byte(channel))
 	index += channelLen
 	copy(buf[index:], auth)
-	buf[start+dlen] = util.CRC8(buf[start : start+dlen])
-	return c.WriteRaw(buf)
-}
-
-// 从连接中读取信息
-func (c *Tcp2) ReadRaw(deadline int) ([]byte, error) {
-	buf := make([]byte, 2)
-	err := c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
-	if err != nil {
-		return nil, err
-	}
-	// 读取数据流长度
-	_, err = io.ReadFull(c.conn, buf)
-	if err != nil {
-		return nil, err
-	}
-	// 将读出来的数据进行解密
-	if c.cipher != nil {
-		c.cipher.Decrypt(buf, buf)
-	}
-	dlen := uint32(binary.BigEndian.Uint16(buf))
-	if dlen < 2 {
-		return nil, errors.New("length is less to 2")
-	}
-	if dlen >= MAX_LENGTH {
-		// 数据包比较大，通过后面的4位长度来表示实际长度
-		buf = make([]byte, 4)
-		_, err := io.ReadFull(c.conn, buf)
-		if err != nil {
-			return nil, err
-		}
-		if c.cipher != nil {
-			c.cipher.Decrypt(buf, buf)
-		}
-		dlen = binary.BigEndian.Uint32(buf)
-		if dlen < MAX_LENGTH {
-			return nil, errors.New("wrong length in read message")
-		}
-	}
-	// 读取指定长度的数据
-	buf = make([]byte, dlen+1) // 最后一个是crc的值
-	_, err = io.ReadFull(c.conn, buf)
-	if err != nil {
-		return nil, err
-	}
-	if c.cipher != nil {
-		c.cipher.Decrypt(buf, buf)
-	}
-	// 检查CRC8
-	if util.CRC8(buf[:dlen]) != buf[dlen] {
-		return nil, errors.New("CRC error")
-	}
-	return buf[:dlen], nil
+	return c.WriteRawPackage(buf, true)
 }
 
 // 获取Auth信息
@@ -338,10 +355,13 @@ func (c *Tcp2) ReadAuthInfo() (proto string, version uint8, channel string, auth
 			return
 		}
 	}()
-	msg, err := c.ReadRaw(c.cf.ReadWait)
+	msg, recycle, err := c.ReadRawPackage(c.cf.ReadWait)
 	if err != nil {
 		return
 	}
+	if recycle {
+		defer pool.Put(msg)
+	}
 	msgLen := len(msg)
 	if msgLen < 9 {
 		err = errors.New("wrong message length")
@@ -387,13 +407,17 @@ func (c *Tcp2) ReadAuthInfo() (proto string, version uint8, channel string, auth
 	channel = string(msg[index : index+channelLen])
 	index += channelLen
 
-	auth = msg[index:]
+	if recycle {
+		auth = slices.Clone(msg[index:])
+	} else {
+		auth = msg[index:]
+	}
 	return
 }
 
 // 发送请求数据包到网络
 func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
-	// 为了区分请求还是响应包，命令字符串不能超过127个字节，如果超过则截断
+	// 为了区分请求还是响应包，命令字符串不能超过127个字节
 	cmdLen := len(cmd)
 	if cmdLen > 0x7F {
 		return errors.New("command length is more than 0x7F")
@@ -401,9 +425,10 @@ func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
 	dlen := len(data)
 	if c.compress && dlen > 0 {
 		compressedData, ok, _ := util.CompressData(data)
+		// id(uint16)+cmdLen(byte)+len(cmd)+option(byte)+len(data)
 		ddlen := 2 + 1 + cmdLen + 1 + len(compressedData)
-		buf, start := c.writeDataLen(ddlen)
-		index := start
+		buf := pool.Get(ddlen)
+		index := 0
 		binary.BigEndian.PutUint16(buf[index:index+2], id)
 		index += 2
 
@@ -423,12 +448,11 @@ func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
 		index++
 
 		copy(buf[index:], compressedData)
-		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.WriteRaw(buf)
+		return c.WriteRawPackage(buf, true)
 	} else {
 		ddlen := 2 + 1 + cmdLen + dlen
-		buf, start := c.writeDataLen(ddlen)
-		index := start
+		buf := pool.Get(ddlen)
+		index := 0
 		binary.BigEndian.PutUint16(buf[index:index+2], id)
 		index += 2
 		buf[index] = byte(cmdLen)
@@ -436,8 +460,7 @@ func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
 		copy(buf[index:index+cmdLen], cmd)
 		index += cmdLen
 		copy(buf[index:], data)
-		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.WriteRaw(buf)
+		return c.WriteRawPackage(buf, true)
 	}
 }
 
@@ -448,8 +471,8 @@ func (c *Tcp2) WriteResponse(id uint16, state uint8, data []byte) error {
 	if c.compress && dlen > 0 {
 		compressedData, ok, _ := util.CompressData(data)
 		ddlen := 2 + 1 + 1 + len(compressedData)
-		buf, start := c.writeDataLen(ddlen)
-		index := start
+		buf := pool.Get(ddlen)
+		index := 0
 		binary.BigEndian.PutUint16(buf[index:index+2], id)
 		index += 2
 
@@ -467,38 +490,37 @@ func (c *Tcp2) WriteResponse(id uint16, state uint8, data []byte) error {
 		index++
 
 		copy(buf[index:], compressedData)
-		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.WriteRaw(buf)
+		return c.WriteRawPackage(buf, true)
 	} else {
 		ddlen := 2 + 1 + len(data)
-		buf, start := c.writeDataLen(ddlen)
-		index := start
+		buf := pool.Get(ddlen)
+		index := 0
 		binary.BigEndian.PutUint16(buf[index:index+2], id)
 		index += 2
 		buf[index] = state | 0x80
 		index++
 		copy(buf[index:], data)
-		buf[start+ddlen] = util.CRC8(buf[start : start+ddlen])
-		return c.WriteRaw(buf)
+		return c.WriteRawPackage(buf, true)
 	}
 }
 
 // 发送ping包
 func (c *Tcp2) WritePing(id uint16) error {
-	dlen := 2
-	buf, start := c.writeDataLen(dlen)
-	index := start
+	buf := pool.Get(2)
+	index := 0
 	binary.BigEndian.PutUint16(buf[index:index+2], id)
-	buf[start+dlen] = util.CRC8(buf[start : start+dlen])
-	return c.WriteRaw(buf)
+	return c.WriteRawPackage(buf, true)
 }
 
 // 获取信息
 func (c *Tcp2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
-	msg, err := c.ReadRaw(deadline)
+	msg, recycle, err := c.ReadRawPackage(deadline)
 	if err != nil {
 		return
 	}
+	if recycle {
+		defer pool.Put(msg)
+	}
 	msgLen := len(msg)
 	id = binary.BigEndian.Uint16(msg[0:2])
 	// ping信息
@@ -529,9 +551,14 @@ func (c *Tcp2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd s
 	if c.compress && dataLen > 1 {
 		var ok bool
 		data, ok, err = util.UncompressData(data)
+		if !ok && recycle {
+			data = slices.Clone(data)
+		}
 		if ok && c.cf.PrintMsg {
 			log.Println("[UncompressData]", dataLen, "->", len(data))
 		}
+	} else if recycle && dataLen > 0 {
+		data = slices.Clone(data)
 	}
 	return
 }

conn/type.go

@@ -35,13 +35,14 @@ func (t MsgType) String() string {
 
 // 连接接口，代表一个原始的连接
 type Connect interface {
-	// io.ReadWriteCloser
-	WriteAuthInfo(channel string, auth []byte) (err error)
 	ReadAuthInfo() (proto string, version uint8, channel string, auth []byte, err error)
+	ReadMessage(deadline int) (msgType MsgType, id uint16, cmd string, state uint8, data []byte, err error)
+	ReadRawPackage(deadline int) (buf []byte, recycle bool, err error)
+	WriteRawPackage(raw []byte, recycle bool) (err error)
+	WriteAuthInfo(channel string, auth []byte) (err error)
 	WriteRequest(id uint16, cmd string, data []byte) error
 	WriteResponse(id uint16, state uint8, data []byte) error
 	WritePing(id uint16) error
-	ReadMessage(deadline int) (msgType MsgType, id uint16, cmd string, state uint8, data []byte, err error)
 	RemoteIP() net.IP
 	LocalIP() net.IP
 	Close() error

conn/ws2/ws2.go

@@ -15,6 +15,7 @@ import (
 	"git.me9.top/git/tinymq/config"
 	"git.me9.top/git/tinymq/conn"
 	"git.me9.top/git/tinymq/util"
+	"git.me9.top/git/tinymq/util/pool"
 	"github.com/gorilla/websocket"
 )
 
@@ -208,7 +209,10 @@ func Dial(cf *config.Config, scheme string, addr string, path string, hash strin
 
 // 发送数据到网络
 // 如果有加密函数的话会直接修改源数据
-func (c *Ws2) writeMessage(buf []byte) (err error) {
+func (c *Ws2) WriteRawPackage(buf []byte, recycle bool) (err error) {
+	if recycle {
+		defer pool.Put(buf)
+	}
 	if c.cipher != nil {
 		c.cipher.Encrypt(buf, buf)
 	}
@@ -216,6 +220,17 @@ func (c *Ws2) writeMessage(buf []byte) (err error) {
 	return c.conn.WriteMessage(websocket.BinaryMessage, buf)
 }
 
+// 从连接中读取数据包
+// recycle 指示是否需要手动将缓存放回内存池
+func (c *Ws2) ReadRawPackage(deadline int) (msg []byte, recycle bool, err error) {
+	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
+	if err != nil {
+		return
+	}
+	_, msg, err = c.conn.ReadMessage()
+	return
+}
+
 // 发送Auth信息
 // 建立连接后第一个发送的消息
 func (c *Ws2) WriteAuthInfo(channel string, auth []byte) (err error) {
@@ -231,7 +246,8 @@ func (c *Ws2) WriteAuthInfo(channel string, auth []byte) (err error) {
 	dlen := 2 + 1 + protoLen + 1 + 1 + 2 + channelLen + len(auth)
 	index := 0
 
-	buf := make([]byte, dlen)
+	// buf := make([]byte, dlen)
+	buf := pool.Get(dlen)
 	binary.BigEndian.PutUint16(buf[index:index+2], config.ID_AUTH)
 	index += 2
 
@@ -256,7 +272,7 @@ func (c *Ws2) WriteAuthInfo(channel string, auth []byte) (err error) {
 	index += channelLen
 
 	copy(buf[index:], auth)
-	return c.writeMessage(buf)
+	return c.WriteRawPackage(buf, true)
 }
 
 // 获取Auth信息
@@ -268,11 +284,13 @@ func (c *Ws2) ReadAuthInfo() (proto string, version uint8, channel string, auth
 			return
 		}
 	}()
-	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(c.cf.ReadWait)))
-	if err != nil {
-		return
-	}
-	_, msg, err := c.conn.ReadMessage()
+	// 由于这里返回的数据都不需要回收，故不做判断处理
+	msg, _, err := c.ReadRawPackage(c.cf.ReadWait)
+	// err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(c.cf.ReadWait)))
+	// if err != nil {
+	// 	return
+	// }
+	// _, msg, err := c.conn.ReadMessage()
 	if err != nil {
 		return
 	}
@@ -340,7 +358,8 @@ func (c *Ws2) WriteRequest(id uint16, cmd string, data []byte) error {
 	if c.compress && dlen > 0 {
 		compressedData, ok, _ := util.CompressData(data)
 		ddlen := 2 + 1 + cmdLen + 1 + len(compressedData)
-		buf := make([]byte, ddlen) // 申请内存
+		// buf := make([]byte, ddlen) // 申请内存
+		buf := pool.Get(ddlen)
 		index := 0
 		binary.BigEndian.PutUint16(buf[index:index+2], id)
 		index += 2
@@ -361,15 +380,15 @@ func (c *Ws2) WriteRequest(id uint16, cmd string, data []byte) error {
 		index++
 
 		copy(buf[index:], compressedData)
-		return c.writeMessage(buf)
+		return c.WriteRawPackage(buf, true)
 	} else {
 		ddlen := 2 + 1 + cmdLen + dlen
-		buf := make([]byte, ddlen) // 申请内存
+		buf := pool.Get(ddlen)
 		binary.BigEndian.PutUint16(buf[0:2], id)
 		buf[2] = byte(cmdLen)
 		copy(buf[3:], cmd)
 		copy(buf[3+cmdLen:], data)
-		return c.writeMessage(buf)
+		return c.WriteRawPackage(buf, true)
 	}
 }
 
@@ -380,7 +399,7 @@ func (c *Ws2) WriteResponse(id uint16, state uint8, data []byte) error {
 	if c.compress && dlen > 0 {
 		compressedData, ok, _ := util.CompressData(data)
 		ddlen := 2 + 1 + 1 + len(compressedData)
-		buf := make([]byte, ddlen)
+		buf := pool.Get(ddlen)
 		binary.BigEndian.PutUint16(buf[0:2], id)
 		buf[2] = state | 0x80
 		if ok {
@@ -392,31 +411,32 @@ func (c *Ws2) WriteResponse(id uint16, state uint8, data []byte) error {
 			buf[3] = 0
 		}
 		copy(buf[4:], compressedData)
-		return c.writeMessage(buf)
+		return c.WriteRawPackage(buf, true)
 	} else {
 		ddlen := 2 + 1 + dlen
-		buf := make([]byte, ddlen)
+		buf := pool.Get(ddlen)
 		binary.BigEndian.PutUint16(buf[0:2], id)
 		buf[2] = state | 0x80
 		copy(buf[3:], data)
-		return c.writeMessage(buf)
+		return c.WriteRawPackage(buf, true)
 	}
 }
 
 // 发送ping包
 func (c *Ws2) WritePing(id uint16) error {
-	buf := make([]byte, 2)
+	buf := pool.Get(2)
 	binary.BigEndian.PutUint16(buf[0:2], id)
-	return c.writeMessage(buf)
+	return c.WriteRawPackage(buf, true)
 }
 
 // 获取信息
 func (c *Ws2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
-	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
-	if err != nil {
-		return
-	}
-	_, msg, err := c.conn.ReadMessage()
+	msg, _, err := c.ReadRawPackage(deadline)
+	// err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
+	// if err != nil {
+	// 	return
+	// }
+	// _, msg, err := c.conn.ReadMessage()
 	if err != nil {
 		return
 	}

util/buf/pool.go

@@ -1,12 +0,0 @@
-package buf
-
-func Get(size int) []byte {
-	if size == 0 {
-		return nil
-	}
-	return DefaultAllocator.Get(size)
-}
-
-func Put(buf []byte) error {
-	return DefaultAllocator.Put(buf)
-}

util/encrypt.go

@@ -13,8 +13,18 @@ const (
 
 // 定义一些加密方法
 var CipherMethod = map[string]*CipherInfo{
-	"plain": {0, 0, NewPlainStream, NewPlainStream},
-	"xor":   {XOR_KEY_LEN, XOR_IV_LEN, NewXorStream, NewXorStream},
+	"plain": {
+		KeyLen:           0,
+		IvLen:            0,
+		NewEncryptStream: NewPlainStream,
+		NewDecryptStream: NewPlainStream,
+	},
+	"xor": {
+		KeyLen:           XOR_KEY_LEN,
+		IvLen:            XOR_IV_LEN,
+		NewEncryptStream: NewXorStream,
+		NewDecryptStream: NewXorStream,
+	},
 }
 
 // 普通的xor加密，没有破解的难度，只是简单的让明文不可见而已

util/buf/alloc.go → util/pool/alloc.go

@@ -1,4 +1,4 @@
-package buf
+package pool
 
 // Inspired by https://github.com/xtaci/smux/blob/master/alloc.go
 

util/pool/pool.go

@@ -0,0 +1,20 @@
+package pool
+
+func Get(size int) []byte {
+	if size == 0 {
+		return nil
+	}
+	// 大块内存直接分配好了
+	if size > 65536 {
+		return make([]byte, size)
+	}
+	return DefaultAllocator.Get(size)
+}
+
+func Put(buf []byte) error {
+	// 大内存不重新利用，直接由系统回收
+	if cap(buf) > 65536 {
+		return nil
+	}
+	return DefaultAllocator.Put(buf)
+}