builder-next: track layers and graphdrivers with leases

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

builder/builder-next/adapters/snapshot/leasemanager.go

@@ -0,0 +1,123 @@
+package snapshot
+
+import (
+	"context"
+	"sync"
+
+	"github.com/containerd/containerd/leases"
+	"github.com/sirupsen/logrus"
+)
+
+type sLM struct {
+	manager leases.Manager
+	s       *snapshotter
+
+	mu         sync.Mutex
+	byLease    map[string]map[string]struct{}
+	bySnapshot map[string]map[string]struct{}
+}
+
+func newLeaseManager(s *snapshotter, lm leases.Manager) *sLM {
+	return &sLM{
+		s:       s,
+		manager: lm,
+
+		byLease:    map[string]map[string]struct{}{},
+		bySnapshot: map[string]map[string]struct{}{},
+	}
+}
+
+func (l *sLM) Create(ctx context.Context, opts ...leases.Opt) (leases.Lease, error) {
+	return l.manager.Create(ctx, opts...)
+}
+
+func (l *sLM) Delete(ctx context.Context, lease leases.Lease, opts ...leases.DeleteOpt) error {
+	if err := l.manager.Delete(ctx, lease, opts...); err != nil {
+		return err
+	}
+	l.mu.Lock()
+	if snaps, ok := l.byLease[lease.ID]; ok {
+		for sID := range snaps {
+			l.delRef(lease.ID, sID)
+		}
+	}
+	l.mu.Unlock()
+	return nil
+}
+
+func (l *sLM) List(ctx context.Context, filters ...string) ([]leases.Lease, error) {
+	return l.manager.List(ctx, filters...)
+}
+
+func (l *sLM) AddResource(ctx context.Context, lease leases.Lease, resource leases.Resource) error {
+	if err := l.manager.AddResource(ctx, lease, resource); err != nil {
+		return err
+	}
+	if resource.Type == "snapshots/default" {
+		l.mu.Lock()
+		l.addRef(lease.ID, resource.ID)
+		l.mu.Unlock()
+	}
+	return nil
+}
+
+func (l *sLM) DeleteResource(ctx context.Context, lease leases.Lease, resource leases.Resource) error {
+	if err := l.manager.DeleteResource(ctx, lease, resource); err != nil {
+		return err
+	}
+	if resource.Type == "snapshots/default" {
+		l.mu.Lock()
+		l.delRef(lease.ID, resource.ID)
+		l.mu.Unlock()
+	}
+	return nil
+}
+
+func (l *sLM) ListResources(ctx context.Context, lease leases.Lease) ([]leases.Resource, error) {
+	return l.manager.ListResources(ctx, lease)
+}
+
+func (l *sLM) addRef(lID, sID string) {
+	load := false
+	snapshots, ok := l.byLease[lID]
+	if !ok {
+		snapshots = map[string]struct{}{}
+		l.byLease[lID] = snapshots
+	}
+	if _, ok := snapshots[sID]; !ok {
+		snapshots[sID] = struct{}{}
+	}
+	leases, ok := l.bySnapshot[sID]
+	if !ok {
+		leases = map[string]struct{}{}
+		l.byLease[sID] = leases
+		load = true
+	}
+	if _, ok := leases[lID]; !ok {
+		leases[lID] = struct{}{}
+	}
+
+	if load {
+		l.s.getLayer(sID, true)
+	}
+}
+
+func (l *sLM) delRef(lID, sID string) {
+	snapshots, ok := l.byLease[lID]
+	if !ok {
+		delete(snapshots, sID)
+		if len(snapshots) == 0 {
+			delete(l.byLease, lID)
+		}
+	}
+	leases, ok := l.bySnapshot[sID]
+	if !ok {
+		delete(leases, lID)
+		if len(leases) == 0 {
+			delete(l.bySnapshot, sID)
+			if err := l.s.remove(context.TODO(), sID); err != nil {
+				logrus.Warnf("failed to remove snapshot %v", sID)
+			}
+		}
+	}
+}

builder/builder-next/adapters/snapshot/snapshot.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/mount"
 	"github.com/containerd/containerd/snapshots"
 	"github.com/docker/docker/daemon/graphdriver"
