From 1abc95d4d5e43f23d66a7b5d64c96a9d4c00e8d3 Mon Sep 17 00:00:00 2001 From: Adam Hamrick Date: Wed, 17 Jun 2026 19:39:12 -0400 Subject: [PATCH 1/3] add graceful cancellation mid-diagnose run --- internal/cmd/diagnose.go | 10 +- internal/cmd/root.go | 9 +- internal/runner/diagnose_cancel.go | 130 ++++++++++++ internal/runner/diagnose_cancel_test.go | 63 ++++++ internal/runner/diagnose_output.go | 28 +++ internal/runner/runner.go | 79 +++++++- internal/runner/runner_test.go | 259 ++++++++++++++++++++++++ 7 files changed, 561 insertions(+), 17 deletions(-) create mode 100644 internal/runner/diagnose_cancel.go create mode 100644 internal/runner/diagnose_cancel_test.go diff --git a/internal/cmd/diagnose.go b/internal/cmd/diagnose.go index 71484c9..90ae2e2 100644 --- a/internal/cmd/diagnose.go +++ b/internal/cmd/diagnose.go @@ -1,6 +1,7 @@ package cmd import ( + "context" "time" "github.com/spf13/cobra" @@ -54,8 +55,11 @@ Do not pass go test -trace; use diagnose --trace instead. With --shuffle-seed, a iterSetup := hooks.BuildIterationHook(runnerOpts, shell, hooks.PhaseSetup) iterTeardown := hooks.BuildIterationHook(runnerOpts, shell, hooks.PhaseTeardown) + runCtx, stopRun := runner.NewDiagnoseRunContext(context.WithoutCancel(cmd.Context())) + defer stopRun() + state, start, resultsDir, runErr := runner.RunIterations( - cmd.Context(), + runCtx, conf, out, args, @@ -66,12 +70,12 @@ Do not pass go test -trace; use diagnose --trace instead. With --shuffle-seed, a finishResourceCleanup(cmd, &runErr, cleanup) if runErr != nil { - if cmd.Context().Err() == nil { + if runCtx.Err() == nil { return runErr } } - return runner.FinishDiagnoseAnalysis(cmd.Context(), conf, out, args, state, start, resultsDir) + return runner.FinishDiagnoseAnalysis(runCtx, conf, out, args, state, start, resultsDir) }), } diff --git a/internal/cmd/root.go b/internal/cmd/root.go index e377481..e591b43 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -66,10 +66,11 @@ func NewRootCommand(runnerOpts hooks.RunOptions) *cobra.Command { return rootCmd } -// Execute runs the root command. A SIGINT or SIGTERM cancels the context so -// long-running subcommands (notably `diagnose`) can stop cleanly and still write -// their post-run analysis. A second signal hits the default handler and -// force-exits. +// Execute runs the root command. A SIGINT or SIGTERM cancels cmd.Context() on +// the first signal so generic subcommands (go test, gotestsum) stop promptly. +// diagnose installs its own two-stage context: first signal finishes in-flight +// iterations; second signal hard-cancels. A second root signal hits the default +// handler and force-exits when diagnose is not running. func Execute(opts ...hooks.Option) { if err := runExecute(opts...); err != nil { os.Exit(1) diff --git a/internal/runner/diagnose_cancel.go b/internal/runner/diagnose_cancel.go new file mode 100644 index 0000000..1c29704 --- /dev/null +++ b/internal/runner/diagnose_cancel.go @@ -0,0 +1,130 @@ +package runner + +import ( + "context" + "os" + "os/signal" + "sync" + "sync/atomic" + "syscall" +) + +type diagnoseCancelKey struct{} + +type diagnoseCancelState struct { + softStop atomic.Bool + hardCancel context.CancelFunc + softStopCh chan struct{} + softOnce sync.Once +} + +// NewDiagnoseRunContext returns a context for diagnose iteration runs with +// two-stage cancellation. The first SIGINT/SIGTERM requests a graceful stop +// (finish in-flight iterations, do not enqueue new ones). The second signal +// hard-cancels the context. Pass context.WithoutCancel(cmd.Context()) as parent +// so the root CLI signal handler does not cancel iteration execution on the +// first press. +func NewDiagnoseRunContext(parent context.Context) (context.Context, func()) { + return newDiagnoseRunContext(parent, true) +} + +// NewDiagnoseRunContextForTest is like NewDiagnoseRunContext but without OS +// signal handling. Use in unit and synctest tests with RequestDiagnoseGracefulStop +// and RequestDiagnoseHardCancel. +func NewDiagnoseRunContextForTest(parent context.Context) (context.Context, func()) { + return newDiagnoseRunContext(parent, false) +} + +func newDiagnoseRunContext(parent context.Context, listenSignals bool) (context.Context, func()) { + ctx, hardCancel := context.WithCancel(parent) + state := &diagnoseCancelState{ + hardCancel: hardCancel, + softStopCh: make(chan struct{}), + } + ctx = context.WithValue(ctx, diagnoseCancelKey{}, state) + + stop := func() { + hardCancel() + } + + if !listenSignals { + return ctx, stop + } + + sigCh := make(chan os.Signal, 2) + signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM) + + stop = func() { + signal.Stop(sigCh) + hardCancel() + } + + go func() { + defer signal.Stop(sigCh) + for { + select { + case <-ctx.Done(): + return + case <-sigCh: + state.requestGracefulStop() + select { + case <-sigCh: + hardCancel() + return + case <-ctx.Done(): + return + } + } + } + }() + + return ctx, stop +} + +func (s *diagnoseCancelState) requestGracefulStop() { + s.softOnce.Do(func() { + s.softStop.Store(true) + close(s.softStopCh) + }) +} + +func diagnoseCancelFromContext(ctx context.Context) *diagnoseCancelState { + if ctx == nil { + return nil + } + state, _ := ctx.Value(diagnoseCancelKey{}).(*diagnoseCancelState) + return state +} + +// DiagnoseGracefulStopChan returns a channel closed on the first graceful stop +// request, or nil when ctx is not from NewDiagnoseRunContext. +func DiagnoseGracefulStopChan(ctx context.Context) <-chan struct{} { + state := diagnoseCancelFromContext(ctx) + if state == nil { + return nil + } + return state.softStopCh +} + +// DiagnoseGracefulStopRequested reports whether the user requested a graceful +// stop (first Ctrl+C) on a context from NewDiagnoseRunContext. +func DiagnoseGracefulStopRequested(ctx context.Context) bool { + state := diagnoseCancelFromContext(ctx) + return state != nil && state.softStop.Load() +} + +// RequestDiagnoseGracefulStop simulates the first interrupt for tests. +func RequestDiagnoseGracefulStop(ctx context.Context) { + state := diagnoseCancelFromContext(ctx) + if state != nil { + state.requestGracefulStop() + } +} + +// RequestDiagnoseHardCancel simulates the second interrupt for tests. +func RequestDiagnoseHardCancel(ctx context.Context) { + state := diagnoseCancelFromContext(ctx) + if state != nil { + state.hardCancel() + } +} diff --git a/internal/runner/diagnose_cancel_test.go b/internal/runner/diagnose_cancel_test.go new file mode 100644 index 0000000..1d6a41f --- /dev/null +++ b/internal/runner/diagnose_cancel_test.go @@ -0,0 +1,63 @@ +package runner + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDiagnoseGracefulStopRequested(t *testing.T) { + t.Parallel() + + t.Run("background", func(t *testing.T) { + t.Parallel() + assert.False(t, DiagnoseGracefulStopRequested(context.Background())) + }) + + t.Run("diagnose_context_before_signal", func(t *testing.T) { + t.Parallel() + ctx, stop := NewDiagnoseRunContext(context.Background()) + defer stop() + assert.False(t, DiagnoseGracefulStopRequested(ctx)) + }) + + t.Run("diagnose_context_after_graceful_request", func(t *testing.T) { + t.Parallel() + ctx, stop := NewDiagnoseRunContext(context.Background()) + defer stop() + RequestDiagnoseGracefulStop(ctx) + assert.True(t, DiagnoseGracefulStopRequested(ctx)) + }) +} + +func TestDiagnoseHardCancel(t *testing.T) { + t.Parallel() + ctx, stop := NewDiagnoseRunContext(context.Background()) + defer stop() + + RequestDiagnoseHardCancel(ctx) + require.Error(t, ctx.Err()) + assert.False(t, DiagnoseGracefulStopRequested(ctx)) +} + +func TestDiagnoseGracefulThenHardCancel(t *testing.T) { + t.Parallel() + ctx, stop := NewDiagnoseRunContext(context.Background()) + defer stop() + + RequestDiagnoseGracefulStop(ctx) + assert.True(t, DiagnoseGracefulStopRequested(ctx)) + require.NoError(t, ctx.Err()) + + RequestDiagnoseHardCancel(ctx) + require.Error(t, ctx.Err()) +} + +func TestDiagnoseRunContextStopTearsDown(t *testing.T) { + t.Parallel() + ctx, stop := NewDiagnoseRunContext(context.Background()) + stop() + require.Error(t, ctx.Err()) +} diff --git a/internal/runner/diagnose_output.go b/internal/runner/diagnose_output.go index 5322944..70e166b 100644 --- a/internal/runner/diagnose_output.go +++ b/internal/runner/diagnose_output.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "runtime" "time" "charm.land/lipgloss/v2" @@ -154,3 +155,30 @@ func diagnoseCSVPath(resultsDir string, rep *Report) string { } return filepath.Join(resultsDir, "report.csv") } + +func diagnoseInterruptKeyHint() string { + if runtime.GOOS == "darwin" { + return "⌘C" + } + return "Ctrl+C" +} + +func printDiagnoseGracefulStopNotice(out *output.Printer, completed, total int) { + if out == nil { + return + } + if out.AIOutput() { + out.Stderrf("stop_graceful completed=%d total=%d\n", completed, total) + return + } + out.ClearInline() + hint := diagnoseInterruptKeyHint() + out.HumanStderr( + termstyle.Accent.Render( + fmt.Sprintf("Stopping diagnose run after current iteration — %d/%d completed.", completed, total), + ) + "\n" + + termstyle.Muted.Render( + fmt.Sprintf("Press %s again to cancel immediately.", hint), + ), + ) +} diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 62c091e..5e0e288 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -16,6 +16,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/smartcontractkit/testrig/internal/config" @@ -61,6 +62,7 @@ type DiagnoseRunState struct { failedFast bool failedFastReason string failedFastIteration int // 0-based diagnose iteration index; -1 if unset + GracefulStop bool iterDurations []time.Duration shuffleSeeds map[int]int64 liveProgress bool @@ -254,16 +256,30 @@ func FinishDiagnoseAnalysis( } func printDiagnoseInterruptions(ctx context.Context, conf *config.App, out *output.Printer, state *DiagnoseRunState) { - interrupted := ctx.Err() != nil - if interrupted && !out.AIOutput() { - out.HumanStderr( - termstyle.Accent.Render( - fmt.Sprintf("interrupted after %d/%d iterations", state.completed, conf.Iterations), - ) + - termstyle.Muted.Render( - " — analyzing partial results…", - ), - ) + if state.GracefulStop { + if !out.AIOutput() { + out.HumanStderr( + termstyle.Accent.Render( + fmt.Sprintf("stopped early after %d/%d iterations", state.completed, conf.Iterations), + ) + + termstyle.Muted.Render( + " — analyzing partial results…", + ), + ) + } + } else if interrupted := ctx.Err() != nil; interrupted { + if out.AIOutput() { + out.Stderrf("interrupted completed=%d total=%d\n", state.completed, conf.Iterations) + } else { + out.HumanStderr( + termstyle.Accent.Render( + fmt.Sprintf("interrupted after %d/%d iterations", state.completed, conf.Iterations), + ) + + termstyle.Muted.Render( + " — analyzing partial results…", + ), + ) + } } if state.failedFast && !out.AIOutput() { @@ -543,6 +559,7 @@ func runDiagnoseIterations( serialProgressMu = new(sync.Mutex) } + var inFlight atomic.Int32 worker := diagnoseWorker{ conf: conf, out: out, @@ -557,6 +574,7 @@ func runDiagnoseIterations( jobs: jobs, results: results, cancel: cancel, + inFlight: &inFlight, } var wg sync.WaitGroup for _, resource := range resources { @@ -567,10 +585,24 @@ func runDiagnoseIterations( wg.Go(func() { defer close(jobs) + gracefulStopCh := DiagnoseGracefulStopChan(ctx) for i := range conf.Iterations { + if DiagnoseGracefulStopRequested(ctx) { + return + } + if gracefulStopCh == nil { + select { + case <-runCtx.Done(): + return + case jobs <- i: + } + continue + } select { case <-runCtx.Done(): return + case <-gracefulStopCh: + return case jobs <- i: } } @@ -582,6 +614,17 @@ func runDiagnoseIterations( }() var firstErr error + var gracefulStopNotice sync.Once + maybePrintGracefulStopNotice := func() { + if !DiagnoseGracefulStopRequested(ctx) || inFlight.Load() > 0 { + return + } + gracefulStopNotice.Do(func() { + state.GracefulStop = true + printDiagnoseGracefulStopNotice(out, state.completed, conf.Iterations) + }) + } + for result := range results { if result.fatalErr != nil { if firstErr == nil { @@ -591,7 +634,9 @@ func runDiagnoseIterations( continue } recordDiagnoseIterationResult(state, result, conf, out, parallelProgress, serialProgressMu) + maybePrintGracefulStopNotice() } + maybePrintGracefulStopNotice() return state, firstErr } @@ -612,16 +657,20 @@ type diagnoseWorker struct { jobs <-chan int results chan<- diagnoseIterationResult cancel context.CancelFunc + inFlight *atomic.Int32 } func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationResource) { used := false for iteration := range w.jobs { + stopInFlight := w.trackInFlight() if runCtx.Err() != nil { + stopInFlight() return } if used && resource.Reset != nil { if err := resource.Reset(runCtx); err != nil { + stopInFlight() if runCtx.Err() != nil { return } @@ -632,6 +681,7 @@ func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationR used = true if w.hooks.iterationSetup != nil { if err := w.hooks.iterationSetup(runCtx); err != nil { + stopInFlight() w.sendFatal(runCtx, iteration, fmt.Errorf("iteration setup %d: %w", iteration, err)) return } @@ -682,6 +732,7 @@ func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationR digest: digest, digestErr: digestErr, } + stopInFlight() select { case w.results <- result: case <-runCtx.Done(): @@ -693,6 +744,14 @@ func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationR } } +func (w *diagnoseWorker) trackInFlight() func() { + if w.inFlight == nil { + return func() {} + } + w.inFlight.Add(1) + return func() { w.inFlight.Add(-1) } +} + // sendFatal posts a fatal iteration result and cancels the run. The select // guards against ctx cancellation racing the send. func (w *diagnoseWorker) sendFatal(runCtx context.Context, iteration int, err error) { diff --git a/internal/runner/runner_test.go b/internal/runner/runner_test.go index 49d2b1a..9a9b5f4 100644 --- a/internal/runner/runner_test.go +++ b/internal/runner/runner_test.go @@ -12,6 +12,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/assert" @@ -1327,6 +1328,264 @@ func TestRunDiagnoseIterationsCallsIterationHooks(t *testing.T) { assert.True(t, setupBeforeTeardown.Load(), "setup must be called before teardown") } +func writePassIterationJSONL(t *testing.T, resultsDir string, iteration int) { + t.Helper() + require.NoError( + t, + os.WriteFile( + filepath.Join(resultsDir, "iteration-"+strconv.Itoa(iteration)+".log.jsonl"), + []byte(`{"Action":"pass","Package":"p","Test":"T","Elapsed":0.01}`+"\n"), + 0600, + ), + ) +} + +func TestRunDiagnoseIterations_gracefulStop_finishesInFlight(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + ctx, stop := NewDiagnoseRunContextForTest(context.Background()) + defer stop() + + resultsDir := t.TempDir() + conf := &config.App{ + RepoRoot: makeTestRepoRoot(t), + AIOutput: true, + Iterations: 5, + } + out := output.NewForTest(true, io.Discard, io.Discard, false) + + var mu sync.Mutex + started := make(map[int]struct{}) + hooks := diagnoseRunHooks{ + runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { + mu.Lock() + started[p.Iteration] = struct{}{} + mu.Unlock() + if p.Iteration == 0 { + RequestDiagnoseGracefulStop(runCtx) + time.Sleep(1 * time.Second) + } + writePassIterationJSONL(t, p.ResultsDir, p.Iteration) + return nil + }, + } + + state, err := runDiagnoseIterations( + ctx, + conf, + out, + resultsDir, + []string{"./pkg"}, + []diagnoseIterationResource{{}}, + hooks, + ) + require.NoError(t, err) + assert.Equal(t, 1, state.completed) + assert.True(t, state.GracefulStop) + assert.Len(t, started, 1) + }) +} + +func TestRunDiagnoseIterations_gracefulStop_printsNoticeAfterDigest(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + ctx, stop := NewDiagnoseRunContextForTest(context.Background()) + defer stop() + + resultsDir := t.TempDir() + conf := &config.App{ + RepoRoot: makeTestRepoRoot(t), + AIOutput: false, + Iterations: 3, + } + var stderr strings.Builder + out := output.NewForTest(false, io.Discard, &stderr, false) + + hooks := diagnoseRunHooks{ + runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { + if p.Iteration == 0 { + RequestDiagnoseGracefulStop(runCtx) + } + writePassIterationJSONL(t, p.ResultsDir, p.Iteration) + return nil + }, + } + + state, err := runDiagnoseIterations( + ctx, + conf, + out, + resultsDir, + []string{"./pkg"}, + []diagnoseIterationResource{{}}, + hooks, + ) + require.NoError(t, err) + assert.True(t, state.GracefulStop) + assert.Contains(t, stderr.String(), "Stopping diagnose run after current iteration") + assert.Contains(t, stderr.String(), "again to cancel immediately") + }) +} + +func TestRunDiagnoseIterations_gracefulStop_idleBetweenIterations(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + ctx, stop := NewDiagnoseRunContextForTest(context.Background()) + defer stop() + + resultsDir := t.TempDir() + conf := &config.App{ + RepoRoot: makeTestRepoRoot(t), + AIOutput: true, + Iterations: 3, + } + out := output.NewForTest(true, io.Discard, io.Discard, false) + + iter0Done := make(chan struct{}) + iter0Release := make(chan struct{}) + hooks := diagnoseRunHooks{ + runIteration: func(_ context.Context, p diagnoseIterationParams) error { + if p.Iteration == 0 { + writePassIterationJSONL(t, p.ResultsDir, p.Iteration) + close(iter0Done) + <-iter0Release + return nil + } + writePassIterationJSONL(t, p.ResultsDir, p.Iteration) + return nil + }, + } + + var state *DiagnoseRunState + var runErr error + go func() { + state, runErr = runDiagnoseIterations( + ctx, + conf, + out, + resultsDir, + []string{"./pkg"}, + []diagnoseIterationResource{{}}, + hooks, + ) + }() + + <-iter0Done + RequestDiagnoseGracefulStop(ctx) + close(iter0Release) + synctest.Wait() + + require.NoError(t, runErr) + require.NotNil(t, state) + assert.Equal(t, 1, state.completed) + assert.True(t, state.GracefulStop) + }) +} + +func TestRunDiagnoseIterations_hardCancel_abortsInFlight(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + ctx, stop := NewDiagnoseRunContextForTest(context.Background()) + defer stop() + + resultsDir := t.TempDir() + conf := &config.App{ + RepoRoot: makeTestRepoRoot(t), + AIOutput: true, + Iterations: 3, + } + out := output.NewForTest(true, io.Discard, io.Discard, false) + + iter0Blocked := make(chan struct{}) + hooks := diagnoseRunHooks{ + runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { + if p.Iteration == 0 { + close(iter0Blocked) + <-runCtx.Done() + return runCtx.Err() + } + writePassIterationJSONL(t, p.ResultsDir, p.Iteration) + return nil + }, + } + + var state *DiagnoseRunState + var runErr error + go func() { + state, runErr = runDiagnoseIterations( + ctx, + conf, + out, + resultsDir, + []string{"./pkg"}, + []diagnoseIterationResource{{}}, + hooks, + ) + }() + + <-iter0Blocked + RequestDiagnoseHardCancel(ctx) + synctest.Wait() + + require.NoError(t, runErr) + require.NotNil(t, state) + assert.Less(t, state.completed, conf.Iterations) + assert.False(t, state.GracefulStop) + require.Error(t, ctx.Err()) + }) +} + +func TestDiagnoseGracefulStop_writesPartialReport(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + repoRoot := makeTestRepoRoot(t) + conf := &config.App{ + RepoRoot: repoRoot, + AIOutput: true, + Iterations: 4, + } + ctx, stop := NewDiagnoseRunContextForTest(context.Background()) + defer stop() + out := output.NewForTest(conf.AIOutput, io.Discard, io.Discard, false) + + resultsDir := t.TempDir() + start := time.Now() + hooks := diagnoseRunHooks{ + runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { + if p.Iteration == 1 { + RequestDiagnoseGracefulStop(runCtx) + } + writePassIterationJSONL(t, p.ResultsDir, p.Iteration) + return nil + }, + } + + state, err := runDiagnoseIterations( + ctx, + conf, + out, + resultsDir, + []string{"./pkg"}, + []diagnoseIterationResource{{}}, + hooks, + ) + require.NoError(t, err) + require.True(t, state.GracefulStop) + require.Equal(t, 2, state.completed) + + require.NoError(t, writeRunState(resultsDir, conf, []string{"./pkg"}, state, start)) + require.NoError(t, FinishDiagnoseAnalysis(ctx, conf, out, []string{"./pkg"}, state, start, resultsDir)) + + //nolint:gosec // G304: path from filepath.Join + reportBytes, err := os.ReadFile(filepath.Join(resultsDir, "report.json")) + require.NoError(t, err) + var rep Report + require.NoError(t, json.Unmarshal(reportBytes, &rep)) + assert.Equal(t, 2, rep.Iterations) + assert.Len(t, rep.IterationSummaries, 2) + }) +} + func makeTestRepoRoot(t *testing.T) string { d := t.TempDir() err := os.WriteFile(filepath.Join(d, "go.mod"), []byte("module test\n"), 0600) From b3d1c7eecb0444846eefe9af99573ffc25ab4c0d Mon Sep 17 00:00:00 2001 From: Adam Hamrick Date: Wed, 17 Jun 2026 19:43:10 -0400 Subject: [PATCH 2/3] DRY --- gracefulCancel_main_review.xml | 786 +++++++++++++++++++++++++++++++++ internal/runner/runner.go | 149 ++++--- 2 files changed, 865 insertions(+), 70 deletions(-) create mode 100644 gracefulCancel_main_review.xml diff --git a/gracefulCancel_main_review.xml b/gracefulCancel_main_review.xml new file mode 100644 index 0000000..94275c3 --- /dev/null +++ b/gracefulCancel_main_review.xml @@ -0,0 +1,786 @@ + + + + + + + package cmd + + import ( ++ "context" + "time" + + "github.com/spf13/cobra" + + + iterSetup := hooks.BuildIterationHook(runnerOpts, shell, hooks.PhaseSetup) + iterTeardown := hooks.BuildIterationHook(runnerOpts, shell, hooks.PhaseTeardown) + ++ runCtx, stopRun := runner.NewDiagnoseRunContext(context.WithoutCancel(cmd.Context())) ++ defer stopRun() ++ + state, start, resultsDir, runErr := runner.RunIterations( +- cmd.Context(), ++ runCtx, + conf, + out, + args, + + + + finishResourceCleanup(cmd, &runErr, cleanup) + if runErr != nil { +- if cmd.Context().Err() == nil { ++ if runCtx.Err() == nil { + return runErr + } + } + +- return runner.FinishDiagnoseAnalysis(cmd.Context(), conf, out, args, state, start, resultsDir) ++ return runner.FinishDiagnoseAnalysis(runCtx, conf, out, args, state, start, resultsDir) + }), + } + + + + + + + return rootCmd + } + +-// Execute runs the root command. A SIGINT or SIGTERM cancels the context so +-// long-running subcommands (notably `diagnose`) can stop cleanly and still write +-// their post-run analysis. A second signal hits the default handler and +-// force-exits. ++// Execute runs the root command. A SIGINT or SIGTERM cancels cmd.Context() on ++// the first signal so generic subcommands (go test, gotestsum) stop promptly. ++// diagnose installs its own two-stage context: first signal finishes in-flight ++// iterations; second signal hard-cancels. A second root signal hits the default ++// handler and force-exits when diagnose is not running. + func Execute(opts ...hooks.Option) { + if err := runExecute(opts...); err != nil { + os.Exit(1) + + + + + ++package runner ++ ++import ( ++ "context" ++ "os" ++ "os/signal" ++ "sync" ++ "sync/atomic" ++ "syscall" ++) ++ ++type diagnoseCancelKey struct{} ++ ++type diagnoseCancelState struct { ++ softStop atomic.Bool ++ hardCancel context.CancelFunc ++ softStopCh chan struct{} ++ softOnce sync.Once ++} ++ ++// NewDiagnoseRunContext returns a context for diagnose iteration runs with ++// two-stage cancellation. The first SIGINT/SIGTERM requests a graceful stop ++// (finish in-flight iterations, do not enqueue new ones). The second signal ++// hard-cancels the context. Pass context.WithoutCancel(cmd.Context()) as parent ++// so the root CLI signal handler does not cancel iteration execution on the ++// first press. ++func NewDiagnoseRunContext(parent context.Context) (context.Context, func()) { ++ return newDiagnoseRunContext(parent, true) ++} ++ ++// NewDiagnoseRunContextForTest is like NewDiagnoseRunContext but without OS ++// signal handling. Use in unit and synctest tests with RequestDiagnoseGracefulStop ++// and RequestDiagnoseHardCancel. ++func NewDiagnoseRunContextForTest(parent context.Context) (context.Context, func()) { ++ return newDiagnoseRunContext(parent, false) ++} ++ ++func newDiagnoseRunContext(parent context.Context, listenSignals bool) (context.Context, func()) { ++ ctx, hardCancel := context.WithCancel(parent) ++ state := &diagnoseCancelState{ ++ hardCancel: hardCancel, ++ softStopCh: make(chan struct{}), ++ } ++ ctx = context.WithValue(ctx, diagnoseCancelKey{}, state) ++ ++ stop := func() { ++ hardCancel() ++ } ++ ++ if !listenSignals { ++ return ctx, stop ++ } ++ ++ sigCh := make(chan os.Signal, 2) ++ signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM) ++ ++ stop = func() { ++ signal.Stop(sigCh) ++ hardCancel() ++ } ++ ++ go func() { ++ defer signal.Stop(sigCh) ++ for { ++ select { ++ case <-ctx.Done(): ++ return ++ case <-sigCh: ++ state.requestGracefulStop() ++ select { ++ case <-sigCh: ++ hardCancel() ++ return ++ case <-ctx.Done(): ++ return ++ } ++ } ++ } ++ }() ++ ++ return ctx, stop ++} ++ ++func (s *diagnoseCancelState) requestGracefulStop() { ++ s.softOnce.Do(func() { ++ s.softStop.Store(true) ++ close(s.softStopCh) ++ }) ++} ++ ++func diagnoseCancelFromContext(ctx context.Context) *diagnoseCancelState { ++ if ctx == nil { ++ return nil ++ } ++ state, _ := ctx.Value(diagnoseCancelKey{}).(*diagnoseCancelState) ++ return state ++} ++ ++// DiagnoseGracefulStopChan returns a channel closed on the first graceful stop ++// request, or nil when ctx is not from NewDiagnoseRunContext. ++func DiagnoseGracefulStopChan(ctx context.Context) <-chan struct{} { ++ state := diagnoseCancelFromContext(ctx) ++ if state == nil { ++ return nil ++ } ++ return state.softStopCh ++} ++ ++// DiagnoseGracefulStopRequested reports whether the user requested a graceful ++// stop (first Ctrl+C) on a context from NewDiagnoseRunContext. ++func DiagnoseGracefulStopRequested(ctx context.Context) bool { ++ state := diagnoseCancelFromContext(ctx) ++ return state != nil && state.softStop.Load() ++} ++ ++// RequestDiagnoseGracefulStop simulates the first interrupt for tests. ++func RequestDiagnoseGracefulStop(ctx context.Context) { ++ state := diagnoseCancelFromContext(ctx) ++ if state != nil { ++ state.requestGracefulStop() ++ } ++} ++ ++// RequestDiagnoseHardCancel simulates the second interrupt for tests. ++func RequestDiagnoseHardCancel(ctx context.Context) { ++ state := diagnoseCancelFromContext(ctx) ++ if state != nil { ++ state.hardCancel() ++ } ++} + + + + + ++package runner ++ ++import ( ++ "context" ++ "testing" ++ ++ "github.com/stretchr/testify/assert" ++ "github.com/stretchr/testify/require" ++) ++ ++func TestDiagnoseGracefulStopRequested(t *testing.T) { ++ t.Parallel() ++ ++ t.Run("background", func(t *testing.T) { ++ t.Parallel() ++ assert.False(t, DiagnoseGracefulStopRequested(context.Background())) ++ }) ++ ++ t.Run("diagnose_context_before_signal", func(t *testing.T) { ++ t.Parallel() ++ ctx, stop := NewDiagnoseRunContext(context.Background()) ++ defer stop() ++ assert.False(t, DiagnoseGracefulStopRequested(ctx)) ++ }) ++ ++ t.Run("diagnose_context_after_graceful_request", func(t *testing.T) { ++ t.Parallel() ++ ctx, stop := NewDiagnoseRunContext(context.Background()) ++ defer stop() ++ RequestDiagnoseGracefulStop(ctx) ++ assert.True(t, DiagnoseGracefulStopRequested(ctx)) ++ }) ++} ++ ++func TestDiagnoseHardCancel(t *testing.T) { ++ t.Parallel() ++ ctx, stop := NewDiagnoseRunContext(context.Background()) ++ defer stop() ++ ++ RequestDiagnoseHardCancel(ctx) ++ require.Error(t, ctx.Err()) ++ assert.False(t, DiagnoseGracefulStopRequested(ctx)) ++} ++ ++func TestDiagnoseGracefulThenHardCancel(t *testing.T) { ++ t.Parallel() ++ ctx, stop := NewDiagnoseRunContext(context.Background()) ++ defer stop() ++ ++ RequestDiagnoseGracefulStop(ctx) ++ assert.True(t, DiagnoseGracefulStopRequested(ctx)) ++ require.NoError(t, ctx.Err()) ++ ++ RequestDiagnoseHardCancel(ctx) ++ require.Error(t, ctx.Err()) ++} ++ ++func TestDiagnoseRunContextStopTearsDown(t *testing.T) { ++ t.Parallel() ++ ctx, stop := NewDiagnoseRunContext(context.Background()) ++ stop() ++ require.Error(t, ctx.Err()) ++} + + + + + + "fmt" + "os" + "path/filepath" ++ "runtime" + "time" + + "charm.land/lipgloss/v2" + + + } + return filepath.Join(resultsDir, "report.csv") + } ++ ++func diagnoseInterruptKeyHint() string { ++ if runtime.GOOS == "darwin" { ++ return "⌘C" ++ } ++ return "Ctrl+C" ++} ++ ++func printDiagnoseGracefulStopNotice(out *output.Printer, completed, total int) { ++ if out == nil { ++ return ++ } ++ if out.AIOutput() { ++ out.Stderrf("stop_graceful completed=%d total=%d\n", completed, total) ++ return ++ } ++ out.ClearInline() ++ hint := diagnoseInterruptKeyHint() ++ out.HumanStderr( ++ termstyle.Accent.Render( ++ fmt.Sprintf("Stopping diagnose run after current iteration — %d/%d completed.", completed, total), ++ ) + "\n" + ++ termstyle.Muted.Render( ++ fmt.Sprintf("Press %s again to cancel immediately.", hint), ++ ), ++ ) ++} + + + + + + "strconv" + "strings" + "sync" ++ "sync/atomic" + "time" + + "github.com/smartcontractkit/testrig/internal/config" + + + failedFast bool + failedFastReason string + failedFastIteration int // 0-based diagnose iteration index; -1 if unset ++ GracefulStop bool + iterDurations []time.Duration + shuffleSeeds map[int]int64 + liveProgress bool + + + } + + func printDiagnoseInterruptions(ctx context.Context, conf *config.App, out *output.Printer, state *DiagnoseRunState) { +- interrupted := ctx.Err() != nil +- if interrupted && !out.AIOutput() { +- out.HumanStderr( +- termstyle.Accent.Render( +- fmt.Sprintf("interrupted after %d/%d iterations", state.completed, conf.Iterations), +- ) + +- termstyle.Muted.Render( +- " — analyzing partial results…", +- ), +- ) ++ if state.GracefulStop { ++ if !out.AIOutput() { ++ out.HumanStderr( ++ termstyle.Accent.Render( ++ fmt.Sprintf("stopped early after %d/%d iterations", state.completed, conf.Iterations), ++ ) + ++ termstyle.Muted.Render( ++ " — analyzing partial results…", ++ ), ++ ) ++ } ++ } else if interrupted := ctx.Err() != nil; interrupted { ++ if out.AIOutput() { ++ out.Stderrf("interrupted completed=%d total=%d\n", state.completed, conf.Iterations) ++ } else { ++ out.HumanStderr( ++ termstyle.Accent.Render( ++ fmt.Sprintf("interrupted after %d/%d iterations", state.completed, conf.Iterations), ++ ) + ++ termstyle.Muted.Render( ++ " — analyzing partial results…", ++ ), ++ ) ++ } + } + + if state.failedFast && !out.AIOutput() { + + + serialProgressMu = new(sync.Mutex) + } + ++ var inFlight atomic.Int32 + worker := diagnoseWorker{ + conf: conf, + out: out, + + + jobs: jobs, + results: results, + cancel: cancel, ++ inFlight: &inFlight, + } + var wg sync.WaitGroup + for _, resource := range resources { + + + + wg.Go(func() { + defer close(jobs) ++ gracefulStopCh := DiagnoseGracefulStopChan(ctx) + for i := range conf.Iterations { ++ if DiagnoseGracefulStopRequested(ctx) { ++ return ++ } ++ if gracefulStopCh == nil { ++ select { ++ case <-runCtx.Done(): ++ return ++ case jobs <- i: ++ } ++ continue ++ } + select { + case <-runCtx.Done(): + return ++ case <-gracefulStopCh: ++ return + case jobs <- i: + } + } + + + }() + + var firstErr error ++ var gracefulStopNotice sync.Once ++ maybePrintGracefulStopNotice := func() { ++ if !DiagnoseGracefulStopRequested(ctx) || inFlight.Load() > 0 { ++ return ++ } ++ gracefulStopNotice.Do(func() { ++ state.GracefulStop = true ++ printDiagnoseGracefulStopNotice(out, state.completed, conf.Iterations) ++ }) ++ } ++ + for result := range results { + if result.fatalErr != nil { + if firstErr == nil { + + + continue + } + recordDiagnoseIterationResult(state, result, conf, out, parallelProgress, serialProgressMu) ++ maybePrintGracefulStopNotice() + } ++ maybePrintGracefulStopNotice() + return state, firstErr + } + + + + jobs <-chan int + results chan<- diagnoseIterationResult + cancel context.CancelFunc ++ inFlight *atomic.Int32 + } + + func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationResource) { + used := false + for iteration := range w.jobs { ++ stopInFlight := w.trackInFlight() + if runCtx.Err() != nil { ++ stopInFlight() + return + } + if used && resource.Reset != nil { + if err := resource.Reset(runCtx); err != nil { ++ stopInFlight() + if runCtx.Err() != nil { + return + } + + + used = true + if w.hooks.iterationSetup != nil { + if err := w.hooks.iterationSetup(runCtx); err != nil { ++ stopInFlight() + w.sendFatal(runCtx, iteration, fmt.Errorf("iteration setup %d: %w", iteration, err)) + return + } + + + digest: digest, + digestErr: digestErr, + } ++ stopInFlight() + select { + case w.results <- result: + case <-runCtx.Done(): + + + } + } + ++func (w *diagnoseWorker) trackInFlight() func() { ++ if w.inFlight == nil { ++ return func() {} ++ } ++ w.inFlight.Add(1) ++ return func() { w.inFlight.Add(-1) } ++} ++ + // sendFatal posts a fatal iteration result and cancels the run. The select + // guards against ctx cancellation racing the send. + func (w *diagnoseWorker) sendFatal(runCtx context.Context, iteration int, err error) { + + + + + + "sync" + "sync/atomic" + "testing" ++ "testing/synctest" + "time" + + "github.com/stretchr/testify/assert" + + + assert.True(t, setupBeforeTeardown.Load(), "setup must be called before teardown") + } + ++func writePassIterationJSONL(t *testing.T, resultsDir string, iteration int) { ++ t.Helper() ++ require.NoError( ++ t, ++ os.WriteFile( ++ filepath.Join(resultsDir, "iteration-"+strconv.Itoa(iteration)+".log.jsonl"), ++ []byte(`{"Action":"pass","Package":"p","Test":"T","Elapsed":0.01}`+"\n"), ++ 0600, ++ ), ++ ) ++} ++ ++func TestRunDiagnoseIterations_gracefulStop_finishesInFlight(t *testing.T) { ++ t.Parallel() ++ synctest.Test(t, func(t *testing.T) { ++ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) ++ defer stop() ++ ++ resultsDir := t.TempDir() ++ conf := &config.App{ ++ RepoRoot: makeTestRepoRoot(t), ++ AIOutput: true, ++ Iterations: 5, ++ } ++ out := output.NewForTest(true, io.Discard, io.Discard, false) ++ ++ var mu sync.Mutex ++ started := make(map[int]struct{}) ++ hooks := diagnoseRunHooks{ ++ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { ++ mu.Lock() ++ started[p.Iteration] = struct{}{} ++ mu.Unlock() ++ if p.Iteration == 0 { ++ RequestDiagnoseGracefulStop(runCtx) ++ time.Sleep(1 * time.Second) ++ } ++ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) ++ return nil ++ }, ++ } ++ ++ state, err := runDiagnoseIterations( ++ ctx, ++ conf, ++ out, ++ resultsDir, ++ []string{"./pkg"}, ++ []diagnoseIterationResource{{}}, ++ hooks, ++ ) ++ require.NoError(t, err) ++ assert.Equal(t, 1, state.completed) ++ assert.True(t, state.GracefulStop) ++ assert.Len(t, started, 1) ++ }) ++} ++ ++func TestRunDiagnoseIterations_gracefulStop_printsNoticeAfterDigest(t *testing.T) { ++ t.Parallel() ++ synctest.Test(t, func(t *testing.T) { ++ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) ++ defer stop() ++ ++ resultsDir := t.TempDir() ++ conf := &config.App{ ++ RepoRoot: makeTestRepoRoot(t), ++ AIOutput: false, ++ Iterations: 3, ++ } ++ var stderr strings.Builder ++ out := output.NewForTest(false, io.Discard, &stderr, false) ++ ++ hooks := diagnoseRunHooks{ ++ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { ++ if p.Iteration == 0 { ++ RequestDiagnoseGracefulStop(runCtx) ++ } ++ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) ++ return nil ++ }, ++ } ++ ++ state, err := runDiagnoseIterations( ++ ctx, ++ conf, ++ out, ++ resultsDir, ++ []string{"./pkg"}, ++ []diagnoseIterationResource{{}}, ++ hooks, ++ ) ++ require.NoError(t, err) ++ assert.True(t, state.GracefulStop) ++ assert.Contains(t, stderr.String(), "Stopping diagnose run after current iteration") ++ assert.Contains(t, stderr.String(), "again to cancel immediately") ++ }) ++} ++ ++func TestRunDiagnoseIterations_gracefulStop_idleBetweenIterations(t *testing.T) { ++ t.Parallel() ++ synctest.Test(t, func(t *testing.T) { ++ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) ++ defer stop() ++ ++ resultsDir := t.TempDir() ++ conf := &config.App{ ++ RepoRoot: makeTestRepoRoot(t), ++ AIOutput: true, ++ Iterations: 3, ++ } ++ out := output.NewForTest(true, io.Discard, io.Discard, false) ++ ++ iter0Done := make(chan struct{}) ++ iter0Release := make(chan struct{}) ++ hooks := diagnoseRunHooks{ ++ runIteration: func(_ context.Context, p diagnoseIterationParams) error { ++ if p.Iteration == 0 { ++ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) ++ close(iter0Done) ++ <-iter0Release ++ return nil ++ } ++ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) ++ return nil ++ }, ++ } ++ ++ var state *DiagnoseRunState ++ var runErr error ++ go func() { ++ state, runErr = runDiagnoseIterations( ++ ctx, ++ conf, ++ out, ++ resultsDir, ++ []string{"./pkg"}, ++ []diagnoseIterationResource{{}}, ++ hooks, ++ ) ++ }() ++ ++ <-iter0Done ++ RequestDiagnoseGracefulStop(ctx) ++ close(iter0Release) ++ synctest.Wait() ++ ++ require.NoError(t, runErr) ++ require.NotNil(t, state) ++ assert.Equal(t, 1, state.completed) ++ assert.True(t, state.GracefulStop) ++ }) ++} ++ ++func TestRunDiagnoseIterations_hardCancel_abortsInFlight(t *testing.T) { ++ t.Parallel() ++ synctest.Test(t, func(t *testing.T) { ++ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) ++ defer stop() ++ ++ resultsDir := t.TempDir() ++ conf := &config.App{ ++ RepoRoot: makeTestRepoRoot(t), ++ AIOutput: true, ++ Iterations: 3, ++ } ++ out := output.NewForTest(true, io.Discard, io.Discard, false) ++ ++ iter0Blocked := make(chan struct{}) ++ hooks := diagnoseRunHooks{ ++ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { ++ if p.Iteration == 0 { ++ close(iter0Blocked) ++ <-runCtx.Done() ++ return runCtx.Err() ++ } ++ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) ++ return nil ++ }, ++ } ++ ++ var state *DiagnoseRunState ++ var runErr error ++ go func() { ++ state, runErr = runDiagnoseIterations( ++ ctx, ++ conf, ++ out, ++ resultsDir, ++ []string{"./pkg"}, ++ []diagnoseIterationResource{{}}, ++ hooks, ++ ) ++ }() ++ ++ <-iter0Blocked ++ RequestDiagnoseHardCancel(ctx) ++ synctest.Wait() ++ ++ require.NoError(t, runErr) ++ require.NotNil(t, state) ++ assert.Less(t, state.completed, conf.Iterations) ++ assert.False(t, state.GracefulStop) ++ require.Error(t, ctx.Err()) ++ }) ++} ++ ++func TestDiagnoseGracefulStop_writesPartialReport(t *testing.T) { ++ t.Parallel() ++ synctest.Test(t, func(t *testing.T) { ++ repoRoot := makeTestRepoRoot(t) ++ conf := &config.App{ ++ RepoRoot: repoRoot, ++ AIOutput: true, ++ Iterations: 4, ++ } ++ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) ++ defer stop() ++ out := output.NewForTest(conf.AIOutput, io.Discard, io.Discard, false) ++ ++ resultsDir := t.TempDir() ++ start := time.Now() ++ hooks := diagnoseRunHooks{ ++ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { ++ if p.Iteration == 1 { ++ RequestDiagnoseGracefulStop(runCtx) ++ } ++ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) ++ return nil ++ }, ++ } ++ ++ state, err := runDiagnoseIterations( ++ ctx, ++ conf, ++ out, ++ resultsDir, ++ []string{"./pkg"}, ++ []diagnoseIterationResource{{}}, ++ hooks, ++ ) ++ require.NoError(t, err) ++ require.True(t, state.GracefulStop) ++ require.Equal(t, 2, state.completed) ++ ++ require.NoError(t, writeRunState(resultsDir, conf, []string{"./pkg"}, state, start)) ++ require.NoError(t, FinishDiagnoseAnalysis(ctx, conf, out, []string{"./pkg"}, state, start, resultsDir)) ++ ++ //nolint:gosec // G304: path from filepath.Join ++ reportBytes, err := os.ReadFile(filepath.Join(resultsDir, "report.json")) ++ require.NoError(t, err) ++ var rep Report ++ require.NoError(t, json.Unmarshal(reportBytes, &rep)) ++ assert.Equal(t, 2, rep.Iterations) ++ assert.Len(t, rep.IterationSummaries, 2) ++ }) ++} ++ + func makeTestRepoRoot(t *testing.T) string { + d := t.TempDir() + err := os.WriteFile(filepath.Join(d, "go.mod"), []byte("module test\n"), 0600) + + + + + diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 5e0e288..c88ad64 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -663,85 +663,94 @@ type diagnoseWorker struct { func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationResource) { used := false for iteration := range w.jobs { - stopInFlight := w.trackInFlight() - if runCtx.Err() != nil { - stopInFlight() + done := w.runIterationJob(runCtx, resource, iteration, &used) + if done { return } - if used && resource.Reset != nil { - if err := resource.Reset(runCtx); err != nil { - stopInFlight() - if runCtx.Err() != nil { - return - } - w.sendFatal(runCtx, iteration, fmt.Errorf("reset database before iteration %d: %w", iteration, err)) - return - } - } - used = true - if w.hooks.iterationSetup != nil { - if err := w.hooks.iterationSetup(runCtx); err != nil { - stopInFlight() - w.sendFatal(runCtx, iteration, fmt.Errorf("iteration setup %d: %w", iteration, err)) - return + } +} + +func (w *diagnoseWorker) runIterationJob( + runCtx context.Context, + resource diagnoseIterationResource, + iteration int, + used *bool, +) bool { + defer w.trackInFlight()() + if runCtx.Err() != nil { + return true + } + if *used && resource.Reset != nil { + if err := resource.Reset(runCtx); err != nil { + if runCtx.Err() != nil { + return true } + w.sendFatal(runCtx, iteration, fmt.Errorf("reset database before iteration %d: %w", iteration, err)) + return true } - var seed int64 - if w.conf.Shuffle { - seed = w.hooks.seed() - } - iterStart := time.Now() - iterErr := w.hooks.runIteration(runCtx, diagnoseIterationParams{ - Conf: w.conf, - Out: w.out, - ResultsDir: w.resultsDir, - GoTestArgs: w.goTestArgs, - ModuleDir: w.moduleDir, - Iteration: iteration, - ShuffleSeed: seed, - Env: resource.Env, - LiveProgress: w.parallel == 1, - ParallelProgress: w.parallelProgress, - DiagnoseRunStart: w.diagnoseRunStart, - SerialProgressMu: w.serialProgressMu, - }) - iterDur := time.Since(iterStart) - if w.hooks.iterationTeardown != nil { - if tdErr := w.hooks.iterationTeardown(runCtx); tdErr != nil && iterErr == nil { - iterErr = fmt.Errorf("iteration teardown %d: %w", iteration, tdErr) - } + } + *used = true + if w.hooks.iterationSetup != nil { + if err := w.hooks.iterationSetup(runCtx); err != nil { + w.sendFatal(runCtx, iteration, fmt.Errorf("iteration setup %d: %w", iteration, err)) + return true } - var dumpErr error - if resource.DumpDiagnostics != nil { - dumpErr = resource.DumpDiagnostics(runCtx, w.resultsDir, iteration) + } + var seed int64 + if w.conf.Shuffle { + seed = w.hooks.seed() + } + iterStart := time.Now() + iterErr := w.hooks.runIteration(runCtx, diagnoseIterationParams{ + Conf: w.conf, + Out: w.out, + ResultsDir: w.resultsDir, + GoTestArgs: w.goTestArgs, + ModuleDir: w.moduleDir, + Iteration: iteration, + ShuffleSeed: seed, + Env: resource.Env, + LiveProgress: w.parallel == 1, + ParallelProgress: w.parallelProgress, + DiagnoseRunStart: w.diagnoseRunStart, + SerialProgressMu: w.serialProgressMu, + }) + iterDur := time.Since(iterStart) + if w.hooks.iterationTeardown != nil { + if tdErr := w.hooks.iterationTeardown(runCtx); tdErr != nil && iterErr == nil { + iterErr = fmt.Errorf("iteration teardown %d: %w", iteration, tdErr) } - digest, digestErr := loadIterationDigest(w.resultsDir, iteration, w.conf.SlowThreshold) - failedFast, failReason := shouldFailFastIteration(w.conf, iterErr, digest, digestErr) - failedFast = failedFast && runCtx.Err() == nil + } + var dumpErr error + if resource.DumpDiagnostics != nil { + dumpErr = resource.DumpDiagnostics(runCtx, w.resultsDir, iteration) + } + digest, digestErr := loadIterationDigest(w.resultsDir, iteration, w.conf.SlowThreshold) + failedFast, failReason := shouldFailFastIteration(w.conf, iterErr, digest, digestErr) + failedFast = failedFast && runCtx.Err() == nil + if failedFast { + w.cancel() + } + result := diagnoseIterationResult{ + iteration: iteration, + duration: iterDur, + shuffle: seed, + iterErr: iterErr, + dumpErr: dumpErr, + failedFast: failedFast, + failReason: failReason, + digest: digest, + digestErr: digestErr, + } + select { + case w.results <- result: + case <-runCtx.Done(): if failedFast { - w.cancel() - } - result := diagnoseIterationResult{ - iteration: iteration, - duration: iterDur, - shuffle: seed, - iterErr: iterErr, - dumpErr: dumpErr, - failedFast: failedFast, - failReason: failReason, - digest: digest, - digestErr: digestErr, - } - stopInFlight() - select { - case w.results <- result: - case <-runCtx.Done(): - if failedFast { - w.results <- result - } - return + w.results <- result } + return true } + return false } func (w *diagnoseWorker) trackInFlight() func() { From e1aa2e691091f4a1aa74d44bef3afa49e0854f44 Mon Sep 17 00:00:00 2001 From: Adam Hamrick Date: Wed, 17 Jun 2026 19:44:06 -0400 Subject: [PATCH 3/3] delete --- gracefulCancel_main_review.xml | 786 --------------------------------- 1 file changed, 786 deletions(-) delete mode 100644 gracefulCancel_main_review.xml diff --git a/gracefulCancel_main_review.xml b/gracefulCancel_main_review.xml deleted file mode 100644 index 94275c3..0000000 --- a/gracefulCancel_main_review.xml +++ /dev/null @@ -1,786 +0,0 @@ - - - - - - - package cmd - - import ( -+ "context" - "time" - - "github.com/spf13/cobra" - - - iterSetup := hooks.BuildIterationHook(runnerOpts, shell, hooks.PhaseSetup) - iterTeardown := hooks.BuildIterationHook(runnerOpts, shell, hooks.PhaseTeardown) - -+ runCtx, stopRun := runner.NewDiagnoseRunContext(context.WithoutCancel(cmd.Context())) -+ defer stopRun() -+ - state, start, resultsDir, runErr := runner.RunIterations( -- cmd.Context(), -+ runCtx, - conf, - out, - args, - - - - finishResourceCleanup(cmd, &runErr, cleanup) - if runErr != nil { -- if cmd.Context().Err() == nil { -+ if runCtx.Err() == nil { - return runErr - } - } - -- return runner.FinishDiagnoseAnalysis(cmd.Context(), conf, out, args, state, start, resultsDir) -+ return runner.FinishDiagnoseAnalysis(runCtx, conf, out, args, state, start, resultsDir) - }), - } - - - - - - - return rootCmd - } - --// Execute runs the root command. A SIGINT or SIGTERM cancels the context so --// long-running subcommands (notably `diagnose`) can stop cleanly and still write --// their post-run analysis. A second signal hits the default handler and --// force-exits. -+// Execute runs the root command. A SIGINT or SIGTERM cancels cmd.Context() on -+// the first signal so generic subcommands (go test, gotestsum) stop promptly. -+// diagnose installs its own two-stage context: first signal finishes in-flight -+// iterations; second signal hard-cancels. A second root signal hits the default -+// handler and force-exits when diagnose is not running. - func Execute(opts ...hooks.Option) { - if err := runExecute(opts...); err != nil { - os.Exit(1) - - - - - -+package runner -+ -+import ( -+ "context" -+ "os" -+ "os/signal" -+ "sync" -+ "sync/atomic" -+ "syscall" -+) -+ -+type diagnoseCancelKey struct{} -+ -+type diagnoseCancelState struct { -+ softStop atomic.Bool -+ hardCancel context.CancelFunc -+ softStopCh chan struct{} -+ softOnce sync.Once -+} -+ -+// NewDiagnoseRunContext returns a context for diagnose iteration runs with -+// two-stage cancellation. The first SIGINT/SIGTERM requests a graceful stop -+// (finish in-flight iterations, do not enqueue new ones). The second signal -+// hard-cancels the context. Pass context.WithoutCancel(cmd.Context()) as parent -+// so the root CLI signal handler does not cancel iteration execution on the -+// first press. -+func NewDiagnoseRunContext(parent context.Context) (context.Context, func()) { -+ return newDiagnoseRunContext(parent, true) -+} -+ -+// NewDiagnoseRunContextForTest is like NewDiagnoseRunContext but without OS -+// signal handling. Use in unit and synctest tests with RequestDiagnoseGracefulStop -+// and RequestDiagnoseHardCancel. -+func NewDiagnoseRunContextForTest(parent context.Context) (context.Context, func()) { -+ return newDiagnoseRunContext(parent, false) -+} -+ -+func newDiagnoseRunContext(parent context.Context, listenSignals bool) (context.Context, func()) { -+ ctx, hardCancel := context.WithCancel(parent) -+ state := &diagnoseCancelState{ -+ hardCancel: hardCancel, -+ softStopCh: make(chan struct{}), -+ } -+ ctx = context.WithValue(ctx, diagnoseCancelKey{}, state) -+ -+ stop := func() { -+ hardCancel() -+ } -+ -+ if !listenSignals { -+ return ctx, stop -+ } -+ -+ sigCh := make(chan os.Signal, 2) -+ signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM) -+ -+ stop = func() { -+ signal.Stop(sigCh) -+ hardCancel() -+ } -+ -+ go func() { -+ defer signal.Stop(sigCh) -+ for { -+ select { -+ case <-ctx.Done(): -+ return -+ case <-sigCh: -+ state.requestGracefulStop() -+ select { -+ case <-sigCh: -+ hardCancel() -+ return -+ case <-ctx.Done(): -+ return -+ } -+ } -+ } -+ }() -+ -+ return ctx, stop -+} -+ -+func (s *diagnoseCancelState) requestGracefulStop() { -+ s.softOnce.Do(func() { -+ s.softStop.Store(true) -+ close(s.softStopCh) -+ }) -+} -+ -+func diagnoseCancelFromContext(ctx context.Context) *diagnoseCancelState { -+ if ctx == nil { -+ return nil -+ } -+ state, _ := ctx.Value(diagnoseCancelKey{}).(*diagnoseCancelState) -+ return state -+} -+ -+// DiagnoseGracefulStopChan returns a channel closed on the first graceful stop -+// request, or nil when ctx is not from NewDiagnoseRunContext. -+func DiagnoseGracefulStopChan(ctx context.Context) <-chan struct{} { -+ state := diagnoseCancelFromContext(ctx) -+ if state == nil { -+ return nil -+ } -+ return state.softStopCh -+} -+ -+// DiagnoseGracefulStopRequested reports whether the user requested a graceful -+// stop (first Ctrl+C) on a context from NewDiagnoseRunContext. -+func DiagnoseGracefulStopRequested(ctx context.Context) bool { -+ state := diagnoseCancelFromContext(ctx) -+ return state != nil && state.softStop.Load() -+} -+ -+// RequestDiagnoseGracefulStop simulates the first interrupt for tests. -+func RequestDiagnoseGracefulStop(ctx context.Context) { -+ state := diagnoseCancelFromContext(ctx) -+ if state != nil { -+ state.requestGracefulStop() -+ } -+} -+ -+// RequestDiagnoseHardCancel simulates the second interrupt for tests. -+func RequestDiagnoseHardCancel(ctx context.Context) { -+ state := diagnoseCancelFromContext(ctx) -+ if state != nil { -+ state.hardCancel() -+ } -+} - - - - - -+package runner -+ -+import ( -+ "context" -+ "testing" -+ -+ "github.com/stretchr/testify/assert" -+ "github.com/stretchr/testify/require" -+) -+ -+func TestDiagnoseGracefulStopRequested(t *testing.T) { -+ t.Parallel() -+ -+ t.Run("background", func(t *testing.T) { -+ t.Parallel() -+ assert.False(t, DiagnoseGracefulStopRequested(context.Background())) -+ }) -+ -+ t.Run("diagnose_context_before_signal", func(t *testing.T) { -+ t.Parallel() -+ ctx, stop := NewDiagnoseRunContext(context.Background()) -+ defer stop() -+ assert.False(t, DiagnoseGracefulStopRequested(ctx)) -+ }) -+ -+ t.Run("diagnose_context_after_graceful_request", func(t *testing.T) { -+ t.Parallel() -+ ctx, stop := NewDiagnoseRunContext(context.Background()) -+ defer stop() -+ RequestDiagnoseGracefulStop(ctx) -+ assert.True(t, DiagnoseGracefulStopRequested(ctx)) -+ }) -+} -+ -+func TestDiagnoseHardCancel(t *testing.T) { -+ t.Parallel() -+ ctx, stop := NewDiagnoseRunContext(context.Background()) -+ defer stop() -+ -+ RequestDiagnoseHardCancel(ctx) -+ require.Error(t, ctx.Err()) -+ assert.False(t, DiagnoseGracefulStopRequested(ctx)) -+} -+ -+func TestDiagnoseGracefulThenHardCancel(t *testing.T) { -+ t.Parallel() -+ ctx, stop := NewDiagnoseRunContext(context.Background()) -+ defer stop() -+ -+ RequestDiagnoseGracefulStop(ctx) -+ assert.True(t, DiagnoseGracefulStopRequested(ctx)) -+ require.NoError(t, ctx.Err()) -+ -+ RequestDiagnoseHardCancel(ctx) -+ require.Error(t, ctx.Err()) -+} -+ -+func TestDiagnoseRunContextStopTearsDown(t *testing.T) { -+ t.Parallel() -+ ctx, stop := NewDiagnoseRunContext(context.Background()) -+ stop() -+ require.Error(t, ctx.Err()) -+} - - - - - - "fmt" - "os" - "path/filepath" -+ "runtime" - "time" - - "charm.land/lipgloss/v2" - - - } - return filepath.Join(resultsDir, "report.csv") - } -+ -+func diagnoseInterruptKeyHint() string { -+ if runtime.GOOS == "darwin" { -+ return "⌘C" -+ } -+ return "Ctrl+C" -+} -+ -+func printDiagnoseGracefulStopNotice(out *output.Printer, completed, total int) { -+ if out == nil { -+ return -+ } -+ if out.AIOutput() { -+ out.Stderrf("stop_graceful completed=%d total=%d\n", completed, total) -+ return -+ } -+ out.ClearInline() -+ hint := diagnoseInterruptKeyHint() -+ out.HumanStderr( -+ termstyle.Accent.Render( -+ fmt.Sprintf("Stopping diagnose run after current iteration — %d/%d completed.", completed, total), -+ ) + "\n" + -+ termstyle.Muted.Render( -+ fmt.Sprintf("Press %s again to cancel immediately.", hint), -+ ), -+ ) -+} - - - - - - "strconv" - "strings" - "sync" -+ "sync/atomic" - "time" - - "github.com/smartcontractkit/testrig/internal/config" - - - failedFast bool - failedFastReason string - failedFastIteration int // 0-based diagnose iteration index; -1 if unset -+ GracefulStop bool - iterDurations []time.Duration - shuffleSeeds map[int]int64 - liveProgress bool - - - } - - func printDiagnoseInterruptions(ctx context.Context, conf *config.App, out *output.Printer, state *DiagnoseRunState) { -- interrupted := ctx.Err() != nil -- if interrupted && !out.AIOutput() { -- out.HumanStderr( -- termstyle.Accent.Render( -- fmt.Sprintf("interrupted after %d/%d iterations", state.completed, conf.Iterations), -- ) + -- termstyle.Muted.Render( -- " — analyzing partial results…", -- ), -- ) -+ if state.GracefulStop { -+ if !out.AIOutput() { -+ out.HumanStderr( -+ termstyle.Accent.Render( -+ fmt.Sprintf("stopped early after %d/%d iterations", state.completed, conf.Iterations), -+ ) + -+ termstyle.Muted.Render( -+ " — analyzing partial results…", -+ ), -+ ) -+ } -+ } else if interrupted := ctx.Err() != nil; interrupted { -+ if out.AIOutput() { -+ out.Stderrf("interrupted completed=%d total=%d\n", state.completed, conf.Iterations) -+ } else { -+ out.HumanStderr( -+ termstyle.Accent.Render( -+ fmt.Sprintf("interrupted after %d/%d iterations", state.completed, conf.Iterations), -+ ) + -+ termstyle.Muted.Render( -+ " — analyzing partial results…", -+ ), -+ ) -+ } - } - - if state.failedFast && !out.AIOutput() { - - - serialProgressMu = new(sync.Mutex) - } - -+ var inFlight atomic.Int32 - worker := diagnoseWorker{ - conf: conf, - out: out, - - - jobs: jobs, - results: results, - cancel: cancel, -+ inFlight: &inFlight, - } - var wg sync.WaitGroup - for _, resource := range resources { - - - - wg.Go(func() { - defer close(jobs) -+ gracefulStopCh := DiagnoseGracefulStopChan(ctx) - for i := range conf.Iterations { -+ if DiagnoseGracefulStopRequested(ctx) { -+ return -+ } -+ if gracefulStopCh == nil { -+ select { -+ case <-runCtx.Done(): -+ return -+ case jobs <- i: -+ } -+ continue -+ } - select { - case <-runCtx.Done(): - return -+ case <-gracefulStopCh: -+ return - case jobs <- i: - } - } - - - }() - - var firstErr error -+ var gracefulStopNotice sync.Once -+ maybePrintGracefulStopNotice := func() { -+ if !DiagnoseGracefulStopRequested(ctx) || inFlight.Load() > 0 { -+ return -+ } -+ gracefulStopNotice.Do(func() { -+ state.GracefulStop = true -+ printDiagnoseGracefulStopNotice(out, state.completed, conf.Iterations) -+ }) -+ } -+ - for result := range results { - if result.fatalErr != nil { - if firstErr == nil { - - - continue - } - recordDiagnoseIterationResult(state, result, conf, out, parallelProgress, serialProgressMu) -+ maybePrintGracefulStopNotice() - } -+ maybePrintGracefulStopNotice() - return state, firstErr - } - - - - jobs <-chan int - results chan<- diagnoseIterationResult - cancel context.CancelFunc -+ inFlight *atomic.Int32 - } - - func (w *diagnoseWorker) run(runCtx context.Context, resource diagnoseIterationResource) { - used := false - for iteration := range w.jobs { -+ stopInFlight := w.trackInFlight() - if runCtx.Err() != nil { -+ stopInFlight() - return - } - if used && resource.Reset != nil { - if err := resource.Reset(runCtx); err != nil { -+ stopInFlight() - if runCtx.Err() != nil { - return - } - - - used = true - if w.hooks.iterationSetup != nil { - if err := w.hooks.iterationSetup(runCtx); err != nil { -+ stopInFlight() - w.sendFatal(runCtx, iteration, fmt.Errorf("iteration setup %d: %w", iteration, err)) - return - } - - - digest: digest, - digestErr: digestErr, - } -+ stopInFlight() - select { - case w.results <- result: - case <-runCtx.Done(): - - - } - } - -+func (w *diagnoseWorker) trackInFlight() func() { -+ if w.inFlight == nil { -+ return func() {} -+ } -+ w.inFlight.Add(1) -+ return func() { w.inFlight.Add(-1) } -+} -+ - // sendFatal posts a fatal iteration result and cancels the run. The select - // guards against ctx cancellation racing the send. - func (w *diagnoseWorker) sendFatal(runCtx context.Context, iteration int, err error) { - - - - - - "sync" - "sync/atomic" - "testing" -+ "testing/synctest" - "time" - - "github.com/stretchr/testify/assert" - - - assert.True(t, setupBeforeTeardown.Load(), "setup must be called before teardown") - } - -+func writePassIterationJSONL(t *testing.T, resultsDir string, iteration int) { -+ t.Helper() -+ require.NoError( -+ t, -+ os.WriteFile( -+ filepath.Join(resultsDir, "iteration-"+strconv.Itoa(iteration)+".log.jsonl"), -+ []byte(`{"Action":"pass","Package":"p","Test":"T","Elapsed":0.01}`+"\n"), -+ 0600, -+ ), -+ ) -+} -+ -+func TestRunDiagnoseIterations_gracefulStop_finishesInFlight(t *testing.T) { -+ t.Parallel() -+ synctest.Test(t, func(t *testing.T) { -+ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) -+ defer stop() -+ -+ resultsDir := t.TempDir() -+ conf := &config.App{ -+ RepoRoot: makeTestRepoRoot(t), -+ AIOutput: true, -+ Iterations: 5, -+ } -+ out := output.NewForTest(true, io.Discard, io.Discard, false) -+ -+ var mu sync.Mutex -+ started := make(map[int]struct{}) -+ hooks := diagnoseRunHooks{ -+ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { -+ mu.Lock() -+ started[p.Iteration] = struct{}{} -+ mu.Unlock() -+ if p.Iteration == 0 { -+ RequestDiagnoseGracefulStop(runCtx) -+ time.Sleep(1 * time.Second) -+ } -+ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) -+ return nil -+ }, -+ } -+ -+ state, err := runDiagnoseIterations( -+ ctx, -+ conf, -+ out, -+ resultsDir, -+ []string{"./pkg"}, -+ []diagnoseIterationResource{{}}, -+ hooks, -+ ) -+ require.NoError(t, err) -+ assert.Equal(t, 1, state.completed) -+ assert.True(t, state.GracefulStop) -+ assert.Len(t, started, 1) -+ }) -+} -+ -+func TestRunDiagnoseIterations_gracefulStop_printsNoticeAfterDigest(t *testing.T) { -+ t.Parallel() -+ synctest.Test(t, func(t *testing.T) { -+ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) -+ defer stop() -+ -+ resultsDir := t.TempDir() -+ conf := &config.App{ -+ RepoRoot: makeTestRepoRoot(t), -+ AIOutput: false, -+ Iterations: 3, -+ } -+ var stderr strings.Builder -+ out := output.NewForTest(false, io.Discard, &stderr, false) -+ -+ hooks := diagnoseRunHooks{ -+ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { -+ if p.Iteration == 0 { -+ RequestDiagnoseGracefulStop(runCtx) -+ } -+ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) -+ return nil -+ }, -+ } -+ -+ state, err := runDiagnoseIterations( -+ ctx, -+ conf, -+ out, -+ resultsDir, -+ []string{"./pkg"}, -+ []diagnoseIterationResource{{}}, -+ hooks, -+ ) -+ require.NoError(t, err) -+ assert.True(t, state.GracefulStop) -+ assert.Contains(t, stderr.String(), "Stopping diagnose run after current iteration") -+ assert.Contains(t, stderr.String(), "again to cancel immediately") -+ }) -+} -+ -+func TestRunDiagnoseIterations_gracefulStop_idleBetweenIterations(t *testing.T) { -+ t.Parallel() -+ synctest.Test(t, func(t *testing.T) { -+ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) -+ defer stop() -+ -+ resultsDir := t.TempDir() -+ conf := &config.App{ -+ RepoRoot: makeTestRepoRoot(t), -+ AIOutput: true, -+ Iterations: 3, -+ } -+ out := output.NewForTest(true, io.Discard, io.Discard, false) -+ -+ iter0Done := make(chan struct{}) -+ iter0Release := make(chan struct{}) -+ hooks := diagnoseRunHooks{ -+ runIteration: func(_ context.Context, p diagnoseIterationParams) error { -+ if p.Iteration == 0 { -+ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) -+ close(iter0Done) -+ <-iter0Release -+ return nil -+ } -+ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) -+ return nil -+ }, -+ } -+ -+ var state *DiagnoseRunState -+ var runErr error -+ go func() { -+ state, runErr = runDiagnoseIterations( -+ ctx, -+ conf, -+ out, -+ resultsDir, -+ []string{"./pkg"}, -+ []diagnoseIterationResource{{}}, -+ hooks, -+ ) -+ }() -+ -+ <-iter0Done -+ RequestDiagnoseGracefulStop(ctx) -+ close(iter0Release) -+ synctest.Wait() -+ -+ require.NoError(t, runErr) -+ require.NotNil(t, state) -+ assert.Equal(t, 1, state.completed) -+ assert.True(t, state.GracefulStop) -+ }) -+} -+ -+func TestRunDiagnoseIterations_hardCancel_abortsInFlight(t *testing.T) { -+ t.Parallel() -+ synctest.Test(t, func(t *testing.T) { -+ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) -+ defer stop() -+ -+ resultsDir := t.TempDir() -+ conf := &config.App{ -+ RepoRoot: makeTestRepoRoot(t), -+ AIOutput: true, -+ Iterations: 3, -+ } -+ out := output.NewForTest(true, io.Discard, io.Discard, false) -+ -+ iter0Blocked := make(chan struct{}) -+ hooks := diagnoseRunHooks{ -+ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { -+ if p.Iteration == 0 { -+ close(iter0Blocked) -+ <-runCtx.Done() -+ return runCtx.Err() -+ } -+ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) -+ return nil -+ }, -+ } -+ -+ var state *DiagnoseRunState -+ var runErr error -+ go func() { -+ state, runErr = runDiagnoseIterations( -+ ctx, -+ conf, -+ out, -+ resultsDir, -+ []string{"./pkg"}, -+ []diagnoseIterationResource{{}}, -+ hooks, -+ ) -+ }() -+ -+ <-iter0Blocked -+ RequestDiagnoseHardCancel(ctx) -+ synctest.Wait() -+ -+ require.NoError(t, runErr) -+ require.NotNil(t, state) -+ assert.Less(t, state.completed, conf.Iterations) -+ assert.False(t, state.GracefulStop) -+ require.Error(t, ctx.Err()) -+ }) -+} -+ -+func TestDiagnoseGracefulStop_writesPartialReport(t *testing.T) { -+ t.Parallel() -+ synctest.Test(t, func(t *testing.T) { -+ repoRoot := makeTestRepoRoot(t) -+ conf := &config.App{ -+ RepoRoot: repoRoot, -+ AIOutput: true, -+ Iterations: 4, -+ } -+ ctx, stop := NewDiagnoseRunContextForTest(context.Background()) -+ defer stop() -+ out := output.NewForTest(conf.AIOutput, io.Discard, io.Discard, false) -+ -+ resultsDir := t.TempDir() -+ start := time.Now() -+ hooks := diagnoseRunHooks{ -+ runIteration: func(runCtx context.Context, p diagnoseIterationParams) error { -+ if p.Iteration == 1 { -+ RequestDiagnoseGracefulStop(runCtx) -+ } -+ writePassIterationJSONL(t, p.ResultsDir, p.Iteration) -+ return nil -+ }, -+ } -+ -+ state, err := runDiagnoseIterations( -+ ctx, -+ conf, -+ out, -+ resultsDir, -+ []string{"./pkg"}, -+ []diagnoseIterationResource{{}}, -+ hooks, -+ ) -+ require.NoError(t, err) -+ require.True(t, state.GracefulStop) -+ require.Equal(t, 2, state.completed) -+ -+ require.NoError(t, writeRunState(resultsDir, conf, []string{"./pkg"}, state, start)) -+ require.NoError(t, FinishDiagnoseAnalysis(ctx, conf, out, []string{"./pkg"}, state, start, resultsDir)) -+ -+ //nolint:gosec // G304: path from filepath.Join -+ reportBytes, err := os.ReadFile(filepath.Join(resultsDir, "report.json")) -+ require.NoError(t, err) -+ var rep Report -+ require.NoError(t, json.Unmarshal(reportBytes, &rep)) -+ assert.Equal(t, 2, rep.Iterations) -+ assert.Len(t, rep.IterationSummaries, 2) -+ }) -+} -+ - func makeTestRepoRoot(t *testing.T) string { - d := t.TempDir() - err := os.WriteFile(filepath.Join(d, "go.mod"), []byte("module test\n"), 0600) - - - - -