Implemented snapshot loading capabilities

Signed-off-by: David Son <davbson@amazon.com>

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)

handlers.go

@@ -34,7 +34,8 @@ const (
 	LinkFilesToRootFSHandlerName       = "fcinit.LinkFilesToRootFS"
 	SetupNetworkHandlerName            = "fcinit.SetupNetwork"
 	SetupKernelArgsHandlerName         = "fcinit.SetupKernelArgs"
-	CreateBalloonHandlerName           = "fcint.CreateBalloon"
+	CreateBalloonHandlerName           = "fcinit.CreateBalloon"
+	LoadSnapshotHandlerName            = "fcinit.LoadSnapshot"
 
 	ValidateCfgHandlerName        = "validate.Cfg"
 	ValidateJailerCfgHandlerName  = "validate.JailerCfg"
@@ -280,6 +281,17 @@ func NewCreateBalloonHandler(amountMib int64, deflateOnOom bool, StatsPollingInt
 	}
 }
 
+// NewLoadSnapshotHandler is a named handler that loads a snapshot
+// from the specified filepath
+func NewLoadSnapshotHandler(memFilePath, snapshotPath string, opts ...LoadSnapshotOpt) Handler {
+	return Handler{
+		Name: LoadSnapshotHandlerName,
+		Fn: func(ctx context.Context, m *Machine) error {
+			return m.loadSnapshot(ctx, memFilePath, snapshotPath, opts...)
+		},
+	}
+}
+
 var defaultFcInitHandlerList = HandlerList{}.Append(
 	SetupNetworkHandler,
 	SetupKernelArgsHandler,

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,10 @@ func (m *Machine) Start(ctx context.Context) error {
 		}
 	}()
 
+	for _, opt := range opts {
+		opt(m)
+	}
+
 	err = m.Handlers.Run(ctx, m)
 	if err != nil {
 		return err
@@ -1093,6 +1097,20 @@ func (m *Machine) CreateSnapshot(ctx context.Context, memFilePath, snapshotPath
 	return nil
 }
 
+func (m *Machine) loadSnapshot(ctx context.Context, memFilePath, snapshotPath string, opts ...LoadSnapshotOpt) 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
+}
+
 // CreateBalloon creates a balloon device if one does not exist
 func (m *Machine) CreateBalloon(ctx context.Context, amountMib int64, deflateOnOom bool, statsPollingIntervals int64, opts ...PutBalloonOpt) error {
 	balloon := models.Balloon{

machine_test.go

@@ -1728,6 +1728,113 @@ func TestCreateSnapshot(t *testing.T) {
 	}
 }
 
+func TestLoadSnapshot(t *testing.T) {
+	fctesting.RequiresKVM(t)
+	fctesting.RequiresRoot(t)
+
+	dir, err := ioutil.TempDir("", t.Name())
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	cases := []struct {
+		name           string
+		createSnapshot func(ctx context.Context, machineLogger *logrus.Logger, socketPath, memPath, snapPath string)
+		loadSnapshot   func(ctx context.Context, machineLogger *logrus.Logger, socketPath, memPath, snapPath string)
+	}{
+		{
+			name: "TestLoadSnapshot",
+			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) {
+				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(memPath, snapPath))
+				require.NoError(t, err)
+
+				t.Errorf("%v", m.Handlers)
+
+				err = m.ResumeVM(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) {
+
+			},
+
+			loadSnapshot: func(ctx context.Context, machineLogger *logrus.Logger, socketPath, memPath, snapPath string) {
+				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(memPath, snapPath))
+				require.Error(t, err)
+			},
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			ctx := context.Background()
+
+			// 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)
+
+			c.createSnapshot(ctx, machineLogger, socketPath, snapPath, memPath)
+			c.loadSnapshot(ctx, machineLogger, socketPath, snapPath, memPath)
+		})
+	}
+
+}
+
 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

@@ -21,6 +21,7 @@ import (
 
 // Opt represents a functional option to help modify functionality of a Machine.
 type Opt func(*Machine)
+type StartOpt func(*Machine)
 
 // 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 +48,11 @@ func WithProcessRunner(cmd *exec.Cmd) Opt {
 		machine.cmd = cmd
 	}
 }
+
+// WithSnapshot will allow for the machine to start using a given snapshot.
+func WithSnapshot(memFilePath, snapshotPath string, opts ...LoadSnapshotOpt) StartOpt {
+	return func(m *Machine) {
+		m.Handlers.FcInit.AppendAfter("SetupKernelArgsHandler",
+			NewLoadSnapshotHandler(memFilePath, snapshotPath, opts...))
+	}
+}