refactor(perf): optimize cases with no child referrers in recursiveCopy (#1729)

Signed-off-by: Billy Zha <[email protected]>

Author: qweeah

cmd/oras/root/cp.go

@@ -225,43 +225,63 @@ func doCopy(ctx context.Context, copyHandler status.CopyHandler, src oras.ReadOn
 // recursiveCopy copies an artifact and its referrers from one target to another.
 // If the artifact is a manifest list or index, referrers of its manifests are copied as well.
 func recursiveCopy(ctx context.Context, src oras.ReadOnlyGraphTarget, dst oras.Target, dstRef string, root ocispec.Descriptor, opts oras.ExtendedCopyOptions) error {
-	if root.MediaType == ocispec.MediaTypeImageIndex || root.MediaType == docker.MediaTypeManifestList {
-		fetched, err := content.FetchAll(ctx, src, root)
-		if err != nil {
-			return err
-		}
-		var index ocispec.Index
-		if err = json.Unmarshal(fetched, &index); err != nil {
-			return nil
-		}
-
-		referrers, err := graph.FindPredecessors(ctx, src, index.Manifests, opts)
-		if err != nil {
-			return err
-		}
-		referrers = slices.DeleteFunc(referrers, func(desc ocispec.Descriptor) bool {
-			return content.Equal(desc, root)
-		})
-
-		findPredecessor := opts.FindPredecessors
-		opts.FindPredecessors = func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
-			descs, err := findPredecessor(ctx, src, desc)
-			if err != nil {
-				return nil, err
-			}
-			if content.Equal(desc, root) {
-				// make sure referrers of child manifests are copied by pointing them to root
-				descs = append(descs, referrers...)
-			}
-			return descs, nil
-		}
+	var err error
+	if opts, err = prepareCopyOption(ctx, src, dst, root, opts); err != nil {
+		return err
 	}
 
-	var err error
 	if dstRef == "" || dstRef == root.Digest.String() {
 		err = oras.ExtendedCopyGraph(ctx, src, dst, root, opts.ExtendedCopyGraphOptions)
 	} else {
 		_, err = oras.ExtendedCopy(ctx, src, root.Digest.String(), dst, dstRef, opts)
 	}
 	return err
 }
+
+func prepareCopyOption(ctx context.Context, src oras.ReadOnlyGraphTarget, dst oras.Target, root ocispec.Descriptor, opts oras.ExtendedCopyOptions) (oras.ExtendedCopyOptions, error) {
+	if root.MediaType != ocispec.MediaTypeImageIndex && root.MediaType != docker.MediaTypeManifestList {
+		return opts, nil
+	}
+
+	fetched, err := content.FetchAll(ctx, src, root)
+	if err != nil {
+		return opts, err
+	}
+	var index ocispec.Index
+	if err = json.Unmarshal(fetched, &index); err != nil {
+		return opts, err
+	}
+
+	if len(index.Manifests) == 0 {
+		// no child manifests, thus no child referrers
+		return opts, nil
+	}
+
+	referrers, err := graph.FindPredecessors(ctx, src, index.Manifests, opts)
+	if err != nil {
+		return opts, err
+	}
+
+	referrers = slices.DeleteFunc(referrers, func(desc ocispec.Descriptor) bool {
+		return content.Equal(desc, root)
+	})
+
+	if len(referrers) == 0 {
+		// no child referrers
+		return opts, nil
+	}
+
+	findPredecessor := opts.FindPredecessors
+	opts.FindPredecessors = func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		descs, err := findPredecessor(ctx, src, desc)
+		if err != nil {
+			return nil, err
+		}
+		if content.Equal(desc, root) {
+			// make sure referrers of child manifests are copied by pointing them to root
+			descs = append(descs, referrers...)
+		}
+		return descs, nil
+	}
+	return opts, nil
+}

cmd/oras/root/cp_test.go

@@ -20,15 +20,20 @@ package root
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"oras.land/oras-go/v2"
+	"oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/content/memory"
 	"oras.land/oras-go/v2/registry/remote"
 	"oras.land/oras/cmd/oras/internal/display/status"
@@ -197,3 +202,143 @@ func Test_doCopy_mounted(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func Test_prepareCopyOption_nonIndex(t *testing.T) {
+	ctx := context.Background()
+	root := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageManifest,
+	}
+	if _, err := prepareCopyOption(ctx, nil, nil, root, oras.ExtendedCopyOptions{}); err != nil {
+		t.Errorf("prepareCopyOption() error = %v, wantErr false", err)
+	}
+}
+
+var errMockedFetch = fmt.Errorf("fetch error")
+
+// fetchFailingReadOnlyGraphTarget is a mock implementation of oras.ReadOnlyGraphTarget
+type fetchFailingReadOnlyGraphTarget struct {
+	oras.ReadOnlyGraphTarget
+}
+
+// Fetch simulates a failure when fetching content from the source.
+func (m *fetchFailingReadOnlyGraphTarget) Fetch(ctx context.Context, target ocispec.Descriptor) (io.ReadCloser, error) {
+	return nil, errMockedFetch
+}
+
+func Test_prepareCopyOption_fetchFailure(t *testing.T) {
+	ctx := context.Background()
+	src := &fetchFailingReadOnlyGraphTarget{}
+	dst := memory.New()
+	root := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageIndex,
+		Digest:    digest.FromString("nonexistent"),
+		Size:      int64(len("nonexistent")),
+	}
+
+	if _, err := prepareCopyOption(ctx, src, dst, root, oras.ExtendedCopyOptions{}); err != errMockedFetch {
+		t.Errorf("prepareCopyOption() error = %v, want %v", err, errMockedFetch)
+	}
+}
+
+func Test_recursiveCopy_prepareCopyOptionFailure(t *testing.T) {
+	ctx := context.Background()
+	src := &fetchFailingReadOnlyGraphTarget{}
+	dst := memory.New()
+	root := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageIndex,
+		Digest:    digest.FromString("nonexistent"),
+		Size:      int64(len("nonexistent")),
+	}
+
+	if _, err := prepareCopyOption(ctx, src, dst, root, oras.ExtendedCopyOptions{}); err != errMockedFetch {
+		t.Errorf("prepareCopyOption() error = %v, want %v", err, errMockedFetch)
+	}
+}
+
+// invalidJSONReadOnlyGraphTarget is a mock implementation of oras.ReadOnlyGraphTarget
+// that returns invalid JSON data to simulate a JSON unmarshalling failure.
+type invalidJSONReadOnlyGraphTarget struct {
+	oras.ReadOnlyGraphTarget
+}
+
+// Fetch simulates a successful fetch of invalid JSON data.
+func (m *invalidJSONReadOnlyGraphTarget) Fetch(ctx context.Context, target ocispec.Descriptor) (io.ReadCloser, error) {
+	// Return invalid JSON data
+	return io.NopCloser(strings.NewReader("invalid-json")), nil
+}
+
+func Test_prepareCopyOption_jsonUnmarshalFailure(t *testing.T) {
+	ctx := context.Background()
+	src := &invalidJSONReadOnlyGraphTarget{}
+	dst := memory.New()
+	root := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageIndex,
+		Digest:    digest.FromString("invalid-json"),
+		Size:      int64(len("invalid-json")),
+	}
+	_, err := prepareCopyOption(ctx, src, dst, root, oras.ExtendedCopyOptions{})
+	if _, ok := err.(*json.SyntaxError); !ok {
+		t.Errorf("prepareCopyOption() error = %v, want json.SyntaxError", err)
+	}
+}
+
+// mockReferrersFailingSource is a mock implementation of oras.ReadOnlyGraphTarget
+// that simulates a failure when fetching referrers.
+type mockReferrersFailingSource struct {
+	oras.ReadOnlyGraphTarget
+	indexContent string
+}
+
+// Fetch simulates successful fetching of index content.
+func (m *mockReferrersFailingSource) Fetch(ctx context.Context, target ocispec.Descriptor) (io.ReadCloser, error) {
+	// Return valid JSON data to pass the fetch step
+	return io.NopCloser(strings.NewReader(m.indexContent)), nil
+}
+
+func Test_prepareCopyOption_referrersFailure(t *testing.T) {
+
+	ctx := context.Background()
+	mockedIndex := `{"schemaVersion":2,"manifests":[{"mediaType":"application/vnd.oci.image.manifest.v1+json","digest":"sha256:44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a","size":2}]}`
+	src := &mockReferrersFailingSource{indexContent: mockedIndex}
+	dst := memory.New()
+	root := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageIndex,
+		Digest:    digest.FromString(mockedIndex),
+		Size:      int64(len(mockedIndex)),
+	}
+	errMockedReferrers := fmt.Errorf("failed to get referrers")
+	opts := oras.ExtendedCopyOptions{
+		ExtendedCopyGraphOptions: oras.ExtendedCopyGraphOptions{
+			FindPredecessors: func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				return nil, errMockedReferrers
+			},
+		},
+	}
+
+	if _, err := prepareCopyOption(ctx, src, dst, root, opts); err != errMockedReferrers {
+		t.Errorf("prepareCopyOption() error = %v, wantErr %v", err, errMockedReferrers)
+	}
+}
+
+func Test_prepareCopyOption_noReferrers(t *testing.T) {
+	ctx := context.Background()
+	mockedIndex := `{"schemaVersion":2,"manifests":[{"mediaType":"application/vnd.oci.image.manifest.v1+json","digest":"sha256:44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a","size":2}]}`
+	src := &mockReferrersFailingSource{indexContent: mockedIndex}
+	dst := memory.New()
+	root := ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageIndex,
+		Digest:    digest.FromString(mockedIndex),
+		Size:      int64(len(mockedIndex)),
+	}
+	opts := oras.ExtendedCopyOptions{
+		ExtendedCopyGraphOptions: oras.ExtendedCopyGraphOptions{
+			FindPredecessors: func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+				return nil, nil
+			},
+		},
+	}
+
+	if _, err := prepareCopyOption(ctx, src, dst, root, opts); err != nil {
+		t.Errorf("prepareCopyOption() error = %v, wantErr false", err)
+	}
+}