Change pool, change mux connection close ,change copy buffer

lib/common/netpackager.go

@@ -4,7 +4,8 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/json"
-	"github.com/cnlh/nps/lib/pool"
+	"errors"
+	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
 	"io"
 	"strings"
 )
@@ -20,7 +21,6 @@ type BasePackager struct {
 }
 
 func (Self *BasePackager) NewPac(contents ...interface{}) (err error) {
-	Self.Content = pool.CopyBuff.Get()
 	Self.clean()
 	for _, content := range contents {
 		switch content.(type) {
@@ -50,7 +50,8 @@ func (Self *BasePackager) appendByte(data []byte) (err error) {
 		copy(Self.Content[m:n], data)
 		return nil
 	} else {
-		return bytes.ErrTooLarge
+		logs.Warn(len(data), len(Self.Content), cap(Self.Content))
+		return errors.New("pack content too large")
 	}
 }
 
@@ -61,20 +62,22 @@ func (Self *BasePackager) Pack(writer io.Writer) (err error) {
 		return
 	}
 	err = binary.Write(writer, binary.LittleEndian, Self.Content)
-	pool.CopyBuff.Put(Self.Content)
 	return
 }
 
 //Unpack 会导致传入的数字类型转化成float64！！
 //主要原因是json unmarshal并未传入正确的数据类型
 func (Self *BasePackager) UnPack(reader io.Reader) (err error) {
-	Self.Content = pool.CopyBuff.Get()
 	Self.clean()
 	err = binary.Read(reader, binary.LittleEndian, &Self.Length)
 	if err != nil {
 		return
 	}
-	Self.Content = Self.Content[:Self.Length]
+	if int(Self.Length) > cap(Self.Content) {
+		logs.Warn("unpack", cap(Self.Content))
+		err = errors.New("unpack err, content length too large")
+	}
+	Self.Content = Self.Content[:int(Self.Length)]
 	//n, err := io.ReadFull(reader, Self.Content)
 	//if n != int(Self.Length) {
 	//	err = io.ErrUnexpectedEOF
@@ -177,6 +180,7 @@ func (Self *MuxPackager) Pack(writer io.Writer) (err error) {
 }
 
 func (Self *MuxPackager) UnPack(reader io.Reader) (err error) {
+	Self.BasePackager.clean() // also clean the content
 	err = binary.Read(reader, binary.LittleEndian, &Self.Flag)
 	if err != nil {
 		return

lib/common/pool.go

@@ -0,0 +1,146 @@
+package common
+
+import (
+	"bytes"
+	"sync"
+)
+
+const PoolSize = 64 * 1024
+const PoolSizeSmall = 100
+const PoolSizeUdp = 1472
+const PoolSizeCopy = 32 << 10
+
+var BufPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, PoolSize)
+	},
+}
+
+var BufPoolUdp = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, PoolSizeUdp)
+	},
+}
+var BufPoolMax = 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)
+	},
+}
+
+func PutBufPoolUdp(buf []byte) {
+	if cap(buf) == PoolSizeUdp {
+		BufPoolUdp.Put(buf[:PoolSizeUdp])
+	}
+}
+
+func PutBufPoolCopy(buf []byte) {
+	if cap(buf) == PoolSizeCopy {
+		BufPoolCopy.Put(buf[:PoolSizeCopy])
+	}
+}
+
+func GetBufPoolCopy() []byte {
+	return (BufPoolCopy.Get().([]byte))[:PoolSizeCopy]
+}
+
+func PutBufPoolMax(buf []byte) {
+	if cap(buf) == PoolSize {
+		BufPoolMax.Put(buf[:PoolSize])
+	}
+}
+
+type CopyBufferPool struct {
+	pool sync.Pool
+}
+
+func (Self *CopyBufferPool) New() {
+	Self.pool = sync.Pool{
+		New: func() interface{} {
+			return make([]byte, PoolSizeCopy, PoolSizeCopy)
+		},
+	}
+}
+
+func (Self *CopyBufferPool) Get() []byte {
+	buf := Self.pool.Get().([]byte)
+	return buf[:PoolSizeCopy] // just like make a new slice, but data may not be 0
+}
+
+func (Self *CopyBufferPool) Put(x []byte) {
+	if len(x) == PoolSizeCopy {
+		Self.pool.Put(x)
+	} else {
+		x = nil // buf is not full, maybe truncated by gc in pool, not allowed
+	}
+}
+
+type BufferPool struct {
+	pool sync.Pool
+}
+
+func (Self *BufferPool) New() {
+	Self.pool = sync.Pool{
+		New: func() interface{} {
+			return new(bytes.Buffer)
+		},
+	}
+}
+
+func (Self *BufferPool) Get() *bytes.Buffer {
+	return Self.pool.Get().(*bytes.Buffer)
+}
+
+func (Self *BufferPool) Put(x *bytes.Buffer) {
+	x.Reset()
+	Self.pool.Put(x)
+}
+
+type MuxPackagerPool struct {
+	pool sync.Pool
+}
+
+func (Self *MuxPackagerPool) New() {
+	Self.pool = sync.Pool{
+		New: func() interface{} {
+			pack := MuxPackager{}
+			return &pack
+		},
+	}
+}
+
+func (Self *MuxPackagerPool) Get() *MuxPackager {
+	pack := Self.pool.Get().(*MuxPackager)
+	buf := CopyBuff.Get()
+	pack.Content = buf
+	return pack
+}
+
+func (Self *MuxPackagerPool) Put(pack *MuxPackager) {
+	CopyBuff.Put(pack.Content)
+	Self.pool.Put(pack)
+}
+
+var once = sync.Once{}
+var BuffPool = BufferPool{}
+var CopyBuff = CopyBufferPool{}
+var MuxPack = MuxPackagerPool{}
+
+func newPool() {
+	BuffPool.New()
+	CopyBuff.New()
+	MuxPack.New()
+}
+
+func init() {
+	once.Do(newPool)
+}

lib/common/util.go

@@ -4,8 +4,8 @@ import (
 	"bytes"
 	"encoding/base64"
 	"encoding/binary"
+	"errors"
 	"github.com/cnlh/nps/lib/crypt"
-	"github.com/cnlh/nps/lib/pool"
 	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
 	"html/template"
 	"io"
@@ -264,34 +264,96 @@ func GetPortByAddr(addr string) int {
 	return p
 }
 
-func CopyBuffer(dst io.Writer, src io.Reader, connId int32) (written int64, err error) {
-	buf := pool.CopyBuff.Get()
-	defer pool.CopyBuff.Put(buf)
-	for {
-		nr, er := src.Read(buf)
-		if er != nil {
-			if er != io.EOF {
-				err = er
-			}
-			break
+type ConnCopy struct {
+	dst    net.Conn
+	src    net.Conn
+	buf    []byte
+	connId int32
+}
+
+func (Self *ConnCopy) New(dst net.Conn, src net.Conn, connId int32) {
+	Self.dst = dst
+	Self.src = src
+	Self.buf = CopyBuff.Get()
+	Self.connId = connId
+}
+
+func (Self *ConnCopy) copyBufferOnce() (written int64, err error) {
+	nr, er := Self.src.Read(Self.buf)
+	if nr > 0 {
+		//logs.Warn("write", Self.connId, nr, string(buf[0:10]))
+		nw, ew := Self.dst.Write(Self.buf[0:nr])
+		if nw > 0 {
+			written = int64(nw)
 		}
-		if nr > 0 {
-			logs.Warn("write", connId, nr, string(buf[0:10]))
-			nw, ew := dst.Write(buf[0:nr])
-			if nw > 0 {
-				written += int64(nw)
-			}
-			if ew != nil {
-				err = ew
-				break
-			}
-			if nr != nw {
-				err = io.ErrShortWrite
-				break
-			}
+		if ew != nil {
+			//logs.Warn("write err ew id nw", ew, Self.connId, nw)
+			err = ew
+			return
+		}
+		if nr != nw {
+			err = io.ErrShortWrite
+			return
+		}
+		if nw == 0 {
+			err = errors.New("io: write on closed pipe")
+			//logs.Warn("write buffer", err)
+			return
 		}
 	}
-	return written, err
+	if nr == 0 && er == nil {
+		err = errors.New("io: read on closed pipe")
+		//logs.Warn("read buffer", err)
+		return
+	}
+	if er != nil {
+		err = er
+		return
+	}
+	return
+}
+
+func (Self *ConnCopy) copyBuffer() (written int64, err error) {
+	var write int64
+	write, err = Self.copyBufferOnce() // first copy, if written is zero and err is io.EOF
+	// means conn already closed, so need to close all the conn
+	written += write
+	if err == io.EOF && written == 0 {
+		err = errors.New("io: read on closed pipe")
+		return
+	} else if err == io.EOF && written > 0 {
+		err = nil
+		return
+	}
+	for {
+		write, err = Self.copyBufferOnce()
+		written += write
+		if err != nil {
+			if err == io.EOF {
+				err = nil
+			}
+			return
+		}
+	}
+}
+
+func (Self *ConnCopy) CopyConn() (written int64, err error) {
+	defer CopyBuff.Put(Self.buf)
+	if Self.dst != nil && Self.src != nil {
+		written, err = Self.copyBuffer()
+	} else {
+		return 0, errors.New("copy conn nil src or dst")
+	}
+	if err != nil { // copyBuffer do not return io.EOF ,close all conn
+		logs.Warn("close by copy conn ", Self.connId, err)
+		if Self.dst != nil {
+			Self.dst.Close()
+		}
+		if Self.src != nil {
+			Self.src.Close()
+		}
+	}
+	return
 }
 
 //send this ip forget to get a local udp port