From a766c61dcced814de06c337aea822749bd8e4f5c Mon Sep 17 00:00:00 2001
From: V2Ray <admin@v2ray.com>
Date: Wed, 16 Sep 2015 23:56:27 +0200
Subject: [PATCH] reuse buffers in transport, benchmark shows it improves
 performance by 10 times

---
 net/transport.go      | 26 ++++++++++--
 net/transport_test.go | 94 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+), 3 deletions(-)
 create mode 100644 net/transport_test.go

diff --git a/net/transport.go b/net/transport.go
index 63d5919d..69fd56a0 100644
--- a/net/transport.go
+++ b/net/transport.go
@@ -2,17 +2,36 @@ package net
 
 import (
 	"io"
-
-	_ "github.com/v2ray/v2ray-core/log"
 )
 
 const (
 	bufferSize = 32 * 1024
 )
 
+var (
+	dirtyBuffers = make(chan []byte, 1024)
+)
+
+func getBuffer() []byte {
+	var buffer []byte
+	select {
+	case buffer = <-dirtyBuffers:
+	default:
+		buffer = make([]byte, bufferSize)
+	}
+	return buffer
+}
+
+func putBuffer(buffer []byte) {
+	select {
+	case dirtyBuffers <- buffer:
+	default:
+	}
+}
+
 func ReaderToChan(stream chan<- []byte, reader io.Reader) error {
 	for {
-		buffer := make([]byte, bufferSize)
+		buffer := getBuffer()
 		nBytes, err := reader.Read(buffer)
 		if nBytes > 0 {
 			stream <- buffer[:nBytes]
@@ -27,6 +46,7 @@ func ReaderToChan(stream chan<- []byte, reader io.Reader) error {
 func ChanToWriter(writer io.Writer, stream <-chan []byte) error {
 	for buffer := range stream {
 		_, err := writer.Write(buffer)
+		putBuffer(buffer)
 		if err != nil {
 			return err
 		}
diff --git a/net/transport_test.go b/net/transport_test.go
new file mode 100644
index 00000000..767f87d1
--- /dev/null
+++ b/net/transport_test.go
@@ -0,0 +1,94 @@
+package net
+
+import (
+	"bytes"
+	"crypto/rand"
+	"io"
+	"io/ioutil"
+	"testing"
+
+	"github.com/v2ray/v2ray-core/testing/unit"
+)
+
+func TestReaderAndWrite(t *testing.T) {
+	assert := unit.Assert(t)
+
+	size := 1024 * 1024
+	buffer := make([]byte, size)
+	nBytes, err := rand.Read(buffer)
+	assert.Int(nBytes).Equals(len(buffer))
+	assert.Error(err).IsNil()
+
+	readerBuffer := bytes.NewReader(buffer)
+	writerBuffer := bytes.NewBuffer(make([]byte, 0, size))
+
+	transportChan := make(chan []byte, size/bufferSize*10)
+
+	err = ReaderToChan(transportChan, readerBuffer)
+	assert.Error(err).Equals(io.EOF)
+	close(transportChan)
+
+	err = ChanToWriter(writerBuffer, transportChan)
+	assert.Error(err).IsNil()
+
+	assert.Bytes(buffer).Equals(writerBuffer.Bytes())
+}
+
+type StaticReader struct {
+	total   int
+	current int
+}
+
+func (reader *StaticReader) Read(b []byte) (size int, err error) {
+	size = len(b)
+	if size > reader.total-reader.current {
+		size = reader.total - reader.current
+	}
+	//rand.Read(b[:size])
+	reader.current += size
+	if reader.current == reader.total {
+		err = io.EOF
+	}
+	return
+}
+
+func BenchmarkTransport(b *testing.B) {
+	size := 1024 * 1024
+
+	for i := 0; i < b.N; i++ {
+		transportChanA := make(chan []byte, 128)
+		transportChanB := make(chan []byte, 128)
+
+		readerA := &StaticReader{size, 0}
+		readerB := &StaticReader{size, 0}
+
+		writerA := ioutil.Discard
+		writerB := ioutil.Discard
+
+		finishA := make(chan bool)
+		finishB := make(chan bool)
+
+		go func() {
+			ChanToWriter(writerA, transportChanA)
+			close(finishA)
+		}()
+
+		go func() {
+			ReaderToChan(transportChanA, readerA)
+			close(transportChanA)
+		}()
+
+		go func() {
+			ChanToWriter(writerB, transportChanB)
+			close(finishB)
+		}()
+
+		go func() {
+			ReaderToChan(transportChanB, readerB)
+			close(transportChanB)
+		}()
+
+		<-transportChanA
+		<-transportChanB
+	}
+}