1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 01:57:12 -05:00

fix buffered writer

This commit is contained in:
Darien Raymond 2016-12-04 21:34:22 +01:00
parent e46bad3f18
commit abe8ffda68
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
2 changed files with 40 additions and 5 deletions

View File

@ -17,7 +17,7 @@ type BufferedWriter struct {
func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
return &BufferedWriter{
writer: rawWriter,
buffer: alloc.NewBuffer().Clear(),
buffer: alloc.NewSmallBuffer().Clear(),
cached: true,
}
}
@ -55,11 +55,22 @@ func (v *BufferedWriter) Write(b []byte) (int, error) {
if !v.cached {
return v.writer.Write(b)
}
nBytes, _ := v.buffer.Write(b)
if v.buffer.IsFull() {
v.FlushWithoutLock()
nBytes, err := v.buffer.Write(b)
if err != nil {
return 0, err
}
return nBytes, nil
if v.buffer.IsFull() {
err := v.FlushWithoutLock()
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
}
func (v *BufferedWriter) Flush() error {

View File

@ -1,6 +1,7 @@
package io_test
import (
"crypto/rand"
"testing"
"v2ray.com/core/common/alloc"
@ -27,3 +28,26 @@ func TestBufferedWriter(t *testing.T) {
writer.SetCached(false)
assert.Int(content.Len()).Equals(16)
}
func TestBufferedWriterLargePayload(t *testing.T) {
assert := assert.On(t)
content := alloc.NewLocalBuffer(128 * 1024).Clear()
writer := NewBufferedWriter(content)
assert.Bool(writer.Cached()).IsTrue()
payload := make([]byte, 64*1024)
rand.Read(payload)
nBytes, err := writer.Write(payload[:1024])
assert.Int(nBytes).Equals(1024)
assert.Error(err).IsNil()
assert.Bool(content.IsEmpty()).IsTrue()
nBytes, err = writer.Write(payload[1024:])
assert.Error(err).IsNil()
assert.Int(nBytes).Equals(63 * 1024)
assert.Bytes(content.Value).Equals(payload)
}