github.com/eagleql/xray-core@v1.4.4/common/buf/readv_reader.go (about)

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