From ad5321c6e31dd6b723935e195830789efdda12c5 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Thu, 16 May 2024 16:45:47 +0200
Subject: [PATCH] pkg/fuzzer/queue: refactor DynamicSource

Use a simpler implementation.
Don't assume the nested Source may be nil.
---
 pkg/fuzzer/queue/queue.go    | 23 +++++++++++------------
 syz-manager/rpc.go           | 19 +++++++++----------
 tools/syz-runtest/runtest.go |  4 ++--
 3 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/pkg/fuzzer/queue/queue.go b/pkg/fuzzer/queue/queue.go
index 68552083c29..0fd87f7c723 100644
--- a/pkg/fuzzer/queue/queue.go
+++ b/pkg/fuzzer/queue/queue.go
@@ -343,24 +343,23 @@ func (pq *PriorityQueue) Next() *Request {
 	return pq.ops.Pop()
 }
 
-type DynamicSource struct {
-	value atomic.Pointer[wrapSource]
+type DynamicSourceCtl struct {
+	value atomic.Pointer[Source]
 }
 
-type wrapSource struct {
-	source Source
+// DynamicSource is assumed never to point to nil.
+func DynamicSource(source Source) *DynamicSourceCtl {
+	var ret DynamicSourceCtl
+	ret.Store(source)
+	return &ret
 }
 
-func (ds *DynamicSource) Store(source Source) {
-	ds.value.Store(&wrapSource{source})
+func (ds *DynamicSourceCtl) Store(source Source) {
+	ds.value.Store(&source)
 }
 
-func (ds *DynamicSource) Next() *Request {
-	val := ds.value.Load()
-	if val == nil || val.source == nil {
-		return nil
-	}
-	return val.source.Next()
+func (ds *DynamicSourceCtl) Next() *Request {
+	return (*ds.value.Load()).Next()
 }
 
 // Deduplicator() keeps track of the previously run requests to avoid re-running them.
diff --git a/syz-manager/rpc.go b/syz-manager/rpc.go
index b54466b463e..81c8b0d9819 100644
--- a/syz-manager/rpc.go
+++ b/syz-manager/rpc.go
@@ -34,12 +34,10 @@ type RPCServer struct {
 	checker *vminfo.Checker
 	port    int
 
-	infoDone      bool
-	checkDone     atomic.Bool
-	checkFailures int
-
-	checkerSource    *queue.DynamicSource
-	baseSource       *queue.DynamicSource
+	infoDone         bool
+	checkDone        atomic.Bool
+	checkFailures    int
+	baseSource       *queue.DynamicSourceCtl
 	enabledFeatures  flatrpc.Feature
 	setupFeatures    flatrpc.Feature
 	modules          []cover.KernelModule
@@ -100,14 +98,15 @@ type RPCManagerView interface {
 }
 
 func startRPCServer(mgr *Manager) (*RPCServer, error) {
-	var baseSource queue.DynamicSource
+	checker := vminfo.New(mgr.cfg)
+	baseSource := queue.DynamicSource(checker)
 	serv := &RPCServer{
 		mgr:        mgr,
 		cfg:        mgr.cfg,
 		target:     mgr.target,
-		checker:    vminfo.New(mgr.cfg),
-		baseSource: &baseSource,
-		execSource: queue.Retry(&baseSource),
+		checker:    checker,
+		baseSource: baseSource,
+		execSource: queue.Retry(baseSource),
 		statExecs:  mgr.statExecs,
 		statExecRetries: stats.Create("exec retries",
 			"Number of times a test program was restarted because the first run failed",
diff --git a/tools/syz-runtest/runtest.go b/tools/syz-runtest/runtest.go
index 878ccad3add..b1caef092d0 100644
--- a/tools/syz-runtest/runtest.go
+++ b/tools/syz-runtest/runtest.go
@@ -66,6 +66,7 @@ func main() {
 		pending:      make(map[string]map[int64]bool),
 	}
 	mgr.checkFiles = mgr.checker.RequiredFiles()
+	mgr.source = queue.DynamicSource(mgr.checker)
 	s, err := rpctype.NewRPCServer(cfg.RPC, "Manager", mgr)
 	if err != nil {
 		log.Fatalf("failed to create rpc server: %v", err)
@@ -96,7 +97,6 @@ func main() {
 		}()
 	}
 	checkResult := <-mgr.checkResultC
-	mgr.source.Store(mgr.checker)
 	calls, _, features, err := mgr.checker.Run(checkResult.Files, checkResult.Features)
 	if err != nil {
 		log.Fatalf("failed to detect enabled syscalls: %v", err)
@@ -146,7 +146,7 @@ type Manager struct {
 	vmStop       chan bool
 	port         int
 	debug        bool
-	source       queue.DynamicSource
+	source       *queue.DynamicSourceCtl
 
 	reqMu   sync.Mutex
 	reqSeq  int64