added ability to iterate over all indexes and use index.Iterate() instead of ReadDir() to walk over the graph

Signed-off-by: Roman Strashkin <roman.strashkin@gmail.com>

daemon/daemon.go

@@ -884,10 +884,7 @@ func (daemon *Daemon) ContainerGraph() *graphdb.Database {
 
 func (daemon *Daemon) ImageGetCached(imgID string, config *runconfig.Config) (*graph.Image, error) {
 	// Retrieve all images
-	images, err := daemon.Graph().Map()
-	if err != nil {
-		return nil, err
-	}
+	images := daemon.Graph().Map()
 
 	// Store the tree in a map of map (map[parentId][childId])
 	imageMap := make(map[string]map[string]struct{})

daemon/image_delete.go

@@ -55,10 +55,7 @@ func (daemon *Daemon) imgDeleteHelper(name string, list *[]types.ImageDelete, fi
 		tag = ""
 	}
 
-	byParents, err := daemon.Graph().ByParent()
-	if err != nil {
-		return err
-	}
+	byParents := daemon.Graph().ByParent()
 
 	repos := daemon.Repositories().ByID()[img.ID]
 

daemon/info.go

@@ -17,7 +17,7 @@ import (
 )
 
 func (daemon *Daemon) SystemInfo() (*types.Info, error) {
-	images, _ := daemon.Graph().Map()
+	images := daemon.Graph().Map()
 	var imgcount int
 	if images == nil {
 		imgcount = 0

graph/graph.go

@@ -327,42 +327,33 @@ func (graph *Graph) Delete(name string) error {
 }
 
 // Map returns a list of all images in the graph, addressable by ID.
-func (graph *Graph) Map() (map[string]*Image, error) {
+func (graph *Graph) Map() map[string]*Image {
 	images := make(map[string]*Image)
-	err := graph.walkAll(func(image *Image) {
+	graph.walkAll(func(image *Image) {
 		images[image.ID] = image
 	})
-	if err != nil {
-		return nil, err
-	}
-	return images, nil
+	return images
 }
 
 // walkAll iterates over each image in the graph, and passes it to a handler.
 // The walking order is undetermined.
-func (graph *Graph) walkAll(handler func(*Image)) error {
-	files, err := ioutil.ReadDir(graph.root)
-	if err != nil {
-		return err
-	}
-	for _, st := range files {
-		if img, err := graph.Get(st.Name()); err != nil {
-			// Skip image
-			continue
+func (graph *Graph) walkAll(handler func(*Image)) {
+	graph.idIndex.Iterate(func(id string) {
+		if img, err := graph.Get(id); err != nil {
+			return
 		} else if handler != nil {
 			handler(img)
 		}
-	}
-	return nil
+	})
 }
 
 // ByParent returns a lookup table of images by their parent.
 // If an image of id ID has 3 children images, then the value for key ID
 // will be a list of 3 images.
 // If an image has no children, it will not have an entry in the table.
-func (graph *Graph) ByParent() (map[string][]*Image, error) {
+func (graph *Graph) ByParent() map[string][]*Image {
 	byParent := make(map[string][]*Image)
-	err := graph.walkAll(func(img *Image) {
+	graph.walkAll(func(img *Image) {
 		parent, err := graph.Get(img.Parent)
 		if err != nil {
 			return
@@ -373,25 +364,22 @@ func (graph *Graph) ByParent() (map[string][]*Image, error) {
 			byParent[parent.ID] = []*Image{img}
 		}
 	})
-	return byParent, err
+	return byParent
 }
 
 // Heads returns all heads in the graph, keyed by id.
 // A head is an image which is not the parent of another image in the graph.
-func (graph *Graph) Heads() (map[string]*Image, error) {
+func (graph *Graph) Heads() map[string]*Image {
 	heads := make(map[string]*Image)
-	byParent, err := graph.ByParent()
-	if err != nil {
-		return nil, err
-	}
-	err = graph.walkAll(func(image *Image) {
+	byParent := graph.ByParent()
+	graph.walkAll(func(image *Image) {
 		// If it's not in the byParent lookup table, then
 		// it's not a parent -> so it's a head!
 		if _, exists := byParent[image.ID]; !exists {
 			heads[image.ID] = image
 		}
 	})
-	return heads, err
+	return heads
 }
 
 func (graph *Graph) imageRoot(id string) string {

graph/graph_test.go

@@ -56,9 +56,8 @@ func TestInit(t *testing.T) {
 		t.Fatal(err)
 	}
 	// Map() should be empty
-	if l, err := graph.Map(); err != nil {
-		t.Fatal(err)
-	} else if len(l) != 0 {
+	l := graph.Map()
+	if len(l) != 0 {
 		t.Fatalf("len(Map()) should return %d, not %d", 0, len(l))
 	}
 }
@@ -110,10 +109,8 @@ func TestGraphCreate(t *testing.T) {
 	if img.DockerVersion != dockerversion.VERSION {
 		t.Fatalf("Wrong docker_version: should be '%s', not '%s'", dockerversion.VERSION, img.DockerVersion)
 	}
-	images, err := graph.Map()
-	if err != nil {
-		t.Fatal(err)
-	} else if l := len(images); l != 1 {
+	images := graph.Map()
+	if l := len(images); l != 1 {
 		t.Fatalf("Wrong number of images. Should be %d, not %d", 1, l)
 	}
 	if images[img.ID] == nil {
@@ -137,9 +134,8 @@ func TestRegister(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if images, err := graph.Map(); err != nil {
-		t.Fatal(err)
-	} else if l := len(images); l != 1 {
+	images := graph.Map()
+	if l := len(images); l != 1 {
 		t.Fatalf("Wrong number of images. Should be %d, not %d", 1, l)
 	}
 	if resultImg, err := graph.Get(image.ID); err != nil {
@@ -254,10 +250,7 @@ func TestByParent(t *testing.T) {
 	_ = graph.Register(childImage1, archive2)
 	_ = graph.Register(childImage2, archive3)
 
-	byParent, err := graph.ByParent()
-	if err != nil {
-		t.Fatal(err)
-	}
+	byParent := graph.ByParent()
 	numChildren := len(byParent[parentImage.ID])
 	if numChildren != 2 {
 		t.Fatalf("Expected 2 children, found %d", numChildren)
@@ -277,9 +270,8 @@ func createTestImage(graph *Graph, t *testing.T) *Image {
 }
 
 func assertNImages(graph *Graph, t *testing.T, n int) {
-	if images, err := graph.Map(); err != nil {
-		t.Fatal(err)
-	} else if actualN := len(images); actualN != n {
+	images := graph.Map()
+	if actualN := len(images); actualN != n {
 		t.Fatalf("Expected %d images, found %d", n, actualN)
 	}
 }

graph/list.go

@@ -58,12 +58,9 @@ func (s *TagStore) Images(config *ImagesConfig) ([]*types.Image, error) {
 	_, filtLabel = imageFilters["label"]
 
 	if config.All && filtTagged {
-		allImages, err = s.graph.Map()
+		allImages = s.graph.Map()
 	} else {
-		allImages, err = s.graph.Heads()
-	}
-	if err != nil {
-		return nil, err
+		allImages = s.graph.Heads()
 	}
 
 	lookup := make(map[string]*types.Image)

graph/load.go

@@ -31,10 +31,7 @@ func (s *TagStore) Load(inTar io.ReadCloser, outStream io.Writer) error {
 	if err := os.Mkdir(repoDir, os.ModeDir); err != nil {
 		return err
 	}
-	images, err := s.graph.Map()
-	if err != nil {
-		return err
-	}
+	images := s.graph.Map()
 	excludes := make([]string, len(images))
 	i := 0
 	for k := range images {

pkg/truncindex/truncindex.go

@@ -108,3 +108,13 @@ func (idx *TruncIndex) Get(s string) (string, error) {
 	}
 	return "", fmt.Errorf("no such id: %s", s)
 }
+
+// Iterates over all stored IDs, and passes each of them to the given handler
+func (idx *TruncIndex) Iterate(handler func(id string)) {
+	idx.RLock()
+	defer idx.RUnlock()
+	idx.trie.Visit(func(prefix patricia.Prefix, item patricia.Item) error {
+		handler(string(prefix))
+		return nil
+	})
+}

pkg/truncindex/truncindex_test.go

@@ -96,6 +96,29 @@ func TestTruncIndex(t *testing.T) {
 	assertIndexGet(t, index, id[:7], id, false)
 	assertIndexGet(t, index, id[:15], id, false)
 	assertIndexGet(t, index, id, id, false)
+
+	assertIndexIterate(t)
+}
+
+func assertIndexIterate(t *testing.T) {
+	ids := []string{
+		"19b36c2c326ccc11e726eee6ee78a0baf166ef96",
+		"28b36c2c326ccc11e726eee6ee78a0baf166ef96",
+		"37b36c2c326ccc11e726eee6ee78a0baf166ef96",
+		"46b36c2c326ccc11e726eee6ee78a0baf166ef96",
+	}
+
+	index := NewTruncIndex(ids)
+
+	index.Iterate(func(targetId string) {
+		for _, id := range ids {
+			if targetId == id {
+				return
+			}
+		}
+
+		t.Fatalf("An unknown ID '%s'", targetId)
+	})
 }
 
 func assertIndexGet(t *testing.T, index *TruncIndex, input, expectedResult string, expectError bool) {