Add pgalloc.DelayedEvictionManual.

PiperOrigin-RevId: 247667272
Change-Id: I16b04e11bb93f50b7e05e888992303f730e4a877

3 files changed, 24 insertions, 13 deletions

diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index 91889b573..85d73ace2 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/pkg/sentry/kernel/kernel.go
@@ -304,10 +304,11 @@ func (k *Kernel) SaveTo(w io.Writer) error {
 	defer k.resumeTimeLocked()
 
 	// Evict all evictable MemoryFile allocations.
-	k.mf.FlushEvictions()
+	k.mf.StartEvictions()
+	k.mf.WaitForEvictions()
 
 	// Flush write operations on open files so data reaches backing storage.
-	// This must come after k.mf.FlushEvictions() since eviction may cause file
+	// This must come after MemoryFile eviction since eviction may cause file
 	// writes.
 	if err := k.tasks.flushWritesToFiles(ctx); err != nil {
 		return err
diff --git a/pkg/sentry/pgalloc/pgalloc.go b/pkg/sentry/pgalloc/pgalloc.go
index 9c1313f6f..2b9924ad7 100644
--- a/pkg/sentry/pgalloc/pgalloc.go
+++ b/pkg/sentry/pgalloc/pgalloc.go
@@ -190,6 +190,11 @@ const (
 	// reclaimer goroutine is out of work (pages to reclaim), then evicts all
 	// pending evictable allocations immediately.
 	DelayedEvictionEnabled
+
+	// DelayedEvictionManual requires that evictable allocations are only
+	// evicted when MemoryFile.StartEvictions() is called. This is extremely
+	// dangerous outside of tests.
+	DelayedEvictionManual
 )
 
 // usageInfo tracks usage information.
@@ -264,7 +269,7 @@ func NewMemoryFile(file *os.File, opts MemoryFileOpts) (*MemoryFile, error) {
 	switch opts.DelayedEviction {
 	case DelayedEvictionDefault:
 		opts.DelayedEviction = DelayedEvictionEnabled
-	case DelayedEvictionDisabled, DelayedEvictionEnabled:
+	case DelayedEvictionDisabled, DelayedEvictionEnabled, DelayedEvictionManual:
 	default:
 		return nil, fmt.Errorf("invalid MemoryFileOpts.DelayedEviction: %v", opts.DelayedEviction)
 	}
@@ -1075,6 +1080,14 @@ func (f *MemoryFile) markReclaimed(fr platform.FileRange) {
 	}
 }
 
+// StartEvictions requests that f evict all evictable allocations. It does not
+// wait for eviction to complete; for this, see MemoryFile.WaitForEvictions.
+func (f *MemoryFile) StartEvictions() {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.startEvictionsLocked()
+}
+
 // Preconditions: f.mu must be locked.
 func (f *MemoryFile) startEvictionsLocked() {
 	for user, info := range f.evictable {
@@ -1122,6 +1135,12 @@ func (f *MemoryFile) startEvictionGoroutineLocked(user EvictableMemoryUser, info
 	}()
 }
 
+// WaitForEvictions blocks until f is no longer evicting any evictable
+// allocations.
+func (f *MemoryFile) WaitForEvictions() {
+	f.evictionWG.Wait()
+}
+
 type usageSetFunctions struct{}
 
 func (usageSetFunctions) MinKey() uint64 {
diff --git a/pkg/sentry/pgalloc/save_restore.go b/pkg/sentry/pgalloc/save_restore.go
index 9534d1aed..d4ba384b1 100644
--- a/pkg/sentry/pgalloc/save_restore.go
+++ b/pkg/sentry/pgalloc/save_restore.go
@@ -28,15 +28,6 @@ import (
 	"gvisor.googlesource.com/gvisor/pkg/state"
 )
 
-// FlushEvictions blocks until f has finished evicting all evictable
-// allocations.
-func (f *MemoryFile) FlushEvictions() {
-	f.mu.Lock()
-	f.startEvictionsLocked()
-	f.mu.Unlock()
-	f.evictionWG.Wait()
-}
-
 // SaveTo writes f's state to the given stream.
 func (f *MemoryFile) SaveTo(w io.Writer) error {
 	// Wait for reclaim.
@@ -51,7 +42,7 @@ func (f *MemoryFile) SaveTo(w io.Writer) error {
 
 	// Ensure that there are no pending evictions.
 	if len(f.evictable) != 0 {
-		panic(fmt.Sprintf("evictions still pending for %d users; call FlushEvictions before SaveTo", len(f.evictable)))
+		panic(fmt.Sprintf("evictions still pending for %d users; call StartEvictions and WaitForEvictions before SaveTo", len(f.evictable)))
 	}
 
 	// Ensure that all pages that contain data have knownCommitted set, since