From b5384d6cdb7d3b34093c57fbf47b462219f281a0 Mon Sep 17 00:00:00 2001
From: Jesse Duffield <jessedduffield@gmail.com>
Date: Wed, 11 May 2022 00:34:20 +1000
Subject: [PATCH] retry upon losing connection to docker

---
 pkg/commands/docker.go        | 53 ++++++++++++++++++++---------------
 pkg/gui/confirmation_panel.go | 16 +++++------
 pkg/gui/gui.go                | 34 ++++++++++++++++++++--
 3 files changed, 70 insertions(+), 33 deletions(-)

diff --git a/pkg/commands/docker.go b/pkg/commands/docker.go
index a44ba65a..d1383ca0 100644
--- a/pkg/commands/docker.go
+++ b/pkg/commands/docker.go
@@ -130,37 +130,44 @@ func (c *DockerCommand) MonitorContainerStats() {
 
 // MonitorCLIContainerStats monitors a stream of container stats and updates the containers as each new stats object is received
 func (c *DockerCommand) MonitorCLIContainerStats() {
-	command := `docker stats --all --no-trunc --format '{{json .}}'`
-	cmd := c.OSCommand.RunCustomCommand(command)
-
-	r, err := cmd.StdoutPipe()
-	if err != nil {
+	onError := func(err error) {
 		c.ErrorChan <- err
-		return
+		time.Sleep(2 * time.Second)
 	}
 
-	_ = cmd.Start()
+	for {
+		command := `docker stats --all --no-trunc --format '{{json .}}'`
+		cmd := c.OSCommand.RunCustomCommand(command)
 
-	scanner := bufio.NewScanner(r)
-	scanner.Split(bufio.ScanLines)
-	for scanner.Scan() {
-		var stats ContainerCliStat
-		// need to strip ANSI codes because uses escape sequences to clear the screen with each refresh
-		cleanString := stripansi.Strip(scanner.Text())
-		if err := json.Unmarshal([]byte(cleanString), &stats); err != nil {
-			c.ErrorChan <- err
-			return
+		r, err := cmd.StdoutPipe()
+		if err != nil {
+			onError(err)
+			continue
 		}
-		c.ContainerMutex.Lock()
-		for _, container := range c.Containers {
-			if container.ID == stats.ID {
-				container.CLIStats = stats
+
+		_ = cmd.Start()
+
+		scanner := bufio.NewScanner(r)
+		scanner.Split(bufio.ScanLines)
+		for scanner.Scan() {
+			var stats ContainerCliStat
+			// need to strip ANSI codes because uses escape sequences to clear the screen with each refresh
+			cleanString := stripansi.Strip(scanner.Text())
+			if err := json.Unmarshal([]byte(cleanString), &stats); err != nil {
+				onError(err)
+				continue
 			}
+			c.ContainerMutex.Lock()
+			for _, container := range c.Containers {
+				if container.ID == stats.ID {
+					container.CLIStats = stats
+				}
+			}
+			c.ContainerMutex.Unlock()
 		}
-		c.ContainerMutex.Unlock()
-	}
 
-	_ = cmd.Wait()
+		_ = cmd.Wait()
+	}
 }
 
 // MonitorClientContainerStats is a function
diff --git a/pkg/gui/confirmation_panel.go b/pkg/gui/confirmation_panel.go
index f38affda..9a8a5504 100644
--- a/pkg/gui/confirmation_panel.go
+++ b/pkg/gui/confirmation_panel.go
@@ -20,20 +20,20 @@ func (gui *Gui) wrappedConfirmationFunction(function func(*gocui.Gui, *gocui.Vie
 				return err
 			}
 		}
-		return gui.closeConfirmationPrompt(g)
+		return gui.closeConfirmationPrompt()
 	}
 }
 
-func (gui *Gui) closeConfirmationPrompt(g *gocui.Gui) error {
-	view, err := g.View("confirmation")
+func (gui *Gui) closeConfirmationPrompt() error {
+	view, err := gui.g.View("confirmation")
 	if err != nil {
 		return nil // if it's already been closed we can just return
 	}
-	if err := gui.returnFocus(g, view); err != nil {
-		panic(err)
+	if err := gui.returnFocus(gui.g, view); err != nil {
+		return err
 	}
-	g.DeleteViewKeybindings("confirmation")
-	return g.DeleteView("confirmation")
+	gui.g.DeleteViewKeybindings("confirmation")
+	return gui.g.DeleteView("confirmation")
 }
 
 func (gui *Gui) getMessageHeight(wrap bool, message string, width int) int {
@@ -110,7 +110,7 @@ func (gui *Gui) createPopupPanel(g *gocui.Gui, currentView *gocui.View, title, p
 	g.Update(func(g *gocui.Gui) error {
 		// delete the existing confirmation panel if it exists
 		if view, _ := g.View("confirmation"); view != nil {
-			if err := gui.closeConfirmationPrompt(g); err != nil {
+			if err := gui.closeConfirmationPrompt(); err != nil {
 				gui.Log.Error(err.Error())
 			}
 		}
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 16ebee38..0d6fe6a9 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -302,8 +302,38 @@ func (gui *Gui) refresh() {
 }
 
 func (gui *Gui) listenForEvents(finish chan struct{}, refresh func()) {
+	errorCount := 0
+
+	onError := func(err error) {
+		if err != nil {
+			gui.ErrorChan <- errors.Errorf("Docker event stream returned error: %s\nRetry count: %d", err.Error(), errorCount)
+		}
+		errorCount++
+		time.Sleep(time.Second * 2)
+	}
+
+outer:
 	for {
 		messageChan, errChan := gui.DockerCommand.Client.Events(context.Background(), types.EventsOptions{})
+
+		if errorCount > 0 {
+			select {
+			case err := <-errChan:
+				onError(err)
+				continue outer
+			default:
+				// If we're here then we lost connection to docker and we just got it back.
+				// The reason we do this refresh explicitly is because successfully
+				// reconnecting with docker does not mean it's going to send us a new
+				// event any time soon.
+
+				// Assuming the confirmation prompt currently holds the given error
+				_ = gui.closeConfirmationPrompt()
+				refresh()
+				errorCount = 0
+			}
+		}
+
 		for {
 			select {
 			case <-finish:
@@ -316,8 +346,8 @@ func (gui *Gui) listenForEvents(finish chan struct{}, refresh func()) {
 
 				gui.Log.Infof("received event of type: %s", message.Type)
 			case err := <-errChan:
-				gui.ErrorChan <- errors.Errorf("Docker event stream returned error: %s", err.Error())
-				break
+				onError(err)
+				continue outer
 			}
 		}
 	}