Merge branch 'demo_branch2'

This commit merges #348 from Gabriel Ionescu as is.

Signed-off-by: Kazuyoshi Kato <katokazu@amazon.com>

Author: kzys

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

@@ -380,7 +380,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 ...Opt) error {
 	m.logger.Debug("Called Machine.Start()")
 	alreadyStarted := true
 	m.startOnce.Do(func() {
@@ -401,6 +401,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
@@ -1092,6 +1096,22 @@ func (m *Machine) CreateSnapshot(ctx context.Context, memFilePath, snapshotPath
 	return nil
 }
 
+// LoadSnapshot load a snapshot
+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 {
+		m.logger.Errorf("failed to load a snapshot for VM: %v", err)
+		return 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

@@ -71,6 +71,9 @@ var (
 	testBalloonDeflateOnOom      = true
 	testStatsPollingIntervals    = int64(1)
 	testNewStatsPollingIntervals = int64(6)
+
+	// How long to wait for the socket to appear.
+	firecrackerSocketWait = int64(10)
 )
 
 func envOrDefault(k, empty string) string {
@@ -1728,6 +1731,76 @@ func TestCreateSnapshot(t *testing.T) {
 	}
 }
 
+func TestLoadSnapshot(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, ...Opt) error
 	StopVMM() error
 	Shutdown(context.Context) error
 	Wait(context.Context) error

opts.go

@@ -14,6 +14,7 @@
 package firecracker
 
 import (
+	"context"
 	"os/exec"
 
 	"github.com/sirupsen/logrus"
@@ -47,3 +48,13 @@ 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) Opt {
+	return func(machine *Machine) {
+		err := machine.LoadSnapshot(ctx, memFilePath, snapshotPath)
+		if err != nil {
+			machine.logger.Errorf("LoadSnapshot failed with %s", err)
+		}
+	}
+}