Adjust MemoryBackend API and add test

Signed-off-by: Brandon Duffany <[email protected]>

Author: bduffany

machine.go

@@ -172,7 +172,7 @@ type Config struct {
 }
 
 func (cfg *Config) hasSnapshot() bool {
-	return cfg.Snapshot.GetMemFilePath() != "" || cfg.Snapshot.SnapshotPath != ""
+	return cfg.Snapshot.GetMemBackendPath() != "" || cfg.Snapshot.SnapshotPath != ""
 }
 
 // Validate will ensure that the required fields are set and that
@@ -235,7 +235,7 @@ func (cfg *Config) ValidateLoadSnapshot() error {
 		return fmt.Errorf("socket %s already exists", cfg.SocketPath)
 	}
 
-	if _, err := os.Stat(cfg.Snapshot.GetMemFilePath()); err != nil {
+	if _, err := os.Stat(cfg.Snapshot.GetMemBackendPath()); err != nil {
 		return err
 	}
 

machine_test.go

@@ -2125,6 +2125,62 @@ func TestLoadSnapshot(t *testing.T) {
 				require.NoError(t, err)
 			},
 		},
+		{
+			name: "TestLoadSnapshotWithMemoryBackend",
+			createSnapshot: func(ctx context.Context, machineLogger *logrus.Logger, socketPath, memPath, snapPath string) {
+				// Create a snapshot
+				cfg := createValidConfig(t, socketPath+".create")
+				m, err := NewMachine(ctx, cfg, func(m *Machine) {
+					// Rewriting m.cmd partially wouldn't work since Cmd has
+					// some unexported members
+					args := m.cmd.Args[1:]
+					m.cmd = exec.Command(getFirecrackerBinaryPath(), args...)
+				}, WithLogger(logrus.NewEntry(machineLogger)))
+				require.NoError(t, err)
+
+				err = m.Start(ctx)
+				require.NoError(t, err)
+
+				err = m.PauseVM(ctx)
+				require.NoError(t, err)
+
+				err = m.CreateSnapshot(ctx, memPath, snapPath)
+				require.NoError(t, err)
+
+				err = m.StopVMM()
+				require.NoError(t, err)
+			},
+
+			loadSnapshot: func(ctx context.Context, machineLogger *logrus.Logger, socketPath, memPath, snapPath string) {
+				// Note that many fields are not necessary when loading a snapshot
+				cfg := Config{
+					SocketPath: socketPath + ".load",
+					Drives: []models.Drive{
+						{
+							DriveID:      String("root"),
+							IsRootDevice: Bool(true),
+							IsReadOnly:   Bool(true),
+							PathOnHost:   String(testRootfs),
+						},
+					},
+				}
+
+				m, err := NewMachine(ctx, cfg, func(m *Machine) {
+					// Rewriting m.cmd partially wouldn't work since Cmd has
+					// some unexported members
+					args := m.cmd.Args[1:]
+					m.cmd = exec.Command(getFirecrackerBinaryPath(), args...)
+				}, WithLogger(logrus.NewEntry(machineLogger)), WithSnapshot("", snapPath, WithMemoryBackend("File", memPath)))
+				require.NoError(t, err)
+				require.Equal(t, m.Cfg.Snapshot.ResumeVM, true)
+
+				err = m.Start(ctx)
+				require.NoError(t, err)
+
+				err = m.StopVMM()
+				require.NoError(t, err)
+			},
+		},
 		{
 			name: "TestLoadSnapshot without create",
 			createSnapshot: func(ctx context.Context, machineLogger *logrus.Logger, socketPath, memPath, snapPath string) {

opts.go

@@ -53,9 +53,15 @@ func WithProcessRunner(cmd *exec.Cmd) Opt {
 // to be passed to LoadSnapshot
 type WithSnapshotOpt func(*SnapshotConfig)
 
-// WithSnapshot will allow for the machine to start using a given snapshot. A
-// UFFD socket file path may be used as memFilePath if a MemoryBackendType opt
-// is passed setting the backend type to models.MemoryBackendBackendTypeUffd.
+// WithSnapshot will allow for the machine to start using a given snapshot.
+//
+// If using the UFFD memory backend, the memFilePath may be empty (it is
+// ignored), and instead the UFFD socket should be specified using
+// MemoryBackendType, as in the following example:
+//
+//	WithSnapshot(
+//	  "", snapshotPath,
+//	  WithMemoryBackend(models.MemoryBackendBackendTypeUffd, "uffd.sock"))
 func WithSnapshot(memFilePath, snapshotPath string, opts ...WithSnapshotOpt) Opt {
 	return func(m *Machine) {
 		m.Cfg.Snapshot.MemFilePath = memFilePath
@@ -70,16 +76,18 @@ func WithSnapshot(memFilePath, snapshotPath string, opts ...WithSnapshotOpt) Opt
 	}
 }
 
-// MemoryBackendType sets the memory backend type to the given value.
-func MemoryBackendType(typ string) WithSnapshotOpt {
+// WithMemoryBackend sets the memory backend to the given type, using the given
+// backing file path (a regular file for "File" type, or a UFFD socket path for
+// "Uffd" type).
+//
+// Note that if MemFilePath is already configured for the snapshot config, it
+// will be ignored, and the backendPath specified here will be used instead.
+func WithMemoryBackend(backendType, backendPath string) WithSnapshotOpt {
 	return func(cfg *SnapshotConfig) {
-		if cfg.MemBackend == nil {
-			cfg.MemBackend = &models.MemoryBackend{
-				BackendPath: String(cfg.MemFilePath),
-			}
+		cfg.MemBackend = &models.MemoryBackend{
+			BackendType: String(backendType),
+			BackendPath: String(backendPath),
 		}
-		cfg.MemBackend.BackendType = String(typ)
-		// Clear MemFilePath since only one of MemFilePath or MemBackend is allowed.
 		cfg.MemFilePath = ""
 	}
 }

snapshot.go

@@ -23,7 +23,7 @@ type SnapshotConfig struct {
 	ResumeVM            bool
 }
 
-func (cfg *SnapshotConfig) GetMemFilePath() string {
+func (cfg *SnapshotConfig) GetMemBackendPath() string {
 	if cfg.MemBackend != nil && cfg.MemBackend.BackendPath != nil {
 		return *cfg.MemBackend.BackendPath
 	}