pgalloc: no-op MemoryFile.UpdateUsage() during saving

During checkpointing, MemoryFile.SaveTo() narrows the set of pages known by
memory accounting to be "committed" to those containing non-zero bytes, in
order to avoid saving zero pages and therefore bloating the checkpoint. In the
process of doing so, it needs to touch pages in order to determine whether they
contain non-zero bytes, and does so without holding MemoryFile.mu (in a
MemoryFile.updateUsageLocked() callback). Thus, a concurrent call to
MemoryFile.UpdateUsage() => MemoryFile.updateUsageLocked() can racily observe
that touched zero pages are committed (via mincore()) and mark them
known-committed accordingly, causing them to be unintentionally saved in the
checkpoint.

When SaveOpts.ExcludeCommittedZeroPages is set, MemoryFile.SaveTo() does not
decommit previously-known-committed zero pages, since doing so would cost time;
the motivation for decommitting zero pages is to avoid increasing (real, host)
memory usage during checkpointing, but previously-known-committed pages must
have been using memory even before being touched. However, this significantly
widens the race window described above, since any future call to
MemoryFile.UpdateUsage() will observe said pages to be committed and mark them
known-committed again, effectively negating SaveOpts.ExcludeCommittedZeroPages.

To fix this, inhibit MemoryFile.UpdateUsage() during MemoryFile.SaveTo(); in
essence, when MemoryFile.SaveTo() is in progress, it exclusively defines what
pages are known-committed.

PiperOrigin-RevId: 733876220

pkg/sentry/pgalloc/pgalloc.go

