From 354c241041ec67d4c500cccfda476dc73435d38e Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Fri, 31 Aug 2018 18:28:25 -0700
Subject: [PATCH] builder: implement ref checker

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
Signed-off-by: Tibor Vass <tibor@docker.com>
---
 .../adapters/snapshot/snapshot.go             |  4 +
 builder/builder-next/controller.go            | 18 +++-
 .../builder-next/imagerefchecker/checker.go   | 96 +++++++++++++++++++
 3 files changed, 114 insertions(+), 4 deletions(-)
 create mode 100644 builder/builder-next/imagerefchecker/checker.go

diff --git a/builder/builder-next/adapters/snapshot/snapshot.go b/builder/builder-next/adapters/snapshot/snapshot.go
index e1a12a4563..d8efb759c2 100644
--- a/builder/builder-next/adapters/snapshot/snapshot.go
+++ b/builder/builder-next/adapters/snapshot/snapshot.go
@@ -110,6 +110,10 @@ func (s *snapshotter) chainID(key string) (layer.ChainID, bool) {
 	return "", false
 }
 
+func (s *snapshotter) GetLayer(key string) (layer.Layer, error) {
+	return s.getLayer(key, true)
+}
+
 func (s *snapshotter) getLayer(key string, withCommitted bool) (layer.Layer, error) {
 	s.mu.Lock()
 	l, ok := s.refs[key]
diff --git a/builder/builder-next/controller.go b/builder/builder-next/controller.go
index 1d0ed2d496..df0f3d48e5 100644
--- a/builder/builder-next/controller.go
+++ b/builder/builder-next/controller.go
@@ -9,6 +9,7 @@ import (
 	"github.com/docker/docker/builder/builder-next/adapters/containerimage"
 	"github.com/docker/docker/builder/builder-next/adapters/snapshot"
 	containerimageexp "github.com/docker/docker/builder/builder-next/exporter"
+	"github.com/docker/docker/builder/builder-next/imagerefchecker"
 	mobyworker "github.com/docker/docker/builder/builder-next/worker"
 	"github.com/docker/docker/daemon/graphdriver"
 	"github.com/moby/buildkit/cache"
@@ -69,11 +70,20 @@ func newController(rt http.RoundTripper, opt Opt) (*control.Controller, error) {
 		MetadataStore: md,
 	})
 
+	layerGetter, ok := sbase.(imagerefchecker.LayerGetter)
+	if !ok {
+		return nil, errors.Errorf("snapshotter does not implement layergetter")
+	}
+
+	refChecker := imagerefchecker.New(imagerefchecker.Opt{
+		ImageStore:  dist.ImageStore,
+		LayerGetter: layerGetter,
+	})
+
 	cm, err := cache.NewManager(cache.ManagerOpt{
-		Snapshotter:   snapshotter,
-		MetadataStore: md,
-		// TODO: implement PruneRefChecker to correctly mark cache objects as "Shared"
-		PruneRefChecker: nil,
+		Snapshotter:     snapshotter,
+		MetadataStore:   md,
+		PruneRefChecker: refChecker,
 	})
 	if err != nil {
 		return nil, err
diff --git a/builder/builder-next/imagerefchecker/checker.go b/builder/builder-next/imagerefchecker/checker.go
new file mode 100644
index 0000000000..052391d589
--- /dev/null
+++ b/builder/builder-next/imagerefchecker/checker.go
@@ -0,0 +1,96 @@
+package imagerefchecker
+
+import (
+	"sync"
+
+	"github.com/docker/docker/image"
+	"github.com/docker/docker/layer"
+	"github.com/moby/buildkit/cache"
+)
+
+// LayerGetter abstracts away the snapshotter
+type LayerGetter interface {
+	GetLayer(string) (layer.Layer, error)
+}
+
+// Opt represents the options needed to create a refchecker
+type Opt struct {
+	LayerGetter LayerGetter
+	ImageStore  image.Store
+}
+
+// New creates new image reference checker that can be used to see if a reference
+// is being used by any of the images in the image store
+func New(opt Opt) cache.ExternalRefCheckerFunc {
+	return func() (cache.ExternalRefChecker, error) {
+		return &checker{opt: opt, layers: lchain{}, cache: map[string]bool{}}, nil
+	}
+}
+
+type lchain map[layer.DiffID]lchain
+
+func (c lchain) add(ids []layer.DiffID) {
+	if len(ids) == 0 {
+		return
+	}
+	id := ids[0]
+	ch, ok := c[id]
+	if !ok {
+		ch = lchain{}
+		c[id] = ch
+	}
+	ch.add(ids[1:])
+}
+
+func (c lchain) has(ids []layer.DiffID) bool {
+	if len(ids) == 0 {
+		return true
+	}
+	ch, ok := c[ids[0]]
+	return ok && ch.has(ids[1:])
+}
+
+type checker struct {
+	opt    Opt
+	once   sync.Once
+	layers lchain
+	cache  map[string]bool
+}
+
+func (c *checker) Exists(key string) bool {
+	if c.opt.ImageStore == nil {
+		return false
+	}
+
+	c.once.Do(c.init)
+
+	if b, ok := c.cache[key]; ok {
+		return b
+	}
+
+	l, err := c.opt.LayerGetter.GetLayer(key)
+	if err != nil || l == nil {
+		c.cache[key] = false
+		return false
+	}
+
+	ok := c.layers.has(diffIDs(l))
+	c.cache[key] = ok
+	return ok
+}
+
+func (c *checker) init() {
+	imgs := c.opt.ImageStore.Map()
+
+	for _, img := range imgs {
+		c.layers.add(img.RootFS.DiffIDs)
+	}
+}
+
+func diffIDs(l layer.Layer) []layer.DiffID {
+	p := l.Parent()
+	if p == nil {
+		return []layer.DiffID{l.DiffID()}
+	}
+	return append(diffIDs(p), l.DiffID())
+}