move rw to opt; make snapshot opt; move to NewContainerOpts

Signed-off-by: Evan Hazlett <ejhazlett@gmail.com>

client.go

@@ -538,11 +538,7 @@ func (c *Client) Restore(ctx context.Context, id string, checkpoint Image, opts
 
 	copts := []NewContainerOpts{}
 	for _, o := range opts {
-		co, err := o(ctx, id, c, checkpoint, index)
+		copts = append(copts, o(ctx, id, c, checkpoint, index))
-		if err != nil {
-			return nil, err
-		}
-		copts = append(copts, co...)
 	}
 
 	ctr, err := c.NewContainer(ctx, id, copts...)
@@ -550,26 +546,6 @@ func (c *Client) Restore(ctx context.Context, id string, checkpoint Image, opts
 		return nil, err
 	}
 
-	// apply rw layer
-	info, err := ctr.Info(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	rw, err := GetIndexByMediaType(index, ocispec.MediaTypeImageLayerGzip)
-	if err != nil {
-		return nil, err
-	}
-
-	mounts, err := c.SnapshotService(info.Snapshotter).Mounts(ctx, info.SnapshotKey)
-	if err != nil {
-		return nil, err
-	}
-
-	if _, err := c.DiffService().Apply(ctx, *rw, mounts); err != nil {
-		return nil, err
-	}
-
 	return ctr, nil
 }
 

cmd/ctr/commands/containers/containers.go

@@ -318,7 +318,10 @@ var checkpointCommand = cli.Command{
 			return err
 		}
 		defer cancel()
-		opts := []containerd.CheckpointOpts{}
+		opts := []containerd.CheckpointOpts{
+			containerd.WithCheckpointRuntime,
+		}
+
 		if context.Bool("image") {
 			opts = append(opts, containerd.WithCheckpointImage)
 		}
@@ -338,10 +341,13 @@ var checkpointCommand = cli.Command{
 				return err
 			}
 		}
