Add Truncindex for container, sandbox and image

pkg/server/container_attach.go

@@ -52,7 +52,7 @@ func (c *criContainerdService) attachContainer(ctx context.Context, id string, s
 	// Get container from our container store.
 	cntr, err := c.containerStore.Get(id)
 	if err != nil {
-		return fmt.Errorf("failed to find container in store: %v", err)
+		return fmt.Errorf("failed to find container %q in store: %v", id, err)
 	}
 	id = cntr.ID
 

pkg/server/container_exec.go

@@ -27,7 +27,7 @@ import (
 func (c *criContainerdService) Exec(ctx context.Context, r *runtime.ExecRequest) (*runtime.ExecResponse, error) {
 	cntr, err := c.containerStore.Get(r.GetContainerId())
 	if err != nil {
-		return nil, fmt.Errorf("failed to find container in store: %v", err)
+		return nil, fmt.Errorf("failed to find container %q in store: %v", r.GetContainerId(), err)
 	}
 	state := cntr.Status.Get().State()
 	if state != runtime.ContainerState_CONTAINER_RUNNING {

pkg/server/container_execsync.go

@@ -78,7 +78,7 @@ func (c *criContainerdService) execInContainer(ctx context.Context, id string, o
 	// Get container from our container store.
 	cntr, err := c.containerStore.Get(id)
 	if err != nil {
-		return nil, fmt.Errorf("failed to find container in store: %v", err)
+		return nil, fmt.Errorf("failed to find container %q in store: %v", id, err)
 	}
 	id = cntr.ID
 

pkg/server/image_pull.go

@@ -142,7 +142,9 @@ func (c *criContainerdService) PullImage(ctx context.Context, r *runtime.PullIma
 		img.RepoTags = []string{repoTag}
 	}
 
-	c.imageStore.Add(img)
+	if err := c.imageStore.Add(img); err != nil {
+		return nil, fmt.Errorf("failed to add image %q into store: %v", img.ID, err)
+	}
 
 	// NOTE(random-liu): the actual state in containerd is the source of truth, even we maintain
 	// in-memory image store, it's only for in-memory indexing. The image could be removed

pkg/server/restart.go

@@ -101,8 +101,10 @@ func (c *criContainerdService) recover(ctx context.Context) error {
 		return fmt.Errorf("failed to load images: %v", err)
 	}
 	for _, image := range images {
-		c.imageStore.Add(image)
 		glog.V(4).Infof("Loaded image %+v", image)
+		if err := c.imageStore.Add(image); err != nil {
+			return fmt.Errorf("failed to add image %q to store: %v", image.ID, err)
+		}
 	}
 
 	// It's possible that containerd containers are deleted unexpectedly. In that case,

pkg/server/sandbox_portforward.go

@@ -35,7 +35,7 @@ func (c *criContainerdService) PortForward(ctx context.Context, r *runtime.PortF
 	// TODO(random-liu): Run a socat container inside the sandbox to do portforward.
 	sandbox, err := c.sandboxStore.Get(r.GetPodSandboxId())
 	if err != nil {
-		return nil, fmt.Errorf("failed to find sandbox: %v", err)
+		return nil, fmt.Errorf("failed to find sandbox %q: %v", r.GetPodSandboxId(), err)
 	}
 
 	t, err := sandbox.Container.Task(ctx, nil)
@@ -59,7 +59,7 @@ func (c *criContainerdService) PortForward(ctx context.Context, r *runtime.PortF
 func (c *criContainerdService) portForward(id string, port int32, stream io.ReadWriteCloser) error {
 	s, err := c.sandboxStore.Get(id)
 	if err != nil {
-		return fmt.Errorf("failed to find sandbox in store: %v", err)
+		return fmt.Errorf("failed to find sandbox %q in store: %v", id, err)
 	}
 	t, err := s.Container.Task(context.Background(), nil)
 	if err != nil {

pkg/store/container/container.go

@@ -20,6 +20,7 @@ import (
 	"sync"
 
 	"github.com/containerd/containerd"
+	"github.com/docker/docker/pkg/truncindex"
 
 	cio "github.com/kubernetes-incubator/cri-containerd/pkg/server/io"
 	"github.com/kubernetes-incubator/cri-containerd/pkg/store"
@@ -92,7 +93,7 @@ func (c *Container) Delete() error {
 type Store struct {
 	lock       sync.RWMutex
 	containers map[string]Container
-	// TODO(random-liu): Add trunc index.
+	idIndex    *truncindex.TruncIndex
 }
 
 // LoadStore loads containers from runtime.
@@ -101,7 +102,10 @@ func LoadStore() *Store { return nil }
 
 // NewStore creates a container store.
 func NewStore() *Store {
-	return &Store{containers: make(map[string]Container)}
+	return &Store{
+		containers: make(map[string]Container),
+		idIndex:    truncindex.NewTruncIndex([]string{}),
+	}
 }
 
 // Add a container into the store. Returns store.ErrAlreadyExist if the
@@ -112,6 +116,9 @@ func (s *Store) Add(c Container) error {
 	if _, ok := s.containers[c.ID]; ok {
 		return store.ErrAlreadyExist
 	}
+	if err := s.idIndex.Add(c.ID); err != nil {
+		return err
+	}
 	s.containers[c.ID] = c
 	return nil
 }
@@ -121,6 +128,13 @@ func (s *Store) Add(c Container) error {
 func (s *Store) Get(id string) (Container, error) {
 	s.lock.RLock()
 	defer s.lock.RUnlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		if err == truncindex.ErrNotExist {
+			err = store.ErrNotExist
+		}
+		return Container{}, err
+	}
 	if c, ok := s.containers[id]; ok {
 		return c, nil
 	}
@@ -142,5 +156,12 @@ func (s *Store) List() []Container {
 func (s *Store) Delete(id string) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		// Note: The idIndex.Delete and delete doesn't handle truncated index.
+		// So we need to return if there are error.
+		return
+	}
+	s.idIndex.Delete(id) // nolint: errcheck
 	delete(s.containers, id)
 }

pkg/store/container/container_test.go

@@ -28,7 +28,6 @@ import (
 )
 
 func TestContainerStore(t *testing.T) {
-	ids := []string{"1", "2", "3"}
 	metadatas := map[string]Metadata{
 		"1": {
 			ID:        "1",
@@ -43,32 +42,44 @@ func TestContainerStore(t *testing.T) {
 			ImageRef: "TestImage-1",
 			LogPath:  "/test/log/path/1",
 		},
-		"2": {
-			ID:        "2",
-			Name:      "Container-2",
-			SandboxID: "Sandbox-2",
+		"2abcd": {
+			ID:        "2abcd",
+			Name:      "Container-2abcd",
+			SandboxID: "Sandbox-2abcd",
 			Config: &runtime.ContainerConfig{
 				Metadata: &runtime.ContainerMetadata{
-					Name:    "TestPod-2",
+					Name:    "TestPod-2abcd",
 					Attempt: 2,
 				},
 			},
 			ImageRef: "TestImage-2",
 			LogPath:  "/test/log/path/2",
 		},
-		"3": {
-			ID:        "3",
-			Name:      "Container-3",
-			SandboxID: "Sandbox-3",
+		"4a333": {
+			ID:        "4a333",
+			Name:      "Container-4a333",
+			SandboxID: "Sandbox-4a333",
 			Config: &runtime.ContainerConfig{
 				Metadata: &runtime.ContainerMetadata{
-					Name:    "TestPod-3",
+					Name:    "TestPod-4a333",
 					Attempt: 3,
 				},
 			},
 			ImageRef: "TestImage-3",
 			LogPath:  "/test/log/path/3",
 		},
+		"4abcd": {
+			ID:        "4abcd",
+			Name:      "Container-4abcd",
+			SandboxID: "Sandbox-4abcd",
+			Config: &runtime.ContainerConfig{
+				Metadata: &runtime.ContainerMetadata{
+					Name:    "TestPod-4abcd",
+					Attempt: 1,
+				},
+			},
+			ImageRef: "TestImage-4abcd",
+		},
 	}
 	statuses := map[string]Status{
 		"1": {
@@ -80,29 +91,39 @@ func TestContainerStore(t *testing.T) {
 			Reason:     "TestReason-1",
 			Message:    "TestMessage-1",
 		},
-		"2": {
+		"2abcd": {
 			Pid:        2,
 			CreatedAt:  time.Now().UnixNano(),
 			StartedAt:  time.Now().UnixNano(),
 			FinishedAt: time.Now().UnixNano(),
 			ExitCode:   2,
-			Reason:     "TestReason-2",
-			Message:    "TestMessage-2",
+			Reason:     "TestReason-2abcd",
+			Message:    "TestMessage-2abcd",
 		},
-		"3": {
+		"4a333": {
 			Pid:        3,
 			CreatedAt:  time.Now().UnixNano(),
 			StartedAt:  time.Now().UnixNano(),
 			FinishedAt: time.Now().UnixNano(),
 			ExitCode:   3,
-			Reason:     "TestReason-3",
-			Message:    "TestMessage-3",
+			Reason:     "TestReason-4a333",
+			Message:    "TestMessage-4a333",
+			Removing:   true,
+		},
+		"4abcd": {
+			Pid:        4,
+			CreatedAt:  time.Now().UnixNano(),
+			StartedAt:  time.Now().UnixNano(),
+			FinishedAt: time.Now().UnixNano(),
+			ExitCode:   4,
+			Reason:     "TestReason-4abcd",
+			Message:    "TestMessage-4abcd",
 			Removing:   true,
 		},
 	}
 	assert := assertlib.New(t)
 	containers := map[string]Container{}
-	for _, id := range ids {
+	for id := range metadatas {
 		container, err := NewContainer(
 			metadatas[id],
 			WithFakeStatus(statuses[id]),
@@ -119,29 +140,35 @@ func TestContainerStore(t *testing.T) {
 	}
 
 	t.Logf("should be able to get container")
+	genTruncIndex := func(normalName string) string { return normalName[:(len(normalName)+1)/2] }
 	for id, c := range containers {
-		got, err := s.Get(id)
+		got, err := s.Get(genTruncIndex(id))
 		assert.NoError(err)
 		assert.Equal(c, got)
 	}
 
 	t.Logf("should be able to list containers")
 	cs := s.List()
-	assert.Len(cs, 3)
+	assert.Len(cs, len(containers))
+
+	cntrNum := len(containers)
+	for testID, v := range containers {
+		truncID := genTruncIndex(testID)
 
-	testID := "2"
-	t.Logf("add should return already exists error for duplicated container")
-	assert.Equal(store.ErrAlreadyExist, s.Add(containers[testID]))
+		t.Logf("add should return already exists error for duplicated container")
+		assert.Equal(store.ErrAlreadyExist, s.Add(v))
 
-	t.Logf("should be able to delete container")
-	s.Delete(testID)
-	cs = s.List()
-	assert.Len(cs, 2)
+		t.Logf("should be able to delete container")
+		s.Delete(truncID)
+		cntrNum--
+		cs = s.List()
+		assert.Len(cs, cntrNum)
 
-	t.Logf("get should return not exist error after deletion")
-	c, err := s.Get(testID)
-	assert.Equal(Container{}, c)
-	assert.Equal(store.ErrNotExist, err)
+		t.Logf("get should return not exist error after deletion")
+		c, err := s.Get(truncID)
+		assert.Equal(Container{}, c)
+		assert.Equal(store.ErrNotExist, err)
+	}
 }
 
 func TestWithContainerIO(t *testing.T) {

pkg/store/image/image.go

@@ -20,6 +20,7 @@ import (
 	"sync"
 
 	"github.com/containerd/containerd"
+	"github.com/docker/docker/pkg/truncindex"
 	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
 
 	"github.com/kubernetes-incubator/cri-containerd/pkg/store"
@@ -46,9 +47,9 @@ type Image struct {
 
 // Store stores all images.
 type Store struct {
-	lock   sync.RWMutex
-	images map[string]Image
-	// TODO(random-liu): Add trunc index.
+	lock    sync.RWMutex
+	images  map[string]Image
+	idIndex *truncindex.TruncIndex
 }
 
 // LoadStore loads images from runtime.
@@ -57,30 +58,50 @@ func LoadStore() *Store { return nil }
 
 // NewStore creates an image store.
 func NewStore() *Store {
-	return &Store{images: make(map[string]Image)}
+	return &Store{
+		images:  make(map[string]Image),
+		idIndex: truncindex.NewTruncIndex([]string{}),
+	}
 }
 
 // Add an image into the store.
-func (s *Store) Add(img Image) {
+func (s *Store) Add(img Image) error {
 	s.lock.Lock()
 	defer s.lock.Unlock()
+	if _, err := s.idIndex.Get(img.ID); err != nil {
+		if err != truncindex.ErrNotExist {
+			return err
+		}
+		if err := s.idIndex.Add(img.ID); err != nil {
+			return err
+		}
+	}
+
 	i, ok := s.images[img.ID]
 	if !ok {
 		// If the image doesn't exist, add it.
 		s.images[img.ID] = img
-		return
+		return nil
 	}
 	// Or else, merge the repo tags/digests.
 	i.RepoTags = mergeStringSlices(i.RepoTags, img.RepoTags)
 	i.RepoDigests = mergeStringSlices(i.RepoDigests, img.RepoDigests)
 	s.images[img.ID] = i
+	return nil
 }
 
 // Get returns the image with specified id. Returns store.ErrNotExist if the
 // image doesn't exist.
 func (s *Store) Get(id string) (Image, error) {
 	s.lock.RLock()
 	defer s.lock.RUnlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		if err == truncindex.ErrNotExist {
+			err = store.ErrNotExist
+		}
+		return Image{}, err
+	}
 	if i, ok := s.images[id]; ok {
 		return i, nil
 	}
@@ -102,6 +123,13 @@ func (s *Store) List() []Image {
 func (s *Store) Delete(id string) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
+	id, err := s.idIndex.Get(id)
+	if err != nil {
+		// Note: The idIndex.Delete and delete doesn't handle truncated index.
+		// So we need to return if there are error.
+		return
+	}
+	s.idIndex.Delete(id) // nolint: errcheck
 	delete(s.images, id)
 }