netbird/gvisor
0
0
Fork 0
mirror of https://github.com/netbirdio/gvisor.git synced 2026-05-22 17:12:49 -07:00
Code Issues Packages Projects Releases Wiki Activity

Plumb safemem.ReaderFunc through pgalloc.Allocate() to avoid heap allocations.

Browse Source 
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
This commit is contained in:
Ayush Ranjan
2024-02-23 12:32:25 -08:00
committed by gVisor bot
parent 5b6d6e6044
commit 2759f79fce
6 changed files with 30 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/safemem/io.go
+10 -20
View File
@@ -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.
pkg/safemem/io_test.go
+2 -2
View File
@@ -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 {
pkg/sentry/fsutil/file_range_set.go
+2 -2
View File
@@ -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)
pkg/sentry/kernel/kcov.go
+3 -3
View File
@@ -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
}
pkg/sentry/mm/pma.go
+5 -4
View File
@@ -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
pkg/sentry/pgalloc/pgalloc.go
+8 -8
View File
@@ -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