netstack: don't allocate interfaces when copying data in and out

In a redis-benchmark PING_INLINE test, this reduces allocations by 32%.

PiperOrigin-RevId: 618248114

pkg/buffer/buffer.go

@@ -492,6 +492,18 @@ func (b *Buffer) Merge(other *Buffer) {
 // WriteFromReader writes to the buffer from an io.Reader. A maximum read size
 // of MaxChunkSize is enforced to prevent allocating views from the heap.
 func (b *Buffer) WriteFromReader(r io.Reader, count int64) (int64, error) {
+	return b.WriteFromReaderAndLimitedReader(r, count, nil)
+}
+
+// WriteFromReaderAndLimitedReader is the same as WriteFromReader, but
+// optimized to avoid allocations if a LimitedReader is passed in.
+//
+// This function clobbers the values of lr.
+func (b *Buffer) WriteFromReaderAndLimitedReader(r io.Reader, count int64, lr *io.LimitedReader) (int64, error) {
+	if lr == nil {
+		lr = &io.LimitedReader{}
+	}
+
 	var done int64
 	for done < count {
 		vsize := count - done
@@ -499,8 +511,9 @@ func (b *Buffer) WriteFromReader(r io.Reader, count int64) (int64, error) {
 			vsize = MaxChunkSize
 		}
 		v := NewView(int(vsize))
-		lr := io.LimitedReader{R: r, N: vsize}
-		n, err := io.Copy(v, &lr)
+		lr.R = r
+		lr.N = vsize
+		n, err := io.Copy(v, lr)
 		b.Append(v)
 		done += n
 		if err == io.EOF {

pkg/sentry/socket/netstack/netstack.go

@@ -366,6 +366,11 @@ type sock struct {
 
 	namespace *inet.Namespace
 
+	mu sync.Mutex `state:"nosave"`
+	// readWriter is an optimization to avoid allocations.
+	// +checklocks:mu
+	readWriter usermem.IOSequenceReadWriter `state:"nosave"`
+
 	// readMu protects access to the below fields.
 	readMu sync.Mutex `state:"nosave"`
 
@@ -482,8 +487,17 @@ func (s *sock) Write(ctx context.Context, src usermem.IOSequence, opts vfs.Write
 		return 0, linuxerr.EOPNOTSUPP
 	}
 
-	r := src.Reader(ctx)
-	n, err := s.Endpoint.Write(r, tcpip.WriteOptions{})
+	var n int64
+	var err tcpip.Error
+	switch s.Endpoint.(type) {
+	case *tcp.Endpoint:
+		s.mu.Lock()
+		s.readWriter.Init(ctx, src)
+		n, err = s.Endpoint.Write(&s.readWriter, tcpip.WriteOptions{})
+		s.mu.Unlock()
+	default:
+		n, err = s.Endpoint.Write(src.Reader(ctx), tcpip.WriteOptions{})
+	}
 	if _, ok := err.(*tcpip.ErrWouldBlock); ok {
 		return 0, linuxerr.ErrWouldBlock
 	}
@@ -2670,19 +2684,30 @@ func (s *sock) nonBlockingRead(ctx context.Context, dst usermem.IOSequence, peek
 	// bytes of data to be discarded, rather than passed back in a
 	// caller-supplied  buffer.
 	var w io.Writer
+	var res tcpip.ReadResult
+	var err tcpip.Error
+
+	s.readMu.Lock()
+	defer s.readMu.Unlock()
+
 	if !isPacket && trunc {
 		w = &tcpip.LimitedWriter{
 			W: ioutil.Discard,
 			N: dst.NumBytes(),
 		}
+		res, err = s.Endpoint.Read(w, readOptions)
 	} else {
-		w = dst.Writer(ctx)
+		switch s.Endpoint.(type) {
+		case *tcp.Endpoint:
+			s.mu.Lock()
+			s.readWriter.Init(ctx, dst)
+			res, err = s.Endpoint.Read(&s.readWriter, readOptions)
+			s.mu.Unlock()
+		default:
+			res, err = s.Endpoint.Read(dst.Writer(ctx), readOptions)
+		}
 	}
 
-	s.readMu.Lock()
-	defer s.readMu.Unlock()
-
-	res, err := s.Endpoint.Read(w, readOptions)
 	if _, ok := err.(*tcpip.ErrBadBuffer); ok && dst.NumBytes() == 0 {
 		err = nil
 	}

pkg/tcpip/transport/tcp/endpoint.go

@@ -348,7 +348,7 @@ type Endpoint struct {
 	stack.TransportEndpointInfo
 	tcpip.DefaultSocketOptionsHandler
 
-	// endpointEntry is used to queue endpoints for processing to the
+	// EndpointEntry is used to queue endpoints for processing to the
 	// a given tcp processor goroutine.
 	//
 	// Precondition: epQueue.mu must be held to read/write this field..
@@ -595,6 +595,11 @@ type Endpoint struct {
 	// listenCtx is used by listening endpoints to store state used while listening for
 	// connections. Nil otherwise.
 	listenCtx *listenContext `state:"nosave"`
+
+	// limRdr is reused to avoid allocations.
+	//
+	// +checklocks:mu
+	limRdr *io.LimitedReader `state:"nosave"`
 }
 
 // UniqueID implements stack.TransportEndpoint.UniqueID.
@@ -864,6 +869,7 @@ func newEndpoint(s *stack.Stack, protocol *protocol, netProto tcpip.NetworkProto
 		txHash:        s.InsecureRNG().Uint32(),
 		windowClamp:   DefaultReceiveBufferSize,
 		maxSynRetries: DefaultSynRetries,
+		limRdr:        &io.LimitedReader{},
 	}
 	e.ops.InitHandler(e, e.stack, GetTCPSendBufferLimits, GetTCPReceiveBufferLimits)
 	e.ops.SetMulticastLoop(true)
@@ -1560,7 +1566,13 @@ func (e *Endpoint) readFromPayloader(p tcpip.Payloader, opts tcpip.WriteOptions,
 	// This is not possible if atomic is set, because we can't allow the
 	// available buffer space to be consumed by some other caller while we
 	// are copying data in.
+	limRdr := e.limRdr
 	if !opts.Atomic {
+		defer func() {
+			e.limRdr = limRdr
+		}()
+		e.limRdr = nil
+
 		e.sndQueueInfo.sndQueueMu.Unlock()
 		defer e.sndQueueInfo.sndQueueMu.Lock()
 
@@ -1576,7 +1588,7 @@ func (e *Endpoint) readFromPayloader(p tcpip.Payloader, opts tcpip.WriteOptions,
 	if avail == 0 {
 		return payload, nil
 	}
-	if _, err := payload.WriteFromReader(p, int64(avail)); err != nil {
+	if _, err := payload.WriteFromReaderAndLimitedReader(p, int64(avail), limRdr); err != nil {
 		payload.Release()
 		return buffer.Buffer{}, &tcpip.ErrBadBuffer{}
 	}

pkg/usermem/usermem.go

@@ -539,6 +539,12 @@ type IOSequenceReadWriter struct {
 	s   IOSequence
 }
 
+// Init initializes the IOSequence.
+func (rw *IOSequenceReadWriter) Init(ctx context.Context, src IOSequence) {
+	rw.ctx = ctx
+	rw.s = src
+}
+
 // Read implements io.Reader.Read.
 func (rw *IOSequenceReadWriter) Read(dst []byte) (int, error) {
 	n, err := rw.s.CopyIn(rw.ctx, dst)