refactor: optimize filters for oras.ExtendedCopy

extendedcopy.go

@@ -17,16 +17,13 @@ package oras
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"regexp"
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
-	artifactspec "github.com/oras-project/artifacts-spec/specs-go/v1"
 	"oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/internal/copyutil"
 	"oras.land/oras-go/v2/internal/descriptor"
-	"oras.land/oras-go/v2/internal/docker"
 	"oras.land/oras-go/v2/registry"
 )
 
@@ -182,11 +179,31 @@ func findRoots(ctx context.Context, storage content.ReadOnlyGraphStorage, node o
 // For performance consideration, when using both FilterArtifactType and
 // FilterAnnotation, it's recommended to call FilterArtifactType first.
 func (opts *ExtendedCopyGraphOptions) FilterAnnotation(key string, regex *regexp.Regexp) {
+	filter := func(predecessors []ocispec.Descriptor, filtered []ocispec.Descriptor) []ocispec.Descriptor {
+		for _, p := range predecessors {
+			if value, ok := p.Annotations[key]; ok && (regex == nil || regex.MatchString(value)) {
+				filtered = append(filtered, p)
+			}
+		}
+		return filtered
+	}
+
 	fp := opts.FindPredecessors
 	opts.FindPredecessors = func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
 		var predecessors []ocispec.Descriptor
 		var err error
 		if fp == nil {
+			if rf, ok := src.(registry.ReferrerFinder); ok {
+				// if src is a ReferrerFinder, use Referrers() for possible memory saving
+				if err := rf.Referrers(ctx, desc, "", func(referrers []ocispec.Descriptor) error {
+					// for each page of the results, filter the referrers
+					predecessors = filter(referrers, predecessors)
+					return nil
+				}); err != nil {
+					return nil, err
+				}
+				return predecessors, nil
+			}
 			predecessors, err = src.Predecessors(ctx, desc)
 		} else {
 			predecessors, err = fp(ctx, src, desc)
@@ -195,40 +212,7 @@ func (opts *ExtendedCopyGraphOptions) FilterAnnotation(key string, regex *regexp
 			return nil, err
 		}
 		var filtered []ocispec.Descriptor
-		for _, p := range predecessors {
-			if p.Annotations == nil {
-				switch p.MediaType {
-				case docker.MediaTypeManifest, ocispec.MediaTypeImageManifest,
-					docker.MediaTypeManifestList, ocispec.MediaTypeImageIndex,
-					artifactspec.MediaTypeArtifactManifest:
-					if err = func() error {
-						rc, err := src.Fetch(ctx, p)
-						if err != nil {
-							return err
-						}
-						defer rc.Close()
-						var manifest struct {
-							Annotations map[string]string `json:"annotations"`
-						}
-						if err := json.NewDecoder(rc).Decode(&manifest); err != nil {
-							return err
-						}
-						if manifest.Annotations == nil {
-							p.Annotations = map[string]string{}
-						} else {
-							p.Annotations = manifest.Annotations
-						}
-						return nil
-					}(); err != nil {
-						return nil, err
-					}
-				}
-			}
-			if value, ok := p.Annotations[key]; ok && (regex == nil || regex.MatchString(value)) {
-				filtered = append(filtered, p)
-			}
-		}
-		return filtered, nil
+		return filter(predecessors, filtered), nil
 	}
 }
 
@@ -241,14 +225,30 @@ func (opts *ExtendedCopyGraphOptions) FilterArtifactType(regex *regexp.Regexp) {
 	if regex == nil {
 		return
 	}
+	filter := func(predecessors []ocispec.Descriptor, filtered []ocispec.Descriptor) []ocispec.Descriptor {
+		for _, p := range predecessors {
+			if regex.MatchString(p.ArtifactType) {
+				filtered = append(filtered, p)
+			}
+		}
+		return filtered
+	}
+
 	fp := opts.FindPredecessors
 	opts.FindPredecessors = func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
 		var predecessors []ocispec.Descriptor
 		var err error
 		if fp == nil {
-			// if src is a ReferrerFinder, use Referrers() to filter the predecessors.
 			if rf, ok := src.(registry.ReferrerFinder); ok {
-				return findReferrersAndFilter(rf, ctx, desc, regex)
+				// if src is a ReferrerFinder, use Referrers() for possible memory saving
+				if err := rf.Referrers(ctx, desc, "", func(referrers []ocispec.Descriptor) error {
+					// for each page of the results, filter the referrers
+					predecessors = filter(referrers, predecessors)
+					return nil
+				}); err != nil {
+					return nil, err
+				}
+				return predecessors, nil
 			}
 			predecessors, err = src.Predecessors(ctx, desc)
 		} else {
@@ -258,45 +258,6 @@ func (opts *ExtendedCopyGraphOptions) FilterArtifactType(regex *regexp.Regexp) {
 			return nil, err
 		}
 		var filtered []ocispec.Descriptor
-		// for each predecessor, decode the manifest and check its artifact type.
-		for _, p := range predecessors {
-			if p.MediaType == artifactspec.MediaTypeArtifactManifest {
-				if err = func() error {
-					rc, err := src.Fetch(ctx, p)
-					if err != nil {
-						return err
-					}
-					defer rc.Close()
-					var manifest artifactspec.Manifest
-					if err := json.NewDecoder(rc).Decode(&manifest); err != nil {
-						return err
-					}
-					if regex.MatchString(manifest.ArtifactType) {
-						filtered = append(filtered, p)
-					}
-					return nil
-				}(); err != nil {
-					return nil, err
-				}
-			}
-		}
-		return filtered, nil
-	}
-}
-
-// findReferrersAndFilter filters the predecessors with Referrers.
-func findReferrersAndFilter(rf registry.ReferrerFinder, ctx context.Context, desc ocispec.Descriptor, regex *regexp.Regexp) ([]ocispec.Descriptor, error) {
-	var predecessors []ocispec.Descriptor
-	if err := rf.Referrers(ctx, desc, "", func(referrers []ocispec.Descriptor) error {
-		// for each page of the results, do the following:
-		for _, referrer := range referrers {
-			if regex.MatchString(referrer.ArtifactType) {
-				predecessors = append(predecessors, referrer)
-			}
-		}
-		return nil
-	}); err != nil {
-		return nil, err
+		return filter(predecessors, filtered), nil
 	}
-	return predecessors, nil
 }