Add image pull concurrency limit.

Signed-off-by: Lantao Liu <[email protected]>

Author: Random-Liu

client.go

@@ -61,6 +61,7 @@ import (
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/pkg/errors"
+	"golang.org/x/sync/semaphore"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
 )
@@ -292,6 +293,9 @@ type RemoteContext struct {
 	// platforms will be used to create a PlatformMatcher with no ordering
 	// preference.
 	Platforms []string
+
+	// MaxConcurrentDownloads is the max concurrent content downloads for each pull.
+	MaxConcurrentDownloads int
 }
 
 func defaultRemoteContext() *RemoteContext {
@@ -407,6 +411,7 @@ func (c *Client) fetch(ctx context.Context, rCtx *RemoteContext, ref string, lim
 
 		isConvertible bool
 		converterFunc func(context.Context, ocispec.Descriptor) (ocispec.Descriptor, error)
+		limiter       *semaphore.Weighted
 	)
 
 	if desc.MediaType == images.MediaTypeDockerSchema1Manifest && rCtx.ConvertSchema1 {
@@ -453,7 +458,10 @@ func (c *Client) fetch(ctx context.Context, rCtx *RemoteContext, ref string, lim
 		}
 	}
 
-	if err := images.Dispatch(ctx, handler, desc); err != nil {
+	if rCtx.MaxConcurrentDownloads > 0 {
+		limiter = semaphore.NewWeighted(int64(rCtx.MaxConcurrentDownloads))
+	}
+	if err := images.Dispatch(ctx, handler, limiter, desc); err != nil {
 		return images.Image{}, err
 	}
 

client_opts.go

@@ -178,3 +178,11 @@ func WithImageHandler(h images.Handler) RemoteOpt {
 		return nil
 	}
 }
+
+// WithMaxConcurrentDownloads sets max concurrent download limit.
+func WithMaxConcurrentDownloads(max int) RemoteOpt {
+	return func(client *Client, c *RemoteContext) error {
+		c.MaxConcurrentDownloads = max
+		return nil
+	}
+}

client_test.go

@@ -327,6 +327,23 @@ func TestImagePullSchema1(t *testing.T) {
 	}
 }
 
+func TestImagePullWithConcurrencyLimit(t *testing.T) {
+	client, err := newClient(t, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	ctx, cancel := testContext()
+	defer cancel()
+	_, err = client.Pull(ctx, testImage,
+		WithPlatformMatcher(platforms.Default()),
+		WithMaxConcurrentDownloads(2))
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestClientReconnect(t *testing.T) {
 	t.Parallel()
 

images/handlers.go

@@ -26,6 +26,7 @@ import (
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
+	"golang.org/x/sync/semaphore"
 )
 
 var (
@@ -108,19 +109,30 @@ func Walk(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) err
 // handler may return `ErrSkipDesc` to signal to the dispatcher to not traverse
 // any children.
 //
+// A concurrency limiter can be passed in to limit the number of concurrent
+// handlers running. When limiter is nil, there is no limit.
+//
 // Typically, this function will be used with `FetchHandler`, often composed
 // with other handlers.
 //
 // If any handler returns an error, the dispatch session will be canceled.
-func Dispatch(ctx context.Context, handler Handler, descs ...ocispec.Descriptor) error {
+func Dispatch(ctx context.Context, handler Handler, limiter *semaphore.Weighted, descs ...ocispec.Descriptor) error {
 	eg, ctx := errgroup.WithContext(ctx)
 	for _, desc := range descs {
 		desc := desc
 
+		if limiter != nil {
+			if err := limiter.Acquire(ctx, 1); err != nil {
+				return err
+			}
+		}
 		eg.Go(func() error {
 			desc := desc
 
 			children, err := handler.Handle(ctx, desc)
+			if limiter != nil {
+				limiter.Release(1)
+			}
 			if err != nil {
 				if errors.Cause(err) == ErrSkipDesc {
 					return nil // don't traverse the children.
@@ -129,7 +141,7 @@ func Dispatch(ctx context.Context, handler Handler, descs ...ocispec.Descriptor)
 			}
 
 			if len(children) > 0 {
-				return Dispatch(ctx, handler, children...)
+				return Dispatch(ctx, handler, limiter, children...)
 			}
 
 			return nil

remotes/handlers.go

@@ -181,7 +181,7 @@ func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, pr
 		pushHandler,
 	)
 
-	if err := images.Dispatch(ctx, images.Handlers(handlers...), desc); err != nil {
+	if err := images.Dispatch(ctx, images.Handlers(handlers...), nil, desc); err != nil {
 		return err
 	}
 

vendor.conf

@@ -29,7 +29,7 @@ github.com/pkg/errors v0.8.0
 github.com/opencontainers/go-digest c9281466c8b2f606084ac71339773efd177436e7
 golang.org/x/sys 1b2967e3c290b7c545b3db0deeda16e9be4f98a2 https://github.com/golang/sys
 github.com/opencontainers/image-spec v1.0.1
-golang.org/x/sync 450f422ab23cf9881c94e2db30cac0eb1b7cf80c
+golang.org/x/sync 42b317875d0fa942474b76e1b46a6060d720ae6e
 github.com/BurntSushi/toml a368813c5e648fee92e5f6c30e3944ff9d5e8895
 github.com/grpc-ecosystem/go-grpc-prometheus 6b7015e65d366bf3f19b2b2a000a831940f0f7e0
 github.com/Microsoft/go-winio v0.4.11