Add support for saving PMAs referencing tmpfs filestore files.

The rootfs is such a mount (with default overlay2=root:self configuration).
It should be possible to checkpoint this mount when application has active
VMAs that have allocated PMAs from the backing filestore file.

2d90b66af1 ("Add checkpoint/restore support for tmpfs with file backend.")
added support for saving such a tmpfs filestore file during checkpoint and
restoring it correctly. Hence, after restore, the VMAs and PMAs should be
valid.

PiperOrigin-RevId: 605766485

pkg/sentry/kernel/kernel.go

@@ -603,6 +603,11 @@ func (k *Kernel) SaveTo(ctx context.Context, w wire.Writer) error {
 	if err := k.vfs.PrepareSave(vfsCtx); err != nil {
 		return err
 	}
+	// Mark all to-be-saved MemoryFiles as savable to inform kernel save below.
+	k.mf.MarkSavable()
+	for _, mf := range mfsToSave {
+		mf.MarkSavable()
+	}
 
 	// Save the CPUID FeatureSet before the rest of the kernel so we can
 	// verify its compatibility on restore before attempting to restore the

pkg/sentry/mm/save_restore.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 )
 
 // InvalidateUnsavable invokes memmap.Mappable.InvalidateUnsavable on all
@@ -38,7 +39,11 @@ func (mm *MemoryManager) InvalidateUnsavable(ctx context.Context) error {
 // beforeSave is invoked by stateify.
 func (mm *MemoryManager) beforeSave() {
 	for pseg := mm.pmas.FirstSegment(); pseg.Ok(); pseg = pseg.NextSegment() {
-		if pma := pseg.ValuePtr(); pma.file != mm.mf {
+		if pma := pseg.ValuePtr(); pma.file != nil {
+			if mf, ok := pma.file.(*pgalloc.MemoryFile); ok && mf.IsSavable() {
+				// If the MemoryFile will be saved, then its PMAs are preserved.
+				continue
+			}
 			// InvalidateUnsavable should have caused all such pmas to be
 			// invalidated.
 			panic(fmt.Sprintf("Can't save pma %#v with non-MemoryFile of type %T:\n%s", pseg.Range(), pma.file, mm))

pkg/sentry/pgalloc/pgalloc.go

@@ -183,6 +183,10 @@ type MemoryFile struct {
 	// notifications used to drive eviction. stopNotifyPressure is
 	// immutable.
 	stopNotifyPressure func()
+
+	// savable is true if this MemoryFile will be saved via SaveTo() during
+	// the kernel's SaveTo operation. savable is protected by mu.
+	savable bool
 }
 
 // MemoryFileOpts provides options to NewMemoryFile.

pkg/sentry/pgalloc/save_restore.go

@@ -115,6 +115,20 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w wire.Writer) error {
 	return nil
 }
 
+// MarkSavable marks f as savable.
+func (f *MemoryFile) MarkSavable() {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.savable = true
+}
+
+// IsSavable returns true if f is savable.
+func (f *MemoryFile) IsSavable() bool {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.savable
+}
+
 // LoadFrom loads MemoryFile state from the given stream.
 func (f *MemoryFile) LoadFrom(ctx context.Context, r wire.Reader) error {
 	// Load metadata.