diff --git a/cmd/av/adopt.go b/cmd/av/adopt.go index 9053a612..f41f90a2 100644 --- a/cmd/av/adopt.go +++ b/cmd/av/adopt.go @@ -70,14 +70,16 @@ The command will adopt the stack of pull requests starting from the specified br if err != nil { return err } - return uiutils.RunBubbleTea(&remoteAdoptViewModel{ + return uiutils.RunBubbleTeaWithContext(ctx, &remoteAdoptViewModel{ + ctx: func() context.Context { return ctx }, repo: repo, db: db, ghClient: client, branchName: adoptFlags.RemoteBranchName, }) } - return uiutils.RunBubbleTea(&adoptViewModel{ + return uiutils.RunBubbleTeaWithContext(ctx, &adoptViewModel{ + ctx: func() context.Context { return ctx }, repo: repo, db: db, currentHEADBranch: plumbing.NewBranchReferenceName(currentBranch), @@ -143,6 +145,7 @@ func adoptForceAdoption( } type adoptViewModel struct { + ctx func() context.Context repo *git.Repo db meta.DB currentHEADBranch plumbing.ReferenceName @@ -162,6 +165,7 @@ func (vm *adoptViewModel) Init() tea.Cmd { func (vm *adoptViewModel) initModel() tea.Cmd { return vm.AddView( actions.NewFindAdoptableLocalBranchesModel( + vm.ctx(), vm.repo, vm.db, vm.initTreeSelector, @@ -247,6 +251,7 @@ func (vm *adoptViewModel) initAdoption(chosenTargets []plumbing.ReferenceName) t } return vm.AddView( actions.NewAdoptBranchesModel( + vm.ctx(), vm.db, branches, func() tea.Cmd { return tea.Quit }, @@ -262,6 +267,7 @@ func (vm *adoptViewModel) ExitError() error { } type remoteAdoptViewModel struct { + ctx func() context.Context repo *git.Repo db meta.DB ghClient *gh.Client @@ -281,6 +287,7 @@ func (vm *remoteAdoptViewModel) Init() tea.Cmd { func (vm *remoteAdoptViewModel) initModel() tea.Cmd { return vm.AddView( actions.NewGetRemoteStackedPRModel( + vm.ctx(), vm.db.ReadTx().Repository(), vm.ghClient, vm.branchName, @@ -369,7 +376,7 @@ func (vm *remoteAdoptViewModel) initGitFetch(prs []actions.RemotePRInfo, chosenT refspecs = append(refspecs, fmt.Sprintf("refs/heads/%s:refs/remotes/%s/%s", target.Short(), remote, target.Short())) } return vm.AddView( - actions.NewGitFetchModel(vm.repo, refspecs, func() tea.Cmd { + actions.NewGitFetchModel(vm.ctx(), vm.repo, refspecs, func() tea.Cmd { return vm.initAdoption(prs, chosenTargets) }), ) @@ -405,6 +412,7 @@ func (vm *remoteAdoptViewModel) initAdoption(prs []actions.RemotePRInfo, chosenT } return vm.AddView( actions.NewAdoptBranchesModel( + vm.ctx(), vm.db, branches, func() tea.Cmd { @@ -417,7 +425,7 @@ func (vm *remoteAdoptViewModel) initAdoption(prs []actions.RemotePRInfo, chosenT for _, ab := range branches { if !hasChild[ab.Name] { // Check out the leaf branch. This can fail if the workspace is dirty. In that case, just quietly exit. - _, _ = vm.repo.Git(context.Background(), "switch", ab.Name, "--quiet") + _, _ = vm.repo.Git(vm.ctx(), "switch", ab.Name, "--quiet") break } } diff --git a/cmd/av/main.go b/cmd/av/main.go index 70db5051..f27d1929 100644 --- a/cmd/av/main.go +++ b/cmd/av/main.go @@ -6,10 +6,12 @@ import ( "fmt" "os" "os/exec" + "os/signal" "path/filepath" "slices" "strings" "sync" + "syscall" "time" "emperror.dev/errors" @@ -136,10 +138,14 @@ func run() int { // runtime and various packages (e.g., package init functions). startTime := time.Now() colors.SetupBackgroundColorTypeFromEnv() + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() defer func() { _ = cachedRepo.Close() }() - err := rootCmd.Execute() + err := rootCmd.ExecuteContext(ctx) logrus.WithField("duration", time.Since(startTime)).Debug("command exited") - checkCliVersion() + if ctx.Err() == nil { + checkCliVersion() + } var exitSilently actions.ErrExitSilently if errors.As(err, &exitSilently) { return exitSilently.ExitCode diff --git a/e2e_tests/sigint_test.go b/e2e_tests/sigint_test.go new file mode 100644 index 00000000..feb64a03 --- /dev/null +++ b/e2e_tests/sigint_test.go @@ -0,0 +1,98 @@ +package e2e_tests + +// TestAdoptExitsOnSIGINT lives outside the testscript harness because +// testscript v1.10.0 has no public API for sending arbitrary signals to a +// specifically named background process - it only sends os.Interrupt to +// every backgrounded command at end-of-test, which is not granular enough +// to verify that `av adopt` returns promptly after a single SIGINT while +// the merge-base step is still running. +// +// Once the project upgrades to a testscript version with a public +// signal-by-name primitive (or grows a custom command for it in +// testscript_test.go), this can move into testdata/script as a .txtar. + +import ( + "bytes" + "context" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/aviator-co/av/internal/git/gittest" + "github.com/stretchr/testify/require" +) + +func TestAdoptExitsOnSIGINT(t *testing.T) { + if _, err := exec.LookPath("sh"); err != nil { + t.Skip("requires /bin/sh shim for the fake git binary") + } + + repo := gittest.NewTempRepo(t) + + repo.Git(t, "checkout", "-b", "stack-1") + repo.CommitFile(t, "my-file", "1a\n", gittest.WithMessage("Commit 1a")) + repo.Git(t, "switch", "main") + + realGit, err := exec.LookPath("git") + require.NoError(t, err) + + fakeGitDir := t.TempDir() + markerPath := filepath.Join(t.TempDir(), "merge-base-started") + fakeGitPath := filepath.Join(fakeGitDir, "git") + fakeGit := `#!/bin/sh +if [ "$1" = "merge-base" ]; then + if [ -n "$AV_TEST_GIT_MERGE_BASE_STARTED" ]; then + : > "$AV_TEST_GIT_MERGE_BASE_STARTED" + fi + exec sleep 30 +fi +exec "$REAL_GIT" "$@" +` + require.NoError(t, os.WriteFile(fakeGitPath, []byte(fakeGit), 0o755)) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + cmd := exec.CommandContext(ctx, avCmdPath, "--debug", "adopt", "--dry-run") + cmd.Dir = repo.RepoDir + cmd.Env = append( + os.Environ(), + "PATH="+fakeGitDir+string(os.PathListSeparator)+os.Getenv("PATH"), + "REAL_GIT="+realGit, + "AV_TEST_GIT_MERGE_BASE_STARTED="+markerPath, + ) + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + require.NoError(t, cmd.Start()) + waitForFile(t, markerPath, 2*time.Second) + + require.NoError(t, cmd.Process.Signal(os.Interrupt)) + done := make(chan error, 1) + go func() { + done <- cmd.Wait() + }() + + select { + case <-done: + // av returned within the deadline; that's the success path. The exit + // code is non-zero because the run was cancelled, which is fine. + case <-time.After(5 * time.Second): + _ = cmd.Process.Kill() + t.Fatalf("av adopt did not exit within 5s of SIGINT\nstdout:\n%s\nstderr:\n%s", stdout.String(), stderr.String()) + } +} + +func waitForFile(t *testing.T, path string, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if _, err := os.Stat(path); err == nil { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatalf("file %s did not appear within %s", path, timeout) +} diff --git a/internal/actions/adopt_branches.go b/internal/actions/adopt_branches.go index 0f788586..abe102e4 100644 --- a/internal/actions/adopt_branches.go +++ b/internal/actions/adopt_branches.go @@ -1,6 +1,7 @@ package actions import ( + "context" "strings" "charm.land/bubbles/v2/spinner" @@ -17,11 +18,13 @@ type AdoptingBranch struct { } func NewAdoptBranchesModel( + ctx context.Context, db meta.DB, branches []AdoptingBranch, onDone func() tea.Cmd, ) tea.Model { return &AdoptBranchesModel{ + ctx: func() context.Context { return ctx }, db: db, spinner: spinner.New(spinner.WithSpinner(spinner.Dot)), branches: branches, @@ -30,6 +33,7 @@ func NewAdoptBranchesModel( } type AdoptBranchesModel struct { + ctx func() context.Context db meta.DB spinner spinner.Model branches []AdoptingBranch @@ -68,6 +72,9 @@ func (m *AdoptBranchesModel) View() tea.View { func (m *AdoptBranchesModel) adoptBranches() tea.Msg { tx := m.db.WriteTx() for _, branch := range m.branches { + if err := m.ctx().Err(); err != nil { + return err + } bi, _ := tx.Branch(branch.Name) bi.Parent = branch.Parent bi.PullRequest = branch.PullRequest @@ -76,6 +83,9 @@ func (m *AdoptBranchesModel) adoptBranches() tea.Msg { } tx.SetBranch(bi) } + if err := m.ctx().Err(); err != nil { + return err + } if err := tx.Commit(); err != nil { return err } diff --git a/internal/actions/find_adoptable_local_branches.go b/internal/actions/find_adoptable_local_branches.go index d80c08fd..6d82b16b 100644 --- a/internal/actions/find_adoptable_local_branches.go +++ b/internal/actions/find_adoptable_local_branches.go @@ -15,6 +15,7 @@ import ( ) func NewFindAdoptableLocalBranchesModel( + ctx context.Context, repo *git.Repo, db meta.DB, onDone func( @@ -24,6 +25,7 @@ func NewFindAdoptableLocalBranchesModel( ) tea.Cmd, ) tea.Model { return &FindAdoptableLocalBranchesModel{ + ctx: func() context.Context { return ctx }, repo: repo, db: db, spinner: spinner.New(spinner.WithSpinner(spinner.Dot)), @@ -32,6 +34,7 @@ func NewFindAdoptableLocalBranchesModel( } type FindAdoptableLocalBranchesModel struct { + ctx func() context.Context repo *git.Repo db meta.DB spinner spinner.Model @@ -49,7 +52,7 @@ func (m *FindAdoptableLocalBranchesModel) Init() tea.Cmd { if err != nil { return err } - branches, err := treedetector.DetectBranches(context.Background(), m.repo, unmanagedBranches) + branches, err := treedetector.DetectBranches(m.ctx(), m.repo, unmanagedBranches) if err != nil { return err } diff --git a/internal/actions/get_remote_stacked_pr.go b/internal/actions/get_remote_stacked_pr.go index fc263184..59326cc2 100644 --- a/internal/actions/get_remote_stacked_pr.go +++ b/internal/actions/get_remote_stacked_pr.go @@ -11,16 +11,17 @@ import ( "github.com/aviator-co/av/internal/gh" "github.com/aviator-co/av/internal/meta" "github.com/aviator-co/av/internal/utils/colors" - "github.com/aviator-co/av/internal/utils/uiutils" ) func NewGetRemoteStackedPRModel( + ctx context.Context, repo meta.Repository, ghClient *gh.Client, initialBranchName string, onDone func([]RemotePRInfo) tea.Cmd, ) tea.Model { return &GetRemoteStackedPRModel{ + ctx: func() context.Context { return ctx }, repo: repo, ghClient: ghClient, spinner: spinner.New(spinner.WithSpinner(spinner.Dot)), @@ -38,6 +39,7 @@ type RemotePRInfo struct { } type GetRemoteStackedPRModel struct { + ctx func() context.Context repo meta.Repository ghClient *gh.Client spinner spinner.Model @@ -48,43 +50,60 @@ type GetRemoteStackedPRModel struct { prs []RemotePRInfo } +// remoteStackedPRFailedMsg signals that PR fetching failed. State is mutated +// in Update() so it stays on the Bubble Tea main thread; the Init() goroutine +// only emits messages. +type remoteStackedPRFailedMsg struct{ err error } + +// remoteStackedPRDoneMsg carries the fully-collected list, delivered to +// Update() once the goroutine finishes. onDone runs in Update() because it +// may change the view stack. +type remoteStackedPRDoneMsg struct { + prs []RemotePRInfo +} + func (m *GetRemoteStackedPRModel) Init() tea.Cmd { return tea.Batch(m.spinner.Tick, func() tea.Msg { - ctx := context.Background() + var collected []RemotePRInfo nextPRNumber := int64(0) for { + if err := m.ctx().Err(); err != nil { + return remoteStackedPRFailedMsg{err: err} + } var pr *gh.PullRequest if nextPRNumber == 0 { // Initial branch is searched based on the branch name. - page, err := m.ghClient.GetPullRequests(ctx, gh.GetPullRequestsInput{ + page, err := m.ghClient.GetPullRequests(m.ctx(), gh.GetPullRequestsInput{ Owner: m.repo.Owner, Repo: m.repo.Name, HeadRefName: m.initialBranchName, }) if err != nil { - m.failed = true - return err + return remoteStackedPRFailedMsg{err: err} } if len(page.PullRequests) == 0 { - m.failed = true - return errors.New("cannot find PR for branch " + m.initialBranchName) + return remoteStackedPRFailedMsg{ + err: errors.New("cannot find PR for branch " + m.initialBranchName), + } } if len(page.PullRequests) > 1 { - m.failed = true - return errors.New("multiple PRs found for branch " + m.initialBranchName) + return remoteStackedPRFailedMsg{ + err: errors.New("multiple PRs found for branch " + m.initialBranchName), + } } pr = &page.PullRequests[0] } else { // Otherwise, we can fetch based on the PR number. var err error - pr, err = m.ghClient.GetPullRequestByNumber(ctx, gh.GetPullRequestByNumberInput{ + pr, err = m.ghClient.GetPullRequestByNumber(m.ctx(), gh.GetPullRequestByNumberInput{ Owner: m.repo.Owner, Repo: m.repo.Name, Number: nextPRNumber, }) if err != nil { - m.failed = true - return errors.Wrapf(err, "failed to get PR %d", nextPRNumber) + return remoteStackedPRFailedMsg{ + err: errors.Wrapf(err, "failed to get PR %d", nextPRNumber), + } } } prMeta, err := ReadPRMetadata(pr.Body) @@ -96,8 +115,9 @@ func (m *GetRemoteStackedPRModel) Init() tea.Cmd { ParentPull: 0, } } else if err != nil { - m.failed = true - return errors.Wrapf(err, "failed to read metadata for PR %d", pr.Number) + return remoteStackedPRFailedMsg{ + err: errors.Wrapf(err, "failed to read metadata for PR %d", pr.Number), + } } remotePRInfo := RemotePRInfo{ Name: strings.TrimPrefix(pr.HeadRefName, "refs/heads/"), @@ -115,14 +135,13 @@ func (m *GetRemoteStackedPRModel) Init() tea.Cmd { MergeCommit: pr.GetMergeCommit(), Title: pr.Title, } - m.prs = append(m.prs, remotePRInfo) + collected = append(collected, remotePRInfo) if remotePRInfo.Parent.Trunk { break } nextPRNumber = prMeta.ParentPull } - m.done = true - return uiutils.SimpleCommandMsg{Cmd: m.onDone(m.prs)} + return remoteStackedPRDoneMsg{prs: collected} }) } @@ -132,6 +151,13 @@ func (m *GetRemoteStackedPRModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { var cmd tea.Cmd m.spinner, cmd = m.spinner.Update(msg) return m, cmd + case remoteStackedPRFailedMsg: + m.failed = true + return m, func() tea.Msg { return msg.err } + case remoteStackedPRDoneMsg: + m.prs = msg.prs + m.done = true + return m, m.onDone(msg.prs) } return m, nil } diff --git a/internal/actions/git_fetch.go b/internal/actions/git_fetch.go index 3ba4495a..26689f6c 100644 --- a/internal/actions/git_fetch.go +++ b/internal/actions/git_fetch.go @@ -15,6 +15,7 @@ import ( ) type GitFetchModel struct { + ctx func() context.Context repo *git.Repo spinner spinner.Model refspecs []string @@ -27,11 +28,13 @@ type GitFetchModel struct { } func NewGitFetchModel( + ctx context.Context, repo *git.Repo, refspecs []string, onDone func() tea.Cmd, ) tea.Model { return &GitFetchModel{ + ctx: func() context.Context { return ctx }, repo: repo, spinner: spinner.New(spinner.WithSpinner(spinner.Dot)), refspecs: refspecs, @@ -80,7 +83,7 @@ func (m *GitFetchModel) fetchBranches() tea.Msg { m.repo.GetRemoteName(), } args = append(args, m.refspecs...) - cmd := exec.CommandContext(context.Background(), "git", args...) + cmd := exec.CommandContext(m.ctx(), "git", args...) cmd.Stdout = &m.stdout cmd.Stderr = &m.stderr if err := cmd.Run(); err != nil { diff --git a/internal/treedetector/detector.go b/internal/treedetector/detector.go index 041fdc1e..cef01700 100644 --- a/internal/treedetector/detector.go +++ b/internal/treedetector/detector.go @@ -39,6 +39,9 @@ func DetectBranches( ret := map[plumbing.ReferenceName]*BranchPiece{} for _, bn := range unmanagedBranches { + if err := ctx.Err(); err != nil { + return nil, err + } currentHash := refToHashMap[bn] nearestTrunkCommit, err := getNearestTrunkCommit(ctx, repo, bn) if err != nil { @@ -49,7 +52,7 @@ func DetectBranches( // don't have to adopt it. continue } - bp, err := traverseUntilTrunk(repo, bn, nearestTrunkCommit, hashToRefMap, refToHashMap) + bp, err := traverseUntilTrunk(ctx, repo, bn, nearestTrunkCommit, hashToRefMap, refToHashMap) if err != nil { return nil, err } @@ -59,6 +62,7 @@ func DetectBranches( } func traverseUntilTrunk( + ctx context.Context, repo *avgit.Repo, branch plumbing.ReferenceName, nearestTrunkCommit plumbing.Hash, @@ -75,6 +79,9 @@ func traverseUntilTrunk( // Do a commit traversal. We can stop the traversal when we hit the trunk or if we find a // commit that has multiple parents. err = object.NewCommitPreorderIter(commit, nil, nil).ForEach(func(c *object.Commit) error { + if err := ctx.Err(); err != nil { + return err + } if c.Hash == nearestTrunkCommit { trunk := repo.DefaultBranch() ret.Parent = plumbing.NewBranchReferenceName(trunk) diff --git a/internal/utils/uiutils/tea.go b/internal/utils/uiutils/tea.go index 0b6de9cd..9b765dd8 100644 --- a/internal/utils/uiutils/tea.go +++ b/internal/utils/uiutils/tea.go @@ -1,6 +1,7 @@ package uiutils import ( + "context" "os" "strings" @@ -19,6 +20,14 @@ type BubbleTeaModelWithExitHandling interface { } func RunBubbleTea(model BubbleTeaModelWithExitHandling) error { + return runBubbleTea(model) +} + +func RunBubbleTeaWithContext(ctx context.Context, model BubbleTeaModelWithExitHandling) error { + return runBubbleTea(model, tea.WithContext(ctx)) +} + +func runBubbleTea(model BubbleTeaModelWithExitHandling, extra ...tea.ProgramOption) error { // Handle input and output independently: the renderer writes to stdout, while // input is read from stdin. stdinIsTTY := isatty.IsTerminal(os.Stdin.Fd()) @@ -36,6 +45,7 @@ func RunBubbleTea(model BubbleTeaModelWithExitHandling) error { // output stays clean and complete. opts = append(opts, tea.WithoutRenderer()) } + opts = append(opts, extra...) p := tea.NewProgram(model, opts...) finalModel, err := p.Run() if err != nil {