-		if err := task.Pause(ctx); err != nil {
+		// pause if running
-			return err
+		if task != nil {
+			if err := task.Pause(ctx); err != nil {
+				return err
+			}
+			defer task.Resume(ctx)
 		}
-		defer task.Resume(ctx)
 
 		if _, err := container.Checkpoint(ctx, ref, opts...); err != nil {
 			return err
@@ -360,6 +366,10 @@ var restoreCommand = cli.Command{
 			Name:  "live",
 			Usage: "restore the runtime and memory data from the checkpoint",
 		},
+		cli.BoolFlag{
+			Name:  "rw",
+			Usage: "restore the rw layer from the checkpoint",
+		},
 	},
 	Action: func(context *cli.Context) error {
 		id := context.Args().First()
@@ -394,6 +404,9 @@ var restoreCommand = cli.Command{
 			containerd.WithRestoreSnapshot,
 			containerd.WithRestoreRuntime,
 		}
+		if context.Bool("rw") {
+			opts = append(opts, containerd.WithRestoreRW)
+		}
 
 		ctr, err := client.Restore(ctx, id, checkpoint, opts...)
 		if err != nil {
@@ -404,6 +417,7 @@ var restoreCommand = cli.Command{
 		if context.Bool("live") {
 			topts = append(topts, containerd.WithTaskCheckpoint(checkpoint))
 		}
+
 		task, err := ctr.NewTask(ctx, cio.NewCreator(cio.WithStdio), topts...)
 		if err != nil {
 			return err

container.go

@@ -17,12 +17,10 @@
 package containerd
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"os"
 	"path/filepath"
-	"runtime"
 	"strings"
 
 	"github.com/containerd/containerd/api/services/tasks/v1"
@@ -40,6 +38,11 @@ import (
 	"github.com/pkg/errors"
 )
 
+const (
+	checkpointImageNameLabel   = "image.name"
+	checkpointRuntimeNameLabel = "runtime.name"
+)
+
 // Container is a metadata object for container resources and task creation
 type Container interface {
 	// ID identifies the container
@@ -312,45 +315,9 @@ func (c *container) Checkpoint(ctx context.Context, ref string, opts ...Checkpoi
 	defer done(ctx)
 
 	// add image name to manifest
-	ir := bytes.NewReader([]byte(img.Name()))
+	index.Annotations[checkpointImageNameLabel] = img.Name()
-	idesc, err := writeContent(ctx, c.client.ContentStore(), images.MediaTypeContainerd1CheckpointImageName, info.ID+"-image-name", ir)
-	if err != nil {
-		return nil, err
-	}
-	idesc.Platform = &ocispec.Platform{
-		OS:           runtime.GOOS,
-		Architecture: runtime.GOARCH,
-	}
-	index.Manifests = append(index.Manifests, idesc)
-
 	// add runtime info to index
-	rr := bytes.NewReader([]byte(info.Runtime.Name))
+	index.Annotations[checkpointRuntimeNameLabel] = info.Runtime.Name
-	rdesc, err := writeContent(ctx, c.client.ContentStore(), images.MediaTypeContainerd1CheckpointRuntimeName, info.ID+"-runtime-name", rr)
-	if err != nil {
-		return nil, err
-	}
-	rdesc.Platform = &ocispec.Platform{
-		OS:           runtime.GOOS,
-		Architecture: runtime.GOARCH,
-	}
-	index.Manifests = append(index.Manifests, rdesc)
-
-	if info.Runtime.Options != nil {
-		data, err := info.Runtime.Options.Marshal()
-		if err != nil {
-			return nil, err
-		}
-		r := bytes.NewReader(data)
-		desc, err := writeContent(ctx, c.client.ContentStore(), images.MediaTypeContainerd1CheckpointRuntimeOptions, info.ID+"-runtime-options", r)
-		if err != nil {
-			return nil, err
-		}
-		desc.Platform = &ocispec.Platform{
-			OS:           runtime.GOOS,
-			Architecture: runtime.GOARCH,
-		}
-		index.Manifests = append(index.Manifests, desc)
-	}
 
 	// process remaining opts
 	for _, o := range opts {

container_checkpoint_opts.go

@@ -17,6 +17,7 @@
 package containerd
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"runtime"
@@ -24,6 +25,7 @@ import (
 	tasks "github.com/containerd/containerd/api/services/tasks/v1"
 	"github.com/containerd/containerd/containers"
 	"github.com/containerd/containerd/diff"
+	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/rootfs"
 	"github.com/containerd/containerd/runtime/v2/runc/options"
 	"github.com/containerd/typeurl"
@@ -78,23 +80,35 @@ func WithCheckpointTask(ctx context.Context, client *Client, c *containers.Conta
 	return nil
 }
 
+// WithCheckpointRuntime includes the container runtime info
+func WithCheckpointRuntime(ctx context.Context, client *Client, c *containers.Container, index *imagespec.Index, copts *options.CheckpointOptions) error {
+	if c.Runtime.Options != nil {
+		data, err := c.Runtime.Options.Marshal()
+		if err != nil {
+			return err
+		}
+		r := bytes.NewReader(data)
+		desc, err := writeContent(ctx, client.ContentStore(), images.MediaTypeContainerd1CheckpointRuntimeOptions, c.ID+"-runtime-options", r)
+		if err != nil {
+			return err
+		}
+		desc.Platform = &imagespec.Platform{
+			OS:           runtime.GOOS,
+			Architecture: runtime.GOARCH,
+		}
+		index.Manifests = append(index.Manifests, desc)
+	}
+	return nil
+}
+
 // WithCheckpointRW includes the rw in the checkpoint
 func WithCheckpointRW(ctx context.Context, client *Client, c *containers.Container, index *imagespec.Index, copts *options.CheckpointOptions) error {
-	cnt, err := client.LoadContainer(ctx, c.ID)
-	if err != nil {
-		return err
-	}
-	info, err := cnt.Info(ctx)
-	if err != nil {
-		return err
-	}
-
 	diffOpts := []diff.Opt{
-		diff.WithReference(fmt.Sprintf("checkpoint-rw-%s", info.SnapshotKey)),
+		diff.WithReference(fmt.Sprintf("checkpoint-rw-%s", c.SnapshotKey)),
 	}
 	rw, err := rootfs.CreateDiff(ctx,
-		info.SnapshotKey,
+		c.SnapshotKey,
-		client.SnapshotService(info.Snapshotter),
+		client.SnapshotService(c.Snapshotter),
 		client.DiffService(),
 		diffOpts...,
 	)

container_restore_opts.go

@@ -19,11 +19,10 @@ package containerd
 import (
 	"context"
 
+	"github.com/containerd/containerd/containers"
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/images"
-	"github.com/containerd/containerd/oci"
 	"github.com/containerd/containerd/platforms"
-	"github.com/containerd/typeurl"
 	"github.com/gogo/protobuf/proto"
 	ptypes "github.com/gogo/protobuf/types"
 	"github.com/opencontainers/image-spec/identity"
@@ -31,135 +30,130 @@ import (
 	"github.com/pkg/errors"
 )
 
+var (
+	// ErrImageNameNotFoundInIndex is returned when the image name is not found in the index
+	ErrImageNameNotFoundInIndex = errors.New("image name not found in index")
+	// ErrRuntimeNameNotFoundInIndex is returned when the runtime name is not found in the index
+	ErrRuntimeNameNotFoundInIndex = errors.New("runtime name not found in index")
+)
+
 // RestoreOpts are options to manage the restore operation
-type RestoreOpts func(context.Context, string, *Client, Image, *imagespec.Index) ([]NewContainerOpts, error)
+type RestoreOpts func(context.Context, string, *Client, Image, *imagespec.Index) NewContainerOpts
 
 // WithRestoreImage restores the image for the container
-func WithRestoreImage(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, error) {
+func WithRestoreImage(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) NewContainerOpts {
-	store := client.ContentStore()
+	return func(ctx context.Context, client *Client, c *containers.Container) error {
-	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointImageName)
+		name, ok := index.Annotations[checkpointImageNameLabel]
-	if err != nil {
+		if !ok || name == "" {
-		if err != ErrMediaTypeNotFound {
+			return ErrRuntimeNameNotFoundInIndex
-			return nil, err
 		}
-	}
+		i, err := client.GetImage(ctx, name)
-	imageName := ""
-	if m != nil {
-		data, err := content.ReadBlob(ctx, store, *m)
 		if err != nil {
-			return nil, err
+			return err
 		}
-		imageName = string(data)
-	}
-	i, err := client.GetImage(ctx, imageName)
-	if err != nil {
-		return nil, err
-	}
 
-	return []NewContainerOpts{
+		c.Image = i.Name()
-		WithImage(i),
+		return nil
-	}, nil
+	}
 }
 
 // WithRestoreRuntime restores the runtime for the container
-func WithRestoreRuntime(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, error) {
+func WithRestoreRuntime(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) NewContainerOpts {
-	store := client.ContentStore()
+	return func(ctx context.Context, client *Client, c *containers.Container) error {
-	n, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointRuntimeName)
+		name, ok := index.Annotations[checkpointRuntimeNameLabel]
-	if err != nil {
+		if !ok {
-		if err != ErrMediaTypeNotFound {
+			return ErrRuntimeNameNotFoundInIndex
-			return nil, err
 		}
-	}
-	runtimeName := ""
-	if n != nil {
-		data, err := content.ReadBlob(ctx, store, *n)
-		if err != nil {
-			return nil, err
-		}
-		runtimeName = string(data)
-	}
 
-	// restore options if present
+		// restore options if present
-	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointRuntimeOptions)
+		m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointRuntimeOptions)
-	if err != nil {
-		if err != ErrMediaTypeNotFound {
-			return nil, err
-		}
-	}
-
-	var options *ptypes.Any
-	if m != nil {
-		data, err := content.ReadBlob(ctx, store, *m)
 		if err != nil {
-			return nil, errors.Wrap(err, "unable to read checkpoint runtime")
+			if err != ErrMediaTypeNotFound {
+				return err
+			}
 		}
-		if err := proto.Unmarshal(data, options); err != nil {
+		var options *ptypes.Any
-			return nil, err
+		if m != nil {
+			store := client.ContentStore()
+			data, err := content.ReadBlob(ctx, store, *m)
+			if err != nil {
+				return errors.Wrap(err, "unable to read checkpoint runtime")
+			}
+			if err := proto.Unmarshal(data, options); err != nil {
+				return err
+			}
 		}
+
+		c.Runtime = containers.RuntimeInfo{
+			Name:    name,
+			Options: options,
+		}
+		return nil
 	}
-	return []NewContainerOpts{
-		WithRuntime(runtimeName, options),
-	}, nil
 }
 
 // WithRestoreSpec restores the spec from the checkpoint for the container
-func WithRestoreSpec(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, error) {
+func WithRestoreSpec(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) NewContainerOpts {
-	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointConfig)
+	return func(ctx context.Context, client *Client, c *containers.Container) error {
-	if err != nil {
+		m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointConfig)
-		return nil, err
+		if err != nil {
+			return err
+		}
+		store := client.ContentStore()
+		data, err := content.ReadBlob(ctx, store, *m)
+		if err != nil {
+			return errors.Wrap(err, "unable to read checkpoint config")
+		}
+		var any ptypes.Any
+		if err := proto.Unmarshal(data, &any); err != nil {
+			return err
+		}
+		c.Spec = &any
+		return nil
 	}
-	store := client.ContentStore()
-	data, err := content.ReadBlob(ctx, store, *m)
-	if err != nil {
-		return nil, errors.Wrap(err, "unable to read checkpoint config")
-	}
-	var any ptypes.Any
-	if err := proto.Unmarshal(data, &any); err != nil {
-		return nil, err
-	}
-
-	v, err := typeurl.UnmarshalAny(&any)
-	if err != nil {
-		return nil, err
-	}
-	spec := v.(*oci.Spec)
-	return []NewContainerOpts{
-		WithSpec(spec),
-	}, nil
 }
 
 // WithRestoreSnapshot restores the snapshot from the checkpoint for the container
-func WithRestoreSnapshot(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, error) {
+func WithRestoreSnapshot(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) NewContainerOpts {
-	imageName := ""
+	return func(ctx context.Context, client *Client, c *containers.Container) error {
-	store := client.ContentStore()
+		imageName, ok := index.Annotations[checkpointImageNameLabel]
-	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointImageName)
+		if !ok || imageName == "" {
-	if err != nil {
+			return ErrRuntimeNameNotFoundInIndex
-		if err != ErrMediaTypeNotFound {
-			return nil, err
 		}
-	}
+		i, err := client.GetImage(ctx, imageName)
-	if m != nil {
-		data, err := content.ReadBlob(ctx, store, *m)
 		if err != nil {
-			return nil, err
+			return err
 		}
-		imageName = string(data)
+		diffIDs, err := i.(*image).i.RootFS(ctx, client.ContentStore(), platforms.Default())
+		if err != nil {
+			return err
+		}
+		parent := identity.ChainID(diffIDs).String()
+		if _, err := client.SnapshotService(DefaultSnapshotter).Prepare(ctx, id, parent); err != nil {
+			return err
+		}
+		c.SnapshotKey = id
+		c.Snapshotter = DefaultSnapshotter
+		return nil
+	}
+}
+
+// WithRestoreSnapshot restores the snapshot from the checkpoint for the container
+func WithRestoreRW(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) NewContainerOpts {
+	return func(ctx context.Context, client *Client, c *containers.Container) error {
+		// apply rw layer
+		rw, err := GetIndexByMediaType(index, imagespec.MediaTypeImageLayerGzip)
+		if err != nil {
+			return err
+		}
+		mounts, err := client.SnapshotService(c.Snapshotter).Mounts(ctx, c.SnapshotKey)
+		if err != nil {
+			return err
+		}
+
+		if _, err := client.DiffService().Apply(ctx, *rw, mounts); err != nil {
+			return err
+		}
+		return nil
 	}
-	i, err := client.GetImage(ctx, imageName)
-	if err != nil {
-		return nil, err
-	}
-
-	diffIDs, err := i.(*image).i.RootFS(ctx, client.ContentStore(), platforms.Default())
-	if err != nil {
-		return nil, err
-	}
-	parent := identity.ChainID(diffIDs).String()
-	if _, err := client.SnapshotService(DefaultSnapshotter).Prepare(ctx, id, parent); err != nil {
-		return nil, err
-	}
-	return []NewContainerOpts{
-		WithImage(i),
-		WithSnapshot(id),
-	}, nil
 }