cache: Remove ImageRef from DescHandlers

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

cache/manager.go

@@ -157,6 +157,9 @@ func (cm *cacheManager) GetByBlob(ctx context.Context, desc ocispec.Descriptor,
 		if p != nil {
 			releaseParent = true
 		}
+		if err := setImageRefMetadata(ref, opts...); err != nil {
+			return nil, errors.Wrapf(err, "failed to append image ref metadata to ref %s", ref.ID())
+		}
 		return ref, nil
 	}
 
@@ -234,6 +237,10 @@ func (cm *cacheManager) GetByBlob(ctx context.Context, desc ocispec.Descriptor,
 		return nil, err
 	}
 
+	if err := setImageRefMetadata(rec, opts...); err != nil {
+		return nil, errors.Wrapf(err, "failed to append image ref metadata to ref %s", rec.ID())
+	}
+
 	queueDiffID(rec.md, diffID.String())
 	queueBlob(rec.md, desc.Digest.String())
 	queueChainID(rec.md, chainID.String())
@@ -429,6 +436,10 @@ func (cm *cacheManager) getRecord(ctx context.Context, id string, opts ...RefOpt
 		return nil, err
 	}
 
+	if err := setImageRefMetadata(rec, opts...); err != nil {
+		return nil, errors.Wrapf(err, "failed to append image ref metadata to ref %s", rec.ID())
+	}
+
 	cm.records[id] = rec
 	if err := checkLazyProviders(rec); err != nil {
 		return nil, err
@@ -515,6 +526,10 @@ func (cm *cacheManager) New(ctx context.Context, s ImmutableRef, opts ...RefOpti
 		return nil, err
 	}
 
+	if err := setImageRefMetadata(rec, opts...); err != nil {
+		return nil, errors.Wrapf(err, "failed to append image ref metadata to ref %s", rec.ID())
+	}
+
 	cm.mu.Lock()
 	defer cm.mu.Unlock()
 
@@ -1004,6 +1019,31 @@ func WithCreationTime(tm time.Time) RefOption {
 	}
 }
 