@@ -21,6 +22,7 @@ import (
 
 var keyParent = []byte("parent")
 var keyCommitted = []byte("committed")
+var keyIsCommitted = []byte("iscommitted")
 var keyChainID = []byte("chainid")
 var keySize = []byte("size")
 
@@ -52,16 +54,16 @@ type snapshotter struct {
 }
 
 // NewSnapshotter creates a new snapshotter
-func NewSnapshotter(opt Opt) (snapshot.Snapshotter, error) {
+func NewSnapshotter(opt Opt, prevLM leases.Manager) (snapshot.Snapshotter, leases.Manager, error) {
 	dbPath := filepath.Join(opt.Root, "snapshots.db")
 	db, err := bolt.Open(dbPath, 0600, nil)
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to open database file %s", dbPath)
+		return nil, nil, errors.Wrapf(err, "failed to open database file %s", dbPath)
 	}
 
 	reg, ok := opt.LayerStore.(graphIDRegistrar)
 	if !ok {
-		return nil, errors.Errorf("layerstore doesn't support graphID registration")
+		return nil, nil, errors.Errorf("layerstore doesn't support graphID registration")
 	}
 
 	s := &snapshotter{
@@ -70,7 +72,28 @@ func NewSnapshotter(opt Opt) (snapshot.Snapshotter, error) {
 		refs: map[string]layer.Layer{},
 		reg:  reg,
 	}
-	return s, nil
+
+	lm := newLeaseManager(s, prevLM)
+
+	// TODO: temp-leases
+
+	ll, err := lm.List(context.TODO())
+	if err != nil {
+		return nil, nil, err
+	}
+	for _, l := range ll {
+		rr, err := lm.ListResources(context.TODO(), l)
+		if err != nil {
+			return nil, nil, err
+		}
+		for _, r := range rr {
+			if r.Type == "snapshots/default" {
+				lm.addRef(l.ID, r.ID)
+			}
+		}
+	}
+
+	return s, lm, nil
 }
 
 func (s *snapshotter) Name() string {
@@ -295,6 +318,10 @@ func (s *snapshotter) Mounts(ctx context.Context, key string) (snapshot.Mountabl
 }
 
 func (s *snapshotter) Remove(ctx context.Context, key string) error {
+	return errors.Errorf("calling snapshot.remove is forbidden")
+}
+
+func (s *snapshotter) remove(ctx context.Context, key string) error {
 	l, err := s.getLayer(key, true)
 	if err != nil {
 		return err
@@ -303,8 +330,17 @@ func (s *snapshotter) Remove(ctx context.Context, key string) error {
 	id, _ := s.getGraphDriverID(key)
 
 	var found bool
+	var alreadyCommitted bool
 	if err := s.db.Update(func(tx *bolt.Tx) error {
-		found = tx.Bucket([]byte(key)) != nil
+		b := tx.Bucket([]byte(key))
+		found = b != nil
+
+		if b != nil {
+			if b.Get(keyIsCommitted) != nil {
+				alreadyCommitted = true
+				return nil
+			}
+		}
 		if found {
 			tx.DeleteBucket([]byte(key))
 			if id != key {
@@ -316,6 +352,10 @@ func (s *snapshotter) Remove(ctx context.Context, key string) error {
 		return err
 	}
 
+	if alreadyCommitted {
+		return nil
+	}
+
 	if l != nil {
 		s.mu.Lock()
 		delete(s.refs, key)
@@ -337,7 +377,17 @@ func (s *snapshotter) Commit(ctx context.Context, name, key string, opts ...snap
 		if err != nil {
 			return err
 		}
-		return b.Put(keyCommitted, []byte(key))
+		if err := b.Put(keyCommitted, []byte(key)); err != nil {
+			return err
+		}
+		b, err = tx.CreateBucketIfNotExists([]byte(key))
+		if err != nil {
+			return err
+		}
+		if err := b.Put(keyIsCommitted, []byte{}); err != nil {
+			return err
+		}
+		return nil
 	})
 }
 

builder/builder-next/controller.go

@@ -59,16 +59,6 @@ func newController(rt http.RoundTripper, opt Opt) (*control.Controller, error) {
 		return nil, errors.Errorf("could not access graphdriver")
 	}
 
-	snapshotter, err := snapshot.NewSnapshotter(snapshot.Opt{
-		GraphDriver:     driver,
-		LayerStore:      dist.LayerStore,
-		Root:            root,
-		IdentityMapping: opt.IdentityMapping,
-	})
-	if err != nil {
-		return nil, err
-	}
-
 	store, err := local.NewStore(filepath.Join(root, "content"))
 	if err != nil {
 		return nil, err
@@ -83,6 +73,18 @@ func newController(rt http.RoundTripper, opt Opt) (*control.Controller, error) {
 
 	store = containerdsnapshot.NewContentStore(mdb.ContentStore(), "buildkit")
 
+	lm := leaseutil.WithNamespace(ctdmetadata.NewLeaseManager(mdb), "buildkit")
+
+	snapshotter, lm, err := snapshot.NewSnapshotter(snapshot.Opt{
+		GraphDriver:     driver,
+		LayerStore:      dist.LayerStore,
+		Root:            root,
+		IdentityMapping: opt.IdentityMapping,
+	}, lm)
+	if err != nil {
+		return nil, err
+	}
+
 	md, err := metadata.NewStore(filepath.Join(root, "metadata.db"))
 	if err != nil {
 		return nil, err
@@ -102,7 +104,7 @@ func newController(rt http.RoundTripper, opt Opt) (*control.Controller, error) {
 		Snapshotter:     snapshotter,
 		MetadataStore:   md,
 		PruneRefChecker: refChecker,
-		LeaseManager:    leaseutil.WithNamespace(ctdmetadata.NewLeaseManager(mdb), "buildkit"),
+		LeaseManager:    lm,
 		ContentStore:    store,
 	})
 	if err != nil {