Finish snapshot support.

Signed-off-by: Lantao Liu <lantaol@google.com>

pkg/server/container_create.go

@@ -20,9 +20,7 @@ import (
 	"fmt"
 	"time"
 
-	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
 	"github.com/golang/glog"
-	imagedigest "github.com/opencontainers/go-digest"
 	"golang.org/x/net/context"
 	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
 
@@ -79,15 +77,22 @@ func (c *criContainerdService) CreateContainer(ctx context.Context, r *runtime.C
 	if imageMeta == nil {
 		return nil, fmt.Errorf("image %q not found", image)
 	}
-	if _, err := c.snapshotService.Prepare(ctx, &snapshotapi.PrepareRequest{
+	if config.GetLinux().GetSecurityContext().GetReadonlyRootfs() {
-		Key: id,
+		if _, err := c.snapshotService.View(ctx, id, imageMeta.ChainID); err != nil {
-		// We are sure that ChainID must be a digest.
+			return nil, fmt.Errorf("failed to view container rootfs %q: %v", imageMeta.ChainID, err)
-		Parent: imagedigest.Digest(imageMeta.ChainID).String(),
+		}
-		//Readonly: config.GetLinux().GetSecurityContext().GetReadonlyRootfs(),
+	} else {
-	}); err != nil {
+		if _, err := c.snapshotService.Prepare(ctx, id, imageMeta.ChainID); err != nil {
 			return nil, fmt.Errorf("failed to prepare container rootfs %q: %v", imageMeta.ChainID, err)
 		}
-	// TODO(random-liu): [P0] Cleanup snapshot on failure after switching to new snapshot api.
+	}
+	defer func() {
+		if retErr != nil {
+			if err := c.snapshotService.Remove(ctx, id); err != nil {
+				glog.Errorf("Failed to remove container snapshot %q: %v", id, err)
+			}
+		}
+	}()
 	meta.ImageRef = imageMeta.ID
 
 	// Create container root directory.

pkg/server/container_create_test.go

@@ -134,7 +134,7 @@ func TestCreateContainer(t *testing.T) {
 	} {
 		t.Logf("TestCase %q", desc)
 		c := newTestCRIContainerdService()
-		fakeSnapshotClient := c.snapshotService.(*servertesting.FakeSnapshotClient)
+		fakeSnapshotClient := WithFakeSnapshotClient(c)
 		fakeOS := c.os.(*ostesting.FakeOS)
 		if test.sandboxMetadata != nil {
 			assert.NoError(t, c.sandboxStore.Create(*test.sandboxMetadata))
@@ -177,6 +177,7 @@ func TestCreateContainer(t *testing.T) {
 				assert.NoError(t, c.containerNameIndex.Reserve(containerName, "random id"),
 					"container name should be released")
 			}
+			assert.Empty(t, fakeSnapshotClient.ListMounts(), "snapshot should be cleaned up")
 			metas, err := c.containerStore.List()
 			assert.NoError(t, err)
 			assert.Empty(t, metas, "container metadata should not be created")

pkg/server/container_remove.go

@@ -19,9 +19,9 @@ package server
 import (
 	"fmt"
 
+	"github.com/containerd/containerd/snapshot"
 	"github.com/golang/glog"
 	"golang.org/x/net/context"
-
 	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
 
 	"github.com/kubernetes-incubator/cri-containerd/pkg/metadata"
@@ -69,7 +69,13 @@ func (c *criContainerdService) RemoveContainer(ctx context.Context, r *runtime.R
 	// kubelet implementation, we'll never start a container once we decide to remove it,
 	// so we don't need the "Dead" state for now.
 
-	// TODO(random-liu): [P0] Cleanup snapshot after switching to new snapshot api.
+	// Remove container snapshot.
+	if err := c.snapshotService.Remove(ctx, id); err != nil {
+		if !snapshot.IsNotExist(err) {
+			return nil, fmt.Errorf("failed to remove container snapshot %q: %v", id, err)
+		}
+		glog.V(5).Infof("Remove called for snapshot %q that does not exist", id)
+	}
 
 	// Cleanup container root directory.
 	containerRootDir := getContainerRootDir(c.rootDir, id)

pkg/server/container_remove_test.go

@@ -21,6 +21,8 @@ import (
 	"testing"
 	"time"
 
+	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
+	"github.com/containerd/containerd/api/types/mount"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context"
@@ -28,6 +30,7 @@ import (
 
 	"github.com/kubernetes-incubator/cri-containerd/pkg/metadata"
 	ostesting "github.com/kubernetes-incubator/cri-containerd/pkg/os/testing"
+	servertesting "github.com/kubernetes-incubator/cri-containerd/pkg/server/testing"
 )
 
 // TestSetContainerRemoving tests setContainerRemoving sets removing
@@ -87,8 +90,16 @@ func TestSetContainerRemoving(t *testing.T) {
 func TestRemoveContainer(t *testing.T) {
 	testID := "test-id"
 	testName := "test-name"
+	testContainerMetadata := &metadata.ContainerMetadata{
+		ID:         testID,
+		CreatedAt:  time.Now().UnixNano(),
+		StartedAt:  time.Now().UnixNano(),
+		FinishedAt: time.Now().UnixNano(),
+	}
+
 	for desc, test := range map[string]struct {
 		metadata            *metadata.ContainerMetadata
+		removeSnapshotErr   error
 		removeDirErr        error
 		expectErr           bool
 		expectUnsetRemoving bool
@@ -113,31 +124,33 @@ func TestRemoveContainer(t *testing.T) {
 		},
 		"should not return error if container does not exist": {
 			metadata:          nil,
+			removeSnapshotErr: servertesting.SnapshotNotExistError,
 			expectErr:         false,
 		},
-		"should return error if remove container root fails": {
+		"should not return error if snapshot does not exist": {
-			metadata: &metadata.ContainerMetadata{
+			metadata:          testContainerMetadata,
-				ID:         testID,
+			removeSnapshotErr: servertesting.SnapshotNotExistError,
-				CreatedAt:  time.Now().UnixNano(),
+			expectErr:         false,
-				StartedAt:  time.Now().UnixNano(),
-				FinishedAt: time.Now().UnixNano(),
 		},
+		"should return error if remove snapshot fails": {
+			metadata:          testContainerMetadata,
+			removeSnapshotErr: errors.New("random error"),
+			expectErr:         true,
+		},
+		"should return error if remove container root fails": {
+			metadata:            testContainerMetadata,
 			removeDirErr:        errors.New("random error"),
 			expectErr:           true,
 			expectUnsetRemoving: true,
 		},
 		"should be able to remove container successfully": {
-			metadata: &metadata.ContainerMetadata{
+			metadata:  testContainerMetadata,
-				ID:         testID,
-				CreatedAt:  time.Now().UnixNano(),
-				StartedAt:  time.Now().UnixNano(),
-				FinishedAt: time.Now().UnixNano(),
-			},
 			expectErr: false,
 		},
 	} {
 		t.Logf("TestCase %q", desc)
 		c := newTestCRIContainerdService()
+		fakeSnapshotClient := WithFakeSnapshotClient(c)
 		fakeOS := c.os.(*ostesting.FakeOS)
 		if test.metadata != nil {
 			assert.NoError(t, c.containerNameIndex.Reserve(testName, testID))
@@ -147,6 +160,17 @@ func TestRemoveContainer(t *testing.T) {
 			assert.Equal(t, getContainerRootDir(c.rootDir, testID), path)
 			return test.removeDirErr
 		}
+		if test.removeSnapshotErr == nil {
+			fakeSnapshotClient.SetFakeMounts(testID, []*mount.Mount{
+				{
+					Type:   "bind",
+					Source: "/test/source",
+					Target: "/test/target",
+				},
+			})
+		} else {
+			fakeSnapshotClient.InjectError("remove", test.removeSnapshotErr)
+		}
 		resp, err := c.RemoveContainer(context.Background(), &runtime.RemoveContainerRequest{
 			ContainerId: testID,
 		})
@@ -171,5 +195,8 @@ func TestRemoveContainer(t *testing.T) {
 		assert.Nil(t, meta, "container metadata should be removed")
 		assert.NoError(t, c.containerNameIndex.Reserve(testName, testID),
 			"container name should be released")
+		mountsResp, err := fakeSnapshotClient.Mounts(context.Background(), &snapshotapi.MountsRequest{Key: testID})
+		assert.Equal(t, servertesting.SnapshotNotExistError, err, "snapshot should be removed")
+		assert.Nil(t, mountsResp)
 	}
 }

pkg/server/container_start.go

@@ -26,7 +26,7 @@ import (
 	"time"
 
 	"github.com/containerd/containerd/api/services/execution"
-	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
+	"github.com/containerd/containerd/api/types/mount"
 	"github.com/containerd/containerd/api/types/task"
 	"github.com/golang/glog"
 	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -176,15 +176,23 @@ func (c *criContainerdService) startContainer(ctx context.Context, id string, me
 	}
 
 	// Get rootfs mounts.
-	mountsResp, err := c.snapshotService.Mounts(ctx, &snapshotapi.MountsRequest{Key: id})
+	rootfsMounts, err := c.snapshotService.Mounts(ctx, id)
 	if err != nil {
 		return fmt.Errorf("failed to get rootfs mounts %q: %v", id, err)
 	}
+	var rootfs []*mount.Mount
+	for _, m := range rootfsMounts {
+		rootfs = append(rootfs, &mount.Mount{
+			Type:    m.Type,
+			Source:  m.Source,
+			Options: m.Options,
+		})
+	}
 
 	// Create containerd container.
 	createOpts := &execution.CreateRequest{
 		ContainerID: id,
-		Rootfs:      mountsResp.Mounts,
+		Rootfs:      rootfs,
 		Stdin:       stdin,
 		Stdout:      stdout,
 		Stderr:      stderr,

pkg/server/container_start_test.go

@@ -590,7 +590,7 @@ func TestStartContainer(t *testing.T) {
 		c := newTestCRIContainerdService()
 		fake := c.containerService.(*servertesting.FakeExecutionClient)
 		fakeOS := c.os.(*ostesting.FakeOS)
-		fakeSnapshotClient := c.snapshotService.(*servertesting.FakeSnapshotClient)
+		fakeSnapshotClient := WithFakeSnapshotClient(c)
 		if test.containerMetadata != nil {
 			assert.NoError(t, c.containerStore.Create(*test.containerMetadata))
 		}

pkg/server/image_pull.go

@@ -24,11 +24,11 @@ import (
 	"sync"
 	"time"
 
-	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
 	"github.com/containerd/containerd/content"
 	containerdimages "github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/remotes/docker"
+	containerdrootfs "github.com/containerd/containerd/rootfs"
 	"github.com/golang/glog"
 	imagedigest "github.com/opencontainers/go-digest"
 	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -187,6 +187,7 @@ func (r *resourceSet) all() map[string]struct{} {
 // pullImage pulls image and returns image id (config digest) and manifest digest.
 // The ref should be normalized image reference.
 func (c *criContainerdService) pullImage(ctx context.Context, ref string) (
+	// TODO(random-liu): Replace with client.Pull.
 	imagedigest.Digest, imagedigest.Digest, error) {
 	// Resolve the image reference to get descriptor and fetcher.
 	resolver := docker.NewResolver(docker.ResolverOptions{
@@ -250,6 +251,8 @@ func (c *criContainerdService) pullImage(ctx context.Context, ref string) (
 	}
 	glog.V(4).Infof("Finished downloading resources for image %q", ref)
 
+	// TODO(random-liu): Replace with image.Unpack.
+	// Unpack the image layers into snapshots.
 	image, err := c.imageStoreService.Get(ctx, ref)
 	if err != nil {
 		return "", "", fmt.Errorf("failed to get image %q from containerd image store: %v", ref, err)
@@ -265,18 +268,24 @@ func (c *criContainerdService) pullImage(ctx context.Context, ref string) (
 		return "", "", fmt.Errorf("unmarshal blob to manifest failed for manifest digest %q: %v",
 			manifestDigest, err)
 	}
-
+	diffIDs, err := image.RootFS(ctx, c.contentStoreService)
-	// Unpack the image layers into snapshots.
-	/* snapshotUnpacker := snapshotservice.NewUnpackerFromClient(c.snapshotService)
-	if _, err = snapshotUnpacker.Unpack(ctx, manifest.Layers); err != nil {
-		return "", "", fmt.Errorf("unpack failed for manifest layers %+v: %v", manifest.Layers, err)
-	} TODO(mikebrow): WIP replacing the commented Unpack with the below Prepare request */
-	_, err = image.RootFS(ctx, c.contentStoreService)
 	if err != nil {
 		return "", "", err
 	}
-	if _, err = c.snapshotService.Prepare(ctx, &snapshotapi.PrepareRequest{Key: ref, Parent: ""}); err != nil {
+	if len(diffIDs) != len(manifest.Layers) {
-		return "", "", err
+		return "", "", fmt.Errorf("mismatched image rootfs and manifest layers")
+	}
+	layers := make([]containerdrootfs.Layer, len(diffIDs))
+	for i := range diffIDs {
+		layers[i].Diff = imagespec.Descriptor{
+			// TODO: derive media type from compressed type
+			MediaType: imagespec.MediaTypeImageLayer,
+			Digest:    diffIDs[i],
+		}
+		layers[i].Blob = manifest.Layers[i]
+	}
+	if _, err = containerdrootfs.ApplyLayers(ctx, layers, c.snapshotService, c.diffService); err != nil {
+		return "", "", fmt.Errorf("failed to apply layers %+v: %v", layers, err)
 	}
 
 	// TODO(random-liu): Considering how to deal with the disk usage of content.

pkg/server/sandbox_remove.go

@@ -19,11 +19,10 @@ package server
 import (
 	"fmt"
 
+	"github.com/containerd/containerd/api/services/execution"
+	"github.com/containerd/containerd/snapshot"
 	"github.com/golang/glog"
 	"golang.org/x/net/context"
-
-	"github.com/containerd/containerd/api/services/execution"
-
 	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
 
 	"github.com/kubernetes-incubator/cri-containerd/pkg/metadata"
@@ -65,7 +64,14 @@ func (c *criContainerdService) RemovePodSandbox(ctx context.Context, r *runtime.
 		return nil, fmt.Errorf("sandbox container %q is not fully stopped", id)
 	}
 
-	// TODO(random-liu): [P0] Cleanup snapshot after switching to new snapshot api.
+	// Remove sandbox container snapshot.
+	if err := c.snapshotService.Remove(ctx, id); err != nil {
+		if !snapshot.IsNotExist(err) {
+			return nil, fmt.Errorf("failed to remove sandbox container snapshot %q: %v", id, err)
+		}
+		glog.V(5).Infof("Remove called for snapshot %q that does not exist", id)
+	}
+
 	// TODO(random-liu): [P0] Cleanup shm created in RunPodSandbox.
 	// TODO(random-liu): [P1] Remove permanent namespace once used.
 

pkg/server/sandbox_remove_test.go

@@ -1,6 +1,4 @@
-/*
+/* Copyright 2017 The Kubernetes Authors.
-Copyright 2017 The Kubernetes 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
@@ -20,16 +18,16 @@ import (
 	"fmt"
 	"testing"
 
+	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
+	"github.com/containerd/containerd/api/types/mount"
+	"github.com/containerd/containerd/api/types/task"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/net/context"
-
+	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
-	"github.com/containerd/containerd/api/types/task"
 
 	"github.com/kubernetes-incubator/cri-containerd/pkg/metadata"
 	ostesting "github.com/kubernetes-incubator/cri-containerd/pkg/os/testing"
 	servertesting "github.com/kubernetes-incubator/cri-containerd/pkg/server/testing"
-
-	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
 )
 
 func TestRemovePodSandbox(t *testing.T) {
@@ -42,6 +40,7 @@ func TestRemovePodSandbox(t *testing.T) {
 	for desc, test := range map[string]struct {
 		sandboxContainers   []task.Task
 		injectMetadata      bool
+		removeSnapshotErr   error
 		injectContainerdErr error
 		injectFSErr         error
 		expectErr           bool
@@ -50,9 +49,22 @@ func TestRemovePodSandbox(t *testing.T) {
 	}{
 		"should not return error if sandbox does not exist": {
 			injectMetadata:    false,
+			removeSnapshotErr: servertesting.SnapshotNotExistError,
 			expectErr:         false,
 			expectCalls:       []string{},
 		},
+		"should not return error if snapshot does not exist": {
+			injectMetadata:    true,
+			removeSnapshotErr: servertesting.SnapshotNotExistError,
+			expectRemoved:     getSandboxRootDir(testRootDir, testID),
+			expectCalls:       []string{"info"},
+		},
+		"should return error if remove snapshot fails": {
+			injectMetadata:    true,
+			removeSnapshotErr: fmt.Errorf("arbitrary error"),
+			expectErr:         true,
+			expectCalls:       []string{"info"},
+		},
 		"should return error when sandbox container is not deleted": {
 			injectMetadata:    true,
 			sandboxContainers: []task.Task{{ID: testID}},
@@ -82,12 +94,24 @@ func TestRemovePodSandbox(t *testing.T) {
 		c := newTestCRIContainerdService()
 		fake := c.containerService.(*servertesting.FakeExecutionClient)
 		fakeOS := c.os.(*ostesting.FakeOS)
+		fakeSnapshotClient := WithFakeSnapshotClient(c)
 		fake.SetFakeContainers(test.sandboxContainers)
 		if test.injectMetadata {
 			c.sandboxNameIndex.Reserve(testName, testID)
 			c.sandboxIDIndex.Add(testID)
 			c.sandboxStore.Create(testMetadata)
 		}
+		if test.removeSnapshotErr == nil {
+			fakeSnapshotClient.SetFakeMounts(testID, []*mount.Mount{
+				{
+					Type:   "bind",
+					Source: "/test/source",
+					Target: "/test/target",
+				},
+			})
+		} else {
+			fakeSnapshotClient.InjectError("remove", test.removeSnapshotErr)
+		}
 		if test.injectContainerdErr != nil {
 			fake.InjectError("info", test.injectContainerdErr)
 		}
@@ -114,6 +138,9 @@ func TestRemovePodSandbox(t *testing.T) {
 		assert.Error(t, err)
 		assert.True(t, metadata.IsNotExistError(err))
 		assert.Nil(t, meta, "sandbox metadata should be removed")
+		mountsResp, err := fakeSnapshotClient.Mounts(context.Background(), &snapshotapi.MountsRequest{Key: testID})
+		assert.Equal(t, servertesting.SnapshotNotExistError, err, "snapshot should be removed")
+		assert.Nil(t, mountsResp)
 		res, err = c.RemovePodSandbox(context.Background(), &runtime.RemovePodSandboxRequest{
 			PodSandboxId: testID,
 		})

pkg/server/sandbox_run.go

@@ -24,9 +24,8 @@ import (
 	"time"
 
 	"github.com/containerd/containerd/api/services/execution"
-	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
+	"github.com/containerd/containerd/api/types/mount"
 	"github.com/golang/glog"
-	imagedigest "github.com/opencontainers/go-digest"
 	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
 	runtimespec "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/opencontainers/runtime-tools/generate"
@@ -86,17 +85,25 @@ func (c *criContainerdService) RunPodSandbox(ctx context.Context, r *runtime.Run
 	if err != nil {
 		return nil, fmt.Errorf("failed to get sandbox image %q: %v", defaultSandboxImage, err)
 	}
-	prepareResp, err := c.snapshotService.Prepare(ctx, &snapshotapi.PrepareRequest{
+	rootfsMounts, err := c.snapshotService.View(ctx, id, imageMeta.ChainID)
-		Key: id,
-		// We are sure that ChainID must be a digest.
-		Parent: imagedigest.Digest(imageMeta.ChainID).String(),
-		//Readonly: true,
-	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to prepare sandbox rootfs %q: %v", imageMeta.ChainID, err)
 	}
-	// TODO(random-liu): [P0] Cleanup snapshot on failure after switching to new snapshot api.
+	defer func() {
-	rootfsMounts := prepareResp.Mounts
+		if retErr != nil {
+			if err := c.snapshotService.Remove(ctx, id); err != nil {
+				glog.Errorf("Failed to remove sandbox container snapshot %q: %v", id, err)
+			}
+		}
+	}()
+	var rootfs []*mount.Mount
+	for _, m := range rootfsMounts {
+		rootfs = append(rootfs, &mount.Mount{
+			Type:    m.Type,
+			Source:  m.Source,
+			Options: m.Options,
+		})
+	}
 
 	// Create sandbox container root directory.
 	// Prepare streaming named pipe.
@@ -156,7 +163,7 @@ func (c *criContainerdService) RunPodSandbox(ctx context.Context, r *runtime.Run
 	glog.V(4).Infof("Sandbox container spec: %+v", spec)
 	createOpts := &execution.CreateRequest{
 		ContainerID: id,
-		Rootfs:      rootfsMounts,
+		Rootfs:      rootfs,
 		// No stdin for sandbox container.
 		Stdout: stdout,
 		Stderr: stderr,

pkg/server/sandbox_run_test.go

@@ -268,7 +268,7 @@ options timeout:1
 func TestRunPodSandbox(t *testing.T) {
 	config, imageConfig, _ := getRunPodSandboxTestData() // TODO: declare and test specCheck see below
 	c := newTestCRIContainerdService()
-	fakeSnapshotClient := c.snapshotService.(*servertesting.FakeSnapshotClient)
+	fakeSnapshotClient := WithFakeSnapshotClient(c)
 	fakeExecutionClient := c.containerService.(*servertesting.FakeExecutionClient)
 	fakeCNIPlugin := c.netPlugin.(*servertesting.FakeCNIPlugin)
 	fakeOS := c.os.(*ostesting.FakeOS)
@@ -292,7 +292,7 @@ func TestRunPodSandbox(t *testing.T) {
 	}
 	// Insert sandbox image metadata.
 	assert.NoError(t, c.imageMetadataStore.Create(imageMetadata))
-	expectSnapshotClientCalls := []string{"prepare"}
+	expectSnapshotClientCalls := []string{"view"}
 	expectExecutionClientCalls := []string{"create", "start"}
 
 	res, err := c.RunPodSandbox(context.Background(), &runtime.RunPodSandboxRequest{Config: config})

pkg/server/service.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 
 	contentapi "github.com/containerd/containerd/api/services/content"
+	diffapi "github.com/containerd/containerd/api/services/diff"
 	"github.com/containerd/containerd/api/services/execution"
 	imagesapi "github.com/containerd/containerd/api/services/images"
 	snapshotapi "github.com/containerd/containerd/api/services/snapshot"
@@ -27,7 +28,10 @@ import (
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/images"
 	contentservice "github.com/containerd/containerd/services/content"
+	diffservice "github.com/containerd/containerd/services/diff"
 	imagesservice "github.com/containerd/containerd/services/images"
+	snapshotservice "github.com/containerd/containerd/services/snapshot"
+	"github.com/containerd/containerd/snapshot"
 	"github.com/docker/docker/pkg/truncindex"
 	"github.com/kubernetes-incubator/cri-o/pkg/ocicni"
 	"google.golang.org/grpc"
@@ -78,7 +82,9 @@ type criContainerdService struct {
 	// contentStoreService is the containerd content service client.
 	contentStoreService content.Store
 	// snapshotService is the containerd snapshot service client.
-	snapshotService snapshotapi.SnapshotClient
+	snapshotService snapshot.Snapshotter
+	// diffService is the containerd diff service client.
+	diffService diffservice.DiffService
 	// imageStoreService is the containerd service to store and track
 	// image metadata.
 	imageStoreService images.Store
@@ -111,7 +117,8 @@ func NewCRIContainerdService(conn *grpc.ClientConn, rootDir, networkPluginBinDir
 		containerService:    execution.NewTasksClient(conn),
 		imageStoreService:   imagesservice.NewStoreFromClient(imagesapi.NewImagesClient(conn)),
 		contentStoreService: contentservice.NewStoreFromClient(contentapi.NewContentClient(conn)),
-		snapshotService:     snapshotapi.NewSnapshotClient(conn),
+		snapshotService:     snapshotservice.NewSnapshotterFromClient(snapshotapi.NewSnapshotClient(conn)),
+		diffService:         diffservice.NewDiffServiceFromClient(diffapi.NewDiffClient(conn)),
 		versionService:      versionapi.NewVersionClient(conn),
 		healthService:       healthapi.NewHealthClient(conn),
 		agentFactory:        agents.NewAgentFactory(),

pkg/server/service_test.go

@@ -22,10 +22,13 @@ import (
 	"testing"
 
 	"github.com/containerd/containerd/api/services/execution"
+	snapshotservice "github.com/containerd/containerd/services/snapshot"
 	"github.com/docker/docker/pkg/truncindex"
+	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context"
+	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
 
 	"github.com/kubernetes-incubator/cri-containerd/pkg/metadata"
 	"github.com/kubernetes-incubator/cri-containerd/pkg/metadata/store"
@@ -33,9 +36,6 @@ import (
 	"github.com/kubernetes-incubator/cri-containerd/pkg/registrar"
 	agentstesting "github.com/kubernetes-incubator/cri-containerd/pkg/server/agents/testing"
 	servertesting "github.com/kubernetes-incubator/cri-containerd/pkg/server/testing"
-	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
-
-	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/v1alpha1"
 )
 
 type nopReadWriteCloser struct{}
@@ -66,20 +66,25 @@ func newTestCRIContainerdService() *criContainerdService {
 		containerStore:     metadata.NewContainerStore(store.NewMetadataStore()),
 		containerNameIndex: registrar.NewRegistrar(),
 		containerService:   servertesting.NewFakeExecutionClient(),
-		snapshotService:    servertesting.NewFakeSnapshotClient(),
 		netPlugin:          servertesting.NewFakeCNIPlugin(),
 		agentFactory:       agentstesting.NewFakeAgentFactory(),
 	}
 }
 
+// WithFakeSnapshotClient add and return fake snapshot client.
+func WithFakeSnapshotClient(c *criContainerdService) *servertesting.FakeSnapshotClient {
+	fake := servertesting.NewFakeSnapshotClient()
+	c.snapshotService = snapshotservice.NewSnapshotterFromClient(fake)
+	return fake
+}
+
 // Test all sandbox operations.
 func TestSandboxOperations(t *testing.T) {
 	c := newTestCRIContainerdService()
 	fake := c.containerService.(*servertesting.FakeExecutionClient)
-	// TODO(random-liu): Clean this up if needed.
-	// fakeSnapshotClient := c.snapshotService.(*servertesting.FakeSnapshotClient)
 	fakeOS := c.os.(*ostesting.FakeOS)
 	fakeCNIPlugin := c.netPlugin.(*servertesting.FakeCNIPlugin)
+	WithFakeSnapshotClient(c)
 	fakeOS.OpenFifoFn = func(ctx context.Context, fn string, flag int, perm os.FileMode) (io.ReadWriteCloser, error) {
 		return nopReadWriteCloser{}, nil
 	}

pkg/server/testing/fake_snapshot_client.go

@@ -25,8 +25,12 @@ import (
 	google_protobuf1 "github.com/golang/protobuf/ptypes/empty"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 )
 
+// SnapshotNotExistError is the fake error returned when snapshot does not exist.
+var SnapshotNotExistError = grpc.Errorf(codes.NotFound, "snapshot does not exist")
+
 // FakeSnapshotClient is a simple fake snapshot client, so that cri-containerd
 // can be run for testing without requiring a real containerd setup.
 type FakeSnapshotClient struct {
@@ -109,6 +113,17 @@ func (f *FakeSnapshotClient) SetFakeMounts(name string, mounts []*mount.Mount) {
 	f.MountList[name] = mounts
 }
 
+// ListMounts lists all the fake mounts.
+func (f *FakeSnapshotClient) ListMounts() [][]*mount.Mount {
+	f.Lock()
+	defer f.Unlock()
+	var ms [][]*mount.Mount
+	for _, m := range f.MountList {
+		ms = append(ms, m)
+	}
+	return ms
+}
+
 // Prepare is a test implementation of snapshot.Prepare
 func (f *FakeSnapshotClient) Prepare(ctx context.Context, prepareOpts *snapshot.PrepareRequest, opts ...grpc.CallOption) (*snapshot.MountsResponse, error) {
 	f.Lock()
@@ -141,7 +156,7 @@ func (f *FakeSnapshotClient) Mounts(ctx context.Context, mountsOpts *snapshot.Mo
 	}
 	mounts, ok := f.MountList[mountsOpts.Key]
 	if !ok {
-		return nil, fmt.Errorf("mounts not exist")
+		return nil, SnapshotNotExistError
 	}
 	return &snapshot.MountsResponse{
 		Mounts: mounts,
@@ -154,13 +169,40 @@ func (f *FakeSnapshotClient) Commit(ctx context.Context, in *snapshot.CommitRequ
 }
 
 // View is a test implementation of snapshot.View
-func (f *FakeSnapshotClient) View(ctx context.Context, in *snapshot.PrepareRequest, opts ...grpc.CallOption) (*snapshot.MountsResponse, error) {
+func (f *FakeSnapshotClient) View(ctx context.Context, viewOpts *snapshot.PrepareRequest, opts ...grpc.CallOption) (*snapshot.MountsResponse, error) {
-	return nil, nil
+	f.Lock()
+	defer f.Unlock()
+	f.appendCalled("view", viewOpts)
+	if err := f.getError("view"); err != nil {
+		return nil, err
+	}
+	_, ok := f.MountList[viewOpts.Key]
+	if ok {
+		return nil, fmt.Errorf("mounts already exist")
+	}
+	f.MountList[viewOpts.Key] = []*mount.Mount{{
+		Type:   "bind",
+		Source: viewOpts.Key,
+		// TODO(random-liu): Fake options based on Readonly option.
+	}}
+	return &snapshot.MountsResponse{
+		Mounts: f.MountList[viewOpts.Key],
+	}, nil
 }
 
 // Remove is a test implementation of snapshot.Remove
-func (f *FakeSnapshotClient) Remove(ctx context.Context, in *snapshot.RemoveRequest, opts ...grpc.CallOption) (*google_protobuf1.Empty, error) {
+func (f *FakeSnapshotClient) Remove(ctx context.Context, removeOpts *snapshot.RemoveRequest, opts ...grpc.CallOption) (*google_protobuf1.Empty, error) {
-	return nil, nil
+	f.Lock()
+	defer f.Unlock()
+	f.appendCalled("remove", removeOpts)
+	if err := f.getError("remove"); err != nil {
+		return nil, err
+	}
+	if _, ok := f.MountList[removeOpts.Key]; !ok {
+		return nil, SnapshotNotExistError
+	}
+	delete(f.MountList, removeOpts.Key)
+	return &google_protobuf1.Empty{}, nil
 }
 
 // Stat is a test implementation of snapshot.Stat