github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/buf/readv_reader.go (about)

     1  //go:build !wasm
     2  // +build !wasm
     3  
     4  package buf
     5  
     6  import (
     7  	"io"
     8  	"runtime"
     9  	"syscall"
    10  
    11  	"github.com/v2fly/v2ray-core/v5/common/platform"
    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 *= 4
    25  	} else {
    26  		s.current = n
    27  	}
    28  
    29  	if s.current > 32 {
    30  		s.current = 32
    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  }
    59  
    60  // NewReadVReader creates a new ReadVReader.
    61  func NewReadVReader(reader io.Reader, rawConn syscall.RawConn) *ReadVReader {
    62  	return &ReadVReader{
    63  		Reader:  reader,
    64  		rawConn: rawConn,
    65  		alloc: allocStrategy{
    66  			current: 1,
    67  		},
    68  		mr: newMultiReader(),
    69  	}
    70  }
    71  
    72  func (r *ReadVReader) readMulti() (MultiBuffer, error) {
    73  	bs := r.alloc.Alloc()
    74  
    75  	r.mr.Init(bs)
    76  	var nBytes int32
    77  	err := r.rawConn.Read(func(fd uintptr) bool {
    78  		n := r.mr.Read(fd)
    79  		if n < 0 {
    80  			return false
    81  		}
    82  
    83  		nBytes = n
    84  		return true
    85  	})
    86  	r.mr.Clear()
    87  
    88  	if err != nil {
    89  		ReleaseMulti(MultiBuffer(bs))
    90  		return nil, err
    91  	}
    92  
    93  	if nBytes == 0 {
    94  		ReleaseMulti(MultiBuffer(bs))
    95  		return nil, io.EOF
    96  	}
    97  
    98  	nBuf := 0
    99  	for nBuf < len(bs) {
   100  		if nBytes <= 0 {
   101  			break
   102  		}
   103  		end := nBytes
   104  		if end > Size {
   105  			end = Size
   106  		}
   107  		bs[nBuf].end = end
   108  		nBytes -= end
   109  		nBuf++
   110  	}
   111  
   112  	for i := nBuf; i < len(bs); i++ {
   113  		bs[i].Release()
   114  		bs[i] = nil
   115  	}
   116  
   117  	return MultiBuffer(bs[:nBuf]), nil
   118  }
   119  
   120  // ReadMultiBuffer implements Reader.
   121  func (r *ReadVReader) ReadMultiBuffer() (MultiBuffer, error) {
   122  	if r.alloc.Current() == 1 {
   123  		b, err := ReadBuffer(r.Reader)
   124  		if b.IsFull() {
   125  			r.alloc.Adjust(1)
   126  		}
   127  		return MultiBuffer{b}, err
   128  	}
   129  
   130  	mb, err := r.readMulti()
   131  	if err != nil {
   132  		return nil, err
   133  	}
   134  	r.alloc.Adjust(uint32(len(mb)))
   135  	return mb, nil
   136  }
   137  
   138  var useReadv = false
   139  
   140  func init() {
   141  	const defaultFlagValue = "NOT_DEFINED_AT_ALL"
   142  	value := platform.NewEnvFlag("v2ray.buf.readv").GetValue(func() string { return defaultFlagValue })
   143  	switch value {
   144  	case defaultFlagValue, "auto":
   145  		if (runtime.GOARCH == "386" || runtime.GOARCH == "amd64" || runtime.GOARCH == "s390x") && (runtime.GOOS == "linux" || runtime.GOOS == "darwin" || runtime.GOOS == "windows") {
   146  			useReadv = true
   147  		}
   148  	case "enable":
   149  		useReadv = true
   150  	}
   151  }