Refactor code to extract a common task runner

Signed-off-by: cd1989 <chende@caicloud.io>

src/common/utils/passports.go

@@ -1,5 +1,12 @@
 package utils
 
+import (
+	"context"
+	"sync"
+
+	"github.com/goharbor/harbor/src/common/utils/log"
+)
+
 // PassportsPool holds a given number of passports, they can be applied or be revoked. PassportsPool
 // is used to control the concurrency of tasks, the pool size determine the max concurrency. When users
 // want to start a goroutine to perform some task, they must apply a passport firstly, and after finish
@@ -48,3 +55,74 @@ func (p *passportsPool) Revoke() bool {
 		return false
 	}
 }
+
+// LimitedConcurrentRunner is used to run tasks, but limit the max concurrency.
+type LimitedConcurrentRunner interface {
+	// AddTask adds a task to run
+	AddTask(task func() error)
+	// Wait waits all the tasks to be finished
+	Wait()
+	// Cancel cancels all tasks, tasks that already started will continue to run
+	Cancel()
+	// IsCancelled checks whether context is cancelled. This happens when some task encountered
+	// critical errors.
+	IsCancelled() bool
+}
+
+type limitedConcurrentRunner struct {
+	wg            *sync.WaitGroup
+	ctx           context.Context
+	cancel        context.CancelFunc
+	passportsPool PassportsPool
+}
+
+// NewLimitedConcurrentRunner creates a runner
+func NewLimitedConcurrentRunner(limit int) LimitedConcurrentRunner {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &limitedConcurrentRunner{
+		wg:            new(sync.WaitGroup),
+		ctx:           ctx,
+		cancel:        cancel,
+		passportsPool: NewPassportsPool(limit, ctx.Done()),
+	}
+}
+
+// AddTask adds a task to run
+func (r *limitedConcurrentRunner) AddTask(task func() error) {
+	r.wg.Add(1)
+	go func() {
+		defer func() {
+			r.wg.Done()
+		}()
+
+		// Return false means no passport acquired, and no valid passport will be dispatched any more.
+		// For example, some crucial errors happened and all tasks should be cancelled.
+		if ok := r.passportsPool.Apply(); !ok {
+			return
+		}
+		defer func() {
+			r.passportsPool.Revoke()
+		}()
+
+		err := task()
+		if err != nil {
+			log.Errorf("%v", err)
+			r.cancel()
+		}
+	}()
+}
+
+// Wait waits all the tasks to be finished
+func (r *limitedConcurrentRunner) Wait() {
+	r.wg.Wait()
+}
+
+// Cancel cancels all tasks, tasks that already started will continue to run
+func (r *limitedConcurrentRunner) Cancel() {
+	r.cancel()
+}
+
+// IsCancelled checks whether context is cancelled. This happens when some task encountered critical errors.
+func (r *limitedConcurrentRunner) IsCancelled() bool {
+	return r.ctx.Err() != nil
+}

src/replication/adapter/aliacr/adapter.go

@@ -1,14 +1,12 @@
 package aliacr
 
 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
 	"path/filepath"
 	"regexp"
-	"sync"
 
 	"github.com/aliyun/alibaba-cloud-sdk-go/services/cr"
 	"github.com/goharbor/harbor/src/common/utils"
@@ -159,33 +157,18 @@ func (a *adapter) FetchImages(filters []*model.Filter) (resources []*model.Resou
 	}
 	log.Debugf("FetchImages.repositories: %#v\n", repositories)
 
-	rawResources := make([]*model.Resource, len(repositories))
-	var wg = new(sync.WaitGroup)
-	ctx, cancel := context.WithCancel(context.Background())
-	var passportsPool = utils.NewPassportsPool(adp.MaxConcurrency, ctx.Done())
+	var rawResources = make([]*model.Resource, len(repositories))
+	runner := utils.NewLimitedConcurrentRunner(adp.MaxConcurrency)
+	defer runner.Cancel()
 
 	for i, r := range repositories {
-		wg.Add(1)
-		go func(index int, repo aliRepo) {
-			defer func() {
-				wg.Done()
-			}()
-
-			// Return false means no passport acquired, and no valid passport will be dispatched any more.
-			// For example, some crucial errors happened and all tasks should be cancelled.
-			if ok := passportsPool.Apply(); !ok {
-				return
-			}
-			defer func() {
-				passportsPool.Revoke()
-			}()
-
+		index := i
+		repo := r
+		runner.AddTask(func() error {
 			var tags []string
 			tags, err = a.getTags(repo, client)
 			if err != nil {
-				log.Errorf("List tags for repo '%s' error: %v", repo.RepoName, err)
-				cancel()
-				return
+				return fmt.Errorf("List tags for repo '%s' error: %v", repo.RepoName, err)
 			}
 
 			var filterTags []string
@@ -194,9 +177,7 @@ func (a *adapter) FetchImages(filters []*model.Filter) (resources []*model.Resou
 					var ok bool
 					ok, err = util.Match(tagsPattern, tag)
 					if err != nil {
-						log.Errorf("Match tag '%s' error: %v", tag, err)
-						cancel()
-						return
+						return fmt.Errorf("Match tag '%s' error: %v", tag, err)
 					}
 					if ok {
 						filterTags = append(filterTags, tag)
@@ -218,16 +199,14 @@ func (a *adapter) FetchImages(filters []*model.Filter) (resources []*model.Resou
 						Labels: []string{},
 					},
 				}
-			} else {
-				rawResources[index] = nil
 			}
-		}(i, r)
-	}
-	wg.Wait()
 
-	err = ctx.Err()
-	cancel()
-	if err != nil {
+			return nil
+		})
+	}
+	runner.Wait()
+
+	if runner.IsCancelled() {
 		return nil, fmt.Errorf("FetchImages error when collect tags for repos")
 	}
 

