refactor checkpoint and restore to client

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

client.go

@@ -17,7 +17,9 @@
 package containerd
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"runtime"
@@ -36,6 +38,7 @@ import (
 	snapshotsapi "github.com/containerd/containerd/api/services/snapshots/v1"
 	"github.com/containerd/containerd/api/services/tasks/v1"
 	versionservice "github.com/containerd/containerd/api/services/version/v1"
+	"github.com/containerd/containerd/cio"
 	"github.com/containerd/containerd/containers"
 	"github.com/containerd/containerd/content"
 	contentproxy "github.com/containerd/containerd/content/proxy"
@@ -46,6 +49,7 @@ import (
 	"github.com/containerd/containerd/leases"
 	leasesproxy "github.com/containerd/containerd/leases/proxy"
 	"github.com/containerd/containerd/namespaces"
+	"github.com/containerd/containerd/oci"
 	"github.com/containerd/containerd/pkg/dialer"
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/containerd/plugin"
@@ -520,6 +524,107 @@ func (c *Client) ListImages(ctx context.Context, filters ...string) ([]Image, er
 	return images, nil
 }
 
+// Restore restores a container from a checkpoint
+func (c *Client) Restore(ctx context.Context, id, ref string, opts ...RestoreOpts) error {
+	checkpoint, err := c.GetImage(ctx, ref)
+	if err != nil {
+		if !errdefs.IsNotFound(err) {
+			return err
+		}
+		ck, err := c.Fetch(ctx, ref)
+		if err != nil {
+			return err
+		}
+		checkpoint = NewImage(c, ck)
+	}
+
+	store := c.ContentStore()
+	index, err := decodeIndex(ctx, store, checkpoint.Target())
+	if err != nil {
+		return err
+	}
+
+	// get image from annotation
+	imageName, ok := index.Annotations["image.name"]
+	if !ok {
+		return ErrCheckpointIndexImageNameNotFound
+	}
+
+	image, err := c.Pull(ctx, imageName, WithPullUnpack)
+	if err != nil {
+		return err
+	}
+
+	ctx, done, err := c.WithLease(ctx)
+	if err != nil {
+		return err
+	}
+	defer done(ctx)
+
+	// container options
+	copts := []NewContainerOpts{
+		WithNewSpec(oci.WithImageConfig(image)),
+		WithNewSnapshot(id, image),
+	}
+	topts := []NewTaskOpts{}
+	for _, o := range opts {
+		co, to, err := o(ctx, c, checkpoint, index)
+		if err != nil {
+			return err
+		}
+		copts = append(copts, co...)
+		topts = append(topts, to...)
+	}
+
+	ctr, err := c.NewContainer(ctx, id, copts...)
+	if err != nil {
+		return err
+	}
+
+	// apply rw layer
+	info, err := ctr.Info(ctx)
+	if err != nil {
+		return err
+	}
+
+	rw, err := GetIndexByMediaType(index, ocispec.MediaTypeImageLayerGzip)
+	if err != nil {
+		return err
+	}
+
+	mounts, err := c.SnapshotService(info.Snapshotter).Mounts(ctx, info.SnapshotKey)
+	if err != nil {
+		return err
+	}
+
+	if _, err := c.DiffService().Apply(ctx, *rw, mounts); err != nil {
+		return err
+	}
+
+	task, err := ctr.NewTask(ctx, cio.NewCreator(cio.WithStdio), topts...)
+	if err != nil {
+		return err
+	}
+
+	if err := task.Start(ctx); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func writeIndex(ctx context.Context, index *ocispec.Index, client *Client, ref string) (d ocispec.Descriptor, err error) {
+	labels := map[string]string{}
+	for i, m := range index.Manifests {
+		labels[fmt.Sprintf("containerd.io/gc.ref.content.%d", i)] = m.Digest.String()
+	}
+	data, err := json.Marshal(index)
+	if err != nil {
+		return ocispec.Descriptor{}, err
+	}
+	return writeContent(ctx, client.ContentStore(), ocispec.MediaTypeImageIndex, ref, bytes.NewReader(data), content.WithLabels(labels))
+}
+
 // Subscribe to events that match one or more of the provided filters.
 //
 // Callers should listen on both the envelope and errs channels. If the errs

cmd/ctr/commands/containers/containers.go

@@ -45,6 +45,8 @@ var Command = cli.Command{
 		infoCommand,
 		listCommand,
 		setLabelsCommand,
+		checkpointCommand,
+		restoreCommand,
 	},
 }
 
@@ -282,3 +284,96 @@ var infoCommand = cli.Command{
 		return nil
 	},
 }
+
+var checkpointCommand = cli.Command{
+	Name:      "checkpoint",
+	Usage:     "checkpoint a container",
+	ArgsUsage: "CONTAINER REF [flags]",
+	Flags: []cli.Flag{
+		cli.BoolFlag{
+			Name:  "rw",
+			Usage: "include the rw layer in the checkpoint",
+		},
+		cli.BoolFlag{
+			Name:  "image",
+			Usage: "include the image in the checkpoint",
+		},
+		cli.BoolFlag{
+			Name:  "task",
+			Usage: "checkpoint container task",
+		},
+	},
+	Action: func(context *cli.Context) error {
+		id := context.Args().First()
+		if id == "" {
+			return errors.New("container id must be provided")
+		}
+		ref := context.Args()[1]
+		if ref == "" {
+			return errors.New("ref must be provided")
+		}
+		client, ctx, cancel, err := commands.NewClient(context)
+		if err != nil {
+			return err
+		}
+		defer cancel()
+		opts := []containerd.CheckpointOpts{}
+		if context.Bool("image") {
+			opts = append(opts, containerd.WithCheckpointImage)
+		}
+		if context.Bool("rw") {
+			opts = append(opts, containerd.WithCheckpointRW)
+		}
+		if context.Bool("task") {
+			opts = append(opts, containerd.WithCheckpointTask)
+		}
+		container, err := client.LoadContainer(ctx, id)
+		if err != nil {
+			return err
+		}
+		if _, err := container.Checkpoint(ctx, ref, opts...); err != nil {
+			return err
+		}
+
+		return nil
+	},
+}
+
+var restoreCommand = cli.Command{
+	Name:      "restore",
+	Usage:     "restore a container from checkpoint",
+	ArgsUsage: "CONTAINER REF [flags]",
+	Flags: []cli.Flag{
+		cli.BoolFlag{
+			Name:  "live",
+			Usage: "restore the runtime and memory data from the checkpoint",
+		},
+	},
+	Action: func(context *cli.Context) error {
+		id := context.Args().First()
+		if id == "" {
+			return errors.New("container id must be provided")
+		}
+		ref := context.Args()[1]
+		if ref == "" {
+			return errors.New("ref must be provided")
+		}
+		opts := []containerd.RestoreOpts{
+			containerd.WithRestoreSpec,
+			containerd.WithRestoreRuntime,
+		}
+		if context.Bool("live") {
+			opts = append(opts, containerd.WithRestoreLive)
+		}
+		client, ctx, cancel, err := commands.NewClient(context)
+		if err != nil {
+			return err
+		}
+		defer cancel()
+		if err := client.Restore(ctx, id, ref, opts...); err != nil {
+			return err
+		}
+
+		return nil
+	},
+}

container.go

@@ -17,10 +17,12 @@
 package containerd
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
 
 	"github.com/containerd/containerd/api/services/tasks/v1"
@@ -28,9 +30,13 @@ import (
 	"github.com/containerd/containerd/cio"
 	"github.com/containerd/containerd/containers"
 	"github.com/containerd/containerd/errdefs"
+	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/oci"
+	"github.com/containerd/containerd/runtime/v2/runc/options"
 	"github.com/containerd/typeurl"
 	prototypes "github.com/gogo/protobuf/types"
+	ver "github.com/opencontainers/image-spec/specs-go"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 )
 
@@ -64,6 +70,8 @@ type Container interface {
 	Extensions(context.Context) (map[string]prototypes.Any, error)
 	// Update a container
 	Update(context.Context, ...UpdateContainerOpts) error
+	// Checkpoint creates a checkpoint image of the current container
+	Checkpoint(context.Context, string, ...CheckpointOpts) (Image, error)
 }
 
 func containerFromRecord(client *Client, c containers.Container) *container {
@@ -272,6 +280,98 @@ func (c *container) Update(ctx context.Context, opts ...UpdateContainerOpts) err
 	return nil
 }
 
+func (c *container) Checkpoint(ctx context.Context, ref string, opts ...CheckpointOpts) (Image, error) {
+	index := &ocispec.Index{
+		Versioned: ver.Versioned{
+			SchemaVersion: 2,
+		},
+		Annotations: make(map[string]string),
+	}
+	copts := &options.CheckpointOptions{
+		Exit:                false,
+		OpenTcp:             false,
+		ExternalUnixSockets: false,
+		Terminal:            false,
+		FileLocks:           true,
+		EmptyNamespaces:     nil,
+	}
+	img, err := c.Image(ctx)
+	if err != nil {
+		return nil, err
+	}
+	index.Annotations["image.name"] = img.Name()
+
+	ctx, done, err := c.client.WithLease(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer done(ctx)
+
+	// pause task to checkpoint
+	if err := func(ctx context.Context) error {
+		task, err := c.Task(ctx, nil)
+		if err != nil {
+			if errdefs.IsNotFound(err) {
+				return nil
+			}
+			return err
+		}
+		if err := task.Pause(ctx); err != nil {
+			return err
+		}
+		defer task.Resume(ctx)
+
+		info, err := c.Info(ctx)
+		if err != nil {
+			return err
+		}
+
+		// add runtime info to index
+		index.Annotations["runtime.name"] = info.Runtime.Name
+		if info.Runtime.Options != nil {
+			data, err := info.Runtime.Options.Marshal()
+			if err != nil {
+				return err
+			}
+			r := bytes.NewReader(data)
+			desc, err := writeContent(ctx, c.client.ContentStore(), images.MediaTypeContainerd1CheckpointRuntimeOptions, info.ID+"-runtime-options", r)
+			if err != nil {
+				return err
+			}
+			desc.Platform = &ocispec.Platform{
+				OS:           runtime.GOOS,
+				Architecture: runtime.GOARCH,
+			}
+			index.Manifests = append(index.Manifests, desc)
+		}
+
+		// process remaining opts
+		for _, o := range opts {
+			if err := o(ctx, c.client, &info, index, copts); err != nil {
+				return err
+			}
+		}
+		return nil
+	}(ctx); err != nil {
+		return nil, err
+	}
+
+	desc, err := writeIndex(ctx, index, c.client, c.ID()+"index")
+	if err != nil {
+		return nil, err
+	}
+	i := images.Image{
+		Name:   ref,
+		Target: desc,
+	}
+	checkpoint, err := c.client.ImageService().Create(ctx, i)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewImage(c.client, checkpoint), nil
+}
+
 func (c *container) loadTask(ctx context.Context, ioAttach cio.Attach) (Task, error) {
 	response, err := c.client.TaskService().Get(ctx, &tasks.GetRequest{
 		ContainerID: c.id,

container_checkpoint_opts.go

@@ -0,0 +1,125 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package containerd
+
+import (
+	"context"
+	"fmt"
+	"runtime"
+
+	tasks "github.com/containerd/containerd/api/services/tasks/v1"
+	"github.com/containerd/containerd/containers"
+	"github.com/containerd/containerd/diff"
+	"github.com/containerd/containerd/rootfs"
+	"github.com/containerd/containerd/runtime/v2/runc/options"
+	"github.com/containerd/typeurl"
+	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/pkg/errors"
+)
+
+var (
+	// ErrCheckpointRWUnsupported is returned if the container runtime does not support checkpoint
+	ErrCheckpointRWUnsupported = errors.New("rw checkpoint is only supported on v2 runtimes")
+	// ErrMediaTypeNotFound returns an error when a media type in the manifest is unknown
+	ErrMediaTypeNotFound = errors.New("media type not found")
+	// ErrCheckpointIndexImageNameNotFound is returned when the checkpoint image name is not present in the index
+	ErrCheckpointIndexImageNameNotFound = errors.New("image name not present in index")
+	// ErrCheckpointIndexRuntimeNameNotFound is returned when the checkpoint runtime name is not present in the index
+	ErrCheckpointIndexRuntimeNameNotFound = errors.New("runtime name not present in index")
+)
+
+// CheckpointOpts are options to manage the checkpoint operation
+type CheckpointOpts func(context.Context, *Client, *containers.Container, *imagespec.Index, *options.CheckpointOptions) error
+
+// WithCheckpointImage includes the container image in the checkpoint
+func WithCheckpointImage(ctx context.Context, client *Client, c *containers.Container, index *imagespec.Index, copts *options.CheckpointOptions) error {
+	ir, err := client.ImageService().Get(ctx, c.Image)
+	if err != nil {
+		return err
+	}
+	index.Manifests = append(index.Manifests, ir.Target)
+	return nil
+}
+
+// WithCheckpointTask includes the running task
+func WithCheckpointTask(ctx context.Context, client *Client, c *containers.Container, index *imagespec.Index, copts *options.CheckpointOptions) error {
+	any, err := typeurl.MarshalAny(copts)
+	if err != nil {
+		return nil
+	}
+	task, err := client.TaskService().Checkpoint(ctx, &tasks.CheckpointTaskRequest{
+		ContainerID: c.ID,
+		Options:     any,
+	})
+	if err != nil {
+		return err
+	}
+	for _, d := range task.Descriptors {
+		index.Manifests = append(index.Manifests, imagespec.Descriptor{
+			MediaType: d.MediaType,
+			Size:      d.Size_,
+			Digest:    d.Digest,
+			Platform: &imagespec.Platform{
+				OS:           runtime.GOOS,
+				Architecture: runtime.GOARCH,
+			},
+		})
+	}
+	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)),
+	}
+	rw, err := rootfs.CreateDiff(ctx,
+		info.SnapshotKey,
+		client.SnapshotService(info.Snapshotter),
+		client.DiffService(),
+		diffOpts...,
+	)
+	if err != nil {
+		return err
+
+	}
+	rw.Platform = &imagespec.Platform{
+		OS:           runtime.GOOS,
+		Architecture: runtime.GOARCH,
+	}
+	index.Manifests = append(index.Manifests, rw)
+	return nil
+}
+
+// GetIndexByMediaType returns the index in a manifest for the specified media type
+func GetIndexByMediaType(index *imagespec.Index, mt string) (*imagespec.Descriptor, error) {
+	for _, d := range index.Manifests {
+		if d.MediaType == mt {
+			return &d, nil
+		}
+	}
+	return nil, ErrMediaTypeNotFound
+}

container_restore_opts.go

@@ -0,0 +1,95 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package containerd
+
+import (
+	"context"
+
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/oci"
+	"github.com/containerd/typeurl"
+	"github.com/gogo/protobuf/proto"
+	ptypes "github.com/gogo/protobuf/types"
+	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/pkg/errors"
+)
+
+// RestoreOpts are options to manage the restore operation
+type RestoreOpts func(context.Context, *Client, Image, *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error)
+
+// WithRestoreLive restores the runtime and memory data for the container
+func WithRestoreLive(ctx context.Context, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+	return nil, []NewTaskOpts{
+		WithTaskCheckpoint(checkpoint),
+	}, nil
+}
+
+// WithRestoreRuntime restores the runtime for the container
+func WithRestoreRuntime(ctx context.Context, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+	runtimeName, ok := index.Annotations["runtime.name"]
+	if !ok {
+		return nil, nil, ErrCheckpointIndexRuntimeNameNotFound
+	}
+	// restore options if present
+	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointRuntimeOptions)
+	if err != nil {
+		if err != ErrMediaTypeNotFound {
+			return nil, nil, err
+		}
+	}
+	var options *ptypes.Any
+	if m != nil {
+		store := client.ContentStore()
+		data, err := content.ReadBlob(ctx, store, *m)
+		if err != nil {
+			return nil, nil, errors.Wrap(err, "unable to read checkpoint runtime")
+		}
+		if err := proto.Unmarshal(data, options); err != nil {
+			return nil, nil, err
+		}
+	}
+	return []NewContainerOpts{
+		WithRuntime(runtimeName, options),
+	}, nil, nil
+}
+
+// WithRestoreSpec restores the spec from the checkpoint for the container
+func WithRestoreSpec(ctx context.Context, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointConfig)
+	if err != nil {
+		return nil, nil, err
+	}
+	store := client.ContentStore()
+	data, err := content.ReadBlob(ctx, store, *m)
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "unable to read checkpoint config")
+	}
+	var any ptypes.Any
+	if err := proto.Unmarshal(data, &any); err != nil {
+		return nil, nil, err
+	}
+
+	v, err := typeurl.UnmarshalAny(&any)
+	if err != nil {
+		return nil, nil, err
+	}
+	spec := v.(*oci.Spec)
+	return []NewContainerOpts{
+		WithSpec(spec),
+	}, nil, nil
+}

images/mediatypes.go

@@ -34,6 +34,7 @@ const (
 	MediaTypeContainerd1Resource                 = "application/vnd.containerd.container.resource.tar"
 	MediaTypeContainerd1RW                       = "application/vnd.containerd.container.rw.tar"
 	MediaTypeContainerd1CheckpointConfig         = "application/vnd.containerd.container.checkpoint.config.v1+proto"
+	MediaTypeContainerd1CheckpointRuntimeOptions = "application/vnd.containerd.container.checkpoint.runtime.options+proto"
 	// Legacy Docker schema1 manifest
 	MediaTypeDockerSchema1Manifest = "application/vnd.docker.distribution.manifest.v1+prettyjws"
 )