From 3f8ecf023c123f64d97be3d91151289a52a7b46b Mon Sep 17 00:00:00 2001 From: Kevin Krakauer Date: Fri, 22 Mar 2024 11:58:16 -0700 Subject: [PATCH] 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 | 17 +++++++++-- pkg/sentry/socket/netstack/netstack.go | 39 +++++++++++++++++++++----- pkg/tcpip/transport/tcp/endpoint.go | 16 +++++++++-- pkg/usermem/usermem.go | 6 ++++ 4 files changed, 67 insertions(+), 11 deletions(-) diff --git a/pkg/buffer/buffer.go b/pkg/buffer/buffer.go index 79151c64b..3e6bc6dd1 100644 --- a/pkg/buffer/buffer.go +++ b/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 { diff --git a/pkg/sentry/socket/netstack/netstack.go b/pkg/sentry/socket/netstack/netstack.go index e1cfbf61c..43a8c651e 100644 --- a/pkg/sentry/socket/netstack/netstack.go +++ b/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 } diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go index 4300cb57d..cf0b0093b 100644 --- a/pkg/tcpip/transport/tcp/endpoint.go +++ b/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{} } diff --git a/pkg/usermem/usermem.go b/pkg/usermem/usermem.go index 2594294a8..b0782155c 100644 --- a/pkg/usermem/usermem.go +++ b/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)