zotregistry.io/zot@v1.4.4-0.20231124084042-02a8ed785457/pkg/extensions/search/cve/update.go (about)

     1  package cveinfo
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"zotregistry.io/zot/pkg/log"
     9  	"zotregistry.io/zot/pkg/scheduler"
    10  )
    11  
    12  type state int
    13  
    14  const (
    15  	pending state = iota
    16  	running
    17  	done
    18  )
    19  
    20  func NewDBUpdateTaskGenerator(
    21  	interval time.Duration,
    22  	scanner Scanner,
    23  	log log.Logger,
    24  ) scheduler.TaskGenerator {
    25  	generator := &DBUpdateTaskGenerator{
    26  		interval,
    27  		scanner,
    28  		log,
    29  		pending,
    30  		0,
    31  		time.Now(),
    32  		&sync.Mutex{},
    33  	}
    34  
    35  	return generator
    36  }
    37  
    38  type DBUpdateTaskGenerator struct {
    39  	interval     time.Duration
    40  	scanner      Scanner
    41  	log          log.Logger
    42  	status       state
    43  	waitTime     time.Duration
    44  	lastTaskTime time.Time
    45  	lock         *sync.Mutex
    46  }
    47  
    48  func (gen *DBUpdateTaskGenerator) Next() (scheduler.Task, error) {
    49  	var newTask scheduler.Task
    50  
    51  	gen.lock.Lock()
    52  
    53  	if gen.status == pending && time.Since(gen.lastTaskTime) >= gen.waitTime {
    54  		newTask = newDBUpdadeTask(gen.interval, gen.scanner, gen, gen.log)
    55  		gen.status = running
    56  	}
    57  	gen.lock.Unlock()
    58  
    59  	return newTask, nil
    60  }
    61  
    62  func (gen *DBUpdateTaskGenerator) IsDone() bool {
    63  	gen.lock.Lock()
    64  	status := gen.status
    65  	gen.lock.Unlock()
    66  
    67  	return status == done
    68  }
    69  
    70  func (gen *DBUpdateTaskGenerator) IsReady() bool {
    71  	return true
    72  }
    73  
    74  func (gen *DBUpdateTaskGenerator) Reset() {
    75  	gen.lock.Lock()
    76  	gen.status = pending
    77  	gen.waitTime = 0
    78  	gen.lock.Unlock()
    79  }
    80  
    81  type dbUpdateTask struct {
    82  	interval  time.Duration
    83  	scanner   Scanner
    84  	generator *DBUpdateTaskGenerator
    85  	log       log.Logger
    86  }
    87  
    88  func newDBUpdadeTask(interval time.Duration, scanner Scanner,
    89  	generator *DBUpdateTaskGenerator, log log.Logger,
    90  ) *dbUpdateTask {
    91  	return &dbUpdateTask{interval, scanner, generator, log}
    92  }
    93  
    94  func (dbt *dbUpdateTask) DoWork(ctx context.Context) error {
    95  	dbt.log.Info().Msg("updating the CVE database")
    96  
    97  	err := dbt.scanner.UpdateDB(ctx)
    98  	if err != nil {
    99  		dbt.generator.lock.Lock()
   100  		dbt.generator.status = pending
   101  
   102  		if dbt.generator.waitTime == 0 {
   103  			dbt.generator.waitTime = time.Second
   104  		}
   105  
   106  		dbt.generator.waitTime *= 2
   107  		dbt.generator.lastTaskTime = time.Now()
   108  		dbt.generator.lock.Unlock()
   109  
   110  		return err
   111  	}
   112  
   113  	dbt.generator.lock.Lock()
   114  	dbt.generator.lastTaskTime = time.Now()
   115  	dbt.generator.status = done
   116  	dbt.generator.lock.Unlock()
   117  	dbt.log.Info().Str("DB update completed, next update scheduled after", dbt.interval.String()).Msg("")
   118  
   119  	return nil
   120  }