Enable stargz snapshotter to use session for snapshot auth

Signed-off-by: ktock <ktokunaga.mail@gmail.com>

cache/manager.go

@@ -513,7 +513,16 @@ func (cm *cacheManager) New(ctx context.Context, s ImmutableRef, sess session.Gr
 		return nil, errors.Wrapf(err, "failed to add snapshot %s to lease", id)
 	}
 
-	if err := cm.Snapshotter.Prepare(ctx, id, parentSnapshotID); err != nil {
+	if parent != nil {
+		if rerr := parent.withRemoteSnapshotLabels(ctx, sess, func() {
+			err = cm.Snapshotter.Prepare(ctx, id, parentSnapshotID)
+		}); rerr != nil {
+			return nil, rerr
+		}
+	} else {
+		err = cm.Snapshotter.Prepare(ctx, id, parentSnapshotID)
+	}
+	if err != nil {
 		return nil, errors.Wrapf(err, "failed to prepare %s", id)
 	}
 

cache/refs.go

@@ -147,11 +147,23 @@ func (cr *cacheRecord) isLazy(ctx context.Context) (bool, error) {
 	if !getBlobOnly(cr.md) {
 		return false, nil
 	}
+
+	// If the content hasn't been downloaded, this layer is lazy.
 	_, err := cr.cm.ContentStore.Info(ctx, digest.Digest(getBlob(cr.md)))
 	if errors.Is(err, errdefs.ErrNotFound) {
 		return true, nil
+	} else if err != nil {
+		return false, err
 	}
-	return false, err
+
+	// If the snapshot is a remote snapshot, this layer is lazy.
+	if info, err := cr.cm.Snapshotter.Stat(ctx, getSnapshotID(cr.md)); err == nil {
+		if _, ok := info.Labels["containerd.io/snapshot/remote"]; ok {
+			return true, nil
+		}
+	}
+
+	return false, nil
 }
 
 func (cr *cacheRecord) IdentityMapping() *idtools.IdentityMapping {
@@ -369,7 +381,7 @@ func (sr *immutableRef) parentRefChain() []*immutableRef {
 	return refs
 }
 
-func (sr *immutableRef) Mount(ctx context.Context, readonly bool, s session.Group) (snapshot.Mountable, error) {
+func (sr *immutableRef) Mount(ctx context.Context, readonly bool, s session.Group) (m snapshot.Mountable, rerr error) {
 	if getBlobOnly(sr.md) {
 		if err := sr.Extract(ctx, s); err != nil {
 			return nil, err
@@ -378,7 +390,13 @@ func (sr *immutableRef) Mount(ctx context.Context, readonly bool, s session.Grou
 
 	sr.mu.Lock()
 	defer sr.mu.Unlock()
-	return sr.mount(ctx, readonly)
+
+	if err := sr.withRemoteSnapshotLabels(ctx, s, func() {
+		m, rerr = sr.mount(ctx, readonly)
+	}); err != nil {
+		return nil, err
+	}
+	return
 }
 
 func (sr *immutableRef) Extract(ctx context.Context, s session.Group) (rerr error) {
@@ -392,62 +410,152 @@ func (sr *immutableRef) Extract(ctx context.Context, s session.Group) (rerr erro
 		ctx = winlayers.UseWindowsLayerMode(ctx)
 	}
 
-	if _, err := sr.prepareRemoteSnapshots(ctx, sr.descHandlers); err != nil {
+	if err := sr.withRemoteSnapshotLabels(ctx, s, func() {
+		if rerr = sr.prepareRemoteSnapshots(ctx, s); rerr != nil {
+			return
+		}
+		rerr = sr.extract(ctx, sr.descHandlers, s)
+	}); err != nil {
 		return err
 	}
-
-	return sr.extract(ctx, sr.descHandlers, s)
+	return
 }
 
-func (sr *immutableRef) prepareRemoteSnapshots(ctx context.Context, dhs DescHandlers) (bool, error) {
-	ok, err := sr.sizeG.Do(ctx, sr.ID()+"-prepare-remote-snapshot", func(ctx context.Context) (_ interface{}, rerr error) {
-		snapshotID := getSnapshotID(sr.md)
-		if _, err := sr.cm.Snapshotter.Stat(ctx, snapshotID); err == nil {
-			return true, nil
+func (sr *immutableRef) withRemoteSnapshotLabels(ctx context.Context, s session.Group, f func()) error {
+	dhs := sr.descHandlers
+	for _, r := range sr.parentRefChain() {
+		r := r
+		info, err := r.cm.Snapshotter.Stat(ctx, getSnapshotID(r.md))
+		if err != nil && !errdefs.IsNotFound(err) {
+			return err
+		} else if errdefs.IsNotFound(err) {
+			continue // This snpashot doesn't exist; skip
+		} else if _, ok := info.Labels["containerd.io/snapshot/remote"]; !ok {
+			continue // This isn't a remote snapshot; skip
 		}
-		desc, err := sr.ociDesc()
+		desc, err := r.ociDesc()
 		if err != nil {
-			return false, err
+			return err
 		}
 		dh := dhs[desc.Digest]
 		if dh == nil {
-			return false, nil
+			continue // no info psased; skip
 		}
 
-		parentID := ""
-		if sr.parent != nil {
-			if ok, err := sr.parent.prepareRemoteSnapshots(ctx, dhs); !ok {
-				return false, err
+		// Append temporary labels (based on dh.SnapshotLabels) as hints for remote snapshots.
+		// For avoiding collosion among calls, keys of these tmp labels contain an unique ID.
+		flds, labels := makeTmpLabels(snapshots.FilterInheritedLabels(dh.SnapshotLabels), s)
+		info.Labels = labels
+		if _, err := r.cm.Snapshotter.Update(ctx, info, flds...); err != nil {
+			return errors.Wrapf(err, "failed to add tmp remote labels for remote snapshot")
+		}
+		defer func() {
+			for k := range info.Labels {
+				info.Labels[k] = "" // Remove labels appended in this call
+			}
+			if _, err := r.cm.Snapshotter.Update(ctx, info, flds...); err != nil {
+				logrus.Warn(errors.Wrapf(err, "failed to remove tmp remote labels"))
+			}
+		}()
+
+		continue
+	}
+
+	f()
+
+	return nil
+}
+
+func (sr *immutableRef) prepareRemoteSnapshots(ctx context.Context, s session.Group) error {
+	_, err := sr.sizeG.Do(ctx, sr.ID()+"-prepare-remote-snapshot", func(ctx context.Context) (_ interface{}, rerr error) {
+		dhs := sr.descHandlers
+		for _, r := range sr.parentRefChain() {
+			r := r
+			snapshotID := getSnapshotID(r.md)
+			if _, err := r.cm.Snapshotter.Stat(ctx, snapshotID); err == nil {
+				continue
+			}
+
+			desc, err := r.ociDesc()
+			if err != nil {
+				return nil, err
 			}
-			parentID = getSnapshotID(sr.parent.md)
-		}
-
-		// Hint labels to the snapshotter
-		labels := dh.SnapshotLabels
-		if labels == nil {
-			labels = make(map[string]string)
-		}
-		labels["containerd.io/snapshot.ref"] = snapshotID
-		opt := snapshots.WithLabels(labels)
-
-		// Try to preapre the remote snapshot
-		key := fmt.Sprintf("tmp-%s %s", identity.NewID(), sr.Info().ChainID)
-		if err = sr.cm.Snapshotter.Prepare(ctx, key, parentID, opt); err != nil {
-			if errdefs.IsAlreadyExists(err) {
-				// Check if the targeting snapshot ID has been prepared as a remote
-				// snapshot in the snapshotter.
-				if _, err := sr.cm.Snapshotter.Stat(ctx, snapshotID); err == nil {
-					// We can use this remote snapshot without unlazying.
-					// Try the next layer as well.
-					return true, nil
+			dh := dhs[desc.Digest]
+			if dh == nil {
+				// We cannot prepare remote snapshots without descHandler.
+				return nil, nil
+			}
+
+			// tmpLabels contains dh.SnapshotLabels + session IDs. All keys contain
+			// an unique ID for avoiding the collision among snapshotter API calls to
+			// this snapshot. tmpLabels will be removed at the end of this function.
+			defaultLabels := snapshots.FilterInheritedLabels(dh.SnapshotLabels)
+			if defaultLabels == nil {
+				defaultLabels = make(map[string]string)
+			}
+			tmpFields, tmpLabels := makeTmpLabels(defaultLabels, s)
+			defaultLabels["containerd.io/snapshot.ref"] = snapshotID
+
+			// Prepare remote snapshots
+			var (
+				key  = fmt.Sprintf("tmp-%s %s", identity.NewID(), r.Info().ChainID)
+				opts = []snapshots.Opt{
+					snapshots.WithLabels(defaultLabels),
+					snapshots.WithLabels(tmpLabels),
 				}
+			)
+			parentID := ""
+			if r.parent != nil {
+				parentID = getSnapshotID(r.parent.md)
 			}
+			if err = r.cm.Snapshotter.Prepare(ctx, key, parentID, opts...); err != nil {
+				if errdefs.IsAlreadyExists(err) {
+					// Check if the targeting snapshot ID has been prepared as
+					// a remote snapshot in the snapshotter.
+					info, err := r.cm.Snapshotter.Stat(ctx, snapshotID)
+					if err == nil { // usable as remote snapshot without unlazying.
+						defer func() {
+							// Remove tmp labels appended in this func
+							for k := range tmpLabels {
+								info.Labels[k] = ""
+							}
+							if _, err := r.cm.Snapshotter.Update(ctx, info, tmpFields...); err != nil {
+								logrus.Warn(errors.Wrapf(err,
+									"failed to remove tmp remote labels after prepare"))
+							}
+						}()
+
+						// Try the next layer as well.
+						continue
+					}
+				}
+			}
+
+			// This layer and all upper layers cannot be prepared without unlazying.
+			break
 		}
 
-		// This layer cannot be prepared without unlazying.
-		return false, nil
+		return nil, nil
 	})
-	return ok.(bool), err
+
+	return err
+}
+
+func makeTmpLabels(labels map[string]string, s session.Group) (fields []string, res map[string]string) {
+	res = make(map[string]string)
+	// Append unique ID to labels for avoiding collision of labels among calls
+	id := identity.NewID()
+	for k, v := range labels {
+		tmpKey := k + "." + id
+		fields = append(fields, "labels."+tmpKey)
+		res[tmpKey] = v
+	}
+	for i, sid := range session.AllSessionIDs(s) {
+		sidKey := "containerd.io/snapshot/remote/stargz.session." + fmt.Sprintf("%d", i) + "." + id
+		fields = append(fields, "labels."+sidKey)
+		res[sidKey] = sid
+	}
+	return
 }
 
 func (sr *immutableRef) extract(ctx context.Context, dhs DescHandlers, s session.Group) error {
@@ -705,11 +813,21 @@ func (sr *mutableRef) commit(ctx context.Context) (*immutableRef, error) {
 	return ref, nil
 }
 
-func (sr *mutableRef) Mount(ctx context.Context, readonly bool, s session.Group) (snapshot.Mountable, error) {
+func (sr *mutableRef) Mount(ctx context.Context, readonly bool, s session.Group) (m snapshot.Mountable, rerr error) {
 	sr.mu.Lock()
 	defer sr.mu.Unlock()
 
-	return sr.mount(ctx, readonly)
+	if sr.parent == nil {
+		return sr.mount(ctx, readonly)
+	}
+
+	if err := sr.parent.withRemoteSnapshotLabels(ctx, s, func() {
+		m, rerr = sr.mount(ctx, readonly)
+	}); err != nil {
+		return nil, err
+	}
+
+	return
 }
 
 func (sr *mutableRef) Commit(ctx context.Context) (ImmutableRef, error) {

cmd/buildkitd/main.go

@@ -69,6 +69,7 @@ func init() {
 type workerInitializerOpt struct {
 	config         *config.Config
 	configMetaData *toml.MetaData
+	sessionManager *session.Manager
 }
 
 type workerInitializer struct {
@@ -590,6 +591,7 @@ func newController(c *cli.Context, cfg *config.Config, md *toml.MetaData) (*cont
 	wc, err := newWorkerController(c, workerInitializerOpt{
 		config:         cfg,
 		configMetaData: md,
+		sessionManager: sessionManager,
 	})
 	if err != nil {
 		return nil, err