From e39ac23f694548bae6a23c5edfb49c0769515953 Mon Sep 17 00:00:00 2001
From: Rob Watson <rob@netflux.io>
Date: Sun, 13 Apr 2025 10:04:53 +0200
Subject: [PATCH] fixup! refactor(container): restart handling

---
 internal/container/container.go      | 35 +++++++++++++++++-----------
 internal/container/container_test.go |  3 ---
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/internal/container/container.go b/internal/container/container.go
index 4ae9f92..778355c 100644
--- a/internal/container/container.go
+++ b/internal/container/container.go
@@ -377,23 +377,30 @@ func (a *Client) runContainerLoop(
 			respC, errC := a.apiClient.ContainerWait(ctx, containerID, container.WaitConditionNextExit)
 			select {
 			case resp := <-respC:
-				if shouldRestartFunc != nil {
-					shouldRestart, err := shouldRestartFunc(resp.StatusCode, restartCount, time.Since(startedWaitingAt))
-					if shouldRestart && err != nil {
-						panic(fmt.Errorf("shouldRestart must return nil error if restarting, but returned: %w", err))
-					}
-					if !shouldRestart {
-						a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "false", "exit_code", resp.StatusCode, "restart_count", restartCount)
-						containerRespC <- containerWaitResponse{
-							WaitResponse: resp,
-							restarting:   false,
-							restartCount: restartCount,
-							err:          err,
-						}
-						return
+				exit := func(err error) {
+					a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "false", "exit_code", resp.StatusCode, "restart_count", restartCount)
+					containerRespC <- containerWaitResponse{
+						WaitResponse: resp,
+						restarting:   false,
+						restartCount: restartCount,
+						err:          err,
 					}
 				}
 
+				if shouldRestartFunc == nil {
+					exit(nil)
+					return
+				}
+
+				shouldRestart, err := shouldRestartFunc(resp.StatusCode, restartCount, time.Since(startedWaitingAt))
+				if shouldRestart && err != nil {
+					panic(fmt.Errorf("shouldRestart must return nil error if restarting, but returned: %w", err))
+				}
+				if !shouldRestart {
+					exit(err)
+					return
+				}
+
 				a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "true", "exit_code", resp.StatusCode, "restart_count", restartCount)
 				timer.Reset(restartInterval)
 
diff --git a/internal/container/container_test.go b/internal/container/container_test.go
index 108dc9a..3e54622 100644
--- a/internal/container/container_test.go
+++ b/internal/container/container_test.go
@@ -94,9 +94,6 @@ func TestClientRunContainer(t *testing.T) {
 				Mode:    0755,
 			},
 		},
-		ShouldRestart: func(int64, int, time.Duration) (bool, error) {
-			return false, nil
-		},
 	})
 
 	done := make(chan struct{})