Add support for snapshot memory backend (#495)

* Add support for memory backend

- Pull in upstream swagger definitions for SnapshotLoadParams and MemoryBackend
- Run `go generate`
- Update Snapshot config to allow setting MemBackend
- Forward the MemBackend through to the Firecracker client

Signed-off-by: Brandon Duffany <brandon@buildbuddy.io>

* Adjust MemoryBackend API and add test

Signed-off-by: Brandon Duffany <brandon@buildbuddy.io>

* Bump firecracker version to v1.4.0

Signed-off-by: Brandon Duffany <brandon@buildbuddy.io>

* Address PR feedback

Signed-off-by: Brandon Duffany <brandon@buildbuddy.io>

* Fix test

Signed-off-by: Brandon Duffany <brandon@buildbuddy.io>

---------

Signed-off-by: Brandon Duffany <brandon@buildbuddy.io>

Author: bduffany

Makefile

@@ -28,7 +28,7 @@ JAILER_BIN=$(FC_TEST_DATA_PATH)/jailer-main
 UID = $(shell id -u)
 GID = $(shell id -g)
 
-firecracker_version=v1.0.0
+firecracker_version=v1.4.0
 
 # The below files are needed and can be downloaded from the internet
 release_url=https://github.com/firecracker-microvm/firecracker/releases/download/$(firecracker_version)/firecracker-$(firecracker_version)-$(arch).tgz

client/models/memory_backend.go

@@ -953,6 +953,25 @@ definitions:
         maximum: 32
         description: Number of vCPUs (either 1 or an even number)
 
+  MemoryBackend:
+    type: object
+    required:
+      - backend_type
+      - backend_path
+    properties:
+      backend_type:
+        type: string
+        enum:
+          - File
+          - Uffd
+      backend_path:
+        type: string
+        description: Based on 'backend_type' it is either
+          1) Path to the file that contains the guest memory to be loaded
+          2) Path to the UDS where a process is listening for a UFFD initialization
+          control payload and open file descriptor that it can use to serve this
+          process's guest memory page faults
+
   Metrics:
     type: object
     description:
@@ -1090,8 +1109,10 @@ definitions:
 
   SnapshotLoadParams:
     type: object
+    description:
+      Defines the configuration used for handling snapshot resume. Exactly one of
+      the two `mem_*` fields must be present in the body of the request.
     required:
-      - mem_file_path
       - snapshot_path
     properties:
       enable_diff_snapshots:
@@ -1100,7 +1121,16 @@ definitions:
           Enable support for incremental (diff) snapshots by tracking dirty guest pages.
       mem_file_path:
         type: string
-        description: Path to the file that contains the guest memory to be loaded.
+        description:
+          Path to the file that contains the guest memory to be loaded.
+          This parameter has been deprecated and is only allowed if
+          `mem_backend` is not present.
+      mem_backend:
+        $ref: "#/definitions/MemoryBackend"
+        description:
+          Configuration for the backend that handles memory load. If this field
+          is specified, `mem_file_path` is forbidden. Either `mem_backend` or
+          `mem_file_path` must be present at a time.
       snapshot_path:
         type: string
         description: Path to the file that contains the microVM state to be loaded.

client/models/snapshot_load_params.go

@@ -172,7 +172,7 @@ type Config struct {
 }
 
 func (cfg *Config) hasSnapshot() bool {
-	return cfg.Snapshot.MemFilePath != "" || 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.MemFilePath); err != nil {
+	if _, err := os.Stat(cfg.Snapshot.GetMemBackendPath()); err != nil {
 		return err
 	}
 
@@ -649,7 +649,7 @@ func (m *Machine) startVMM(ctx context.Context) error {
 	return nil
 }
 
-//StopVMM stops the current VMM.
+// StopVMM stops the current VMM.
 func (m *Machine) StopVMM() error {
 	return m.stopVMM()
 }
@@ -1171,7 +1171,8 @@ func (m *Machine) CreateSnapshot(ctx context.Context, memFilePath, snapshotPath
 // loadSnapshot loads a snapshot of the VM
 func (m *Machine) loadSnapshot(ctx context.Context, snapshot *SnapshotConfig) error {
 	snapshotParams := &models.SnapshotLoadParams{
-		MemFilePath:         &snapshot.MemFilePath,
+		MemFilePath:         snapshot.MemFilePath,
+		MemBackend:          snapshot.MemBackend,
 		SnapshotPath:        &snapshot.SnapshotPath,
 		EnableDiffSnapshots: snapshot.EnableDiffSnapshots,
 		ResumeVM:            snapshot.ResumeVM,

client/swagger.yaml

@@ -2125,6 +2125,61 @@ 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)
+
+				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) {

machine.go

@@ -16,6 +16,7 @@ package firecracker
 import (
 	"os/exec"
 
+	"github.com/firecracker-microvm/firecracker-go-sdk/client/models"
 	"github.com/sirupsen/logrus"
 )
 
@@ -53,6 +54,14 @@ func WithProcessRunner(cmd *exec.Cmd) Opt {
 type WithSnapshotOpt func(*SnapshotConfig)
 
 // 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
@@ -66,3 +75,18 @@ func WithSnapshot(memFilePath, snapshotPath string, opts ...WithSnapshotOpt) Opt
 		m.Handlers.FcInit = loadSnapshotHandlerList
 	}
 }
+
+// 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) {
+		cfg.MemBackend = &models.MemoryBackend{
+			BackendType: String(backendType),
+			BackendPath: String(backendPath),
+		}
+	}
+}