From 2759f79fcee177b353012103b67ec68f5b867018 Mon Sep 17 00:00:00 2001 From: Ayush Ranjan Date: Fri, 23 Feb 2024 12:28:56 -0800 Subject: [PATCH] Plumb safemem.ReaderFunc through pgalloc.Allocate() to avoid heap allocations. The safemem.Reader interface receiver was causing the implementation to escape to heap (as at compile time, the compiler can not prove anything about the implementation). Using generics for io.ReadFullToBlocks() does not help in avoiding the heap allocation. With this change, we avoid at least these allocations: - rw variable in kcov.Kcov.TaskWork(). - safemem.BlockSeqReader struct in in mm.MemoryManager.getPMAsInternalLocked(). - gr variable in fsutil.FileRangeSet.Fill(). - h variable in gofer.dentry.Translate(). - h variable in gofer.dentryReadWriter.ReadToBlocks(). Some of these are on hot paths for IO workloads. PiperOrigin-RevId: 609805848 --- pkg/safemem/io.go | 30 ++++++++++------------------- pkg/safemem/io_test.go | 4 ++-- pkg/sentry/fsutil/file_range_set.go | 4 ++-- pkg/sentry/kernel/kcov.go | 6 +++--- pkg/sentry/mm/pma.go | 9 +++++---- pkg/sentry/pgalloc/pgalloc.go | 16 +++++++-------- 6 files changed, 30 insertions(+), 39 deletions(-) diff --git a/pkg/safemem/io.go b/pkg/safemem/io.go index 9551ca853..ecbfdf39b 100644 --- a/pkg/safemem/io.go +++ b/pkg/safemem/io.go @@ -46,12 +46,13 @@ type Writer interface { WriteFromBlocks(srcs BlockSeq) (uint64, error) } -// ReadFullToBlocks repeatedly invokes r.ReadToBlocks until dsts.NumBytes() -// bytes have been read or ReadToBlocks returns an error. -func ReadFullToBlocks(r Reader, dsts BlockSeq) (uint64, error) { +// ReadFullToBlocks repeatedly invokes r until dsts.NumBytes() bytes have been +// read or r returns an error. Note that we avoid a Reader interface receiver +// to avoid heap allocation. +func ReadFullToBlocks(r ReaderFunc, dsts BlockSeq) (uint64, error) { var done uint64 for !dsts.IsEmpty() { - n, err := r.ReadToBlocks(dsts) + n, err := r(dsts) done += n if err != nil { return done, err @@ -61,12 +62,13 @@ func ReadFullToBlocks(r Reader, dsts BlockSeq) (uint64, error) { return done, nil } -// WriteFullFromBlocks repeatedly invokes w.WriteFromBlocks until -// srcs.NumBytes() bytes have been written or WriteFromBlocks returns an error. -func WriteFullFromBlocks(w Writer, srcs BlockSeq) (uint64, error) { +// WriteFullFromBlocks repeatedly invokes w until srcs.NumBytes() bytes have +// been written or w returns an error. Note that we avoid a Writer interface +// receiver to avoid heap allocation. +func WriteFullFromBlocks(w WriterFunc, srcs BlockSeq) (uint64, error) { var done uint64 for !srcs.IsEmpty() { - n, err := w.WriteFromBlocks(srcs) + n, err := w(srcs) done += n if err != nil { return done, err @@ -144,18 +146,6 @@ func (r ToIOReader) Read(dst []byte) (int, error) { return int(n), err } -// ToIOWriter implements io.Writer for a (safemem.)Writer. -type ToIOWriter struct { - Writer Writer -} - -// Write implements io.Writer.Write. -func (w ToIOWriter) Write(src []byte) (int, error) { - // io.Writer does not permit partial writes. - n, err := WriteFullFromBlocks(w.Writer, BlockSeqOf(BlockFromSafeSlice(src))) - return int(n), err -} - // FromIOReader implements Reader for an io.Reader by repeatedly invoking // io.Reader.Read until it returns an error or partial read. This is not // thread-safe. diff --git a/pkg/safemem/io_test.go b/pkg/safemem/io_test.go index 629741bee..db4cd6352 100644 --- a/pkg/safemem/io_test.go +++ b/pkg/safemem/io_test.go @@ -100,7 +100,7 @@ func TestSingleByteReader(t *testing.T) { func TestReadFullToBlocks(t *testing.T) { r := FromIOReader{singleByteReader{bytes.NewBufferString("foobar")}} dsts := makeBlocks(make([]byte, 3), make([]byte, 3)) - n, err := ReadFullToBlocks(r, BlockSeqFromSlice(dsts)) + n, err := ReadFullToBlocks(r.ReadToBlocks, BlockSeqFromSlice(dsts)) // ReadFullToBlocks should call into FromIOReader => singleByteReader // repeatedly until dsts is exhausted. if wantN := uint64(6); n != wantN || err != nil { @@ -187,7 +187,7 @@ func TestWriteFullToBlocks(t *testing.T) { srcs := makeBlocks([]byte("foo"), []byte("bar")) var dst bytes.Buffer w := FromIOWriter{singleByteWriter{&dst}} - n, err := WriteFullFromBlocks(w, BlockSeqFromSlice(srcs)) + n, err := WriteFullFromBlocks(w.WriteFromBlocks, BlockSeqFromSlice(srcs)) // WriteFullToBlocks should call into FromIOWriter => singleByteWriter // repeatedly until srcs is exhausted. if wantN := uint64(6); n != wantN || err != nil { diff --git a/pkg/sentry/fsutil/file_range_set.go b/pkg/sentry/fsutil/file_range_set.go index f6a8f1e62..41477b4a0 100644 --- a/pkg/sentry/fsutil/file_range_set.go +++ b/pkg/sentry/fsutil/file_range_set.go @@ -129,7 +129,7 @@ func (s *FileRangeSet) Fill(ctx context.Context, required, optional memmap.Mappa MemCgID: memCgID, } if readAt != nil { - opts.Reader = safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) { + opts.ReaderFunc = func(dsts safemem.BlockSeq) (uint64, error) { var done uint64 for !dsts.IsEmpty() { n, err := func() (uint64, error) { @@ -167,7 +167,7 @@ func (s *FileRangeSet) Fill(ctx context.Context, required, optional memmap.Mappa } } return done, nil - }) + } } fr, err := mf.Allocate(gr.Length(), opts) diff --git a/pkg/sentry/kernel/kcov.go b/pkg/sentry/kernel/kcov.go index 27c3515e0..19b81ad7d 100644 --- a/pkg/sentry/kernel/kcov.go +++ b/pkg/sentry/kernel/kcov.go @@ -99,7 +99,7 @@ func (kcov *Kcov) TaskWork(t *Task) { } // Read in the PC count. - if _, err := safemem.ReadFullToBlocks(rw, kcov.countBlock()); err != nil { + if _, err := safemem.ReadFullToBlocks(rw.ReadToBlocks, kcov.countBlock()); err != nil { panic(fmt.Sprintf("Internal error reading count from kcov area: %v", err)) } @@ -111,7 +111,7 @@ func (kcov *Kcov) TaskWork(t *Task) { // output. kcov.count += uint64(n / 8) rw.off = 0 - if _, err := safemem.WriteFullFromBlocks(rw, kcov.countBlock()); err != nil { + if _, err := safemem.WriteFullFromBlocks(rw.WriteFromBlocks, kcov.countBlock()); err != nil { panic(fmt.Sprintf("Internal error writing count to kcov area: %v", err)) } @@ -337,6 +337,6 @@ type kcovIOWriter struct { // Write implements io.Writer.Write. func (w *kcovIOWriter) Write(p []byte) (int, error) { bs := safemem.BlockSeqOf(safemem.BlockFromSafeSlice(p)) - n, err := safemem.WriteFullFromBlocks(w.rw, bs) + n, err := safemem.WriteFullFromBlocks(w.rw.WriteFromBlocks, bs) return int(n), err } diff --git a/pkg/sentry/mm/pma.go b/pkg/sentry/mm/pma.go index 81d06a01e..8a1e02a3c 100644 --- a/pkg/sentry/mm/pma.go +++ b/pkg/sentry/mm/pma.go @@ -371,11 +371,12 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter return pstart, pseg.PrevGap(), err } // Copy contents. + reader := safemem.BlockSeqReader{Blocks: mm.internalMappingsLocked(pseg, copyAR)} fr, err := mm.mf.Allocate(uint64(copyAR.Length()), pgalloc.AllocOpts{ - Kind: usage.Anonymous, - Mode: pgalloc.AllocateAndWritePopulate, - MemCgID: memCgID, - Reader: &safemem.BlockSeqReader{mm.internalMappingsLocked(pseg, copyAR)}, + Kind: usage.Anonymous, + Mode: pgalloc.AllocateAndWritePopulate, + MemCgID: memCgID, + ReaderFunc: reader.ReadToBlocks, }) if _, ok := err.(safecopy.BusError); ok { // If we got SIGBUS during the copy, deliver SIGBUS to diff --git a/pkg/sentry/pgalloc/pgalloc.go b/pkg/sentry/pgalloc/pgalloc.go index 186b87e92..ee72c3f9f 100644 --- a/pkg/sentry/pgalloc/pgalloc.go +++ b/pkg/sentry/pgalloc/pgalloc.go @@ -465,12 +465,12 @@ type AllocOpts struct { // that will fill the allocated memory by invoking host system calls should // pass AllocateOnly. Mode AllocationMode - // If Reader is provided, the allocated memory is filled by calling - // ReadToBlocks() repeatedly until either length bytes are read or a non-nil - // error is returned. It returns the allocated memory, truncated down to the - // nearest page. If this is shorter than length bytes due to an error - // returned by ReadToBlocks(), it returns the partially filled fr and error. - Reader safemem.Reader + // If ReaderFunc is provided, the allocated memory is filled by calling it + // repeatedly until either length bytes are read or a non-nil error is + // returned. It returns the allocated memory, truncated down to the nearest + // page. If this is shorter than length bytes due to an error returned by + // ReaderFunc, it returns the partially filled fr and error. + ReaderFunc safemem.ReaderFunc } // Allocate returns a range of initially-zeroed pages of the given length with @@ -514,7 +514,7 @@ func (f *MemoryFile) Allocate(length uint64, opts AllocOpts) (memmap.FileRange, default: panic(fmt.Sprintf("unknown allocation mode: %d", opts.Mode)) } - if opts.Reader != nil { + if opts.ReaderFunc != nil { if dsts.IsEmpty() { dsts, err = f.MapInternal(fr, hostarch.Write) if err != nil { @@ -522,7 +522,7 @@ func (f *MemoryFile) Allocate(length uint64, opts AllocOpts) (memmap.FileRange, return memmap.FileRange{}, err } } - n, err := safemem.ReadFullToBlocks(opts.Reader, dsts) + n, err := safemem.ReadFullToBlocks(opts.ReaderFunc, dsts) un := uint64(hostarch.Addr(n).RoundDown()) if un < length { // Free unused memory and update fr to contain only the memory that is