1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-06-10 01:40:44 +00:00

rewrite SliceBySize

This commit is contained in:
Darien Raymond 2018-11-18 20:16:14 +01:00
parent bcd5d026fe
commit 5c4e33f759
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
5 changed files with 39 additions and 28 deletions

View File

@ -3,7 +3,6 @@ package buf
import (
"io"
"v2ray.com/core/common"
"v2ray.com/core/common/errors"
"v2ray.com/core/common/serial"
)
@ -127,6 +126,39 @@ func SplitFirst(mb MultiBuffer) (MultiBuffer, *Buffer) {
return mb, b
}
// SplitSize splits the beginning of the MultiBuffer into another one, for at most size bytes.
func SplitSize(mb MultiBuffer, size int32) (MultiBuffer, MultiBuffer) {
if len(mb) == 0 {
return mb, nil
}
if mb[0].Len() > size {
b := New()
copy(b.Extend(size), mb[0].BytesTo(size))
mb[0].Advance(size)
return mb, MultiBuffer{b}
}
totalBytes := int32(0)
var r MultiBuffer
endIndex := 0
for i := range mb {
if totalBytes+mb[i].Len() > size {
endIndex = i
break
}
r = append(r, mb[i])
mb[i] = nil
}
if endIndex == len(mb) {
// To reuse mb array
mb = mb[:0]
} else {
mb = mb[endIndex:]
}
return mb, r
}
// Len returns the total number of bytes in the MultiBuffer.
func (mb MultiBuffer) Len() int32 {
if mb == nil {
@ -159,29 +191,6 @@ func (mb MultiBuffer) String() string {
return serial.Concat(v...)
}
// SliceBySize splits the beginning of this MultiBuffer into another one, for at most size bytes.
func (mb *MultiBuffer) SliceBySize(size int32) MultiBuffer {
slice := make(MultiBuffer, 0, 10)
sliceSize := int32(0)
endIndex := len(*mb)
for i, b := range *mb {
if b.Len()+sliceSize > size {
endIndex = i
break
}
sliceSize += b.Len()
slice = append(slice, b)
(*mb)[i] = nil
}
*mb = (*mb)[endIndex:]
if endIndex == 0 && len(*mb) > 0 {
b := New()
common.Must2(b.ReadFullFrom((*mb)[0], size))
return MultiBuffer{b}
}
return slice
}
// MultiBufferContainer is a ReadWriteCloser wrapper over MultiBuffer.
type MultiBufferContainer struct {
MultiBuffer

View File

@ -44,6 +44,6 @@ func TestMultiBufferSliceBySizeLarge(t *testing.T) {
mb := MergeBytes(nil, lb)
mb2 := mb.SliceBySize(1024)
mb, mb2 := SplitSize(mb, 1024)
assert(mb2.Len(), Equals, int32(1024))
}

View File

@ -87,7 +87,8 @@ func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
r.Buffer = mb
}
mb := r.Buffer.SliceBySize(size)
rb, mb := SplitSize(r.Buffer, size)
r.Buffer = rb
if r.Buffer.IsEmpty() {
r.Buffer = nil
}

View File

@ -143,7 +143,8 @@ func (w *ChunkStreamWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
mb2Write := make(buf.MultiBuffer, 0, mbLen/buf.Size+mbLen/sliceSize+2)
for {
slice := mb.SliceBySize(sliceSize)
mb2, slice := buf.SplitSize(mb, sliceSize)
mb = mb2
b := buf.New()
w.sizeEncoder.Encode(uint16(slice.Len()), b.Extend(w.sizeEncoder.SizeBytes()))

View File

@ -94,7 +94,7 @@ func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
for !mb.IsEmpty() {
var chunk buf.MultiBuffer
if w.transferType == protocol.TransferTypeStream {
chunk = mb.SliceBySize(8 * 1024)
mb, chunk = buf.SplitSize(mb, 8*1024)
} else {
mb2, b := buf.SplitFirst(mb)
mb = mb2