add new file

lib/pool/goroutine.go

@@ -0,0 +1,76 @@
+package pool
+
+import (
+	"ehang.io/nps/lib/common"
+	"github.com/panjf2000/ants/v2"
+	"io"
+	"net"
+	"sync"
+)
+
+var connBp = NewBufferPool(MaxReadSize)
+var packetBp = NewBufferPool(1500)
+
+const MaxReadSize = 32 * 1024
+
+var CopyConnGoroutinePool *ants.PoolWithFunc
+var CopyPacketGoroutinePool *ants.PoolWithFunc
+
+type CopyConnGpParams struct {
+	Writer io.Writer
+	Reader io.Reader
+	Wg     *sync.WaitGroup
+}
+
+type CopyPacketGpParams struct {
+	RPacket net.PacketConn
+	WPacket net.PacketConn
+	Wg      *sync.WaitGroup
+}
+
+func init() {
+	var err error
+	CopyConnGoroutinePool, err = ants.NewPoolWithFunc(1000000, func(i interface{}) {
+		gpp, ok := i.(*CopyConnGpParams)
+		if !ok {
+			return
+		}
+		buf := connBp.Get()
+		_, _ = common.CopyBuffer(gpp.Writer, gpp.Reader, buf)
+		connBp.Put(buf)
+		gpp.Wg.Done()
+		if v, ok := gpp.Reader.(*net.TCPConn); ok {
+			_ = v.CloseWrite()
+		}
+		if v, ok := gpp.Writer.(*net.TCPConn); ok {
+			_ = v.CloseRead()
+		}
+	})
+	if err != nil {
+		panic(err)
+	}
+	CopyPacketGoroutinePool, err = ants.NewPoolWithFunc(1000000, func(i interface{}) {
+		cpp, ok := i.(*CopyPacketGpParams)
+		if !ok {
+			return
+		}
+		buf := connBp.Get()
+		for {
+			n, addr, err := cpp.RPacket.ReadFrom(buf)
+			if err != nil {
+				break
+			}
+			_, err = cpp.WPacket.WriteTo(buf[:n], addr)
+			if err != nil {
+				break
+			}
+		}
+		connBp.Put(buf)
+		_ = cpp.RPacket.Close()
+		_ = cpp.WPacket.Close()
+		cpp.Wg.Done()
+	})
+	if err != nil {
+		panic(err)
+	}
+}

lib/pool/pool.go

@@ -0,0 +1,32 @@
+package pool
+
+import "sync"
+
+type BufferPool struct {
+	pool     sync.Pool
+	poolSize int
+}
+
+func NewBufferPool(poolSize int) *BufferPool {
+	bp := &BufferPool{}
+	bp.pool = sync.Pool{
+		New: func() interface{} {
+			return make([]byte, poolSize, poolSize)
+		},
+	}
+	bp.poolSize = poolSize
+	return bp
+}
+
+func (bp *BufferPool) Get() []byte {
+	buf := bp.pool.Get().([]byte)
+	return buf[:bp.poolSize] // just like make a new slice, but data may not be 0
+}
+
+func (bp *BufferPool) Put(x []byte) {
+	if len(x) == bp.poolSize {
+		bp.pool.Put(x)
+	} else {
+		x = nil // buf is not full, not allowed, New method returns a full buf
+	}
+}