diff --git a/common/alloc/buffer_pool.go b/common/alloc/buffer_pool.go index 019992f86..36e72e3a1 100644 --- a/common/alloc/buffer_pool.go +++ b/common/alloc/buffer_pool.go @@ -49,6 +49,11 @@ func (p *BufferPool) Free(buffer *Buffer) { } } +const ( + BufferSize = 8*1024 - defaultOffset + LargeBufferSize = 64*1024 - defaultOffset +) + var smallPool = NewBufferPool(1024, 64) var mediumPool = NewBufferPool(8*1024, 128) var largePool = NewBufferPool(64*1024, 64) diff --git a/common/io/reader.go b/common/io/reader.go index 8a2abf032..ef4ae7be5 100644 --- a/common/io/reader.go +++ b/common/io/reader.go @@ -29,7 +29,6 @@ type Reader interface { type AdaptiveReader struct { reader io.Reader allocate func() *alloc.Buffer - isLarge bool } // NewAdaptiveReader creates a new AdaptiveReader. @@ -38,7 +37,6 @@ func NewAdaptiveReader(reader io.Reader) *AdaptiveReader { return &AdaptiveReader{ reader: reader, allocate: alloc.NewBuffer, - isLarge: false, } } @@ -46,12 +44,10 @@ func NewAdaptiveReader(reader io.Reader) *AdaptiveReader { func (this *AdaptiveReader) Read() (*alloc.Buffer, error) { buffer, err := ReadFrom(this.reader, this.allocate()) - if buffer.IsFull() && !this.isLarge { + if buffer.Len() >= alloc.BufferSize { this.allocate = alloc.NewLargeBuffer - this.isLarge = true - } else if !buffer.IsFull() { + } else { this.allocate = alloc.NewBuffer - this.isLarge = false } if err != nil { diff --git a/common/io/reader_test.go b/common/io/reader_test.go new file mode 100644 index 000000000..0ff53475a --- /dev/null +++ b/common/io/reader_test.go @@ -0,0 +1,28 @@ +package io_test + +import ( + "bytes" + "testing" + + "github.com/v2ray/v2ray-core/common/alloc" + . "github.com/v2ray/v2ray-core/common/io" + v2testing "github.com/v2ray/v2ray-core/testing" + "github.com/v2ray/v2ray-core/testing/assert" +) + +func TestAdaptiveReader(t *testing.T) { + v2testing.Current(t) + + rawContent := make([]byte, 1024*1024) + + reader := NewAdaptiveReader(bytes.NewBuffer(rawContent)) + b1, err := reader.Read() + assert.Error(err).IsNil() + assert.Bool(b1.IsFull()).IsTrue() + assert.Int(b1.Len()).Equals(alloc.BufferSize) + + b2, err := reader.Read() + assert.Error(err).IsNil() + assert.Bool(b2.IsFull()).IsTrue() + assert.Int(b2.Len()).Equals(alloc.LargeBufferSize) +} diff --git a/common/io/writer.go b/common/io/writer.go index 2df487ce8..85e651e43 100644 --- a/common/io/writer.go +++ b/common/io/writer.go @@ -28,12 +28,15 @@ func NewAdaptiveWriter(writer io.Writer) *AdaptiveWriter { // Write implements Writer.Write(). Write() takes ownership of the given buffer. func (this *AdaptiveWriter) Write(buffer *alloc.Buffer) error { - nBytes, err := this.writer.Write(buffer.Value) - if nBytes < buffer.Len() { - _, err = this.writer.Write(buffer.Value[nBytes:]) + defer buffer.Release() + for !buffer.IsEmpty() { + nBytes, err := this.writer.Write(buffer.Value) + if err != nil { + return err + } + buffer.SliceFrom(nBytes) } - buffer.Release() - return err + return nil } func (this *AdaptiveWriter) Release() { diff --git a/common/io/writer_test.go b/common/io/writer_test.go new file mode 100644 index 000000000..1bf8be4c7 --- /dev/null +++ b/common/io/writer_test.go @@ -0,0 +1,26 @@ +package io_test + +import ( + "bytes" + "crypto/rand" + "testing" + + "github.com/v2ray/v2ray-core/common/alloc" + . "github.com/v2ray/v2ray-core/common/io" + v2testing "github.com/v2ray/v2ray-core/testing" + "github.com/v2ray/v2ray-core/testing/assert" +) + +func TestAdaptiveWriter(t *testing.T) { + v2testing.Current(t) + + lb := alloc.NewLargeBuffer() + rand.Read(lb.Value) + + writeBuffer := make([]byte, 0, 1024*1024) + + writer := NewAdaptiveWriter(NewBufferedWriter(bytes.NewBuffer(writeBuffer))) + err := writer.Write(lb) + assert.Error(err).IsNil() + assert.Bytes(lb.Bytes()).Equals(writeBuffer) +}