Improve: pool buffer alloc

tunnel/connection.go

@@ -61,9 +61,9 @@ func handleHTTP(request *adapters.HTTPAdapter, outbound net.Conn) {
 		}
 
 		// even if resp.Write write body to the connection, but some http request have to Copy to close it
-		buf := pool.BufPool.Get().([]byte)
+		buf := pool.Get(pool.RelayBufferSize)
 		_, err = io.CopyBuffer(request, resp.Body, buf)
-		pool.BufPool.Put(buf[:cap(buf)])
+		pool.Put(buf)
 		if err != nil && err != io.EOF {
 			break
 		}
@@ -90,8 +90,8 @@ func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata
 }
 
 func handleUDPToLocal(packet C.UDPPacket, pc net.PacketConn, key string, fAddr net.Addr) {
-	buf := pool.BufPool.Get().([]byte)
-	defer pool.BufPool.Put(buf[:cap(buf)])
+	buf := pool.Get(pool.RelayBufferSize)
+	defer pool.Put(buf)
 	defer natTable.Delete(key)
 	defer pc.Close()
 
@@ -122,16 +122,16 @@ func relay(leftConn, rightConn net.Conn) {
 	ch := make(chan error)
 
 	go func() {
-		buf := pool.BufPool.Get().([]byte)
+		buf := pool.Get(pool.RelayBufferSize)
 		_, err := io.CopyBuffer(leftConn, rightConn, buf)
-		pool.BufPool.Put(buf[:cap(buf)])
+		pool.Put(buf)
 		leftConn.SetReadDeadline(time.Now())
 		ch <- err
 	}()
 
-	buf := pool.BufPool.Get().([]byte)
+	buf := pool.Get(pool.RelayBufferSize)
 	io.CopyBuffer(rightConn, leftConn, buf)
-	pool.BufPool.Put(buf[:cap(buf)])
+	pool.Put(buf)
 	rightConn.SetReadDeadline(time.Now())
 	<-ch
 }