+// Need a separate type for imageRef because it needs to be called outside
+// initializeMetadata while still being a RefOption, so wrapping it in a
+// different type ensures initializeMetadata won't catch it too and duplicate
+// setting the metadata.
+type imageRefOption func(m withMetadata) error
+
+// WithImageRef appends the given imageRef to the cache ref's metadata
+func WithImageRef(imageRef string) RefOption {
+	return imageRefOption(func(m withMetadata) error {
+		return appendImageRef(m.Metadata(), imageRef)
+	})
+}
+
+func setImageRefMetadata(m withMetadata, opts ...RefOption) error {
+	md := m.Metadata()
+	for _, opt := range opts {
+		if fn, ok := opt.(imageRefOption); ok {
+			if err := fn(m); err != nil {
+				return err
+			}
+		}
+	}
+	return md.Commit()
+}
+
 func initializeMetadata(m withMetadata, parent string, opts ...RefOption) error {
 	md := m.Metadata()
 	if tm := GetCreatedAt(md); !tm.IsZero() {

cache/metadata.go

@@ -28,6 +28,7 @@ const keyBlob = "cache.blob"
 const keySnapshot = "cache.snapshot"
 const keyBlobOnly = "cache.blobonly"
 const keyMediaType = "cache.mediatype"
+const keyImageRefs = "cache.imageRefs"
 
 // BlobSize is the packed blob size as specified in the oci descriptor
 const keyBlobSize = "cache.blobsize"
@@ -310,6 +311,40 @@ func getSize(si *metadata.StorageItem) int64 {
 	return size
 }
 
+func appendImageRef(si *metadata.StorageItem, s string) error {
+	return si.GetAndSetValue(keyImageRefs, func(v *metadata.Value) (*metadata.Value, error) {
+		var imageRefs []string
+		if v != nil {
+			if err := v.Unmarshal(&imageRefs); err != nil {
+				return nil, err
+			}
+		}
+		for _, existing := range imageRefs {
+			if existing == s {
+				return nil, metadata.ErrSkipSetValue
+			}
+		}
+		imageRefs = append(imageRefs, s)
+		v, err := metadata.NewValue(imageRefs)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to create imageRefs value")
+		}
+		return v, nil
+	})
+}
+
+func getImageRefs(si *metadata.StorageItem) []string {
+	v := si.Get(keyImageRefs)
+	if v == nil {
+		return nil
+	}
+	var refs []string
+	if err := v.Unmarshal(&refs); err != nil {
+		return nil
+	}
+	return refs
+}
+
 func queueBlobSize(si *metadata.StorageItem, s int64) error {
 	v, err := metadata.NewValue(s)
 	if err != nil {

cache/metadata/metadata.go

@@ -372,6 +372,22 @@ func (s *StorageItem) SetValue(b *bolt.Bucket, key string, v *Value) error {
 	return nil
 }
 
+var ErrSkipSetValue = errors.New("skip setting metadata value")
+
+func (s *StorageItem) GetAndSetValue(key string, fn func(*Value) (*Value, error)) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.Update(func(b *bolt.Bucket) error {
+		v, err := fn(s.values[key])
+		if errors.Is(err, ErrSkipSetValue) {
+			return nil
+		} else if err != nil {
+			return err
+		}
+		return s.SetValue(b, key, v)
+	})
+}
+
 type Value struct {
 	Value json.RawMessage `json:"value,omitempty"`
 	Index string          `json:"index,omitempty"`

cache/opts.go

@@ -10,7 +10,6 @@ import (
 
 type DescHandler struct {
 	Provider content.Provider
-	ImageRef string
 	Progress progress.Controller
 }
 

cache/remote.go

@@ -58,26 +58,46 @@ func (sr *immutableRef) GetRemote(ctx context.Context, createIfNeeded bool, comp
 			}
 		}
 
-		dh := sr.descHandlers[desc.Digest]
+		// update distribution source annotation for lazy-refs (non-lazy refs
+		// will already have their dsl stored in the content store, which is
+		// used by the push handlers)
+		if isLazy, err := ref.isLazy(ctx); err != nil {
+			return nil, err
+		} else if isLazy {
+			imageRefs := getImageRefs(ref.md)
+			for _, imageRef := range imageRefs {
+				refspec, err := reference.Parse(imageRef)
+				if err != nil {
+					return nil, err
+				}
 
-		// update distribution source annotation
-		if dh != nil && dh.ImageRef != "" {
-			refspec, err := reference.Parse(dh.ImageRef)
-			if err != nil {
-				return nil, err
-			}
+				u, err := url.Parse("dummy://" + refspec.Locator)
+				if err != nil {
+					return nil, err
+				}
 
-			u, err := url.Parse("dummy://" + refspec.Locator)
-			if err != nil {
-				return nil, err
-			}
+				source, repo := u.Hostname(), strings.TrimPrefix(u.Path, "/")
+				if desc.Annotations == nil {
+					desc.Annotations = make(map[string]string)
+				}
+				dslKey := fmt.Sprintf("%s.%s", "containerd.io/distribution.source", source)
 
-			source, repo := u.Hostname(), strings.TrimPrefix(u.Path, "/")
-			if desc.Annotations == nil {
-				desc.Annotations = make(map[string]string)
+				var existingRepos []string
+				if existings, ok := desc.Annotations[dslKey]; ok {
+					existingRepos = strings.Split(existings, ",")
+				}
+				addNewRepo := true
+				for _, existing := range existingRepos {
+					if existing == repo {
+						addNewRepo = false
+						break
+					}
+				}
+				if addNewRepo {
+					existingRepos = append(existingRepos, repo)
+				}
+				desc.Annotations[dslKey] = strings.Join(existingRepos, ",")
 			}
-			dslKey := fmt.Sprintf("%s.%s", "containerd.io/distribution.source", source)
-			desc.Annotations[dslKey] = repo
 		}
 
 		remote.Descriptors = append(remote.Descriptors, desc)
@@ -161,9 +181,11 @@ func (p lazyRefProvider) Unlazy(ctx context.Context) error {
 			return nil, err
 		}
 
-		if p.dh.ImageRef != "" {
+		if imageRefs := getImageRefs(p.ref.md); len(imageRefs) > 0 {
+			// just use the first image ref, it's arbitrary
+			imageRef := imageRefs[0]
 			if GetDescription(p.ref.md) == "" {
-				queueDescription(p.ref.md, "pulled from "+p.dh.ImageRef)
+				queueDescription(p.ref.md, "pulled from "+imageRef)
 				err := p.ref.md.Commit()
 				if err != nil {
 					return nil, err

source/containerimage/pull.go

@@ -212,7 +212,6 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 
 			descHandler := &cache.DescHandler{
 				Provider: p.manifest.Remote.Provider,
-				ImageRef: p.manifest.Ref,
 				Progress: progressController,
 			}
 
@@ -275,7 +274,8 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (ir cache.Immuta
 	var parent cache.ImmutableRef
 	for _, layerDesc := range p.manifest.Remote.Descriptors {
 		parent = current
-		current, err = p.CacheAccessor.GetByBlob(ctx, layerDesc, parent, p.descHandlers)
+		current, err = p.CacheAccessor.GetByBlob(ctx, layerDesc, parent,
+			p.descHandlers, cache.WithImageRef(p.manifest.Ref))
 		if parent != nil {
 			parent.Release(context.TODO())
 		}