github/zot
1
0
Fork 0
mirror of https://github.com/project-zot/zot.git synced 2026-06-17 04:48:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Lisca Ana-Roberta <ana.kagome@yahoo.com>
This commit is contained in:
Lisca Ana-Roberta
2023-07-04 11:03:29 +03:00
committed by GitHub
parent 7881ce32b2
commit d4f200c2e1
14 changed files with 101 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/scheduler/scheduler.go
+23 -8
View File
@@ -3,9 +3,11 @@ package scheduler
import (
"container/heap"
"context"
"runtime"
"sync"
"time"
"zotregistry.io/zot/pkg/api/config"
"zotregistry.io/zot/pkg/log"
)
@@ -55,7 +57,7 @@ func (pq *generatorsPriorityQueue) Pop() any {
const (
rateLimiterScheduler = 400
rateLimit = 5 * time.Second
numWorkers = 3
numWorkersMultiplier = 4
)
type Scheduler struct {
@@ -68,13 +70,15 @@ type Scheduler struct {
log log.Logger
stopCh chan struct{}
RateLimit time.Duration
NumWorkers int
}
func NewScheduler(logC log.Logger) *Scheduler {
func NewScheduler(cfg *config.Config, logC log.Logger) *Scheduler {
chLow := make(chan Task, rateLimiterScheduler)
chMedium := make(chan Task, rateLimiterScheduler)
chHigh := make(chan Task, rateLimiterScheduler)
generatorPQ := make(generatorsPriorityQueue, 0)
numWorkers := getNumWorkers(cfg)
sublogger := logC.With().Str("component", "scheduler").Logger()
heap.Init(&generatorPQ)
@@ -88,7 +92,8 @@ func NewScheduler(logC log.Logger) *Scheduler {
log: log.Logger{Logger: sublogger},
stopCh: make(chan struct{}),
// default value
RateLimit: rateLimit,
RateLimit: rateLimit,
NumWorkers: numWorkers,
}
}
@@ -110,6 +115,8 @@ func (scheduler *Scheduler) poolWorker(numWorkers int, tasks chan Task) {
func (scheduler *Scheduler) RunScheduler(ctx context.Context) {
throttle := time.NewTicker(rateLimit).C
numWorkers := scheduler.NumWorkers
tasksWorker := make(chan Task, numWorkers)
// start worker pool
@@ -274,8 +281,8 @@ const (
done
)
type Generator interface {
GenerateTask() (Task, error)
type TaskGenerator interface {
Next() (Task, error)
IsDone() bool
Reset()
}
@@ -285,7 +292,7 @@ type generator struct {
lastRun time.Time
done bool
priority Priority
taskGenerator Generator
taskGenerator TaskGenerator
remainingTask Task
index int
}
@@ -298,7 +305,7 @@ func (gen *generator) generate(sch *Scheduler) {
// in case there is no task already generated, generate a new task
if gen.remainingTask == nil {
nextTask, err := gen.taskGenerator.GenerateTask()
nextTask, err := gen.taskGenerator.Next()
if err != nil {
sch.log.Error().Err(err).Msg("scheduler: error while executing generator")
@@ -347,7 +354,7 @@ func (gen *generator) getState() state {
return ready
}
func (scheduler *Scheduler) SubmitGenerator(taskGenerator Generator, interval time.Duration, priority Priority) {
func (scheduler *Scheduler) SubmitGenerator(taskGenerator TaskGenerator, interval time.Duration, priority Priority) {
newGenerator := &generator{
interval: interval,
done: false,
@@ -362,3 +369,11 @@ func (scheduler *Scheduler) SubmitGenerator(taskGenerator Generator, interval ti
// add generator to the generators priority queue
heap.Push(&scheduler.generators, newGenerator)
}
func getNumWorkers(cfg *config.Config) int {
if cfg.Scheduler != nil && cfg.Scheduler.NumWorkers != 0 {
return cfg.Scheduler.NumWorkers
}
return runtime.NumCPU() * numWorkersMultiplier
}
pkg/scheduler/scheduler_test.go
+27 -8
View File
@@ -5,11 +5,13 @@ import (
"errors"
"fmt"
"os"
"runtime"
"testing"
"time"
. "github.com/smartystreets/goconvey/convey"
"zotregistry.io/zot/pkg/api/config"
"zotregistry.io/zot/pkg/log"
"zotregistry.io/zot/pkg/scheduler"
)
@@ -40,7 +42,7 @@ type generator struct {
step int
}
func (g *generator) GenerateTask() (scheduler.Task, error) {
func (g *generator) Next() (scheduler.Task, error) {
if g.step > 1 {
g.done = true
}
@@ -67,7 +69,7 @@ type shortGenerator struct {
step int
}
func (g *shortGenerator) GenerateTask() (scheduler.Task, error) {
func (g *shortGenerator) Next() (scheduler.Task, error) {
g.done = true
return &task{log: g.log, msg: fmt.Sprintf("executing %s task; index: %d", g.priority, g.index), err: false}, nil
@@ -90,7 +92,7 @@ func TestScheduler(t *testing.T) {
defer os.Remove(logFile.Name()) // clean up
logger := log.NewLogger("debug", logFile.Name())
sch := scheduler.NewScheduler(logger)
sch := scheduler.NewScheduler(config.New(), logger)
genH := &shortGenerator{log: logger, priority: "high priority"}
// interval has to be higher than throttle value to simulate
@@ -114,7 +116,9 @@ func TestScheduler(t *testing.T) {
defer os.Remove(logFile.Name()) // clean up
logger := log.NewLogger("debug", logFile.Name())
sch := scheduler.NewScheduler(logger)
cfg := config.New()
cfg.Scheduler = &config.SchedulerConfig{NumWorkers: 3}
sch := scheduler.NewScheduler(cfg, logger)
genL := &generator{log: logger, priority: "low priority"}
sch.SubmitGenerator(genL, time.Duration(0), scheduler.LowPriority)
@@ -126,6 +130,7 @@ func TestScheduler(t *testing.T) {
sch.SubmitGenerator(genH, time.Duration(0), scheduler.HighPriority)
ctx, cancel := context.WithCancel(context.Background())
sch.RunScheduler(ctx)
time.Sleep(4 * time.Second)
@@ -147,7 +152,7 @@ func TestScheduler(t *testing.T) {
defer os.Remove(logFile.Name()) // clean up
logger := log.NewLogger("debug", logFile.Name())
sch := scheduler.NewScheduler(logger)
sch := scheduler.NewScheduler(config.New(), logger)
t := &task{log: logger, msg: "", err: true}
sch.SubmitTask(t, scheduler.MediumPriority)
@@ -171,7 +176,7 @@ func TestScheduler(t *testing.T) {
defer os.Remove(logFile.Name()) // clean up
logger := log.NewLogger("debug", logFile.Name())
sch := scheduler.NewScheduler(logger)
sch := scheduler.NewScheduler(config.New(), logger)
genL := &generator{log: logger, priority: "low priority"}
sch.SubmitGenerator(genL, 20*time.Millisecond, scheduler.LowPriority)
@@ -195,7 +200,7 @@ func TestScheduler(t *testing.T) {
defer os.Remove(logFile.Name()) // clean up
logger := log.NewLogger("debug", logFile.Name())
sch := scheduler.NewScheduler(logger)
sch := scheduler.NewScheduler(config.New(), logger)
t := &task{log: logger, msg: "", err: false}
sch.SubmitTask(t, -1)
@@ -212,7 +217,7 @@ func TestScheduler(t *testing.T) {
defer os.Remove(logFile.Name()) // clean up
logger := log.NewLogger("debug", logFile.Name())
sch := scheduler.NewScheduler(logger)
sch := scheduler.NewScheduler(config.New(), logger)
ctx, cancel := context.WithCancel(context.Background())
@@ -228,3 +233,17 @@ func TestScheduler(t *testing.T) {
So(string(data), ShouldNotContainSubstring, "scheduler: adding a new task")
})
}
func TestGetNumWorkers(t *testing.T) {
Convey("Test setting the number of workers - default value", t, func() {
sch := scheduler.NewScheduler(config.New(), log.NewLogger("debug", "logFile"))
So(sch.NumWorkers, ShouldEqual, runtime.NumCPU()*4)
})
Convey("Test setting the number of workers - getting the value from config", t, func() {
cfg := config.New()
cfg.Scheduler = &config.SchedulerConfig{NumWorkers: 3}
sch := scheduler.NewScheduler(cfg, log.NewLogger("debug", "logFile"))
So(sch.NumWorkers, ShouldEqual, 3)
})
}