@@ -153,11 +153,15 @@ type MemoryFile struct {
 	// nextCommitScan is the next time at which UpdateUsage() may scan the
 	// backing file for commitment information.
 	//
+	// isSaving is non-zero during f.SaveTo() to prevent concurrent calls to
+	// f.UpdateUsage() from marking pages as committed.
+	//
 	// All of these fields are protected by mu.
 	memAcct             memAcctSet
 	knownCommittedBytes uint64
 	commitSeq           uint64
 	nextCommitScan      time.Time
+	isSaving            uint
 
 	// evictable maps EvictableMemoryUsers to eviction state.
 	//
@@ -1561,6 +1565,11 @@ func (f *MemoryFile) UpdateUsage(memCgIDs map[uint32]struct{}) error {
 		return nil
 	}
 
+	if f.isSaving != 0 {
+		log.Debugf("pgalloc.MemoryFile.UpdateUsage() inhibited during MemoryFile save")
+		return nil
+	}
+
 	// Linux updates usage values at CONFIG_HZ; throttle our scans to the same
 	// frequency.
 	startTime := time.Now()
@@ -1571,7 +1580,11 @@ func (f *MemoryFile) UpdateUsage(memCgIDs map[uint32]struct{}) error {
 		f.nextCommitScan = startTime.Add(time.Second / linux.CLOCKS_PER_SEC)
 	}
 
-	err = f.updateUsageLocked(memCgIDs, false /* alsoScanCommitted */, mincore)
+	err = f.updateUsageLocked(memCgIDs, false /* alsoScanCommitted */, false /* callerIsSaveTo */, mincore)
+	if _, ok := err.(updateUsageDuringSaveErr); ok {
+		log.Debugf("pgalloc.MemoryFile.UpdateUsage() inhibited during MemoryFile save")
+		return nil
+	}
 	if log.IsLogging(log.Debug) {
 		log.Debugf("UpdateUsage: took %v, currentUsage=%d knownCommittedBytes=%d",
 			time.Since(startTime), currentUsage, f.knownCommittedBytes)
@@ -1592,9 +1605,12 @@ func (f *MemoryFile) UpdateUsage(memCgIDs map[uint32]struct{}) error {
 // checkCommitted and false otherwise; wasCommitted can only be true if
 // alsoScanCommitted is true.
 //
+// callerIsSaveTo is true if the caller is f.SaveTo() and false if the caller
+// is f.UpdateUsage().
+//
 // Precondition: f.mu must be held; it may be unlocked and reacquired.
 // +checklocks:f.mu
-func (f *MemoryFile) updateUsageLocked(memCgIDs map[uint32]struct{}, alsoScanCommitted bool, checkCommitted func(bs []byte, committed []byte, off uint64, wasCommitted bool) error) error {
+func (f *MemoryFile) updateUsageLocked(memCgIDs map[uint32]struct{}, alsoScanCommitted, callerIsSaveTo bool, checkCommitted func(bs []byte, committed []byte, off uint64, wasCommitted bool) error) error {
 	// Track if anything changed to elide the merge.
 	changedAny := false
 	defer func() {
@@ -1666,10 +1682,14 @@ func (f *MemoryFile) updateUsageLocked(memCgIDs map[uint32]struct{}, alsoScanCom
 			}
 
 			// Reconcile internal state with buf. Since we temporarily dropped
-			// f.mu, f.memAcct may have changed, and maseg/ma are no longer
-			// valid. If wasCommitted is false, then we are marking ranges that
-			// are now committed; otherwise, we are marking ranges that are now
-			// uncommitted.
+			// f.mu, f.isSaving and f.memAcct may have changed, and maseg/ma
+			// are no longer valid. If wasCommitted is false, then we are
+			// marking ranges that are now committed; otherwise, we are marking
+			// ranges that are now uncommitted.
+			if !callerIsSaveTo && f.isSaving != 0 {
+				checkErr = updateUsageDuringSaveErr{}
+				return false
+			}
 			unchangedVal := byte(0)
 			if wasCommitted {
 				unchangedVal = 1
@@ -1743,6 +1763,13 @@ func (f *MemoryFile) updateUsageLocked(memCgIDs map[uint32]struct{}, alsoScanCom
 	return nil
 }
 
+type updateUsageDuringSaveErr struct{}
+
+// Error implements error.Error.
+func (updateUsageDuringSaveErr) Error() string {
+	return "pgalloc.MemoryFile.UpdateUsage() called during MemoryFile save"
+}
+
 // TotalUsage returns an aggregate usage for all memory statistics except
 // Mapped (which is external to MemoryFile). This is generally much cheaper
 // than UpdateUsage, but will not provide a fine-grained breakdown.

pkg/sentry/pgalloc/save_restore.go

@@ -78,12 +78,21 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w io.Writer, pw io.Writer, opts
 
 	// Ensure that all pages that contain non-zero bytes are marked
 	// known-committed, since we only store known-committed pages below.
+	//
+	// f.updateUsageLocked() will unlock f.mu before calling our callback,
+	// allowing concurrent calls to f.UpdateUsage() => f.updateUsageLocked() to
+	// observe pages that we transiently commit (for comparisons to zero) or
+	// leave committed (if opts.ExcludeCommittedZeroPages is true). Bump
+	// f.isSaving to prevent this.
+	f.isSaving++
+	defer func() { f.isSaving-- }()
 	timeScanStart := time.Now()
 	zeroPage := make([]byte, hostarch.PageSize)
 	var (
 		decommitWarnOnce  sync.Once
 		decommitPendingFR memmap.FileRange
 		scanTotal         uint64
+		committedTotal    uint64
 		decommitTotal     uint64
 		decommitCount     uint64
 	)
@@ -124,13 +133,14 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w io.Writer, pw io.Writer, opts
 			decommitPendingFR = memmap.FileRange{}
 		}
 	}
-	err := f.updateUsageLocked(nil, opts.ExcludeCommittedZeroPages, func(bs []byte, committed []byte, off uint64, wasCommitted bool) error {
+	err := f.updateUsageLocked(nil, opts.ExcludeCommittedZeroPages, true /* callerIsSaveTo */, func(bs []byte, committed []byte, off uint64, wasCommitted bool) error {
 		scanTotal += uint64(len(bs))
 		for pgoff := 0; pgoff < len(bs); pgoff += hostarch.PageSize {
 			i := pgoff / hostarch.PageSize
 			pg := bs[pgoff : pgoff+hostarch.PageSize]
 			if !bytes.Equal(pg, zeroPage) {
 				committed[i] = 1
+				committedTotal += hostarch.PageSize
 				continue
 			}
 			committed[i] = 0
@@ -149,7 +159,7 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w io.Writer, pw io.Writer, opts
 	if err != nil {
 		return err
 	}
-	log.Infof("MemoryFile(%p): saving scanned %d bytes, decommitted %d bytes in %d syscalls, %s", f, scanTotal, decommitTotal, decommitCount, time.Since(timeScanStart))
+	log.Infof("MemoryFile(%p): saving scanned %d bytes, saw %d committed bytes (ExcludeCommittedZeroPages=%v), decommitted %d bytes in %d syscalls, %s", f, scanTotal, committedTotal, opts.ExcludeCommittedZeroPages, decommitTotal, decommitCount, time.Since(timeScanStart))
 
 	// Save metadata.
 	timeMetadataStart := time.Now()