src/replication/adapter/dockerhub/adapter.go

@@ -2,14 +2,12 @@ package dockerhub
 
 import (
 	"bytes"
-	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
 	"strings"
-	"sync"
 
 	"github.com/goharbor/harbor/src/common/utils"
 	"github.com/goharbor/harbor/src/common/utils/log"
@@ -250,26 +248,12 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 	}
 
 	var rawResources = make([]*model.Resource, len(repos))
-	var wg = new(sync.WaitGroup)
-	ctx, cancel := context.WithCancel(context.Background())
-	var passportsPool = utils.NewPassportsPool(adp.MaxConcurrency, ctx.Done())
-
+	runner := utils.NewLimitedConcurrentRunner(adp.MaxConcurrency)
+	defer runner.Cancel()
 	for i, r := range repos {
-		wg.Add(1)
-		go func(index int, repo Repo) {
-			defer func() {
-				wg.Done()
-			}()
-
-			// Return false means no passport acquired, and no valid passport will be dispatched any more.
-			// For example, some crucial errors happened and all tasks should be cancelled.
-			if ok := passportsPool.Apply(); !ok {
-				return
-			}
-			defer func() {
-				passportsPool.Revoke()
-			}()
-
+		index := i
+		repo := r
+		runner.AddTask(func() error {
 			name := fmt.Sprintf("%s/%s", repo.Namespace, repo.Name)
 			log.Infof("Routine started to collect tags for repo: %s", name)
 
@@ -277,12 +261,10 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 			if len(nameFilter) != 0 {
 				m, err := util.Match(nameFilter, name)
 				if err != nil {
-					cancel()
-					log.Errorf("match repo name '%s' against pattern '%s' error: %v", name, nameFilter, err)
-					return
+					return fmt.Errorf("match repo name '%s' against pattern '%s' error: %v", name, nameFilter, err)
 				}
 				if !m {
-					return
+					return nil
 				}
 			}
 
@@ -292,18 +274,14 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 			for {
 				pageTags, err := a.getTags(repo.Namespace, repo.Name, page, pageSize)
 				if err != nil {
-					cancel()
-					log.Errorf("get tags for repo '%s/%s' from DockerHub error: %v", repo.Namespace, repo.Name, err)
-					return
+					return fmt.Errorf("get tags for repo '%s/%s' from DockerHub error: %v", repo.Namespace, repo.Name, err)
 				}
 				for _, t := range pageTags.Tags {
 					// If tag filter set, skip tags that don't match the filter pattern.
 					if len(tagFilter) != 0 {
 						m, err := util.Match(tagFilter, t.Name)
 						if err != nil {
-							cancel()
-							log.Errorf("match tag name '%s' against pattern '%s' error: %v", t.Name, tagFilter, err)
-							return
+							return fmt.Errorf("match tag name '%s' against pattern '%s' error: %v", t.Name, tagFilter, err)
 						}
 
 						if !m {
@@ -319,9 +297,7 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 				page++
 			}
 
-			if len(tags) == 0 {
-				rawResources[index] = nil
-			} else {
+			if len(tags) > 0 {
 				rawResources[index] = &model.Resource{
 					Type:     model.ResourceTypeImage,
 					Registry: a.registry,
@@ -333,13 +309,13 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 					},
 				}
 			}
-		}(i, r)
-	}
-	wg.Wait()
 
-	err = ctx.Err()
-	cancel()
-	if err != nil {
+			return nil
+		})
+	}
+	runner.Wait()
+
+	if runner.IsCancelled() {
 		return nil, fmt.Errorf("FetchImages error when collect tags for repos")
 	}
 

src/replication/adapter/harbor/image_registry.go

@@ -15,10 +15,8 @@
 package harbor
 
 import (
-	"context"
 	"fmt"
 	"strings"
-	"sync"
 
 	"github.com/goharbor/harbor/src/common/utils"
 	"github.com/goharbor/harbor/src/common/utils/log"
@@ -48,47 +46,30 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 			}
 		}
 
-		rawResources := make([]*model.Resource, len(repositories))
-		var wg = new(sync.WaitGroup)
-		ctx, cancel := context.WithCancel(context.Background())
-		var passportsPool = utils.NewPassportsPool(adp.MaxConcurrency, ctx.Done())
+		var rawResources = make([]*model.Resource, len(repositories))
+		runner := utils.NewLimitedConcurrentRunner(adp.MaxConcurrency)
+		defer runner.Cancel()
 
 		for i, r := range repositories {
-			wg.Add(1)
-			go func(index int, repo *adp.Repository) {
-				defer func() {
-					wg.Done()
-				}()
-
-				// Return false means no passport acquired, and no valid passport will be dispatched any more.
-				// For example, some crucial errors happened and all tasks should be cancelled.
-				if ok := passportsPool.Apply(); !ok {
-					return
-				}
-				defer func() {
-					passportsPool.Revoke()
-				}()
-
+			index := i
+			repo := r
+			runner.AddTask(func() error {
 				vTags, err := a.getTags(repo.Name)
 				if err != nil {
-					log.Errorf("List tags for repo '%s' error: %v", repo.Name, err)
-					cancel()
-					return
+					return fmt.Errorf("List tags for repo '%s' error: %v", repo.Name, err)
 				}
 				if len(vTags) == 0 {
 					rawResources[index] = nil
-					return
+					return nil
 				}
 				for _, filter := range filters {
 					if err = filter.DoFilter(&vTags); err != nil {
-						log.Errorf("Filter tags %v error: %v", vTags, err)
-						cancel()
-						return
+						return fmt.Errorf("Filter tags %v error: %v", vTags, err)
 					}
 				}
 				if len(vTags) == 0 {
 					rawResources[index] = nil
-					return
+					return nil
 				}
 				tags := []string{}
 				for _, vTag := range vTags {
@@ -105,13 +86,13 @@ func (a *adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 						Vtags: tags,
 					},
 				}
-			}(i, r)
-		}
-		wg.Wait()
 
-		err = ctx.Err()
-		cancel()
-		if err != nil {
+				return nil
+			})
+		}
+		runner.Wait()
+
+		if runner.IsCancelled() {
 			return nil, fmt.Errorf("FetchImages error when collect tags for repos")
 		}
 

src/replication/adapter/native/adapter.go

@@ -15,7 +15,6 @@
 package native
 
 import (
-	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -162,45 +161,28 @@ func (a *Adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 		}
 	}
 
-	rawResources := make([]*model.Resource, len(repositories))
-	var wg = new(sync.WaitGroup)
-	ctx, cancel := context.WithCancel(context.Background())
-	var passportsPool = utils.NewPassportsPool(adp.MaxConcurrency, ctx.Done())
+	var rawResources = make([]*model.Resource, len(repositories))
+	runner := utils.NewLimitedConcurrentRunner(adp.MaxConcurrency)
+	defer runner.Cancel()
 
 	for i, r := range repositories {
-		wg.Add(1)
-		go func(index int, repo *adp.Repository) {
-			defer func() {
-				wg.Done()
-			}()
-
-			// Return false means no passport acquired, and no valid passport will be dispatched any more.
-			// For example, some crucial errors happened and all tasks should be cancelled.
-			if ok := passportsPool.Apply(); !ok {
-				return
-			}
-			defer func() {
-				passportsPool.Revoke()
-			}()
-
+		index := i
+		repo := r
+		runner.AddTask(func() error {
 			vTags, err := a.getVTags(repo.Name)
 			if err != nil {
-				log.Errorf("List tags for repo '%s' error: %v", repo.Name, err)
-				cancel()
-				return
+				return fmt.Errorf("List tags for repo '%s' error: %v", repo.Name, err)
 			}
 			if len(vTags) == 0 {
-				return
+				return nil
 			}
 			for _, filter := range filters {
 				if err = filter.DoFilter(&vTags); err != nil {
-					log.Errorf("Filter tags %v error: %v", vTags, err)
-					cancel()
-					return
+					return fmt.Errorf("Filter tags %v error: %v", vTags, err)
 				}
 			}
 			if len(vTags) == 0 {
-				return
+				return nil
 			}
 			tags := []string{}
 			for _, vTag := range vTags {
@@ -216,13 +198,13 @@ func (a *Adapter) FetchImages(filters []*model.Filter) ([]*model.Resource, error
 					Vtags: tags,
 				},
 			}
-		}(i, r)
-	}
-	wg.Wait()
 
-	err = ctx.Err()
-	cancel()
-	if err != nil {
+			return nil
+		})
+	}
+	runner.Wait()
+
+	if runner.IsCancelled() {
 		return nil, fmt.Errorf("FetchImages error when collect tags for repos")
 	}