zot/pkg/scheduler/scheduler_test.go

package scheduler_test

import (
	"context"
	"errors"
	"fmt"
	"os"
	"runtime"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"

	"zotregistry.io/zot/pkg/api/config"
	"zotregistry.io/zot/pkg/extensions/monitoring"
	"zotregistry.io/zot/pkg/log"
	"zotregistry.io/zot/pkg/scheduler"
)

type task struct {
	log log.Logger
	msg string
	err bool
}

var errInternal = errors.New("task: internal error")

func (t *task) DoWork(ctx context.Context) error {
	if t.err {
		return errInternal
	}

	for idx := 0; idx < 5; idx++ {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		time.Sleep(100 * time.Millisecond)
	}

	t.log.Info().Msg(t.msg)

	return nil
}

func (t *task) String() string {
	return t.Name()
}

func (t *task) Name() string {
	return "TestTask"
}

type generator struct {
	log      log.Logger
	priority string
	done     bool
	index    int
	step     int
}

func (g *generator) Next() (scheduler.Task, error) {
	if g.step > 100 {
		g.done = true
	}
	g.step++
	g.index++

	if g.step%11 == 0 {
		return nil, nil
	}

	if g.step%13 == 0 {
		return nil, errInternal
	}

	return &task{log: g.log, msg: fmt.Sprintf("executing %s task; index: %d", g.priority, g.index), err: false}, nil
}

func (g *generator) IsDone() bool {
	return g.done
}

func (g *generator) IsReady() bool {
	return true
}

func (g *generator) Reset() {
	g.done = false
	g.step = 0
}

type shortGenerator struct {
	log      log.Logger
	priority string
	done     bool
	index    int
	step     int
}

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
}

func (g *shortGenerator) IsDone() bool {
	return g.done
}

func (g *shortGenerator) IsReady() bool {
	return true
}

func (g *shortGenerator) Reset() {
	g.done = true
	g.step = 0
}

func TestScheduler(t *testing.T) {
	Convey("Test active to waiting periodic generator", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)

		genH := &shortGenerator{log: logger, priority: "high priority"}
		// interval has to be higher than throttle value to simulate
		sch.SubmitGenerator(genH, 1*time.Second, scheduler.HighPriority)

		sch.RunScheduler()
		time.Sleep(2 * time.Second)
		sch.Shutdown()

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)
		So(string(data), ShouldContainSubstring, "waiting generator is ready, pushing to ready generators")
	})

	Convey("Test order of generators in queue", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		cfg := config.New()
		cfg.Scheduler = &config.SchedulerConfig{NumWorkers: 3}
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(cfg, metrics, logger)
		sch.RateLimit = 5 * time.Second

		genL := &generator{log: logger, priority: "low priority"}
		sch.SubmitGenerator(genL, time.Duration(0), scheduler.LowPriority)

		genM := &generator{log: logger, priority: "medium priority"}
		sch.SubmitGenerator(genM, time.Duration(0), scheduler.MediumPriority)

		genH := &generator{log: logger, priority: "high priority"}
		sch.SubmitGenerator(genH, time.Duration(0), scheduler.HighPriority)

		sch.RunScheduler()
		time.Sleep(4 * time.Second)
		sch.Shutdown()

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)

		So(string(data), ShouldContainSubstring, "executing high priority task; index: 1")
		So(string(data), ShouldContainSubstring, "executing high priority task; index: 2")
		So(string(data), ShouldNotContainSubstring, "executing medium priority task; index: 1")
		So(string(data), ShouldNotContainSubstring, "failed to execute task")
	})

	Convey("Test task returning an error", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)

		t := &task{log: logger, msg: "", err: true}
		sch.SubmitTask(t, scheduler.MediumPriority)

		sch.RunScheduler()
		time.Sleep(500 * time.Millisecond)
		sch.Shutdown()

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)
		So(string(data), ShouldContainSubstring, "adding a new task")
		So(string(data), ShouldContainSubstring, "failed to execute task")
	})

	Convey("Test resubmit generator", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)

		genL := &generator{log: logger, priority: "low priority"}
		sch.SubmitGenerator(genL, 20*time.Millisecond, scheduler.LowPriority)

		sch.RunScheduler()
		time.Sleep(1 * time.Second)
		sch.Shutdown()

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)
		So(string(data), ShouldContainSubstring, "executing low priority task; index: 1")
		So(string(data), ShouldContainSubstring, "executing low priority task; index: 2")
	})

	Convey("Try to add a task with wrong priority", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)

		t := &task{log: logger, msg: "", err: false}
		sch.SubmitTask(t, -1)

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)
		So(string(data), ShouldNotContainSubstring, "adding a new task")
	})

	Convey("Test adding a new task when context is done", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)

		sch.RunScheduler()
		sch.Shutdown()
		time.Sleep(500 * time.Millisecond)

		t := &task{log: logger, msg: "", err: false}
		sch.SubmitTask(t, scheduler.LowPriority)

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)
		So(string(data), ShouldNotContainSubstring, "adding a new task")
	})

	Convey("Test stopping scheduler by calling Shutdown()", t, func() {
		logFile, err := os.CreateTemp("", "zot-log*.txt")
		So(err, ShouldBeNil)

		defer os.Remove(logFile.Name()) // clean up

		logger := log.NewLogger("debug", logFile.Name())
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)

		genL := &generator{log: logger, priority: "medium priority"}
		sch.SubmitGenerator(genL, 20*time.Millisecond, scheduler.MediumPriority)

		sch.RunScheduler()
		time.Sleep(1 * time.Second)
		sch.Shutdown()

		data, err := os.ReadFile(logFile.Name())
		So(err, ShouldBeNil)
		So(string(data), ShouldContainSubstring, "executing medium priority task; index: 1")
		So(string(data), ShouldContainSubstring, "executing medium priority task; index: 2")
		So(string(data), ShouldContainSubstring, "received stop signal, gracefully shutting down...")
	})

	Convey("Test scheduler Priority.String() method", t, func() {
		var p scheduler.Priority //nolint: varnamelen
		// test invalid priority
		p = 6238734
		So(p.String(), ShouldEqual, "invalid")
		p = scheduler.LowPriority
		So(p.String(), ShouldEqual, "low")
		p = scheduler.MediumPriority
		So(p.String(), ShouldEqual, "medium")
		p = scheduler.HighPriority
		So(p.String(), ShouldEqual, "high")
	})

	Convey("Test scheduler State.String() method", t, func() {
		var s scheduler.State //nolint: varnamelen
		// test invalid state
		s = -67
		So(s.String(), ShouldEqual, "invalid")
		s = scheduler.Ready
		So(s.String(), ShouldEqual, "ready")
		s = scheduler.Waiting
		So(s.String(), ShouldEqual, "waiting")
		s = scheduler.Done
		So(s.String(), ShouldEqual, "done")
	})
}

func TestGetNumWorkers(t *testing.T) {
	Convey("Test setting the number of workers - default value", t, func() {
		logger := log.NewLogger("debug", "logFile")
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(config.New(), metrics, logger)
		defer os.Remove("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}
		logger := log.NewLogger("debug", "logFile")
		metrics := monitoring.NewMetricsServer(true, logger)
		sch := scheduler.NewScheduler(cfg, metrics, logger)
		defer os.Remove("logFile")
		So(sch.NumWorkers, ShouldEqual, 3)
	})
}