refactor io package

common/buf/io.go

@@ -1,6 +1,24 @@
 package buf
 
-import "io"
+import (
+	"io"
+
+	"v2ray.com/core/common/errors"
+)
+
+// Reader extends io.Reader with alloc.Buffer.
+type Reader interface {
+	Release()
+	// Read reads content from underlying reader, and put it into an alloc.Buffer.
+	Read() (*Buffer, error)
+}
+
+// Writer extends io.Writer with alloc.Buffer.
+type Writer interface {
+	Release()
+	// Write writes an alloc.Buffer into underlying writer.
+	Write(*Buffer) error
+}
 
 func ReadFrom(reader io.Reader) Supplier {
 	return func(b []byte) (int, error) {
@@ -13,3 +31,60 @@ func ReadFullFrom(reader io.Reader, size int) Supplier {
 		return io.ReadFull(reader, b[:size])
 	}
 }
+
+// Pipe dumps all content from reader to writer, until an error happens.
+func Pipe(reader Reader, writer Writer) error {
+	for {
+		buffer, err := reader.Read()
+		if err != nil {
+			return err
+		}
+
+		if buffer.IsEmpty() {
+			buffer.Release()
+			continue
+		}
+
+		err = writer.Write(buffer)
+		if err != nil {
+			buffer.Release()
+			return err
+		}
+	}
+}
+
+// PipeUntilEOF behaves the same as Pipe(). The only difference is PipeUntilEOF returns nil on EOF.
+func PipeUntilEOF(reader Reader, writer Writer) error {
+	err := Pipe(reader, writer)
+	if err != nil && errors.Cause(err) != io.EOF {
+		return err
+	}
+	return nil
+}
+
+// NewReader creates a new Reader.
+// The Reader instance doesn't take the ownership of reader.
+func NewReader(reader io.Reader) Reader {
+	return &BytesToBufferReader{
+		reader: reader,
+	}
+}
+
+func NewBytesReader(stream Reader) *BufferToBytesReader {
+	return &BufferToBytesReader{
+		stream: stream,
+	}
+}
+
+// NewWriter creates a new Writer.
+func NewWriter(writer io.Writer) Writer {
+	return &BufferToBytesWriter{
+		writer: writer,
+	}
+}
+
+func NewBytesWriter(writer Writer) *BytesToBufferWriter {
+	return &BytesToBufferWriter{
+		writer: writer,
+	}
+}

common/buf/reader.go

@@ -0,0 +1,102 @@
+package buf
+
+import (
+	"io"
+	"sync"
+)
+
+// BytesToBufferReader is a Reader that adjusts its reading speed automatically.
+type BytesToBufferReader struct {
+	reader      io.Reader
+	largeBuffer *Buffer
+	highVolumn  bool
+}
+
+// Read implements Reader.Read().
+func (v *BytesToBufferReader) Read() (*Buffer, error) {
+	if v.highVolumn && v.largeBuffer.IsEmpty() {
+		if v.largeBuffer == nil {
+			v.largeBuffer = NewLocal(32 * 1024)
+		}
+		err := v.largeBuffer.AppendSupplier(ReadFrom(v.reader))
+		if err != nil {
+			return nil, err
+		}
+		if v.largeBuffer.Len() < Size {
+			v.highVolumn = false
+		}
+	}
+
+	buffer := New()
+	if !v.largeBuffer.IsEmpty() {
+		buffer.AppendSupplier(ReadFrom(v.largeBuffer))
+		return buffer, nil
+	}
+
+	err := buffer.AppendSupplier(ReadFrom(v.reader))
+	if err != nil {
+		buffer.Release()
+		return nil, err
+	}
+
+	if buffer.IsFull() {
+		v.highVolumn = true
+	}
+
+	return buffer, nil
+}
+
+// Release implements Releasable.Release().
+func (v *BytesToBufferReader) Release() {
+	v.reader = nil
+}
+
+type BufferToBytesReader struct {
+	sync.Mutex
+	stream  Reader
+	current *Buffer
+	eof     bool
+}
+
+
+
+// Private: Visible for testing.
+func (v *BufferToBytesReader) Fill() {
+	b, err := v.stream.Read()
+	v.current = b
+	if err != nil {
+		v.eof = true
+		v.current = nil
+	}
+}
+
+func (v *BufferToBytesReader) Read(b []byte) (int, error) {
+	if v.eof {
+		return 0, io.EOF
+	}
+
+	v.Lock()
+	defer v.Unlock()
+	if v.current == nil {
+		v.Fill()
+		if v.eof {
+			return 0, io.EOF
+		}
+	}
+	nBytes, err := v.current.Read(b)
+	if v.current.IsEmpty() {
+		v.current.Release()
+		v.current = nil
+	}
+	return nBytes, err
+}
+
+func (v *BufferToBytesReader) Release() {
+	v.Lock()
+	defer v.Unlock()
+
+	v.eof = true
+	v.current.Release()
+	v.current = nil
+	v.stream = nil
+}

common/buf/reader_test.go

@@ -0,0 +1,28 @@
+package buf_test
+
+import (
+	"bytes"
+	"testing"
+
+	. "v2ray.com/core/common/buf"
+	"v2ray.com/core/testing/assert"
+)
+
+func TestAdaptiveReader(t *testing.T) {
+	assert := assert.On(t)
+
+	rawContent := make([]byte, 1024*1024)
+	buffer := bytes.NewBuffer(rawContent)
+
+	reader := NewReader(buffer)
+	b1, err := reader.Read()
+	assert.Error(err).IsNil()
+	assert.Bool(b1.IsFull()).IsTrue()
+	assert.Int(b1.Len()).Equals(Size)
+	assert.Int(buffer.Len()).Equals(cap(rawContent) - Size)
+
+	b2, err := reader.Read()
+	assert.Error(err).IsNil()
+	assert.Bool(b2.IsFull()).IsTrue()
+	assert.Int(buffer.Len()).Equals(1007616)
+}

common/buf/writer.go

@@ -0,0 +1,68 @@
+package buf
+
+import (
+	"io"
+	"sync"
+)
+
+// BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer.
+type BufferToBytesWriter struct {
+	writer io.Writer
+}
+
+// Write implements Writer.Write(). Write() takes ownership of the given buffer.
+func (v *BufferToBytesWriter) Write(buffer *Buffer) error {
+	defer buffer.Release()
+	for {
+		nBytes, err := v.writer.Write(buffer.Bytes())
+		if err != nil {
+			return err
+		}
+		if nBytes == buffer.Len() {
+			break
+		}
+		buffer.SliceFrom(nBytes)
+	}
+	return nil
+}
+
+// Release implements Releasable.Release().
+func (v *BufferToBytesWriter) Release() {
+	v.writer = nil
+}
+
+type BytesToBufferWriter struct {
+	sync.Mutex
+	writer Writer
+}
+
+func (v *BytesToBufferWriter) Write(payload []byte) (int, error) {
+	v.Lock()
+	defer v.Unlock()
+	if v.writer == nil {
+		return 0, io.ErrClosedPipe
+	}
+
+	bytesWritten := 0
+	size := len(payload)
+	for size > 0 {
+		buffer := New()
+		nBytes, _ := buffer.Write(payload)
+		size -= nBytes
+		payload = payload[nBytes:]
+		bytesWritten += nBytes
+		err := v.writer.Write(buffer)
+		if err != nil {
+			return bytesWritten, err
+		}
+	}
+
+	return bytesWritten, nil
+}
+
+func (v *BytesToBufferWriter) Release() {
+	v.Lock()
+	v.writer.Release()
+	v.writer = nil
+	v.Unlock()
+}

common/buf/writer_test.go

@@ -0,0 +1,27 @@
+package buf_test
+
+import (
+	"bytes"
+	"crypto/rand"
+	"testing"
+
+	. "v2ray.com/core/common/buf"
+	"v2ray.com/core/common/bufio"
+	"v2ray.com/core/testing/assert"
+)
+
+func TestWriter(t *testing.T) {
+	assert := assert.On(t)
+
+	lb := New()
+	lb.AppendSupplier(ReadFrom(rand.Reader))
+
+	expectedBytes := append([]byte(nil), lb.Bytes()...)
+
+	writeBuffer := bytes.NewBuffer(make([]byte, 0, 1024*1024))
+
+	writer := NewWriter(bufio.NewWriter(writeBuffer))
+	err := writer.Write(lb)
+	assert.Error(err).IsNil()
+	assert.Bytes(expectedBytes).Equals(writeBuffer.Bytes())
+}