github.com/xraypb/Xray-core@v1.8.1/common/buf/readv_reader.go (about)

     1  //go:build !wasm
     2  // +build !wasm
     3  
     4  package buf
     5  
     6  import (
     7  	"io"
     8  	"syscall"
     9  
    10  	"github.com/xraypb/Xray-core/common/platform"
    11  	"github.com/xraypb/Xray-core/features/stats"
    12  )
    13  
    14  type allocStrategy struct {
    15  	current uint32
    16  }
    17  
    18  func (s *allocStrategy) Current() uint32 {
    19  	return s.current
    20  }
    21  
    22  func (s *allocStrategy) Adjust(n uint32) {
    23  	if n >= s.current {
    24  		s.current *= 2
    25  	} else {
    26  		s.current = n
    27  	}
    28  
    29  	if s.current > 8 {
    30  		s.current = 8
    31  	}
    32  
    33  	if s.current == 0 {
    34  		s.current = 1
    35  	}
    36  }
    37  
    38  func (s *allocStrategy) Alloc() []*Buffer {
    39  	bs := make([]*Buffer, s.current)
    40  	for i := range bs {
    41  		bs[i] = New()
    42  	}
    43  	return bs
    44  }
    45  
    46  type multiReader interface {
    47  	Init([]*Buffer)
    48  	Read(fd uintptr) int32
    49  	Clear()
    50  }
    51  
    52  // ReadVReader is a Reader that uses readv(2) syscall to read data.
    53  type ReadVReader struct {
    54  	io.Reader
    55  	rawConn syscall.RawConn
    56  	mr      multiReader
    57  	alloc   allocStrategy
    58  	counter stats.Counter
    59  }
    60  
    61  // NewReadVReader creates a new ReadVReader.
    62  func NewReadVReader(reader io.Reader, rawConn syscall.RawConn, counter stats.Counter) *ReadVReader {
    63  	return &ReadVReader{
    64  		Reader:  reader,
    65  		rawConn: rawConn,
    66  		alloc: allocStrategy{
    67  			current: 1,
    68  		},
    69  		mr:      newMultiReader(),
    70  		counter: counter,
    71  	}
    72  }
    73  
    74  func (r *ReadVReader) readMulti() (MultiBuffer, error) {
    75  	bs := r.alloc.Alloc()
    76  
    77  	r.mr.Init(bs)
    78  	var nBytes int32
    79  	err := r.rawConn.Read(func(fd uintptr) bool {
    80  		n := r.mr.Read(fd)
    81  		if n < 0 {
    82  			return false
    83  		}
    84  
    85  		nBytes = n
    86  		return true
    87  	})
    88  	r.mr.Clear()
    89  
    90  	if err != nil {
    91  		ReleaseMulti(MultiBuffer(bs))
    92  		return nil, err
    93  	}
    94  
    95  	if nBytes == 0 {
    96  		ReleaseMulti(MultiBuffer(bs))
    97  		return nil, io.EOF
    98  	}
    99  
   100  	nBuf := 0
   101  	for nBuf < len(bs) {
   102  		if nBytes <= 0 {
   103  			break
   104  		}
   105  		end := nBytes
   106  		if end > Size {
   107  			end = Size
   108  		}
   109  		bs[nBuf].end = end
   110  		nBytes -= end
   111  		nBuf++
   112  	}
   113  
   114  	for i := nBuf; i < len(bs); i++ {
   115  		bs[i].Release()
   116  		bs[i] = nil
   117  	}
   118  
   119  	return MultiBuffer(bs[:nBuf]), nil
   120  }
   121  
   122  // ReadMultiBuffer implements Reader.
   123  func (r *ReadVReader) ReadMultiBuffer() (MultiBuffer, error) {
   124  	if r.alloc.Current() == 1 {
   125  		b, err := ReadBuffer(r.Reader)
   126  		if b.IsFull() {
   127  			r.alloc.Adjust(1)
   128  		}
   129  		if r.counter != nil && b != nil {
   130  			r.counter.Add(int64(b.Len()))
   131  		}
   132  		return MultiBuffer{b}, err
   133  	}
   134  
   135  	mb, err := r.readMulti()
   136  	if r.counter != nil && mb != nil {
   137  		r.counter.Add(int64(mb.Len()))
   138  	}
   139  	if err != nil {
   140  		return nil, err
   141  	}
   142  	r.alloc.Adjust(uint32(len(mb)))
   143  	return mb, nil
   144  }
   145  
   146  var useReadv bool
   147  
   148  func init() {
   149  	const defaultFlagValue = "NOT_DEFINED_AT_ALL"
   150  	value := platform.NewEnvFlag("xray.buf.readv").GetValue(func() string { return defaultFlagValue })
   151  	switch value {
   152  	case defaultFlagValue, "auto", "enable":
   153  		useReadv = true
   154  	}
   155  }