pgalloc: integrate async page loading

When a pages file is provided to `runsc restore`, reads from that file are
asynchronous (via statefile.AsyncReader) in order to maximize throughput.
However, all such reads must complete before Kernel.LoadFrom() returns, so
applications cannot execute before MemoryFile loading is complete. The main
objective of this CL is to allow reads to continue after Kernel.LoadFrom()
returns, allowing applications to execute while MemoryFile loading is still in
progress. This behavior is user-visible: it affects whether deleting the pages
file frees disk space immediately on POSIX filesystems, may affect whether
deletion is possible on non-POSIX filesystems, and prevents unmounting
regardless. Thus it is flag-guarded as `runsc restore --background`.

MemoryFile ranges that have yet to be loaded, but that are being waited-for by
applications, should be prioritized over ranges for which no application is
waiting. This requires that application requests for data (calls to
MemoryFile.(memmap.File).DataFD/MapInternal()) are able to determine which
ranges have not yet been loaded, request reads for such ranges with elevated
priority, and wait for only those reads to be completed; none of these are
supported by the existing statefile.AsyncReader.

Thus:

- Add //pkg/sentry/pgalloc/aio, which provides an async I/O API that is
  designed to be easily implementable using a goroutine pool, Linux native AIO,
  or io_uring, though only includes a goroutine pool implementation. (io_uring
  is widely disabled due to security vulnerabilities. In my testing, Linux
  native AIO is slower than the goroutine pool, but this may change with lower
  GOMAXPROCS which needs further testing.)

- Move I/O scheduling into pgalloc: introduce an async page loader goroutine
  that is started by MemoryFile.LoadFrom() when async page loading is requested
  (implicitly, via the existence of a pages file), which is responsible for
  driving submission of read requests and handling their completions.

PiperOrigin-RevId: 679321884

pkg/aio/BUILD

