merge bufio into buf

common/buf/buffer.go

@@ -175,17 +175,17 @@ func (b *Buffer) String() string {
 	return string(b.Bytes())
 }
 
-// New creates a Buffer with 8K bytes of arbitrary content.
+// New creates a Buffer with 0 length and 8K capacity.
 func New() *Buffer {
 	return mediumPool.Allocate()
 }
 
-// NewSmall returns a buffer with 2K bytes capacity.
+// NewSmall returns a buffer with 0 length and 2K capacity.
 func NewSmall() *Buffer {
 	return smallPool.Allocate()
 }
 
-// NewLocal creates and returns a buffer on current thread.
+// NewLocal creates and returns a buffer with 0 length and given capacity on current thread.
 func NewLocal(size int) *Buffer {
 	return &Buffer{
 		v:    make([]byte, size),

common/buf/buffered_reader.go

@@ -0,0 +1,54 @@
+package buf
+
+import (
+	"io"
+)
+
+// BufferedReader is a reader with internal cache.
+type BufferedReader struct {
+	reader   io.Reader
+	buffer   *Buffer
+	buffered bool
+}
+
+// NewReader creates a new BufferedReader based on an io.Reader.
+func NewBufferedReader(rawReader io.Reader) *BufferedReader {
+	return &BufferedReader{
+		reader:   rawReader,
+		buffer:   NewLocal(1024),
+		buffered: true,
+	}
+}
+
+// IsBuffered returns true if the internal cache is effective.
+func (v *BufferedReader) IsBuffered() bool {
+	return v.buffered
+}
+
+// SetBuffered is to enable or disable internal cache. If cache is disabled,
+// Read() calls will be delegated to the underlying io.Reader directly.
+func (v *BufferedReader) SetBuffered(cached bool) {
+	v.buffered = cached
+}
+
+// Read implements io.Reader.Read().
+func (v *BufferedReader) Read(b []byte) (int, error) {
+	if !v.buffered || v.buffer == nil {
+		if !v.buffer.IsEmpty() {
+			return v.buffer.Read(b)
+		}
+		return v.reader.Read(b)
+	}
+	if v.buffer.IsEmpty() {
+		err := v.buffer.AppendSupplier(ReadFrom(v.reader))
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	if v.buffer.IsEmpty() {
+		return 0, nil
+	}
+
+	return v.buffer.Read(b)
+}

common/buf/buffered_reader_test.go

@@ -0,0 +1,36 @@
+package buf_test
+
+import (
+	"crypto/rand"
+	"testing"
+
+	. "v2ray.com/core/common/buf"
+	"v2ray.com/core/testing/assert"
+)
+
+func TestBufferedReader(t *testing.T) {
+	assert := assert.On(t)
+
+	content := New()
+	assert.Error(content.AppendSupplier(ReadFrom(rand.Reader))).IsNil()
+
+	len := content.Len()
+
+	reader := NewBufferedReader(content)
+	assert.Bool(reader.IsBuffered()).IsTrue()
+
+	payload := make([]byte, 16)
+
+	nBytes, err := reader.Read(payload)
+	assert.Int(nBytes).Equals(16)
+	assert.Error(err).IsNil()
+
+	len2 := content.Len()
+	assert.Int(len - len2).GreaterThan(16)
+
+	nBytes, err = reader.Read(payload)
+	assert.Int(nBytes).Equals(16)
+	assert.Error(err).IsNil()
+
+	assert.Int(content.Len()).Equals(len2)
+}

common/buf/buffered_writer.go

@@ -0,0 +1,92 @@
+package buf
+
+import (
+	"io"
+
+	"v2ray.com/core/common/errors"
+)
+
+// BufferedWriter is an io.Writer with internal buffer. It writes to underlying writer when buffer is full or on demand.
+// This type is not thread safe.
+type BufferedWriter struct {
+	writer   io.Writer
+	buffer   *Buffer
+	buffered bool
+}
+
+// NewWriter creates a new BufferedWriter.
+func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
+	return &BufferedWriter{
+		writer:   rawWriter,
+		buffer:   NewLocal(1024),
+		buffered: true,
+	}
+}
+
+// ReadFrom implements io.ReaderFrom.ReadFrom().
+func (v *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
+	totalBytes := int64(0)
+	for {
+		oriSize := v.buffer.Len()
+		err := v.buffer.AppendSupplier(ReadFrom(reader))
+		totalBytes += int64(v.buffer.Len() - oriSize)
+		if err != nil {
+			if errors.Cause(err) == io.EOF {
+				return totalBytes, nil
+			}
+			return totalBytes, err
+		}
+		if err := v.Flush(); err != nil {
+			return totalBytes, err
+		}
+	}
+}
+
+func (v *BufferedWriter) Write(b []byte) (int, error) {
+	if !v.buffered || v.buffer == nil {
+		return v.writer.Write(b)
+	}
+	nBytes, err := v.buffer.Write(b)
+	if err != nil {
+		return 0, err
+	}
+	if v.buffer.IsFull() {
+		err := v.Flush()
+		if err != nil {
+			return 0, err
+		}
+		if nBytes < len(b) {
+			if _, err := v.writer.Write(b[nBytes:]); err != nil {
+				return nBytes, err
+			}
+		}
+	}
+	return len(b), nil
+}
+
+// Flush writes all buffered content into underlying writer, if any.
+func (v *BufferedWriter) Flush() error {
+	defer v.buffer.Clear()
+	for !v.buffer.IsEmpty() {
+		nBytes, err := v.writer.Write(v.buffer.Bytes())
+		if err != nil {
+			return err
+		}
+		v.buffer.SliceFrom(nBytes)
+	}
+	return nil
+}
+
+// IsBuffered returns true if this BufferedWriter holds a buffer.
+func (v *BufferedWriter) IsBuffered() bool {
+	return v.buffered
+}
+
+// SetBuffered controls whether the BufferedWriter holds a buffer for writing. If not buffered, any write() calls into underlying writer directly.
+func (v *BufferedWriter) SetBuffered(cached bool) error {
+	v.buffered = cached
+	if !cached && !v.buffer.IsEmpty() {
+		return v.Flush()
+	}
+	return nil
+}

common/buf/buffered_writer_test.go

@@ -0,0 +1,52 @@
+package buf_test
+
+import (
+	"crypto/rand"
+	"testing"
+
+	. "v2ray.com/core/common/buf"
+	"v2ray.com/core/testing/assert"
+)
+
+func TestBufferedWriter(t *testing.T) {
+	assert := assert.On(t)
+
+	content := New()
+
+	writer := NewBufferedWriter(content)
+	assert.Bool(writer.IsBuffered()).IsTrue()
+
+	payload := make([]byte, 16)
+
+	nBytes, err := writer.Write(payload)
+	assert.Int(nBytes).Equals(16)
+	assert.Error(err).IsNil()
+
+	assert.Bool(content.IsEmpty()).IsTrue()
+
+	writer.SetBuffered(false)
+	assert.Int(content.Len()).Equals(16)
+}
+
+func TestBufferedWriterLargePayload(t *testing.T) {
+	assert := assert.On(t)
+
+	content := NewLocal(128 * 1024)
+
+	writer := NewBufferedWriter(content)
+	assert.Bool(writer.IsBuffered()).IsTrue()
+
+	payload := make([]byte, 64*1024)
+	rand.Read(payload)
+
+	nBytes, err := writer.Write(payload[:512])
+	assert.Int(nBytes).Equals(512)
+	assert.Error(err).IsNil()
+
+	assert.Bool(content.IsEmpty()).IsTrue()
+
+	nBytes, err = writer.Write(payload[512:])
+	assert.Error(err).IsNil()
+	assert.Int(nBytes).Equals(64*1024 - 512)
+	assert.Bytes(content.Bytes()).Equals(payload)
+}

common/buf/io.go

@@ -33,7 +33,8 @@ func ReadFullFrom(reader io.Reader, size int) Supplier {
 	}
 }
 
-// Pipe dumps all content from reader to writer, until an error happens.
+// Pipe dumps all payload from reader to writer, until an error occurs.
+// ActivityTimer gets updated as soon as there is a payload.
 func Pipe(timer *signal.ActivityTimer, reader Reader, writer Writer) error {
 	for {
 		buffer, err := reader.Read()
@@ -73,7 +74,8 @@ func NewReader(reader io.Reader) Reader {
 	}
 }
 
-func NewBytesReader(stream Reader) *BufferToBytesReader {
+// ToBytesReader converts a Reaaer to io.Reader.
+func ToBytesReader(stream Reader) io.Reader {
 	return &BufferToBytesReader{
 		stream: stream,
 	}
@@ -86,7 +88,8 @@ func NewWriter(writer io.Writer) Writer {
 	}
 }
 
-func NewBytesWriter(writer Writer) *BytesToBufferWriter {
+// ToBytesWriter converts a Writer to io.Writer
+func ToBytesWriter(writer Writer) io.Writer {
 	return &BytesToBufferWriter{
 		writer: writer,
 	}

common/buf/writer_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 
 	. "v2ray.com/core/common/buf"
-	"v2ray.com/core/common/bufio"
 	"v2ray.com/core/testing/assert"
 )
 
@@ -20,7 +19,7 @@ func TestWriter(t *testing.T) {
 
 	writeBuffer := bytes.NewBuffer(make([]byte, 0, 1024*1024))
 
-	writer := NewWriter(bufio.NewWriter(writeBuffer))
+	writer := NewWriter(NewBufferedWriter(writeBuffer))
 	err := writer.Write(lb)
 	assert.Error(err).IsNil()
 	assert.Bytes(expectedBytes).Equals(writeBuffer.Bytes())