diff --git a/pkg/gui/controllers.go b/pkg/gui/controllers.go index d8ce6766af6..44e97811bf4 100644 --- a/pkg/gui/controllers.go +++ b/pkg/gui/controllers.go @@ -148,6 +148,7 @@ func (gui *Gui) resetHelpersAndControllers() { syncController := controllers.NewSyncController( common, + &gui.branchesBeingPushed, ) submodulesController := controllers.NewSubmodulesController(common) @@ -182,7 +183,7 @@ func (gui *Gui) resetHelpersAndControllers() { renameSimilarityThresholdController := controllers.NewRenameSimilarityThresholdController(common) verticalScrollControllerFactory := controllers.NewVerticalScrollControllerFactory(common, &gui.viewBufferManagerMap) - branchesController := controllers.NewBranchesController(common) + branchesController := controllers.NewBranchesController(common, &gui.branchesBeingPushed) gitFlowController := controllers.NewGitFlowController(common) stashController := controllers.NewStashController(common) commitFilesController := controllers.NewCommitFilesController(common) diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go index 244092681b0..4cb9d63be10 100644 --- a/pkg/gui/controllers/branches_controller.go +++ b/pkg/gui/controllers/branches_controller.go @@ -3,6 +3,7 @@ package controllers import ( "errors" "fmt" + "sync" "github.com/jesseduffield/gocui" "github.com/jesseduffield/lazygit/pkg/commands/git_commands" @@ -12,18 +13,21 @@ import ( "github.com/jesseduffield/lazygit/pkg/gui/types" "github.com/jesseduffield/lazygit/pkg/utils" "github.com/samber/lo" + "github.com/sasha-s/go-deadlock" ) type BranchesController struct { baseController *ListControllerTrait[*models.Branch] - c *ControllerCommon + c *ControllerCommon + branchesBeingPushed *sync.Map } var _ types.IController = &BranchesController{} func NewBranchesController( c *ControllerCommon, + branchesBeingPushed *sync.Map, ) *BranchesController { return &BranchesController{ baseController: baseController{}, @@ -34,6 +38,7 @@ func NewBranchesController( c.Contexts().Branches.GetSelected, c.Contexts().Branches.GetSelectedItems, ), + branchesBeingPushed: branchesBeingPushed, } } @@ -421,11 +426,58 @@ func (self *BranchesController) promptToCheckoutWorktree(worktree *models.Worktr return nil } +func (self *BranchesController) blockForBranchFinishPush(branch *models.Branch) *models.Branch { + if val, ok := self.branchesBeingPushed.Load(branch.Name); ok { + // We only store waitgroups in here, so there isn't much thought into the false case + if wg, ok := val.(*deadlock.WaitGroup); ok { + wg.Wait() + // When we refresh after a branch push, the entire branch list gets replaced, so we must re-retrieve the + // branch pointer. The currently selected item might have changed since we initially requested the pull request, + // but the commit hash should continue to be the same. + updatedBranch, found := lo.Find(self.context().ListViewModel.GetItems(), func(b *models.Branch) bool { + return b.CommitHash == branch.CommitHash + }) + // If it _isn't_ found, something wack is going on, so lets stick with the original + if found { + branch = updatedBranch + } + } + } + return branch +} + func (self *BranchesController) handleCreatePullRequest(selectedBranch *models.Branch) error { - if !selectedBranch.IsTrackingRemote() { - return errors.New(self.c.Tr.PullRequestNoUpstream) + message := utils.ResolvePlaceholderString( + self.c.Tr.PendingCreatePullRequest, map[string]string{ + "branchName": selectedBranch.Name, + }) + err := self.c.WithPendingMessage(message, + func(_ gocui.Task) error { + _ = self.blockForBranchFinishPush(selectedBranch) + return nil + }, + func(_ gocui.Task) error { + // When we refresh after a branch push, the entire branch list gets replaced, so we must re-retrieve the + // branch pointer. The currently selected item might have changed since we initially requested the pull request, + // but the commit hash should continue to be the same. + updatedBranch, found := lo.Find(self.context().ListViewModel.GetItems(), func(b *models.Branch) bool { + return b.CommitHash == selectedBranch.CommitHash + }) + // If it _isn't_ found, something wack is going on, so lets stick with the original + if found { + selectedBranch = updatedBranch + } + + if !selectedBranch.IsTrackingRemote() { + return errors.New(self.c.Tr.PullRequestNoUpstream) + } + + return self.createPullRequest(selectedBranch.UpstreamBranch, "") + }) + if err != nil { + return err } - return self.createPullRequest(selectedBranch.UpstreamBranch, "") + return nil } func (self *BranchesController) handleCreatePullRequestMenu(selectedBranch *models.Branch) error { diff --git a/pkg/gui/controllers/helpers/app_status_helper.go b/pkg/gui/controllers/helpers/app_status_helper.go index 9375186326a..54e48e54245 100644 --- a/pkg/gui/controllers/helpers/app_status_helper.go +++ b/pkg/gui/controllers/helpers/app_status_helper.go @@ -1,6 +1,7 @@ package helpers import ( + "context" "time" "github.com/jesseduffield/gocui" @@ -64,6 +65,18 @@ func (self *AppStatusHelper) WithWaitingStatus(message string, f func(gocui.Task }) } +// WithPendingMessage displays message and begins executing pending. +// After pending finishes, via completion or cancellation, the message will be removed. +// If pending finishes with success, main will begin executing. +func (self *AppStatusHelper) WithPendingMessage( + message string, + pending func(gocui.Task) error, + main func(gocui.Task) error, +) { + pendingTask := self.c.OnWorkerPending(context.Background(), pending, main) + self.WithWaitingStatus(message, func(_ gocui.Task) error { <-pendingTask.Ctx.Done(); return nil }) +} + func (self *AppStatusHelper) WithWaitingStatusImpl(message string, f func(gocui.Task) error, task gocui.Task) error { return self.statusMgr().WithWaitingStatus(message, self.renderAppStatus, func(waitingStatusHandle *status.WaitingStatusHandle) error { return f(appStatusHelperTask{task, waitingStatusHandle}) diff --git a/pkg/gui/controllers/helpers/refresh_helper.go b/pkg/gui/controllers/helpers/refresh_helper.go index cac9310d161..65ba216e87d 100644 --- a/pkg/gui/controllers/helpers/refresh_helper.go +++ b/pkg/gui/controllers/helpers/refresh_helper.go @@ -53,10 +53,6 @@ func NewRefreshHelper( } func (self *RefreshHelper) Refresh(options types.RefreshOptions) error { - if options.Mode == types.ASYNC && options.Then != nil { - panic("RefreshOptions.Then doesn't work with mode ASYNC") - } - t := time.Now() defer func() { self.c.Log.Infof("Refresh took %s", time.Since(t)) @@ -103,7 +99,9 @@ func (self *RefreshHelper) Refresh(options types.RefreshOptions) error { // everything happens fast and it's better to have everything update // in the one frame if !self.c.InDemo() && options.Mode == types.ASYNC { + wg.Add(1) self.c.OnWorker(func(t gocui.Task) error { + defer wg.Done() f() return nil }) @@ -189,11 +187,20 @@ func (self *RefreshHelper) Refresh(options types.RefreshOptions) error { self.refreshStatus() - wg.Wait() - - if options.Then != nil { - if err := options.Then(); err != nil { - return err + if options.Mode == types.ASYNC { + self.c.OnWorker(func(t gocui.Task) error { + wg.Wait() + if options.Then != nil { + return options.Then() + } + return nil + }) + } else { + wg.Wait() + if options.Then != nil { + if err := options.Then(); err != nil { + return err + } } } diff --git a/pkg/gui/controllers/quit_actions.go b/pkg/gui/controllers/quit_actions.go index 000fe1792e9..fd8a84d8ef9 100644 --- a/pkg/gui/controllers/quit_actions.go +++ b/pkg/gui/controllers/quit_actions.go @@ -53,6 +53,13 @@ func (self *QuitActions) confirmQuitDuringUpdate() error { } func (self *QuitActions) Escape() error { + pendingTasks := self.c.State().GetPendingTasks() + if len(pendingTasks) > 0 { + head := pendingTasks[0] + head.CancelFunc() + return nil + } + currentContext := self.c.Context().Current() if listContext, ok := currentContext.(types.IListContext); ok { diff --git a/pkg/gui/controllers/sync_controller.go b/pkg/gui/controllers/sync_controller.go index 66c480b93c4..0e6565924e2 100644 --- a/pkg/gui/controllers/sync_controller.go +++ b/pkg/gui/controllers/sync_controller.go @@ -4,6 +4,7 @@ import ( "errors" "fmt" "strings" + "sync" "github.com/jesseduffield/gocui" "github.com/jesseduffield/lazygit/pkg/commands/git_commands" @@ -11,21 +12,25 @@ import ( "github.com/jesseduffield/lazygit/pkg/gui/context" "github.com/jesseduffield/lazygit/pkg/gui/types" "github.com/jesseduffield/lazygit/pkg/utils" + "github.com/sasha-s/go-deadlock" ) type SyncController struct { baseController - c *ControllerCommon + c *ControllerCommon + branchesBeingPushed *sync.Map } var _ types.IController = &SyncController{} func NewSyncController( common *ControllerCommon, + branchesBeingPushed *sync.Map, ) *SyncController { return &SyncController{ - baseController: baseController{}, - c: common, + baseController: baseController{}, + c: common, + branchesBeingPushed: branchesBeingPushed, } } @@ -194,6 +199,12 @@ type pushOpts struct { func (self *SyncController) pushAux(currentBranch *models.Branch, opts pushOpts) error { return self.c.WithInlineStatus(currentBranch, types.ItemOperationPushing, context.LOCAL_BRANCHES_CONTEXT_KEY, func(task gocui.Task) error { + val, _ := self.branchesBeingPushed.LoadOrStore(currentBranch.Name, &deadlock.WaitGroup{}) + wg, ok := val.(*deadlock.WaitGroup) + wg.Add(1) + if !ok { + self.c.Log.Fatalf("Somehow the WaitGroup we put in, didn't come back out as a WaitGroup! Received %s instead of type %t", val, val) + } self.c.LogAction(self.c.Tr.Actions.Push) err := self.c.Git().Sync.Push( task, @@ -229,7 +240,10 @@ func (self *SyncController) pushAux(currentBranch *models.Branch, opts pushOpts) } return err } - return self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC}) + return self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC, Then: func() error { + wg.Done() + return nil + }}) }) } diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go index cf62011cacc..50afbf1f056 100644 --- a/pkg/gui/gui.go +++ b/pkg/gui/gui.go @@ -95,6 +95,8 @@ type Gui struct { // so that you can return to the superproject RepoPathStack *utils.StringStack + branchesBeingPushed sync.Map // keys: branchName string, value: deadlock.WaitGroup + // this tells us whether our views have been initially set up ViewsSetup bool @@ -156,6 +158,10 @@ func (self *StateAccessor) GetRepoPathStack() *utils.StringStack { return self.gui.RepoPathStack } +func (self *StateAccessor) GetPendingTasks() []*gocui.PendingTask { + return self.gui.g.PendingTasks() +} + func (self *StateAccessor) GetUpdating() bool { return self.gui.Updating } @@ -690,6 +696,9 @@ func NewGui( func() types.Context { return gui.State.ContextMgr.Current() }, gui.createMenu, func(message string, f func(gocui.Task) error) { gui.helpers.AppStatus.WithWaitingStatus(message, f) }, + func(message string, pending func(gocui.Task) error, f func(gocui.Task) error) { + gui.helpers.AppStatus.WithPendingMessage(message, pending, f) + }, func(message string, f func() error) error { return gui.helpers.AppStatus.WithWaitingStatusSync(message, f) }, @@ -1118,6 +1127,10 @@ func (gui *Gui) onWorker(f func(gocui.Task) error) { gui.g.OnWorker(f) } +func (gui *Gui) onWorkerPending(ctx goContext.Context, pending func(gocui.Task) error, main func(gocui.Task) error) *gocui.PendingTask { + return gui.g.OnWorkerPending(ctx, pending, main) +} + func (gui *Gui) getWindowDimensions(informationStr string, appStatus string) map[string]boxlayout.Dimensions { return gui.helpers.WindowArrangement.GetWindowDimensions(informationStr, appStatus) } diff --git a/pkg/gui/gui_common.go b/pkg/gui/gui_common.go index dfa4be52f3d..1a52cc44c88 100644 --- a/pkg/gui/gui_common.go +++ b/pkg/gui/gui_common.go @@ -1,6 +1,8 @@ package gui import ( + "context" + "github.com/jesseduffield/gocui" "github.com/jesseduffield/lazygit/pkg/commands" "github.com/jesseduffield/lazygit/pkg/commands/oscommands" @@ -115,6 +117,10 @@ func (self *guiCommon) OnWorker(f func(gocui.Task) error) { self.gui.onWorker(f) } +func (self *guiCommon) OnWorkerPending(ctx context.Context, pending func(gocui.Task) error, main func(gocui.Task) error) *gocui.PendingTask { + return self.gui.onWorkerPending(ctx, pending, main) +} + func (self *guiCommon) RenderToMainViews(opts types.RefreshMainOpts) { self.gui.refreshMainViews(opts) } diff --git a/pkg/gui/popup/popup_handler.go b/pkg/gui/popup/popup_handler.go index 8ea50de2b9b..06677a8b056 100644 --- a/pkg/gui/popup/popup_handler.go +++ b/pkg/gui/popup/popup_handler.go @@ -18,6 +18,7 @@ type PopupHandler struct { currentContextFn func() types.Context createMenuFn func(types.CreateMenuOptions) error withWaitingStatusFn func(message string, f func(gocui.Task) error) + withPendingMessageFn func(message string, pending func(gocui.Task) error, f func(gocui.Task) error) withWaitingStatusSyncFn func(message string, f func() error) error toastFn func(message string, kind types.ToastKind) getPromptInputFn func() string @@ -34,6 +35,7 @@ func NewPopupHandler( currentContextFn func() types.Context, createMenuFn func(types.CreateMenuOptions) error, withWaitingStatusFn func(message string, f func(gocui.Task) error), + withPendingMessageFn func(message string, pending func(gocui.Task) error, f func(gocui.Task) error), withWaitingStatusSyncFn func(message string, f func() error) error, toastFn func(message string, kind types.ToastKind), getPromptInputFn func() string, @@ -47,6 +49,7 @@ func NewPopupHandler( currentContextFn: currentContextFn, createMenuFn: createMenuFn, withWaitingStatusFn: withWaitingStatusFn, + withPendingMessageFn: withPendingMessageFn, withWaitingStatusSyncFn: withWaitingStatusSyncFn, toastFn: toastFn, getPromptInputFn: getPromptInputFn, @@ -75,6 +78,11 @@ func (self *PopupHandler) WithWaitingStatus(message string, f func(gocui.Task) e return nil } +func (self *PopupHandler) WithPendingMessage(message string, pending func(gocui.Task) error, f func(gocui.Task) error) error { + self.withPendingMessageFn(message, pending, f) + return nil +} + func (self *PopupHandler) WithWaitingStatusSync(message string, f func() error) error { return self.withWaitingStatusSyncFn(message, f) } diff --git a/pkg/gui/types/common.go b/pkg/gui/types/common.go index 82a309fb113..b8d8abf0dbe 100644 --- a/pkg/gui/types/common.go +++ b/pkg/gui/types/common.go @@ -1,6 +1,8 @@ package types import ( + "context" + "github.com/jesseduffield/gocui" "github.com/jesseduffield/lazygit/pkg/commands" "github.com/jesseduffield/lazygit/pkg/commands/git_commands" @@ -68,6 +70,7 @@ type IGuiCommon interface { // Runs a function in a goroutine. Use this whenever you want to run a goroutine and keep track of the fact // that lazygit is still busy. See docs/dev/Busy.md OnWorker(f func(gocui.Task) error) + OnWorkerPending(ctx context.Context, pending func(gocui.Task) error, main func(gocui.Task) error) *gocui.PendingTask // Function to call at the end of our 'layout' function which renders views // For example, you may want a view's line to be focused only after that view is // resized, if in accordion mode. @@ -126,6 +129,7 @@ type IPopupHandler interface { // Shows a popup prompting the user for input. Prompt(opts PromptOpts) WithWaitingStatus(message string, f func(gocui.Task) error) error + WithPendingMessage(message string, pending func(gocui.Task) error, f func(gocui.Task) error) error WithWaitingStatusSync(message string, f func() error) error Menu(opts CreateMenuOptions) error Toast(message string) @@ -353,6 +357,7 @@ type IStateAccessor interface { GetItemOperation(item HasUrn) ItemOperation SetItemOperation(item HasUrn, operation ItemOperation) ClearItemOperation(item HasUrn) + GetPendingTasks() []*gocui.PendingTask } type IRepoStateAccessor interface { diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go index 81620e47b24..c4467c31a48 100644 --- a/pkg/i18n/english.go +++ b/pkg/i18n/english.go @@ -724,6 +724,7 @@ type TranslationSet struct { SelectTargetRemote string NoValidRemoteName string CreatePullRequest string + PendingCreatePullRequest string SelectConfigFile string NoConfigFileFoundErr string LoadingFileSuggestions string @@ -1292,6 +1293,7 @@ func EnglishTranslationSet() *TranslationSet { AllBranchesLogGraph: `Show/cycle all branch logs`, UnsupportedGitService: `Unsupported git service`, CreatePullRequest: `Create pull request`, + PendingCreatePullRequest: `Waiting to create pull request for {{.branchName}}`, CopyPullRequestURL: `Copy pull request URL to clipboard`, NoBranchOnRemote: `This branch doesn't exist on remote. You need to push it to remote first.`, Fetch: `Fetch`, diff --git a/vendor/github.com/jesseduffield/gocui/gui.go b/vendor/github.com/jesseduffield/gocui/gui.go index 03d55912a3f..9c466da49ea 100644 --- a/vendor/github.com/jesseduffield/gocui/gui.go +++ b/vendor/github.com/jesseduffield/gocui/gui.go @@ -260,6 +260,15 @@ func (g *Gui) NewTask() *TaskImpl { return g.taskManager.NewTask() } +// All pending tasks that are currently waiting to complete or be cancelled +func (g *Gui) PendingTasks() []*PendingTask { + return g.taskManager.GetPendingTasks() +} + +func (g *Gui) NewPendingTask(ctx context.Context) *PendingTask { + return g.taskManager.NewPendingTask(ctx) +} + // An idle listener listens for when the program is idle. This is useful for // integration tests which can wait for the program to be idle before taking // the next step in the test. @@ -692,6 +701,50 @@ func (g *Gui) OnWorker(f func(Task) error) { }() } +// A wrapper around [OnWorker] that allows the scheduling of a main function of work +// that is dependent on the pending function completing. +// If pending returns an error, the main function will not be executed. +// +// If the provided ctx, or the CancelFunc on the returned PendingTask is cancelled prior +// to pending completing, main will not be executed. +// +// The created PendingTask is returned. +// A slice of all PendingTasks is also available via [Gui.PendingTasks]. +func (g *Gui) OnWorkerPending(ctx context.Context, pending func(Task) error, main func(Task) error) *PendingTask { + pendingTask := g.NewPendingTask(ctx) + beginMain := make(chan struct{}) + + // Pending task + g.OnWorker(func(t Task) error { + err := pending(t) + closeReason := pendingTask.Ctx.Err() + if err != nil && closeReason == nil { + // We only want to return our error if the user has not cancelled the task + if closeReason != nil { + return nil + } + pendingTask.CancelFunc() + return err + } + close(beginMain) + return nil + }) + + // Main task + g.OnWorker(func(_ Task) error { + select { + case <-beginMain: + g.OnWorker(main) + pendingTask.CancelFunc() + return nil + case <-pendingTask.Ctx.Done(): + return nil + } + }) + + return pendingTask +} + func (g *Gui) onWorkerAux(f func(Task) error, task Task) { panicking := true defer func() { diff --git a/vendor/github.com/jesseduffield/gocui/pending_task.go b/vendor/github.com/jesseduffield/gocui/pending_task.go new file mode 100644 index 00000000000..bb8181e4ba4 --- /dev/null +++ b/vendor/github.com/jesseduffield/gocui/pending_task.go @@ -0,0 +1,9 @@ +package gocui + +import "context" + +type PendingTask struct { + id int + Ctx context.Context + CancelFunc context.CancelFunc +} diff --git a/vendor/github.com/jesseduffield/gocui/task_manager.go b/vendor/github.com/jesseduffield/gocui/task_manager.go index e3c82b4d4c4..260478a96ab 100644 --- a/vendor/github.com/jesseduffield/gocui/task_manager.go +++ b/vendor/github.com/jesseduffield/gocui/task_manager.go @@ -1,6 +1,9 @@ package gocui -import "sync" +import ( + "context" + "sync" +) // Tracks whether the program is busy (i.e. either something is happening on // the main goroutine or a worker goroutine). Used by integration tests @@ -10,7 +13,8 @@ type TaskManager struct { idleListeners []chan struct{} tasks map[int]Task // auto-incrementing id for new tasks - nextId int + nextId int + pendingTasks []*PendingTask mutex sync.Mutex } @@ -22,6 +26,10 @@ func newTaskManager() *TaskManager { } } +func (self *TaskManager) GetPendingTasks() []*PendingTask { + return self.pendingTasks +} + func (self *TaskManager) NewTask() *TaskImpl { self.mutex.Lock() defer self.mutex.Unlock() @@ -36,6 +44,27 @@ func (self *TaskManager) NewTask() *TaskImpl { return task } +func (self *TaskManager) NewPendingTask(ctx context.Context) *PendingTask { + self.mutex.Lock() + defer self.mutex.Unlock() + ctx, cancelFunc := context.WithCancel(ctx) + + self.nextId++ + taskId := self.nextId + go func() { + <-ctx.Done() + self.deletePendingTask(taskId) + }() + + pendingTask := PendingTask{ + id: taskId, + Ctx: ctx, + CancelFunc: cancelFunc, + } + self.pendingTasks = append(self.pendingTasks, &pendingTask) + return &pendingTask +} + func (self *TaskManager) addIdleListener(c chan struct{}) { self.idleListeners = append(self.idleListeners, c) } @@ -65,3 +94,15 @@ func (self *TaskManager) delete(taskId int) { delete(self.tasks, taskId) }) } + +func (self *TaskManager) deletePendingTask(pendingTaskId int) { + self.withMutex(func() { + pendingTasks := make([]*PendingTask, 0, len(self.pendingTasks)-1) + for _, task := range self.pendingTasks { + if task.id != pendingTaskId { + pendingTasks = append(pendingTasks, task) + } + } + self.pendingTasks = pendingTasks + }) +}