Provide stat descriptor for Create method during cross-repo mount (#1857)

* Allow precomputed stats on cross-mounted blobs

Signed-off-by: Michal Minář <miminar@redhat.com>

* Extended cross-repo mount tests

Signed-off-by: Michal Minář <miminar@redhat.com>

blobs.go

@@ -198,6 +198,9 @@ type CreateOptions struct {
 	Mount struct {
 		ShouldMount bool
 		From        reference.Canonical
+		// Stat allows to pass precalculated descriptor to link and return.
+		// Blob access check will be skipped if set.
+		Stat *Descriptor
 	}
 }
 

registry/storage/linkedblobstore.go

@@ -137,7 +137,7 @@ func (lbs *linkedBlobStore) Create(ctx context.Context, options ...distribution.
 	}
 
 	if opts.Mount.ShouldMount {
-		desc, err := lbs.mount(ctx, opts.Mount.From, opts.Mount.From.Digest())
+		desc, err := lbs.mount(ctx, opts.Mount.From, opts.Mount.From.Digest(), opts.Mount.Stat)
 		if err == nil {
 			// Mount successful, no need to initiate an upload session
 			return nil, distribution.ErrBlobMounted{From: opts.Mount.From, Descriptor: desc}
@@ -280,15 +280,22 @@ func (lbs *linkedBlobStore) Enumerate(ctx context.Context, ingestor func(digest.
 	return nil
 }
 
-func (lbs *linkedBlobStore) mount(ctx context.Context, sourceRepo reference.Named, dgst digest.Digest) (distribution.Descriptor, error) {
+func (lbs *linkedBlobStore) mount(ctx context.Context, sourceRepo reference.Named, dgst digest.Digest, sourceStat *distribution.Descriptor) (distribution.Descriptor, error) {
+	var stat distribution.Descriptor
+	if sourceStat == nil {
+		// look up the blob info from the sourceRepo if not already provided
 		repo, err := lbs.registry.Repository(ctx, sourceRepo)
 		if err != nil {
 			return distribution.Descriptor{}, err
 		}
-	stat, err := repo.Blobs(ctx).Stat(ctx, dgst)
+		stat, err = repo.Blobs(ctx).Stat(ctx, dgst)
 		if err != nil {
 			return distribution.Descriptor{}, err
 		}
+	} else {
+		// use the provided blob info
+		stat = *sourceStat
+	}
 
 	desc := distribution.Descriptor{
 		Size: stat.Size,

registry/storage/linkedblobstore_test.go

@@ -1,7 +1,10 @@
 package storage
 
 import (
+	"fmt"
 	"io"
+	"reflect"
+	"strconv"
 	"testing"
 
 	"github.com/docker/distribution"
@@ -16,6 +19,10 @@ func TestLinkedBlobStoreCreateWithMountFrom(t *testing.T) {
 	fooRepoName, _ := reference.ParseNamed("nm/foo")
 	fooEnv := newManifestStoreTestEnv(t, fooRepoName, "thetag")
 	ctx := context.Background()
+	stats, err := mockRegistry(t, fooEnv.registry)
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	// Build up some test layers and add them to the manifest, saving the
 	// readseekers for upload later.
@@ -48,7 +55,6 @@ func TestLinkedBlobStoreCreateWithMountFrom(t *testing.T) {
 
 	// create another repository nm/bar
 	barRepoName, _ := reference.ParseNamed("nm/bar")
-
 	barRepo, err := fooEnv.registry.Repository(ctx, barRepoName)
 	if err != nil {
 		t.Fatalf("unexpected error getting repo: %v", err)
@@ -75,4 +81,137 @@ func TestLinkedBlobStoreCreateWithMountFrom(t *testing.T) {
 			t.Fatalf("expected ErrMountFrom error, not %T: %v", err, err)
 		}
 	}
+	for dgst := range testLayers {
+		fooCanonical, _ := reference.WithDigest(fooRepoName, dgst)
+		count, exists := stats[fooCanonical.String()]
+		if !exists {
+			t.Errorf("expected entry %q not found among handled stat calls", fooCanonical.String())
+		} else if count != 1 {
+			t.Errorf("expected exactly one stat call for entry %q, not %d", fooCanonical.String(), count)
+		}
+	}
+
+	clearStats(stats)
+
+	// create yet another repository nm/baz
+	bazRepoName, _ := reference.ParseNamed("nm/baz")
+	bazRepo, err := fooEnv.registry.Repository(ctx, bazRepoName)
+	if err != nil {
+		t.Fatalf("unexpected error getting repo: %v", err)
+	}
+
+	// cross-repo mount them into a nm/baz and provide a prepopulated blob descriptor
+	for dgst := range testLayers {
+		fooCanonical, _ := reference.WithDigest(fooRepoName, dgst)
+		size, err := strconv.ParseInt("0x"+dgst.Hex()[:8], 0, 64)
+		if err != nil {
+			t.Fatal(err)
+		}
+		prepolutatedDescriptor := distribution.Descriptor{
+			Digest:    dgst,
+			Size:      size,
+			MediaType: "application/octet-stream",
+		}
+		_, err = bazRepo.Blobs(ctx).Create(ctx, WithMountFrom(fooCanonical), &statCrossMountCreateOption{
+			desc: prepolutatedDescriptor,
+		})
+		blobMounted, ok := err.(distribution.ErrBlobMounted)
+		if !ok {
+			t.Errorf("expected ErrMountFrom error, not %T: %v", err, err)
+			continue
+		}
+		if !reflect.DeepEqual(blobMounted.Descriptor, prepolutatedDescriptor) {
+			t.Errorf("unexpected descriptor: %#+v != %#+v", blobMounted.Descriptor, prepolutatedDescriptor)
+		}
+	}
+	// this time no stat calls will be made
+	if len(stats) != 0 {
+		t.Errorf("unexpected number of stats made: %d != %d", len(stats), len(testLayers))
+	}
+}
+
+func clearStats(stats map[string]int) {
+	for k := range stats {
+		delete(stats, k)
+	}
+}
+
+// mockRegistry sets a mock blob descriptor service factory that overrides
+// statter's Stat method to note each attempt to stat a blob in any repository.
+// Returned stats map contains canonical references to blobs with a number of
+// attempts.
+func mockRegistry(t *testing.T, nm distribution.Namespace) (map[string]int, error) {
+	registry, ok := nm.(*registry)
+	if !ok {
+		return nil, fmt.Errorf("not an expected type of registry: %T", nm)
+	}
+	stats := make(map[string]int)
+
+	registry.blobDescriptorServiceFactory = &mockBlobDescriptorServiceFactory{
+		t:     t,
+		stats: stats,
+	}
+
+	return stats, nil
+}
+
+type mockBlobDescriptorServiceFactory struct {
+	t     *testing.T
+	stats map[string]int
+}
+
+func (f *mockBlobDescriptorServiceFactory) BlobAccessController(svc distribution.BlobDescriptorService) distribution.BlobDescriptorService {
+	return &mockBlobDescriptorService{
+		BlobDescriptorService: svc,
+		t:     f.t,
+		stats: f.stats,
+	}
+}
+
+type mockBlobDescriptorService struct {
+	distribution.BlobDescriptorService
+	t     *testing.T
+	stats map[string]int
+}
+
+var _ distribution.BlobDescriptorService = &mockBlobDescriptorService{}
+
+func (bs *mockBlobDescriptorService) Stat(ctx context.Context, dgst digest.Digest) (distribution.Descriptor, error) {
+	statter, ok := bs.BlobDescriptorService.(*linkedBlobStatter)
+	if !ok {
+		return distribution.Descriptor{}, fmt.Errorf("unexpected blob descriptor service: %T", bs.BlobDescriptorService)
+	}
+
+	name := statter.repository.Named()
+	canonical, err := reference.WithDigest(name, dgst)
+	if err != nil {
+		return distribution.Descriptor{}, fmt.Errorf("failed to make canonical reference: %v", err)
+	}
+
+	bs.stats[canonical.String()]++
+	bs.t.Logf("calling Stat on %s", canonical.String())
+
+	return bs.BlobDescriptorService.Stat(ctx, dgst)
+}
+
+// statCrossMountCreateOptions ensures the expected options type is passed, and optionally pre-fills the cross-mount stat info
+type statCrossMountCreateOption struct {
+	desc distribution.Descriptor
+}
+
+var _ distribution.BlobCreateOption = statCrossMountCreateOption{}
+
+func (f statCrossMountCreateOption) Apply(v interface{}) error {
+	opts, ok := v.(*distribution.CreateOptions)
+	if !ok {
+		return fmt.Errorf("Unexpected create options: %#v", v)
+	}
+
+	if !opts.Mount.ShouldMount {
+		return nil
+	}
+
+	opts.Mount.Stat = &f.desc
+
+	return nil
 }