Add //pkg/sentry/stateio and use it for async page loading.

PiperOrigin-RevId: 803052657

Author: nixprime

pkg/abi/linux/fs.go

@@ -14,6 +14,12 @@
 
 package linux
 
+import (
+	"math"
+
+	"gvisor.dev/gvisor/pkg/hostarch"
+)
+
 // Filesystem types used in statfs(2).
 //
 // See linux/magic.h.
@@ -127,3 +133,8 @@ const (
 	WHITEOUT_MODE = 0
 	WHITEOUT_DEV  = 0
 )
+
+// MAX_RW_COUNT is the maximum size in bytes of a single read or write.
+// Reads and writes that exceed this size may be truncated.
+// (Linux: include/linux/fs.h:MAX_RW_COUNT)
+var MAX_RW_COUNT = int(hostarch.PageRoundDown(uint32(math.MaxInt32)))

pkg/sentry/kernel/BUILD

@@ -390,6 +390,7 @@ go_library(
         "//pkg/sentry/seccheck/points:points_go_proto",
         "//pkg/sentry/socket/netlink/port",
         "//pkg/sentry/socket/unix/transport",
+        "//pkg/sentry/state/stateio",
         "//pkg/sentry/time",
         "//pkg/sentry/unimpl",
         "//pkg/sentry/unimpl:unimplemented_syscall_go_proto",

pkg/sentry/kernel/kernel_restore.go

@@ -15,15 +15,14 @@
 package kernel
 
 import (
-	"bufio"
 	"fmt"
 	"io"
 
 	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/context"
-	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
+	"gvisor.dev/gvisor/pkg/sentry/state/stateio"
 	"gvisor.dev/gvisor/pkg/state"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/timing"
@@ -272,7 +271,7 @@ type AsyncMFLoader struct {
 // If timeline is provided, it will be used to track async page loading.
 // It takes ownership of the timeline, and will end it when done loading all
 // pages.
-func NewAsyncMFLoader(pagesMetadata, pagesFile *fd.FD, mainMF *pgalloc.MemoryFile, timeline *timing.Timeline) *AsyncMFLoader {
+func NewAsyncMFLoader(pagesMetadata io.ReadCloser, pagesFile stateio.AsyncReader, mainMF *pgalloc.MemoryFile, timeline *timing.Timeline) *AsyncMFLoader {
 	mfl := &AsyncMFLoader{
 		privateMFsChan: make(chan map[string]*pgalloc.MemoryFile, 1),
 	}
@@ -283,27 +282,25 @@ func NewAsyncMFLoader(pagesMetadata, pagesFile *fd.FD, mainMF *pgalloc.MemoryFil
 	return mfl
 }
 
-func (mfl *AsyncMFLoader) backgroundGoroutine(pagesMetadataFD, pagesFileFD *fd.FD, mainMF *pgalloc.MemoryFile, timeline *timing.Timeline) {
+func (mfl *AsyncMFLoader) backgroundGoroutine(pagesMetadata io.ReadCloser, pagesFile stateio.AsyncReader, mainMF *pgalloc.MemoryFile, timeline *timing.Timeline) {
 	defer timeline.End()
-	defer pagesMetadataFD.Close()
-	defer pagesFileFD.Close()
+	defer pagesMetadata.Close()
 	cu := cleanup.Make(func() {
 		mfl.metadataWg.Done()
 		mfl.loadWg.Done()
 	})
 	defer cu.Clean()
 
-	// //pkg/state/wire reads one byte at a time; buffer these reads to
-	// avoid making one syscall per read. For the "main" state file, this
-	// buffering is handled by statefile.NewReader() => compressio.Reader
-	// or compressio.NewSimpleReader().
-	pagesMetadata := bufio.NewReader(pagesMetadataFD)
-
 	mfl.loadWg.Add(1)
-	apfl := pgalloc.StartAsyncPagesFileLoad(int32(pagesFileFD.FD()), func(err error) {
+	apfl, err := pgalloc.StartAsyncPagesFileLoad(pagesFile, func(err error) {
 		defer mfl.loadWg.Done()
 		mfl.loadErr = err
-	}, timeline)
+	}, timeline) // transfers ownership of pagesFile
+	if err != nil {
+		mfl.loadWg.Done()
+		log.Warningf("Failed to start async page loading: %v", err)
+		return
+	}
 	cu.Add(apfl.MemoryFilesDone)
 
 	opts := pgalloc.LoadOpts{
@@ -314,7 +311,7 @@ func (mfl *AsyncMFLoader) backgroundGoroutine(pagesMetadataFD, pagesFileFD *fd.F
 	timeline.Reached("loading mainMF")
 	log.Infof("Loading metadata for main MemoryFile: %p", mainMF)
 	ctx := context.Background()
-	err := mainMF.LoadFrom(ctx, pagesMetadata, &opts)
+	err = mainMF.LoadFrom(ctx, pagesMetadata, &opts)
 	mfl.metadataErr = err
 	mfl.mainMetadataErr = err
 	mfl.mainMFStartWg.Done()

pkg/sentry/kernel/task_usermem.go

@@ -28,11 +28,6 @@ import (
 
 const iovecLength = 16
 
-// MAX_RW_COUNT is the maximum size in bytes of a single read or write.
-// Reads and writes that exceed this size may be silently truncated.
-// (Linux: include/linux/fs.h:MAX_RW_COUNT)
-var MAX_RW_COUNT = int(hostarch.Addr(math.MaxInt32).RoundDown())
-
 // Activate ensures that the task has an active address space.
 func (t *Task) Activate() {
 	if mm := t.MemoryManager(); mm != nil {
@@ -190,7 +185,7 @@ func copyInIovec(ctx marshal.CopyContext, t *Task, addr hostarch.Addr) (hostarch
 	if err != nil {
 		return hostarch.AddrRangeSeq{}, err
 	}
-	return hostarch.AddrRangeSeqOf(ar).TakeFirst(MAX_RW_COUNT), nil
+	return hostarch.AddrRangeSeqOf(ar).TakeFirst(linux.MAX_RW_COUNT), nil
 }
 
 // copyInIovecs copies an array of numIovecs struct iovecs from the memory
@@ -243,7 +238,7 @@ func copyInIovecs(ctx marshal.CopyContext, t *Task, addr hostarch.Addr, numIovec
 	var total uint64
 	for i := range dst {
 		dstlen := uint64(dst[i].Length())
-		if rem := uint64(MAX_RW_COUNT) - total; rem < dstlen {
+		if rem := uint64(linux.MAX_RW_COUNT) - total; rem < dstlen {
 			dst[i].End -= hostarch.Addr(dstlen - rem)
 			dstlen = rem
 		}
@@ -288,8 +283,8 @@ func makeIovec(ctx marshal.CopyContext, t *Task, addr hostarch.Addr, b []byte) (
 // access_ok() in fs/read_write.c:vfs_read/vfs_write, and overflowing address
 // ranges are truncated to MAX_RW_COUNT by fs/read_write.c:rw_verify_area().)
 func (t *Task) SingleIOSequence(addr hostarch.Addr, length int, opts usermem.IOOpts) (usermem.IOSequence, error) {
-	if length > MAX_RW_COUNT {
-		length = MAX_RW_COUNT
+	if length > linux.MAX_RW_COUNT {
+		length = linux.MAX_RW_COUNT
 	}
 	ar, ok := t.MemoryManager().CheckIORange(addr, int64(length))
 	if !ok {

pkg/sentry/pgalloc/BUILD

@@ -166,7 +166,6 @@ go_library(
     visibility = ["//pkg/sentry:internal"],
     deps = [
         "//pkg/abi/linux",
-        "//pkg/aio",
         "//pkg/atomicbitops",
         "//pkg/bitmap",
         "//pkg/context",
@@ -180,6 +179,7 @@ go_library(
         "//pkg/sentry/arch",
         "//pkg/sentry/hostmm",
         "//pkg/sentry/memmap",
+        "//pkg/sentry/state/stateio",
         "//pkg/sentry/usage",
         "//pkg/state",
         "//pkg/state/wire",