mirror of https://github.com/v2ray/v2ray-core
parent
e46bad3f18
commit
abe8ffda68
|
@ -17,7 +17,7 @@ type BufferedWriter struct {
|
||||||
func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
|
func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
|
||||||
return &BufferedWriter{
|
return &BufferedWriter{
|
||||||
writer: rawWriter,
|
writer: rawWriter,
|
||||||
buffer: alloc.NewBuffer().Clear(),
|
buffer: alloc.NewSmallBuffer().Clear(),
|
||||||
cached: true,
|
cached: true,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -55,11 +55,22 @@ func (v *BufferedWriter) Write(b []byte) (int, error) {
|
||||||
if !v.cached {
|
if !v.cached {
|
||||||
return v.writer.Write(b)
|
return v.writer.Write(b)
|
||||||
}
|
}
|
||||||
nBytes, _ := v.buffer.Write(b)
|
nBytes, err := v.buffer.Write(b)
|
||||||
if v.buffer.IsFull() {
|
if err != nil {
|
||||||
v.FlushWithoutLock()
|
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 {
|
func (v *BufferedWriter) Flush() error {
|
||||||
|
|
|
@ -1,6 +1,7 @@
|
||||||
package io_test
|
package io_test
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"crypto/rand"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"v2ray.com/core/common/alloc"
|
"v2ray.com/core/common/alloc"
|
||||||
|
@ -27,3 +28,26 @@ func TestBufferedWriter(t *testing.T) {
|
||||||
writer.SetCached(false)
|
writer.SetCached(false)
|
||||||
assert.Int(content.Len()).Equals(16)
|
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)
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue