fix: changing default numWorkers, making it customizable and refactoring scheduler (#1563)

Signed-off-by: Lisca Ana-Roberta <ana.kagome@yahoo.com>

pkg/extensions/extension_scrub.go

@@ -57,7 +57,7 @@ type taskGenerator struct {
 	done     bool
 }
 
-func (gen *taskGenerator) GenerateTask() (scheduler.Task, error) {
+func (gen *taskGenerator) Next() (scheduler.Task, error) {
 	repo, err := gen.imgStore.GetNextRepository(gen.lastRepo)
 
 	if err != nil && !errors.Is(err, io.EOF) {

pkg/extensions/extension_search.go

@@ -86,7 +86,7 @@ type TrivyTaskGenerator struct {
 	lock         *sync.Mutex
 }
 
-func (gen *TrivyTaskGenerator) GenerateTask() (scheduler.Task, error) {
+func (gen *TrivyTaskGenerator) Next() (scheduler.Task, error) {
 	var newTask scheduler.Task
 
 	gen.lock.Lock()

pkg/extensions/extension_search_test.go

@@ -5,6 +5,7 @@ package extensions_test
 
 import (
 	"context"
+	"io"
 	"os"
 	"testing"
 	"time"
@@ -12,6 +13,7 @@ import (
 	ispec "github.com/opencontainers/image-spec/specs-go/v1"
 	. "github.com/smartystreets/goconvey/convey"
 
+	"zotregistry.io/zot/pkg/api/config"
 	. "zotregistry.io/zot/pkg/extensions"
 	cveinfo "zotregistry.io/zot/pkg/extensions/search/cve"
 	"zotregistry.io/zot/pkg/log"
@@ -30,8 +32,13 @@ func TestTrivyDBGenerator(t *testing.T) {
 
 		defer os.Remove(logFile.Name()) // clean up
 
-		logger := log.NewLogger("debug", logFile.Name())
-		sch := scheduler.NewScheduler(logger)
+		logger := log.NewLogger("debug", logPath)
+		writers := io.MultiWriter(os.Stdout, logFile)
+		logger.Logger = logger.Output(writers)
+
+		cfg := config.New()
+		cfg.Scheduler = &config.SchedulerConfig{NumWorkers: 3}
+		sch := scheduler.NewScheduler(cfg, logger)
 
 		repoDB := &mocks.RepoDBMock{
 			GetRepoMetaFn: func(repo string) (repodb.RepoMetadata, error) {
@@ -56,13 +63,14 @@ func TestTrivyDBGenerator(t *testing.T) {
 		sch.SubmitGenerator(generator, 12000*time.Millisecond, scheduler.HighPriority)
 
 		ctx, cancel := context.WithCancel(context.Background())
+
 		sch.RunScheduler(ctx)
 
 		defer cancel()
 
 		// Wait for trivy db to download
 		found, err := ReadLogFileAndCountStringOccurence(logPath,
-			"DB update completed, next update scheduled", 90*time.Second, 2)
+			"DB update completed, next update scheduled", 120*time.Second, 2)
 		So(err, ShouldBeNil)
 		So(found, ShouldBeTrue)
 	})

pkg/extensions/sync/sync.go

@@ -88,7 +88,7 @@ func NewTaskGenerator(service Service, log log.Logger) *TaskGenerator {
 	}
 }
 
-func (gen *TaskGenerator) GenerateTask() (scheduler.Task, error) {
+func (gen *TaskGenerator) Next() (scheduler.Task, error) {
 	if err := gen.Service.SetNextAvailableURL(); err != nil {
 		return nil, err
 	}