@@ -0,0 +1,30 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+go_library(
+    name = "aio",
+    srcs = [
+        "aio.go",
+        "aio_unsafe.go",
+    ],
+    visibility = ["//pkg/sentry:internal"],
+    deps = [
+        "//pkg/sync",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+go_test(
+    name = "aio_test",
+    size = "small",
+    srcs = ["aio_test.go"],
+    library = ":aio",
+    deps = [
+        "//pkg/bitmap",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)

pkg/aio/aio.go

@@ -0,0 +1,185 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package aio provides asynchronous I/O on host file descriptors.
+package aio
+
+import (
+	"fmt"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
+// A Queue provides the ability to concurrently execute multiple read/write
+// operations on host file descriptors.
+//
+// Queues are not safe to use concurrently in multiple goroutines.
+type Queue interface {
+	// Destroy cancels all inflight operations and releases resources owned by
+	// the Queue. Destroy waits for cancelation, so the Queue will not access
+	// memory corresponding to inflight operations after Destroy returns.
+	Destroy()
+
+	// Cap returns the Queue's capacity, which is the maximum number of
+	// concurrent operations supported by the Queue.
+	Cap() int
+
+	// Add enqueues an inflight operation.
+	//
+	// Note that some Queue implementations may not begin execution of new
+	// Requests until the following call to Wait.
+	//
+	// Preconditions:
+	// - The current number of inflight operations < Cap().
+	Add(req Request)
+
+	// Wait blocks until at least minCompletions inflight operations have
+	// completed, then appends all completed inflight operations to cs and
+	// returns the updated slice.
+	//
+	// If Wait returns a non-nil error, no Queue methods may be subsequently
+	// called except Destroy.
+	//
+	// Preconditions:
+	// - 0 <= minCompletions <= Cap().
+	Wait(cs []Completion, minCompletions int) ([]Completion, error)
+}
+
+// Request is defined in aio_unsafe.go.
+
+// Op selects an asynchronous I/O operation.
+type Op uint8
+
+// Possible values for Request.Op.
+const (
+	// OpRead represents a read into addresses [Buf, Buf+Len).
+	OpRead Op = iota
+
+	// OpWrite represents a write from addresses [Buf, Buf+Len).
+	OpWrite
+
+	// OpReadv represents a read, where the destination addresses are given by
+	// the struct iovec array at address Buf of length Len.
+	OpReadv
+
+	// OpWritev represents a write, where the source addresses are given by the
+	// struct iovec array at address Buf of length Len.
+	OpWritev
+)
+
+// Completion provides outputs from an asynchronous I/O operation.
+type Completion struct {
+	ID     uint64 // copied from Request.ID in the corresponding Request
+	Result int64  // number of bytes or negative errno
+}
+
+// Err returns an error representing the Completion. If Err returns nil,
+// c.Result is the number of bytes completed by the operation.
+func (c Completion) Err() error {
+	if c.Result >= 0 {
+		return nil
+	}
+	return unix.Errno(-c.Result)
+}
+
+// GoQueue implements Queue using a pool of worker goroutines.
+type GoQueue struct {
+	requests    chan Request
+	completions chan Completion
+	shutdown    chan struct{}
+	workers     sync.WaitGroup
+}
+
+// NewGoQueue returns a new GoQueue with the given capacity.
+func NewGoQueue(cap int) *GoQueue {
+	q := &GoQueue{
+		requests:    make(chan Request, cap),
+		completions: make(chan Completion, cap),
+		shutdown:    make(chan struct{}),
+	}
+	q.workers.Add(cap)
+	for range cap {
+		go q.workerMain()
+	}
+	return q
+}
+
+func (q *GoQueue) workerMain() {
+	defer q.workers.Done()
+	for {
+		select {
+		case <-q.shutdown:
+			return
+		case r := <-q.requests:
+			var sysno uintptr
+			switch r.Op {
+			case OpRead:
+				sysno = unix.SYS_PREAD64
+			case OpWrite:
+				sysno = unix.SYS_PWRITE64
+			case OpReadv:
+				sysno = unix.SYS_PREADV2
+			case OpWritev:
+				sysno = unix.SYS_PWRITEV2
+			default:
+				panic(fmt.Sprintf("unknown op %v", r.Op))
+			}
+			n, _, e := unix.Syscall6(sysno, uintptr(r.FD), uintptr(r.Buf), uintptr(r.Len), uintptr(r.Off), 0 /* pos_h */, 0 /* flags/unused */)
+			c := Completion{
+				ID:     r.ID,
+				Result: int64(n),
+			}
+			if e != 0 {
+				c.Result = -int64(e)
+			}
+			q.completions <- c
+		}
+	}
+}
+
+// Destroy implements Queue.Destroy.
+func (q *GoQueue) Destroy() {
+	close(q.shutdown)
+	q.workers.Wait()
+}
+
+// Cap implements Queue.Cap.
+func (q *GoQueue) Cap() int {
+	return cap(q.requests)
+}
+
+// Add implements Queue.Add.
+func (q *GoQueue) Add(r Request) {
+	q.requests <- r
+}
+
+// Wait implements Queue.Wait.
+func (q *GoQueue) Wait(cs []Completion, minCompletions int) ([]Completion, error) {
+	i := 0
+	for {
+		if i < minCompletions {
+			cs = append(cs, <-q.completions)
+			i++
+		} else {
+			select {
+			case c := <-q.completions:
+				cs = append(cs, c)
+				i++
+			default:
+				return cs, nil
+			}
+		}
+	}
+}

pkg/aio/aio_test.go

@@ -0,0 +1,280 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package aio
+
+import (
+	"bytes"
+	"io"
+	"math/rand"
+	"os"
+	"testing"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/bitmap"
+)
+
+func TestRead(t *testing.T) {
+	// Create a temp file.
+	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_read")
+	if err != nil {
+		t.Fatalf("failed to create temp file: %v", err)
+	}
+	defer testFile.Close()
+	defer os.Remove(testFile.Name())
+
+	// Create random data.
+	const chunkSize = 4096
+	const dataLen = 1024 * chunkSize
+	data := make([]byte, dataLen)
+	_, _ = rand.Read(data)
+
+	// Write data to the file using sync writes.
+	if _, err := testFile.Write(data); err != nil {
+		t.Fatalf("failed to write temp file: %v", err)
+	}
+
+	// Read data from the file using async reads.
+	q := NewGoQueue(8)
+	defer q.Destroy()
+	qavail := q.Cap()
+	off := int64(0)
+	fd := int32(testFile.Fd())
+	buf := make([]byte, dataLen)
+	added := 0
+	done := 0
+	var cs []Completion
+	for done < dataLen {
+		for qavail > 0 && added < dataLen {
+			Read(q, 0 /* id */, fd, off, buf[added:added+chunkSize])
+			qavail--
+			off += chunkSize
+			added += chunkSize
+		}
+		cs, err := q.Wait(cs[:0], 1 /* minCompletions */)
+		if err != nil {
+			t.Fatalf("Queue.Wait failed: %v", err)
+		}
+		for _, c := range cs {
+			if err := c.Err(); err != nil {
+				t.Fatalf("Queue returned completion with error: %v", err)
+			}
+			if c.Result != chunkSize {
+				t.Fatalf("Queue returned completion of %d bytes, want %d", c.Result, chunkSize)
+			}
+			qavail++
+			done += chunkSize
+		}
+	}
+	if bytes.Compare(data, buf) != 0 {
+		t.Errorf("bytes differ")
+	}
+}
+
+func TestReadv(t *testing.T) {
+	// Create a temp file.
+	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_readv")
+	if err != nil {
+		t.Fatalf("failed to create temp file: %v", err)
+	}
+	defer testFile.Close()
+	defer os.Remove(testFile.Name())
+
+	// Create random data.
+	const chunkSize = 4096
+	const dataLen = 1024 * chunkSize
+	data := make([]byte, dataLen)
+	_, _ = rand.Read(data)
+
+	// Write data to the file using sync writes.
+	if _, err := testFile.Write(data); err != nil {
+		t.Fatalf("failed to write temp file: %v", err)
+	}
+
+	// Read data from the file using async vectored reads.
+	q := NewGoQueue(8)
+	defer q.Destroy()
+	qavail := q.Cap()
+	iovecsData := make([][2]unix.Iovec, qavail)
+	iovecsBusy := bitmap.New(uint32(qavail))
+	off := int64(0)
+	fd := int32(testFile.Fd())
+	buf := make([]byte, dataLen)
+	added := 0
+	done := 0
+	var cs []Completion
+	for done < dataLen {
+		for qavail > 0 && added < dataLen {
+			id, err := iovecsBusy.FirstZero(0)
+			if err != nil {
+				t.Fatalf("all iovecs busy with qavail=%d", qavail)
+			}
+			iovecsBusy.Add(id)
+			iovecs := &iovecsData[id]
+			iovecs[0].Base = &buf[added]
+			iovecs[0].Len = chunkSize
+			iovecs[1].Base = &buf[added+chunkSize]
+			iovecs[1].Len = chunkSize
+			Readv(q, uint64(id), fd, off, iovecs[:])
+			qavail--
+			off += 2 * chunkSize
+			added += 2 * chunkSize
+		}
+		cs, err := q.Wait(cs[:0], 1 /* minCompletions */)
+		if err != nil {
+			t.Fatalf("Queue.Wait failed: %v", err)
+		}
+		for _, c := range cs {
+			if err := c.Err(); err != nil {
+				t.Fatalf("Queue returned completion with error: %v", err)
+			}
+			if c.Result != 2*chunkSize {
+				t.Fatalf("Queue returned completion of %d bytes, want %d", c.Result, 2*chunkSize)
+			}
+			qavail++
+			iovecsBusy.Remove(uint32(c.ID))
+			done += 2 * chunkSize
+		}
+	}
+	if bytes.Compare(data, buf) != 0 {
+		t.Errorf("bytes differ")
+	}
+}
+
+func TestWrite(t *testing.T) {
+	// Create a temp file.
+	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_write")
+	if err != nil {
+		t.Fatalf("failed to create temp file: %v", err)
+	}
+	defer testFile.Close()
+	defer os.Remove(testFile.Name())
+
+	// Create random data.
+	const chunkSize = 4096
+	const dataLen = 1024 * chunkSize
+	data := make([]byte, dataLen)
+	_, _ = rand.Read(data)
+
+	// Write data to the file using async writes.
+	q := NewGoQueue(8)
+	defer q.Destroy()
+	qavail := q.Cap()
+	off := int64(0)
+	fd := int32(testFile.Fd())
+	added := 0
+	done := 0
+	var cs []Completion
+	for done < dataLen {
+		for qavail > 0 && added < dataLen {
+			Write(q, 0 /* id */, fd, off, data[added:added+chunkSize])
+			qavail--
+			off += chunkSize
+			added += chunkSize
+		}
+		cs, err := q.Wait(cs[:0], 1 /* minCompletions */)
+		if err != nil {
+			t.Fatalf("Queue.Wait failed: %v", err)
+		}
+		for _, c := range cs {
+			if err := c.Err(); err != nil {
+				t.Fatalf("Queue returned completion with error: %v", err)
+			}
+			if c.Result != chunkSize {
+				t.Fatalf("Queue returned completion of %d bytes, want %d", c.Result, chunkSize)
+			}
+			qavail++
+			done += chunkSize
+		}
+	}
+
+	// Read data from the file using sync reads.
+	buf := make([]byte, dataLen)
+	if n, err := io.ReadFull(testFile, buf); err != nil {
+		t.Fatalf("failed to read temp file after %d bytes: %v", n, err)
+	}
+	if bytes.Compare(data, buf) != 0 {
+		t.Errorf("bytes differ")
+	}
+}
+
+func TestWritev(t *testing.T) {
+	// Create a temp file.
+	testFile, err := os.CreateTemp(t.TempDir(), "aio_test_writev")
+	if err != nil {
+		t.Fatalf("failed to create temp file: %v", err)
+	}
+	defer testFile.Close()
+	defer os.Remove(testFile.Name())
+
+	// Create random data.
+	const chunkSize = 4096
+	const dataLen = 1024 * chunkSize
+	data := make([]byte, dataLen)
+	_, _ = rand.Read(data)
+
+	// Write data to the file using async vectored writes.
+	q := NewGoQueue(8)
+	defer q.Destroy()
+	qavail := q.Cap()
+	iovecsData := make([][2]unix.Iovec, qavail)
+	iovecsBusy := bitmap.New(uint32(qavail))
+	off := int64(0)
+	fd := int32(testFile.Fd())
+	added := 0
+	done := 0
+	var cs []Completion
+	for done < dataLen {
+		for qavail > 0 && added < dataLen {
+			id, err := iovecsBusy.FirstZero(0)
+			if err != nil {
+				t.Fatalf("all iovecs busy with qavail=%d", qavail)
+			}
+			iovecsBusy.Add(id)
+			iovecs := &iovecsData[id]
+			iovecs[0].Base = &data[added]
+			iovecs[0].Len = chunkSize
+			iovecs[1].Base = &data[added+chunkSize]
+			iovecs[1].Len = chunkSize
+			Writev(q, uint64(id), fd, off, iovecs[:])
+			qavail--
+			off += 2 * chunkSize
+			added += 2 * chunkSize
+		}
+		cs, err := q.Wait(cs[:0], 1 /* minCompletions */)
+		if err != nil {
+			t.Fatalf("Queue.Wait failed: %v", err)
+		}
+		for _, c := range cs {
+			if err := c.Err(); err != nil {
+				t.Fatalf("Queue returned completion with error: %v", err)
+			}
+			if c.Result != 2*chunkSize {
+				t.Fatalf("Queue returned completion of %d bytes, want %d", c.Result, 2*chunkSize)
+			}
+			qavail++
+			iovecsBusy.Remove(uint32(c.ID))
+			done += 2 * chunkSize
+		}
+	}
+
+	// Read data from the file using sync reads.
+	buf := make([]byte, dataLen)
+	if n, err := io.ReadFull(testFile, buf); err != nil {
+		t.Fatalf("failed to read temp file after %d bytes: %v", n, err)
+	}
+	if bytes.Compare(data, buf) != 0 {
+		t.Errorf("bytes differ")
+	}
+}

pkg/aio/aio_unsafe.go

@@ -0,0 +1,93 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package aio
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+// Request provides inputs to an asynchronous I/O operation.
+type Request struct {
+	ID  uint64 // copied to Completion.ID in the corresponding Completion
+	Op  Op
+	FD  int32          // host file descriptor
+	Off int64          // offset into FD
+	Buf unsafe.Pointer // depends on Op
+	Len int            // depends on Op
+}
+
+// Read enqueues a read. The caller must ensure that the memory referred to by
+// dst remains valid until the read is complete.
+//
+// Preconditions: As for q.Add().
+func Read(q Queue, id uint64, fd int32, off int64, dst []byte) {
+	q.Add(Request{
+		ID:  id,
+		Op:  OpRead,
+		FD:  fd,
+		Off: off,
+		Buf: unsafe.Pointer(unsafe.SliceData(dst)),
+		Len: len(dst),
+	})
+}
+
+// Write enqueues a write. The caller must ensure that the memory referred to
+// by src remains valid until the write is complete.
+//
+// Preconditions: As for q.Add().
+func Write(q Queue, id uint64, fd int32, off int64, src []byte) {
+	q.Add(Request{
+		ID:  id,
+		Op:  OpWrite,
+		FD:  fd,
+		Off: off,
+		Buf: unsafe.Pointer(unsafe.SliceData(src)),
+		Len: len(src),
+	})
+}
+
+// Readv enqueues a vectored read. The caller must ensure that the struct iovec
+// array referred to by dst, and the memory that those struct iovecs refer to,
+// remain valid until the read is complete.
+//
+// Preconditions: As for q.Add().
+func Readv(q Queue, id uint64, fd int32, off int64, dst []unix.Iovec) {
+	q.Add(Request{
+		ID:  id,
+		Op:  OpReadv,
+		FD:  fd,
+		Off: off,
+		Buf: unsafe.Pointer(unsafe.SliceData(dst)),
+		Len: len(dst),
+	})
+}
+
+// Writev enqueues a vectored write. The caller must ensure that the struct
+// iovec array referred to by src, and the memory that those struct iovecs
+// refer to, remain valid until the write is complete.
+//
+// Preconditions: As for q.Add().
+func Writev(q Queue, id uint64, fd int32, off int64, src []unix.Iovec) {
+	q.Add(Request{
+		ID:  id,
+		Op:  OpWritev,
+		FD:  fd,
+		Off: off,
+		Buf: unsafe.Pointer(unsafe.SliceData(src)),
+		Len: len(src),
+	})
+}

pkg/sentry/kernel/BUILD

@@ -386,7 +386,6 @@ go_library(
         "//pkg/sentry/usage",
         "//pkg/sentry/vfs",
         "//pkg/state",
-        "//pkg/state/statefile",
         "//pkg/state/wire",
         "//pkg/sync",
         "//pkg/sync/locking",

pkg/sentry/kernel/kernel.go

@@ -75,7 +75,6 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/uniqueid"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/state"
-	"gvisor.dev/gvisor/pkg/state/statefile"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 )
@@ -597,7 +596,7 @@ func savePrivateMFs(ctx context.Context, w io.Writer, pw io.Writer, mfsToSave ma
 	return nil
 }
 
-func loadPrivateMFs(ctx context.Context, r io.Reader, pr *statefile.AsyncReader) error {
+func loadPrivateMFs(ctx context.Context, r io.Reader, opts *pgalloc.LoadOpts) error {
 	// Load the metadata.
 	var meta privateMemoryFileMetadata
 	if _, err := state.Load(ctx, r, &meta); err != nil {
@@ -614,7 +613,7 @@ func loadPrivateMFs(ctx context.Context, r io.Reader, pr *statefile.AsyncReader)
 		if !ok {
 			return fmt.Errorf("saved memory file for %q was not configured on restore", fsID)
 		}
-		if err := mf.LoadFrom(ctx, r, pr); err != nil {
+		if err := mf.LoadFrom(ctx, r, opts); err != nil {
 			return err
 		}
 	}
@@ -767,9 +766,16 @@ func (k *Kernel) invalidateUnsavableMappings(ctx context.Context) error {
 }
 
 // LoadFrom returns a new Kernel loaded from args.
-func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pagesFile *fd.FD, timeReady chan struct{}, net inet.Stack, clocks sentrytime.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
+//
+// LoadFrom takes ownership of pagesFile.
+func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pagesFile *fd.FD, background bool, timeReady chan struct{}, net inet.Stack, clocks sentrytime.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
 	loadStart := time.Now()
 
+	defer func() {
+		if pagesFile != nil {
+			pagesFile.Close()
+		}
+	}()
 	var (
 		mfLoadWg  sync.WaitGroup
 		mfLoadErr error
@@ -780,7 +786,8 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 		mfLoadWg.Add(1)
 		go func() {
 			defer mfLoadWg.Done()
-			mfLoadErr = k.loadMemoryFiles(ctx, r, pagesMetadata, pagesFile)
+			mfLoadErr = k.loadMemoryFiles(ctx, r, pagesMetadata, pagesFile, background)
+			pagesFile = nil // transferred to k.loadMemoryFiles()
 		}()
 		// Defer a Wait() so we wait for k.loadMemoryFiles() to complete even if we
 		// error out without reaching the other Wait() below.
@@ -824,7 +831,8 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 	if parallelMfLoad {
 		mfLoadWg.Wait()
 	} else {
-		mfLoadErr = k.loadMemoryFiles(ctx, r, pagesMetadata, pagesFile)
+		mfLoadErr = k.loadMemoryFiles(ctx, r, pagesMetadata, pagesFile, background)
+		pagesFile = nil // transferred to k.loadMemoryFiles()
 	}
 	if mfLoadErr != nil {
 		return mfLoadErr
@@ -869,28 +877,40 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 	return nil
 }
 
-func (k *Kernel) loadMemoryFiles(ctx context.Context, r, pagesMetadata io.Reader, pagesFile *fd.FD) error {
-	// Load the memory files' state.
+// loadMemoryFiles takes ownership of pagesFile.
+func (k *Kernel) loadMemoryFiles(ctx context.Context, r, pagesMetadata io.Reader, pagesFile *fd.FD, background bool) error {
 	memoryStart := time.Now()
 	pmr := r
 	if pagesMetadata != nil {
 		pmr = pagesMetadata
 	}
-	var pr *statefile.AsyncReader
-	if pagesFile != nil {
-		pr = statefile.NewAsyncReader(pagesFile, 0 /* off */)
-		defer pr.Close()
+	var (
+		pagesFileUsers  atomicbitops.Int64
+		asyncPageLoadWG sync.WaitGroup
+	)
+	opts := pgalloc.LoadOpts{
+		PagesFile: pagesFile,
+		OnAsyncPageLoadStart: func() {
+			pagesFileUsers.Add(1)
+			asyncPageLoadWG.Add(1)
+		},
+		OnAsyncPageLoadDone: func(error) {
+			if n := pagesFileUsers.Add(-1); n == 0 {
+				pagesFile.Close()
+			} else if n < 0 {
+				panic("pagesFileUsers < 0")
+			}
+			asyncPageLoadWG.Done()
+		},
 	}
-	if err := k.mf.LoadFrom(ctx, pmr, pr); err != nil {
+	if err := k.mf.LoadFrom(ctx, pmr, &opts); err != nil {
 		return err
 	}
-	if err := loadPrivateMFs(ctx, pmr, pr); err != nil {
+	if err := loadPrivateMFs(ctx, pmr, &opts); err != nil {
 		return err
 	}
-	if pr != nil {
-		if err := pr.Wait(); err != nil {
-			return err
-		}
+	if !background {
+		asyncPageLoadWG.Wait()
 	}
 	log.Infof("Memory files load took [%s].", time.Since(memoryStart))
 	return nil

pkg/sentry/pgalloc/BUILD

@@ -8,10 +8,10 @@ package(
 )
 
 declare_mutex(
-    name = "mappings_mutex",
-    out = "mappings_mutex.go",
+    name = "apl_shared_mutex",
+    out = "apl_shared_mutex.go",
     package = "pgalloc",
-    prefix = "mappings",
+    prefix = "aplShared",
 )
 
 declare_mutex(
@@ -21,6 +21,23 @@ declare_mutex(
     prefix = "memoryFile",
 )
 
+go_template_instance(
+    name = "apl_unloaded_set",
+    out = "apl_unloaded_set.go",
+    imports = {
+        "memmap": "gvisor.dev/gvisor/pkg/sentry/memmap",
+    },
+    package = "pgalloc",
+    prefix = "aplUnloaded",
+    template = "//pkg/segment:generic_set",
+    types = {
+        "Key": "uint64",
+        "Range": "memmap.FileRange",
+        "Value": "aplUnloadedInfo",
+        "Functions": "aplUnloadedSetFunctions",
+    },
+)
+
 go_template_instance(
     name = "evictable_range",
     out = "evictable_range.go",
@@ -51,7 +68,6 @@ go_template_instance(
     out = "memacct_set.go",
     consts = {
         "minDegree": "10",
-        "trackGaps": "1",
     },
     imports = {
         "memmap": "gvisor.dev/gvisor/pkg/sentry/memmap",
@@ -112,10 +128,12 @@ go_template_instance(
 go_library(
     name = "pgalloc",
     srcs = [
+        "apl_shared_mutex.go",
+        "apl_unloaded_set.go",
         "context.go",
+        "debug.go",
         "evictable_range.go",
         "evictable_range_set.go",
-        "mappings_mutex.go",
         "memacct_set.go",
         "memory_file_mutex.go",
         "pgalloc.go",
@@ -127,21 +145,26 @@ go_library(
     visibility = ["//pkg/sentry:internal"],
     deps = [
         "//pkg/abi/linux",
+        "//pkg/aio",
         "//pkg/atomicbitops",
+        "//pkg/bitmap",
         "//pkg/context",
         "//pkg/errors/linuxerr",
+        "//pkg/fd",
+        "//pkg/goid",
         "//pkg/hostarch",
         "//pkg/log",
+        "//pkg/ringdeque",
         "//pkg/safemem",
         "//pkg/sentry/arch",
         "//pkg/sentry/hostmm",
         "//pkg/sentry/memmap",
         "//pkg/sentry/usage",
         "//pkg/state",
-        "//pkg/state/statefile",
         "//pkg/state/wire",
         "//pkg/sync",
         "//pkg/sync/locking",
+        "//pkg/syncevent",
         "//pkg/usermem",
         "@org_golang_x_sys//unix:go_default_library",
     ],

pkg/sentry/pgalloc/debug.go

@@ -0,0 +1,19 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pgalloc
+
+// If logAwaitedLoads is true, events relevant to awaited async page loads will
+// be logged at level Info.
+const logAwaitedLoads = false

pkg/sentry/pgalloc/pgalloc.go

@@ -183,6 +183,10 @@ type MemoryFile struct {
 	// immutable.
 	stopNotifyPressure func()
 
+	// If asyncPageLoad is non-nil, it tracks the state of in-progress or
+	// failed async page loading.
+	asyncPageLoad atomic.Pointer[aplShared]
+
 	// file is the backing file. The file pointer is immutable.
 	file *os.File
 
@@ -1162,7 +1166,11 @@ func (f *MemoryFile) IncRef(fr memmap.FileRange, memCgID uint32) {
 
 	f.mu.Lock()
 	defer f.mu.Unlock()
+	f.incRefLocked(fr)
+}
 
+// Preconditions: f.mu must be locked.
+func (f *MemoryFile) incRefLocked(fr memmap.FileRange) {
 	f.forEachChunk(fr, func(chunk *chunkInfo, chunkFR memmap.FileRange) bool {
 		unfree := &f.unfreeSmall
 		if chunk.huge {
@@ -1219,6 +1227,10 @@ func (f *MemoryFile) DecRef(fr memmap.FileRange) {
 					ma.wasteOrReleasing = true
 					return true
 				})
+				// Cancel any pending async load on waste pages.
+				if apl := f.asyncPageLoad.Load(); apl != nil {
+					apl.cancelWasteLoad(wasteFR)
+				}
 			}
 			return true
 		})
@@ -1418,6 +1430,12 @@ func (f *MemoryFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (s
 		return safemem.BlockSeq{}, linuxerr.EACCES
 	}
 
+	if apl := f.asyncPageLoad.Load(); apl != nil {
+		if err := apl.awaitLoad(f, fr); err != nil {
+			return safemem.BlockSeq{}, err
+		}
+	}
+
 	chunks := ((fr.End + chunkMask) / chunkSize) - (fr.Start / chunkSize)
 	if chunks == 1 {
 		// Avoid an unnecessary slice allocation.
@@ -1753,6 +1771,11 @@ func (f *MemoryFile) File() *os.File {
 
 // DataFD implements memmap.File.DataFD.
 func (f *MemoryFile) DataFD(fr memmap.FileRange) (int, error) {
+	if apl := f.asyncPageLoad.Load(); apl != nil {
+		if err := apl.awaitLoad(f, fr); err != nil {
+			return -1, err
+		}
+	}
 	return f.FD(), nil
 }
 

pkg/sentry/pgalloc/pgalloc_unsafe.go

@@ -36,3 +36,11 @@ func mincore(s []byte, buf []byte, off uint64, wasCommitted bool) error {
 	}
 	return nil
 }
+
+func sliceFromIovec(iov unix.Iovec) []byte {
+	return unsafe.Slice(iov.Base, iov.Len)
+}
+
+func canMergeIovecAndSlice(iov unix.Iovec, bs []byte) bool {
+	return uintptr(unsafe.Pointer(iov.Base))+uintptr(iov.Len) == uintptr(unsafe.Pointer(unsafe.SliceData(bs)))
+}

pkg/sentry/state/state.go

@@ -144,12 +144,25 @@ type LoadOpts struct {
 	// PagesFile is non-nil. Otherwise this content is stored in Source.
 	PagesFile *fd.FD
 
+	// If Background is true, the sentry may read from PagesFile after Load has
+	// returned.
+	Background bool
+
 	// Key is used for state integrity check.
 	Key []byte
 }
 
 // Load loads the given kernel, setting the provided platform and stack.
+//
+// Load takes ownership of (and unsets) opts.PagesFile.
 func (opts LoadOpts) Load(ctx context.Context, k *kernel.Kernel, timeReady chan struct{}, n inet.Stack, clocks time.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
+	defer func() {
+		if opts.PagesFile != nil {
+			opts.PagesFile.Close()
+			opts.PagesFile = nil
+		}
+	}()
+
 	// Open the file.
 	r, m, err := statefile.NewReader(opts.Source, opts.Key)
 	if err != nil {
@@ -167,5 +180,7 @@ func (opts LoadOpts) Load(ctx context.Context, k *kernel.Kernel, timeReady chan
 	previousMetadata = m
 
 	// Restore the Kernel object graph.
-	return k.LoadFrom(ctx, r, pagesMetadata, opts.PagesFile, timeReady, n, clocks, vfsOpts, saveRestoreNet)
+	err = k.LoadFrom(ctx, r, pagesMetadata, opts.PagesFile, opts.Background, timeReady, n, clocks, vfsOpts, saveRestoreNet)
+	opts.PagesFile = nil // transferred to k.LoadFrom()
+	return err
 }

pkg/shim/extension/extension.go

@@ -42,8 +42,9 @@ type Process interface {
 
 // RestoreConfig is the configuration for a restore request.
 type RestoreConfig struct {
-	ImagePath string
-	Direct    bool
+	ImagePath  string
+	Direct     bool
+	Background bool
 }
 
 // TaskServiceExt extends TaskRequest with extra functionality required by the shim.

pkg/shim/proc/init.go

@@ -240,9 +240,10 @@ func (p *Init) start(ctx context.Context, restoreConf *extension.RestoreConfig)
 		}
 	} else {
 		if err := p.runtime.Restore(ctx, p.id, cio, &runsccmd.RestoreOpts{
-			ImagePath: restoreConf.ImagePath,
-			Detach:    true,
-			Direct:    restoreConf.Direct,
+			ImagePath:  restoreConf.ImagePath,
+			Detach:     true,
+			Direct:     restoreConf.Direct,
+			Background: restoreConf.Background,
 		}); err != nil {
 			return p.runtimeError(err, "OCI runtime restore failed")
 		}

pkg/shim/runsccmd/runsc.go

@@ -233,9 +233,10 @@ func (r *Runsc) start(context context.Context, cio runc.IO, cmd *exec.Cmd) error
 
 // RestoreOpts is a set of options to runsc.Restore().
 type RestoreOpts struct {
-	ImagePath string
-	Detach    bool
-	Direct    bool
+	ImagePath  string
+	Detach     bool
+	Direct     bool
+	Background bool
 }
 
 func (o *RestoreOpts) args() []string {
@@ -249,6 +250,9 @@ func (o *RestoreOpts) args() []string {
 	if o.Direct {
 		out = append(out, "--direct")
 	}
+	if o.Background {
+		out = append(out, "--background")
+	}
 	return out
 }
 

pkg/state/statefile/BUILD

@@ -8,28 +8,18 @@ package(
 go_library(
     name = "statefile",
     srcs = [
-        "async_io.go",
         "statefile.go",
     ],
     visibility = ["//:sandbox"],
-    deps = [
-        "//pkg/compressio",
-        "//pkg/fd",
-        "//pkg/sync",
-    ],
+    deps = ["//pkg/compressio"],
 )
 
 go_test(
     name = "statefile_test",
     size = "small",
     srcs = [
-        "async_io_test.go",
         "statefile_test.go",
     ],
     library = ":statefile",
-    deps = [
-        "//pkg/compressio",
-        "//pkg/fd",
-        "@org_golang_x_sys//unix:go_default_library",
-    ],
+    deps = ["//pkg/compressio"],
 )

pkg/state/statefile/async_io.go

@@ -1,94 +0,0 @@
-// Copyright 2024 The gVisor Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package statefile
-
-import (
-	"runtime"
-	"sync/atomic"
-
-	"gvisor.dev/gvisor/pkg/fd"
-	"gvisor.dev/gvisor/pkg/sync"
-)
-
-type chunk struct {
-	dst []byte
-	off int64
-}
-
-// AsyncReader can be used to do reads asynchronously. It does not change the
-// underlying file's offset.
-type AsyncReader struct {
-	// in is the backing file which contains all pages.
-	in *fd.FD
-	// off is the offset being read.
-	off int64
-	// q is the work queue.
-	q chan chunk
-	// err stores the latest IO error that occured during async read.
-	err atomic.Pointer[error]
-	// wg tracks all in flight work.
-	wg sync.WaitGroup
-}
-
-// NewAsyncReader initializes a new AsyncReader.
-func NewAsyncReader(in *fd.FD, off int64) *AsyncReader {
-	workers := runtime.GOMAXPROCS(0)
-	r := &AsyncReader{
-		in:  in,
-		off: off,
-		q:   make(chan chunk, workers),
-	}
-	for i := 0; i < workers; i++ {
-		go r.work()
-	}
-	return r
-}
-
-// ReadAsync schedules a read of len(p) bytes from current offset into p.
-func (r *AsyncReader) ReadAsync(p []byte) {
-	r.wg.Add(1)
-	r.q <- chunk{off: r.off, dst: p}
-	r.off += int64(len(p))
-}
-
-// Wait blocks until all in flight work is complete and then returns any IO
-// errors that occurred since the last call to Wait().
-func (r *AsyncReader) Wait() error {
-	r.wg.Wait()
-	if err := r.err.Swap(nil); err != nil {
-		return *err
-	}
-	return nil
-}
-
-// Close calls Wait() and additionally cleans up all worker goroutines.
-func (r *AsyncReader) Close() error {
-	err := r.Wait()
-	close(r.q)
-	return err
-}
-
-func (r *AsyncReader) work() {
-	for {
-		c := <-r.q
-		if c.dst == nil {
-			return
-		}
-		if _, err := r.in.ReadAt(c.dst, c.off); err != nil {
-			r.err.Store(&err)
-		}
-		r.wg.Done()
-	}
-}

pkg/state/statefile/async_io_test.go

@@ -1,64 +0,0 @@
-// Copyright 2024 The gVisor Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package statefile
-
-import (
-	"bytes"
-	"math/rand"
-	"os"
-	"testing"
-
-	"golang.org/x/sys/unix"
-	"gvisor.dev/gvisor/pkg/fd"
-)
-
-func TestAsyncReader(t *testing.T) {
-	// Create random data.
-	const chunkSize = 4096
-	const dataLen = 1024 * chunkSize
-	data := make([]byte, dataLen)
-	_, _ = rand.Read(data)
-
-	// Create a temp file with the data.
-	testFile, err := os.CreateTemp(t.TempDir(), "source")
-	if err != nil {
-		t.Fatalf("failed to create temp source file: %v", err)
-	}
-	if _, err := testFile.Write(data); err != nil {
-		t.Fatalf("failed to write temp source file: %v", err)
-	}
-	testFilePath := testFile.Name()
-	if err := testFile.Close(); err != nil {
-		t.Fatalf("failed to close temp source file: %v", err)
-	}
-
-	// Read the data from the file using async reads.
-	sourceFD, err := fd.Open(testFilePath, unix.O_RDONLY, 0)
-	if err != nil {
-		t.Fatalf("failed to open source file %q: %v", testFilePath, err)
-	}
-	ar := NewAsyncReader(sourceFD, 0 /* off */)
-	defer ar.Close()
-	p := make([]byte, dataLen)
-	for i := 0; i < dataLen; i += chunkSize {
-		ar.ReadAsync(p[i : i+chunkSize])
-	}
-	if err := ar.Wait(); err != nil {
-		t.Fatalf("AsyncReader.Wait failed: %v", err)
-	}
-	if ret := bytes.Compare(p, data); ret != 0 {
-		t.Errorf("bytes differ")
-	}
-}

runsc/boot/controller.go

@@ -488,6 +488,7 @@ type RestoreOpts struct {
 	urpc.FilePayload
 	HavePagesFile  bool
 	HaveDeviceFile bool
+	Background     bool
 }
 
 // Restore loads a container from a statefile.
@@ -524,7 +525,11 @@ func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 		return fmt.Errorf("statefile cannot be empty")
 	}
 
-	cm.restorer = &restorer{restoreDone: cm.onRestoreDone, stateFile: stateFile}
+	cm.restorer = &restorer{
+		restoreDone: cm.onRestoreDone,
+		stateFile:   stateFile,
+		background:  o.Background,
+	}
 	cm.l.restoreWaiters = sync.NewCond(&cm.l.mu)
 	cm.l.state = restoring
 	// Release `cm.l.mu`.

runsc/boot/restore.go

@@ -76,6 +76,10 @@ type restorer struct {
 	pagesMetadata *fd.FD
 	pagesFile     *fd.FD
 
+	// If background is true, pagesFile may continue to be read after
+	// restorer.restore() returns.
+	background bool
+
 	// deviceFile is the required to start the platform.
 	deviceFile *fd.FD
 
@@ -231,8 +235,15 @@ func (r *restorer) restore(l *Loader) error {
 	ctx = context.WithValue(ctx, devutil.CtxDevGoferClientProvider, l.k)
 
 	// Load the state.
-	loadOpts := state.LoadOpts{Source: r.stateFile, PagesMetadata: r.pagesMetadata, PagesFile: r.pagesFile}
-	if err := loadOpts.Load(ctx, l.k, nil, oldInetStack, time.NewCalibratedClocks(), &vfs.CompleteRestoreOptions{}, l.saveRestoreNet); err != nil {
+	loadOpts := state.LoadOpts{
+		Source:        r.stateFile,
+		PagesMetadata: r.pagesMetadata,
+		PagesFile:     r.pagesFile,
+		Background:    r.background,
+	}
+	err = loadOpts.Load(ctx, l.k, nil, oldInetStack, time.NewCalibratedClocks(), &vfs.CompleteRestoreOptions{}, l.saveRestoreNet)
+	r.pagesFile = nil // transferred to loadOpts.Load()
+	if err != nil {
 		return err
 	}