github.com/google/syzkaller@v0.0.0-20240517125934-c0f1611a36d6/pkg/corpus/prio.go (about)

     1  // Copyright 2024 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package corpus
     5  
     6  import (
     7  	"math/rand"
     8  	"sort"
     9  	"sync"
    10  
    11  	"github.com/google/syzkaller/pkg/signal"
    12  	"github.com/google/syzkaller/prog"
    13  )
    14  
    15  type ProgramsList struct {
    16  	mu       sync.RWMutex
    17  	progs    []*prog.Prog
    18  	sumPrios int64
    19  	accPrios []int64
    20  }
    21  
    22  func (pl *ProgramsList) ChooseProgram(r *rand.Rand) *prog.Prog {
    23  	pl.mu.RLock()
    24  	defer pl.mu.RUnlock()
    25  	if len(pl.progs) == 0 {
    26  		return nil
    27  	}
    28  	randVal := r.Int63n(pl.sumPrios + 1)
    29  	idx := sort.Search(len(pl.accPrios), func(i int) bool {
    30  		return pl.accPrios[i] >= randVal
    31  	})
    32  	return pl.progs[idx]
    33  }
    34  
    35  func (pl *ProgramsList) Programs() []*prog.Prog {
    36  	pl.mu.RLock()
    37  	defer pl.mu.RUnlock()
    38  	return pl.progs
    39  }
    40  
    41  func (pl *ProgramsList) saveProgram(p *prog.Prog, signal signal.Signal) {
    42  	pl.mu.Lock()
    43  	defer pl.mu.Unlock()
    44  	prio := int64(len(signal))
    45  	if prio == 0 {
    46  		prio = 1
    47  	}
    48  	pl.sumPrios += prio
    49  	pl.accPrios = append(pl.accPrios, pl.sumPrios)
    50  	pl.progs = append(pl.progs, p)
    51  }
    52  
    53  func (pl *ProgramsList) replace(other *ProgramsList) {
    54  	pl.mu.Lock()
    55  	defer pl.mu.Unlock()
    56  	pl.sumPrios = other.sumPrios
    57  	pl.accPrios = other.accPrios
    58  	pl.progs = other.progs
    59  }