LoadSnapshot

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

Author: kzys

firecracker.go

@@ -474,3 +474,18 @@ func (f *Client) PatchBalloonStatsInterval(ctx context.Context, balloonStatsUpda
 
 	return f.client.Operations.PatchBalloonStatsInterval(params)
 }
+
+type LoadSnapshotOpt func(*ops.LoadSnapshotParams)
+
+func (f *Client) LoadSnapshot(ctx context.Context, cfg *models.SnapshotLoadParams, opts ...LoadSnapshotOpt) (*ops.LoadSnapshotNoContent, error) {
+	timeout, cancel := context.WithTimeout(ctx, time.Duration(f.firecrackerRequestTimeout)*time.Millisecond)
+	defer cancel()
+
+	params := ops.NewLoadSnapshotParamsWithContext(timeout)
+	params.SetBody(cfg)
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	return f.client.Operations.LoadSnapshot(params)
+}

handlers.go

@@ -286,8 +286,10 @@ var defaultFcInitHandlerList = HandlerList{}.Append(
 	StartVMMHandler,
 	CreateLogFilesHandler,
 	BootstrapLoggingHandler,
+
 	CreateMachineHandler,
 	CreateBootSourceHandler,
+
 	AttachDrivesHandler,
 	CreateNetworkInterfacesHandler,
 	AddVsocksHandler,

machine.go

@@ -150,6 +150,22 @@ type Config struct {
 	// It is possible to use a valid IPv4 link-local address (169.254.0.0/16).
 	// If not provided, the default address (169.254.169.254) will be used.
 	MmdsAddress net.IP
+
+	Snapshot SnapshotConfig
+}
+
+func (cfg *Config) validateSnapshot() error {
+	ss := cfg.Snapshot
+
+	if _, err := os.Stat(ss.MemFilePath); err != nil {
+		return fmt.Errorf("failed to stat snapshot memory file path, %q: %v", cfg.KernelImagePath, err)
+	}
+
+	if _, err := os.Stat(ss.SnapshotPath); err != nil {
+		return fmt.Errorf("failed to stat snapshot file path, %q: %v", cfg.InitrdPath, err)
+	}
+
+	return nil
 }
 
 // Validate will ensure that the required fields are set and that
@@ -159,6 +175,11 @@ func (cfg *Config) Validate() error {
 		return nil
 	}
 
+	ss := cfg.Snapshot
+	if ss.SnapshotPath != "" || ss.MemFilePath != "" {
+		return cfg.validateSnapshot()
+	}
+
 	if _, err := os.Stat(cfg.KernelImagePath); err != nil {
 		return fmt.Errorf("failed to stat kernel image path, %q: %v", cfg.KernelImagePath, err)
 	}
@@ -718,6 +739,17 @@ func (m *Machine) captureFifoToFileWithChannel(ctx context.Context, logger *log.
 }
 
 func (m *Machine) createMachine(ctx context.Context) error {
+	ss := m.Cfg.Snapshot
+	if ss.SnapshotPath != "" || ss.MemFilePath != "" {
+		_, err := m.client.LoadSnapshot(ctx, &models.SnapshotLoadParams{
+			SnapshotPath:        String(ss.SnapshotPath),
+			MemFilePath:         String(ss.MemFilePath),
+			EnableDiffSnapshots: ss.EnableDiffSnapshots,
+			ResumeVM:            ss.ResumeVM,
+		})
+		return err
+	}
+
 	resp, err := m.client.PutMachineConfiguration(ctx, &m.Cfg.MachineCfg)
 	if err != nil {
 		m.logger.Errorf("PutMachineConfiguration returned %s", resp.Error())
@@ -734,6 +766,11 @@ func (m *Machine) createMachine(ctx context.Context) error {
 }
 
 func (m *Machine) createBootSource(ctx context.Context, imagePath, initrdPath, kernelArgs string) error {
+	ss := m.Cfg.Snapshot
+	if ss.SnapshotPath != "" || ss.MemFilePath != "" {
+		return nil
+	}
+
 	bsrc := models.BootSource{
 		KernelImagePath: &imagePath,
 		InitrdPath:      initrdPath,
@@ -841,6 +878,10 @@ func (m *Machine) addVsock(ctx context.Context, dev VsockDevice) error {
 }
 
 func (m *Machine) startInstance(ctx context.Context) error {
+	if m.Cfg.Snapshot.SnapshotPath != "" || m.Cfg.Snapshot.MemFilePath != "" {
+		return nil
+	}
+
 	action := models.InstanceActionInfoActionTypeInstanceStart
 	info := models.InstanceActionInfo{
 		ActionType: &action,

machine_test.go

@@ -1728,6 +1728,64 @@ func TestCreateSnapshot(t *testing.T) {
 	}
 }
 
+func TestSnapshot(t *testing.T) {
+	fctesting.RequiresKVM(t)
+	fctesting.RequiresRoot(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	snapshotDir := t.TempDir()
+	snapPath := filepath.Join(snapshotDir, "snap.out")
+	memPath := filepath.Join(snapshotDir, "mem.out")
+
+	t.Run("create", func(t *testing.T) {
+		cfg := createValidConfig(t, filepath.Join(snapshotDir, "fc.sock"))
+		m, err := NewMachine(ctx, cfg, func(m *Machine) {
+			args := m.cmd.Args[1:]
+			m.cmd = exec.Command(getFirecrackerBinaryPath(), args...)
+		})
+		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)
+	})
+
+	t.Run("load", func(t *testing.T) {
+		dir := t.TempDir()
+		cfg := Config{
+			SocketPath: filepath.Join(dir, "fc.sock"),
+			Snapshot: SnapshotConfig{
+				SnapshotPath: snapPath,
+				MemFilePath:  memPath,
+			},
+		}
+		m, err := NewMachine(ctx, cfg, func(m *Machine) {
+			args := m.cmd.Args[1:]
+			m.cmd = exec.Command(getFirecrackerBinaryPath(), args...)
+		})
+		require.NoError(t, err)
+
+		err = m.Start(ctx)
+		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)