4 files changed, 90 insertions, 0 deletions
diff --git a/pkg/fuzzer/job.go b/pkg/fuzzer/job.go
index 9fd294ca0..d54e6bd48 100644
--- a/pkg/fuzzer/job.go
+++ b/pkg/fuzzer/job.go
@@ -70,6 +70,7 @@ func candidateRequest(fuzzer *Fuzzer, input Candidate) (*queue.Request, ProgType
 		Prog:       input.Prog,
 		NeedSignal: queue.NewSignal,
 		Stat:       fuzzer.statExecCandidate,
+		Important:  true,
 	}, flags
 }
 
@@ -88,6 +89,7 @@ type triageJob struct {
 }
 
 func (job *triageJob) execute(req *queue.Request, opts ...execOpt) *queue.Result {
+	req.Important = true // All triage executions are important.
 	return job.fuzzer.execute(job.queue, req, opts...)
 }
 
diff --git a/pkg/fuzzer/queue/queue.go b/pkg/fuzzer/queue/queue.go
index 14d4ace78..48b912846 100644
--- a/pkg/fuzzer/queue/queue.go
+++ b/pkg/fuzzer/queue/queue.go
@@ -40,11 +40,16 @@ type Request struct {
 	BinaryFile string // If set, it's executed instead of Prog.
 	Repeat     int    // Repeats in addition to the first run.
 
+	// Important requests will be retried even from crashed VMs.
+	Important bool
+
 	// The callback will be called on request completion in the LIFO order.
 	// If it returns false, all further processing will be stopped.
 	// It allows wrappers to intercept Done() requests.
 	callback DoneCallback
 
+	onceCrashed bool
+
 	mu     sync.Mutex
 	result *Result
 	done   chan struct{}
@@ -91,6 +96,11 @@ func (r *Request) Wait(ctx context.Context) *Result {
 	}
 }
 
+// Risky() returns true if there's a substantial risk of the input crashing the VM.
+func (r *Request) Risky() bool {
+	return r.onceCrashed
+}
+
 func (r *Request) hash() hash.Sig {
 	buf := new(bytes.Buffer)
 	if r.ExecOpts != nil {
diff --git a/pkg/fuzzer/queue/retry.go b/pkg/fuzzer/queue/retry.go
index 0b2e02ba5..c59a2c048 100644
--- a/pkg/fuzzer/queue/retry.go
+++ b/pkg/fuzzer/queue/retry.go
@@ -33,5 +33,11 @@ func (r *retryer) done(req *Request, res *Result) bool {
 		r.pq.Submit(req)
 		return false
 	}
+	// Retry important requests from crashed VMs once.
+	if res.Status == Crashed && req.Important && !req.onceCrashed {
+		req.onceCrashed = true
+		r.pq.Submit(req)
+		return false
+	}
 	return true
 }
diff --git a/pkg/fuzzer/queue/retry_test.go b/pkg/fuzzer/queue/retry_test.go
new file mode 100644
index 000000000..8529779a2
--- /dev/null
+++ b/pkg/fuzzer/queue/retry_test.go
@@ -0,0 +1,72 @@
+// Copyright 2024 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package queue
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRetryerOnRestart(t *testing.T) {
+	q := Plain()
+	retryerObj := Retry(q)
+
+	q.Submit(&Request{Important: true})
+	q.Submit(&Request{Important: false})
+
+	// The requests must be retried forever.
+	req1 := retryerObj.Next()
+	req2 := retryerObj.Next()
+	for i := 0; i < 10; i++ {
+		req1.Done(&Result{Status: Restarted})
+		req2.Done(&Result{Status: Restarted})
+		assert.Equal(t, req1, retryerObj.Next())
+		assert.Equal(t, req2, retryerObj.Next())
+	}
+
+	// Once successful, requests should no longer appear.
+	req1.Done(&Result{Status: Success})
+	req2.Done(&Result{Status: Success})
+
+	assert.Equal(t, Success, req1.Wait(context.Background()).Status)
+	assert.Equal(t, Success, req2.Wait(context.Background()).Status)
+
+	assert.Nil(t, retryerObj.Next())
+	assert.Nil(t, retryerObj.Next())
+}
+
+func TestRetryerOnCrash(t *testing.T) {
+	q := Plain()
+	retryerObj := Retry(q)
+
+	// Unimportant requests will not be retried.
+	req := &Request{Important: false}
+	q.Submit(req)
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Nil(t, retryerObj.Next())
+	assert.Equal(t, Crashed, req.Wait(context.Background()).Status)
+
+	// Important requests will be retried once.
+	req = &Request{Important: true}
+	q.Submit(req)
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Success})
+	assert.Nil(t, retryerObj.Next())
+	assert.Equal(t, Success, req.Wait(context.Background()).Status)
+
+	// .. but not more than once.
+	req = &Request{Important: true}
+	q.Submit(req)
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Equal(t, req, retryerObj.Next())
+	req.Done(&Result{Status: Crashed})
+	assert.Nil(t, retryerObj.Next())
+	assert.Equal(t, Crashed, req.Wait(context.Background()).Status)
+}