zotregistry.dev/zot@v1.4.4-0.20240314164342-eec277e14d20/pkg/extensions/search/cve/update.go (about)

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