Maintain an internal buffer pool to accelerate allocation

common/alloc/buffer.go

@@ -0,0 +1,104 @@
+package alloc
+
+import (
+	//"fmt"
+	"time"
+)
+
+type Buffer struct {
+	head  []byte
+	pool  *bufferPool
+	Value []byte
+}
+
+func (b *Buffer) Release() {
+	b.pool.free(b)
+}
+
+func (b *Buffer) Clear() {
+	b.Value = b.Value[:0]
+}
+
+func (b *Buffer) Append(data []byte) {
+	b.Value = append(b.Value, data...)
+}
+
+func (b *Buffer) Slice(from, to int) {
+	b.Value = b.Value[from:to]
+}
+
+func (b *Buffer) SliceFrom(from int) {
+	b.Value = b.Value[from:]
+}
+
+func (b *Buffer) Len() int {
+	return len(b.Value)
+}
+
+type bufferPool struct {
+	chain       chan *Buffer
+	allocator   func(*bufferPool) *Buffer
+	minElements int
+	maxElements int
+}
+
+func newBufferPool(allocator func(*bufferPool) *Buffer, minElements, maxElements int) *bufferPool {
+	pool := &bufferPool{
+		chain:       make(chan *Buffer, maxElements*2),
+		allocator:   allocateSmall,
+		minElements: minElements,
+		maxElements: maxElements,
+	}
+	for i := 0; i < minElements; i++ {
+		pool.chain <- allocator(pool)
+	}
+	go pool.cleanup(time.Tick(1 * time.Second))
+	return pool
+}
+
+func (p *bufferPool) allocate() *Buffer {
+	//fmt.Printf("Pool size: %d\n", len(p.chain))
+	var b *Buffer
+	select {
+	case b = <-p.chain:
+	default:
+		b = p.allocator(p)
+	}
+	b.Value = b.head
+	return b
+}
+
+func (p *bufferPool) free(buffer *Buffer) {
+	select {
+	case p.chain <- buffer:
+	default:
+	}
+	//fmt.Printf("Pool size: %d\n", len(p.chain))
+}
+
+func (p *bufferPool) cleanup(tick <-chan time.Time) {
+	for range tick {
+		pSize := len(p.chain)
+		for delta := pSize - p.minElements; delta > 0; delta-- {
+			p.chain <- p.allocator(p)
+		}
+		for delta := p.maxElements - pSize; delta > 0; delta-- {
+			<-p.chain
+		}
+	}
+}
+
+func allocateSmall(pool *bufferPool) *Buffer {
+	b := &Buffer{
+		head: make([]byte, 8*1024),
+	}
+	b.Value = b.head
+	b.pool = pool
+	return b
+}
+
+var smallPool = newBufferPool(allocateSmall, 256, 1024)
+
+func NewBuffer() *Buffer {
+	return smallPool.allocate()
+}

common/net/packet.go

@@ -1,12 +1,16 @@
 package net
 
+import (
+	"github.com/v2ray/v2ray-core/common/alloc"
+)
+
 type Packet interface {
 	Destination() Destination
-	Chunk() []byte // First chunk of this commnunication
+	Chunk() *alloc.Buffer // First chunk of this commnunication
 	MoreChunks() bool
 }
 
