Add vfs.FilesystemImplSaveRestoreExtension.BeforeResume.

pkg/sentry/control/state.go

@@ -18,8 +18,6 @@ import (
 	"errors"
 	"fmt"
 
-	"gvisor.dev/gvisor/pkg/abi/linux"
-	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/state"
@@ -84,18 +82,7 @@ func (s *State) Save(o *SaveOpts, _ *struct{}) error {
 		Key:                o.Key,
 		Metadata:           o.Metadata,
 		MemoryFileSaveOpts: o.MemoryFileSaveOpts,
-		Callback: func(err error) {
-			if err == nil {
-				log.Infof("Save succeeded: exiting...")
-				s.Kernel.SetSaveSuccess(false /* autosave */)
-			} else {
-				log.Warningf("Save failed: %v", err)
-				s.Kernel.SetSaveError(err)
-			}
-			if !o.Resume {
-				s.Kernel.Kill(linux.WaitStatusExit(0))
-			}
-		},
+		Resume:             o.Resume,
 	}
 	if o.HavePagesFile {
 		saveOpts.PagesMetadata, err = o.ReleaseFD(1)

pkg/sentry/fsimpl/gofer/save_restore.go

@@ -31,6 +31,8 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
 
+var _ vfs.FilesystemImplSaveRestoreExtension = (*filesystem)(nil)
+
 // +stateify savable
 type savedDentryRW struct {
 	read  bool
@@ -134,6 +136,11 @@ func (d *dentry) beforeSave() {
 	}
 }
 
+// BeforeResume implements vfs.FilesystemImplSaveRestoreExtension.BeforeResume.
+func (fs *filesystem) BeforeResume(ctx context.Context) {
+	fs.savedDentryRW = nil
+}
+
 // afterLoad is invoked by stateify.
 func (fs *filesystem) afterLoad(ctx goContext.Context) {
 	fs.mf = pgalloc.MemoryFileFromContext(ctx)

pkg/sentry/fsimpl/tmpfs/save_restore.go

@@ -23,6 +23,8 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
 
+var _ vfs.FilesystemImplSaveRestoreExtension = (*filesystem)(nil)
+
 // saveMf is called by stateify.
 func (fs *filesystem) saveMf() string {
 	if !fs.mf.IsSavable() {
@@ -75,6 +77,9 @@ func (fs *filesystem) PrepareSave(ctx context.Context) error {
 	return nil
 }
 
+// BeforeResume implements vfs.FilesystemImplSaveRestoreExtension.BeforeResume.
+func (fs *filesystem) BeforeResume(ctx context.Context) {}
+
 // CompleteRestore implements
 // vfs.FilesystemImplSaveRestoreExtension.CompleteRestore.
 func (fs *filesystem) CompleteRestore(ctx context.Context, opts vfs.CompleteRestoreOptions) error {

pkg/sentry/kernel/kernel.go

@@ -702,6 +702,11 @@ func (k *Kernel) SaveTo(ctx context.Context, w, pagesMetadata io.Writer, pagesFi
 	return nil
 }
 
+// BeforeResume is called before the kernel is resumed after save.
+func (k *Kernel) BeforeResume(ctx context.Context) {
+	k.vfs.BeforeResume(ctx)
+}
+
 func (k *Kernel) saveMemoryFiles(ctx context.Context, w, pagesMetadata io.Writer, pagesFile *fd.FD, mfsToSave map[string]*pgalloc.MemoryFile, mfOpts pgalloc.SaveOpts) error {
 	// Save the memory files' state.
 	memoryStart := time.Now()

pkg/sentry/state/state.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"io"
 
+	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/fd"
@@ -68,15 +69,18 @@ type SaveOpts struct {
 	// MemoryFileSaveOpts is passed to calls to pgalloc.MemoryFile.SaveTo().
 	MemoryFileSaveOpts pgalloc.SaveOpts
 
-	// Callback is called prior to unpause, with any save error.
-	Callback func(err error)
-
 	// Resume indicates if the statefile is used for save-resume.
 	Resume bool
+
+	// Autosave indicates if the statefile is used for autosave.
+	Autosave bool
 }
 
 // Save saves the system state.
 func (opts SaveOpts) Save(ctx context.Context, k *kernel.Kernel, w *watchdog.Watchdog) error {
+	t, _ := CPUTime()
+	log.Infof("Before save CPU usage: %s", t.String())
+
 	log.Infof("Sandbox save started, pausing all tasks.")
 	k.Pause()
 	k.ReceiveTaskStates()
@@ -127,7 +131,22 @@ func (opts SaveOpts) Save(ctx context.Context, k *kernel.Kernel, w *watchdog.Wat
 			}
 		}
 	}
-	opts.Callback(err)
+
+	t1, _ := CPUTime()
+	log.Infof("Save CPU usage: %s", (t1 - t).String())
+	if err == nil {
+		log.Infof("Save succeeded: exiting...")
+		k.SetSaveSuccess(opts.Autosave)
+	} else {
+		log.Warningf("Save failed: exiting... %v", err)
+		k.SetSaveError(err)
+	}
+	if opts.Resume {
+		k.BeforeResume(ctx)
+	} else {
+		// Kill the sandbox.
+		k.Kill(linux.WaitStatusExit(0))
+	}
 	return err
 }
 

pkg/sentry/vfs/save_restore.go

@@ -54,6 +54,10 @@ type FilesystemImplSaveRestoreExtension interface {
 	// PrepareSave prepares this filesystem for serialization.
 	PrepareSave(ctx context.Context) error
 
+	// BeforeResume is called before the kernel is resumed after save. It can be
+	// used to clean up any state that should be discarded after save.
+	BeforeResume(ctx context.Context)
+
 	// CompleteRestore completes restoration from checkpoint for this
 	// filesystem after deserialization.
 	CompleteRestore(ctx context.Context, opts CompleteRestoreOptions) error
@@ -73,6 +77,17 @@ func (vfs *VirtualFilesystem) PrepareSave(ctx context.Context) error {
 	return nil
 }
 
+// BeforeResume is called before the kernel is resumed after save and allows
+// filesystems to clean up S/R state.
+func (vfs *VirtualFilesystem) BeforeResume(ctx context.Context) {
+	for fs := range vfs.getFilesystems() {
+		if ext, ok := fs.impl.(FilesystemImplSaveRestoreExtension); ok {
+			ext.BeforeResume(ctx)
+		}
+		fs.DecRef(ctx)
+	}
+}
+
 // CompleteRestore completes restoration from checkpoint for all filesystems
 // after deserialization.
 func (vfs *VirtualFilesystem) CompleteRestore(ctx context.Context, opts *CompleteRestoreOptions) error {

runsc/boot/autosave.go

@@ -18,7 +18,6 @@ import (
 	"bytes"
 	"fmt"
 
-	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
@@ -28,42 +27,19 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
-func getSaveOpts(l *Loader, k *kernel.Kernel, isResume bool) state.SaveOpts {
-	t, _ := state.CPUTime()
-	log.Infof("Before save CPU usage: %s", t.String())
-	saveOpts := state.SaveOpts{
-		Key:    nil,
-		Resume: isResume,
-		Callback: func(err error) {
-			t1, _ := state.CPUTime()
-			log.Infof("Save CPU usage: %s", (t1 - t).String())
-			if err == nil {
-				log.Infof("Save succeeded: exiting...")
-				k.SetSaveSuccess(true)
-			} else {
-				log.Warningf("Save failed: exiting... %v", err)
-				k.SetSaveError(err)
-			}
-
-			if !isResume {
-				// Kill the sandbox.
-				k.Kill(linux.WaitStatusExit(0))
-			}
-		},
-	}
-	return saveOpts
-}
-
 func getTargetForSaveResume(l *Loader) func(k *kernel.Kernel) {
 	return func(k *kernel.Kernel) {
 		l.addVersionToCheckpoint()
 		l.addContainerSpecsToCheckpoint()
-		saveOpts := getSaveOpts(l, k, true /* isResume */)
 		// Store the state file contents in a buffer for save-resume.
 		// There is no need to verify the state file, we just need the
 		// sandbox to continue running after save.
 		var buf bytes.Buffer
-		saveOpts.Destination = &buf
+		saveOpts := state.SaveOpts{
+			Autosave:    true,
+			Resume:      true,
+			Destination: &buf,
+		}
 		saveOpts.Save(k.SupervisorContext(), k, l.watchdog)
 	}
 }
@@ -78,8 +54,11 @@ func getTargetForSaveRestore(l *Loader, files []*fd.FD) func(k *kernel.Kernel) {
 		once.Do(func() {
 			l.addVersionToCheckpoint()
 			l.addContainerSpecsToCheckpoint()
-			saveOpts := getSaveOpts(l, k, false /* isResume */)
-			saveOpts.Destination = files[0]
+			saveOpts := state.SaveOpts{
+				Autosave:    true,
+				Resume:      false,
+				Destination: files[0],
+			}
 			if len(files) == 3 {
 				saveOpts.PagesMetadata = files[1]
 				saveOpts.PagesFile = files[2]