feat: add concurrency support for pull and push operations

Signed-off-by: chlins <[email protected]>
CloudNativeAI · Dec 17, 2024 · dc09343 · dc09343
1 parent 44c0c19
commit dc09343
Show file tree

Hide file tree

Showing 8 changed files with 93 additions and 33 deletions.
diff --git a/cmd/login.go b/cmd/login.go
@@ -89,9 +89,8 @@ func runLogin(ctx context.Context, registry string) error {
 
 	fmt.Println("\nLogging In...")
 
-	opts := []backend.Option{}
-	if loginConfig.PlainHTTP {
-		opts = append(opts, backend.WithPlainHTTP())
+	opts := []backend.Option{
+		backend.WithPlainHTTP(loginConfig.PlainHTTP),
 	}
 
 	if err := b.Login(ctx, registry, loginConfig.Username, loginConfig.Password, opts...); err != nil {

diff --git a/cmd/pull.go b/cmd/pull.go
@@ -38,13 +38,18 @@ var pullCmd = &cobra.Command{
 	SilenceUsage:       true,
 	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
 	RunE: func(cmd *cobra.Command, args []string) error {
+		if err := pullConfig.Validate(); err != nil {
+			return err
+		}
+
 		return runPull(context.Background(), args[0])
 	},
 }
 
 // init initializes pull command.
 func init() {
 	flags := pullCmd.Flags()
+	flags.IntVar(&pullConfig.Concurrency, "concurrency", pullConfig.Concurrency, "specify the number of concurrent pull operations (default: 3)")
 	flags.BoolVar(&pullConfig.PlainHTTP, "plain-http", false, "use plain HTTP instead of HTTPS")
 	flags.BoolVar(&pullConfig.Insecure, "insecure", false, "use insecure connection for the pull operation and skip TLS verification")
 	flags.StringVar(&pullConfig.Proxy, "proxy", "", "use proxy for the pull operation")
@@ -66,9 +71,10 @@ func runPull(ctx context.Context, target string) error {
 		return fmt.Errorf("target is required")
 	}
 
-	opts := []backend.Option{backend.WithInsecure(pullConfig.Insecure)}
-	if pullConfig.PlainHTTP {
-		opts = append(opts, backend.WithPlainHTTP())
+	opts := []backend.Option{
+		backend.WithInsecure(pullConfig.Insecure),
+		backend.WithPlainHTTP(pullConfig.PlainHTTP),
+		backend.WithConcurrency(pullConfig.Concurrency),
 	}
 
 	if pullConfig.Proxy != "" {

diff --git a/cmd/push.go b/cmd/push.go
@@ -38,13 +38,18 @@ var pushCmd = &cobra.Command{
 	SilenceUsage:       true,
 	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
 	RunE: func(cmd *cobra.Command, args []string) error {
+		if err := pushConfig.Validate(); err != nil {
+			return err
+		}
+
 		return runPush(context.Background(), args[0])
 	},
 }
 
 // init initializes push command.
 func init() {
 	flags := pushCmd.Flags()
+	flags.IntVar(&pushConfig.Concurrency, "concurrency", pushConfig.Concurrency, "specify the number of concurrent push operations (default: 3)")
 	flags.BoolVar(&pushConfig.PlainHTTP, "plain-http", false, "use plain HTTP instead of HTTPS")
 
 	if err := viper.BindPFlags(flags); err != nil {
@@ -59,9 +64,9 @@ func runPush(ctx context.Context, target string) error {
 		return err
 	}
 
-	opts := []backend.Option{}
-	if pushConfig.PlainHTTP {
-		opts = append(opts, backend.WithPlainHTTP())
+	opts := []backend.Option{
+		backend.WithPlainHTTP(pushConfig.PlainHTTP),
+		backend.WithConcurrency(pushConfig.Concurrency),
 	}
 
 	if err := b.Push(ctx, target, opts...); err != nil {

diff --git a/pkg/backend/options.go b/pkg/backend/options.go
@@ -19,16 +19,24 @@ package backend
 type Option func(*Options)
 
 type Options struct {
-	plainHTTP bool
-	proxy     string
-	insecure  bool
-	output    string
+	concurrency int
+	plainHTTP   bool
+	proxy       string
+	insecure    bool
+	output      string
+}
+
+// WithConcurrency sets the concurrency option.
+func WithConcurrency(concurrency int) Option {
+	return func(opts *Options) {
+		opts.concurrency = concurrency
+	}
 }
 
 // WithPlainHTTP sets the plain HTTP option.
-func WithPlainHTTP() Option {
+func WithPlainHTTP(plainHTTP bool) Option {
 	return func(opts *Options) {
-		opts.plainHTTP = true
+		opts.plainHTTP = plainHTTP
 	}
 }
 

diff --git a/pkg/backend/pull.go b/pkg/backend/pull.go
@@ -26,6 +26,7 @@ import (
 	"net/url"
 
 	"github.com/CloudNativeAI/modctl/pkg/storage"
+	"golang.org/x/sync/errgroup"
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2/registry/remote"
@@ -110,12 +111,15 @@ func (b *backend) Pull(ctx context.Context, target string, opts ...Option) error
 	// note: the order is important, manifest should be pushed at last.
 
 	// copy the layers.
-	// TODO: parallelize the layer copy.
 	dst := b.store
+	g := &errgroup.Group{}
+	g.SetLimit(options.concurrency)
 	for _, layer := range manifest.Layers {
-		if err := pullIfNotExist(ctx, pb, promptCopyingBlob, src, dst, layer, repo, tag); err != nil {
-			return fmt.Errorf("failed to pull blob to local: %w", err)
-		}
+		g.Go(func() error { return pullIfNotExist(ctx, pb, promptCopyingBlob, src, dst, layer, repo, tag) })
+	}
+
+	if err := g.Wait(); err != nil {
+		return fmt.Errorf("failed to pull blob to local: %w", err)
 	}
 
 	// copy the config.

diff --git a/pkg/backend/push.go b/pkg/backend/push.go
@@ -23,6 +23,7 @@ import (
 	"fmt"
 
 	"github.com/CloudNativeAI/modctl/pkg/storage"
+	"golang.org/x/sync/errgroup"
 
 	godigest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -91,11 +92,14 @@ func (b *backend) Push(ctx context.Context, target string, opts ...Option) error
 	// note: the order is important, manifest should be pushed at last.
 
 	// copy the layers.
-	// TODO: parallelize the layer copy.
+	g := &errgroup.Group{}
+	g.SetLimit(options.concurrency)
 	for _, layer := range manifest.Layers {
-		if err := pushIfNotExist(ctx, pb, promptCopyingBlob, src, dst, layer, repo, tag); err != nil {
-			return fmt.Errorf("failed to push blob to remote: %w", err)
-		}
+		g.Go(func() error { return pushIfNotExist(ctx, pb, promptCopyingBlob, src, dst, layer, repo, tag) })
+	}
+
+	if err := g.Wait(); err != nil {
+		return fmt.Errorf("failed to push blob to remote: %w", err)
 	}
 
 	// copy the config.

diff --git a/pkg/config/pull.go b/pkg/config/pull.go
@@ -16,18 +16,35 @@
 
 package config
 
+import "fmt"
+
+const (
+	// defaultPullConcurrency is the default number of concurrent pull operations.
+	defaultPullConcurrency = 3
+)
+
 type Pull struct {
-	PlainHTTP  bool
-	Proxy      string
-	Insecure   bool
-	ExtractDir string
+	Concurrency int
+	PlainHTTP   bool
+	Proxy       string
+	Insecure    bool
+	ExtractDir  string
 }
 
 func NewPull() *Pull {
 	return &Pull{
-		PlainHTTP:  false,
-		Proxy:      "",
-		Insecure:   false,
-		ExtractDir: "",
+		Concurrency: defaultPullConcurrency,
+		PlainHTTP:   false,
+		Proxy:       "",
+		Insecure:    false,
+		ExtractDir:  "",
+	}
+}
+
+func (p *Pull) Validate() error {
+	if p.Concurrency < 1 {
+		return fmt.Errorf("invalid concurrency: %d", p.Concurrency)
 	}
+
+	return nil
 }
diff --git a/pkg/config/push.go b/pkg/config/push.go
@@ -16,12 +16,29 @@
 
 package config
 
+import "fmt"
+
+const (
+	// defaultPushConcurrency is the default number of concurrent push operations.
+	defaultPushConcurrency = 3
+)
+
 type Push struct {
-	PlainHTTP bool
+	Concurrency int
+	PlainHTTP   bool
 }
 
 func NewPush() *Pull {
 	return &Pull{
-		PlainHTTP: false,
+		Concurrency: defaultPushConcurrency,
+		PlainHTTP:   false,
+	}
+}
+
+func (p *Push) Validate() error {
+	if p.Concurrency < 1 {
+		return fmt.Errorf("invalid concurrency: %d", p.Concurrency)
 	}
+
+	return nil
 }