-func NewPacket(dest Destination, firstChunk []byte, moreChunks bool) Packet {
+func NewPacket(dest Destination, firstChunk *alloc.Buffer, moreChunks bool) Packet {
 	return &packetImpl{
 		dest:     dest,
 		data:     firstChunk,
@@ -16,7 +20,7 @@ func NewPacket(dest Destination, firstChunk []byte, moreChunks bool) Packet {
 
 type packetImpl struct {
 	dest     Destination
-	data     []byte
+	data     *alloc.Buffer
 	moreData bool
 }
 
@@ -24,7 +28,7 @@ func (packet *packetImpl) Destination() Destination {
 	return packet.dest
 }
 
-func (packet *packetImpl) Chunk() []byte {
+func (packet *packetImpl) Chunk() *alloc.Buffer {
 	return packet.data
 }
 

common/net/transport.go

@@ -2,63 +2,41 @@ package net
 
 import (
 	"io"
+
+	"github.com/v2ray/v2ray-core/common/alloc"
 )
 
-const (
-	minBufferSizeKilo = 2
-	maxBufferSizeKilo = 128
-)
-
-func ReadFrom(reader io.Reader, sizeInKilo int) ([]byte, error) {
-	buffer := make([]byte, sizeInKilo<<10)
-	nBytes, err := reader.Read(buffer)
-	if nBytes == 0 {
-		return nil, err
+func ReadFrom(reader io.Reader, buffer *alloc.Buffer) (*alloc.Buffer, error) {
+	if buffer == nil {
+		buffer = alloc.NewBuffer()
 	}
-	return buffer[:nBytes], err
-}
-
-func roundUp(size int) int {
-	if size <= minBufferSizeKilo {
-		return minBufferSizeKilo
-	}
-	if size >= maxBufferSizeKilo {
-		return maxBufferSizeKilo
-	}
-	size--
-	size |= size >> 1
-	size |= size >> 2
-	size |= size >> 4
-	return size + 1
+	nBytes, err := reader.Read(buffer.Value)
+	buffer.Slice(0, nBytes)
+	return buffer, err
 }
 
 // ReaderToChan dumps all content from a given reader to a chan by constantly reading it until EOF.
-func ReaderToChan(stream chan<- []byte, reader io.Reader) error {
-	bufferSizeKilo := 2
+func ReaderToChan(stream chan<- *alloc.Buffer, reader io.Reader) error {
 	for {
-		data, err := ReadFrom(reader, bufferSizeKilo)
-		if len(data) > 0 {
-			stream <- data
+		buffer, err := ReadFrom(reader, nil)
+		if buffer.Len() > 0 {
+			stream <- buffer
+		} else {
+			buffer.Release()
+			buffer = nil
 		}
 		if err != nil {
 			return err
 		}
-		if bufferSizeKilo == maxBufferSizeKilo {
-			continue
-		}
-		dataLenKilo := len(data) >> 10
-		if dataLenKilo == bufferSizeKilo {
-			bufferSizeKilo <<= 1
-		} else {
-			bufferSizeKilo = roundUp(dataLenKilo)
-		}
 	}
 }
 
 // ChanToWriter dumps all content from a given chan to a writer until the chan is closed.
-func ChanToWriter(writer io.Writer, stream <-chan []byte) error {
+func ChanToWriter(writer io.Writer, stream <-chan *alloc.Buffer) error {
 	for buffer := range stream {
-		_, err := writer.Write(buffer)
+		_, err := writer.Write(buffer.Value)
+		buffer.Release()
+		buffer = nil
 		if err != nil {
 			return err
 		}

common/net/transport_test.go

@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"testing"
 
+	"github.com/v2ray/v2ray-core/common/alloc"
 	"github.com/v2ray/v2ray-core/testing/unit"
 )
 
@@ -22,7 +23,7 @@ func TestReaderAndWrite(t *testing.T) {
 	readerBuffer := bytes.NewReader(buffer)
 	writerBuffer := bytes.NewBuffer(make([]byte, 0, size))
 
-	transportChan := make(chan []byte, 1024)
+	transportChan := make(chan *alloc.Buffer, 1024)
 
 	err = ReaderToChan(transportChan, readerBuffer)
 	assert.Error(err).Equals(io.EOF)
@@ -44,7 +45,7 @@ func (reader *StaticReader) Read(b []byte) (size int, err error) {
 	if size > reader.total-reader.current {
 		size = reader.total - reader.current
 	}
-	for i := 0; i < len(b); i++ {
+	for i := 0; i < size; i++ {
 		b[i] = byte(i)
 	}
 	//rand.Read(b[:size])
@@ -113,8 +114,8 @@ func BenchmarkTransport10M(b *testing.B) {
 
 func runBenchmarkTransport(size int) {
 
-	transportChanA := make(chan []byte, 16)
-	transportChanB := make(chan []byte, 16)
+	transportChanA := make(chan *alloc.Buffer, 16)
+	transportChanB := make(chan *alloc.Buffer, 16)
 
 	readerA := &StaticReader{size, 0}
 	readerB := &StaticReader{size, 0}