Xray-core/common/buf/readv_reader.go
2023-10-29 15:16:57 -04:00

156 lines
2.6 KiB
Go

//go:build !wasm
// +build !wasm
package buf
import (
"io"
"syscall"
"github.com/xtls/xray-core/common/platform"
"github.com/xtls/xray-core/features/stats"
)
type allocStrategy struct {
current uint32
}
func (s *allocStrategy) Current() uint32 {
return s.current
}
func (s *allocStrategy) Adjust(n uint32) {
if n >= s.current {
s.current *= 2
} else {
s.current = n
}
if s.current > 8 {
s.current = 8
}
if s.current == 0 {
s.current = 1
}
}
func (s *allocStrategy) Alloc() []*Buffer {
bs := make([]*Buffer, s.current)
for i := range bs {
bs[i] = New()
}
return bs
}
type multiReader interface {
Init([]*Buffer)
Read(fd uintptr) int32
Clear()
}
// ReadVReader is a Reader that uses readv(2) syscall to read data.
type ReadVReader struct {
io.Reader
rawConn syscall.RawConn
mr multiReader
alloc allocStrategy
counter stats.Counter
}
// NewReadVReader creates a new ReadVReader.
func NewReadVReader(reader io.Reader, rawConn syscall.RawConn, counter stats.Counter) *ReadVReader {
return &ReadVReader{
Reader: reader,
rawConn: rawConn,
alloc: allocStrategy{
current: 1,
},
mr: newMultiReader(),
counter: counter,
}
}
func (r *ReadVReader) readMulti() (MultiBuffer, error) {
bs := r.alloc.Alloc()
r.mr.Init(bs)
var nBytes int32
err := r.rawConn.Read(func(fd uintptr) bool {
n := r.mr.Read(fd)
if n < 0 {
return false
}
nBytes = n
return true
})
r.mr.Clear()
if err != nil {
ReleaseMulti(MultiBuffer(bs))
return nil, err
}
if nBytes == 0 {
ReleaseMulti(MultiBuffer(bs))
return nil, io.EOF
}
nBuf := 0
for nBuf < len(bs) {
if nBytes <= 0 {
break
}
end := nBytes
if end > Size {
end = Size
}
bs[nBuf].end = end
nBytes -= end
nBuf++
}
for i := nBuf; i < len(bs); i++ {
bs[i].Release()
bs[i] = nil
}
return MultiBuffer(bs[:nBuf]), nil
}
// ReadMultiBuffer implements Reader.
func (r *ReadVReader) ReadMultiBuffer() (MultiBuffer, error) {
if r.alloc.Current() == 1 {
b, err := ReadBuffer(r.Reader)
if b.IsFull() {
r.alloc.Adjust(1)
}
if r.counter != nil && b != nil {
r.counter.Add(int64(b.Len()))
}
return MultiBuffer{b}, err
}
mb, err := r.readMulti()
if r.counter != nil && mb != nil {
r.counter.Add(int64(mb.Len()))
}
if err != nil {
return nil, err
}
r.alloc.Adjust(uint32(len(mb)))
return mb, nil
}
var useReadv bool
func init() {
const defaultFlagValue = "NOT_DEFINED_AT_ALL"
value := platform.NewEnvFlag(platform.UseReadV).GetValue(func() string { return defaultFlagValue })
switch value {
case defaultFlagValue, "auto", "enable":
useReadv = true
}
}