zot/pkg/extensions/search/cve/update.go

package cveinfo

import (
	"context"
	"sync"
	"time"

	"zotregistry.io/zot/pkg/log"
	"zotregistry.io/zot/pkg/scheduler"
)

type state int

const (
	pending state = iota
	running
	done
)

func NewDBUpdateTaskGenerator(
	interval time.Duration,
	scanner Scanner,
	log log.Logger,
) scheduler.TaskGenerator {
	generator := &DBUpdateTaskGenerator{
		interval,
		scanner,
		log,
		pending,
		0,
		time.Now(),
		&sync.Mutex{},
	}

	return generator
}

type DBUpdateTaskGenerator struct {
	interval     time.Duration
	scanner      Scanner
	log          log.Logger
	status       state
	waitTime     time.Duration
	lastTaskTime time.Time
	lock         *sync.Mutex
}

func (gen *DBUpdateTaskGenerator) Next() (scheduler.Task, error) {
	var newTask scheduler.Task

	gen.lock.Lock()

	if gen.status == pending && time.Since(gen.lastTaskTime) >= gen.waitTime {
		newTask = newDBUpdadeTask(gen.interval, gen.scanner, gen, gen.log)
		gen.status = running
	}
	gen.lock.Unlock()

	return newTask, nil
}

func (gen *DBUpdateTaskGenerator) IsDone() bool {
	gen.lock.Lock()
	status := gen.status
	gen.lock.Unlock()

	return status == done
}

func (gen *DBUpdateTaskGenerator) IsReady() bool {
	return true
}

func (gen *DBUpdateTaskGenerator) Reset() {
	gen.lock.Lock()
	gen.status = pending
	gen.waitTime = 0
	gen.lock.Unlock()
}

type dbUpdateTask struct {
	interval  time.Duration
	scanner   Scanner
	generator *DBUpdateTaskGenerator
	log       log.Logger
}

func newDBUpdadeTask(interval time.Duration, scanner Scanner,
	generator *DBUpdateTaskGenerator, log log.Logger,
) *dbUpdateTask {
	return &dbUpdateTask{interval, scanner, generator, log}
}

func (dbt *dbUpdateTask) DoWork(ctx context.Context) error {
	dbt.log.Info().Msg("updating the CVE database")

	err := dbt.scanner.UpdateDB()
	if err != nil {
		dbt.generator.lock.Lock()
		dbt.generator.status = pending

		if dbt.generator.waitTime == 0 {
			dbt.generator.waitTime = time.Second
		}

		dbt.generator.waitTime *= 2
		dbt.generator.lastTaskTime = time.Now()
		dbt.generator.lock.Unlock()

		return err
	}

	dbt.generator.lock.Lock()
	dbt.generator.lastTaskTime = time.Now()
	dbt.generator.status = done
	dbt.generator.lock.Unlock()
	dbt.log.Info().Str("DB update completed, next update scheduled after", dbt.interval.String()).Msg("")

	return nil
}