test cases for v2io

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)

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 {

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)
+}

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 {
+	defer buffer.Release()
+	for !buffer.IsEmpty() {
 		nBytes, err := this.writer.Write(buffer.Value)
-	if nBytes < buffer.Len() {
-		_, err = this.writer.Write(buffer.Value[nBytes:])
-	}
-	buffer.Release()
+		if err != nil {
 			return err
+		}
+		buffer.SliceFrom(nBytes)
+	}
+	return nil
 }
 
 func (this *AdaptiveWriter) Release() {

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)
+}