Implemented snapshot loading capabilities

Signed-off-by: David Son <[email protected]>

Author: sondavidb

firecracker.go

@@ -265,6 +265,23 @@ func (f *Client) CreateSnapshot(ctx context.Context, snapshotParams *models.Snap
 	return f.client.Operations.CreateSnapshot(params)
 }
 
+// LoadSnapshotOpt is a functional option to be used for the
+// LoadSnapshot API in setting any additional optional fields.
+type LoadSnapshotOpt func(*ops.LoadSnapshotParams)
+
+// LoadSnapshot is a wrapper for the swagger generated client to make
+// calling of the API easier.
+func (f *Client) LoadSnapshot(ctx context.Context, snapshotParams *models.SnapshotLoadParams, opts ...LoadSnapshotOpt) (*ops.LoadSnapshotNoContent, error) {
+	params := ops.NewLoadSnapshotParamsWithContext(ctx)
+	params.SetBody(snapshotParams)
+
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	return f.client.Operations.LoadSnapshot(params)
+}
+
 // CreateSyncActionOpt is a functional option to be used for the
 // CreateSyncAction API in setting any additional optional fields.
 type CreateSyncActionOpt func(*ops.CreateSyncActionParams)

machine.go

@@ -381,7 +381,7 @@ func NewMachine(ctx context.Context, cfg Config, opts ...Opt) (*Machine, error)
 // handlers succeed, then this will start the VMM instance.
 // Start may only be called once per Machine.  Subsequent calls will return
 // ErrAlreadyStarted.
-func (m *Machine) Start(ctx context.Context) error {
+func (m *Machine) Start(ctx context.Context, opts ...StartOpt) error {
 	m.logger.Debug("Called Machine.Start()")
 	alreadyStarted := true
 	m.startOnce.Do(func() {
@@ -402,6 +402,12 @@ func (m *Machine) Start(ctx context.Context) error {
 		}
 	}()
 
+	for _, opt := range opts {
+		if err = opt(m); err != nil {
+			return err
+		}
+	}
+
 	err = m.Handlers.Run(ctx, m)
 	if err != nil {
 		return err

machine_test.go

@@ -1728,6 +1728,76 @@ func TestCreateSnapshot(t *testing.T) {
 	}
 }
 
+func TestCreateLoadSnapshot(t *testing.T) {
+	fctesting.RequiresKVM(t)
+	fctesting.RequiresRoot(t)
+
+	ctx := context.Background()
+
+	dir, err := ioutil.TempDir("", t.Name())
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	// Set snap and mem paths
+	socketPath := filepath.Join(dir, fsSafeTestName.Replace(t.Name()))
+	snapPath := socketPath + "SnapFile"
+	memPath := socketPath + "MemFile"
+	defer os.Remove(socketPath)
+	defer os.Remove(snapPath)
+	defer os.Remove(memPath)
+
+	// Tee logs for validation:
+	var logBuffer bytes.Buffer
+	machineLogger := logrus.New()
+	machineLogger.Out = io.MultiWriter(os.Stderr, &logBuffer)
+
+	// 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)
+	}
+
+	// Load a snapshot
+	{
+		cfg := createValidConfig(t, socketPath+".load")
+		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, WithSnapshot(ctx, memPath, snapPath))
+		require.NoError(t, err)
+
+		err = m.ResumeVM(ctx)
+		require.NoError(t, err)
+
+		err = m.StopVMM()
+		require.NoError(t, err)
+	}
+
+}
+
 func testCreateBalloon(ctx context.Context, t *testing.T, m *Machine) {
 	if err := m.CreateBalloon(ctx, testBalloonMemory, testBalloonDeflateOnOom, testStatsPollingIntervals); err != nil {
 		t.Errorf("Create balloon device failed from testAttachBalloon: %s", err)

machineiface.go

@@ -23,7 +23,7 @@ var _ MachineIface = (*Machine)(nil)
 // MachineIface can be used for mocking and testing of the Machine. The Machine
 // is subject to change, meaning this interface would change.
 type MachineIface interface {
-	Start(context.Context) error
+	Start(context.Context, ...StartOpt) error
 	StopVMM() error
 	Shutdown(context.Context) error
 	Wait(context.Context) error

opts.go

@@ -14,13 +14,17 @@
 package firecracker
 
 import (
+	"context"
+	"fmt"
 	"os/exec"
 
+	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
 	"github.com/sirupsen/logrus"
 )
 
 // Opt represents a functional option to help modify functionality of a Machine.
 type Opt func(*Machine)
+type StartOpt func(*Machine) error
 
 // WithClient will use the client in place rather than the client constructed
 // during bootstrapping of the machine. This option is useful for mocking out
@@ -47,3 +51,20 @@ func WithProcessRunner(cmd *exec.Cmd) Opt {
 		machine.cmd = cmd
 	}
 }
+
+// WithSnapshot will allow for the machine to start using a given snapshot.
+func WithSnapshot(ctx context.Context, memFilePath, snapshotPath string, opts ...LoadSnapshotOpt) StartOpt {
+	return func(m *Machine) error {
+		snapshotParams := &models.SnapshotLoadParams{
+			MemFilePath:  String(memFilePath),
+			SnapshotPath: String(snapshotPath),
+		}
+
+		if _, err := m.client.LoadSnapshot(ctx, snapshotParams, opts...); err != nil {
+			return fmt.Errorf("failed to load a snapshot for VM: %v", err)
+		}
+
+		m.logger.Debug("snapshot loaded successfully")
+		return nil
+	}
+}