github.com/twelsh-aw/go/src@v0.0.0-20230516233729-a56fe86a7c81/os/readfrom_linux.go (about)

     1  // Copyright 2020 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package os
     6  
     7  import (
     8  	"internal/poll"
     9  	"io"
    10  	"syscall"
    11  )
    12  
    13  var (
    14  	pollCopyFileRange = poll.CopyFileRange
    15  	pollSplice        = poll.Splice
    16  )
    17  
    18  func (f *File) readFrom(r io.Reader) (written int64, handled bool, err error) {
    19  	// Neither copy_file_range(2) nor splice(2) supports destinations opened with
    20  	// O_APPEND, so don't bother to try zero-copy with these system calls.
    21  	//
    22  	// Visit https://man7.org/linux/man-pages/man2/copy_file_range.2.html#ERRORS and
    23  	// https://man7.org/linux/man-pages/man2/splice.2.html#ERRORS for details.
    24  	if f.appendMode {
    25  		return 0, false, nil
    26  	}
    27  
    28  	written, handled, err = f.copyFileRange(r)
    29  	if handled {
    30  		return
    31  	}
    32  	return f.spliceToFile(r)
    33  }
    34  
    35  func (f *File) spliceToFile(r io.Reader) (written int64, handled bool, err error) {
    36  	var (
    37  		remain int64
    38  		lr     *io.LimitedReader
    39  	)
    40  	if lr, r, remain = tryLimitedReader(r); remain <= 0 {
    41  		return 0, true, nil
    42  	}
    43  
    44  	pfd := getPollFD(r)
    45  	// TODO(panjf2000): run some tests to see if we should unlock the non-streams for splice.
    46  	// Streams benefit the most from the splice(2), non-streams are not even supported in old kernels
    47  	// where splice(2) will just return EINVAL; newer kernels support non-streams like UDP, but I really
    48  	// doubt that splice(2) could help non-streams, cuz they usually send small frames respectively
    49  	// and one splice call would result in one frame.
    50  	// splice(2) is suitable for large data but the generation of fragments defeats its edge here.
    51  	// Therefore, don't bother to try splice if the r is not a streaming descriptor.
    52  	if pfd == nil || !pfd.IsStream {
    53  		return
    54  	}
    55  
    56  	var syscallName string
    57  	written, handled, syscallName, err = pollSplice(&f.pfd, pfd, remain)
    58  
    59  	if lr != nil {
    60  		lr.N = remain - written
    61  	}
    62  
    63  	return written, handled, wrapSyscallError(syscallName, err)
    64  }
    65  
    66  // getPollFD tries to get the poll.FD from the given io.Reader by expecting
    67  // the underlying type of r to be the implementation of syscall.Conn that contains
    68  // a *net.rawConn.
    69  func getPollFD(r io.Reader) *poll.FD {
    70  	sc, ok := r.(syscall.Conn)
    71  	if !ok {
    72  		return nil
    73  	}
    74  	rc, err := sc.SyscallConn()
    75  	if err != nil {
    76  		return nil
    77  	}
    78  	ipfd, ok := rc.(interface{ PollFD() *poll.FD })
    79  	if !ok {
    80  		return nil
    81  	}
    82  	return ipfd.PollFD()
    83  }
    84  
    85  func (f *File) copyFileRange(r io.Reader) (written int64, handled bool, err error) {
    86  	var (
    87  		remain int64
    88  		lr     *io.LimitedReader
    89  	)
    90  	if lr, r, remain = tryLimitedReader(r); remain <= 0 {
    91  		return 0, true, nil
    92  	}
    93  
    94  	src, ok := r.(*File)
    95  	if !ok {
    96  		return 0, false, nil
    97  	}
    98  	if src.checkValid("ReadFrom") != nil {
    99  		// Avoid returning the error as we report handled as false,
   100  		// leave further error handling as the responsibility of the caller.
   101  		return 0, false, nil
   102  	}
   103  
   104  	written, handled, err = pollCopyFileRange(&f.pfd, &src.pfd, remain)
   105  	if lr != nil {
   106  		lr.N -= written
   107  	}
   108  	return written, handled, wrapSyscallError("copy_file_range", err)
   109  }
   110  
   111  // tryLimitedReader tries to assert the io.Reader to io.LimitedReader, it returns the io.LimitedReader,
   112  // the underlying io.Reader and the remaining amount of bytes if the assertion succeeds,
   113  // otherwise it just returns the original io.Reader and the theoretical unlimited remaining amount of bytes.
   114  func tryLimitedReader(r io.Reader) (*io.LimitedReader, io.Reader, int64) {
   115  	var remain int64 = 1<<63 - 1 // by default, copy until EOF
   116  
   117  	lr, ok := r.(*io.LimitedReader)
   118  	if !ok {
   119  		return nil, r, remain
   120  	}
   121  
   122  	remain = lr.N
   123  	return lr, lr.R, remain
   124  }