github.com/google/syzkaller@v0.0.0-20240517125934-c0f1611a36d6/syz-verifier/exectask.go (about)

     1  // Copyright 2021 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 main
     5  
     6  import (
     7  	"container/heap"
     8  	"sync"
     9  	"sync/atomic"
    10  	"time"
    11  
    12  	"github.com/google/syzkaller/pkg/rpctype"
    13  	"github.com/google/syzkaller/prog"
    14  )
    15  
    16  type EnvDescr int64
    17  
    18  const (
    19  	AnyEnvironment EnvDescr = iota
    20  	NewEnvironment
    21  	// TODO: add CleanVMEnvironment support.
    22  
    23  	EnvironmentsCount
    24  )
    25  
    26  // ExecTask is the atomic analysis entity. Once executed, it could trigger the
    27  // pipeline propagation for the program.
    28  type ExecTask struct {
    29  	CreationTime   time.Time
    30  	Program        *prog.Prog
    31  	ID             int64
    32  	ExecResultChan ExecResultChan
    33  
    34  	priority int // The priority of the item in the queue.
    35  	// The index is needed by update and is maintained by the heap.Interface methods.
    36  	index int // The index of the item in the heap.
    37  }
    38  
    39  func (t *ExecTask) ToRPC() *rpctype.ExecTask {
    40  	return &rpctype.ExecTask{
    41  		Prog: t.Program.Serialize(),
    42  		ID:   t.ID,
    43  	}
    44  }
    45  
    46  type ExecTaskFactory struct {
    47  	chanMapMutex           sync.Mutex
    48  	taskIDToExecResultChan map[int64]ExecResultChan
    49  	taskCounter            int64
    50  }
    51  
    52  func MakeExecTaskFactory() *ExecTaskFactory {
    53  	return &ExecTaskFactory{
    54  		taskIDToExecResultChan: make(map[int64]ExecResultChan),
    55  		taskCounter:            -1,
    56  	}
    57  }
    58  
    59  type ExecResultChan chan *ExecResult
    60  
    61  func (factory *ExecTaskFactory) MakeExecTask(prog *prog.Prog) *ExecTask {
    62  	task := &ExecTask{
    63  		CreationTime:   time.Now(),
    64  		Program:        prog,
    65  		ExecResultChan: make(ExecResultChan),
    66  		ID:             atomic.AddInt64(&factory.taskCounter, 1),
    67  	}
    68  
    69  	factory.chanMapMutex.Lock()
    70  	defer factory.chanMapMutex.Unlock()
    71  	factory.taskIDToExecResultChan[task.ID] = task.ExecResultChan
    72  
    73  	return task
    74  }
    75  
    76  func (factory *ExecTaskFactory) ExecTasksQueued() int {
    77  	factory.chanMapMutex.Lock()
    78  	defer factory.chanMapMutex.Unlock()
    79  	return len(factory.taskIDToExecResultChan)
    80  }
    81  
    82  func (factory *ExecTaskFactory) DeleteExecTask(task *ExecTask) {
    83  	factory.chanMapMutex.Lock()
    84  	defer factory.chanMapMutex.Unlock()
    85  	delete(factory.taskIDToExecResultChan, task.ID)
    86  }
    87  
    88  func (factory *ExecTaskFactory) GetExecResultChan(taskID int64) ExecResultChan {
    89  	factory.chanMapMutex.Lock()
    90  	defer factory.chanMapMutex.Unlock()
    91  
    92  	return factory.taskIDToExecResultChan[taskID]
    93  }
    94  
    95  func MakeExecTaskQueue() *ExecTaskQueue {
    96  	return &ExecTaskQueue{
    97  		pq: make(ExecTaskPriorityQueue, 0),
    98  	}
    99  }
   100  
   101  // ExecTaskQueue respects the pq.priority. Internally it is a thread-safe PQ.
   102  type ExecTaskQueue struct {
   103  	pq ExecTaskPriorityQueue
   104  	mu sync.Mutex
   105  }
   106  
   107  // PopTask return false if no tasks are available.
   108  func (q *ExecTaskQueue) PopTask() (*ExecTask, bool) {
   109  	q.mu.Lock()
   110  	defer q.mu.Unlock()
   111  	if q.pq.Len() == 0 {
   112  		return nil, false
   113  	}
   114  	return heap.Pop(&q.pq).(*ExecTask), true
   115  }
   116  
   117  func (q *ExecTaskQueue) PushTask(task *ExecTask) {
   118  	q.mu.Lock()
   119  	defer q.mu.Unlock()
   120  	heap.Push(&q.pq, task)
   121  }
   122  
   123  func (q *ExecTaskQueue) Len() int {
   124  	q.mu.Lock()
   125  	defer q.mu.Unlock()
   126  	return q.pq.Len()
   127  }
   128  
   129  // ExecTaskPriorityQueue reused example from https://pkg.go.dev/container/heap
   130  type ExecTaskPriorityQueue []*ExecTask
   131  
   132  func (pq ExecTaskPriorityQueue) Len() int { return len(pq) }
   133  
   134  func (pq ExecTaskPriorityQueue) Less(i, j int) bool {
   135  	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
   136  	return pq[i].priority > pq[j].priority
   137  }
   138  
   139  func (pq ExecTaskPriorityQueue) Swap(i, j int) {
   140  	pq[i], pq[j] = pq[j], pq[i]
   141  	pq[i].index = i
   142  	pq[j].index = j
   143  }
   144  
   145  func (pq *ExecTaskPriorityQueue) Push(x interface{}) {
   146  	n := len(*pq)
   147  	item := x.(*ExecTask)
   148  	item.index = n
   149  	*pq = append(*pq, item)
   150  }
   151  
   152  func (pq *ExecTaskPriorityQueue) Pop() interface{} {
   153  	old := *pq
   154  	n := len(old)
   155  	item := old[n-1]
   156  	old[n-1] = nil  // avoid memory leak
   157  	item.index = -1 // for safety
   158  	*pq = old[0 : n-1]
   159  	return item
   160  }